Merge branch 'master' into extract-types

1 files changed, 40 insertions, 37 deletions

diff --git a/tools/extract_pyi.py b/tools/extract_pyi.py
index d6309212a..25b676798 100644
--- a/tools/extract_pyi.py
+++ b/tools/extract_pyi.py
@@ -6,47 +6,46 @@ import traceback
 top_level = sys.argv[1].strip("/")
 stub_directory = sys.argv[2]
 
-if top_level.count("/") == 1:
-    top_level, module = top_level.split("/")
-    modules = [module]
-else:
-    modules = os.listdir(top_level)
-    modules = sorted(modules)
-
-ok = 0
-total = 0
-for module in modules:
-    module_path = os.path.join(top_level, module)
-    if not os.path.isdir(module_path):
-        continue
+def convert_folder(top_level, stub_directory):
+    ok = 0
+    total = 0
+    filenames = sorted(os.listdir(top_level))
     pyi_lines = []
-    classes = os.listdir(module_path)
-    classes = [x for x in sorted(classes) if x.endswith(".c")]
-    if classes and classes[-1] == "__init__.c":
-        classes.insert(0, classes.pop())
-    for class_file in classes:
-        class_path = os.path.join(module_path, class_file)
-        with open(class_path, "r") as f:
-            for line in f:
-                if line.startswith("//|"):
-                    if line[3] == " ":
-                        line = line[4:]
-                    elif line[3] == "\n":
-                        line = line[3:]
-                    else:
-                        continue
-                    pyi_lines.append(line)
+    for filename in filenames:
+        full_path = os.path.join(top_level, filename)
+        file_lines = []
+        if os.path.isdir(full_path):
+            mok, mtotal = convert_folder(full_path, os.path.join(stub_directory, filename))
+            ok += mok
+            total += mtotal
+        elif filename.endswith(".c"):
+            with open(full_path, "r") as f:
+                for line in f:
+                    if line.startswith("//|"):
+                        if line[3] == " ":
+                            line = line[4:]
+                        elif line[3] == "\n":
+                            line = line[3:]
+                        else:
+                            continue
+                        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)
 
-    raw_stubs = [x for x in sorted(classes) if x.endswith(".pyi")]
-    if raw_stubs and raw_stubs[-1] == "__init__.pyi":
-        raw_stubs.insert(0, raw_stubs.pop())
-    for raw_stub in raw_stubs:
-        raw_stub_path = os.path.join(module_path, raw_stub)
-        with open(raw_stub_path, "r") as f:
-            pyi_lines.extend(f.readlines())
-    stub_filename = os.path.join(stub_directory, module + ".pyi")
+    if not pyi_lines:
+        return ok, total
+
+    stub_filename = os.path.join(stub_directory, "__init__.pyi")
     print(stub_filename)
     stub_contents = "".join(pyi_lines)
+    os.makedirs(stub_directory, exist_ok=True)
     with open(stub_filename, "w") as f:
         f.write(stub_contents)
 
@@ -76,6 +75,10 @@ for module in modules:
         traceback.print_exception(type(e), e, e.__traceback__)
     except TypeError as err:
         print(err)
+    print()
+    return ok, total
+
+ok, total = convert_folder(top_level, stub_directory)
 
 print(f"{ok} ok out of {total}")