Do some simple Python style cleanup

ocdskingfishercolab/__init__.py

@@ -7,7 +7,7 @@
 
    from ocdskingfishercolab import (create_connection, authenticate_gspread, getResults, saveToCSV, saveToSheets,
                                     saveStraightToSheets, downloadReleases, output_notebook, set_spreadsheet_name,
-                                    authenticate_pydrive, downloadReleases)
+                                    authenticate_pydrive)
 """
 import json
 
@@ -29,16 +29,10 @@ def create_connection(database, user, password, host, port='5432'):
     Creates a connection to the database.
     """
     global conn
-    if conn and conn.closed > 0:
+    if conn and conn.closed:
         reset_connection()
     if not conn:
-        conn = psycopg2.connect(
-            database=database,
-            user=user,
-            password=password,
-            host=host,
-            port=port,
-        )
+        conn = psycopg2.connect(database=database, user=user, password=password, host=host, port=port)
     return conn
 
 
@@ -49,7 +43,7 @@ def reset_connection():
     This does not re-open the connection again.
     """
     global conn
-    if conn is not None and conn.closed != 0:
+    if conn and not conn.closed:
         try:
             conn.cancel()
             conn.reset()
@@ -64,8 +58,7 @@ def authenticate_gspread():
     Authenticates the current user and gives the notebook permission to connect to Google Spreadsheets.
     """
     auth.authenticate_user()
-    gc = gspread.authorize(GoogleCredentials.get_application_default())
-    return gc
+    return gspread.authorize(GoogleCredentials.get_application_default())
 
 
 def authenticate_pydrive():
@@ -75,17 +68,15 @@ def authenticate_pydrive():
     auth.authenticate_user()
     gauth = GoogleAuth()
     gauth.credentials = GoogleCredentials.get_application_default()
-    drive = GoogleDrive(gauth)
-    return drive
+    return GoogleDrive(gauth)
 
 
 def getResults(cur):
     """
     Takes in a database cursor and returns a Pandas DataFrame.
     """
     headers = [desc[0] for desc in cur.description]
-    results = pandas.DataFrame(cur.fetchall(), columns=headers)
-    return results
+    return pandas.DataFrame(cur.fetchall(), columns=headers)
 
 
 def saveToCSV(dataframe, filename):
@@ -104,7 +95,6 @@ def set_spreadsheet_name(name):
     spreadsheet_name = name
 
 
-# option to bypass confirmation in save to sheets
 def saveStraightToSheets(dataframe, sheetname):
     """
     Saves a DataFrame straight to a Google Spreadsheet.
@@ -119,72 +109,66 @@ def saveStraightToSheets(dataframe, sheetname):
     try:
         worksheet = gSheet.add_worksheet(sheetname, dataframe.shape[0], dataframe.shape[1])
     except Exception:
-        newsheetname = input(sheetname + " already exists, enter a different name: ")
+        newsheetname = input(sheetname + ' already exists, enter a different name: ')
         worksheet = gSheet.add_worksheet(newsheetname, dataframe.shape[0], dataframe.shape[1])
 
-    # save dataframe to worksheet
     set_with_dataframe(worksheet, dataframe)
 
 
-# saves dataframe to sheets after user confirms yes
 def saveToSheets(dataframe, sheetname):
     """
     The same as saveStraightToSheets, except the user is prompted if they are sure first.
     """
-    if input("Save to Google Sheets? (y/n)") == 'y':
+    if input('Save to Google Sheets? (y/n)') == 'y':
         saveStraightToSheets(dataframe, sheetname)
 
 
 def downloadReleases(collection_id, ocid, package_type):
     """
     Downloads some releases into a file in the notebook.
     """
+    if package_type not in ('record', 'release'):
+        print("package_type parameter must be either 'release' or 'record'")
+        return
+
+    querystring = """
+    with releases as (
+        select
+            ocid,
+            data
+        from
+            data
+        join
+            release_with_collection on data.id = release_with_collection.data_id
+        where collection_id = %(collection_id)s
+    )
+    SELECT
+      jsonb_build_object('releases', jsonb_agg(data)),
+      jsonb_build_object(
+          'ocid', %(ocid)s,
+          'records', jsonb_build_array(jsonb_build_object('releases', jsonb_agg(data)))
+      )
+    FROM
+      releases
+    WHERE
+      ocid = %(ocid)s
+    """
+
     with conn, conn.cursor() as cur:
-        if package_type != 'release' and package_type != 'record':
-            print("package_type parameter must be either 'release' or 'record'")
-        else:
-
-            querystring = """
-
-          with releases as (
-
-              select
-                  ocid,
-                  data
-              from
-                  data
-              join
-                  release_with_collection on data.id = release_with_collection.data_id
-              where collection_id = %(collection_id)s
-
-          )
-
-          SELECT
-            jsonb_build_object('releases', jsonb_agg(data)),
-            jsonb_build_object(
-                'ocid', %(ocid)s,
-                'records', jsonb_build_array(jsonb_build_object('releases', jsonb_agg(data)))
-            )
-          FROM
-            releases
-          WHERE
-            ocid = %(ocid)s
-
-          """
-
-            cur.execute(querystring,
-                        {"ocid": ocid, "collection_id": collection_id}
-                        )
-
-            file = ocid + '_' + package_type + '_package.json'
-
-            with open(file, 'w') as f:
-                if package_type == 'release':
-                    json.dump(cur.fetchone()[0], f, indent=2, ensure_ascii=False)
-                elif package_type == 'record':
-                    json.dump(cur.fetchone()[1], f, indent=2, ensure_ascii=False)
-
-            files.download(file)
+        cur.execute(querystring, {'ocid': ocid, 'collection_id': collection_id})
+
+        file = ocid + '_' + package_type + '_package.json'
+
+        result = cur.fetchone()
+        if package_type == 'release':
+            index = 0
+        elif package_type == 'record':
+            index = 1
+
+        with open(file, 'w') as f:
+            json.dump(result[index], f, indent=2, ensure_ascii=False)
+
+        files.download(file)
 
 
 def output_notebook(sql, params=None):