WIP

Author: lagru

src/docstub-stubs/_analysis.pyi

@@ -14,6 +14,7 @@ from typing import Any, ClassVar
 import libcst as cst
 import libcst.matchers as cstm
 
+from ._report import Stats
 from ._utils import accumulate_qualname, module_name_from_path, pyfile_checksum
 
 logger: logging.Logger
@@ -83,6 +84,7 @@ class TypeMatcher:
         types: dict[str, PyImport] | None = ...,
         type_prefixes: dict[str, PyImport] | None = ...,
         type_nicknames: dict[str, str] | None = ...,
+        stats: Stats | None = ...,
     ) -> None: ...
     def _resolve_nickname(self, name: str) -> str: ...
     def match(self, search: str) -> tuple[str | None, PyImport | None]: ...

src/docstub-stubs/_app_generate_stubs.pyi

@@ -6,7 +6,6 @@ from collections import Counter
 from collections.abc import Iterable, Sequence
 from contextlib import contextmanager
 from pathlib import Path
-from typing import Literal
 
 from ._analysis import PyImport, TypeCollector, TypeMatcher, common_known_types
 from ._cache import CACHE_DIR_NAME, FileCache
@@ -18,9 +17,8 @@ from ._path_utils import (
     walk_source_and_targets,
     walk_source_package,
 )
-from ._report import setup_logging
+from ._report import Stats, setup_logging
 from ._stubs import Py2StubTransformer, try_format_stub
-from ._utils import update_with_add_values
 from ._version import __version__
 
 logger: logging.Logger

src/docstub-stubs/_cli.pyi

@@ -4,7 +4,6 @@ import logging
 import sys
 from collections.abc import Callable, Sequence
 from pathlib import Path
-from typing import Literal
 
 import click
 from _typeshed import Incomplete

src/docstub-stubs/_docstrings.pyi

@@ -5,28 +5,22 @@ import traceback
 from collections.abc import Generator, Iterable
 from dataclasses import dataclass, field
 from functools import cached_property
-from pathlib import Path
-from typing import Any, ClassVar
 
 import click
 import lark
 import lark.visitors
 import numpydoc.docscrape as npds
 
 from ._analysis import PyImport, TypeMatcher
-from ._report import ContextReporter
-from ._utils import DocstubError, escape_qualname
+from ._doctype import BlacklistedQualname, Expression, Token, TokenKind, parse_doctype
+from ._report import ContextReporter, Stats
+from ._utils import escape_qualname
 
 logger: logging.Logger
 
-here: Path
-grammar_path: Path
-
-with grammar_path.open() as file:
-    _grammar: str
-
-_lark: lark.Lark
-
+def update_qualnames(
+    expr: Expression, *, _parents: tuple[Expression, ...] = ...
+) -> Generator[tuple[tuple[Expression, ...], Token], str]: ...
 def _find_one_token(tree: lark.Tree, *, name: str) -> lark.Token: ...
 @dataclass(frozen=True, slots=True, kw_only=True)
 class Annotation:
@@ -54,51 +48,23 @@ class Annotation:
 
 FallbackAnnotation: Annotation
 
-class QualnameIsKeyword(DocstubError):
-    pass
-
-class DoctypeTransformer(lark.visitors.Transformer):
-    matcher: TypeMatcher
-    stats: dict[str, Any]
-
-    blacklisted_qualnames: ClassVar[frozenset[str]]
-
-    def __init__(
-        self, *, matcher: TypeMatcher | None = ..., **kwargs: dict[Any, Any]
-    ) -> None: ...
-    def doctype_to_annotation(
-        self, doctype: str, *, reporter: ContextReporter | None = ...
-    ) -> tuple[Annotation, list[tuple[str, int, int]]]: ...
-    def qualname(self, tree: lark.Tree) -> lark.Token: ...
-    def rst_role(self, tree: lark.Tree) -> lark.Token: ...
-    def union(self, tree: lark.Tree) -> str: ...
-    def subscription(self, tree: lark.Tree) -> str: ...
-    def natlang_literal(self, tree: lark.Tree) -> str: ...
-    def natlang_container(self, tree: lark.Tree) -> str: ...
-    def natlang_array(self, tree: lark.Tree) -> str: ...
-    def array_name(self, tree: lark.Tree) -> lark.Token: ...
-    def shape(self, tree: lark.Tree) -> lark.visitors._DiscardType: ...
-    def optional_info(self, tree: lark.Tree) -> lark.visitors._DiscardType: ...
-    def __default__(
-        self, data: lark.Token, children: list[lark.Token], meta: lark.tree.Meta
-    ) -> lark.Token | list[lark.Token]: ...
-    def _match_import(self, qualname: str, *, meta: lark.tree.Meta) -> str: ...
-
 def _uncombine_numpydoc_params(
     params: list[npds.Parameter],
 ) -> Generator[npds.Parameter]: ...
