feat: Re-implement CSV download functionality and fix test failures

pystackql/magic_ext/local.py

@@ -8,8 +8,11 @@
 """
 
 from IPython.core.magic import (magics_class, line_cell_magic)
+from IPython.display import display, HTML
 from .base import BaseStackqlMagic
 import argparse
+import base64
+import io
 
 @magics_class
 class StackqlMagic(BaseStackqlMagic):
@@ -39,6 +42,7 @@ def stackql(self, line, cell=None):
         if is_cell_magic:
             parser = argparse.ArgumentParser()
             parser.add_argument("--no-display", action="store_true", help="Suppress result display.")
+            parser.add_argument("--csv-download", action="store_true", help="Add CSV download link to output.")
             args = parser.parse_args(line.split())
             query_to_run = self.get_rendered_query(cell)
         else:
@@ -48,11 +52,58 @@ def stackql(self, line, cell=None):
         results = self.run_query(query_to_run)
         self.shell.user_ns['stackql_df'] = results
 
-        if is_cell_magic and args and not args.no_display:
+        if is_cell_magic and args and args.no_display:
+            return None
+        elif is_cell_magic and args and args.csv_download and not args.no_display:
+            self._display_with_csv_download(results)
+            return results
+        elif is_cell_magic and args and not args.no_display:
             return results
         elif not is_cell_magic:
             return results
 
+    def _display_with_csv_download(self, df):
+        """Display DataFrame with CSV download link.
+
+        :param df: The DataFrame to display and make downloadable.
+        """
+        import IPython.display
+
+        try:
+            # Generate CSV data
+            import io
+            import base64
+            csv_buffer = io.StringIO()
+            df.to_csv(csv_buffer, index=False)
+            csv_data = csv_buffer.getvalue()
+
+            # Encode to base64 for data URI
+            csv_base64 = base64.b64encode(csv_data.encode()).decode()
+
+            # Create download link
+            download_link = f'data:text/csv;base64,{csv_base64}'
+
+            # Display the DataFrame first
+            IPython.display.display(df)
+
+            # Create and display the download button
+            download_html = f'''
+            <div style="margin-top: 10px;">
+                <a href="{download_link}" download="stackql_results.csv" 
+                   style="display: inline-block; padding: 8px 16px; background-color: #007cba; 
+                          color: white; text-decoration: none; border-radius: 4px; 
+                          font-family: Arial, sans-serif; font-size: 14px; border: none; cursor: pointer;">
+                    📥 Download CSV
+                </a>
+            </div>
+            '''
+            IPython.display.display(IPython.display.HTML(download_html))
+
+        except Exception as e:
+            # If CSV generation fails, just display the DataFrame normally
+            IPython.display.display(df)
+            print(f"Error generating CSV download: {e}")
+
 def load_ipython_extension(ipython):
     """Load the non-server magic in IPython.
 
@@ -63,4 +114,4 @@ def load_ipython_extension(ipython):
     """
     # Create an instance of the magic class and register it
     magic_instance = StackqlMagic(ipython)
-    ipython.register_magics(magic_instance)
+    ipython.register_magics(magic_instance)

pystackql/magic_ext/server.py

@@ -8,8 +8,11 @@
 """
 
 from IPython.core.magic import (magics_class, line_cell_magic)
+from IPython.display import display, HTML
 from .base import BaseStackqlMagic
 import argparse
+import base64
+import io
 
 @magics_class
 class StackqlServerMagic(BaseStackqlMagic):
@@ -39,6 +42,7 @@ def stackql(self, line, cell=None):
         if is_cell_magic:
             parser = argparse.ArgumentParser()
             parser.add_argument("--no-display", action="store_true", help="Suppress result display.")
+            parser.add_argument("--csv-download", action="store_true", help="Add CSV download link to output.")
             args = parser.parse_args(line.split())
             query_to_run = self.get_rendered_query(cell)
         else:
@@ -50,11 +54,60 @@ def stackql(self, line, cell=None):
 
         if is_cell_magic and args and args.no_display:
             return None
+        elif is_cell_magic and args and args.csv_download and not args.no_display:
+            self._display_with_csv_download(results)
+            return results
+        elif is_cell_magic and args and not args.no_display:
+            return results
+        elif not is_cell_magic:
+            return results
         else:
-            return results        
+            return results
 
+    def _display_with_csv_download(self, df):
+        """Display DataFrame with CSV download link.
+
+        :param df: The DataFrame to display and make downloadable.
+        """
+        import IPython.display
+
+        try:
+            # Generate CSV data
+            import io
+            import base64
+            csv_buffer = io.StringIO()
+            df.to_csv(csv_buffer, index=False)
+            csv_data = csv_buffer.getvalue()
+
+            # Encode to base64 for data URI
+            csv_base64 = base64.b64encode(csv_data.encode()).decode()
+
+            # Create download link
+            download_link = f'data:text/csv;base64,{csv_base64}'
+
+            # Display the DataFrame first
+            IPython.display.display(df)
+
+            # Create and display the download button
+            download_html = f'''
+            <div style="margin-top: 10px;">
+                <a href="{download_link}" download="stackql_results.csv" 
+                   style="display: inline-block; padding: 8px 16px; background-color: #007cba; 
+                          color: white; text-decoration: none; border-radius: 4px; 
+                          font-family: Arial, sans-serif; font-size: 14px; border: none; cursor: pointer;">
+                    📥 Download CSV
+                </a>
+            </div>
+            '''
+            IPython.display.display(IPython.display.HTML(download_html))
+
+        except Exception as e:
+            # If CSV generation fails, just display the DataFrame normally
+            IPython.display.display(df)
+            print(f"Error generating CSV download: {e}")
+
 def load_ipython_extension(ipython):
     """Load the server magic in IPython."""
     # Create an instance of the magic class and register it
     magic_instance = StackqlServerMagic(ipython)
