[R] de-monkeypatching set_line_info method and introduce StmtMixIn

4 files changed, 80 insertions, 93 deletions

diff --git a/inference.py b/inference.py
index 5bb3e77a..5168c3ea 100644
--- a/inference.py
+++ b/inference.py
@@ -245,10 +245,10 @@ def infer_import(self, context=None, asname=True):
         yield self.do_import_module(name)
 nodes.Import.infer = path_wrapper(infer_import)
 
-def infer_name_module(node, name):
-    context = InferenceContext(node)
+def infer_name_module(self, name):
+    context = InferenceContext(self)
     context.lookupname = name
-    return node.infer(context, asname=False)
+    return self.infer(context, asname=False)
 nodes.Import.infer_name_module = infer_name_module
 
 
diff --git a/lookup.py b/lookup.py
index 2f0576e4..1bdcaf86 100644
--- a/lookup.py
+++ b/lookup.py
@@ -189,8 +189,3 @@ def builtin_lookup(name):
     return builtinastng, stmts
 
 
-def decorators_scope(self):
-    # skip the function node to go directly to the upper level scope
-    return self.parent.parent.scope()
-nodes.Decorators.scope = decorators_scope
-
diff --git a/node_classes.py b/node_classes.py
index 361d2f0f..a445ee24 100644
--- a/node_classes.py
+++ b/node_classes.py
@@ -53,6 +53,11 @@ class ClassNG(object):# (Class, NodeNG)
 class CompareNG(object):# (Compare, StmtMixIn, NodeNG)
     """class representing a Compare node"""
 
+    def get_children(self):
+        """override get_children for tuple fields"""
+        yield self.left
+        for _, comparator in self.ops:
+            yield comparator # we don't want the 'op'
 
 class ComprehensionNG(object):# (Comprehension, StmtMixIn, NodeNG)
     """class representing a Comprehension node"""
@@ -69,6 +74,9 @@ class ContinueNG(object):# (Continue, NodeNG)
 class DecoratorsNG(object):# (Decorators, StmtMixIn, NodeNG)
     """class representing a Decorators node"""
 
+    def scope(self):
+        # skip the function node to go directly to the upper level scope
+        return self.parent.parent.scope()
 
 class DelAttrNG(object):# (DelAttr, StmtMixIn, NodeNG)
     """class representing a DelAttr node"""
@@ -85,6 +93,13 @@ class DeleteNG(object):# (Delete, NodeNG)
 class DictNG(object):# (Dict, StmtMixIn, NodeNG)
     """class representing a Dict node"""
 
+    def get_children(self):
+        """get children of a Dict node"""
+        # overrides get_children
+        for key, value in self.items:
+            yield key
+            yield value
+
 
 class DiscardNG(object):# (Discard, NodeNG)
     """class representing a Discard node"""
@@ -101,6 +116,14 @@ class EmptyNodeNG(object):# (EmptyNode, StmtMixIn, NodeNG)
 class ExceptHandlerNG(object):# (ExceptHandler, NodeNG)
     """class representing an ExceptHandler node"""
 
+    def _blockstart_toline(self):
+        if self.name:
+            return self.name.tolineno
+        elif self.type:
+            return self.type.tolineno
+        else:
+            return self.lineno
+
 
 class ExecNG(object):# (Exec, NodeNG)
     """class representing an Exec node"""
@@ -113,6 +136,9 @@ class ExtSliceNG(object):# (ExtSlice, StmtMixIn, NodeNG)
 class ForNG(object):# (For, NodeNG)
     """class representing a For node"""
 
+    def _blockstart_toline(self):
+        return self.iter.tolineno
+
 
 class FromNG(object):# (From, NodeNG)
     """class representing a From node"""
@@ -137,6 +163,10 @@ class GlobalNG(object):# (Global, NodeNG)
 class IfNG(object):# (If, NodeNG)
     """class representing an If node"""
 
+    def _blockstart_toline(self):
+        return self.test.tolineno
+
+
 
 class IfExpNG(object):# (IfExp, StmtMixIn, NodeNG)
     """class representing an IfExp node"""
@@ -169,6 +199,8 @@ class ListCompNG(object):# (ListComp, StmtMixIn, NodeNG)
 class ModuleNG(object):# (Module, StmtMixIn, NodeNG)
     """class representing a Module node"""
 
+    is_statement = False # override StatementMixin
+
 
 class NameNG(object):# (Name, StmtMixIn, NodeNG)
     """class representing a Name node"""
@@ -201,10 +233,16 @@ class SubscriptNG(object):# (Subscript, StmtMixIn, NodeNG)
 class TryExceptNG(object):# (TryExcept, NodeNG)
     """class representing a TryExcept node"""
 
+    def _blockstart_toline(self):
+        return self.lineno
+
 
 class TryFinallyNG(object):# (TryFinally, NodeNG)
     """class representing a TryFinally node"""
 
+    def _blockstart_toline(self):
+        return self.lineno
+
 
 class TupleNG(object):# (Tuple, StmtMixIn, NodeNG)
     """class representing a Tuple node"""
@@ -217,10 +255,19 @@ class UnaryOpNG(object):# (UnaryOp, StmtMixIn, NodeNG)
 class WhileNG(object):# (While, NodeNG)
     """class representing a While node"""
 
+    def _blockstart_toline(self):
+        return self.test.tolineno
+
 
 class WithNG(object):# (With, NodeNG)
     """class representing a With node"""
 
+    def _blockstart_toline(self):
+        if self.vars:
+            return self.vars.tolineno
+        else:
+            return self.expr.tolineno
+
 
 class YieldNG(object):# (Yield, NodeNG)
     """class representing a Yield node"""
