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

Discussion:

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..c946fef 100644
--- a/pyupgrade.py
+++ b/pyupgrade.py
@@ -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 (
@@ -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