Added header arg to the download method

iakov-kaiumov · Oct 3, 2023 · f693247 · f693247
1 parent 5ec0ead
commit f693247
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 6 deletions.
diff --git a/gsheet_pandas/adapter/connection.py b/gsheet_pandas/adapter/connection.py
@@ -107,16 +107,26 @@ def get_all_files_in_folder(self, folder_id):
             raise e
         return files
 
-    def download(self, drive_table: str, sheet_name: str, range_name: str = DEFAULT_RANGE_NAME) -> pd.DataFrame:
+    def download(self, drive_table: str, sheet_name: str, range_name: str = DEFAULT_RANGE_NAME,
+                 header: int | None = 0) -> pd.DataFrame:
         service = self._get_service()
         sheet = self._get_service().spreadsheets()
         result = sheet.values().get(spreadsheetId=drive_table, range=sheet_name + range_name).execute()
         values = result.get('values', [])
         service.close()
         if not values:
             raise Exception('Empty data')
-        df = pd.DataFrame(values[1:])
-        columns = values[0]
+
+        if header is None:
+            return pd.DataFrame(values)
+
+        columns = values[header]
+        data = values[header + 1:]
+        if len(data) == 0:
+            # Return empty df
+            return pd.DataFrame(columns=columns)
+
+        df = pd.DataFrame(data)
         if len(df.columns) > len(columns):
             columns += [f'Unknown {i}' for i in range(len(df.columns) - len(columns))]
         df.columns = columns

diff --git a/readme.md b/readme.md
@@ -74,13 +74,14 @@ To download dataframe:
 ```python
 df = drive.download(drive_table=table_name, 
                     sheet_name=sheet_name,
-                    range_name='!A1:C100') # Range in Sheets; Optional
+                    range_name='!A1:C100', # Range in Sheets; Optional
+                    header=0) # Column row
 ```
 Default `range_name` is `'!A1:ZZ900000'`.
 
 To upload dataframe:
 ```python
-df = drive.download(df,
+df = drive.upload(df,
                     drive_table=table_name, 
                     sheet_name=sheet_name,
                     range_name='!B1:ZZ900000', # Range in Sheets; Optional

diff --git a/setup.py b/setup.py
@@ -8,7 +8,7 @@
 
 setup(
     name='gsheet-pandas',
-    version='0.2.4',
+    version='0.2.5',
     description='Download and upload pandas dataframes to the Google sheets',
     url='https://github.com/iakov-kaiumov/gsheet-pandas',
     author='Iakov Kaiumov',