Skip to content

Commit 5d1155d

Browse files
jeffreyavenjeffreyaven
committed
output formatting
1 parent 06f9c17 commit 5d1155d

File tree

4 files changed

+176
-16
lines changed

4 files changed

+176
-16
lines changed

‎pystackql/magic_ext/base.py‎

Lines changed: 54 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -52,10 +52,59 @@ def run_query(self, query):
5252

5353
return self.stackql_instance.execute(query)
5454

55+
# def _display_with_csv_download(self, df):
56+
# """Display DataFrame with CSV download link.
57+
58+
# :param df: The DataFrame to display and make downloadable.
59+
# """
60+
# import IPython.display
61+
62+
# try:
63+
# # Generate CSV data
64+
# import io
65+
# import base64
66+
# csv_buffer = io.StringIO()
67+
# df.to_csv(csv_buffer, index=False)
68+
# csv_data = csv_buffer.getvalue()
69+
70+
# # Encode to base64 for data URI
71+
# csv_base64 = base64.b64encode(csv_data.encode()).decode()
72+
73+
# # Create download link
74+
# download_link = f'data:text/csv;base64,{csv_base64}'
75+
76+
# # # Display the DataFrame first
77+
# # IPython.display.display(df)
78+
79+
# # Create and display the download button
80+
# download_html = f'''
81+
# <div style="margin-top: 15px; font-family: system-ui, -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, Oxygen, Ubuntu, Cantarell, 'Open Sans', sans-serif;">
82+
# <a href="{download_link}" download="stackql_results.csv"
83+
# style="display: inline-flex; align-items: center; gap: 8px; padding: 9px 16px;
84+
# background-color: #2196F3; color: white; text-decoration: none;
85+
# border-radius: 4px; font-size: 14px; font-weight: 500;
86+
# box-shadow: 0 2px 4px rgba(0,0,0,0.08); transition: all 0.2s ease;">
87+
# <svg xmlns="http://www.w3.org/2000/svg" width="16" height="16" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2.5" stroke-linecap="round" stroke-linejoin="round">
88+
# <path d="M21 15v4a2 2 0 0 1-2 2H5a2 2 0 0 1-2-2v-4"></path>
89+
# <polyline points="7 10 12 15 17 10"></polyline>
90+
# <line x1="12" y1="15" x2="12" y2="3"></line>
91+
# </svg>
92+
# Download CSV
93+
# </a>
94+
# </div>
95+
# '''
96+
97+
# IPython.display.display(IPython.display.HTML(download_html))
98+
99+
# except Exception as e:
100+
# # If CSV generation fails, just display the DataFrame normally
101+
# IPython.display.display(df)
102+
# print(f"Error generating CSV download: {e}")
103+
55104
def _display_with_csv_download(self, df):
56-
"""Display DataFrame with CSV download link.
105+
"""Display a CSV download link for the DataFrame without displaying the DataFrame again.
57106
58-
:param df: The DataFrame to display and make downloadable.
107+
:param df: The DataFrame to make downloadable.
59108
"""
60109
import IPython.display
61110

@@ -73,10 +122,7 @@ def _display_with_csv_download(self, df):
73122
# Create download link
74123
download_link = f'data:text/csv;base64,{csv_base64}'
75124

76-
# # Display the DataFrame first
77-
# IPython.display.display(df)
78-
79-
# Create and display the download button
125+
# Only display the download button, not the DataFrame
80126
download_html = f'''
81127
<div style="margin-top: 15px; font-family: system-ui, -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, Oxygen, Ubuntu, Cantarell, 'Open Sans', sans-serif;">
82128
<a href="{download_link}" download="stackql_results.csv"
@@ -92,11 +138,9 @@ def _display_with_csv_download(self, df):
92138
Download CSV
93139
</a>
94140
</div>
95-
'''
96-
141+
'''
97142
IPython.display.display(IPython.display.HTML(download_html))
98143

99144
except Exception as e:
100-
# If CSV generation fails, just display the DataFrame normally
101-
IPython.display.display(df)
145+
# If CSV generation fails, just print an error message without displaying anything
102146
print(f"Error generating CSV download: {e}")

‎pystackql/magic_ext/local.py‎

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,43 @@ def __init__(self, shell):
2222
"""
2323
super().__init__(shell, server_mode=False)
2424

