Skip to content

Commit 4ddb831

Browse files
jeffreyavenjeffreyaven
authored
Merge pull request #52 from stackql/claude/issue-18-20250612_000410
feat: Re-implement CSV download functionality and fix test failures
2 parents 926d0b5 + cd0c4ca commit 4ddb831

File tree

4 files changed

+330
-4
lines changed

4 files changed

+330
-4
lines changed

‎pystackql/magic_ext/local.py

Lines changed: 53 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -8,8 +8,11 @@
88
"""
99

1010
from IPython.core.magic import (magics_class, line_cell_magic)
11+
from IPython.display import display, HTML
1112
from .base import BaseStackqlMagic
1213
import argparse
14+
import base64
15+
import io
1316

1417
@magics_class
1518
class StackqlMagic(BaseStackqlMagic):
@@ -39,6 +42,7 @@ def stackql(self, line, cell=None):
3942
if is_cell_magic:
4043
parser = argparse.ArgumentParser()
4144
parser.add_argument("--no-display", action="store_true", help="Suppress result display.")
45+
parser.add_argument("--csv-download", action="store_true", help="Add CSV download link to output.")
4246
args = parser.parse_args(line.split())
4347
query_to_run = self.get_rendered_query(cell)
4448
else:
@@ -48,11 +52,58 @@ def stackql(self, line, cell=None):
4852
results = self.run_query(query_to_run)
4953
self.shell.user_ns['stackql_df'] = results
5054

51-
if is_cell_magic and args and not args.no_display:
55+
if is_cell_magic and args and args.no_display:
56+
return None
57+
elif is_cell_magic and args and args.csv_download and not args.no_display:
58+
self._display_with_csv_download(results)
59+
return results
60+
elif is_cell_magic and args and not args.no_display:
5261
return results
5362
elif not is_cell_magic:
5463
return results
5564

65+
def _display_with_csv_download(self, df):
66+
"""Display DataFrame with CSV download link.
67+
68+
:param df: The DataFrame to display and make downloadable.
69+
"""
70+
import IPython.display
71+
72+
try:
73+
# Generate CSV data
74+
import io
75+
import base64
76+
csv_buffer = io.StringIO()
77+
df.to_csv(csv_buffer, index=False)
78+
csv_data = csv_buffer.getvalue()
79+
80+
# Encode to base64 for data URI
81+
csv_base64 = base64.b64encode(csv_data.encode()).decode()
82+
83+
# Create download link
84+
download_link = f'data:text/csv;base64,{csv_base64}'
85+
86+
# Display the DataFrame first
87+
IPython.display.display(df)
88+
89+
# Create and display the download button
90+
download_html = f'''
91+
<div style="margin-top: 10px;">
92+
<a href="{download_link}" download="stackql_results.csv"
93+
style="display: inline-block; padding: 8px 16px; background-color: #007cba;
94+
color: white; text-decoration: none; border-radius: 4px;
95+
font-family: Arial, sans-serif; font-size: 14px; border: none; cursor: pointer;">
96+
📥 Download CSV
97+
</a>
98+
</div>
99+
'''
100+
IPython.display.display(IPython.display.HTML(download_html))
101+
102+
except Exception as e:
103+
# If CSV generation fails, just display the DataFrame normally
104+
IPython.display.display(df)
105+
print(f"Error generating CSV download: {e}")
106+
56107
def load_ipython_extension(ipython):
57108
"""Load the non-server magic in IPython.
58109
@@ -63,4 +114,4 @@ def load_ipython_extension(ipython):
63114
"""
64115
# Create an instance of the magic class and register it
65116
magic_instance = StackqlMagic(ipython)
66-
ipython.register_magics(magic_instance)
117+
ipython.register_magics(magic_instance)

‎pystackql/magic_ext/server.py

Lines changed: 55 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -8,8 +8,11 @@
88
"""
99

1010
from IPython.core.magic import (magics_class, line_cell_magic)
11+
from IPython.display import display, HTML
1112
from .base import BaseStackqlMagic
1213
import argparse
14+
import base64
15+
import io
1316

