diff --git a/pystackql/magic_ext/local.py b/pystackql/magic_ext/local.py index 0c5952c..ba98d99 100644 --- a/pystackql/magic_ext/local.py +++ b/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''' +
+ + 📥 Download CSV + +
+ ''' + 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) \ No newline at end of file + ipython.register_magics(magic_instance) diff --git a/pystackql/magic_ext/server.py b/pystackql/magic_ext/server.py index d7ccdcf..074603f 100644 --- a/pystackql/magic_ext/server.py +++ b/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''' +
+ + 📥 Download CSV + +
+ ''' + 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) \ No newline at end of file + ipython.register_magics(magic_instance) diff --git a/tests/test_magic.py b/tests/test_magic.py index a32ab7c..aaaac49 100644 --- a/tests/test_magic.py +++ b/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 diff --git a/tests/test_server_magic.py b/tests/test_server_magic.py index fa0a9e2..0296f37 100644 --- a/tests/test_server_magic.py +++ b/tests/test_server_magic.py @@ -104,6 +104,117 @@ def test_cell_magic_query_no_display(self): self.shell.user_ns['stackql_df'].equals(self.expected_result), True, True) + def test_cell_magic_query_csv_download(self): + """Test cell magic with CSV download functionality in server mode.""" + # 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 (server mode)", + result.equals(self.expected_result) and + 'stackql_df' in self.shell.user_ns and + self.stackql_magic._display_with_csv_download.called, + True, True) + + def test_cell_magic_query_csv_download_with_no_display(self): + """Test that --no-display takes precedence over --csv-download in server mode.""" + # 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 (server mode)", + result is None and + 'stackql_df' in self.shell.user_ns and + not self.stackql_magic._display_with_csv_download.called, + True, True) + + def test_display_with_csv_download_method(self): + """Test the _display_with_csv_download method directly in server mode.""" + 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 (server mode)", + mock_display.call_count == 2 and mock_html.called, + True, True) + + def test_display_with_csv_download_error_handling(self): + """Test error handling in _display_with_csv_download method in server mode.""" + 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 (server mode)", + mock_display.called and not mock_html.called and mock_print.called, + True, True) + def test_server_magic_extension_loading(mock_interactive_shell): """Test that server magic extension can be loaded.""" # Test loading server magic