25+
# @line_cell_magic
26+
# def stackql(self, line, cell=None):
27+
# """A Jupyter magic command to run StackQL queries.
28+
29+
# Can be used as both line and cell magic:
30+
# - As a line magic: `%stackql QUERY`
31+
# - As a cell magic: `%%stackql [OPTIONS]` followed by the QUERY in the next line.
32+
33+
# :param line: The arguments and/or StackQL query when used as line magic.
34+
# :param cell: The StackQL query when used as cell magic.
35+
# :return: StackQL query results as a named Pandas DataFrame (`stackql_df`).
36+
# """
37+
# is_cell_magic = cell is not None
38+
39+
# if is_cell_magic:
40+
# parser = argparse.ArgumentParser()
41+
# parser.add_argument("--no-display", action="store_true", help="Suppress result display.")
42+
# parser.add_argument("--csv-download", action="store_true", help="Add CSV download link to output.")
43+
# args = parser.parse_args(line.split())
44+
# query_to_run = self.get_rendered_query(cell)
45+
# else:
46+
# args = None
47+
# query_to_run = self.get_rendered_query(line)
48+
49+
# results = self.run_query(query_to_run)
50+
# self.shell.user_ns['stackql_df'] = results
51+
52+
# if is_cell_magic and args and args.no_display:
53+
# return None
54+
# elif is_cell_magic and args and args.csv_download and not args.no_display:
55+
# self._display_with_csv_download(results)
56+
# return results
57+
# elif is_cell_magic and args and not args.no_display:
58+
# return results
59+
# elif not is_cell_magic:
60+
# return results
61+
2562
@line_cell_magic
2663
def stackql(self, line, cell=None):
2764
"""A Jupyter magic command to run StackQL queries.
@@ -52,12 +89,18 @@ def stackql(self, line, cell=None):
5289
if is_cell_magic and args and args.no_display:
5390
return None
5491
elif is_cell_magic and args and args.csv_download and not args.no_display:
92+
# First display the DataFrame
93+
import IPython.display
94+
IPython.display.display(results)
95+
# Then add the download button without displaying the DataFrame again
5596
self._display_with_csv_download(results)
5697
return results
5798
elif is_cell_magic and args and not args.no_display:
5899
return results
59100
elif not is_cell_magic:
60101
return results
102+
else:
103+
return results
61104

62105
def load_ipython_extension(ipython):
63106
"""Load the non-server magic in IPython.

