scientific-python · lagru · Apr 18, 2025 · Mar 5, 2025 · Mar 5, 2025 · Mar 10, 2025
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
@@ -15,6 +15,12 @@ env:
   # Many color libraries just need this to be set to any value, but at least
   # one distinguishes color depth, where "3" -> "256-bit color".
   FORCE_COLOR: 3
+  MYPYPATH: ${{ github.workspace }}/stubs
+
+defaults:
+  run:
+    # Make sure that bash specific stuff works on Windows
+    shell: bash
 
 jobs:
   lint:
@@ -69,4 +75,8 @@ jobs:
 
       - name: Generate docstub stubs
         run: |
-          python -m docstub -v src/docstub
+          python -m docstub -v src/docstub -o ${MYPYPATH}/docstub
+
+      - name: Check docstub stubs with mypy
+        run: |
+          python -m mypy.stubtest --allowlist stubtest_allow.txt docstub
diff --git a/examples/example_pkg-stubs/__init__.pyi b/examples/example_pkg-stubs/__init__.pyi
@@ -1,4 +1,4 @@
-# Generated with docstub. Manual edits will be overwritten!
+# File generated with docstub
 
 import _numpy as np_
 from _basic import func_contains

diff --git a/examples/example_pkg-stubs/_basic.pyi b/examples/example_pkg-stubs/_basic.pyi
@@ -1,4 +1,5 @@
-# Generated with docstub. Manual edits will be overwritten!
+# File generated with docstub
+
 import configparser
 import logging
 from collections.abc import Sequence

diff --git a/examples/example_pkg-stubs/_numpy.pyi b/examples/example_pkg-stubs/_numpy.pyi
@@ -1,4 +1,5 @@
-# Generated with docstub. Manual edits will be overwritten!
+# File generated with docstub
+
 import numpy as np
 from numpy.typing import ArrayLike, NDArray
 

