feat: Add CSV download functionality and comprehensive tests

pystackql/magic_ext/local.py

@@ -10,6 +10,8 @@
 from IPython.core.magic import (magics_class, line_cell_magic)
 from .base import BaseStackqlMagic
 import argparse
+import base64
+from IPython.display import HTML
 
 @magics_class
 class StackqlMagic(BaseStackqlMagic):
@@ -30,6 +32,10 @@ def stackql(self, line, cell=None):
         - As a line magic: `%stackql QUERY`
         - As a cell magic: `%%stackql [OPTIONS]` followed by the QUERY in the next line.
 
+        Available options for cell magic:
+        - --no-display: Suppress result display
+        - --csv-download: Add CSV download link (only when results are displayed)
+
         :param line: The arguments and/or StackQL query when used as line magic.
         :param cell: The StackQL query when used as cell magic.
         :return: StackQL query results as a named Pandas DataFrame (`stackql_df`).
@@ -39,6 +45,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.")
             args = parser.parse_args(line.split())
             query_to_run = self.get_rendered_query(cell)
         else:
@@ -48,11 +55,62 @@ 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:
+        # Handle display logic and CSV download
+        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:
+            # Display results with CSV download link
+            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, dataframe):
+        """Display DataFrame with CSV download link.
+
+        :param dataframe: The pandas DataFrame to display and make downloadable
+        """
+        try:
+            # Generate CSV content
+            csv_content = dataframe.to_csv(index=False)
+
+            # Encode CSV content to base64 for data URI
+            csv_base64 = base64.b64encode(csv_content.encode('utf-8')).decode('utf-8')
+
+            # Create download link with styled button
+            download_html = f"""
+            <div style="margin-top: 10px;">
+                <a href="data:text/csv;base64,{csv_base64}" 
+                   download="stackql_results.csv"
+                   style="display: inline-block; padding: 8px 16px; background-color: #007cba; 
+                          color: white; text-decoration: none; border-radius: 4px; 
+                          font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, 
+                          'Helvetica Neue', Arial, sans-serif; font-size: 14px; font-weight: 500;
+                          transition: background-color 0.2s;"
+                   onmouseover="this.style.backgroundColor='#005a87'" 
+                   onmouseout="this.style.backgroundColor='#007cba'">
+                    📥 Download CSV
+                </a>
+            </div>
+            """
+
+            # Display the download link
+            from IPython.display import display
+            display(HTML(download_html))
+
+        except Exception as e:
+            # Graceful error handling
+            error_html = f"""
+            <div style="margin-top: 10px; padding: 8px; background-color: #ffebee; 
+                        border: 1px solid #ffcdd2; border-radius: 4px; color: #c62828;">
+                <strong>CSV Download Error:</strong> {str(e)}
+            </div>
+            """
+            from IPython.display import display
+            display(HTML(error_html))
+
 def load_ipython_extension(ipython):
     """Load the non-server magic in IPython.
 

pystackql/magic_ext/server.py

@@ -10,6 +10,8 @@
 from IPython.core.magic import (magics_class, line_cell_magic)
 from .base import BaseStackqlMagic
 import argparse
+import base64
+from IPython.display import HTML
 
 @magics_class
 class StackqlServerMagic(BaseStackqlMagic):
@@ -30,6 +32,10 @@ def stackql(self, line, cell=None):
         - As a line magic: `%stackql QUERY`
         - As a cell magic: `%%stackql [OPTIONS]` followed by the QUERY in the next line.
 
+        Available options for cell magic:
+        - --no-display: Suppress result display
+        - --csv-download: Add CSV download link (only when results are displayed)
+
         :param line: The arguments and/or StackQL query when used as line magic.
         :param cell: The StackQL query when used as cell magic.
         :return: StackQL query results as a named Pandas DataFrame (`stackql_df`).
@@ -39,6 +45,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.")
             args = parser.parse_args(line.split())
             query_to_run = self.get_rendered_query(cell)
         else:
@@ -48,11 +55,60 @@ def stackql(self, line, cell=None):
         results = self.run_query(query_to_run)
         self.shell.user_ns['stackql_df'] = results
 
+        # Handle display logic and CSV download
         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:
+            # Display results with CSV download link
+            self._display_with_csv_download(results)
+            return results
         else:
             return results        
 
+    def _display_with_csv_download(self, dataframe):
+        """Display DataFrame with CSV download link.
+
+        :param dataframe: The pandas DataFrame to display and make downloadable
+        """
+        try:
+            # Generate CSV content
+            csv_content = dataframe.to_csv(index=False)
+
+            # Encode CSV content to base64 for data URI
+            csv_base64 = base64.b64encode(csv_content.encode('utf-8')).decode('utf-8')
+
+            # Create download link with styled button
+            download_html = f"""
+            <div style="margin-top: 10px;">
+                <a href="data:text/csv;base64,{csv_base64}" 
+                   download="stackql_results.csv"
+                   style="display: inline-block; padding: 8px 16px; background-color: #007cba; 
+                          color: white; text-decoration: none; border-radius: 4px; 
+                          font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, 
+                          'Helvetica Neue', Arial, sans-serif; font-size: 14px; font-weight: 500;
+                          transition: background-color 0.2s;"
+                   onmouseover="this.style.backgroundColor='#005a87'" 
+                   onmouseout="this.style.backgroundColor='#007cba'">
+                    📥 Download CSV
+                </a>
+            </div>
+            """
+
+            # Display the download link
+            from IPython.display import display
+            display(HTML(download_html))
+
+        except Exception as e:
+            # Graceful error handling
+            error_html = f"""
+            <div style="margin-top: 10px; padding: 8px; background-color: #ffebee; 
+                        border: 1px solid #ffcdd2; border-radius: 4px; color: #c62828;">
+                <strong>CSV Download Error:</strong> {str(e)}
+            </div>
+            """
+            from IPython.display import display
+            display(HTML(error_html))
+
 def load_ipython_extension(ipython):
     """Load the server magic in IPython."""
     # Create an instance of the magic class and register it

tests/test_magic.py

@@ -106,6 +106,124 @@ 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 a query and --csv-download option."""
+        # 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 track if it was called
+        self.stackql_magic._display_with_csv_download = MagicMock()
+
+        # Execute the magic with our query and --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.shell.user_ns['stackql_df'].equals(self.expected_result) and
+                          self.stackql_magic._display_with_csv_download.called,
+                          False, True)
+
+    def test_cell_magic_query_csv_download_with_no_display(self):
+        """Test cell magic with both --csv-download and --no-display options."""
+        # 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 track if it was called
+        self.stackql_magic._display_with_csv_download = MagicMock()
+
+        # Execute the magic with both options
+        result = self.stackql_magic.stackql(line="--csv-download --no-display", cell=self.query)
+
+        # Validate the outcome - --no-display should take precedence
+        assert result is None, "Result should be None when --no-display is set"
+        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 
+                          self.shell.user_ns['stackql_df'].equals(self.expected_result) 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('pystackql.magic_ext.local.display') as mock_display, \
+             patch('pystackql.magic_ext.local.HTML') as mock_html:
+
+            # Call the method
+            self.stackql_magic._display_with_csv_download(test_df)
+
+            # Verify display was called
+            mock_display.assert_called_once()
+            mock_html.assert_called_once()
+
+            # Get the HTML content that was passed
+            html_content = mock_html.call_args[0][0]
+
+            # Verify the HTML contains expected elements
+            assert "data:text/csv;base64," in html_content, "HTML should contain base64 data URI"
+            assert "download=\"stackql_results.csv\"" in html_content, "HTML should contain download filename"
+            assert "📥 Download CSV" in html_content, "HTML should contain download button text"
+
+            # Verify the CSV content is properly base64 encoded
+            expected_csv = test_df.to_csv(index=False)
+            expected_base64 = base64.b64encode(expected_csv.encode('utf-8')).decode('utf-8')
+            assert expected_base64 in html_content, "HTML should contain correctly encoded CSV data"
+
+        print_test_result("CSV download method test", 
+                          True,  # All assertions passed if we reach here
+                          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('pystackql.magic_ext.local.display') as mock_display, \
+             patch('pystackql.magic_ext.local.HTML') as mock_html:
+
+            # Call the method with the problematic DataFrame
+            self.stackql_magic._display_with_csv_download(mock_df)
+
+            # Verify display was called (for error message)
+            mock_display.assert_called_once()
+            mock_html.assert_called_once()
+
+            # Get the HTML content that was passed
+            html_content = mock_html.call_args[0][0]
+
+            # Verify the HTML contains error message
+            assert "CSV Download Error:" in html_content, "HTML should contain error message"
+            assert "Test CSV error" in html_content, "HTML should contain specific error text"
+
+        print_test_result("CSV download error handling test", 
+                          True,  # All assertions passed if we reach here
+                          False, True)
+
 def test_magic_extension_loading(mock_interactive_shell):
     """Test that non-server magic extension can be loaded."""
     # Test loading non-server magic