‎pystackql/magic_ext/server.py‎

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,45 @@ def __init__(self, shell):
2222
"""
2323
super().__init__(shell, server_mode=True)
2424

25+
# @line_cell_magic
26+
# def stackql(self, line, cell=None):
27+
# """A Jupyter magic command to run StackQL queries.
28+
29+
# Can be used as both line and cell magic:
30+
# - As a line magic: `%stackql QUERY`
31+
# - As a cell magic: `%%stackql [OPTIONS]` followed by the QUERY in the next line.
32+
33+
# :param line: The arguments and/or StackQL query when used as line magic.
34+
# :param cell: The StackQL query when used as cell magic.
35+
# :return: StackQL query results as a named Pandas DataFrame (`stackql_df`).
36+
# """
37+
# is_cell_magic = cell is not None
38+
39+
# if is_cell_magic:
40+
# parser = argparse.ArgumentParser()
41+
# parser.add_argument("--no-display", action="store_true", help="Suppress result display.")
42+
# parser.add_argument("--csv-download", action="store_true", help="Add CSV download link to output.")
43+
# args = parser.parse_args(line.split())
44+
# query_to_run = self.get_rendered_query(cell)
45+
# else:
46+
# args = None
47+
# query_to_run = self.get_rendered_query(line)
48+
49+
# results = self.run_query(query_to_run)
50+
# self.shell.user_ns['stackql_df'] = results
51+
52+
# if is_cell_magic and args and args.no_display:
53+
# return None
54+
# elif is_cell_magic and args and args.csv_download and not args.no_display:
55+
# self._display_with_csv_download(results)
56+
# return results
57+
# elif is_cell_magic and args and not args.no_display:
58+
# return results
59+
# elif not is_cell_magic:
60+
# return results
61+
# else:
62+
# return results
63+
2564
@line_cell_magic
2665
def stackql(self, line, cell=None):
2766
"""A Jupyter magic command to run StackQL queries.
@@ -52,6 +91,10 @@ def stackql(self, line, cell=None):
5291
if is_cell_magic and args and args.no_display:
5392
return None
5493
elif is_cell_magic and args and args.csv_download and not args.no_display:
94+
# First display the DataFrame
95+
import IPython.display
96+
IPython.display.display(results)
97+
# Then add the download button without displaying the DataFrame again
5598
self._display_with_csv_download(results)
5699
return results
57100
elif is_cell_magic and args and not args.no_display:

‎tests/test_magic_base.py‎

Lines changed: 36 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -190,6 +190,36 @@ def test_display_with_csv_download_method(self):
190190
mock_display.call_count == 1 and mock_html.called,
191191
self.is_server_mode, True)
192192

193+
# def test_display_with_csv_download_error_handling(self):
194+
# """Test error handling in _display_with_csv_download method."""
195+
196+
# # Create a mock DataFrame that will raise an exception during to_csv()
197+
# mock_df = MagicMock()
198+
# mock_df.to_csv.side_effect = Exception("Test CSV error")
199+
200+
# # Mock IPython display functionality
201+
# with patch('IPython.display.display') as mock_display, \
202+
# patch('IPython.display.HTML') as mock_html, \
203+
# patch('builtins.print') as mock_print:
204+
205+
# # Call the method with the problematic DataFrame
206+
# self.stackql_magic._display_with_csv_download(mock_df)
207+
208+
# # Verify display was not called (we now only print an error message)
209+
# mock_display.assert_not_called()
210+
211+
# # Verify HTML was not called due to error
212+
# mock_html.assert_not_called()
213+
214+
# # Verify error message was printed
215+
# mock_print.assert_called_once()
216+
# error_message = mock_print.call_args[0][0]
217+
# assert "Error generating CSV download:" in error_message
218+
219+
# print_test_result(f"_display_with_csv_download error handling test{' (server mode)' if self.is_server_mode else ''}",
220+
# not mock_display.called and not mock_html.called and mock_print.called,
221+
# self.is_server_mode, True)
222+
193223
def test_display_with_csv_download_error_handling(self):
194224
"""Test error handling in _display_with_csv_download method."""
195225

@@ -199,16 +229,16 @@ def test_display_with_csv_download_error_handling(self):
199229

200230
# Mock IPython display functionality
201231
with patch('IPython.display.display') as mock_display, \
202-
patch('IPython.display.HTML') as mock_html, \
203-
patch('builtins.print') as mock_print:
232+
patch('IPython.display.HTML') as mock_html, \
233+
patch('builtins.print') as mock_print:
204234

205235
# Call the method with the problematic DataFrame
206236
self.stackql_magic._display_with_csv_download(mock_df)
207237

208-
# Verify display was not called (we now only print an error message)
238+
# Verify display was not called in the error case
209239
mock_display.assert_not_called()
210240

211-
# Verify HTML was not called due to error
241+
# Verify HTML was not called in the error case
212242
mock_html.assert_not_called()
213243

214244
# Verify error message was printed
@@ -217,5 +247,5 @@ def test_display_with_csv_download_error_handling(self):
217247
assert "Error generating CSV download:" in error_message
218248

219249
print_test_result(f"_display_with_csv_download error handling test{' (server mode)' if self.is_server_mode else ''}",
220-
not mock_display.called and not mock_html.called and mock_print.called,
221-
self.is_server_mode, True)
250+
not mock_display.called and not mock_html.called and mock_print.called,
251+
self.is_server_mode, True)

0 commit comments

Comments
 (0)