add body detector, mod pull sheet

1 files changed, 18 insertions, 7 deletions

diff --git a/megapixels/commands/datasets/pull_spreadsheet.py b/megapixels/commands/datasets/pull_spreadsheet.py
index b8b68094..caf5eb43 100644
--- a/megapixels/commands/datasets/pull_spreadsheet.py
+++ b/megapixels/commands/datasets/pull_spreadsheet.py
@@ -21,6 +21,10 @@ from app.utils.logger_utils import Logger
 log = Logger.getLogger()
 
 opt_sheets = ['datasets', 'relationships', 'funding', 'references', 'sources', 'tags', 'citations', 'legal']
+dataset_sheet_keys = ['key', 'name_short', 'name_full', 'url', 'dl_im', 'purpose', 'funded_by',
+  'year_start', 'year_end', 'year_published', 'images', 'videos', 'identities', 
+  'faces_or_persons', 'campus', 'youtube', 'flickr', 'google', 'bing', 'comment']
+
 
 @click.command()
 @click.option('-n', '--name', 'opt_spreadsheets', multiple=True,
@@ -30,11 +34,15 @@ opt_sheets = ['datasets', 'relationships', 'funding', 'references', 'sources', '
 @click.option('--all', 'opt_all', is_flag=True,
   help='Get all sheets')
 @click.option('-o', '--output', 'opt_fp_out', required=True,
+  type=click.Path(file_okay=False, dir_okay=True),
   help='Path to directory or filename')
+@click.option('--share', 'opt_share', required=True,
+  type=click.Choice(['nyt', 'ft']),
+  help='Share filter')
 @click.option('-f', '--force', 'opt_force', is_flag=True,
   help='Force overwrite')
 @click.pass_context
-def cli(ctx, opt_spreadsheets, opt_fp_out, opt_all, opt_force):
+def cli(ctx, opt_spreadsheets, opt_fp_out, opt_all, opt_share, opt_force):
   """Fetch Google spreadsheet"""
   
   import sys
@@ -47,6 +55,12 @@ def cli(ctx, opt_spreadsheets, opt_fp_out, opt_all, opt_force):
   
   for sheet_name in opt_spreadsheets:
     log.info(f'Get spreadsheet: {sheet_name}')
+    fp_out = join(opt_fp_out, f'{sheet_name}.csv')
+    fpp_out = Path(fp_out)
+    if fpp_out.exists() and not opt_force:
+      log.error(f'File "{fpp_out} exists. Use "-f" to overwrite')
+      return
+
     sheet_data = fetch_google_sheet_objects(name=sheet_name)
     df_sheet = pd.DataFrame.from_dict(sheet_data)
     if sheet_name == 'datasets':
@@ -58,22 +72,19 @@ def cli(ctx, opt_spreadsheets, opt_fp_out, opt_all, opt_force):
       fpp_out = fpp_out.parent
     else:
       fpp_out = join(opt_fp_out, f'{sheet_name}.csv')
+    log.info(f'Writing file: {fpp_out}')
     df_sheet.to_csv(fpp_out)
 
 
 def clean_datasets_sheet_ft(df):
   # clean data for FT
   df = df[df['ft_share'] == 'Y']
-  keys = ['key', 'name_short', 'name_full', 'url', 'downloaded', 'purpose', 'wild']
-  keys += ['campus', 'year_start', 'year_end', 'year_published', 'images', 'videos', 'identities', 'faces_or_persons', 'youtube', 'flickr', 'google', 'bing', 'comment']
-  return df[keys]
+  return df[dataset_sheet_keys]
 
 def clean_datasets_sheet_nyt(df):
   # clean data for FT
   df = df[df['ft_share'] == 'Y']
-  keys = ['key', 'name_short', 'name_full', 'url', 'downloaded', 'purpose', 'wild']
-  keys += ['campus', 'year_start', 'year_end', 'year_published', 'images', 'videos', 'identities', 'faces_or_persons', 'youtube', 'flickr', 'google', 'bing', 'comment']
-  return df[keys]
+  return df[dataset_sheet_keys]
 
 def fetch_spreadsheet():
   """Open the Google Spreadsheet, which contains the individual worksheets"""