+def _red_partial_underline(doctype: str, *, start: int, stop: int) -> str: ...
 
 class DocstringAnnotations:
     docstring: str
-    transformer: DoctypeTransformer
+    matcher: TypeMatcher
     reporter: ContextReporter
 
     def __init__(
         self,
         docstring: str,
         *,
-        transformer: DoctypeTransformer,
+        matcher: TypeMatcher,
         reporter: ContextReporter | None = ...,
+        stats: Stats | None = ...,
     ) -> None: ...
     def _doctype_to_annotation(
         self, doctype: str, ds_line: int = ...

src/docstub-stubs/_report.pyi

@@ -2,11 +2,13 @@
 
 import dataclasses
 import logging
+from collections.abc import Hashable, Iterator, Mapping, Sequence
 from pathlib import Path
 from textwrap import indent
-from typing import Any, ClassVar, Literal, Self, TextIO
+from typing import Any, ClassVar, Self, TextIO
 
 import click
+from pre_commit.envcontext import UNSET
 
 from ._cli_help import should_strip_ansi
 
@@ -79,3 +81,22 @@ class LogCounter(logging.NullHandler):
 def setup_logging(
     *, verbosity: Literal[-2, -1, 0, 1, 2, 3], group_errors: bool
 ) -> tuple[ReportHandler, LogCounter]: ...
+def update_with_add_values(
+    *mappings: Mapping[Hashable, int | Sequence], out: dict | None = ...
+) -> dict: ...
+
+class Stats(Mapping):
+    class _UNSET:
+        pass
+
+    def __init__(self, stats: dict[str, list[Any] | str] | None = ...) -> None: ...
+    def __getitem__(self, key: str) -> list[Any] | int: ...
+    def __iter__(self) -> Iterator: ...
+    def __len__(self) -> int: ...
+    def inc_counter(self, key: str, *, inc: int = ...) -> None: ...
+    def append_to_list(self, key: str, value: Any) -> None: ...
+    @classmethod
+    def merge(cls, *stats: Self) -> Self: ...
+    def __repr__(self) -> str: ...
+    def pop(self, key: str, *, default: Any = ...) -> list[Any] | int: ...
+    def pop_all(self) -> dict[str, list[Any] | int]: ...

src/docstub-stubs/_stubs.pyi

@@ -12,14 +12,10 @@ import libcst.matchers as cstm
 from _typeshed import Incomplete
 
 from ._analysis import PyImport, TypeMatcher
-from ._docstrings import (
-    Annotation,
-    DocstringAnnotations,
-    DoctypeTransformer,
-    FallbackAnnotation,
-)
-from ._report import ContextReporter
-from ._utils import module_name_from_path, update_with_add_values
+from ._docstrings import Annotation, DocstringAnnotations, FallbackAnnotation
+from ._doctype import DoctypeTransformer
+from ._report import ContextReporter, Stats
+from ._utils import module_name_from_path
 
 logger: logging.Logger
 
@@ -73,9 +69,6 @@ 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, Hashable, Mapping, Sequence
+from collections.abc import Callable
 from functools import lru_cache, wraps
 from pathlib import Path
 from zlib import crc32
@@ -12,9 +12,6 @@ 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/_analysis.py

@@ -13,6 +13,7 @@
 import libcst.matchers as cstm
 
 from ._utils import accumulate_qualname, module_name_from_path, pyfile_checksum
+from ._report import Stats
 
 logger: logging.Logger = logging.getLogger(__name__)
 
@@ -492,22 +493,21 @@ def __init__(
         types=None,
         type_prefixes=None,
         type_nicknames=None,
+        stats=None,
     ):
         """
         Parameters
         ----------
         types : dict[str, PyImport]
         type_prefixes : dict[str, PyImport]
         type_nicknames : dict[str, str]
+        stats : ~.Stats, optional
         """
         self.types = common_known_types() | (types or {})
         self.type_prefixes = type_prefixes or {}
         self.type_nicknames = type_nicknames or {}
 
-        self.stats = {
-            "matched_type_names": 0,
-            "unknown_type_names": [],
-        }
+        self.stats = stats or Stats()
 
         self.current_file = None
 
@@ -623,8 +623,8 @@ def match(self, search):
             type_name = type_name[type_name.find(py_import.target) :]
 
         if type_name is not None:
-            self.stats["matched_type_names"] += 1
+            self.stats.inc_counter("matched_type_names")
         else:
-            self.stats["unknown_type_names"].append(search)
+            self.stats.append_to_list("unknown_type_names", search)
 
         return type_name, py_import

src/docstub/_app_generate_stubs.py

@@ -24,9 +24,8 @@
     walk_source_and_targets,
     walk_source_package,
 )
