(WIP) add node visitors

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

@@ -16,6 +16,17 @@
 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.visitors import (
+    QueryParamsExtractor,
+    TreePrinter,
+)
 
 
-__all__ = ("Lexer", "ParseError", "Parser", "TokenType")
+__all__ = (
+    "Lexer",
+    "ParseError",
+    "Parser",
+    "QueryParamsExtractor",
+    "TokenType",
+    "TreePrinter",
+)

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

@@ -0,0 +1,240 @@
+# 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.
+
+from __future__ import print_function
+
+
+class NodeVisitor(object):
+    """Base visitor class implementing the dispatch machinery."""
+
+    def visit(self, node):
+        method_name = "visit_{}".format(type(node).__name__)
+        visitor_method = getattr(self, method_name, self.method_missing)
+        return visitor_method(node)
+
+    def method_missing(self, node):
+        raise Exception("No visit_{} method".format(type(node).__name__))
+
+
+class TreePrinter(NodeVisitor):
+    """Print the values of the parsed cell magic arguments.
+
+    Primarily useful for debugging.
+    """
+
+    def __init__(self):
+        self._indent = 0  # the current print indentation
+        self.INDENT_SIZE = 4
+
+    def visit_InputLine(self, node):
+        print("Input:")
+        self._indent += self.INDENT_SIZE
+
+        self.visit(node.destination_var)
+        self.visit(node.option_list)
+
+        self._indent -= self.INDENT_SIZE
+
+    def visit_DestinationVar(self, node):
+        print(" " * self._indent, end="")
+
+        var_name = node.name if node.name is not None else "<NOT_GIVEN>"
+        print("destination_var: ", var_name)
+
+    def visit_CmdOptionList(self, node):
+        print(" " * self._indent, end="")
+        print("Command options:")
+
+        self._indent += self.INDENT_SIZE
+
+        if not node.options:
+            print(" " * self._indent, "<NONE_GIVEN>", end="")
+        else:
+            for opt in node.options:
+                self.visit(opt)
+                print()
+
+        self._indent -= self.INDENT_SIZE
+
+    def visit_CmdOption(self, node):
+        print(" " * self._indent, end="")
+        print("--", node.name, " ", sep="", end="")
+
+        if node.value is not None:
+            self.visit(node.value)
+
+    def visit_CmdOptionValue(self, node):
+        print(node.value, end="")
+
+    def visit_ParamsOption(self, node):
+        print(" " * self._indent, end="")
+        print("--", node.name, " ", sep="", end="")
+
+        self.visit(node.value)
+
+    def visit_PyVarExpansion(self, node):
+        print(node.raw_value, end="")
+
+    def visit_PyDict(self, node):
+        print("{", end="")
+
+        for i, item in enumerate(node.items):
+            if i > 0:
+                print(", ", end="")
+            self.visit(item)
+
+        print("}", end="")
+
+    def visit_PyDictItem(self, node):
+        self.visit(node.key)
+        print(": ", end="")
+        self.visit(node.value)
+
+    def visit_PyDictKey(self, node):
+        print(node.key_value, end="")
+
+    def visit_PyScalarValue(self, node):
+        print(node.raw_value, end="")
+
+    def visit_PyTuple(self, node):
+        print("(", end="")
+
+        for i, item in enumerate(node.items):
+            if i > 0:
+                print(", ", end="")
+            self.visit(item)
+
+        print(")", end="")
+
+    def visit_PyList(self, node):
+        print("[", end="")
+
+        for i, item in enumerate(node.items):
+            if i > 0:
+                print(", ", end="")
+            self.visit(item)
+
+        print("]", end="")
+
+
+class QueryParamsExtractor(NodeVisitor):
+    """A visitor that extracts the "--params <...>" part from input line arguments.
+    """
+
+    def visit_InputLine(self, node):
+        params_dict_parts = []
+        other_parts = []
+
+        dest_var_parts = self.visit(node.destination_var)
+        params, other_options = self.visit(node.option_list)
+
+        if dest_var_parts:
+            other_parts.extend(dest_var_parts)
+
+        if dest_var_parts and other_options:
+            other_parts.append(" ")
+        other_parts.extend(other_options)
+
+        params_dict_parts.extend(params)
+
+        return "".join(params_dict_parts), "".join(other_parts)
+
+    def visit_DestinationVar(self, node):
+        return [node.name] if node.name is not None else []
+
+    def visit_CmdOptionList(self, node):
+        params_opt_parts = []
+        other_parts = []
+
+        # TODO: iterate directly and flatten the result?
+        # TODO: more genrally, use generators and only combine the result at end?
+
+        for i, opt in enumerate(node.options):
+            option_parts = self.visit(opt)
+            list_to_extend = params_opt_parts if opt.name == "params" else other_parts
+
+            if list_to_extend:
+                list_to_extend.append(" ")
+            list_to_extend.extend(option_parts)
+
+        return params_opt_parts, other_parts
+
+    def visit_CmdOption(self, node):
+        result = ["--{}".format(node.name)]
+
+        if node.value is not None:
+            result.append(" ")
+            value_parts = self.visit(node.value)
+            result.extend(value_parts)
+
+        return result
+
+    def visit_CmdOptionValue(self, node):
+        return [node.value]
+
+    def visit_ParamsOption(self, node):
+        value_parts = self.visit(node.value)
+        return value_parts
+
+    def visit_PyVarExpansion(self, node):
+        return [node.raw_value]
+
+    def visit_PyDict(self, node):
+        result = ["{"]
+
+        for i, item in enumerate(node.items):
+            if i > 0:
+                result.append(", ")
+            item_parts = self.visit(item)
+            result.extend(item_parts)
+
+        result.append("}")
+        return result
+
+    def visit_PyDictItem(self, node):
+        result = self.visit(node.key)  # key parts
+        result.append(": ")
+        value_parts = self.visit(node.value)
+        result.extend(value_parts)
+        return result
+
+    def visit_PyDictKey(self, node):
+        return [node.key_value]  # TODO: does this contain quotes? it should...
+
+    def visit_PyScalarValue(self, node):
+        return [node.raw_value]
+
+    def visit_PyTuple(self, node):
+        result = ["("]
+
+        for i, item in enumerate(node.items):
+            if i > 0:
+                result.append(", ")
+            item_parts = self.visit(item)
+            result.extend(item_parts)
+
+        result.append(")")
+        return result
+
+    def visit_PyList(self, node):
+        result = ["["]
+
+        for i, item in enumerate(node.items):
+            if i > 0:
+                result.append(", ")
+            item_parts = self.visit(item)
+            result.extend(item_parts)
+
+        result.append("]")
+        return result