add circuitpython code

1 files changed, 260 insertions, 0 deletions

diff --git a/circuitpython/tools/extract_pyi.py b/circuitpython/tools/extract_pyi.py
new file mode 100644
index 0000000..7b87404
--- /dev/null
+++ b/circuitpython/tools/extract_pyi.py
@@ -0,0 +1,260 @@
+# SPDX-FileCopyrightText: 2014 MicroPython & CircuitPython contributors (https://github.com/adafruit/circuitpython/graphs/contributors)
+#
+# SPDX-License-Identifier: MIT
+
+# Run with 'python tools/extract_pyi.py shared-bindings/ path/to/stub/dir
+# You can also test a specific library in shared-bindings by putting the path
+# to that directory instead
+
+import ast
+import os
+import re
+import sys
+import traceback
+import types
+import typing
+
+import isort
+import black
+
+import circuitpython_typing
+import circuitpython_typing.socket
+
+PATHS_IGNORE = frozenset({"shared-bindings/__future__"})
+
+TYPE_MODULE_IMPORTS_IGNORE = frozenset(
+    {
+        "array",
+        "bool",
+        "buffer",
+        "bytearray",
+        "bytes",
+        "dict",
+        "file",
+        "float",
+        "int",
+        "list",
+        "range",
+        "set",
+        "slice",
+        "str",
+        "struct_time",
+        "tuple",
+    }
+)
+
+# Include all definitions in these type modules, minus some name conflicts.
+AVAILABLE_TYPE_MODULE_IMPORTS = {
+    "types": frozenset(types.__all__),
+    # Conflicts: countio.Counter, canio.Match
+    "typing": frozenset(typing.__all__) - set(["Counter", "Match"]),
+    "circuitpython_typing": frozenset(circuitpython_typing.__all__),
+    "circuitpython_typing.socket": frozenset(circuitpython_typing.socket.__all__),
+}
+
+
+def is_typed(node, allow_any=False):
+    if node is None:
+        return False
+    if allow_any:
+        return True
+    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 find_stub_issues(tree):
+    for node in ast.walk(tree):
+        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):
+            allargs = list(node.args + node.kwonlyargs)
+            if sys.version_info >= (3, 8):
+                allargs.extend(node.posonlyargs)
+            for arg_node in allargs:
+                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()
+    used_type_module_imports = {module: set() for module in AVAILABLE_TYPE_MODULE_IMPORTS.keys()}
+
+    def collect_annotations(anno_tree):
+        if anno_tree is None:
+            return
+        for node in ast.walk(anno_tree):
+            if isinstance(node, ast.Name):
+                if node.id in TYPE_MODULE_IMPORTS_IGNORE:
+                    continue
+                for module, imports in AVAILABLE_TYPE_MODULE_IMPORTS.items():
+                    if node.id in imports:
+                        used_type_module_imports[module].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):
+        if isinstance(node, (ast.AnnAssign, ast.arg)):
+            collect_annotations(node.annotation)
+        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 isinstance(deco, ast.Name) and (
+                    deco.id in AVAILABLE_TYPE_MODULE_IMPORTS["typing"]
+                ):
+                    used_type_module_imports["typing"].add(deco.id)
+
+    return (modules, used_type_module_imports)
+
+
+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))
+    stub_fragments = []
+    references = set()
+
+    for filename in filenames:
+        full_path = os.path.join(top_level, filename)
+        if full_path in PATHS_IGNORE:
+            continue
+
+        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", encoding="utf-8") as f:
+                for line in f:
+                    line = line.rstrip()
+                    if line.startswith("//|"):
+                        if len(line) == 3:
+                            line = ""
+                        elif line[3] == " ":
+                            line = line[4:]
+                        else:
+                            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(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 = "\n\n".join(stub_fragments)
+
+    # Validate the stub code.
+    try:
+        tree = ast.parse(stub_contents)
+    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, type_imports = extract_imports(tree)
+    import_lines = ["from __future__ import annotations"]
+    for type_module, used_types in type_imports.items():
+        import_lines.append(f"from {type_module} import {', '.join(sorted(used_types))}")
+    import_lines.extend(f"import {m}" for m in sorted(imports))
+    import_body = "\n".join(import_lines)
+    m = re.match(r'(\s*""".*?""")', stub_contents, flags=re.DOTALL)
+    if m:
+        stub_contents = m.group(1) + "\n\n" + import_body + "\n\n" + stub_contents[m.end() :]
+    else:
+        stub_contents = import_body + "\n\n" + stub_contents
+
+    # Code formatting
+    stub_contents = isort.code(stub_contents)
+    stub_contents = black.format_str(stub_contents, mode=black.FileMode(is_pyi=True))
+
+    os.makedirs(stub_directory, exist_ok=True)
+    with open(stub_filename, "w") as f:
+        f.write(stub_contents)
+
+    return (ok, total)
+
+
+if __name__ == "__main__":
+    top_level = sys.argv[1].strip("/")
+    stub_directory = sys.argv[2]
+
+    (ok, total) = convert_folder(top_level, stub_directory)
+
+    print(f"Parsing .pyi files: {total - ok} failed, {ok} passed")
+
+    if ok != total:
+        sys.exit(total - ok)