Add support for specifing sheet from Excel file to be converted

table2sql/converters/from_excel.py

@@ -1,16 +1,16 @@
-from typing import Any, List, Tuple
+from typing import Any, List, Optional, Tuple, cast
 
 from openpyxl import load_workbook
+from openpyxl.worksheet.worksheet import Worksheet
 
-ListOfTuples = List[Tuple[Any, ...]]
 
-
-def _get_sheet_content(path_to_file: str) -> ListOfTuples:
-    wb = load_workbook(filename=path_to_file)
-    sheet = wb.active
+def get_list_of_tuples_from_excel(
+    path_to_file: str, sheet_name: Optional[str] = None
+) -> List[Tuple[Any, ...]]:
+    workbook = load_workbook(filename=path_to_file)
+    if sheet_name:
+        sheet = cast(Worksheet, workbook[sheet_name])
+    else:
+        sheet = workbook.active
 
     return list(sheet.values)
-
-
-def get_list_of_tuples_from_excel(path_to_file: str) -> ListOfTuples:
-    return _get_sheet_content(path_to_file)

table2sql/main.py

@@ -64,7 +64,7 @@ def _get_file_extension(file_path: str):
 
 
 def convert_table_file_to_insert_statement(
-    path_to_file: str, output_table: str, delimiter=str, has_types_row=bool
+    path_to_file: str, output_table: str, delimiter=str, has_types_row=bool, sheet_name=None
 ):
     """Converts CSV file to SQL insert statements.
 
@@ -78,8 +78,10 @@ def convert_table_file_to_insert_statement(
         path_to_file (str): Name of file containing data to be converted to SQL insert statements.
         output_table (str): Name of table into which data should be inserted.
         delimiter (str): Delimiter of given CSV file.
-        types_row (bool, optional): If second row of table file contains row with types
+        has_types_row (bool, optional): If second row of table file contains row with types
           (from `TYPES_MAP`). Defaults to False.
+        sheet_name (str, optional): If file is Excel - pass sheet name which should be converted
+          into insert statements.
 
     Returns:
         str: SQL insert statement
@@ -90,7 +92,7 @@ def convert_table_file_to_insert_statement(
     if file_extension == "csv":
         rows = get_list_of_tuples_from_csv(path_to_file=path_to_file, delimiter=delimiter)
     elif file_extension in ("xlsx", "xlsm", "xltx", "xltm"):
-        rows = get_list_of_tuples_from_excel(path_to_file=path_to_file)
+        rows = get_list_of_tuples_from_excel(path_to_file=path_to_file, sheet_name=sheet_name)
     else:
         raise NotImplementedError(f"'.{file_extension}' file extension is not supported")
 

tests/helpers.py

@@ -1,8 +1,9 @@
 import os
 import tempfile
-from typing import Any, Iterable, Tuple
+from typing import Any, Iterable, Optional, Tuple, cast
 
 import openpyxl
+from openpyxl.worksheet.worksheet import Worksheet
 
 
 def save_to_csv(csv_str: str):
@@ -14,17 +15,32 @@ def save_to_csv(csv_str: str):
     return new_filename
 
 
-def save_to_excel(rows: Iterable[Tuple[Any]], file_extension: str = "xlsx"):
-    wb = openpyxl.Workbook()
-    ws = wb.active
-    ws.title = "range names"
+def save_to_excel(
+    rows: Iterable[Tuple[Any]],
+    file_extension: str = "xlsx",
+    sheet_name: Optional[str] = None,
+    additional_sheets: Optional[Iterable[str]] = None,
+):
+    workbook = openpyxl.Workbook()
+
+    if sheet_name:
+        workbook.create_sheet(sheet_name)
+        sheet = cast(Worksheet, workbook[sheet_name])
+    else:
+        sheet = workbook.active
+
+    if additional_sheets:
+        for additional_sheet in additional_sheets:
+            workbook.create_sheet(additional_sheet)
+
+    sheet.title = sheet_name or "some sheet"
 
     for r_idx, row in enumerate(rows, start=1):
         for c_idx, value in enumerate(row, start=1):
-            _ = ws.cell(column=c_idx, row=r_idx, value=value)
+            _ = sheet.cell(column=c_idx, row=r_idx, value=value)
 
     with tempfile.NamedTemporaryFile(delete=False) as f:
         new_filename = f.name + "." + file_extension
         os.rename(f.name, new_filename)
-        wb.save(filename=new_filename)
+        workbook.save(filename=new_filename)
     return new_filename

tests/test_main.py

@@ -171,3 +171,33 @@ def test_convert_table_file_to_insert_statement_from_excel(
     os.remove(test_filename)
 
     assert result_insert_statement == expected
+
+
+def test_convert_table_file_to_insert_statement_from_excel_with_multiple_sheets():
+    test_data = [
+        ("a", "b", "c", "d"),
+        ("int", "str", "float", "sql"),
+        (1, 2, 3, "(SELECT 1)"),
+        (5, 6, 7, "(SELECT 1)"),
+    ]
+
+    sheet_name = "test sheet"
+
+    expected = """
+INSERT INTO test.table (a, b, c, d)
+VALUES (1, '2', 3.0, (SELECT 1)), (5, '6', 7.0, (SELECT 1));""".strip()
+
+    test_filename = save_to_excel(
+        test_data, sheet_name=sheet_name, additional_sheets=["1", "2", "3"]
+    )
+
+    result_insert_statement = convert_table_file_to_insert_statement(
+        path_to_file=test_filename,
+        output_table=TEST_TABLE_NAME,
+        has_types_row=True,
+        sheet_name=sheet_name,
+    )
+
+    os.remove(test_filename)
+
+    assert result_insert_statement == expected