google/cloud/bigquery/ipython_magics/line_arg_parser/__init__.py

@@ -12,11 +12,21 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from google.cloud.bigquery.ipython_magics.line_arg_parser.exceptions import ParseError
 from google.cloud.bigquery.ipython_magics.line_arg_parser.lexer import Lexer
+from google.cloud.bigquery.ipython_magics.line_arg_parser.lexer import TokenType
 from google.cloud.bigquery.ipython_magics.line_arg_parser.parser import Parser
-from google.cloud.bigquery.ipython_magics.line_arg_parser import (
-    visitors,
-)  # TODO: import all
+from google.cloud.bigquery.ipython_magics.line_arg_parser.visitors import (
+    QueryParamsExtractor,
+    TreePrinter,
+)
 
 
-__all__ = ("Lexer", "Parser", "visitors")
+__all__ = (
+    "Lexer",
+    "ParseError",
+    "Parser",
+    "QueryParamsExtractor",
+    "TokenType",
+    "TreePrinter",
+)

google/cloud/bigquery/ipython_magics/line_arg_parser/exceptions.py

@@ -0,0 +1,17 @@
+# Copyright 2020 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+class ParseError(Exception):
+    pass

google/cloud/bigquery/ipython_magics/line_arg_parser/lexer.py

@@ -32,8 +32,9 @@
     state_2=OrderedDict(
         GOTO_STATE_3=r"(?P<GOTO_STATE_3>(?=--params(?=\s|$)))",  # the --params option
         OPTION_SPEC=r"(?P<OPTION_SPEC>--\w+)",
-        # NOTE: currently the only valid value for a non "--params" option is project ID
-        OPT_VAL=r"(?P<OPT_VAL>[^_\d\W](?:\w|\.)+)",
+        # NOTE: currently the only valid value for a non "--params" option is
+        # either a project ID or an integer (e.g. max_results)
+        OPT_VAL=r"(?P<OPT_VAL>\d+|[^_\d\W](?:\w|\.)+)",
     ),
     state_3=OrderedDict(
         PY_STRING=r"(?P<PY_STRING>(?:{})|(?:{}))".format(
@@ -43,9 +44,7 @@
         GOTO_STATE_2=r"(?P<GOTO_STATE_2>(?=--\w+))",  # found another option spec
         PY_BOOL=r"(?P<PY_BOOL>True|False)",
         DOLLAR_PY_ID=r"(?P<DOLLAR_PY_ID>\$[^\d\W]\w*)",
-        PY_ID=r"(?P<PY_ID>[^\d\W]\w*)",
-        # TODO: supporting only ints or floats, add floats in scientific notation, too?
-        PY_NUMBER=r"(?P<PY_NUMBER>-?[1-9]\d*(?:\.\d+)?)",
+        PY_NUMBER=r"(?P<PY_NUMBER>-?[1-9]\d*(?:\.\d+)?(:?[e|E][+-]?\d+)?)",
         SQUOTE=r"(?P<SQUOTE>')",
         DQUOTE=r'(?P<DQUOTE>")',
         COLON=r"(?P<COLON>:)",
@@ -61,7 +60,7 @@
         WS=r"(?P<WS>\s+)",
         EOL=r"(?P<EOL>$)",
         UNKNOWN=r"(?P<UNKNOWN>\S+)",  # anything not a whitespace or matched by something else
-    )
+    ),
 )
 
 
@@ -73,9 +72,10 @@ def _generate_next_value_(name, start, count, last_values):
 TokenType = AutoStrEnum(
     "TokenType",
     [
-        name for name in itertools.chain.from_iterable(token_types.values())
+        name
+        for name in itertools.chain.from_iterable(token_types.values())
         if not name.startswith("GOTO_STATE")
-    ]
+    ],
 )
 
 
@@ -115,18 +115,14 @@ class Lexer(object):
 
     def __init__(self, input_text):
         self._text = input_text
-        self._state_handlers = {
-            LexerState.STATE_1: self._state_1,
-            LexerState.STATE_2: self._state_2,
-            LexerState.STATE_3: self._state_3,
-        }
 
     def __iter__(self):
         # Since re.scanner does not seem to support manipulating inner scanner states,
         # we need to implement lexer state transitions manually using special
         # non-capturing lookahead token patterns to signal when a state transition
         # should be made.
-        # Each state is then processed by a dedicated state handler method.
+        # Since we don't have "nested" states, we don't really need a stack and
+        # this simple mechanism is sufficient.
         state = LexerState.STATE_1
         offset = 0  # the number of characters processed so far
 
@@ -147,43 +143,36 @@ def __iter__(self):
                     break
 
     def _get_state_token_generator(self, state, current_offset):
-        """TODO: explain... we need to create the canner and pick the state handler
-        and return that
+        """Return token generator for the current state starting at ``current_offset``.
+
+        Args:
+            state (LexerState): The current lexer state.
+            current_offset (int): The offset in the input text, i.e. the number
+                of characters already scanned so far.
+
+        Returns:
+            The next ``Token`` or ``StateTransition`` instance.
         """
-        state_handler = self._state_handlers[state]
         pattern = self._GRAND_PATTERNS[state]
         scanner = pattern.scanner(self._text, pos=current_offset)
-        return state_handler(scanner)
-
-    def _state_1(self, scanner):
-        for match in iter(scanner.match, None):
-            token_type = match.lastgroup
-
-            if token_type == "GOTO_STATE_2":
-                yield StateTransition(
-                    new_state=LexerState.STATE_2, total_offset=match.start(),
-                )
+        return self._scan_for_tokens(scanner)
 
-            yield Token(token_type, match.group(), match.start())
+    def _scan_for_tokens(self, scanner):
+        """Yield tokens produced by the scanner or state transition objects.
 
-    def _state_2(self, scanner):
-        for match in iter(scanner.match, None):
-            token_type = match.lastgroup
+        Args:
+            scanner (SRE_Scanner): The text tokenizer.
 
-            if token_type == "GOTO_STATE_3":
-                yield StateTransition(
-                    new_state=LexerState.STATE_3, total_offset=match.start(),
-                )
-
-            yield Token(token_type, match.group(), match.start())
-
-    def _state_3(self, scanner):
+        Yields:
+            The next ``Token`` or ``StateTransition`` instance.
+        """
         for match in iter(scanner.match, None):
             token_type = match.lastgroup
 
-            if token_type == "GOTO_STATE_2":
+            if token_type.startswith("GOTO_STATE"):
                 yield StateTransition(
-                    new_state=LexerState.STATE_2, total_offset=match.start(),
+                    new_state=getattr(LexerState, token_type[5:]),  # w/o "GOTO_" prefix
+                    total_offset=match.start(),
                 )
 
             yield Token(token_type, match.group(), match.start())