Repository: asottile/pyupgrade

Issue:
rewrite typing.NamedTuple and typing.TypedDict to class form in --py36-plus
```python
Employee = NamedTuple('Employee', [('name', str), ('id', int)])
```

becomes

```python
class Employee(NamedTuple):
    name: str
    id: int
```

Discussion:

Patch:
diff --git a/README.md b/README.md
index 8178ec6..4e3a01a 100644
--- a/README.md
+++ b/README.md
@@ -387,6 +387,34 @@ except OSError:
     raise
 ```
 
+### `typing.NamedTuple` / `typing.TypedDict` py36+ syntax
+
+Availability:
+- `--py36-plus` is passed on the commandline.
+
+```python
+# input
+NT = typing.NamedTuple('NT', [('a': int, 'b': Tuple[str, ...])])
+
+D1 = typing.TypedDict('D1', a=int, b=str)
+
+D2 = typing.TypedDict('D2', {'a': int, 'b': str})
+
+# output
+
+class NT(typing.NamedTuple):
+    a: int
+    b: Tuple[str, ...]
+
+class D1(typing.TypedDict):
+    a: int
+    b: str
+
+class D2(typing.TypedDict):
+    a: int
+    b: str
+```
+
 ### f-strings
 
 Availability:
diff --git a/pyupgrade.py b/pyupgrade.py
index bcd4250..e6511be 100644
--- a/pyupgrade.py
+++ b/pyupgrade.py
@@ -54,6 +54,8 @@ SyncFunctionDef = Union[ast.FunctionDef, ast.Lambda]
 
 _stdlib_parse_format = string.Formatter().parse
 
+_KEYWORDS = frozenset(keyword.kwlist)
+
 
 def parse_format(s: str) -> Tuple[DotFormatPart, ...]:
     """Makes the empty string not a special case.  In the stdlib, there's
@@ -1000,7 +1002,7 @@ def _fix_percent_format_dict(
         elif not k.s.isidentifier():
             return
         # a keyword
-        elif k.s in keyword.kwlist:
+        elif k.s in _KEYWORDS:
             return
         seen_keys.add(k.s)
         keys[_ast_to_offset(k)] = k
@@ -2160,9 +2162,35 @@ def _format_params(call: ast.Call) -> Dict[str, str]:
     return params
 
 
-class FindSimpleFormats(ast.NodeVisitor):
+class FindPy36Plus(ast.NodeVisitor):
     def __init__(self) -> None:
-        self.found: Dict[Offset, ast.Call] = {}
+        self.fstrings: Dict[Offset, ast.Call] = {}
+        self.named_tuples: Dict[Offset, ast.Call] = {}
+        self.dict_typed_dicts: Dict[Offset, ast.Call] = {}
+        self.kw_typed_dicts: Dict[Offset, ast.Call] = {}
+        self._from_imports: Dict[str, Set[str]] = collections.defaultdict(set)
+
+    def visit_ImportFrom(self, node: ast.ImportFrom) -> None:
+        if node.module in {'typing', 'typing_extensions'}:
+            for name in node.names:
+                if not name.asname:
+                    self._from_imports[node.module].add(name.name)
+        self.generic_visit(node)
+
+    def _is_attr(self, node: ast.AST, mods: Set[str], name: str) -> bool:
+        return (
+            (
+                isinstance(node, ast.Name) and
+                node.id == name and
+                any(name in self._from_imports[mod] for mod in mods)
+            ) or
+            (
+                isinstance(node, ast.Attribute) and
+                node.attr == name and
+                isinstance(node.value, ast.Name) and
+                node.value.id in mods
+            )
+        )
 
     def _parse(self, node: ast.Call) -> Optional[Tuple[DotFormatPart, ...]]:
         if not (
@@ -2207,7 +2235,64 @@ class FindSimpleFormats(ast.NodeVisitor):
                     if not candidate:
                         i += 1
             else:
-                self.found[_ast_to_offset(node)] = node
+                self.fstrings[_ast_to_offset(node)] = node
+
+        self.generic_visit(node)
+
+    def visit_Assign(self, node: ast.Assign) -> None:
+        if (
+                # NT = ...("NT", ...)
+                len(node.targets) == 1 and
+                isinstance(node.targets[0], ast.Name) and
+                isinstance(node.value, ast.Call) and
+                len(node.value.args) >= 1 and
+                isinstance(node.value.args[0], ast.Str) and
+                node.targets[0].id == node.value.args[0].s and
+                not _starargs(node.value)
+        ):
+            if (
+                    self._is_attr(
+                        node.value.func, {'typing'}, 'NamedTuple',
+                    ) and
+                    len(node.value.args) == 2 and
+                    not node.value.keywords and
+                    isinstance(node.value.args[1], (ast.List, ast.Tuple)) and
+                    len(node.value.args[1].elts) > 0 and
+                    all(
+                        isinstance(tup, ast.Tuple) and
+                        len(tup.elts) == 2 and
+                        isinstance(tup.elts[0], ast.Str) and
+                        tup.elts[0].s not in _KEYWORDS
+                        for tup in node.value.args[1].elts
+                    )
+            ):
+                self.named_tuples[_ast_to_offset(node)] = node.value
+            elif (
+                    self._is_attr(
+                        node.value.func,
+                        {'typing', 'typing_extensions'},
+                        'TypedDict',
+                    ) and
+                    len(node.value.args) == 1 and
+                    len(node.value.keywords) > 0
+            ):
+                self.kw_typed_dicts[_ast_to_offset(node)] = node.value
+            elif (
+                    self._is_attr(
+                        node.value.func,
+                        {'typing', 'typing_extensions'},
+                        'TypedDict',
+                    ) and
+                    len(node.value.args) == 2 and
+                    not node.value.keywords and
+                    isinstance(node.value.args[1], ast.Dict) and
+                    node.value.args[1].keys and
+                    all(
+                        isinstance(k, ast.Str)
+                        for k in node.value.args[1].keys
+                    )
+            ):
+                self.dict_typed_dicts[_ast_to_offset(node)] = node.value
 
         self.generic_visit(node)
 
@@ -2219,6 +2304,23 @@ def _unparse(node: ast.expr) -> str:
         return ''.join((_unparse(node.value), '.', node.attr))
     elif isinstance(node, ast.Call):
         return '{}()'.format(_unparse(node.func))
+    elif isinstance(node, ast.Subscript):
+        assert isinstance(node.slice, ast.Index), ast.dump(node)
+        if isinstance(node.slice.value, ast.Name):
+            slice_s = _unparse(node.slice.value)
+        elif isinstance(node.slice.value, ast.Tuple):
+            slice_s = ', '.join(_unparse(elt) for elt in node.slice.value.elts)
+        else:
+            raise NotImplementedError(ast.dump(node))
+        return '{}[{}]'.format(_unparse(node.value), slice_s)
+    elif isinstance(node, ast.Str):
+        return repr(node.s)
+    elif isinstance(node, ast.Ellipsis):
+        return '...'
+    elif isinstance(node, ast.List):
+        return '[{}]'.format(', '.join(_unparse(elt) for elt in node.elts))
+    elif isinstance(node, ast.NameConstant):
+        return repr(node.value)
     else:
         raise NotImplementedError(ast.dump(node))
 
@@ -2238,16 +2340,43 @@ def _to_fstring(src: str, call: ast.Call) -> str:
     return unparse_parsed_string(parts)
 
 
-def _fix_fstrings(contents_text: str) -> str:
+def _replace_typed_class(
+        tokens: List[Token],
+        i: int,
+        call: ast.Call,
+        types: Dict[str, ast.expr],
+) -> None:
+    if i > 0 and tokens[i - 1].name in {'INDENT', UNIMPORTANT_WS}:
+        indent = f'{tokens[i - 1].src}{" " * 4}'
+    else:
+        indent = ' ' * 4
+
+    # NT = NamedTuple("nt", [("a", int)])
+    # ^i                                 ^end
+    end = i + 1
+    while end < len(tokens) and tokens[end].name != 'NEWLINE':
+        end += 1
+
+    attrs = '\n'.join(f'{indent}{k}: {_unparse(v)}' for k, v in types.items())
+    src = f'class {tokens[i].src}({_unparse(call.func)}):\n{attrs}'
+    tokens[i:end] = [Token('CODE', src)]
+
+
+def _fix_py36_plus(contents_text: str) -> str:
     try:
         ast_obj = ast_parse(contents_text)
     except SyntaxError:
         return contents_text
 
-    visitor = FindSimpleFormats()
+    visitor = FindPy36Plus()
     visitor.visit(ast_obj)
 
-    if not visitor.found:
+    if not any((
+            visitor.fstrings,
+            visitor.named_tuples,
+            visitor.dict_typed_dicts,
+            visitor.kw_typed_dicts,
+    )):
         return contents_text
 
     try:
@@ -2255,27 +2384,50 @@ def _fix_fstrings(contents_text: str) -> str:
     except tokenize.TokenError:  # pragma: no cover (bpo-2180)
         return contents_text
     for i, token in reversed_enumerate(tokens):
-        node = visitor.found.get(token.offset)
-        if node is None:
-            continue
+        if token.offset in visitor.fstrings:
+            node = visitor.fstrings[token.offset]
 
-        # TODO: handle \N escape sequences
-        if r'\N' in token.src:
-            continue
+            # TODO: handle \N escape sequences
+            if r'\N' in token.src:
+                continue
 
-        paren = i + 3
-        if tokens_to_src(tokens[i + 1:paren + 1]) != '.format(':
-            continue
+            paren = i + 3
+            if tokens_to_src(tokens[i + 1:paren + 1]) != '.format(':
+                continue
 
-        # we don't actually care about arg position, so we pass `node`
-        victims = _victims(tokens, paren, node, gen=False)
-        end = victims.ends[-1]
-        # if it spans more than one line, bail
-        if tokens[end].line != token.line:
-            continue
+            # we don't actually care about arg position, so we pass `node`
+            victims = _victims(tokens, paren, node, gen=False)
+            end = victims.ends[-1]
+            # if it spans more than one line, bail
+            if tokens[end].line != token.line:
+                continue
 
-        tokens[i] = token._replace(src=_to_fstring(token.src, node))
-        del tokens[i + 1:end + 1]
+            tokens[i] = token._replace(src=_to_fstring(token.src, node))
+            del tokens[i + 1:end + 1]
+        elif token.offset in visitor.named_tuples and token.name == 'NAME':
+            call = visitor.named_tuples[token.offset]
+            types: Dict[str, ast.expr] = {
+                tup.elts[0].s: tup.elts[1]  # type: ignore  # (checked above)
+                for tup in call.args[1].elts  # type: ignore  # (checked above)
+            }
+            _replace_typed_class(tokens, i, call, types)
+        elif token.offset in visitor.kw_typed_dicts and token.name == 'NAME':
+            call = visitor.kw_typed_dicts[token.offset]
+            types = {
+                arg.arg: arg.value  # type: ignore  # (checked above)
+                for arg in call.keywords
+            }
+            _replace_typed_class(tokens, i, call, types)
+        elif token.offset in visitor.dict_typed_dicts and token.name == 'NAME':
+            call = visitor.dict_typed_dicts[token.offset]
+            types = {
+                k.s: v  # type: ignore  # (checked above)
+                for k, v in zip(
+                    call.args[1].keys,  # type: ignore  # (checked above)
+                    call.args[1].values,  # type: ignore  # (checked above)
+                )
+            }
+            _replace_typed_class(tokens, i, call, types)
 
     return tokens_to_src(tokens)
 
@@ -2300,7 +2452,7 @@ def _fix_file(filename: str, args: argparse.Namespace) -> int:
     if args.min_version >= (3,):
         contents_text = _fix_py3_plus(contents_text)
     if args.min_version >= (3, 6):
-        contents_text = _fix_fstrings(contents_text)
+        contents_text = _fix_py36_plus(contents_text)
 
     if filename == '-':
         print(contents_text, end='')


Test patch:
diff --git a/tests/fstrings_test.py b/tests/fstrings_test.py
index 2090930..fcc6a92 100644
--- a/tests/fstrings_test.py
+++ b/tests/fstrings_test.py
@@ -1,6 +1,6 @@
 import pytest
 
-from pyupgrade import _fix_fstrings
+from pyupgrade import _fix_py36_plus
 
 
 @pytest.mark.parametrize(
@@ -33,7 +33,7 @@ from pyupgrade import _fix_fstrings
     ),
 )
 def test_fix_fstrings_noop(s):
-    assert _fix_fstrings(s) == s
+    assert _fix_py36_plus(s) == s
 
 
 @pytest.mark.parametrize(
@@ -55,4 +55,4 @@ def test_fix_fstrings_noop(s):
     ),
 )
 def test_fix_fstrings(s, expected):
-    assert _fix_fstrings(s) == expected
+    assert _fix_py36_plus(s) == expected
diff --git a/tests/typing_named_tuple_test.py b/tests/typing_named_tuple_test.py
new file mode 100644
index 0000000..83a41ef
--- /dev/null
+++ b/tests/typing_named_tuple_test.py
@@ -0,0 +1,155 @@
+import pytest
+
+from pyupgrade import _fix_py36_plus
+
+
+@pytest.mark.parametrize(
+    's',
+    (
+        '',
+
+        pytest.param(
+            'from typing import NamedTuple as wat\n'
+            'C = wat("C", ("a", int))\n',
+            id='currently not following as imports',
+        ),
+
+        pytest.param('C = typing.NamedTuple("C", ())', id='no types'),
+        pytest.param('C = typing.NamedTuple("C")', id='not enough args'),
+        pytest.param(
+            'C = typing.NamedTuple("C", (), nonsense=1)',
+            id='namedtuple with named args',
+        ),
+        pytest.param(
+            'C = typing.NamedTuple("C", {"foo": int, "bar": str})',
+            id='namedtuple without a list/tuple',
+        ),
+        pytest.param(
+            'C = typing.NamedTuple("C", [["a", str], ["b", int]])',
+            id='namedtuple without inner tuples',
+        ),
+        pytest.param(
+            'C = typing.NamedTuple("C", [(), ()])',
+            id='namedtuple but inner tuples are incorrect length',
+        ),
+        pytest.param(
+            'C = typing.NamedTuple("C", [(not_a_str, str)])',
+            id='namedtuple but attribute name is not a string',
+        ),
+
+        pytest.param(
+            'C = typing.NamedTuple("C", [("def", int)])',
+            id='uses keyword',
+        ),
+        pytest.param(
+            'C = typing.NamedTuple("C", *types)',
+            id='NamedTuple starargs',
+        ),
+    ),
+)
+def test_typing_named_tuple_noop(s):
+    assert _fix_py36_plus(s) == s
+
+
+@pytest.mark.parametrize(
+    ('s', 'expected'),
+    (
+        pytest.param(
+            'from typing import NamedTuple\n'
+            'C = NamedTuple("C", [("a", int), ("b", str)])\n',
+
+            'from typing import NamedTuple\n'
+            'class C(NamedTuple):\n'
+            '    a: int\n'
+            '    b: str\n',
+            id='typing from import',
+        ),
+        pytest.param(
+            'import typing\n'
+            'C = typing.NamedTuple("C", [("a", int), ("b", str)])\n',
+
+            'import typing\n'
+            'class C(typing.NamedTuple):\n'
+            '    a: int\n'
+            '    b: str\n',
+
+            id='import typing',
+        ),
+        pytest.param(
+            'C = typing.NamedTuple("C", [("a", List[int])])',
+
+            'class C(typing.NamedTuple):\n'
+            '    a: List[int]',
+
+            id='generic attribute types',
+        ),
+        pytest.param(
+            'C = typing.NamedTuple("C", [("a", Mapping[int, str])])',
+
+            'class C(typing.NamedTuple):\n'
+            '    a: Mapping[int, str]',
+
+            id='generic attribute types with multi types',
+        ),
+        pytest.param(
+            'C = typing.NamedTuple("C", [("a", "Queue[int]")])',
+
+            'class C(typing.NamedTuple):\n'
+            "    a: 'Queue[int]'",
+
+            id='quoted type names',
+        ),
+        pytest.param(
+            'C = typing.NamedTuple("C", [("a", Tuple[int, ...])])',
+
+            'class C(typing.NamedTuple):\n'
+            '    a: Tuple[int, ...]',
+
+            id='type with ellipsis',
+        ),
+        pytest.param(
+            'C = typing.NamedTuple("C", [("a", Callable[[Any], None])])',
+
+            'class C(typing.NamedTuple):\n'
+            '    a: Callable[[Any], None]',
+
+            id='type containing a list',
+        ),
+        pytest.param(
+            'if False:\n'
+            '    pass\n'
+            'C = typing.NamedTuple("C", [("a", int)])\n',
+
+            'if False:\n'
+            '    pass\n'
+            'class C(typing.NamedTuple):\n'
+            '    a: int\n',
+
+            id='class directly after block',
+        ),
+        pytest.param(
+            'if True:\n'
+            '    C = typing.NamedTuple("C", [("a", int)])\n',
+
+            'if True:\n'
+            '    class C(typing.NamedTuple):\n'
+            '        a: int\n',
+
+            id='indented',
+        ),
+        pytest.param(
+            'if True:\n'
+            '    ...\n'
+            '    C = typing.NamedTuple("C", [("a", int)])\n',
+
+            'if True:\n'
+            '    ...\n'
+            '    class C(typing.NamedTuple):\n'
+            '        a: int\n',
+
+            id='indented, but on next line',
+        ),
+    ),
+)
+def test_fix_typing_named_tuple(s, expected):
+    assert _fix_py36_plus(s) == expected
diff --git a/tests/typing_typed_dict_test.py b/tests/typing_typed_dict_test.py
new file mode 100644
index 0000000..af4e4e7
--- /dev/null
+++ b/tests/typing_typed_dict_test.py
@@ -0,0 +1,96 @@
+import pytest
+
+from pyupgrade import _fix_py36_plus
+
+
+@pytest.mark.parametrize(
+    's',
+    (
+        pytest.param(
+            'from wat import TypedDict\n'
+            'Q = TypedDict("Q")\n',
+            id='from imported from elsewhere',
+        ),
+        pytest.param('D = typing.TypedDict("D")', id='no typed kwargs'),
+        pytest.param('D = typing.TypedDict("D", {})', id='no typed args'),
+        pytest.param('D = typing.TypedDict("D", {}, a=int)', id='both'),
+        pytest.param('D = typing.TypedDict("D", 1)', id='not a dict'),
+        pytest.param(
+            'D = typing.TypedDict("D", {1: str})',
+            id='key is not a string',
+        ),
+        pytest.param(
+            'D = typing.TypedDict("D", {**d, "a": str})',
+            id='dictionary splat operator',
+        ),
+        pytest.param(
+            'C = typing.TypedDict("C", *types)',
+            id='starargs',
+        ),
+        pytest.param(
+            'D = typing.TypedDict("D", **types)',
+            id='starstarkwargs',
+        ),
+    ),
+)
+def test_typing_typed_dict_noop(s):
+    assert _fix_py36_plus(s) == s
+
+
+@pytest.mark.parametrize(
+    ('s', 'expected'),
+    (
+        pytest.param(
+            'from typing import TypedDict\n'
+            'D = TypedDict("D", a=int)\n',
+
+            'from typing import TypedDict\n'
+            'class D(TypedDict):\n'
+            '    a: int\n',
+
+            id='keyword TypedDict from imported',
+        ),
+        pytest.param(
+            'import typing\n'
+            'D = typing.TypedDict("D", a=int)\n',
+
+            'import typing\n'
+            'class D(typing.TypedDict):\n'
+            '    a: int\n',
+
+            id='keyword TypedDict from attribute',
+        ),
+        pytest.param(
+            'import typing\n'
+            'D = typing.TypedDict("D", {"a": int})\n',
+
+            'import typing\n'
+            'class D(typing.TypedDict):\n'
+            '    a: int\n',
+
+            id='TypedDict from dict literal',
+        ),
+        pytest.param(
+            'from typing_extensions import TypedDict\n'
+            'D = TypedDict("D", a=int)\n',
+
+            'from typing_extensions import TypedDict\n'
+            'class D(TypedDict):\n'
+            '    a: int\n',
+
+            id='keyword TypedDict from typing_extensions',
+        ),
+        pytest.param(
+            'import typing_extensions\n'
+            'D = typing_extensions.TypedDict("D", {"a": int})\n',
+
+            'import typing_extensions\n'
+            'class D(typing_extensions.TypedDict):\n'
+            '    a: int\n',
+
+            id='keyword TypedDict from typing_extensions',
+        ),
+    ),
+)
+def test_typing_typed_dict(s, expected):
+    assert _fix_py36_plus(s) == expected