11# `%load_ext pystackql.magics` - loads the stackql magic with server_mode=True
2- from IPython .core .magic import magics_class
2+ from IPython .core .magic import ( magics_class , line_cell_magic )
33from .base_stackql_magic import BaseStackqlMagic
4+ import argparse
45
56@magics_class
67class StackqlServerMagic (BaseStackqlMagic ):
78 def __init__ (self , shell ):
89 super ().__init__ (shell , server_mode = True )
910
11+ @line_cell_magic
12+ def stackql (self , line , cell = None ):
13+ """A Jupyter magic command to run StackQL queries.
14+
15+ Can be used as both line and cell magic:
16+ - As a line magic: `%stackql QUERY`
17+ - As a cell magic: `%%stackql [OPTIONS]` followed by the QUERY in the next line.
18+
19+ :param line: The arguments and/or StackQL query when used as line magic.
20+ :param cell: The StackQL query when used as cell magic.
21+ :return: StackQL query results as a named Pandas DataFrame (`stackql_df`).
22+ """
23+ is_cell_magic = cell is not None
24+
25+ if is_cell_magic :
26+ parser = argparse .ArgumentParser ()
27+ parser .add_argument ("--no-display" , action = "store_true" , help = "Suppress result display." )
28+ args = parser .parse_args (line .split ())
29+ query_to_run = self .get_rendered_query (cell )
30+ else :
31+ args = None
32+ query_to_run = self .get_rendered_query (line )
33+
34+ results = self .run_query (query_to_run )
35+ self .shell .user_ns ['stackql_df' ] = results
36+
37+ if is_cell_magic and args and not args .no_display :
38+ return results
39+ elif not is_cell_magic :
40+ return results
41+
1042def load_ipython_extension (ipython ):
1143 """Load the extension in IPython."""
1244 ipython .register_magics (StackqlServerMagic )
0 commit comments