Merge pull request #3 from uzairmukadam/v2.0.1b

Fixed result export issue requiring an exisiting file rather than creating a new one.

Author: uzairmukadam

sheet_ql.py

@@ -4,7 +4,7 @@
 from collections import deque
 import pandas as pd
 import duckdb
-from typing import Any, Optional
+from typing import Any, Optional, List, Tuple, Dict
 
 from rich.console import Console
 from rich.table import Table
@@ -39,7 +39,7 @@ def __init__(self) -> None:
         """Initializes the SheetQL tool."""
         self.console = Console()
         self.db_connection: Optional[duckdb.DuckDBPyConnection] = None
-        self.results_to_save: dict[str, pd.DataFrame] = {}
+        self.results_to_save: Dict[str, pd.DataFrame] = {}
         self.history: deque[str] = deque(maxlen=self.HISTORY_MAX_LEN)
 
     def run_interactive(self) -> None:
@@ -116,9 +116,9 @@ def _display_welcome(self) -> None:
             )
 
     def _prompt_for_paths(
-        self, title: str, filetypes: list[tuple[str, str]], allow_multiple: bool
-    ) -> Optional[list[str]]:
-        """Generic method to get file paths from the user via GUI or CLI."""
+        self, title: str, filetypes: List[Tuple[str, str]], allow_multiple: bool
+    ) -> Optional[List[str]]:
+        """Generic method to get existing file paths from the user."""
         if TKINTER_AVAILABLE:
             root = tk.Tk()
             root.withdraw()
@@ -147,7 +147,39 @@ def _prompt_for_paths(
 
         return valid_paths
 
-    def _load_data(self, file_paths: list[str]) -> dict[str, pd.DataFrame]:
+    def _prompt_for_save_path(self) -> Optional[str]:
+        """Prompts the user for a new file save path."""
+        if TKINTER_AVAILABLE:
+            root = tk.Tk()
+            root.withdraw()
+            save_path = filedialog.asksaveasfilename(
+                title="Select Save Location",
+                initialfile=self.DEFAULT_EXPORT_FILENAME,
+                defaultextension=".xlsx",
+                filetypes=[("Excel Files", "*.xlsx")],
+            )
+            root.destroy()
+            return save_path if save_path else None
+
+        self.console.print("\n[cyan]Please enter a save path for the export.[/cyan]")
+        save_path_input = self.console.input(
+            f"[bold]Save path (default: {self.DEFAULT_EXPORT_FILENAME}): [/bold]"
+        )
+        if not save_path_input:
+            save_path_input = self.DEFAULT_EXPORT_FILENAME
+
+        directory = os.path.dirname(save_path_input)
+        if directory and not os.path.exists(directory):
+            try:
+                os.makedirs(directory)
+            except OSError as e:
+                self.console.print(
+                    f"[red]Error: Could not create directory '{directory}'. {e}[/red]"
+                )
+                return None
+        return save_path_input
+
+    def _load_data(self, file_paths: List[str]) -> Dict[str, pd.DataFrame]:
         """Loads all supported files into pandas DataFrames."""
         all_dataframes = {}
         with self.console.status("[bold green]Loading data files...[/bold green]"):
@@ -196,7 +228,7 @@ def _load_data(self, file_paths: list[str]) -> dict[str, pd.DataFrame]:
         )
         return all_dataframes
 
-    def _register_dataframes(self, dataframes: dict[str, pd.DataFrame]) -> None:
+    def _register_dataframes(self, dataframes: Dict[str, pd.DataFrame]) -> None:
         """Registers new DataFrames as views in the existing DuckDB connection."""
         if not self.db_connection:
             return
@@ -240,7 +272,7 @@ def _handle_meta_command(self, command_str: str) -> bool:
         """Handles meta-commands. Returns True if the app should exit."""
         parts = command_str.split()
         command = parts[0].lower()
