Improve walking of the source directory

src/docstub/_cli.py

@@ -17,8 +17,8 @@
 from ._config import Config
 from ._path_utils import (
     STUB_HEADER_COMMENT,
-    walk_python_package,
     walk_source_and_targets,
+    walk_source_package,
 )
 from ._stubs import Py2StubTransformer, try_format_stub
 from ._utils import ErrorReporter, GroupedErrorReporter, module_name_from_path
@@ -100,7 +100,7 @@ def _collect_type_info(root_path, *, ignore=()):
     type_prefixes = {}
 
     if root_path.is_dir():
-        for source_path in walk_python_package(root_path, ignore=ignore):
+        for source_path in walk_source_package(root_path, ignore=ignore):
 
             module = module_name_from_path(source_path)
             module = module.replace(".", "/")
@@ -228,8 +228,8 @@ def run(root_path, out_dir, config_paths, ignore, group_errors, allow_errors, ve
     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."
+            "Running docstub on a single file. Relative imports "
+            "or type references outside this file won't work."
         )
 
     config = _load_configuration(config_paths)

src/docstub/_path_utils.py

@@ -17,27 +17,68 @@
 STUB_HEADER_COMMENT = "# File generated with docstub"
 
 
-def is_docstub_generated(path):
+def is_docstub_generated(stub_path):
     """Check if the stub file was generated by docstub.
 
     Parameters
     ----------
-    path : Path
+    stub_path : Path
+        Path to a stub file.
 
     Returns
     -------
     is_generated : bool
+
+    Examples
+    --------
+    >>> from pathlib import Path
+    >>> from docstub import _version
+    >>> is_docstub_generated(Path(_version.__file__).with_suffix(".pyi"))
+    False
+
+    >>> is_docstub_generated(Path(__file__))
+    Traceback (most recent call last):
+        ...
+    TypeError: expected stub file (ending with '.pyi'), ...
     """
-    assert path.suffix == ".pyi"
-    with path.open("r") as fo:
+    if stub_path.suffix != ".pyi":
+        raise TypeError(f"expected stub file (ending with '.pyi'), got {stub_path}")
+    with stub_path.open("r") as fo:
         content = fo.read()
     if re.match(f"^{re.escape(STUB_HEADER_COMMENT)}", content):
         return True
     return False
 
 
-def is_python_package(path):
+def is_python_or_stub_file(path):
+    """Check whether `path` is a Python source file.
+
+    Parameters
+    ----------
+    path : Path
+
+    Returns
+    -------
+    is_python_or_stub_file : bool
+
+    See Also
+    --------
+    is_python_package_dir
+
+    Examples
+    --------
+    >>> from pathlib import Path
+    >>> is_python_or_stub_file(Path(__file__))
+    True
+    >>> is_python_or_stub_file(Path(__file__).parent)
+    False
     """
+    return path.is_file() and path.suffix in (".py", ".pyi")
+
+
+def is_python_package_dir(path):
+    """Check whether `path` is a valid Python package and a directory.
+
     Parameters
     ----------
     path : Path
@@ -46,14 +87,18 @@ def is_python_package(path):
     -------
     is_package : bool
 
+    See Also
+    --------
+    is_python_or_stub_file
+
     Examples
     --------
     >>> from pathlib import Path
-    >>> is_python_package(Path(__file__))
+    >>> is_python_package_dir(Path(__file__))
     False
-    >>> is_python_package(Path(__file__).parent)
+    >>> is_python_package_dir(Path(__file__).parent)
     True
-    >>> is_python_package(Path(__file__).parent.parent)
+    >>> is_python_package_dir(Path(__file__).parent.parent)
     False
     """
     has_init = (path / "__init__.py").is_file() or (path / "__init__.pyi").is_file()
@@ -84,7 +129,7 @@ def find_package_root(path):
         root = root.parent
 
     for _ in range(2**16):
-        if not is_python_package(root):
+        if not is_python_package_dir(root):
             logger.debug("detected %s as the package root of %s", root, path)
             return root
         root = root.parent
@@ -95,7 +140,7 @@ def find_package_root(path):
 
 @lru_cache(maxsize=10)
 def glob_patterns_to_regex(patterns, relative_to=None):
-    r"""Combine glob-style patterns into a single regex.
+    r"""Combine glob-style patterns into a single regex [1].
 
     Parameters
     ----------
@@ -106,6 +151,10 @@ def glob_patterns_to_regex(patterns, relative_to=None):
     -------
     regex : re.Pattern | None
 
+    References
+    ----------
+    .. [1] https://docs.python.org/3/library/glob.html#glob.translate
+
     Examples
     --------
     >>> from pathlib import Path
@@ -143,17 +192,59 @@ def prefix(pattern):
     return regex
 
 
