Fixing make stub warnings and mypy issuesmak

1 files changed, 8 insertions, 1 deletions

diff --git a/tools/extract_pyi.py b/tools/extract_pyi.py
index 651216e11..b7ce584a1 100644
--- a/tools/extract_pyi.py
+++ b/tools/extract_pyi.py
@@ -17,7 +17,8 @@ import black
 
 
 IMPORTS_IGNORE = frozenset({'int', 'float', 'bool', 'str', 'bytes', 'tuple', 'list', 'set', 'dict', 'bytearray', 'slice', 'file', 'buffer', 'range', 'array', 'struct_time'})
-IMPORTS_TYPING = frozenset({'Any', 'Optional', 'Union', 'Tuple', 'List', 'Sequence', 'NamedTuple', 'Iterable', 'Iterator', 'Callable', 'AnyStr', 'overload'})
+IMPORTS_TYPING = frozenset({'Any', 'Optional', 'Union', 'Tuple', 'List', 'Sequence', 'NamedTuple', 'Iterable', 'Iterator', 'Callable', 'AnyStr', 'overload', 'Type'})
+IMPORTS_TYPES = frozenset({'TracebackType'})
 CPY_TYPING = frozenset({'ReadableBuffer', 'WriteableBuffer', 'AudioSample', 'FrameBuffer'})
 
 
@@ -63,6 +64,7 @@ def find_stub_issues(tree):
 def extract_imports(tree):
     modules = set()
     typing = set()
+    types = set()
     cpy_typing = set()
 
     def collect_annotations(anno_tree):
@@ -74,6 +76,8 @@ def extract_imports(tree):
                     continue
                 elif node.id in IMPORTS_TYPING:
                     typing.add(node.id)
+                elif node.id in IMPORTS_TYPES:
+                    types.add(node.id)
                 elif node.id in CPY_TYPING:
                     cpy_typing.add(node.id)
             elif isinstance(node, ast.Attribute):
@@ -94,6 +98,7 @@ def extract_imports(tree):
     return {
         "modules": sorted(modules),
         "typing": sorted(typing),
+        "types": sorted(types),
         "cpy_typing": sorted(cpy_typing),
     }
 
@@ -181,6 +186,8 @@ def convert_folder(top_level, stub_directory):
     # Add import statements
     imports = extract_imports(tree)
     import_lines = ["from __future__ import annotations"]
+    if imports["types"]:
+        import_lines.append("from types import " + ", ".join(imports["types"]))
     if imports["typing"]:
         import_lines.append("from typing import " + ", ".join(imports["typing"]))
     if imports["cpy_typing"]: