diff options
Diffstat (limited to 'circuitpython/tools/extract_pyi.py')
| -rw-r--r-- | circuitpython/tools/extract_pyi.py | 260 |
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) |