-def walk_python_package(root_dir, *, ignore=()):
+def _walk_source_package(path, *, ignore_regex):
     """Iterate source files in a Python package.
 
+    .. note::
+        Inner function of :func:`walk_source_package`. See that function
+        for more details.
+
+    Parameters
+    ----------
+    path : Path
+        Root directory of a Python package. Can also be a single Python or stub
+        file.
+    ignore_regex : re.Pattern
+        Don't yield files matching this regex-compiled glob-like pattern.
+
+    Yields
+    ------
+    source_path : Path
+        Either a Python file or a stub file that takes precedence.
+    """
+    if ignore_regex and ignore_regex.match(str(path)):
+        logger.info("ignoring %s", path)
+        return
+
+    if is_python_package_dir(path):
+        for sub_path in path.iterdir():
+            yield from _walk_source_package(sub_path, ignore_regex=ignore_regex)
+
+    elif is_python_or_stub_file(path):
+        stub_path = path.with_suffix(".pyi")
+        if stub_path == path or not stub_path.is_file():
+            # If `path` is a stub file return it. If it is a regular Python
+            # file, only return it if no corresponding stub file exists.
+            yield path
+
+    elif path.is_dir():
+        logger.debug("skipping directory %s which isn't a Python package", path)
+
+    elif path.is_file():
+        logger.debug("skipping non-Python file %s", path)
+
+
+def walk_source_package(path, *, ignore=()):
+    """Iterate over a source package for docstub.
+
     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
     ----------
-    root_dir : Path
-        Root directory of a Python package.
+    path : Path
+        A Python package, either a directory or a single file.
     ignore : Sequence[str], optional
         Don't yield files matching these glob-like patterns. The pattern is
         interpreted relative to the root of the Python package unless it starts
@@ -163,36 +254,61 @@ def walk_python_package(root_dir, *, ignore=()):
     Yields
     ------
     source_path : Path
-        Either a Python file or a stub file that takes precedence.
+        Either a Python file or a stub file that takes precedence. Note that
+        stub files generated by docstub itself are not returned.
+
+    Raises
+    ------
+    TypeError
+        If `path` is not a valid Python package. Note that a single
+        Python file is considered a "package".
+
+    See Also
+    --------
+    walk_source_and_targets
+
+    Examples
+    --------
+    >>> from pathlib import Path
+    >>> this_file = Path(__file__)
+
+    Walk `path` to current file
+    >>> package_files = sorted(walk_source_package(this_file))
+    >>> len(package_files)
+    1
+    >>> package_files[0].as_posix()
+    '.../docstub/_path_utils.py'
+
+    Walk `path` to directory of current file
+    >>> package_files = walk_source_package(this_file.parent)
+    >>> sorted(package_files)
+    [.../docstub/__init__.py'), ...]
+
+    Ignoring all files ending with '.py' will return nothing
+    >>> next(walk_source_package(this_file.parent, ignore=("*.py")))
+    Traceback (most recent call last):
+        ...
+    StopIteration
     """
-    package_root = find_package_root(root_dir)
-    regex = glob_patterns_to_regex(tuple(ignore), relative_to=package_root)
+    if not is_python_package_dir(path) and not is_python_or_stub_file(path):
+        raise TypeError(f"{path} must be a Python file or package")
 
-    if regex and regex.match(str(root_dir)):
-        logger.info("ignoring %s", root_dir)
-        return
+    regex = glob_patterns_to_regex(tuple(ignore), relative_to=path)
 
-    for path in root_dir.iterdir():
-        if regex and regex.match(str(path.resolve())):
-            logger.info("ignoring %s", path)
-            continue
-        if path.is_dir():
-            if is_python_package(path):
-                yield from walk_python_package(path, ignore=ignore)
-            else:
-                logger.debug("skipping directory %s which isn't a Python package", path)
-            continue
-
-        assert path.is_file()
-        suffix = path.suffix.lower()
-
-        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
+    if is_python_or_stub_file(path):
+        stub_file = path.with_suffix(".pyi")
+        if (
+            stub_file != path
+            and stub_file.is_file()
+            and not is_docstub_generated(stub_file)
+        ):
+            # Special case: `path` is a Python file for which a stub file
+            # exists, we want to return that one while taking into account
+            # `ignore` and other logic. A simple way to do so is to just pass
+            # the stub file instead of `path`.
+            path = stub_file
+
+    yield from _walk_source_package(path, ignore_regex=regex)
 
 
 def walk_source_and_targets(root_path, target_dir, *, ignore=()):
@@ -216,12 +332,37 @@ def walk_source_and_targets(root_path, target_dir, *, ignore=()):
         Either a Python file or a stub file that takes precedence.
     stub_path : Path
         Target stub file.
+
+    Raises
+    ------
+    TypeError
+        If `root_path` is not a valid Python package. Note that a single
+        Python file is considered a "package".
+
+    See Also
+    --------
+    walk_source_package
+
+    Examples
+    --------
+    >>> from pathlib import Path
+    >>> current_root = Path(__file__).parent
+    >>> sources_n_targets = sorted(
+    ...     walk_source_and_targets(current_root, target_dir=current_root)
+    ... )
+    >>> source_path, stub_path = sources_n_targets[0]
+    >>> source_path.as_posix()
+    '.../docstub/__init__.py'
+    >>> stub_path.as_posix()
+    '.../docstub/__init__.pyi'
+    >>> stub_path.is_file()
+    False
     """
     if root_path.is_file():
         stub_path = target_dir / root_path.with_suffix(".pyi").name
         yield root_path, stub_path
         return
 
-    for source_path in walk_python_package(root_path, ignore=ignore):
+    for source_path in walk_source_package(root_path, ignore=ignore):
         stub_path = target_dir / source_path.with_suffix(".pyi").relative_to(root_path)
         yield source_path, stub_path