From d85c54dc4708b4b2a85b5322809b6a19633f43a3 Mon Sep 17 00:00:00 2001 From: Jukka Lehtosalo Date: Thu, 15 Jul 2021 11:22:28 +0100 Subject: [PATCH] Add simple tool for creating baseline stubs (#5777) This tool is a simple wrapper around stubgen and other tools that simplifies the creation of "baseline" stubs with minimal annotation coverage. These stubs allow basic type checking (e.g. number of arguments, whether X is valid as a base class), and once we have baseline stubs, it's easy to contribute incremental changes that add type annotations. Here's a summary of what the tool does: 1. Check that the package is installed 2. Run stubgen 3. Copy generated stubs to the correct directory under `stubs/` 4. Run black 5. Run isort 6. Generate basic metadata with correct package version 7. Update pyright exclusions (needed since there are missing types) 8. Print suggestions about next steps needed to contribute stubs to typeshed For example, to generate stubs for `iso8601`, you can run it like this: ``` python scripts/create_baseline_stubs.py iso8601 ``` Sometimes the project name is different from the runtime package name. In this case run the tool with `--package `. --- scripts/create_baseline_stubs.py | 173 +++++++++++++++++++++++++++++++ 1 file changed, 173 insertions(+) create mode 100644 scripts/create_baseline_stubs.py diff --git a/scripts/create_baseline_stubs.py b/scripts/create_baseline_stubs.py new file mode 100644 index 000000000..e95536da4 --- /dev/null +++ b/scripts/create_baseline_stubs.py @@ -0,0 +1,173 @@ +"""Script to generate unannotated baseline stubs using stubgen. + +Basic usage: +$ python3 scripts/create_baseline_stubs.py + +Run with -h for more help. +""" + +import argparse +import os +import re +import shutil +import subprocess +import sys +from typing import Optional, Tuple + +PYRIGHT_CONFIG = "pyrightconfig.stricter.json" + + +def search_pip_freeze_output(project: str, output: str) -> Optional[Tuple[str, str]]: + # Look for lines such as "typed-ast==1.4.2". '-' matches '_' and + # '_' matches '-' in project name, so that "typed_ast" matches + # "typed-ast", and vice versa. + regex = "^(" + re.sub(r"[-_]", "[-_]", project) + ")==(.*)" + m = re.search(regex, output, flags=re.IGNORECASE | re.MULTILINE) + if not m: + return None + return m.group(1), m.group(2) + + +def get_installed_package_info(project: str) -> Optional[Tuple[str, str]]: + """Find package information from pip freeze output. + + Match project name somewhat fuzzily (case sensitive; '-' matches '_', and + vice versa). + + Return (normalized project name, installed version) if successful. + """ + r = subprocess.run(["pip", "freeze"], capture_output=True, text=True, check=True) + return search_pip_freeze_output(project, r.stdout) + + +def run_stubgen(package: str) -> None: + print(f"Running stubgen: stubgen -p {package}") + subprocess.run(["python", "-m", "mypy.stubgen", "-p", package], check=True) + + +def copy_stubs(src_base_dir: str, package: str, stub_dir: str) -> None: + """Copy generated stubs to the target directory under stub_dir/.""" + print(f"Copying stubs to {stub_dir}") + if not os.path.isdir(stub_dir): + os.mkdir(stub_dir) + src_dir = os.path.join(src_base_dir, package) + if os.path.isdir(src_dir): + shutil.copytree(src_dir, os.path.join(stub_dir, package)) + else: + src_file = os.path.join("out", package + ".pyi") + if not os.path.isfile(src_file): + sys.exit("Error: Cannot find generated stubs") + shutil.copy(src_file, stub_dir) + + +def run_black(stub_dir: str) -> None: + print(f"Running black: black {stub_dir}") + subprocess.run(["black", stub_dir]) + + +def run_isort(stub_dir: str) -> None: + print(f"Running isort: isort --recursive {stub_dir}") + subprocess.run(["isort", "--recursive", stub_dir]) + + +def create_metadata(stub_dir: str, version: str) -> None: + """Create a METADATA.toml file.""" + m = re.match(r"[0-9]+.[0-9]+", version) + if m is None: + sys.exit(f"Error: Cannot parse version number: {version}") + fnam = os.path.join(stub_dir, "METADATA.toml") + version = m.group(0) + assert not os.path.exists(fnam) + print(f"Writing {fnam}") + with open(fnam, "w") as f: + f.write(f'version = "{version}"\n') + + +def add_pyright_exclusion(stub_dir: str) -> None: + """Exclude stub_dir from strict pyright checks.""" + with open(PYRIGHT_CONFIG) as f: + lines = f.readlines() + i = 0 + while i < len(lines) and not lines[i].strip().startswith('"exclude": ['): + i += 1 + assert i < len(lines), f"Error parsing {PYRIGHT_CONFIG}" + while not lines[i].strip().startswith("]"): + i += 1 + line_to_add = f' "{stub_dir}",' + initial = i - 1 + while lines[i].lower() > line_to_add.lower(): + i -= 1 + if lines[i + 1].strip().rstrip(",") == line_to_add.strip().rstrip(","): + print(f"{PYRIGHT_CONFIG} already up-to-date") + return + if i == initial: + # Special case: when adding to the end of the list, commas need tweaking + line_to_add = line_to_add.rstrip(",") + lines[i] = lines[i].rstrip() + ",\n" + lines.insert(i + 1, line_to_add + "\n") + print(f"Updating {PYRIGHT_CONFIG}") + with open(PYRIGHT_CONFIG, "w") as f: + f.writelines(lines) + + +def main() -> None: + parser = argparse.ArgumentParser( + description="""Generate baseline stubs automatically for an installed pip package + using stubgen. Also run black and isort. If the name of + the project is different from the runtime Python package name, you must + also use --package (example: --package yaml PyYAML).""" + ) + parser.add_argument("project", help="name of PyPI project for which to generate stubs under stubs/") + parser.add_argument("--package", help="generate stubs for this Python package (defaults to project)") + args = parser.parse_args() + project = args.project + package = args.package + + if not re.match(r"[a-zA-Z0-9-_.]+$", project): + sys.exit(f"Invalid character in project name: {project!r}") + + if not package: + package = project # TODO: infer from installed files + + if not os.path.isdir("stubs") or not os.path.isdir("stdlib"): + sys.exit("Error: Current working directory must be the root of typeshed repository") + + # Get normalized project name and version of installed package. + info = get_installed_package_info(project) + if info is None: + print(f'Error: "{project}" is not installed', file=sys.stderr) + print("", file=sys.stderr) + print(f'Suggestion: Run "python3 -m pip install {project}" and try again', file=sys.stderr) + sys.exit(1) + project, version = info + + stub_dir = os.path.join("stubs", project) + if os.path.exists(stub_dir): + sys.exit(f"Error: {stub_dir} already exists (delete it first)") + + run_stubgen(package) + + # Stubs were generated under out/. Copy them to stubs/. + copy_stubs("out", package, stub_dir) + + run_black(stub_dir) + run_isort(stub_dir) + + create_metadata(stub_dir, version) + + # Since the generated stubs won't have many type annotations, we + # have to exclude them from strict pyright checks. + add_pyright_exclusion(stub_dir) + + print("\nDone!\n\nSuggested next steps:") + print(f" 1. Manually review the generated stubs in {stub_dir}") + print(" 2. Run stubtest to test the generated stubs against runtime definitions") + print(f' 3. Run "flake8 {stub_dir}" to check code style') + print(f' 4. Run "mypy {stub_dir}" to check for errors') + print(f' 5. Run "black {stub_dir}" (if you\'ve made code changes)') + print(" 6. Create branch in the typeshed repo and commit the stubs (and other changes)") + print(" 7. Create typeshed PR") + + +if __name__ == "__main__": + main()