Add method to create new sheet and list existing sheets

gsheet_pandas/__init__.py

@@ -1,5 +1,5 @@
 from .adapter.connection import DriveConnection, setup
 
 
-__version__ = '0.1.8'
+__version__ = '0.2.6'
 __author__ = 'Iakov Kaiumov'

gsheet_pandas/adapter/connection.py

@@ -1,10 +1,10 @@
-from __future__ import print_function
-
 import datetime
+import logging
 import os.path
 import socket
 from pathlib import Path
 
+import googleapiclient
 import pandas as pd
 from pandas import Timestamp
 from pandas._libs.lib import Decimal
@@ -16,7 +16,7 @@
 from googleapiclient.discovery import build
 from googleapiclient.errors import HttpError
 
-
+logger = logging.getLogger('gsheet-pandas')
 timeout_in_sec = 60 * 1
 socket.setdefaulttimeout(timeout_in_sec)
 
@@ -82,7 +82,7 @@ def _get_service(self, service_name='sheets', service_version='v4'):
             service = build(service_name, service_version, credentials=creds)
             return service
         except HttpError as err:
-            print(err)
+            logger.error(err)
             return None
 
     def get_all_files_in_folder(self, folder_id):
@@ -103,15 +103,25 @@ def get_all_files_in_folder(self, folder_id):
                     break
 
         except HttpError as e:
-            print(f'get_all_files_in_folder: An error occurred - {e}')
             raise e
         return files
 
