From c83302ad67e68902c2f19db7f36637a8f0d5af87 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Lars=20Gr=C3=BCter?= Date: Wed, 23 Apr 2025 00:14:12 +0200 Subject: [PATCH] Support and default to inplace stub generation --- pyproject.toml | 1 + src/docstub/_cli.py | 31 +++++++++++--------- src/docstub/_stubs.py | 66 +++++++++++++++++++++++++++++++------------ 3 files changed, 67 insertions(+), 31 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index e5e2e4a..e562caf 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -90,6 +90,7 @@ ignore = [ "RET504", # Assignment before `return` statement facilitates debugging "PTH123", # Using builtin open() instead of Path.open() is fine "SIM108", # Terniary operator is always more readable + "SIM103", # Don't recommend returning the condition directly ] diff --git a/src/docstub/_cli.py b/src/docstub/_cli.py index 2a29a83..c22820e 100644 --- a/src/docstub/_cli.py +++ b/src/docstub/_cli.py @@ -16,9 +16,10 @@ from ._cache import FileCache from ._config import Config from ._stubs import ( + STUB_HEADER_COMMENT, Py2StubTransformer, try_format_stub, - walk_source, + walk_python_package, walk_source_and_targets, ) from ._utils import ErrorReporter, GroupedErrorReporter @@ -27,9 +28,6 @@ logger = logging.getLogger(__name__) -STUB_HEADER_COMMENT = "# File generated with docstub" - - def _load_configuration(config_path=None): """Load and merge configuration from CWD and optional files. @@ -78,13 +76,13 @@ def _setup_logging(*, verbose): ) -def _build_import_map(config, source_dir): +def _build_import_map(config, root_path): """Build a map of known imports. Parameters ---------- config : ~.Config - source_dir : Path + root_path : Path Returns ------- @@ -98,10 +96,11 @@ def _build_import_map(config, source_dir): cache_dir=Path.cwd() / ".docstub_cache", name=f"{__version__}/collected_types", ) - for source_path in walk_source(source_dir): - logger.info("collecting types in %s", source_path) - known_imports_in_source = collect_cached_types(source_path) - known_imports.update(known_imports_in_source) + if root_path.is_dir(): + for source_path in walk_python_package(root_path): + logger.info("collecting types in %s", source_path) + known_imports_in_source = collect_cached_types(source_path) + known_imports.update(known_imports_in_source) known_imports.update(KnownImport.many_from_config(config.known_imports)) @@ -138,7 +137,7 @@ def report_execution_time(): "--out-dir", type=click.Path(file_okay=False), metavar="PATH", - help="Set output directory explicitly.", + help="Set output directory explicitly. Otherwise, stubs are generated inplace.", ) @click.option( "--config", @@ -176,7 +175,7 @@ def main(root_path, out_dir, config_path, group_errors, allow_errors, verbose): Parameters ---------- - source_dir : Path + root_path : Path out_dir : Path config_path : Path group_errors : bool @@ -189,6 +188,12 @@ def main(root_path, out_dir, config_path, group_errors, allow_errors, verbose): _setup_logging(verbose=verbose) root_path = Path(root_path) + if root_path.is_file(): + logger.warning( + "Running docstub on a single file is experimental. Relative imports " + "or type references won't work." + ) + config = _load_configuration(config_path) known_imports = _build_import_map(config, root_path) @@ -204,7 +209,7 @@ def main(root_path, out_dir, config_path, group_errors, allow_errors, verbose): if root_path.is_file(): out_dir = root_path.parent else: - out_dir = root_path.parent / (root_path.name + "-stubs") + out_dir = root_path out_dir = Path(out_dir) out_dir.mkdir(parents=True, exist_ok=True) diff --git a/src/docstub/_stubs.py b/src/docstub/_stubs.py index 4b8d9da..46854cf 100644 --- a/src/docstub/_stubs.py +++ b/src/docstub/_stubs.py @@ -1,7 +1,13 @@ -"""Transform Python source files to typed stub files.""" +"""Transform Python source files to typed stub files. + +Attributes +---------- +STUB_HEADER_COMMENT : Final[str] +""" import enum import logging +import re from dataclasses import dataclass from functools import wraps from typing import ClassVar @@ -16,7 +22,10 @@ logger = logging.getLogger(__name__) -def _is_python_package(path): +STUB_HEADER_COMMENT = "# File generated with docstub" + + +def is_python_package(path): """ Parameters ---------- @@ -30,8 +39,31 @@ def _is_python_package(path): return is_package -def walk_source(root_dir): - """Iterate modules in a Python package and its target stub files. +def is_docstub_generated(path): + """Check if the stub file was generated by docstub. + + Parameters + ---------- + path : Path + + Returns + ------- + is_generated : bool + """ + assert path.suffix == ".pyi" + with path.open("r") as fo: + content = fo.read() + if re.match(f"^{re.escape(STUB_HEADER_COMMENT)}", content): + return True + return False + + +def walk_python_package(root_dir): + """Iterate source files in a Python package. + + Given a Python package, yield the path of contained Python modules. If an + alternate stub file already exists and isn't generated by docstub, it is + returned instead. Parameters ---------- @@ -43,26 +75,24 @@ def walk_source(root_dir): source_path : Path Either a Python file or a stub file that takes precedence. """ - queue = [root_dir] - while queue: - path = queue.pop(0) - + for path in root_dir.iterdir(): if path.is_dir(): - if _is_python_package(path): - queue.extend(path.iterdir()) + if is_python_package(path): + yield from walk_python_package(path) else: - logger.debug("skipping directory %s", path) + logger.debug("skipping directory %s which isn't a Python package", path) continue assert path.is_file() - suffix = path.suffix.lower() - if suffix not in {".py", ".pyi"}: - continue - if suffix == ".py" and path.with_suffix(".pyi").exists(): - continue # Stub file already exists and takes precedence - yield path + if suffix == ".py": + stub = path.with_suffix(".pyi") + if stub.exists() and not is_docstub_generated(stub): + # Non-generated stub file already exists and takes precedence + yield stub + else: + yield path def walk_source_and_targets(root_path, target_dir): @@ -87,7 +117,7 @@ def walk_source_and_targets(root_path, target_dir): yield root_path, stub_path return - for source_path in walk_source(root_path): + for source_path in walk_python_package(root_path): stub_path = target_dir / source_path.with_suffix(".pyi").relative_to(root_path) yield source_path, stub_path