Fix remove_class

pytest-dev · Nov 28, 2020 · 0ef2386 · 0ef2386
1 parent 39cdfaa
commit 0ef2386
Show file tree

Hide file tree

Showing 3 changed files with 61 additions and 13 deletions.
diff --git a/tests/fixtures/remove_class/assertEqual_in.py b/tests/fixtures/remove_class/assertEqual_in.py
@@ -6,4 +6,22 @@ def test_you(self):
 class TestAssertNotEqual(TestCase):
 
     def test_you(self):
-        self.assertNotEqual(abc, 'xxx')
+        """
+        hello
+        """
+        pass
+        if True:
+            if False:
+                pass
+
+    @property
+    def test_me(self):
+        [
+            item
+            for item in sequence
+            if item
+        ]
+        (
+            1,
+            2,
+        )
diff --git a/tests/fixtures/remove_class/assertEqual_out.py b/tests/fixtures/remove_class/assertEqual_out.py
@@ -4,4 +4,22 @@ def test_you(self):
 
 
 def test_you(self):
-    self.assertNotEqual(abc, 'xxx')
+    """
+        hello
+        """
+    pass
+    if True:
+        if False:
+            pass
+
+@property
+def test_me(self):
+    [
+        item
+        for item in sequence
+        if item
+    ]
+    (
+        1,
+        2,
+    )
diff --git a/unittest2pytest/fixes/fix_remove_class.py b/unittest2pytest/fixes/fix_remove_class.py
@@ -57,17 +57,29 @@ class FixRemoveClass(BaseFix):
     """
 
     def dedent(self, suite, dedent):
-        self.line_num = suite.get_lineno()
-        for kid in suite.leaves():
-            if kid.type in (token.INDENT, token.DEDENT):
-                self.line_num = kid.get_lineno()
-                # todo: handle tabs
-                kid.value = kid.value[dedent:]
-                self.current_indent = kid.value
-            elif kid.get_lineno() != self.line_num:
-                # todo: handle tabs
-                if len(kid.prefix) > len(self.current_indent):
-                    kid.prefix = self.current_indent
+        def safe_dedent(s):
+            for i, c in enumerate(s):
+                if c not in "\r\n":
+                    break
+            else:
+                i = len(s)
+            print(repr(s[:i]), repr(s[i:]))
+            return s[:i] + s[i:-dedent]
+
+        leaves = suite.leaves()
+        for leaf in leaves:
+            if leaf.type == token.NEWLINE:
+                leaf = next(leaves, None)
+                if leaf is None:
+                    return
+                if leaf.type == token.INDENT:
+                    leaf.value = safe_dedent(leaf.value)
+                else:
+                    leaf.prefix = safe_dedent(leaf.prefix)
+            elif leaf.type == token.INDENT:
+                leaf.value = leaf.value[:-dedent]
+            elif leaf.prefix[:1] in "\r\n":
+                leaf.prefix = leaf.prefix[:-dedent]
 
 
     def transform(self, node, results):