diff --git a/pyproject.toml b/pyproject.toml
@@ -45,6 +45,7 @@ dev = [
 test = [
   "pytest >=5.0.0",
   "pytest-cov >= 5.0.0",
+  "mypy",
 ]
 
 [project.urls]
@@ -80,6 +81,7 @@ extend-select = [
   "UP",       # pyupgrade
   "YTT",      # flake8-2020
   "EXE",      # flake8-executable
+#  "PYI",      # flake8-pyi
 ]
 ignore = [
   "PLR09",    # Too many <...>

diff --git a/src/docstub/_analysis.py b/src/docstub/_analysis.py
@@ -71,10 +71,12 @@ class KnownImport:
     <KnownImport 'from numpy import uint8 as ui8'>
     """
 
-    import_name: str = None
-    import_path: str = None
-    import_alias: str = None
-    builtin_name: str = None
+    # docstub: off
+    import_name: str | None = None
+    import_path: str | None = None
+    import_alias: str | None = None
+    builtin_name: str | None = None
+    # docstub: on
 
     @classmethod
     @cache
@@ -194,15 +196,15 @@ def __post_init__(self):
         elif self.import_name is None:
             raise ValueError("non builtin must at least define an `import_name`")
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         if self.builtin_name:
             info = f"{self.target} (builtin)"
         else:
             info = f"{self.format_import()!r}"
         out = f"<{type(self).__name__} {info}>"
         return out
 
-    def __str__(self):
+    def __str__(self) -> str:
         out = self.format_import()
         return out
 
@@ -406,7 +408,10 @@ class TypesDatabase:
 
     Attributes
     ----------
-    current_source : ~.PackageFile | None
+    current_source : Path | None
+    source_pkgs : list[Path]
+    known_imports: dict[str, KnownImport]
+    stats : dict[str, Any]
 
     Examples
     --------
@@ -427,11 +432,13 @@ def __init__(
         ----------
         source_pkgs: list[Path], optional
         known_imports: dict[str, KnownImport], optional
+            If not provided, defaults to imports returned by
+            :func:`common_known_imports`.
         """
         if source_pkgs is None:
             source_pkgs = []
         if known_imports is None:
-            known_imports = {}
+            known_imports = common_known_imports()
 
         self.current_source = None
         self.source_pkgs = source_pkgs
@@ -524,6 +531,6 @@ def query(self, search_name):
 
         return annotation_name, known_import
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         repr = f"{type(self).__name__}({self.source_pkgs})"
         return repr
diff --git a/src/docstub/_cache.py b/src/docstub/_cache.py
@@ -54,7 +54,7 @@ def create_cache(path):
 
     gitignore_path = path / ".gitignore"
     gitignore_content = (
-        "# This file is a cache directory tag automatically created by docstub.\n" "*\n"
+        "# This file is a cache directory automatically created by docstub.\n" "*\n"
     )
     if not gitignore_path.is_file():
         with open(gitignore_path, "w") as fp:

diff --git a/src/docstub/_cli.py b/src/docstub/_cli.py
@@ -14,12 +14,20 @@
 )
 from ._cache import FileCache
 from ._config import Config
-from ._stubs import Py2StubTransformer, walk_source, walk_source_and_targets
+from ._stubs import (
+    Py2StubTransformer,
+    try_format_stub,
+    walk_source,
+    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.
 
@@ -139,6 +147,14 @@ def report_execution_time():
 @click.help_option("-h", "--help")
 @report_execution_time()
 def main(source_dir, out_dir, config_path, verbose):
+    """
+    Parameters
+    ----------
+    source_dir : Path
+    out_dir : Path
+    config_path : Path
+    verbose : str
+    """
     _setup_logging(verbose=verbose)
 
     source_dir = Path(source_dir)
@@ -171,6 +187,8 @@ def main(source_dir, out_dir, config_path, verbose):
                 stub_content = stub_transformer.python_to_stub(
                     py_content, module_path=source_path
                 )
+                stub_content = f"{STUB_HEADER_COMMENT}\n\n{stub_content}"
+                stub_content = try_format_stub(stub_content)
             except (SystemExit, KeyboardInterrupt):
                 raise
             except Exception as e:
@@ -185,14 +203,14 @@ def main(source_dir, out_dir, config_path, verbose):
     successful_queries = types_db.stats["successful_queries"]
     click.secho(f"{successful_queries} matched annotations", fg="green")
 
-    grammar_errors = stub_transformer.transformer.stats["grammar_errors"]
-    if grammar_errors:
-        click.secho(f"{grammar_errors} grammar violations", fg="red")
+    grammar_error_count = stub_transformer.transformer.stats["grammar_errors"]
+    if grammar_error_count:
+        click.secho(f"{grammar_error_count} grammar violations", fg="red")
 
     unknown_doctypes = types_db.stats["unknown_doctypes"]
     if unknown_doctypes:
         click.secho(f"{len(unknown_doctypes)} unknown doctypes:", fg="red")
         click.echo("  " + "\n  ".join(set(unknown_doctypes)))
 
-    if unknown_doctypes or grammar_errors:
+    if unknown_doctypes or grammar_error_count:
         sys.exit(1)
diff --git a/src/docstub/_config.py b/src/docstub/_config.py
@@ -18,8 +18,17 @@ class Config:
     _source: tuple[Path, ...] = ()
 
     @classmethod
-    def from_toml(cls, path: Path | str) -> "Config":
-        """Return configuration options in local TOML file if they exist."""
+    def from_toml(cls, path):
+        """Return configuration options in local TOML file if they exist.
+
+        Parameters
+        ----------
+        path : Path or str
+
+        Returns
+        -------
+        config : Self
+        """
         path = Path(path)
         with open(path, "rb") as fp:
             raw = tomllib.load(fp)
@@ -29,11 +38,26 @@ def from_toml(cls, path: Path | str) -> "Config":
 
     @classmethod
     def from_default(cls):
+        """Create a configuration with default values.
+
+        Returns
+        -------
+        config : Self
+        """
         config = cls.from_toml(cls.DEFAULT_CONFIG_PATH)
         return config
 
     def merge(self, other):
-        """Merge contents with other and return a new Config instance."""
+        """Merge contents with other and return a new Config instance.
+
+        Parameters
+        ----------
+        other : Self
+
+        Returns
+        -------
+        merged : Self
+        """
         if not isinstance(other, type(self)):
             return NotImplemented
         new = Config(
@@ -56,7 +80,7 @@ def __post_init__(self):
         if not isinstance(self.replace_doctypes, dict):
             raise TypeError("replace_doctypes must be a string")
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         sources = " | ".join(str(s) for s in self._source)
         formatted = f"<{type(self).__name__}: {sources}>"
         return formatted
diff --git a/src/docstub/_docstrings.py b/src/docstub/_docstrings.py
@@ -10,9 +10,9 @@
 import click
 import lark
 import lark.visitors
-from numpydoc.docscrape import NumpyDocString
+from numpydoc.docscrape import NumpyDocString  # type: ignore[import-untyped]
 
-from ._analysis import KnownImport
+from ._analysis import KnownImport, TypesDatabase
 from ._utils import ContextFormatter, DocstubError, accumulate_qualname, escape_qualname
 
 logger = logging.getLogger(__name__)
@@ -150,7 +150,10 @@ class DoctypeTransformer(lark.visitors.Transformer):
 
     Attributes
     ----------
-    blacklisted_qualnames : frozenset[str]
+    types_db : ~.TypesDatabase
+    replace_doctypes : dict[str, str]
+    stats : dict[str, Any]
+    blacklisted_qualnames : ClassVar[frozenset[str]]
         All Python keywords [1]_ are blacklisted from use in qualnames except for ``True``
         ``False`` and ``None``.
 
@@ -161,11 +164,13 @@ class DoctypeTransformer(lark.visitors.Transformer):
     Examples
     --------
     >>> transformer = DoctypeTransformer()
-    >>> annotation, unknown_names = transformer.doctype_to_annotation("tuple of int")
+    >>> annotation, unknown_names = transformer.doctype_to_annotation(
+    ...     "tuple of (int or ndarray)"
+    ... )
     >>> annotation.value
-    'tuple[int]'
+    'tuple[int | ndarray]'
     >>> unknown_names
-    [('tuple', 0, 5), ('int', 9, 12)]
+    [('ndarray', 17, 24)]
     """
 
     blacklisted_qualnames = frozenset(
@@ -209,15 +214,19 @@ def __init__(self, *, types_db=None, replace_doctypes=None, **kwargs):
         """
         Parameters
         ----------
-        types_db : ~.TypesDatabase
-            A static database of collected types usable as an annotation.
+        types_db : ~.TypesDatabase, optional
+            A static database of collected types usable as an annotation. If
+            not given, defaults to a database with common types from the
+            standard library (see :func:`~.common_known_imports`).
         replace_doctypes : dict[str, str], optional
             Replacements for human-friendly aliases.
         kwargs : dict[Any, Any], optional
             Keyword arguments passed to the init of the parent class.
         """
         if replace_doctypes is None:
             replace_doctypes = {}
+        if types_db is None:
+            types_db = TypesDatabase()
 
         self.types_db = types_db
         self.replace_doctypes = replace_doctypes
@@ -272,14 +281,14 @@ def __default__(self, data, children, meta):
         ----------
         data : lark.Token
             The rule-token of the current node.
-        children : list[lark.Token, ...]
+        children : list[lark.Token]
             The children of the current node.
         meta : lark.tree.Meta
             Meta information for the current node.
 
         Returns
         -------
-        out : lark.Token or list[lark.Token, ...]
+        out : lark.Token or list[lark.Token]
             Either a token or list of tokens.
         """
         if isinstance(children, list) and len(children) == 1: