Enable generating stubs in parallel (multiprocessing)

.github/scripts/assert-unchanged.sh

@@ -15,8 +15,10 @@ UNTRACKED=$(git ls-files --others --exclude-standard "$CHECK_DIR")
 echo "$UNTRACKED" | xargs -I _ git --no-pager diff /dev/null _ || true
 
 # Display changes in tracked files and capture non-zero exit code if so
-git diff --exit-code HEAD "$CHECK_DIR" || true
+set +e
+git diff --exit-code HEAD "$CHECK_DIR"
 GIT_DIFF_HEAD_EXIT_CODE=$?
+set -e
 
 # Display changes in tracked files and capture exit status
 if [ $GIT_DIFF_HEAD_EXIT_CODE -ne 0 ] ||  [ -n "$UNTRACKED" ]; then

.github/workflows/ci.yml

@@ -89,14 +89,21 @@ jobs:
               examples/example_pkg
           .github/scripts/assert-unchanged.sh examples/
 
-      - name: Check docstub-stubs
+      - name: Check docstub-stubs (single process)
         # Check that stubs for docstub are up-to-date by regenerating them
         # with docstub and looking for differences.
         run: |
           rm -rf src/docstub-stubs
           python -m docstub run -v src/docstub -o src/docstub-stubs
           .github/scripts/assert-unchanged.sh src/docstub-stubs/
 
+      - name: Check docstub-stubs (multiprocess)
+        # Repeat test with multiprocessing enabled
+        run: |
+          rm -rf src/docstub-stubs
+          python -m docstub run -v src/docstub -o src/docstub-stubs --workers 2
+          .github/scripts/assert-unchanged.sh src/docstub-stubs/
+
       - name: Check with mypy.stubtest
         run: |
           python -m mypy.stubtest \

REMINDERS.md

@@ -0,0 +1,10 @@
+# Reminders
+
+## With Python >=3.13
+
+Remove vendored `glob_translate` in `docstub._vendored.stdlib`.
+
+
+## With Python >=3.14
+
+Remove vendored `ProcessPoolExecutor` in `docstub._vendored.stdlib`.

docs/command_line.md

@@ -49,10 +49,6 @@ Options:
           Set output directory explicitly. Stubs will be directly written into
           that directory while preserving the directory structure under
           PACKAGE_PATH. Otherwise, stubs are generated inplace.
-      --config PATH
-          Set one or more configuration file(s) explicitly. Otherwise, it will
-          look for a `pyproject.toml` or `docstub.toml` in the current
-          directory.
       --ignore GLOB
           Ignore files matching this glob-style pattern. Can be used multiple
           times.
@@ -67,8 +63,15 @@ Options:
   -W, --fail-on-warning
           Return non-zero exit code when a warning is raised. Will add to
           --allow-errors.
+      --workers INT
+          Experimental: Process files in parallel with the desired number of
+          workers. By default, no multiprocessing is used.  [default: 1]
       --no-cache
           Ignore pre-existing cache and don't create a new one.
+      --config PATH
+          Set one or more configuration file(s) explicitly. Otherwise, it will
+          look for a `pyproject.toml` or `docstub.toml` in the current
+          directory.
   -v, --verbose
           Print more details. Use once to show information messages. Use -vv to
           print debug messages.

examples/example_pkg-stubs/_basic.pyi

@@ -17,7 +17,7 @@ __all__ = [
     "func_empty",
 ]
 
