Commits

Anonymous committed df09a22

Improved support for empty super() - looks up base class

  • Participants
  • Parent commits 7f336b5

Comments (0)

Files changed (2)

lib3to2/fixes/fix_super.py

             # Probably a '*'
             return None
 
-def insert_args(name, rparen):
+
+def get_class_name(super_node):
+    parent = super_node.parent
+    while parent.type != syms.classdef and parent.parent:
+        parent = parent.parent
+
+    if parent.type != syms.classdef:
+        # super() called without arguments outside of a classdef
+        return None
+
+    children = parent.children
+    assert len(children) > 2
+    class_name = children[1]
+    assert class_name.type == token.NAME
+    return class_name.value
+
+
+def insert_args(name, class_name, rparen):
     parent = rparen.parent
+    
+    if class_name:
+        class_node = Node(syms.power, [Name(class_name)])
+    else:
+        class_node = Node(syms.power, [Name(name), dot_class.clone()])
+
     idx = parent.children.index(rparen)
-    parent.insert_child(idx, Name(name, prefix=" "))
+    parent.insert_child(idx, Name(name, prefix=u" "))
     parent.insert_child(idx, Comma())
-    parent.insert_child(idx, Node(syms.power, [Name(name), dot_class.clone()]))
+    parent.insert_child(idx, class_node)
 
 
 class FixSuper(fixer_base.BaseFix):
 
-    PATTERN = "power< 'super' trailer< '(' rparen=')' > any* >"
+    PATTERN = u"power< 'super' trailer< '(' rparen=')' > any* >"
 
     def transform(self, node, results):
         param = get_firstparam(node)
         if param is None:
-            self.cannot_convert(node, "super() with no arguments must be called inside a function that has at least one parameter")
+            self.cannot_convert(node, u"super() with no arguments must be called inside a function that has at least one parameter")
             return
-        rparen = results["rparen"]
-        insert_args(param, rparen)
+        class_name = get_class_name(node)
+        rparen = results[u"rparen"]
+        insert_args(param, class_name, rparen)

lib3to2/tests/test_super.py

 
         s = "def m():\n    super()"
         self.unchanged(s, ignore_warnings=True)
+
+    def test_class_noargs(self):
+
+        b = "class c:\n    def m(self):\n        super()"
+        a = "class c:\n    def m(self):\n        super(c, self)"
+        self.check(b, a)
+
+    def test_class_other_params(self):
+
+        b = "class c:\n    def m(a, self=None):\n        super()"
+        a = "class c:\n    def m(a, self=None):\n        super(c, a)"
+        self.check(b, a)
+
+    def test_class_no_with_stars(self):
+
+        s = "class c:\n    def m(*args, **kwargs):\n        super()"
+        self.unchanged(s, ignore_warnings=True)
+
+    def test_class_no_with_noargs(self):
+
+        s = "class c:\n    def m():\n        super()"
+        self.unchanged(s, ignore_warnings=True)