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

@@ -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)
 

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