Skip to content

Commit 4e30e4c

Browse files
committed
Use a custom parser for cell magic arguments
1 parent 07d7651 commit 4e30e4c

File tree

2 files changed

+39
-14
lines changed

2 files changed

+39
-14
lines changed

google/cloud/bigquery/ipython_magics/magics.py

+37-13
Original file line numberDiff line numberDiff line change
@@ -65,13 +65,6 @@
6565
the variable name (ex. ``$my_dict_var``). See ``In[6]`` and ``In[7]``
6666
in the Examples section below.
6767
68-
.. note::
69-
70-
Due to the way IPython argument parser works, negative numbers in
71-
dictionaries are incorrectly "recognized" as additional arguments,
72-
resulting in an error ("unrecognized arguments"). To get around this,
73-
pass such dictionary as a JSON string variable.
74-
7568
* ``<query>`` (required, cell argument):
7669
SQL query to run. If the query does not contain any whitespace (aside
7770
from leading and trailing whitespace), it is assumed to represent a
@@ -159,13 +152,15 @@
159152
except ImportError: # pragma: NO COVER
160153
raise ImportError("This module can only be loaded in IPython.")
161154

155+
import six
156+
162157
from google.api_core import client_info
163158
from google.api_core.exceptions import NotFound
164159
import google.auth
165160
from google.cloud import bigquery
166161
import google.cloud.bigquery.dataset
167162
from google.cloud.bigquery.dbapi import _helpers
168-
import six
163+
from google.cloud.bigquery.ipython_magics import line_arg_parser as lap
169164

170165

171166
IPYTHON_USER_AGENT = "ipython-{}".format(IPython.__version__)
@@ -473,7 +468,11 @@ def _cell_magic(line, query):
473468
Returns:
474469
pandas.DataFrame: the query results.
475470
"""
476-
args = magic_arguments.parse_argstring(_cell_magic, line)
471+
# The built-in parser does not recognize Python structures such as dicts, thus
472+
# we extract the "--params" option and inteprpret it separately.
473+
params_option_value, rest_of_args = _split_args_line(line)
474+
475+
args = magic_arguments.parse_argstring(_cell_magic, rest_of_args)
477476

478477
if args.use_bqstorage_api is not None:
479478
warnings.warn(
@@ -484,11 +483,17 @@ def _cell_magic(line, query):
484483
use_bqstorage_api = not args.use_rest_api
485484

486485
params = []
487-
if args.params is not None:
488-
try:
489-
params = _helpers.to_query_parameters(
490-
ast.literal_eval("".join(args.params))
486+
if params_option_value:
487+
# A non-existing params variable is not expanded and ends up in the input
488+
# in its raw form, e.g. "$query_params".
489+
if params_option_value.startswith("$"):
490+
msg = 'Parameter expansion failed, undefined variable "{}".'.format(
491+
params_option_value[1:]
491492
)
493+
raise NameError(msg)
494+
495+
try:
496+
params = _helpers.to_query_parameters(ast.literal_eval(params_option_value))
492497
except Exception:
493498
raise SyntaxError(
494499
"--params is not a correctly formatted JSON string or a JSON "
@@ -598,6 +603,25 @@ def _cell_magic(line, query):
598603
close_transports()
599604

600605

606+
def _split_args_line(line):
607+
"""Split out the --params option value from the input line arguments.
608+
609+
Args:
610+
line (str): The line arguments passed to the cell magic.
611+
612+
Returns:
613+
Tuple[str, str]
614+
"""
615+
lexer = lap.Lexer(line)
616+
scanner = lap.Parser(lexer)
617+
tree = scanner.input_line()
618+
619+
extractor = lap.QueryParamsExtractor()
620+
params_option_value, rest_of_args = extractor.visit(tree)
621+
622+
return params_option_value, rest_of_args
623+
624+
601625
def _make_bqstorage_client(use_bqstorage_api, credentials):
602626
if not use_bqstorage_api:
603627
return None

tests/unit/test_magics.py

+2-1
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,7 @@
4343
from google.cloud import bigquery
4444
from google.cloud.bigquery import job
4545
from google.cloud.bigquery import table
46+
from google.cloud.bigquery.ipython_magics import line_arg_parser as lap
4647
from google.cloud.bigquery.ipython_magics import magics
4748
from tests.unit.helpers import make_connection
4849
from test_utils.imports import maybe_fail_import
@@ -1244,7 +1245,7 @@ def test_bigquery_magic_with_improperly_formatted_params():
12441245

12451246
sql = "SELECT @num AS num"
12461247

1247-
with pytest.raises(SyntaxError):
1248+
with pytest.raises(lap.exceptions.QueryParamsParseError):
12481249
ip.run_cell_magic("bigquery", "--params {17}", sql)
12491250

12501251

0 commit comments

Comments
 (0)