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 cb68943..28152c8 100644 --- a/src/docstub/_cli.py +++ b/src/docstub/_cli.py @@ -15,9 +15,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 ._version import __version__ @@ -25,9 +26,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. @@ -96,7 +94,7 @@ 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): + for source_path in walk_python_package(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) @@ -135,7 +133,7 @@ def report_execution_time(): "-o", "--out-dir", type=click.Path(file_okay=False), - help="Set output directory explicitly.", + help="Set output directory explicitly. Otherwise, stubs are generated inplace.", ) @click.option( "--config", @@ -170,7 +168,7 @@ def main(source_dir, out_dir, config_path, verbose): ) if not out_dir: - out_dir = source_dir.parent / (source_dir.name + "-stubs") + out_dir = source_dir 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 57c81da..dba7db5 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_dir, target_dir): @@ -75,18 +105,14 @@ def walk_source_and_targets(root_dir, target_dir): target_dir : Path Root directory in which a matching stub package will be created. - Returns - ------- + Yields + ------ source_path : Path Either a Python file or a stub file that takes precedence. stub_path : Path Target stub file. - - Notes - ----- - Files starting with "test_" are skipped entirely for now. """ - for source_path in walk_source(root_dir): + for source_path in walk_python_package(root_dir): stub_path = target_dir / source_path.with_suffix(".pyi").relative_to(root_dir) yield source_path, stub_path