circuitpython/tools/extract_pyi.py

Ignoring revisions in .git-blame-ignore-revs. Click here to bypass and see the normal blame view.

257 lines
8.7 KiB
Python
Raw Normal View History

2020-06-03 18:40:05 -04:00
# 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
2020-07-21 11:37:22 -04:00
import ast
2020-04-27 17:36:14 -04:00
import os
2020-07-21 11:37:22 -04:00
import re
2020-04-27 17:36:14 -04:00
import sys
import traceback
import types
import typing
2020-04-27 17:36:14 -04:00
2020-07-21 11:37:22 -04:00
import isort
2020-07-23 14:22:41 -04:00
import black
2020-07-21 11:37:22 -04:00
import circuitpython_typing
import circuitpython_typing.socket
2020-07-21 11:37:22 -04:00
2021-03-15 09:57:36 -04:00
IMPORTS_IGNORE = frozenset(
{
"array",
2021-03-15 09:57:36 -04:00
"bool",
"buffer",
"bytearray",
2021-03-15 09:57:36 -04:00
"bytes",
"dict",
"file",
"float",
"int",
"list",
2021-03-15 09:57:36 -04:00
"range",
"set",
"slice",
"str",
2021-03-15 09:57:36 -04:00
"struct_time",
"tuple",
2021-03-15 09:57:36 -04:00
}
)
# 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__),
}
2020-07-21 11:37:22 -04:00
2020-08-03 00:35:43 -04:00
def is_typed(node, allow_any=False):
2020-07-21 11:37:22 -04:00
if node is None:
2020-08-03 00:35:43 -04:00
return False
if allow_any:
2020-07-21 11:37:22 -04:00
return True
2020-08-03 00:35:43 -04:00
elif isinstance(node, ast.Name) and node.id == "Any":
return False
2021-03-15 09:57:36 -04:00
elif (
isinstance(node, ast.Attribute)
and type(node.value) == ast.Name
and node.value.id == "typing"
and node.attr == "Any"
):
2020-08-03 00:35:43 -04:00
return False
return True
2020-07-21 11:37:22 -04:00
2020-08-03 00:35:43 -04:00
def find_stub_issues(tree):
2020-07-21 11:37:22 -04:00
for node in ast.walk(tree):
2020-08-03 00:35:43 -04:00
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):
2020-08-07 12:33:24 -04:00
allargs = list(node.args + node.kwonlyargs)
if sys.version_info >= (3, 8):
allargs.extend(node.posonlyargs)
for arg_node in allargs:
2021-03-15 09:57:36 -04:00
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}",
)
2020-08-03 00:35:43 -04:00
if node.vararg and not is_typed(node.vararg.annotation, allow_any=True):
2021-03-15 09:57:36 -04:00
yield (
"WARN",
f"Missing argument type: *{node.vararg.arg} on line {node.vararg.lineno}",
)
2020-08-03 00:35:43 -04:00
if node.kwarg and not is_typed(node.kwarg.annotation, allow_any=True):
2021-03-15 09:57:36 -04:00
yield (
"WARN",
f"Missing argument type: **{node.kwarg.arg} on line {node.kwarg.lineno}",
)
2020-08-03 00:35:43 -04:00
elif isinstance(node, ast.FunctionDef):
if not is_typed(node.returns):
yield ("WARN", f"Missing return type: {node.name} on line {node.lineno}")
2020-07-21 11:37:22 -04:00
def extract_imports(tree):
modules = set()
used_type_module_imports = {module: set() for module in AVAILABLE_TYPE_MODULE_IMPORTS.keys()}
2020-07-21 11:37:22 -04:00
def collect_annotations(anno_tree):
if anno_tree is None:
return
for node in ast.walk(anno_tree):
2020-08-03 00:35:43 -04:00
if isinstance(node, ast.Name):
2020-07-21 11:37:22 -04:00
if node.id in 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)
2020-08-03 00:35:43 -04:00
elif isinstance(node, ast.Attribute):
if isinstance(node.value, ast.Name):
2020-07-23 14:22:41 -04:00
modules.add(node.value.id)
2020-07-21 11:37:22 -04:00
for node in ast.walk(tree):
2020-08-03 00:35:43 -04:00
if isinstance(node, (ast.AnnAssign, ast.arg)):
2020-07-21 11:37:22 -04:00
collect_annotations(node.annotation)
2020-08-03 00:35:43 -04:00
elif isinstance(node, ast.Assign):
collect_annotations(node.value)
elif isinstance(node, ast.FunctionDef):
2020-07-21 11:37:22 -04:00
collect_annotations(node.returns)
2020-07-23 14:22:41 -04:00
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)
2020-07-21 11:37:22 -04:00
2020-04-27 17:36:14 -04:00
2020-08-03 00:35:43 -04:00
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
2020-05-14 18:57:35 -04:00
def convert_folder(top_level, stub_directory):
ok = 0
total = 0
filenames = sorted(os.listdir(top_level))
2020-08-03 00:35:43 -04:00
stub_fragments = []
references = set()
2020-07-21 11:37:22 -04:00
2020-05-14 18:57:35 -04:00
for filename in filenames:
full_path = os.path.join(top_level, filename)
file_lines = []
if os.path.isdir(full_path):
2020-07-21 11:37:22 -04:00
(mok, mtotal) = convert_folder(full_path, os.path.join(stub_directory, filename))
2020-05-14 18:57:35 -04:00
ok += mok
total += mtotal
elif filename.endswith(".c"):
2020-08-03 00:35:43 -04:00
with open(full_path, "r", encoding="utf-8") as f:
2020-05-14 18:57:35 -04:00
for line in f:
2020-08-03 00:35:43 -04:00
line = line.rstrip()
2020-05-14 18:57:35 -04:00
if line.startswith("//|"):
2020-08-03 00:35:43 -04:00
if len(line) == 3:
line = ""
elif line[3] == " ":
2020-05-14 18:57:35 -04:00
line = line[4:]
else:
2020-08-03 00:35:43 -04:00
line = line[3:]
print("[WARN] There must be at least one space after '//|'")
2020-05-14 18:57:35 -04:00
file_lines.append(line)
elif filename.endswith(".pyi"):
with open(full_path, "r") as f:
2020-08-03 00:35:43 -04:00
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:
2020-07-21 11:37:22 -04:00
return (ok, total)
2020-05-14 18:57:35 -04:00
stub_filename = os.path.join(stub_directory, "__init__.pyi")
2020-04-27 17:36:14 -04:00
print(stub_filename)
2020-08-03 00:35:43 -04:00
stub_contents = "\n\n".join(stub_fragments)
2020-04-27 17:36:14 -04:00
2020-08-03 00:35:43 -04:00
# Validate the stub code.
2020-04-27 17:36:14 -04:00
try:
2020-07-21 11:37:22 -04:00
tree = ast.parse(stub_contents)
except SyntaxError as e:
2020-04-27 17:36:14 -04:00
traceback.print_exception(type(e), e, e.__traceback__)
2020-07-21 11:37:22 -04:00
return (ok, total)
2020-08-03 00:35:43 -04:00
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
2020-07-21 11:37:22 -04:00
# Add import statements
imports, type_imports = extract_imports(tree)
2020-07-21 11:37:22 -04:00
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))
2020-07-21 11:37:22 -04:00
import_body = "\n".join(import_lines)
m = re.match(r'(\s*""".*?""")', stub_contents, flags=re.DOTALL)
if m:
2021-03-15 09:57:36 -04:00
stub_contents = m.group(1) + "\n\n" + import_body + "\n\n" + stub_contents[m.end() :]
2020-07-21 11:37:22 -04:00
else:
stub_contents = import_body + "\n\n" + stub_contents
2020-07-23 14:22:41 -04:00
# Code formatting
stub_contents = isort.code(stub_contents)
stub_contents = black.format_str(stub_contents, mode=black.FileMode(is_pyi=True))
2020-07-21 11:37:22 -04:00
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]
2020-05-14 18:57:35 -04:00
2020-07-21 11:37:22 -04:00
(ok, total) = convert_folder(top_level, stub_directory)
2020-04-27 17:36:14 -04:00
2020-07-21 11:37:22 -04:00
print(f"Parsing .pyi files: {total - ok} failed, {ok} passed")
2020-07-21 11:37:22 -04:00
if ok != total:
sys.exit(total - ok)