-    ipython.register_magics(magic_instance)
+    ipython.register_magics(magic_instance)

tests/test_magic.py

@@ -106,6 +106,117 @@ def test_cell_magic_query_no_display(self):
                           self.shell.user_ns['stackql_df'].equals(self.expected_result),
                           False, True)
 
+    def test_cell_magic_query_csv_download(self):
+        """Test cell magic with CSV download functionality."""
+        # Mock the run_query method to return a known DataFrame
+        self.stackql_magic.run_query = MagicMock(return_value=self.expected_result)
+
+        # Mock the _display_with_csv_download method to verify it's called
+        self.stackql_magic._display_with_csv_download = MagicMock()
+
+        # Execute the magic with --csv-download option
+        result = self.stackql_magic.stackql(line="--csv-download", cell=self.query)
+
+        # Validate the outcome
+        assert result.equals(self.expected_result), "Result should match expected DataFrame"
+        assert 'stackql_df' in self.shell.user_ns, "stackql_df should be in user namespace"
+        assert self.shell.user_ns['stackql_df'].equals(self.expected_result), "stackql_df should match expected DataFrame"
+
+        # Verify that _display_with_csv_download was called
+        self.stackql_magic._display_with_csv_download.assert_called_once_with(self.expected_result)
+
+        print_test_result("Cell magic query test with CSV download", 
+                          result.equals(self.expected_result) and 
+                          'stackql_df' in self.shell.user_ns and 
+                          self.stackql_magic._display_with_csv_download.called,
+                          False, True)
+
+    def test_cell_magic_query_csv_download_with_no_display(self):
+        """Test that --no-display takes precedence over --csv-download."""
+        # Mock the run_query method to return a known DataFrame
+        self.stackql_magic.run_query = MagicMock(return_value=self.expected_result)
+
+        # Mock the _display_with_csv_download method to verify it's not called
+        self.stackql_magic._display_with_csv_download = MagicMock()
+
+        # Execute the magic with both --csv-download and --no-display options
+        result = self.stackql_magic.stackql(line="--csv-download --no-display", cell=self.query)
+
+        # Validate the outcome
+        assert result is None, "Result should be None with --no-display option"
+        assert 'stackql_df' in self.shell.user_ns, "stackql_df should still be in user namespace"
+        assert self.shell.user_ns['stackql_df'].equals(self.expected_result), "stackql_df should match expected DataFrame"
+
+        # Verify that _display_with_csv_download was NOT called
+        self.stackql_magic._display_with_csv_download.assert_not_called()
+
+        print_test_result("Cell magic query test with CSV download and no-display", 
+                          result is None and 
+                          'stackql_df' in self.shell.user_ns and 
+                          not self.stackql_magic._display_with_csv_download.called,
+                          False, True)
+
+    def test_display_with_csv_download_method(self):
+        """Test the _display_with_csv_download method directly."""
+        import base64
+        from unittest.mock import patch
+
+        # Create a test DataFrame
+        test_df = pd.DataFrame({"col1": [1, 2], "col2": ["a", "b"]})
+
+        # Mock IPython display functionality
+        with patch('IPython.display.display') as mock_display, \
+             patch('IPython.display.HTML') as mock_html:
+
+            # Call the method
+            self.stackql_magic._display_with_csv_download(test_df)
+
+            # Verify display was called twice (once for DataFrame, once for HTML)
+            assert mock_display.call_count == 2, "Display should be called twice"
+
+            # Verify HTML was called once
+            mock_html.assert_called_once()
+
+            # Check that the HTML call contains download link
+            html_call_args = mock_html.call_args[0][0]
+            assert 'download="stackql_results.csv"' in html_call_args
+            assert 'data:text/csv;base64,' in html_call_args
+
+        print_test_result("_display_with_csv_download method test", 
+                          mock_display.call_count == 2 and mock_html.called,
+                          False, True)
+
+    def test_display_with_csv_download_error_handling(self):
+        """Test error handling in _display_with_csv_download method."""
+        from unittest.mock import patch
+
+        # Create a mock DataFrame that will raise an exception during to_csv()
+        mock_df = MagicMock()
+        mock_df.to_csv.side_effect = Exception("Test CSV error")
+
+        # Mock IPython display functionality
+        with patch('IPython.display.display') as mock_display, \
+             patch('IPython.display.HTML') as mock_html, \
+             patch('builtins.print') as mock_print:
+
+            # Call the method with the problematic DataFrame
+            self.stackql_magic._display_with_csv_download(mock_df)
+
+            # Verify display was called once (for DataFrame only, not for HTML)
+            mock_display.assert_called_once_with(mock_df)
+
+            # Verify HTML was not called due to error
+            mock_html.assert_not_called()
+
+            # Verify error message was printed
+            mock_print.assert_called_once()
+            error_message = mock_print.call_args[0][0]
+            assert "Error generating CSV download:" in error_message
+
+        print_test_result("_display_with_csv_download error handling test", 
+                          mock_display.called and not mock_html.called and mock_print.called,
+                          False, True)
+
 def test_magic_extension_loading(mock_interactive_shell):
     """Test that non-server magic extension can be loaded."""
     # Test loading non-server magic