Repository: asottile/pyupgrade

Issue:
Remove `__metaclass__ = type` at module scope in --py3-plus mode
this was a way to trigger new-style classes in python2

Comments:

Patch:
diff --git a/README.md b/README.md
index 1a2f764..f6f6ffc 100644
--- a/README.md
+++ b/README.md
@@ -170,11 +170,19 @@ class C(Base):
 Availability:
 - `--py3-plus` is passed on the commandline.
 
+#### rewrites class declaration
+
 ```python
 class C(object): pass     # class C: pass
 class C(B, object): pass  # class C(B): pass
 ```
 
+#### removes `__metaclass__ = type` declaration
+
+```diff
+-__metaclass__ = type
+```
+
 ### forced `str("native")` literals
 
 Availability:
diff --git a/pyupgrade.py b/pyupgrade.py
index 8913496..29fbb8f 100644
--- a/pyupgrade.py
+++ b/pyupgrade.py
@@ -242,8 +242,8 @@ def _victims(
     return Victims(starts, sorted(set(ends)), first_comma_index, arg_index)
 
 
-def _find_token(tokens: List[Token], i: int, token: Token) -> int:
-    while tokens[i].src != token:
+def _find_token(tokens: List[Token], i: int, src: str) -> int:
+    while tokens[i].src != src:
         i += 1
     return i
 
@@ -454,7 +454,7 @@ ESCAPE_RE = re.compile(r'\\.', re.DOTALL)
 NAMED_ESCAPE_NAME = re.compile(r'\{[^}]+\}')
 
 
-def _fix_escape_sequences(token: Token) -> str:
+def _fix_escape_sequences(token: Token) -> Token:
     prefix, rest = parse_string_literal(token.src)
     actual_prefix = prefix.lower()
 
@@ -1204,6 +1204,7 @@ class FindPy3Plus(ast.NodeVisitor):
         self.if_py3_blocks: Set[Offset] = set()
         self.if_py2_blocks_else: Set[Offset] = set()
         self.if_py3_blocks_else: Set[Offset] = set()
+        self.metaclass_type_assignments: Set[Offset] = set()
 
         self.native_literals: Set[Offset] = set()
 
@@ -1486,6 +1487,19 @@ class FindPy3Plus(ast.NodeVisitor):
 
         self.generic_visit(node)
 
+    def visit_Assign(self, node: ast.Assign) -> None:
+        if (
+                len(node.targets) == 1 and
+                isinstance(node.targets[0], ast.Name) and
+                node.targets[0].col_offset == 0 and
+                node.targets[0].id == '__metaclass__' and
+                isinstance(node.value, ast.Name) and
+                node.value.id == 'type'
+        ):
+            self.metaclass_type_assignments.add(_ast_to_offset(node))
+
+        self.generic_visit(node)
+
     @staticmethod
     def _eq(test: ast.Compare, n: int) -> bool:
         return (
@@ -1609,7 +1623,7 @@ class Block(NamedTuple):
     colon: int
     block: int
     end: int
-    line: int
+    line: bool
 
     def _initial_indent(self, tokens: List[Token]) -> int:
         if tokens[self.start].src.isspace():
@@ -1869,6 +1883,7 @@ def _fix_py3_plus(contents_text: str) -> str:
             visitor.if_py2_blocks_else,
             visitor.if_py3_blocks,
             visitor.if_py3_blocks_else,
+            visitor.metaclass_type_assignments,
             visitor.native_literals,
             visitor.io_open_calls,
             visitor.open_mode_calls,
@@ -1943,6 +1958,16 @@ def _fix_py3_plus(contents_text: str) -> str:
                 if_block, else_block = _find_if_else_block(tokens, j)
                 del tokens[if_block.end:else_block.end]
                 if_block.replace_condition(tokens, [Token('NAME', 'else')])
+        elif token.offset in visitor.metaclass_type_assignments:
+            j = _find_token(tokens, i, 'type')
+            while tokens[j].name not in {'NEWLINE', 'ENDMARKER'}:
+                j += 1
+            # depending on the version of python, some will not emit
+            # NEWLINE('') at the end of a file which does not end with a
+            # newline (for example 3.6.5)
+            if tokens[j].name == 'ENDMARKER':  # pragma: no cover
+                j -= 1
+            del tokens[i:j + 1]
         elif token.offset in visitor.native_literals:
             j = _find_open_paren(tokens, i)
             func_args, end = _parse_call_args(tokens, j)


Test patch:
diff --git a/tests/metaclass_type_test.py b/tests/metaclass_type_test.py
new file mode 100644
index 0000000..51b92c2
--- /dev/null
+++ b/tests/metaclass_type_test.py
@@ -0,0 +1,57 @@
+import pytest
+
+from pyupgrade import _fix_py3_plus
+
+
+@pytest.mark.parametrize(
+    's',
+    (
+        pytest.param(
+            'x = type\n'
+            '__metaclass__ = x\n',
+            id='not rewriting "type" rename',
+        ),
+        pytest.param(
+            'def foo():\n'
+            '    __metaclass__ = type\n',
+            id='not rewriting function scope',
+        ),
+        pytest.param(
+            'class Foo:\n'
+            '    __metaclass__ = type\n',
+            id='not rewriting class scope',
+        ),
+        pytest.param(
+            '__metaclass__, __meta_metaclass__ = type, None\n',
+            id='not rewriting multiple assignment',
+        ),
+    ),
+)
+def test_metaclass_type_assignment_noop(s):
+    assert _fix_py3_plus(s) == s
+
+
+@pytest.mark.parametrize(
+    ('s', 'expected'),
+    (
+        pytest.param(
+            '__metaclass__ = type',
+            '',
+            id='module-scope assignment',
+        ),
+        pytest.param(
+            '__metaclass__  =   type',
+            '',
+            id='module-scope assignment with extra whitespace',
+        ),
+        pytest.param(
+            '__metaclass__ = (\n'
+            '   type\n'
+            ')\n',
+            '',
+            id='module-scope assignment across newline',
+        ),
+    ),
+)
+def test_fix_metaclass_type_assignment(s, expected):
+    assert _fix_py3_plus(s) == expected