diff --git a/pystackql/magic_ext/local.py b/pystackql/magic_ext/local.py index 0c5952c..ca29de8 100644 --- a/pystackql/magic_ext/local.py +++ b/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""" +
+ + 📥 Download CSV + +
+ """ + + # Display the download link + from IPython.display import display + display(HTML(download_html)) + + except Exception as e: + # Graceful error handling + error_html = f""" +
+ CSV Download Error: {str(e)} +
+ """ + from IPython.display import display + display(HTML(error_html)) + def load_ipython_extension(ipython): """Load the non-server magic in IPython. diff --git a/pystackql/magic_ext/server.py b/pystackql/magic_ext/server.py index d7ccdcf..16b88a5 100644 --- a/pystackql/magic_ext/server.py +++ b/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""" +
+ + 📥 Download CSV + +
+ """ + + # Display the download link + from IPython.display import display + display(HTML(download_html)) + + except Exception as e: + # Graceful error handling + error_html = f""" +
+ CSV Download Error: {str(e)} +
+ """ + 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 diff --git a/tests/test_magic.py b/tests/test_magic.py index a32ab7c..c9b148d 100644 --- a/tests/test_magic.py +++ b/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 diff --git a/tests/test_server_magic.py b/tests/test_server_magic.py index fa0a9e2..ed9b7de 100644 --- a/tests/test_server_magic.py +++ b/tests/test_server_magic.py @@ -104,6 +104,124 @@ 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 a query and --csv-download option 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 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, server mode)", + 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, + True, True) + + def test_cell_magic_query_csv_download_with_no_display(self): + """Test cell magic with both --csv-download and --no-display options 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 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, server mode)", + 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, + 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('pystackql.magic_ext.server.display') as mock_display, \ + patch('pystackql.magic_ext.server.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 (server mode)", + True, # All assertions passed if we reach here + 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('pystackql.magic_ext.server.display') as mock_display, \ + patch('pystackql.magic_ext.server.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 (server mode)", + True, # All assertions passed if we reach here + True, True) + def test_server_magic_extension_loading(mock_interactive_shell): """Test that server magic extension can be loaded.""" # Test loading server magic