1417
@magics_class
1518
class StackqlServerMagic(BaseStackqlMagic):
@@ -39,6 +42,7 @@ def stackql(self, line, cell=None):
3942
if is_cell_magic:
4043
parser = argparse.ArgumentParser()
4144
parser.add_argument("--no-display", action="store_true", help="Suppress result display.")
45+
parser.add_argument("--csv-download", action="store_true", help="Add CSV download link to output.")
4246
args = parser.parse_args(line.split())
4347
query_to_run = self.get_rendered_query(cell)
4448
else:
@@ -50,11 +54,60 @@ def stackql(self, line, cell=None):
5054

5155
if is_cell_magic and args and args.no_display:
5256
return None
57+
elif is_cell_magic and args and args.csv_download and not args.no_display:
58+
self._display_with_csv_download(results)
59+
return results
60+
elif is_cell_magic and args and not args.no_display:
61+
return results
62+
elif not is_cell_magic:
63+
return results
5364
else:
54-
return results
65+
return results
5566

67+
def _display_with_csv_download(self, df):
68+
"""Display DataFrame with CSV download link.
69+
70+
:param df: The DataFrame to display and make downloadable.
71+
"""
72+
import IPython.display
73+
74+
try:
75+
# Generate CSV data
76+
import io
77+
import base64
78+
csv_buffer = io.StringIO()
79+
df.to_csv(csv_buffer, index=False)
80+
csv_data = csv_buffer.getvalue()
81+
82+
# Encode to base64 for data URI
83+
csv_base64 = base64.b64encode(csv_data.encode()).decode()
84+
85+
# Create download link
86+
download_link = f'data:text/csv;base64,{csv_base64}'
87+
88+
# Display the DataFrame first
89+
IPython.display.display(df)
90+
91+
# Create and display the download button
92+
download_html = f'''
93+
<div style="margin-top: 10px;">
94+
<a href="{download_link}" download="stackql_results.csv"
95+
style="display: inline-block; padding: 8px 16px; background-color: #007cba;
96+
color: white; text-decoration: none; border-radius: 4px;
97+
font-family: Arial, sans-serif; font-size: 14px; border: none; cursor: pointer;">
98+
📥 Download CSV
99+
</a>
100+
</div>
101+
'''
102+
IPython.display.display(IPython.display.HTML(download_html))
103+
104+
except Exception as e:
105+
# If CSV generation fails, just display the DataFrame normally
106+
IPython.display.display(df)
107+
print(f"Error generating CSV download: {e}")
108+
56109
def load_ipython_extension(ipython):
57110
"""Load the server magic in IPython."""
58111
# Create an instance of the magic class and register it
59112
magic_instance = StackqlServerMagic(ipython)
60-
ipython.register_magics(magic_instance)
113+
ipython.register_magics(magic_instance)

‎tests/test_magic.py

Lines changed: 111 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -106,6 +106,117 @@ def test_cell_magic_query_no_display(self):
106106
self.shell.user_ns['stackql_df'].equals(self.expected_result),
107107
False, True)
108108

