Support and default to inplace stub generation

Author: lagru

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
 ]
 
 

src/docstub/_cli.py

@@ -15,19 +15,17 @@
 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__
 
 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)
 

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