Moved the convert function away from pytree

jedi/parser/__init__.py

@@ -50,6 +50,34 @@ class Parser(object):
 
         if not source.endswith('\n'):
             source += '\n'
+
+        _ast_mapping = {
+            'expr_stmt': pr.ExprStmt,
+            'classdef': pr.Class,
+            'funcdef': pr.Function,
+            'file_input': pr.SubModule,
+            'import_name': pr.Import,
+            'import_from': pr.Import,
+            'break_stmt': pr.KeywordStatement,
+            'continue_stmt': pr.KeywordStatement,
+            'return_stmt': pr.ReturnStmt,
+            'raise_stmt': pr.KeywordStatement,
+            'yield_expr': pr.YieldExpr,
+            'del_stmt': pr.KeywordStatement,
+            'pass_stmt': pr.KeywordStatement,
+            'global_stmt': pr.GlobalStmt,
+            'nonlocal_stmt': pr.KeywordStatement,
+            'assert_stmt': pr.KeywordStatement,
+            'if_stmt': pr.IfStmt,
+            'with_stmt': pr.WithStmt,
+            'for_stmt': pr.ForStmt,
+            'while_stmt': pr.WhileStmt,
+            'try_stmt': pr.TryStmt,
+        }
+
+        self._ast_mapping = dict((getattr(pytree.python_symbols, k), v)
+                                 for k, v in _ast_mapping.items())
+
         self.global_names = []
         #if self.options["print_function"]:
         #    python_grammar = pygram.python_grammar_no_print_statement
@@ -68,7 +96,7 @@ class Parser(object):
         self.module.set_global_names(self.global_names)
 
     def convert_node(self, grammar, raw_node):
-        new_node = pytree.convert(grammar, raw_node)
+        new_node = self._convert(grammar, raw_node)
         # We need to check raw_node always, because the same node can be
         # returned by convert multiple times.
         if raw_node[0] == pytree.python_symbols.global_stmt:
@@ -88,9 +116,24 @@ class Parser(object):
             new_node.names_dict = scope_names
         return new_node
 
+    def _convert(self, grammar, raw_node):
+        """
+        Convert raw node information to a Node or Leaf instance.
+
+        This is passed to the parser driver which calls it whenever a reduction of a
+        grammar rule produces a new complete node, so that the tree is build
+        strictly bottom-up.
+        """
+        type, value, context, children = raw_node
+        #print(raw_node, pytree.type_repr(type))
+        try:
+            return self._ast_mapping[type](children)
+        except KeyError:
+            return pr.Node(type, children)
+
     def convert_leaf(self, grammar, raw_node):
         type, value, context, children = raw_node
-        #print('leaf', raw_node, type_repr(type))
+        #print('leaf', raw_node, pytree.type_repr(type))
         prefix, start_pos = context
         if type == tokenize.NAME:
             if value in grammar.keywords:

jedi/parser/pgen2/parse.py

@@ -179,6 +179,9 @@ class Parser(object):
         """Pop a nonterminal.  (Internal)"""
         popdfa, popstate, popnode = self.stack.pop()
         children = popnode[3]
+        # If there's exactly one child, return that child instead of creating a
+        # new node.  We still create expr_stmt and file_input though, because a
+        # lot of Jedi depends on its logic.
         if len(children) != 1 or popnode[0] in (self.grammar.symbol2number['expr_stmt'],
                                                 self.grammar.symbol2number['file_input']):
             newnode = self.convert_node(self.grammar, popnode)

jedi/parser/pytree.py

@@ -57,64 +57,3 @@ def type_repr(type_num):
     return _type_reprs.setdefault(type_num, type_num)
 
 
-def gen_ast_mapping():
-    from jedi.parser import representation as pr
-    _ast_mapping = {
-        'expr_stmt': pr.ExprStmt,
-        'classdef': pr.Class,
-        'funcdef': pr.Function,
-        'file_input': pr.SubModule,
-        'import_name': pr.Import,
-        'import_from': pr.Import,
-        'break_stmt': pr.KeywordStatement,
-        'continue_stmt': pr.KeywordStatement,
-        'return_stmt': pr.ReturnStmt,
-        'raise_stmt': pr.KeywordStatement,
-        'yield_expr': pr.YieldExpr,
-        'del_stmt': pr.KeywordStatement,
-        'pass_stmt': pr.KeywordStatement,
-        'global_stmt': pr.GlobalStmt,
-        'nonlocal_stmt': pr.KeywordStatement,
-        'assert_stmt': pr.KeywordStatement,
-        'if_stmt': pr.IfStmt,
-        'with_stmt': pr.WithStmt,
-        'for_stmt': pr.ForStmt,
-        'while_stmt': pr.WhileStmt,
-        'try_stmt': pr.TryStmt,
-    }
-
-    global ast_mapping, pr
-    ast_mapping = dict((getattr(python_symbols, k), v) for k, v in _ast_mapping.items())
-
-
-def convert(grammar, raw_node):
-    """
-    Convert raw node information to a Node or Leaf instance.
-
-    This is passed to the parser driver which calls it whenever a reduction of a
-    grammar rule produces a new complete node, so that the tree is build
-    strictly bottom-up.
-    """
-    try:
-        ast_mapping
-    except NameError:
-        gen_ast_mapping()
-
-    #import pdb; pdb.set_trace()
-    type, value, context, children = raw_node
-    if type in grammar.number2symbol:
-        # If there's exactly one child, return that child instead of
-        # creating a new node.
-        # We still create expr_stmt and file_input though, because a lot of
-        # Jedi depends on its logic.
-        if len(children) == 1 and type not in (python_symbols.expr_stmt,
-                                               python_symbols.file_input):
-            return children[0]
-        #print(raw_node, type_repr(type))
-        #import pdb; pdb.set_trace()
-        try:
-            return ast_mapping[type](children)
-        except KeyError:
-            return pr.Node(type, children)
-    else:
-        raise NotImplementedError