correctly compute indent for typeddict after dedent

asottile · Aug 2, 2021 · a9a451c · a9a451c
1 parent 42be346
commit a9a451c
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 0 deletions.
diff --git a/pyupgrade/_main.py b/pyupgrade/_main.py
@@ -746,6 +746,8 @@ def _typed_class_replacement(
         call: ast.Call,
         types: Dict[str, ast.expr],
 ) -> Tuple[int, str]:
+    while i > 0 and tokens[i - 1].name == 'DEDENT':
+        i -= 1
     if i > 0 and tokens[i - 1].name in {'INDENT', UNIMPORTANT_WS}:
         indent = f'{tokens[i - 1].src}{" " * 4}'
     else:

diff --git a/tests/features/typing_typed_dict_test.py b/tests/features/typing_typed_dict_test.py
@@ -134,6 +134,22 @@ def test_typing_typed_dict_noop(s):
 
             id='index unparse error',
         ),
+        pytest.param(
+            'import typing\n'
+            'if True:\n'
+            '    if False:\n'
+            '        pass\n'
+            '    D = typing.TypedDict("D", a=int)\n',
+
+            'import typing\n'
+            'if True:\n'
+            '    if False:\n'
+            '        pass\n'
+            '    class D(typing.TypedDict):\n'
+            '        a: int\n',
+
+            id='right after a dedent',
+        ),
     ),
 )
 def test_typing_typed_dict(s, expected):