-from ._report import setup_logging
+from ._report import setup_logging, Stats
 from ._stubs import Py2StubTransformer, try_format_stub
-from ._utils import update_with_add_values
 from ._version import __version__
 
 logger: logging.Logger = logging.getLogger(__name__)
@@ -234,7 +233,9 @@ def _generate_single_stub(task):
         logger.info("Wrote %s", stub_path)
         fo.write(stub_content)
 
-    stats = stub_transformer.collect_stats()
+    stats = Stats.merge(
+        stub_transformer.stats.pop_all(), stub_transformer.matcher.stats.pop_all()
+    )
 
     return stats
 
@@ -350,7 +351,7 @@ def generate_stubs(
         stats_per_task = executor.map(
             _generate_single_stub, task_args, chunksize=chunk_size
         )
-        stats = update_with_add_values(*stats_per_task)
+        stats = Stats.merge(*stats_per_task)
 
     py_typed_out = out_dir / "py.typed"
     if not py_typed_out.exists():
@@ -368,24 +369,26 @@ def generate_stubs(
     total_warnings = error_counter.warning_count
     total_errors = error_counter.error_count
 
-    logger.info("Recognized type names: %i", stats["matched_type_names"])
-    logger.info("Transformed doctypes: %i", stats["transformed_doctypes"])
+    logger.info("Recognized type names: %i", stats.pop("matched_type_names", default=0))
+    logger.info("Transformed doctypes: %i", stats.pop("transformed_doctypes", default=0))
     if total_warnings:
         logger.warning("Warnings: %i", total_warnings)
-    if stats["doctype_syntax_errors"]:
+    if "doctype_syntax_errors" in stats:
         assert total_errors
-        logger.warning("Syntax errors: %i", stats["doctype_syntax_errors"])
-    if stats["unknown_type_names"]:
+        logger.warning("Syntax errors: %i", stats.pop("doctype_syntax_errors"))
+    if "unknown_type_names" in stats:
         assert total_errors
         logger.warning(
             "Unknown type names: %i (locations: %i)",
             len(set(stats["unknown_type_names"])),
             len(stats["unknown_type_names"]),
-            extra={"details": _format_unknown_names(stats["unknown_type_names"])},
+            extra={"details": _format_unknown_names(stats.pop("unknown_type_names"))},
         )
     if total_errors:
         logger.error("Total errors: %i", total_errors)
 
+    assert len(stats) == 0
+
     total_fails = total_errors
     if fail_on_warning:
         total_fails += total_warnings