-        commands: dict[str, Any] = {
+        commands: Dict[str, Any] = {
             ".exit": lambda: True,
             ".quit": lambda: True,
             ".help": self._show_help,
@@ -354,14 +386,8 @@ def _export_results(self) -> None:
             self.console.print("[yellow]No results are staged for export.[/yellow]")
             return
 
-        save_path = self._prompt_for_paths(
-            title="Select Save Location",
-            filetypes=[("Excel Files", "*.xlsx")],
-            allow_multiple=False,
-        )
-
-        if save_path and save_path[0]:
-            self._save_to_excel(save_path[0])
+        if save_path := self._prompt_for_save_path():
+            self._save_to_excel(save_path)
         else:
             self.console.print("[yellow]Save operation cancelled.[/yellow]")
 
@@ -395,11 +421,14 @@ def _format_excel_sheets(self, writer: pd.ExcelWriter) -> None:
                 cell.fill = header_fill
 
             for column_cells in worksheet.columns:
-                max_length = max(
-                    len(str(cell.value))
-                    for cell in column_cells
-                    if cell.value is not None
-                )
+                try:
+                    max_length = max(
+                        len(str(cell.value))
+                        for cell in column_cells
+                        if cell.value is not None
+                    )
+                except ValueError:
+                    max_length = 0
                 worksheet.column_dimensions[column_cells[0].column_letter].width = (
                     max_length + 2
                 )
@@ -437,7 +466,7 @@ def _add_new_files(self) -> None:
                 self.console.print("[green]✔ New files loaded and registered.[/green]")
                 self._list_tables()
 
-    def _describe_table(self, command_parts: list[str]) -> None:
+    def _describe_table(self, command_parts: List[str]) -> None:
         """Shows the schema for a given table."""
         if len(command_parts) != 2:
             self.console.print("[red]Usage: .schema <table_name>[/red]")
@@ -462,7 +491,7 @@ def _describe_table(self, command_parts: list[str]) -> None:
         except Exception as e:
             self.console.print(f"[bold red]❌ Error describing table: {e}[/bold red]")
 
-    def _rename_table(self, command_parts: list[str]) -> None:
+    def _rename_table(self, command_parts: List[str]) -> None:
         """Renames a view in the database."""
         if len(command_parts) != 3:
             self.console.print(
@@ -510,7 +539,7 @@ def _handle_history_rerun(self, command_str: str) -> None:
                 "[red]Invalid history command. Use !N where N is a number.[/red]"
             )
 
-    def _run_script_interactive(self, command_parts: list[str]) -> None:
+    def _run_script_interactive(self, command_parts: List[str]) -> None:
         """Handles the .runscript meta-command."""
         script_path = command_parts[1] if len(command_parts) == 2 else None
         if not script_path:
@@ -546,7 +575,7 @@ def _run_script_interactive(self, command_parts: list[str]) -> None:
         except yaml.YAMLError as e:
             self.console.print(f"[bold red]❌ Error parsing YAML file: {e}[/bold red]")
 
-    def _execute_yaml_script(self, config: dict[str, Any]) -> None:
+    def _execute_yaml_script(self, config: Dict[str, Any]) -> None:
         """Processes the actions defined in a parsed YAML script."""
         if "inputs" in config:
             self._process_yaml_inputs(config.get("inputs", []))
@@ -555,7 +584,7 @@ def _execute_yaml_script(self, config: dict[str, Any]) -> None:
         if "export" in config:
             self._process_yaml_export(config.get("export", {}))
 
-    def _process_yaml_inputs(self, inputs: list[dict[str, str]]) -> None:
+    def _process_yaml_inputs(self, inputs: List[Dict[str, str]]) -> None:
         """Loads and aliases data files from a YAML script's 'inputs' block."""
         self.console.print("\n[bold]--- 1. Loading Input Files ---[/bold]")
         if not inputs:
@@ -590,7 +619,7 @@ def _process_yaml_inputs(self, inputs: list[dict[str, str]]) -> None:
                 self.db_connection.register(final_name, df)
         self.console.print("[green]✔ All specified inputs loaded and aliased.[/green]")
 
-    def _process_yaml_tasks(self, tasks: list[dict[str, str]]) -> None:
+    def _process_yaml_tasks(self, tasks: List[Dict[str, str]]) -> None:
         """Executes queries from a YAML script's 'tasks' block."""
         self.console.print("\n[bold]--- 2. Executing Tasks ---[/bold]")
         if not tasks:
@@ -613,7 +642,7 @@ def _process_yaml_tasks(self, tasks: list[dict[str, str]]) -> None:
                     f"[bold red]  ❌ Error in task '{name}': {e}[/bold red]"
                 )
 
-    def _process_yaml_export(self, export_config: dict[str, str]) -> None:
+    def _process_yaml_export(self, export_config: Dict[str, str]) -> None:
         """Exports staged results based on a YAML script's 'export' block."""
         self.console.print("\n[bold]--- 3. Exporting Results ---[/bold]")
         if not export_config:

tests/test_sheet_ql.py

@@ -149,19 +149,18 @@ def test_07_load_command(self):
         self.assertIn("sample_json", final_tables)
 
     @patch("sheet_ql.SheetQL._format_excel_sheets")
-    @patch("sheet_ql.SheetQL._prompt_for_paths")
-    def test_08_export_command(self, mock_prompt, mock_format):
+    @patch("sheet_ql.SheetQL._prompt_for_save_path")
+    def test_08_export_command(self, mock_prompt_save, mock_format):
         """Test the '.export' command workflow."""
         self.tool.results_to_save["my_results"] = pd.DataFrame({"a": [1]})
         self.assertTrue(self.tool.results_to_save)
 
-        # Use the safe, temporary test directory
         save_path = os.path.join(self.test_dir, "report.xlsx")
-        mock_prompt.return_value = [save_path]
+        mock_prompt_save.return_value = save_path
 
         self.tool._export_results()
 
-        mock_prompt.assert_called_once()
+        mock_prompt_save.assert_called_once()
         mock_format.assert_called_once()
         self.assertFalse(
             self.tool.results_to_save,