Modify some Python stubs

1 files changed, 95 insertions, 57 deletions

diff --git a/tools/extract_pyi.py b/tools/extract_pyi.py
index d29c3444f..4d82e9e5d 100644
--- a/tools/extract_pyi.py
+++ b/tools/extract_pyi.py
@@ -17,79 +17,99 @@ 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', 'Iterable', 'Iterator', 'overload'})
-IMPORTS_TYPESHED = frozenset({'ReadableBuffer', 'WriteableBuffer'})
+IMPORTS_TYPING = frozenset({'Any', 'Optional', 'Union', 'Tuple', 'List', 'Sequence', 'NamedTuple', 'Iterable', 'Iterator', 'Callable', 'AnyStr', 'overload'})
+CPY_TYPING = frozenset({'ReadableBuffer', 'WriteableBuffer', 'AudioSample', 'FrameBuffer'})
 
 
-def is_any(node):
-    node_type = type(node)
+def is_typed(node, allow_any=False):
     if node is None:
+        return False
+    if allow_any:
         return True
-    if node_type == ast.Name and node.id == "Any":
-        return True
-    if (node_type == ast.Attribute and type(node.value) == ast.Name
-            and node.value.id == "typing" and node.attr == "Any"):
-        return True
-    return False
+    elif isinstance(node, ast.Name) and node.id == "Any":
+        return False
+    elif isinstance(node, ast.Attribute) and type(node.value) == ast.Name \
+            and node.value.id == "typing" and node.attr == "Any":
+        return False
+    return True
 
 
-def report_missing_annotations(tree):
+def find_stub_issues(tree):
     for node in ast.walk(tree):
-        node_type = type(node)
-        if node_type == ast.AnnAssign:
-            if is_any(node.annotation):
-                print(f"Missing attribute type on line {node.lineno}")
-        elif node_type == ast.arg:
-            if is_any(node.annotation) and node.arg != "self":
-                print(f"Missing argument type: {node.arg} on line {node.lineno}")
-        elif node_type == ast.FunctionDef:
-            if is_any(node.returns) and node.name != "__init__":
-                print(f"Missing return type: {node.name} on line {node.lineno}")
+        if isinstance(node, ast.AnnAssign):
+            if not is_typed(node.annotation):
+                yield ("WARN", f"Missing attribute type on line {node.lineno}")
+            if isinstance(node.value, ast.Constant) and node.value.value == Ellipsis:
+                yield ("WARN", f"Unnecessary Ellipsis assignment (= ...) on line {node.lineno}.")
+        elif isinstance(node, ast.Assign):
+            if isinstance(node.value, ast.Constant) and node.value.value == Ellipsis:
+                yield ("WARN", f"Unnecessary Ellipsis assignment (= ...) on line {node.lineno}.")
+        elif isinstance(node, ast.arguments):
+            for arg_node in (node.args + node.posonlyargs + node.kwonlyargs):
+                if not is_typed(arg_node.annotation) and (arg_node.arg != "self" and arg_node.arg != "cls"):
+                    yield ("WARN", f"Missing argument type: {arg_node.arg} on line {arg_node.lineno}")
+            if node.vararg and not is_typed(node.vararg.annotation, allow_any=True):
+                yield ("WARN", f"Missing argument type: *{node.vararg.arg} on line {node.vararg.lineno}")
+            if node.kwarg and not is_typed(node.kwarg.annotation, allow_any=True):
+                yield ("WARN", f"Missing argument type: **{node.kwarg.arg} on line {node.kwarg.lineno}")
+        elif isinstance(node, ast.FunctionDef):
+            if not is_typed(node.returns):
+                yield ("WARN", f"Missing return type: {node.name} on line {node.lineno}")
 
 
 def extract_imports(tree):
     modules = set()
     typing = set()
-    typeshed = set()
+    cpy_typing = set()
 
     def collect_annotations(anno_tree):
         if anno_tree is None:
             return
         for node in ast.walk(anno_tree):
-            node_type = type(node)
-            if node_type == ast.Name:
+            if isinstance(node, ast.Name):
                 if node.id in IMPORTS_IGNORE:
                     continue
                 elif node.id in IMPORTS_TYPING:
                     typing.add(node.id)
-                elif node.id in IMPORTS_TYPESHED:
-                    typeshed.add(node.id)
-            if node_type == ast.Attribute:
-                if type(node.value) == ast.Name:
+                elif node.id in CPY_TYPING:
+                    cpy_typing.add(node.id)
+            elif isinstance(node, ast.Attribute):
+                if isinstance(node.value, ast.Name):
                     modules.add(node.value.id)
 
     for node in ast.walk(tree):
-        node_type = type(node)
-        if (node_type == ast.AnnAssign) or (node_type == ast.arg):
+        if isinstance(node, (ast.AnnAssign, ast.arg)):
             collect_annotations(node.annotation)