109+
def test_cell_magic_query_csv_download(self):
110+
"""Test cell magic with CSV download functionality."""
111+
# Mock the run_query method to return a known DataFrame
112+
self.stackql_magic.run_query = MagicMock(return_value=self.expected_result)
113+
114+
# Mock the _display_with_csv_download method to verify it's called
115+
self.stackql_magic._display_with_csv_download = MagicMock()
116+
117+
# Execute the magic with --csv-download option
118+
result = self.stackql_magic.stackql(line="--csv-download", cell=self.query)
119+
120+
# Validate the outcome
121+
assert result.equals(self.expected_result), "Result should match expected DataFrame"
122+
assert 'stackql_df' in self.shell.user_ns, "stackql_df should be in user namespace"
123+
assert self.shell.user_ns['stackql_df'].equals(self.expected_result), "stackql_df should match expected DataFrame"
124+
125+
# Verify that _display_with_csv_download was called
126+
self.stackql_magic._display_with_csv_download.assert_called_once_with(self.expected_result)
127+
128+
print_test_result("Cell magic query test with CSV download",
129+
result.equals(self.expected_result) and
130+
'stackql_df' in self.shell.user_ns and
131+
self.stackql_magic._display_with_csv_download.called,
132+
False, True)
133+
134+
def test_cell_magic_query_csv_download_with_no_display(self):
135+
"""Test that --no-display takes precedence over --csv-download."""
136+
# Mock the run_query method to return a known DataFrame
137+
self.stackql_magic.run_query = MagicMock(return_value=self.expected_result)
138+
139+
# Mock the _display_with_csv_download method to verify it's not called
140+
self.stackql_magic._display_with_csv_download = MagicMock()
141+
142+
# Execute the magic with both --csv-download and --no-display options
143+
result = self.stackql_magic.stackql(line="--csv-download --no-display", cell=self.query)
144+
145+
# Validate the outcome
146+
assert result is None, "Result should be None with --no-display option"
147+
assert 'stackql_df' in self.shell.user_ns, "stackql_df should still be in user namespace"
148+
assert self.shell.user_ns['stackql_df'].equals(self.expected_result), "stackql_df should match expected DataFrame"
149+
150+
# Verify that _display_with_csv_download was NOT called
151+
self.stackql_magic._display_with_csv_download.assert_not_called()
152+
153+
print_test_result("Cell magic query test with CSV download and no-display",
154+
result is None and
155+
'stackql_df' in self.shell.user_ns and
156+
not self.stackql_magic._display_with_csv_download.called,
157+
False, True)
158+
159+
def test_display_with_csv_download_method(self):
160+
"""Test the _display_with_csv_download method directly."""
161+
import base64
162+
from unittest.mock import patch
163+
164+
# Create a test DataFrame
165+
test_df = pd.DataFrame({"col1": [1, 2], "col2": ["a", "b"]})
166+
167+
# Mock IPython display functionality
168+
with patch('IPython.display.display') as mock_display, \
169+
patch('IPython.display.HTML') as mock_html:
170+
171+
# Call the method
172+
self.stackql_magic._display_with_csv_download(test_df)
173+
174+
# Verify display was called twice (once for DataFrame, once for HTML)
175+
assert mock_display.call_count == 2, "Display should be called twice"
176+
177+
# Verify HTML was called once
178+
mock_html.assert_called_once()
179+
180+
# Check that the HTML call contains download link
181+
html_call_args = mock_html.call_args[0][0]
182+
assert 'download="stackql_results.csv"' in html_call_args
183+
assert 'data:text/csv;base64,' in html_call_args
184+
185+
print_test_result("_display_with_csv_download method test",
186+
mock_display.call_count == 2 and mock_html.called,
187+
False, True)
188+
189+
def test_display_with_csv_download_error_handling(self):
190+
"""Test error handling in _display_with_csv_download method."""
191+
from unittest.mock import patch
192+
193+
# Create a mock DataFrame that will raise an exception during to_csv()
194+
mock_df = MagicMock()
195+
mock_df.to_csv.side_effect = Exception("Test CSV error")
196+
197+
# Mock IPython display functionality
198+
with patch('IPython.display.display') as mock_display, \
199+
patch('IPython.display.HTML') as mock_html, \
200+
patch('builtins.print') as mock_print:
201+
202+
# Call the method with the problematic DataFrame
203+
self.stackql_magic._display_with_csv_download(mock_df)
204+
205+
# Verify display was called once (for DataFrame only, not for HTML)
206+
mock_display.assert_called_once_with(mock_df)
207+
208+
# Verify HTML was not called due to error
209+
mock_html.assert_not_called()
210+
211+
# Verify error message was printed
212+
mock_print.assert_called_once()
213+
error_message = mock_print.call_args[0][0]
214+
assert "Error generating CSV download:" in error_message
215+
216+
print_test_result("_display_with_csv_download error handling test",
217+
mock_display.called and not mock_html.called and mock_print.called,
218+
False, True)
219+
109220
def test_magic_extension_loading(mock_interactive_shell):
110221
"""Test that non-server magic extension can be loaded."""
111222
# Test loading non-server magic

0 commit comments

Comments
 (0)