-def func_empty(a1: Incomplete, a2: Incomplete, a3) -> None: ...
+def func_empty(a1: Incomplete, a2: Incomplete, a3: Incomplete) -> None: ...
 def func_contains(
     a1: list[float],
     a2: dict[str, Union[int, str]],

pyproject.toml

@@ -124,6 +124,9 @@ xfail_strict = true
 filterwarnings = ["error"]
 log_cli_level = "info"
 testpaths = ["src", "tests"]
+markers = [
+    "slow: marks tests as slow (deselect with `-m 'not slow'`)",
+]
 
 
 [tool.coverage]

src/docstub-stubs/_cli.pyi

@@ -16,6 +16,7 @@ from _typeshed import Incomplete
 from ._analysis import PyImport, TypeCollector, TypeMatcher, common_known_types
 from ._cache import CACHE_DIR_NAME, FileCache, validate_cache
 from ._cli_help import HelpFormatter
+from ._concurrency import LoggingProcessExecutor, guess_concurrency_params
 from ._config import Config
 from ._path_utils import (
     STUB_HEADER_COMMENT,
@@ -25,6 +26,7 @@ from ._path_utils import (
 )
 from ._report import setup_logging
 from ._stubs import Py2StubTransformer, try_format_stub
+from ._utils import update_with_add_values
 from ._version import __version__
 
 logger: logging.Logger
@@ -45,6 +47,9 @@ click.Context.formatter_class = HelpFormatter
 @click.group()
 def cli() -> None: ...
 def _add_verbosity_options(func: Callable) -> Callable: ...
+def _transform_to_stub(
+    task: tuple[Path, Path, Py2StubTransformer],
+) -> dict[str, int | list[str]]: ...
 @cli.command()
 def run(
     *,
@@ -55,6 +60,7 @@ def run(
     group_errors: bool,
     allow_errors: int,
     fail_on_warning: bool,
+    desired_worker_count: int,
     no_cache: bool,
     verbose: int,
     quiet: int,

src/docstub-stubs/_concurrency.pyi

@@ -0,0 +1,48 @@
+# File generated with docstub
+
+import logging
+import logging.handlers
+import math
+import multiprocessing
+import os
+from collections.abc import Callable, Iterator
+from concurrent.futures import Executor
+from dataclasses import dataclass
+from types import TracebackType
+from typing import Any
+
+from ._vendored.stdlib import ProcessPoolExecutor
+
+logger: logging.Logger
+
+class MockPoolExecutor(Executor):
+    def map[T](
+        self, fn: Callable[..., T], *iterables: Any, **__: Any
+    ) -> Iterator[T]: ...
+
+@dataclass(kw_only=True)
+class LoggingProcessExecutor:
+
+    max_workers: int | None = ...
+    logging_handlers: tuple[logging.Handler, ...] = ...
+    initializer: Callable | None = ...
+    initargs: tuple | None = ...
+
+    @staticmethod
+    def _initialize_worker(
+        queue: multiprocessing.Queue,
+        worker_log_level: int,
+        initializer: Callable,
+        initargs: tuple[Any],
+    ) -> None: ...
+    def __enter__(self) -> ProcessPoolExecutor | MockPoolExecutor: ...
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType,
+    ) -> bool: ...
+
+def guess_concurrency_params(
+    *, task_count: int, desired_worker_count: int | None = ...
+) -> tuple[int, int]: ...

src/docstub-stubs/_report.pyi

@@ -47,6 +47,9 @@ class ContextReporter:
     def error(
         self, short: str, *args: Any, details: str | None = ..., **log_kw: Any
     ) -> None: ...
+    def critical(
+        self, short: str, *args: Any, details: str | None = ..., **log_kw: Any
+    ) -> None: ...
     def __post_init__(self) -> None: ...
     @staticmethod
     def underline(line: str, *, char: str = ...) -> str: ...
@@ -62,9 +65,17 @@ class ReportHandler(logging.StreamHandler):
         self, stream: TextIO | None = ..., group_errors: bool = ...
     ) -> None: ...
     def format(self, record: logging.LogRecord) -> str: ...
-    def emit(self, record: logging.LogRecord) -> None: ...
+    def handle(self, record: logging.LogRecord) -> bool: ...
     def emit_grouped(self) -> None: ...
 
+class LogCounter(logging.NullHandler):
+    critical_count: int
+    error_count: int
+    warning_count: int
+
+    def __init__(self) -> None: ...
+    def handle(self, record: logging.LogRecord) -> bool: ...
+
 def setup_logging(
     *, verbosity: Literal[-2, -1, 0, 1, 2, 3], group_errors: bool
-) -> ReportHandler: ...
+) -> tuple[ReportHandler, LogCounter]: ...

src/docstub-stubs/_stubs.pyi

@@ -19,7 +19,7 @@ from ._docstrings import (
     FallbackAnnotation,
 )
 from ._report import ContextReporter
-from ._utils import module_name_from_path
+from ._utils import module_name_from_path, update_with_add_values
 
 logger: logging.Logger
 
@@ -73,6 +73,9 @@ class Py2StubTransformer(cst.CSTTransformer):
     @property
     def is_inside_function_def(self) -> bool: ...
     def python_to_stub(self, source: str, *, module_path: Path | None = ...) -> str: ...
+    def collect_stats(
+        self, *, reset_after: bool = ...
+    ) -> dict[str, int | list[str]]: ...
     def visit_ClassDef(self, node: cst.ClassDef) -> Literal[True]: ...
     def leave_ClassDef(
         self, original_node: cst.ClassDef, updated_node: cst.ClassDef

src/docstub-stubs/_utils.pyi

@@ -2,7 +2,7 @@
 
 import itertools
 import re
-from collections.abc import Callable
+from collections.abc import Callable, Hashable, Mapping, Sequence
 from functools import lru_cache, wraps
 from pathlib import Path
 from zlib import crc32
@@ -12,6 +12,9 @@ def escape_qualname(name: str) -> str: ...
 def _resolve_path_before_caching(func: Callable) -> Callable: ...
 def module_name_from_path(path: Path) -> str: ...
 def pyfile_checksum(path: Path) -> str: ...
+def update_with_add_values(
+    *mappings: Mapping[Hashable, int | Sequence], out: dict | None = ...
+) -> dict: ...
 
 class DocstubError(Exception):
     pass

src/docstub-stubs/_vendored/stdlib.pyi

@@ -3,6 +3,7 @@
 import os
 import re
 from collections.abc import Sequence
+from concurrent.futures import ProcessPoolExecutor as _ProcessPoolExecutor
 
 def _fnmatch_translate(pat: str, STAR: str, QUESTION_MARK: str) -> str: ...
 def glob_translate(
@@ -12,3 +13,17 @@ def glob_translate(
     include_hidden: bool = ...,
     seps: Sequence[str] | None = ...,
 ) -> str: ...
+
+if not hasattr(_ProcessPoolExecutor, "terminate_workers"):
+    _TERMINATE: str
+    _KILL: str
+
+    _SHUTDOWN_CALLBACK_OPERATION: set[str]
+
+    class ProcessPoolExecutor(_ProcessPoolExecutor):
+        def _force_shutdown(self, operation: str) -> None: ...
+        def terminate_workers(self) -> None: ...
+        def kill_workers(self) -> None: ...
+
+else:
+    ProcessPoolExecutor: _ProcessPoolExecutor  # type: ignore[no-redef]

src/docstub/_analysis.py

@@ -500,12 +500,14 @@ def __init__(
         type_prefixes : dict[str, PyImport]
         type_nicknames : dict[str, str]
         """
-
         self.types = common_known_types() | (types or {})
         self.type_prefixes = type_prefixes or {}
         self.type_nicknames = type_nicknames or {}
-        self.successful_queries = 0
-        self.unknown_qualnames = []
+
+        self.stats = {
+            "matched_type_names": 0,
+            "unknown_type_names": [],
+        }
 
         self.current_file = None
 
@@ -621,8 +623,8 @@ def match(self, search):
             type_name = type_name[type_name.find(py_import.target) :]
 
         if type_name is not None:
-            self.successful_queries += 1
+            self.stats["matched_type_names"] += 1
         else:
-            self.unknown_qualnames.append(search)
+            self.stats["unknown_type_names"].append(search)
 
         return type_name, py_import