-        elif node_type == ast.FunctionDef:
+        elif isinstance(node, ast.Assign):
+            collect_annotations(node.value)
+        elif isinstance(node, ast.FunctionDef):
             collect_annotations(node.returns)
             for deco in node.decorator_list:
-                if deco.id in IMPORTS_TYPING:
+                if isinstance(deco, ast.Name) and (deco.id in IMPORTS_TYPING):
                     typing.add(deco.id)
 
     return {
         "modules": sorted(modules),
         "typing": sorted(typing),
-        "typeshed": sorted(typeshed),
+        "cpy_typing": sorted(cpy_typing),
     }
 
 
+def find_references(tree):
+    for node in ast.walk(tree):
+        if isinstance(node, ast.arguments):
+            for node in ast.walk(node):
+                if isinstance(node, ast.Attribute):
+                    if isinstance(node.value, ast.Name) and node.value.id[0].isupper():
+                        yield node.value.id
+
+
 def convert_folder(top_level, stub_directory):
     ok = 0
     total = 0
     filenames = sorted(os.listdir(top_level))
-    pyi_lines = []
+    stub_fragments = []
+    references = set()
 
     for filename in filenames:
         full_path = os.path.join(top_level, filename)
@@ -99,50 +119,69 @@ def convert_folder(top_level, stub_directory):
             ok += mok
             total += mtotal
         elif filename.endswith(".c"):
-            with open(full_path, "r") as f:
+            with open(full_path, "r", encoding="utf-8") as f:
                 for line in f:
+                    line = line.rstrip()
                     if line.startswith("//|"):
-                        if line[3] == " ":
+                        if len(line) == 3:
+                            line = ""
+                        elif line[3] == " ":
                             line = line[4:]
-                        elif line[3] == "\n":
-                            line = line[3:]
                         else:
-                            continue
+                            line = line[3:]
+                            print("[WARN] There must be at least one space after '//|'")
                         file_lines.append(line)
         elif filename.endswith(".pyi"):
             with open(full_path, "r") as f:
-                file_lines.extend(f.readlines())
-
-        # Always put the contents from an __init__ first.
-        if filename.startswith("__init__."):
-            pyi_lines = file_lines + pyi_lines
-        else:
-            pyi_lines.extend(file_lines)
-
-    if not pyi_lines:
+                file_lines.extend(line.rstrip() for line in f)
+
+        fragment = "\n".join(file_lines).strip()
+        try:
+            tree = ast.parse(fragment)
+        except SyntaxError as e:
+            print(f"[ERROR] Failed to parse a Python stub from {full_path}")
+            traceback.print_exception(type(e), e, e.__traceback__)
+            return (ok, total + 1)
+        references.update(find_references(tree))
+
+        if fragment:
+            name = os.path.splitext(os.path.basename(filename))[0]
+            if name == "__init__" or (name in references):
+                stub_fragments.insert(0, fragment)
+            else:
+                stub_fragments.append(fragment)
+
+    if not stub_fragments:
         return (ok, total)
 
     stub_filename = os.path.join(stub_directory, "__init__.pyi")
     print(stub_filename)
-    stub_contents = "".join(pyi_lines)
+    stub_contents = "\n\n".join(stub_fragments)
 
-    # Validate that the module is a parseable stub.
-    total += 1
+    # Validate the stub code.
     try:
         tree = ast.parse(stub_contents)
-        imports = extract_imports(tree)
-        report_missing_annotations(tree)
-        ok += 1
     except SyntaxError as e:
         traceback.print_exception(type(e), e, e.__traceback__)
         return (ok, total)
 
+    error = False
+    for (level, msg) in find_stub_issues(tree):
+        if level == "ERROR":
+            error = True
+        print(f"[{level}] {msg}")
+
+    total += 1
+    if not error:
+        ok += 1
+
     # Add import statements
+    imports = extract_imports(tree)
     import_lines = ["from __future__ import annotations"]
     if imports["typing"]:
         import_lines.append("from typing import " + ", ".join(imports["typing"]))
-    if imports["typeshed"]:
-        import_lines.append("from _typeshed import " + ", ".join(imports["typeshed"]))
+    if imports["cpy_typing"]:
+        import_lines.append("from _typing import " + ", ".join(imports["cpy_typing"]))
     import_lines.extend(f"import {m}" for m in imports["modules"])
     import_body = "\n".join(import_lines)
     m = re.match(r'(\s*""".*?""")', stub_contents, flags=re.DOTALL)
@@ -159,7 +198,6 @@ def convert_folder(top_level, stub_directory):
     with open(stub_filename, "w") as f:
         f.write(stub_contents)
 
-    print()
     return (ok, total)