diff --git a/nodes.py b/nodes.py
index adcbb35a..12c29739 100644
--- a/nodes.py
+++ b/nodes.py
@@ -369,17 +369,25 @@ def _repr_tree(node, result, indent='', _done=None):
             result.append(  indent + field + " = " )
             _repr_tree(value, result, indent, _done)
 
+class StmtMixIn(object):
+    """StmtMixIn used only for a adding a few attributes"""
+    is_statement = True
 
-def replace_child(self, child, newchild):
-    sequence = self.child_sequence(child)
-    newchild.parent = self
-    child.parent = None
-    sequence[sequence.index(child)] = newchild
+    def replace(self, child, newchild):
+        sequence = self.child_sequence(child)
+        newchild.parent = self
+        child.parent = None
+        sequence[sequence.index(child)] = newchild
+
+
+for cls in ALL_NODES:
+    cls_name = REDIRECT.get(cls.__name__, cls.__name__) + "NG"
+    if cls in STMT_NODES:
+        addons = (NodeNG, StmtMixIn, getattr(node_classes, cls_name))
+    else:
+        addons = (NodeNG, getattr(node_classes, cls_name))
+    extend_class(cls, addons)
 
-for klass in STMT_NODES:
-    klass.is_statement = True
-    klass.replace = replace_child
-Module.replace = replace_child
 
 CONST_CLS = {
     list: List,
@@ -387,74 +395,19 @@ CONST_CLS = {
     dict: Dict,
     }
 
-def const_factory(value):
-    """return an astng node for a python value"""
-    try:
-        # if value is of class list, tuple, dict use specific class, not Const
-        cls = CONST_CLS[value.__class__]
-        node = cls()
-        if isinstance(node, Dict):
-            node.items = ()
-        else:
-            node.elts = ()
-    except KeyError:
-        try:
-            node = Const(value)
-        except KeyError:
-            node = _const_factory(value)
-    return node
-
-def _get_children_nochildren(self):
-    return ()
-
-#  get_children overrides  ####################################################
-
-def _dict_get_children(node): # XXX : distinguish key and value ?
-    """override get_children for Dict"""
-    for key, value in node.items:
-        yield key
-        yield value
-Dict.get_children = _dict_get_children
+# block range overrides #######################################################
 
 
-def _compare_get_children(node):
-    """override get_children for tuple fields"""
-    yield node.left
-    for _, comparator in node.ops:
-        yield comparator # we don't want the 'op'
-Compare.get_children = _compare_get_children
+class BlockRangeMixIn(object):
+    """override block range """
+    def set_line_info(self, lastchild):
+        self.fromlineno = self.lineno
+        self.tolineno = lastchild.tolineno
+        self.blockstart_tolineno = self._blockstart_toline()
 
-# block range overrides #######################################################
+for kls in ExceptHandler, For, If, TryExcept, TryFinally, While, With:
+    extend_class(kls, (BlockRangeMixIn,))
 
-def for_set_line_info(self, lastchild):
-    self.fromlineno = self.lineno
-    self.tolineno = lastchild.tolineno
-    self.blockstart_tolineno = self.iter.tolineno
-For.set_line_info = for_set_line_info
-
-def if_set_line_info(self, lastchild):
-    self.fromlineno = self.lineno
-    self.tolineno = lastchild.tolineno
-    self.blockstart_tolineno = self.test.tolineno
-If.set_line_info = if_set_line_info
-While.set_line_info = if_set_line_info
-
-def try_set_line_info(self, lastchild):
-    self.fromlineno = self.blockstart_tolineno = self.lineno
-    self.tolineno = lastchild.tolineno
-TryExcept.set_line_info = try_set_line_info
-TryFinally.set_line_info = try_set_line_info
-
-def excepthandler_set_line_info(self, lastchild):
-    self.fromlineno = self.lineno
-    if self.name:
-        self.blockstart_tolineno= self.name.tolineno
-    elif self.type:
-        self.blockstart_tolineno= self.type.tolineno
-    else:
-        self.blockstart_tolineno= self.lineno
-    self.tolineno = lastchild.tolineno
-ExceptHandler.set_line_info = excepthandler_set_line_info
 
 def excepthandler_catch(self, exceptions):
     if self.type is None or exceptions is None:
@@ -464,15 +417,6 @@ def excepthandler_catch(self, exceptions):
             return True
 ExceptHandler.catch = excepthandler_catch
 
-def with_set_line_info(self, lastchild):
-    self.fromlineno = self.blockstart_tolineno = self.lineno
-    self.tolineno = lastchild.tolineno
-    if self.vars:
-        self.blockstart_tolineno = self.vars.tolineno
-    else:
-        self.blockstart_tolineno = self.expr.tolineno
-With.set_line_info = with_set_line_info
-
 
 def object_block_range(node, lineno):
     """handle block line numbers range for function/class statements:
@@ -542,10 +486,10 @@ TryFinally.block_range = try_finalbody_block_range
 
 # From and Import #############################################################
 
-def real_name(node, asname):
+def real_name(self, asname):
     """get name from 'as' name"""
-    for index in range(len(node.names)):
-        name, _asname = node.names[index]
+    for index in range(len(self.names)):
+        name, _asname = self.names[index]
         if name == '*':
             return asname
         if not _asname:
@@ -554,6 +498,7 @@ def real_name(node, asname):
         if asname == _asname:
             return name
     raise NotFoundError(asname)
+
 From.real_name = real_name
 Import.real_name = real_name