-    def download(self, drive_table: str, sheet_name: str, range_name: str = DEFAULT_RANGE_NAME,
+    def download(self,
+                 spreadsheet_id: str,
+                 sheet_name: str,
+                 range_name: str = DEFAULT_RANGE_NAME,
                  header: int | None = 0) -> pd.DataFrame:
+        """
+        Downloads Google Spreadsheet as Pandas DataFrame
+        :param spreadsheet_id: spreadsheet id
+        :param sheet_name: sheet name
+        :param range_name: range name (default is !A1:ZZ900000)
+        :param header: index of header row
+        :return dataframe
+        """
         service = self._get_service()
         sheet = self._get_service().spreadsheets()
-        result = sheet.values().get(spreadsheetId=drive_table, range=sheet_name + range_name).execute()
+        result = sheet.values().get(spreadsheetId=spreadsheet_id, range=sheet_name + range_name).execute()
         values = result.get('values', [])
         service.close()
         if not values:
@@ -134,26 +144,74 @@ def download(self, drive_table: str, sheet_name: str, range_name: str = DEFAULT_
 
     def upload(self,
                df: pd.DataFrame,
-               drive_table: str,
+               spreadsheet_id: str,
                sheet_name: str,
                range_name: str = DEFAULT_RANGE_NAME,
-               drop_columns: bool = False):
+               drop_columns: bool = False) -> None:
+        """
+        Uploads Pandas DataFrame to the Google Spreadsheet
+        :param df: Pandas DataFrame
+        :param spreadsheet_id: spreadsheet id
+        :param sheet_name: sheet name
+        :param range_name: range name (default is !A1:ZZ900000)
+        :param drop_columns: whether to drop DataFrame columns or not
+        """
         try:
             df = _fix_dtypes(df)
             values = df.T.reset_index().T.values.tolist()
             if drop_columns:
                 values = df.values.tolist()
             service = self._get_service()
             service.spreadsheets().values().update(
-                spreadsheetId=drive_table,
+                spreadsheetId=spreadsheet_id,
                 valueInputOption='RAW',
                 range=sheet_name + range_name,
                 body=dict(majorDimension='ROWS', values=values),
             ).execute()
             service.close()
         except socket.timeout as e:
-            print(f'pandas_to_sheet: Error: {e}')
             raise e
         except Exception as e:
-            print(f'pandas_to_sheet: Error: {e}')
             raise e
+
+    def get_sheets_names(self, spreadsheet_id: str) -> list[str]:
+        """
+        Get sheets names for spreadsheet
+        :param spreadsheet_id: spreadsheet id
+        :return: list of names
+        """
+        service = self._get_service()
+        sheet_metadata = service.spreadsheets().get(spreadsheetId=spreadsheet_id).execute()
+        sheets = sheet_metadata.get('sheets', '')
+        return [sheet.get("properties", {}).get("title") for sheet in sheets]
+
+    def create_sheet(self, spreadsheet_id: str, sheet_name: str) -> int | None:
+        """
+        Creates new sheet in existing spreadsheet
+        :param spreadsheet_id: spreadsheet id
+        :param sheet_name: new sheet's name
+        :return: new sheet id if success, None if sheet exists
+        """
+        service = self._get_service()
+        batch_update_spreadsheet_request_body = {
+            "requests": [
+                {
+                    "addSheet": {
+                        "properties": {
+                            "title": sheet_name,
+                        }
+                    }
+                }
+            ]
+        }
+
+        request = service.spreadsheets().batchUpdate(spreadsheetId=spreadsheet_id,
+                                                     body=batch_update_spreadsheet_request_body)
+        try:
+            response = request.execute()
+        except googleapiclient.errors.HttpError as e:
+            if e.status_code == 400:
+                return None
+            raise e
+        service.close()
+        return response['replies'][0]['addSheet']['properties']['sheetId']

gsheet_pandas/tests.py

@@ -9,45 +9,59 @@
 from adapter import connection
 import dotenv
 
-
 BASE_DIR = Path(__file__).resolve().parent.parent
 dotenv.load_dotenv(BASE_DIR / '.env')
 
-table_name = os.getenv('table_name')
+spreadsheet_id = os.getenv('table_name')
 sheet_name = os.getenv('sheet_name')
 
 
+data_dir = Path(__file__).resolve().parent.parent / 'data'
+
+
 class TestConnectionMethods(unittest.TestCase):
+    @staticmethod
+    def _get_drive() -> connection.DriveConnection:
+        return connection.DriveConnection(
+            credentials_dir=data_dir / 'credentials.json',
+            token_dir=data_dir / 'token.json'
+        )
+
+    def test_list_sheets(self):
+        drive = self._get_drive()
+        sheets = drive.get_sheets_names(spreadsheet_id)
+        self.assertEqual(sheets, ['test', 'test2'])
+
+    def test_create_sheet(self):
+        drive = self._get_drive()
+        _id = drive.create_sheet(spreadsheet_id=spreadsheet_id, sheet_name='test2')
+        self.assertIsNone(_id)
 
     def test_connection_class(self):
-        data_dir = Path(__file__).resolve().parent.parent / 'data'
-        drive = connection.DriveConnection(credentials_dir=data_dir / 'credentials.json', token_dir=data_dir / 'token.json')
-
-        df = drive.download(drive_table=table_name, sheet_name=sheet_name)
+        drive = self._get_drive()
+        df = drive.download(spreadsheet_id=spreadsheet_id, sheet_name=sheet_name)
 
         new_column_value = str(random.random())
         df['column1'] = new_column_value
 
-        drive.upload(df, drive_table=table_name, sheet_name=sheet_name)
+        drive.upload(df, spreadsheet_id=spreadsheet_id, sheet_name=sheet_name)
 
-        df = drive.download(drive_table=table_name, sheet_name=sheet_name)
+        df = drive.download(spreadsheet_id=spreadsheet_id, sheet_name=sheet_name)
 
         values = df['column1'].values.tolist()
         for value in values:
             self.assertEqual(value, new_column_value)
 
     def test_pandas_extension(self):
-        data_dir = Path(__file__).resolve().parent.parent / 'data'
-
         connection.setup(credentials_dir=data_dir / 'credentials.json', token_dir=data_dir / 'token.json')
 
-        df = pd.from_gsheet(drive_table=table_name, sheet_name=sheet_name)
+        df = pd.from_gsheet(spreadsheet_id=spreadsheet_id, sheet_name=sheet_name)
         new_column_value = str(random.random())
         df['column1'] = new_column_value
 
-        df.to_gsheet(drive_table=table_name, sheet_name=sheet_name)
+        df.to_gsheet(spreadsheet_id=spreadsheet_id, sheet_name=sheet_name)
 
-        df = pd.from_gsheet(drive_table=table_name, sheet_name=sheet_name)
+        df = pd.from_gsheet(spreadsheet_id=spreadsheet_id, sheet_name=sheet_name)
         values = df['column1'].values.tolist()
         for value in values:
             self.assertEqual(value, new_column_value)

readme.md

@@ -47,15 +47,15 @@ To download dataframe:
 ```python
 import pandas as pd
 
-df = pd.from_gsheet(drive_table=table_name, 
+df = pd.from_gsheet(spreadsheet_id, 
                     sheet_name=sheet_name,
                     range_name='!A1:C100') # Range in Sheets; Optional
 ```
 Default `range_name` is `'!A1:ZZ900000'`.
 
 To upload dataframe:
 ```python
-df.to_gsheet(drive_table=table_name, 
+df.to_gsheet(spreadsheet_id, 
              sheet_name=sheet_name,
              range_name='!B1:ZZ900000', # Range in Sheets; Optional
              drop_columns=False) # Upload column names or not; Optional
@@ -72,7 +72,7 @@ drive = DriveConnection(credentials_dir=secret_path / 'credentials.json',
 
 To download dataframe:
 ```python
-df = drive.download(drive_table=table_name, 
+df = drive.download(spreadsheet_id, 
                     sheet_name=sheet_name,
                     range_name='!A1:C100', # Range in Sheets; Optional
                     header=0) # Column row
@@ -82,8 +82,8 @@ Default `range_name` is `'!A1:ZZ900000'`.
 To upload dataframe:
 ```python
 df = drive.upload(df,
-                    drive_table=table_name, 
-                    sheet_name=sheet_name,
-                    range_name='!B1:ZZ900000', # Range in Sheets; Optional
-                    drop_columns=False) # Upload column names or not; Optional
+                  spreadsheet_id, 
+                  sheet_name=sheet_name,
+                  range_name='!B1:ZZ900000', # Range in Sheets; Optional
+                  drop_columns=False) # Upload column names or not; Optional
 ```

setup.py

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