Add --workers option controlling parallel stub generation (#110)

Add the new experimental `--workers` option to `docstub run` to generate stub files
in parallel. This can reduce runtime significantly for projects with many files.
For now, multiprocessing has to be explicitly enabled by requesting a number of
workers.

Also ensure that docstub still returns a non-zero exit code when it is run with output
completely disabled (`-qq`). The same applies when running docstub with 
`--quiet --fail-on-warnings`.

--

* Draft multiprocessing with PoolExecutor

* Count errors in dedicated handler

Splitting these two responsibilities (logging to console and counting)
simplifies both in context of multiprocessing.

* Update how stats are propagated back to the main process

* Add critical method to ContextReporter

* Use "--workers" instead of "--jobs"

The former is what scikit-image agreed on.

* Test error and warning propagation with -qq

Mark these as "slow" (because they are)

* Add reminders to remove vendored stdlib code eventually

* Make docstest independent of available CPUs

* Add tests for `guess_concurrency_params`

* Fix swallowing exit code for diffs

* Improve exit function of LoggingProcessExecutor

* Fix stubtest errors

* Update example_pkg-stubs

* Sync _concurrency.pyi

* Mark --worker option as experimental and default to single process

This is probably safest for now. I don't want to break peoples setup
the new version of docstub suddenly hangs or fails. This way they users
have to opt in explicitly and will now what might have caused it.

* Fix typo in docstest caused by accidental refactor

* Test multiprocessing in CI

* Factor CI multiprocessing step into its own step

* Fix test setup for monkeypatching process_cpu_count

* Skip subprocess tests on Linux 3.12-3.13

* Start logging queue is in same multiprocessing context as pool

This fixes the following:

    RuntimeError: A SemLock created in a fork context is being shared
    with a process in a spawn context. This is not supported. Please use
    the same context to create multiprocessing objects and Process.

* Fix unkown "Queue" in doctype

Author: lagru

.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