1
+ from dataclasses import dataclass
2
+ import typing
1
3
import gspread
2
4
import gspread_formatting
3
5
from enum import Enum
4
- from googleapiclient .discovery import build
5
- import numpy as np
6
+ import pandas as pd
6
7
7
8
FONT_SIZE_PTS = 10
8
9
PTS_PIXELS_RATIO = 4 / 3
@@ -16,20 +17,33 @@ class FILE_OVERRIDE_BEHAVIORS(Enum):
16
17
EXIT_IF_IN_SAME_PLACE = 2
17
18
EXIT_ANYWHERE = 3
18
19
20
+
19
21
class WORKSHEET_OVERRIDE_BEHAVIORS (Enum ):
20
22
OVERRIDE = 1
21
23
EXIT = 2
22
24
25
+
23
26
class COLUMN_FORMAT_OPTIONS (Enum ):
24
27
DEFAULT = 1
25
28
PERCENT_UNCOLORED = 2
26
29
PERCENT_COLORED = 3
30
+ YEAR_MONTH_DATE = 4
31
+
32
+
33
+ class CHART_TYPES (Enum ):
34
+ LINE = "LINE"
27
35
28
36
DEFAULT_SHEET_FORMATTING_OPTIONS = {
29
37
"bold_header" : True ,
30
38
"center_header" : True ,
31
39
"freeze_header" : True ,
32
- "column_widths" : {"justify" : True , "buffer_chars" : DEFAULT_BUFFER_CHARS }
40
+ "column_widths" : {"justify" : True , "buffer_chars" : DEFAULT_BUFFER_CHARS },
41
+ "extra_columns" : 0 ,
42
+ "extra_columns_width" : 50 ,
43
+ }
44
+
45
+ DEFAULT_GSPREAD_UPDATE_ARGS = {
46
+ "value_input_option" : gspread .utils .ValueInputOption .user_entered ,
33
47
}
34
48
35
49
def extract_credentials (authentication_response ):
@@ -41,7 +55,7 @@ def authenticate_gspread(authentication_response):
41
55
gc = gspread .authorize (extract_credentials (authentication_response ))
42
56
return gc
43
57
44
- def authenticate_drive_api (authentication_response ):
58
+ def authenticate_google_api (authentication_response ):
45
59
"""Authenticates the Drive API using the response from api.authenticate"""
46
60
return authentication_response [0 ]
47
61
@@ -107,21 +121,21 @@ def search_for_folder_id(drive_api, folder_name, allow_trashed = False, allow_du
107
121
return [file ["id" ] for file in files_exact_match ]
108
122
109
123
110
- def create_sheet_in_folder (authentication_response , sheet_name , parent_folder_name = None , override_behavior = FILE_OVERRIDE_BEHAVIORS .EXIT_ANYWHERE ):
124
+ def create_sheet_in_folder (drive_authentication_response , sheet_name , parent_folder_name = None , override_behavior = FILE_OVERRIDE_BEHAVIORS .EXIT_ANYWHERE ):
111
125
"""
112
126
Create a new sheet in the project with the given name and parent folder.
113
127
Returns the new sheet.
114
128
115
- :param authentication_response : the service parameters tuple
129
+ :param drive_authentication_response : the service parameters tuple
116
130
:param sheet_name: the name of the new sheet
117
131
:param parent_folder_name: the name of the parent folder for the new sheet
118
132
:param override_behavior: the behavior to take if the sheet already exists
119
133
:returns: the gspread.Spreadsheet object of the new sheet
120
134
:rtype: gspread.Spreadsheet
121
135
"""
122
136
# Build Drive API
123
- gc = authenticate_gspread (authentication_response )
124
- drive_api = authenticate_drive_api ( authentication_response )
137
+ gc = authenticate_gspread (drive_authentication_response )
138
+ drive_api = authenticate_google_api ( drive_authentication_response )
125
139
parent_folder_id = None if parent_folder_name is None else search_for_folder_id (drive_api , parent_folder_name )[0 ]
126
140
127
141
# Check if sheet already exists and handle based on input
@@ -163,8 +177,9 @@ def fill_worksheet_with_df(
163
177
df ,
164
178
worksheet_name ,
165
179
overlapBehavior ,
166
- sheet_formatting_options = DEFAULT_SHEET_FORMATTING_OPTIONS ,
167
- column_formatting_options = {}
180
+ sheet_formatting_options = {},
181
+ column_formatting_options = {},
182
+ ** gspread_update_args
168
183
):
169
184
"""
170
185
Fill a worksheet with the contents of a DataFrame.
@@ -193,38 +208,50 @@ def fill_worksheet_with_df(
193
208
title = worksheet_name , rows = df .shape [0 ], cols = df .shape [1 ]
194
209
)
195
210
211
+ sheet_formatting_options_filled = {** DEFAULT_SHEET_FORMATTING_OPTIONS , ** sheet_formatting_options }
212
+
213
+ # Add extra blank columns to the right of the worksheet
214
+ df_to_insert = pd .concat (
215
+ [df ] + [pd .Series (" " , index = df .index , name = "" )] * sheet_formatting_options_filled ["extra_columns" ],
216
+ axis = 1
217
+ )
196
218
# Add data to worksheet
197
- worksheet .update ([df .columns .values .tolist ()] + df .fillna ("NA" ).values .tolist ())
219
+ worksheet .update (
220
+ [df_to_insert .columns .values .tolist ()] + df_to_insert .fillna ("NA" ).values .tolist (),
221
+ ** {** DEFAULT_GSPREAD_UPDATE_ARGS , ** gspread_update_args }
222
+ )
198
223
199
224
# Format worksheet
200
225
# Justify Column Widths
201
- if "column_widths" not in sheet_formatting_options or sheet_formatting_options ["column_widths" ]["justify" ]:
226
+ if "column_widths" not in sheet_formatting_options_filled or sheet_formatting_options_filled ["column_widths" ]["justify" ]:
202
227
text_widths = df .astype (str ).columns .map (
203
228
lambda column_name : df [column_name ].astype (str ).str .len ().max ()
204
229
)
205
230
header_widths = df .columns .str .len ()
206
231
buffer_chars = (
207
232
DEFAULT_BUFFER_CHARS
208
- if ("column_widths" not in sheet_formatting_options or "buffer_chars" not in sheet_formatting_options ["column_widths" ])
209
- else sheet_formatting_options ["column_widths" ]["buffer_chars" ]
233
+ if ("column_widths" not in sheet_formatting_options_filled or "buffer_chars" not in sheet_formatting_options_filled ["column_widths" ])
234
+ else sheet_formatting_options_filled ["column_widths" ]["buffer_chars" ]
210
235
)
211
- column_widths = [
236
+ data_column_widths = [
212
237
round ((max (len_tuple ) + buffer_chars ) * FONT_SIZE_PTS * 1 / PTS_PIXELS_RATIO )
213
238
for len_tuple in zip (text_widths , header_widths )
214
239
]
240
+ extra_column_widths = [sheet_formatting_options_filled ["extra_columns_width" ]] * sheet_formatting_options_filled ["extra_columns" ]
241
+ combined_column_widths = data_column_widths + extra_column_widths
215
242
column_positions = [
216
- gspread .utils .rowcol_to_a1 (1 , i + 1 )[0 ] for i , _ in enumerate (column_widths )
243
+ gspread .utils .rowcol_to_a1 (1 , i + 1 )[0 ] for i , _ in enumerate (combined_column_widths )
217
244
]
218
- gspread_formatting .set_column_widths (worksheet , zip (column_positions , column_widths ))
245
+ gspread_formatting .set_column_widths (worksheet , zip (column_positions , combined_column_widths ))
219
246
# Freeze Header
220
- if "freeze_header" not in sheet_formatting_options or sheet_formatting_options ["freeze_header" ]:
247
+ if "freeze_header" not in sheet_formatting_options_filled or sheet_formatting_options_filled ["freeze_header" ]:
221
248
gspread_formatting .set_frozen (worksheet , rows = 1 )
222
249
base_format_options = gspread_formatting .CellFormat ()
223
250
# Bold Header
224
- if "bold_header" not in sheet_formatting_options or sheet_formatting_options ["bold_header" ]:
251
+ if "bold_header" not in sheet_formatting_options_filled or sheet_formatting_options_filled ["bold_header" ]:
225
252
base_format_options += gspread_formatting .CellFormat (textFormat = gspread_formatting .TextFormat (bold = True ))
226
253
# Center Header
227
- if "center_header" not in sheet_formatting_options or sheet_formatting_options ["center_header" ]:
254
+ if "center_header" not in sheet_formatting_options_filled or sheet_formatting_options_filled ["center_header" ]:
228
255
base_format_options += gspread_formatting .CellFormat (horizontalAlignment = "CENTER" )
229
256
# Handle column specific formatting
230
257
for column in column_formatting_options :
@@ -269,6 +296,13 @@ def fill_worksheet_with_df(
269
296
column_range ,
270
297
gspread_formatting .CellFormat (numberFormat = gspread_formatting .NumberFormat (type = 'PERCENT' , pattern = '0.0%' ))
271
298
)
299
+ if column_formatting_options [column ] == COLUMN_FORMAT_OPTIONS .YEAR_MONTH_DATE :
300
+ # Apply date format rule
301
+ gspread_formatting .format_cell_range (
302
+ worksheet ,
303
+ column_range ,
304
+ gspread_formatting .CellFormat (numberFormat = gspread_formatting .NumberFormat (type = 'DATE' , pattern = 'yyyy-mm' ))
305
+ )
272
306
273
307
# Apply base formatting options
274
308
gspread_formatting .format_cell_range (
@@ -281,7 +315,7 @@ def fill_worksheet_with_df(
281
315
if "Sheet1" in [i .title for i in sheet .worksheets ()]:
282
316
sheet .del_worksheet (sheet .worksheet ("Sheet1" ))
283
317
284
- def fill_spreadsheet_with_df_dict (sheet , df_dict , overlapBehavior , sheet_formatting_options = {}, column_formatting_options = {}):
318
+ def fill_spreadsheet_with_df_dict (sheet , df_dict , overlapBehavior , sheet_formatting_options = {}, column_formatting_options = {}, ** gspread_update_args ):
285
319
"""
286
320
Fill a sheet with the contents of a dictionary of DataFrames.
287
321
The keys of the dictionary are the names of the worksheets, and the values contain the data to be placed in the sheet.
@@ -307,6 +341,150 @@ def fill_spreadsheet_with_df_dict(sheet, df_dict, overlapBehavior, sheet_formatt
307
341
for worksheet_name , df in df_dict .items ():
308
342
fill_worksheet_with_df (
309
343
sheet , df , worksheet_name , overlapBehavior ,
310
- sheet_formatting_options = sheet_formatting_options .get (worksheet_name , DEFAULT_SHEET_FORMATTING_OPTIONS ),
311
- column_formatting_options = column_formatting_options .get (worksheet_name , {})
312
- )
344
+ sheet_formatting_options = sheet_formatting_options .get (worksheet_name , {}),
345
+ column_formatting_options = column_formatting_options .get (worksheet_name , {}),
346
+ ** gspread_update_args
347
+ )
348
+
349
+ def update_sheet_raw (sheets_authentication_response , sheet , * updates ):
350
+ """
351
+ Directly call the Google Sheets api to update the specified sheet with the optional arguments.
352
+ """
353
+ assert len (updates ) > 0
354
+ sheets_api = authenticate_google_api (sheets_authentication_response )
355
+ sheet_id = sheet .id
356
+ body = {"requests" : list (updates )}
357
+ response = (
358
+ sheets_api .spreadsheets ()
359
+ .batchUpdate (spreadsheetId = sheet_id , body = body )
360
+ .execute ()
361
+ )
362
+ return response
363
+
364
+ REQUIRED_CHART_ARGS = []
365
+
366
+ DEFAULT_CHART_ARGS = {
367
+ "title" : "" ,
368
+ "x_axis_title" : "" ,
369
+ "y_axis_title" : "" ,
370
+ "invert_x_axis" : False ,
371
+ "chart_position" : None , # None means it will be created in a new sheet
372
+ "chart_position_offset_x" : 0 ,
373
+ "chart_position_offset_y" : 0 ,
374
+ "chart_width" : 600 ,
375
+ "chart_height" : 371 ,
376
+ }
377
+
378
+ @dataclass
379
+ class WorksheetRange :
380
+ """
381
+ A dataclass to represent a range of cells in a worksheet in the one-sided interval [top_left, bottom_right).
382
+ :param worksheet: the gspread.worksheet.Worksheet object
383
+ :param top_left: the top left cell of the range. This cell will be included in the range
384
+ :param bottom_right: the bottom right cell of the range. This cell will not be included in the range
385
+ """
386
+ worksheet : gspread .worksheet .Worksheet
387
+ top_left : gspread .cell .Cell
388
+ bottom_right : gspread .cell .Cell
389
+
390
+ @property
391
+ def range_dict (self ):
392
+ """The range as a dictionary for the sources field in the Google Sheets api"""
393
+ return {
394
+ "sheetId" : self .worksheet .id ,
395
+ "startRowIndex" : self .top_left .row - 1 ,
396
+ "endRowIndex" : self .bottom_right .row - 1 ,
397
+ "startColumnIndex" : self .top_left .col - 1 ,
398
+ "endColumnIndex" : self .bottom_right .col - 1 ,
399
+ }
400
+
401
+ def _cell_to_grid_coordinate (cell , worksheet ):
402
+ return {
403
+ "sheetId" : worksheet .id ,
404
+ "rowIndex" : cell .row - 1 ,
405
+ "columnIndex" : cell .col - 1 ,
406
+ }
407
+
408
+ def add_chart_to_sheet (sheets_authentication_response , sheet , worksheet , chart_type , domain , series , ** chart_args ):
409
+ """
410
+ Add a chart to a specified workshet
411
+ :param sheets_authentication_response: the response from ga.authenticate. Must be for the sheets api v4
412
+ :param sheet: the gspread.Spreadsheet object
413
+ :param worksheet: the gspread.Worksheet object
414
+ :param chart_type: the type of chart to add
415
+ :param domain: the domain of the chart as a WorksheetRange. Must contain either one row or one column
416
+ :param series: the series of the chart as a WorksheetRange. Must contain either one row or one column
417
+ :param chart_args: other arguments to create the chart. See DEFAULT_CHART_ARGS
418
+ """
419
+ complete_chart_args = {** DEFAULT_CHART_ARGS , ** chart_args }
420
+ if complete_chart_args ["chart_position" ] is not None :
421
+ position_dict = {
422
+ "overlayPosition" : {
423
+ "anchorCell" : _cell_to_grid_coordinate (complete_chart_args ["chart_position" ], worksheet ),
424
+ "offsetXPixels" : complete_chart_args ["chart_position_offset_x" ],
425
+ "offsetYPixels" : complete_chart_args ["chart_position_offset_y" ],
426
+ "widthPixels" : complete_chart_args ["chart_width" ],
427
+ "heightPixels" : complete_chart_args ["chart_height" ],
428
+ }
429
+ }
430
+ else :
431
+ position_dict = {"newSheet" : True }
432
+ formatted_domains = [
433
+ {
434
+ "domain" : {
435
+ "sourceRange" : {
436
+ "sources" : [
437
+ domain .range_dict
438
+ ],
439
+ },
440
+ },
441
+ "reversed" : complete_chart_args ["invert_x_axis" ],
442
+ },
443
+ ]
444
+
445
+ formatted_series = [
446
+ {
447
+ "series" : {
448
+ "sourceRange" : {
449
+ "sources" : [
450
+ series_source .range_dict
451
+ ],
452
+ },
453
+ },
454
+ "targetAxis" : "LEFT_AXIS" ,
455
+ }
456
+ for series_source in series
457
+ ]
458
+ formatted_axis = []
459
+ if complete_chart_args ["x_axis_title" ]:
460
+ formatted_axis .append ({
461
+ "title" : complete_chart_args ["x_axis_title" ],
462
+ "position" : "BOTTOM_AXIS" ,
463
+ })
464
+ if complete_chart_args ["y_axis_title" ]:
465
+ formatted_axis .append ({
466
+ "title" : complete_chart_args ["y_axis_title" ],
467
+ "position" : "LEFT_AXIS" ,
468
+ })
469
+ request = {
470
+ "addChart" : {
471
+ "chart" : {
472
+ "spec" : {
473
+ "title" : complete_chart_args ["title" ],
474
+ #TODO: insert legend position
475
+ #TODO: insert axis positions
476
+ "basicChart" : {
477
+ "axis" : formatted_axis ,
478
+ "chartType" : chart_type .value ,
479
+ "domains" : formatted_domains ,
480
+ "headerCount" : 1 , #TODO: not sure what this means
481
+ "series" : formatted_series ,
482
+ },
483
+ },
484
+ "position" : position_dict
485
+ },
486
+ },
487
+ }
488
+
489
+ response = update_sheet_raw (sheets_authentication_response , sheet , request )
490
+ return response
0 commit comments