1# Copyright 2017 The Chromium Authors. All rights reserved.
2# Use of this source code is governed by a BSD-style license that can be
3# found in the LICENSE file.
4
5"""Functions to use the drive APIs to manipulate spreadsheets and folders."""
6
7# Pylint warnings occur if you install a different version than in third_party
8# pylint: disable=no-member
9
10import time
11
12import httplib2
13
14from apiclient import discovery  # pylint: disable=import-error
15from oauth2client import service_account  # pylint: disable=no-name-in-module
16
17# Update this to the location you downloaded the keyfile to.
18# See https://developers.google.com/sheets/api/quickstart/python
19_PATH_TO_JSON_KEYFILE = 'PATH_TO/keyfile.json'
20
21
22def ReadSpreadsheet(spreadsheet_id, range_name):
23  """Returns the values in the given range.
24
25  Args:
26    spreadsheet_id: The spreadsheet id, which can be found in the url.
27    range_name: A string range, such as 'A2:D9'
28
29  Returns: A 2D list of values in the given range.
30  """
31  sheets_service = _GetSheetsService()
32  result = sheets_service.spreadsheets().values().get(
33      spreadsheetId=spreadsheet_id, range=range_name).execute()
34  return result.get('values', [])
35
36
37def CreateSpreadsheet(title, sheets, folder_id):
38  """Creates the given spreadsheet in the given folder.
39
40  Args:
41    title: The title of the file to be created.
42    sheets: A list of dicts containing:
43      'name': The name of the tab for the sheet
44      'values': A 2D list of data to fill in
45    folder_id: The id of the folder to put the spreadsheet in, from CreateFolder
46
47  Returns: The url of the created spreadsheet.
48  """
49  file_info = CreateEmptyFile(
50      title, 'application/vnd.google-apps.spreadsheet', folder_id)
51  spreadsheet_id = file_info['id']
52  spreadsheet_url = _AddSheetsToSpreadsheet(sheets, spreadsheet_id)
53  print 'Created Spreadsheet %s' % spreadsheet_url
54  # Limit of 100 actions in 100 seconds; sleep 1s between sheets.
55  time.sleep(10)
56  return spreadsheet_url
57
58def _AddSheetsToSpreadsheet(sheets, spreadsheet_id):
59  # Batch and execute requests to add empty sheets with the correct names.
60  requests = []
61  for sheet in sheets:
62    requests.append({
63        'addSheet': {'properties': {'title': sheet['name']}}
64    })
65  sheets_service = _GetSheetsService()
66  sheets_service.spreadsheets().batchUpdate(
67      spreadsheetId=spreadsheet_id, body={'requests': requests}).execute()
68  sheet_metadata = sheets_service.spreadsheets().get(
69      spreadsheetId=spreadsheet_id).execute()
70
71  sheet_ids = dict((
72      s['properties']['title'],
73      s['properties']['sheetId']) for s in sheet_metadata['sheets'])
74
75  # Delete placeholder 'Sheet1'
76  requests = [{
77      'deleteSheet': {
78          'sheetId': sheet_ids['Sheet1']
79      }
80  }]
81  sheets_service.spreadsheets().batchUpdate(
82      spreadsheetId=spreadsheet_id, body={'requests': requests}).execute()
83
84  # Add the correct data to the new sheets.
85  for sheet in sheets:
86    # Update values
87    r = '%s!A1' % sheet['name']
88    sheets_service.spreadsheets().values().update(
89        spreadsheetId=spreadsheet_id,
90        body={'values': sheet['values']},
91        range=r,
92        valueInputOption='USER_ENTERED').execute()
93
94    # Batch and execute requests to set the formatting on the new sheets
95    # (one frozen bolded row at the top for a header).
96    requests = [{
97        'repeatCell': {
98            'range': {
99                'sheetId': sheet_ids[sheet['name']],
100                'startRowIndex': 0,
101                'endRowIndex': 1
102            },
103            'cell': {
104                'userEnteredFormat': {
105                    'textFormat': {
106                        'bold': True
107                    }
108                }
109            },
110            'fields': 'userEnteredFormat(textFormat)'
111        }
112    }, {
113        'updateSheetProperties': {
114            'properties': {
115                'sheetId': sheet_ids[sheet['name']],
116                'gridProperties': {
117                    'frozenRowCount': 1
118                }
119            },
120            'fields': 'gridProperties.frozenRowCount'
121        }
122    }]
123    sheets_service.spreadsheets().batchUpdate(
124        spreadsheetId=spreadsheet_id, body={'requests': requests}).execute()
125  return sheet_metadata['spreadsheetUrl']
126
127
128def CreateEmptyFile(name, mime_type, folder_id):
129  """Creates an empty drive file and returns the result."""
130  drive_service = _GetDriveService()
131  file_metadata = {
132      'name': name,
133      'mimeType': mime_type,
134      'parents': [folder_id],
135  }
136  print 'Creating file: %s' % file_metadata
137  return drive_service.files().create(body=file_metadata).execute()
138
139
140def CreateFolder(title):
141  """Creates a drive folder with the given title."""
142  drive_service = _GetDriveService()
143  folder_metadata = {
144      'name': title,
145      'mimeType': 'application/vnd.google-apps.folder'
146  }
147  folder = drive_service.files().create(
148      body=folder_metadata, fields='id').execute()
149  folder_id = folder.get('id')
150  # TODO(sullivan): What should permissions look like?
151  user_permission = {
152      'type': 'user',
153      'role': 'writer',
154      'emailAddress': 'sullivan@google.com'
155  }
156  drive_service.permissions().create(
157      fileId=folder_id, body=user_permission, fields='id').execute()
158  return folder_id
159
160def _GetSheetsService():
161  scopes = ['https://www.googleapis.com/auth/spreadsheets']
162  return _BuildDiscoveryService(
163      scopes, 'sheets', 'v4',
164      'https://sheets.googleapis.com/$discovery/rest?version=v4')
165
166def _GetDriveService():
167  scopes = ['https://www.googleapis.com/auth/drive']
168  return _BuildDiscoveryService(scopes, 'drive', 'v3', None)
169
170
171def _BuildDiscoveryService(scopes, api, version, url):
172  """Helper functions to build discovery services for sheets and drive."""
173  creds = service_account.ServiceAccountCredentials.from_json_keyfile_name(
174      _PATH_TO_JSON_KEYFILE, scopes)
175  http = creds.authorize(httplib2.Http())
176  if url:
177    return discovery.build(api, version, http=http, discoveryServiceUrl=url)
178  return discovery.build(api, version, http=http)
179