diff --git a/docs/configuration.md b/docs/configuration.md index d5ac9fb..bff108b 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -1,20 +1,21 @@ -# Configuration files +# Configuration -Docstub will automatically look for configuration files in the current directory. -These files must be named +Docstub will automatically look for the following configuration files in the current directory: - `pyproject.toml` - `docstub.toml` Alternatively, config files can also be passed in the command line with the `--config` option. -In this case, docstub will not look for configuration files in the current directory. -When multiple configuration files are passed explicitly, their content is merged. +In this case, docstub will **not** look for configuration files in the current directory. -Out of the box, docstub makes use of an internal configuration files are always loaded. +When multiple configuration files are found/used, their content is merged in order. +In case of conflicts, the later file takes precedence. + +Docstub also uses internal configuration files. One such file is [`numpy_config.toml`](../src/docstub/numpy_config.toml) which provides defaults for NumPy types. -## Configuration fields in `[tool.docstub]` +## Configuration fields All configuration must be declared inside a `[tool.docstub]` table. @@ -27,7 +28,7 @@ Ignore files and directories matching these [glob-style patterns](https://docs.p Patterns that don't start with "/" are interpreted as relative to the directory that contains the Python package for which stubs are generated. -Example: +**Example:** ```toml [tool.docstub] @@ -39,6 +40,26 @@ ignore_files = [ - Will ignore any directory anywhere that is named `tests`. +(plugin-paths)= +### `plugin_paths` + +[TOML type](https://toml.io/en/latest): array of string(s) + +File paths to Python files that implement a [plugin](create_plugin.md). +Relative paths will be treated as relative to the configuration file's directory. + +**Example:** + +```toml +[tool.docstub] +plugin_paths = [ + "tools/docstub_extend_syntax.py", +] +``` + +- Will load `docstub_extend_syntax.py` as a plugin. + + ### `types` [TOML type](https://toml.io/en/latest): table, mapping string to string @@ -48,7 +69,7 @@ Docstub can't yet automatically discover where to import types from other packag Instead, you can provide this information explicitly. Any type on the left side will be associated with the given "module" on the right side. -Example: +**Example:** ```toml [tool.docstub.types] @@ -69,7 +90,7 @@ Docstub can't yet automatically discover where to import types from other packag Instead, you can provide this information explicitly. Any type in a docstring whose prefix matches the name given on the left side, will be associated with the given "module" on the right side. -Example: +**Example:** ```toml [tool.docstub.type_prefixes] @@ -87,7 +108,7 @@ plt = "matplotlib.pyplot Nicknames for types that can be used in docstrings to describe valid Python types or annotations. -Example: +**Example:** ```toml [tool.docstub.type_nicknames] @@ -97,3 +118,59 @@ buffer = "collections.abc.Buffer" - Will map `func` to the `Callable`. - Will map `buffer` to `collections.abc.Buffer`. + + +### `doctype_grammar` + +[TOML type](https://toml.io/en/latest): string + +Override the default grammar with your own. +You can import, override and extend rules and terminals from "doctype.lark". +Note that if your grammar results in a modified tree, you need to implement a [custom transformer](#doctype_transformer) as well. + +**Example:** + +```toml +[tool.docstub] +grammar = """ +%import doctype (start, WS_INLINE, ARRAY_NAME) + +%extend ARRAY_NAME: "darray" + +%ignore WS_INLINE +""" +``` + +- Will allow using [array expressions](#natlang-array) with the term ``darray`` in addition to existing ones. + + +### `doctype_transformer` + +Python type: {class}`lark.visitors.Transformer` + +Provide your own class to transform the parsed doctype tree into a valid {term}`annotation expression`. + +:::{note} +Only supported in `docstub.py` files! +::: + +**Example:** + +```python +# docstub.py + +import lark.visitors +from docstub._doctype import DoctypeTransformer, Expr + +@lark.visitors.v_args(tree=True) +class MyTransformer(DoctypeTransformer): + def union(self, tree: lark.Tree) -> Expr: + # Sort items in union alphabetically + tree.children = sorted(tree.children, key=str) + out = super().union(tree) + return out + +doctype_transformer = MyTransformer # Declare the configuration field +``` + +- Will sort all members of a union alphabetically, regardless of order in the docstring. diff --git a/docs/create_plugin.md b/docs/create_plugin.md new file mode 100644 index 0000000..ff18bd7 --- /dev/null +++ b/docs/create_plugin.md @@ -0,0 +1,159 @@ +# Create a plugin + +Docstub has experimental support for plugins: standalone Python files that implement a set of hooks to reach directly into Docstub's runtime. +Plugins are loaded via the configuration field [`plugin_paths`](#plugin-paths). + +For a reference of all available hooks see {mod}`docstub.plugins`. + + +## Example: Allow another name in array expressions + +We want to expand docstub such that we can use `my_array` inside [array expressions](#natlang-array) like this: + +- `my_array of dtype int` → `my_array[int]` +- `my_array of shape (4, 4)` → `my_array` + +First we create the plugin: + +```python +# my_array_plugin.py + +from itertools import product + +import lark +from docstub.plugins import hookimpl +from docstub.doctype import DoctypeTransformer, Expr, Term + +my_grammar = """ +%import doctype (start, WS_INLINE, ARRAY_NAME) + +%extend ARRAY_NAME: "my_array" + +%ignore WS_INLINE +""" + + +@hookimpl +def docstub_load_doctype_grammar(): + return my_grammar +``` + +The enable the plugin in the configuration: + +```toml +[tool.docstub] +plugin_paths = ["path/to/my_array_plugin.py"] +``` + + +## Example: Expansion expression for qualnames + +We want to expand docstub such that the following doctype transformations are possible + +- `{Iterable, Sequence} of int(s)` → `Iterable[int] | Sequence[int]` +- `{list, set} of {int, float}` → `list[int] | list[float] | set[int] | set[float]` + +First we create the plugin: + +```python +# expansion_plugin.py + +from itertools import product + +import lark +from docstub.plugins import hookimpl +from docstub.doctype import DoctypeTransformer, Expr, Term + +my_grammar = """ +%import doctype (start, WS_INLINE, qualname) + +%extend qualname: "{" qualname ("," qualname)+ "}" -> expansion + +%ignore WS_INLINE +""" + + +@lark.visitors.v_args(tree=True) +class ExpansionTransformer(DoctypeTransformer): + def start(self, tree): + """Transform "start" rule. + + Parameters + ---------- + tree : lark.Tree + + Returns + ------- + expr : Expr + """ + children = map(self._expand_expression, tree.children) + return Expr(rule="start", children=tuple(children)) + + def expansion(self, tree): + """Transform "expansion" rule. + + Actual expansion happens in `start` when the complete expression with + all potential expansions can be seen. + + Parameters + ---------- + tree : lark.Tree + + Returns + ------- + expr : Expr + """ + return Expr(rule="expansion", children=tuple(tree.children)) + + def _expand_expression(self, expr): + """Join combination of all expansions in a new union. + + Parameters + ---------- + expr : Expr + + Returns + ------- + expanded : Expr + """ + expansions = expr.filter(rule="expansion") + if not expansions: + return expr + + # Create Cartesian product of all expansions where each + # value is prefixed by the expansion it will replace + expansions = [ + [(expansion, value) for value in expansion.children] + for expansion in expansions + ] + combinations = list(product(*expansions)) + + # Create a new expression for each combination + interpolated = [] + for combination in combinations: + new_expr = expr + for expansion, value in combination: + new_expr = new_expr.replace(expansion, value) + interpolated.append(new_expr) + + # Join new expressions as union + expanded = self.union(lark.Tree("union", children=interpolated)) + return expanded + + +@hookimpl +def docstub_load_doctype_grammar(): + return my_grammar + + +@hookimpl +def docstub_load_doctype_transformer(): + return ExpansionTransformer() +``` + +The enable the plugin in the configuration: + +```toml +[tool.docstub] +plugin_paths = ["path/to/expansion_plugin.py"] +``` diff --git a/docs/index.md b/docs/index.md index 185084b..eda73ab 100644 --- a/docs/index.md +++ b/docs/index.md @@ -22,6 +22,7 @@ It does so by supporting widely used readable conventions such as `array of dtyp installation introduction +create_plugin development/index ::: diff --git a/docs/installation.md b/docs/installation.md index f72c02e..aa47abe 100644 --- a/docs/installation.md +++ b/docs/installation.md @@ -19,7 +19,7 @@ In case you want to check out an unreleased version you can install the latest v pip install 'docstub @ git+https://github.com/scientific-python/docstub' ``` -You can append `@COMMIT_SHA` to the repository URL above to intall a specific version other that the `main` branch. +You can append `@COMMIT_SHA` to the repository URL above to install a specific version other that the `main` branch. :::{tip} If you are want to set up a development environment instead, checkout the [Contribution guide](development/contributing.md). diff --git a/docs/introduction.md b/docs/introduction.md index cabbf57..4ac49ba 100644 --- a/docs/introduction.md +++ b/docs/introduction.md @@ -139,7 +139,7 @@ Adopting docstub on a large codebase may initially generate many errors. Two command line options can help addressing these errors gradually: * `--group-errors` will group identical errors together. - This helps identifying common groups of errors that may be addressed in one go. + This helps to identify common groups of errors that may be addressed in one go. * `--allow-errors` puts an upper limit (["ratchet"](https://qntm.org/ratchet)) on the number of allowed errors. This way you can adjust the upper bound of allowed errors as they are addressed. @@ -154,7 +154,7 @@ Feel welcome to [open an issue](https://github.com/scientific-python/docstub/iss ## Dealing with typing problems For various reasons – missing features in docstub, or limitations of Python's typing system – it may not always be possible to correctly type something in a docstring. -In those cases, you docstub provides a few approaches to dealing with this. +In those cases, docstub provides a few approaches to dealing with this. ### Use inline type annotation @@ -200,6 +200,11 @@ class Foo: d: str ``` +### Create a plugin + +Depending on your problem you might be able to deal with it by hooking directly into Docstub's runtime via a plugin. See [](create_plugin.md) for a guide on writing such a plugin. + + ### Write a manual stub file If all of the above does not solve your issue, you can fall back to writing a correct stub file by hand. diff --git a/docs/typing_syntax.md b/docs/typing_syntax.md index 6e787a6..4f2fb64 100644 --- a/docs/typing_syntax.md +++ b/docs/typing_syntax.md @@ -69,6 +69,7 @@ In the future, docstub may warn against or disallow nesting these natural langua ::: +(natlang-array)= ## Shape and dtype syntax for arrays This expression allows adding shape and datatype information for data structures like [NumPy arrays](https://numpy.org/doc/stable/reference/generated/numpy.ndarray.html). diff --git a/pyproject.toml b/pyproject.toml index 86afcd6..f9e3a1c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -40,6 +40,7 @@ dependencies = [ "lark >=1.1.9", "black >=24.4.2", "isort >=5.13.2", + "pluggy", ] [project.urls] diff --git a/src/docstub-stubs/_app_generate_stubs.pyi b/src/docstub-stubs/_app_generate_stubs.pyi index 14d6810..9b36d5b 100644 --- a/src/docstub-stubs/_app_generate_stubs.pyi +++ b/src/docstub-stubs/_app_generate_stubs.pyi @@ -1,5 +1,6 @@ # File generated with docstub +import importlib.util import logging import sys import time @@ -7,8 +8,12 @@ from collections import Counter from collections.abc import Iterable, Sequence from contextlib import contextmanager from pathlib import Path +from types import ModuleType from typing import Literal +import pluggy + +from . import doctype, plugins from ._analysis import PyImport, TypeCollector, TypeMatcher, common_known_types from ._cache import CACHE_DIR_NAME, FileCache from ._concurrency import LoggingProcessExecutor, guess_concurrency_params @@ -21,12 +26,17 @@ from ._path_utils import ( ) from ._report import Stats, setup_logging from ._stubs import Py2StubTransformer, try_format_stub +from ._utils import escape_qualname from ._version import __version__ logger: logging.Logger def cache_dir_in_cwd() -> Path: ... def _load_configuration(config_paths: list[Path] | None = ...) -> Config: ... +def _import_standalone_module_from_path( + path: Path, *, module_name: str | None = ... +) -> ModuleType: ... +def _load_plugins(paths: Iterable[Path] = ...) -> pluggy.PluginManager: ... def _collect_type_info( root_path: Path, *, ignore: Sequence[str] = ..., cache: bool = ... ) -> tuple[dict[str, PyImport], dict[str, PyImport]]: ... diff --git a/src/docstub-stubs/_cache.pyi b/src/docstub-stubs/_cache.pyi index 43c2af2..5651492 100644 --- a/src/docstub-stubs/_cache.pyi +++ b/src/docstub-stubs/_cache.pyi @@ -1,7 +1,8 @@ # File generated with docstub import logging -from collections.abc import Callable +from collections import OrderedDict +from collections.abc import Callable, Hashable, Iterable from functools import cached_property from pathlib import Path from typing import Any, Protocol @@ -47,3 +48,16 @@ class FileCache: @property def cache_sub_dir(self) -> None: ... def __call__(self, *args: Any, **kwargs: Any) -> Any: ... + +class LRUMemoryCache: + class _MISS: + pass + + def __init__(self, *, maxsize: int = ...) -> None: ... + def __contains__(self, key: Hashable) -> bool: ... + def __getitem__(self, key: Hashable) -> Any: ... + def get(self, key: Hashable, *, default: Any | None = ...) -> Any: ... + def __setitem__(self, key: Hashable, value: Any) -> None: ... + def __len__(self) -> int: ... + def __repr__(self) -> str: ... + def __iter__(self) -> Iterable[Any]: ... diff --git a/src/docstub-stubs/_config.pyi b/src/docstub-stubs/_config.pyi index 110be6a..0213fe6 100644 --- a/src/docstub-stubs/_config.pyi +++ b/src/docstub-stubs/_config.pyi @@ -20,6 +20,7 @@ class Config: type_prefixes: dict[str, str] = ... type_nicknames: dict[str, str] = ... ignore_files: list[str] = ... + plugin_paths: list[Path] = ... config_paths: tuple[Path, ...] = ... diff --git a/src/docstub-stubs/_docstrings.pyi b/src/docstub-stubs/_docstrings.pyi index f116ada..937f3e3 100644 --- a/src/docstub-stubs/_docstrings.pyi +++ b/src/docstub-stubs/_docstrings.pyi @@ -13,15 +13,12 @@ import lark.visitors import numpydoc.docscrape as npds from ._analysis import PyImport, TypeMatcher -from ._doctype import BlacklistedQualname, Expr, Term, TermKind, parse_doctype from ._report import ContextReporter, Stats from ._utils import escape_qualname +from .doctype import BlacklistedQualname, DoctypeTranspiler, TermKind logger: logging.Logger -def _update_qualnames( - expr: Expr, *, _parents: tuple[Expr, ...] = ... -) -> Generator[tuple[tuple[Expr, ...], Term], str]: ... @dataclass(frozen=True, slots=True, kw_only=True) class Annotation: @@ -58,6 +55,7 @@ def doctype_to_annotation( matcher: TypeMatcher | None = ..., reporter: ContextReporter | None = ..., stats: Stats | None = ..., + transpiler: DoctypeTranspiler | None = ..., ) -> Annotation: ... class DocstringAnnotations: @@ -72,6 +70,7 @@ class DocstringAnnotations: matcher: TypeMatcher | None = ..., reporter: ContextReporter | None = ..., stats: Stats | None = ..., + doctype_transpiler: DoctypeTranspiler | None = ..., ) -> None: ... @cached_property def attributes(self) -> dict[str, Annotation]: ... diff --git a/src/docstub-stubs/_report.pyi b/src/docstub-stubs/_report.pyi index 6430e76..8138f35 100644 --- a/src/docstub-stubs/_report.pyi +++ b/src/docstub-stubs/_report.pyi @@ -2,7 +2,7 @@ import dataclasses import logging -from collections.abc import Hashable, Iterator, Mapping, Sequence +from collections.abc import Hashable, Mapping, Sequence from pathlib import Path from textwrap import indent from typing import Any, ClassVar, Literal, Self, TextIO @@ -105,18 +105,14 @@ def update_with_add_values( *mappings: Mapping[Hashable, int | Sequence], out: dict | None = ... ) -> dict: ... -class Stats(Mapping): +class Stats(dict): 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]: ... diff --git a/src/docstub-stubs/_stubs.pyi b/src/docstub-stubs/_stubs.pyi index 590d421..ec22fc6 100644 --- a/src/docstub-stubs/_stubs.pyi +++ b/src/docstub-stubs/_stubs.pyi @@ -14,9 +14,9 @@ from _typeshed import Incomplete from ._analysis import PyImport, TypeMatcher from ._docstrings import Annotation, DocstringAnnotations, FallbackAnnotation -from ._doctype import DoctypeTransformer from ._report import ContextReporter, Stats from ._utils import module_name_from_path +from .doctype import DoctypeTransformer, DoctypeTranspiler logger: logging.Logger @@ -65,7 +65,12 @@ class Py2StubTransformer(cst.CSTTransformer): _Annotation_Incomplete: ClassVar[cst.Annotation] _Annotation_None: ClassVar[cst.Annotation] - def __init__(self, *, matcher: TypeMatcher | None = ...) -> None: ... + def __init__( + self, + *, + matcher: TypeMatcher | None = ..., + doctype_transpiler: DoctypeTranspiler | None = ... + ) -> None: ... @property def current_source(self) -> Path: ... @current_source.setter diff --git a/src/docstub-stubs/_doctype.pyi b/src/docstub-stubs/doctype.pyi similarity index 70% rename from src/docstub-stubs/_doctype.pyi rename to src/docstub-stubs/doctype.pyi index 97e462a..cac18dd 100644 --- a/src/docstub-stubs/_doctype.pyi +++ b/src/docstub-stubs/doctype.pyi @@ -3,8 +3,9 @@ import enum import keyword import logging -from collections.abc import Generator, Iterable, Sequence +from collections.abc import Callable, Generator, Iterable, Sequence from dataclasses import dataclass +from functools import cached_property from pathlib import Path from textwrap import indent from typing import Any, Final, Self @@ -13,17 +14,14 @@ import lark import lark.visitors from _typeshed import Incomplete -from ._report import ContextReporter +from ._cache import LRUMemoryCache +from ._report import ContextReporter, Stats from ._utils import DocstubError logger: Final[logging.Logger] -grammar_path: Final[Path] - -with grammar_path.open() as file: - _grammar: Final[str] - -_lark: Final[lark.Lark] +DOCTYPE_GRAMMAR_PATH: Final[Path] +DOCTYPE_GRAMMAR: Final[str] def flatten_recursive(iterable: Iterable[Iterable | str]) -> Generator[str]: ... def insert_between(iterable: Iterable, *, sep: Any) -> list[Any]: ... @@ -46,12 +44,13 @@ class Term(str): def __repr__(self) -> str: ... def __getnewargs_ex__(self) -> tuple[tuple[Any, ...], dict[str, Any]]: ... -@dataclass(slots=True) +@dataclass(slots=True, frozen=True) class Expr: rule: str - children: list[Self | Term] + children: tuple[Self | Term, ...] + def __post_init__(self) -> None: ... @property def terms(self) -> list[Term]: ... @property @@ -64,6 +63,10 @@ class Expr: def __repr__(self) -> str: ... def __str__(self) -> str: ... def as_code(self) -> str: ... + def replace(self, old: Expr | Term, new: Expr | Term) -> Expr: ... + def filter( + self, func: Callable | None = ..., /, **attributes: Any + ) -> Generator[Expr | Term]: ... BLACKLISTED_QUALNAMES: Final[set[str]] @@ -90,6 +93,21 @@ class DoctypeTransformer(lark.visitors.Transformer): def shape(self, tree: lark.Tree) -> lark.visitors._DiscardType: ... def optional_info(self, tree: lark.Tree) -> lark.visitors._DiscardType: ... def extra_info(self, tree: lark.Tree) -> lark.visitors._DiscardType: ... - def _format_subscription(self, sequence: Sequence[str], *, rule: str) -> Expr: ... - -def parse_doctype(doctype: str, *, reporter: ContextReporter | None = ...) -> Expr: ... + @staticmethod + def _format_subscription(sequence: Sequence[str], *, rule: str) -> Expr: ... + +class DoctypeTranspiler: + def __init__( + self, + *, + grammar: str | None = ..., + grammar_import_paths: tuple[Path, ...] = ..., + transformer: lark.visitors.Transformer | None = ..., + reporter: ContextReporter | None = ..., + max_cache_size: int = ..., + ) -> None: ... + def __call__(self, doctype: str) -> Expr: ... + @property + def stats(self) -> None: ... + @cached_property + def _parser(self) -> lark.Lark: ... diff --git a/src/docstub-stubs/plugins.pyi b/src/docstub-stubs/plugins.pyi new file mode 100644 index 0000000..dc885ae --- /dev/null +++ b/src/docstub-stubs/plugins.pyi @@ -0,0 +1,14 @@ +# File generated with docstub + +from typing import Final + +import pluggy + +from .doctype import DoctypeTransformer + +_hookspec: pluggy.HookspecMarker + +hookimpl: Final[pluggy.HookimplMarker] + +def docstub_load_doctype_grammar() -> str: ... +def docstub_load_doctype_transformer() -> DoctypeTransformer: ... diff --git a/src/docstub/_analysis.py b/src/docstub/_analysis.py index 04d7cfb..8fb42b9 100644 --- a/src/docstub/_analysis.py +++ b/src/docstub/_analysis.py @@ -627,4 +627,5 @@ def match(self, search): else: self.stats.append_to_list("unknown_type_names", search) + assert bool(type_name) == bool(py_import) # Both or neither should be `None` return type_name, py_import diff --git a/src/docstub/_app_generate_stubs.py b/src/docstub/_app_generate_stubs.py index 3f98483..2355aff 100644 --- a/src/docstub/_app_generate_stubs.py +++ b/src/docstub/_app_generate_stubs.py @@ -4,6 +4,7 @@ implementation (with slow imports) speeds up the ``--help`` option. """ +import importlib.util import logging import sys import time @@ -11,6 +12,9 @@ from contextlib import contextmanager from pathlib import Path +import pluggy + +from . import doctype, plugins from ._analysis import ( PyImport, TypeCollector, @@ -28,6 +32,7 @@ ) from ._report import Stats, setup_logging from ._stubs import Py2StubTransformer, try_format_stub +from ._utils import escape_qualname from ._version import __version__ logger: logging.Logger = logging.getLogger(__name__) @@ -60,8 +65,8 @@ def _load_configuration(config_paths=None): if config_paths: for path in config_paths: - logger.info("Using %s", path) add_config = Config.from_toml(path) + logger.info("Using %s", path) config = config.merge(add_config) else: @@ -73,13 +78,69 @@ def _load_configuration(config_paths=None): docstub_toml = Path.cwd() / "docstub.toml" if docstub_toml.is_file(): - logger.info("Using %s", docstub_toml) add_config = Config.from_toml(docstub_toml) + logger.info("Using %s", docstub_toml) config = config.merge(add_config) return config +def _import_standalone_module_from_path(path, *, module_name=None): + """Import standalone module at the given file path. + + "Standalone" because relative imports inside the loaded module may not work? + + Parameters + ---------- + path : Path + Path to a Python file. + module_name : str, optional + The name to load the module as. Derived from `path` if not provided. + The functions raises if such a module already exists in :mod:`sys.modules`. + + Returns + ------- + module : ModuleType + """ + module_name = module_name or escape_qualname(str(path)) + if module_name in sys.modules: + raise RuntimeError(f"module {module_name} already defined") + + spec = importlib.util.spec_from_file_location(module_name, path) + module = importlib.util.module_from_spec(spec) + sys.modules[module_name] = module + spec.loader.exec_module(module) + return module + + +def _load_plugins(paths=()): + """Load plugins for docstub. + + Parameters + ---------- + paths : Iterable of Path(s) + File paths to load plugins from. + + Returns + ------- + manager : pluggy.PluginManager + """ + manager = pluggy.PluginManager("docstub") + manager.add_hookspecs(plugins) + + # Loading plugins via entrypoints is intentionally not supported (for now) + # manager.load_setuptools_entrypoints("docstub") + + for path in paths: + plugin_module = _import_standalone_module_from_path(path) + manager.register(plugin_module) + hooks = manager.get_hookcallers(plugin_module) + logger.info("Registered plugin at %s (%i hook(s))", path, len(hooks)) + logger.debug("Hooks implemented by plugin at %s:\n%r", path, hooks) + + return manager + + def _collect_type_info(root_path, *, ignore=(), cache=False): """Collect types. @@ -237,7 +298,9 @@ def _generate_single_stub(task): fo.write(stub_content) stats = Stats.merge( - stub_transformer.stats.pop_all(), stub_transformer.matcher.stats.pop_all() + stub_transformer.stats.pop_all(), + stub_transformer.matcher.stats.pop_all(), + stub_transformer.doctype_transpiler.stats.pop_all(), ) return stats @@ -298,6 +361,9 @@ def generate_stubs( config = config.merge(Config(ignore_files=list(ignore))) logger.debug("Fully loaded config:\n%s", config) + plugin_manager = _load_plugins(config.plugin_paths) + plugin_hooks = plugin_manager.hook + types, type_prefixes = _collect_type_info( root_path, ignore=config.ignore_files, cache=not no_cache ) @@ -318,10 +384,17 @@ def generate_stubs( for prefix, module in config.type_prefixes.items() } + grammar = plugin_hooks.docstub_load_doctype_grammar() + doctype_transformer = plugin_hooks.docstub_load_doctype_transformer() + doctype_transpiler = doctype.DoctypeTranspiler( + grammar=grammar, transformer=doctype_transformer + ) matcher = TypeMatcher( types=types, type_prefixes=type_prefixes, type_nicknames=config.type_nicknames ) - stub_transformer = Py2StubTransformer(matcher=matcher) + stub_transformer = Py2StubTransformer( + matcher=matcher, doctype_transpiler=doctype_transpiler + ) if not out_dir: if root_path.is_file(): @@ -373,13 +446,17 @@ def generate_stubs( output_handler.group_errors = False # Report basic statistics + logger.debug( + "DoctypeTranspiler caching stats: hits=%i, misses=%i, maxsize=%i", + stats.pop("doctype_transpiler_cache_hits", 0), + stats.pop("doctype_transpiler_cache_misses", 0), + max(stats.pop("doctype_transpiler_curr_cache_size", [0])), + ) + logger.info("Recognized type names: %i", stats.pop("matched_type_names", 0)) + logger.info("Transformed doctypes: %i", stats.pop("transformed_doctypes", 0)) + total_warnings = error_counter.warning_count total_errors = error_counter.error_count - - 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 "doctype_syntax_errors" in stats: diff --git a/src/docstub/_cache.py b/src/docstub/_cache.py index e08b78c..c42244c 100644 --- a/src/docstub/_cache.py +++ b/src/docstub/_cache.py @@ -1,6 +1,7 @@ """Caching to disk.""" import logging +from collections import OrderedDict from functools import cached_property from typing import Any, Protocol @@ -226,3 +227,122 @@ def __call__(self, *args, **kwargs): self.cache_misses += 1 return data + + +class LRUMemoryCache: + """Simple LRU memory cache. + + Implements a reasonably performant cache policy [1] based on + :class:`collections.OrderedDict`. Has better ergonomics than + :func:`functools.lru_cache` in some situations. + + Parameters + ---------- + maxsize : int, optional + The maximum size to which the cache is allowed to grow. + + Examples + -------- + >>> cache = LRUMemoryCache(maxsize=3) + >>> cache[1] = True + >>> cache[2] = True + >>> cache[3] = True + >>> cache[1] # Cache hit for key=1 + True + >>> cache[4] = True # Push the oldest item of the cache + >>> list(cache) + [3, 1, 4] + + References + ---------- + .. [1] https://en.wikipedia.org/wiki/Page_replacement_algorithm#Least_recently_used + """ + + class _MISS: + """Sentinel value to signal a miss internally.""" + + def __init__(self, *, maxsize=128): + self._cache = OrderedDict() + self.maxsize = maxsize + + def __contains__(self, key): + """Check if an item is cached at `key`. + + Parameters + ---------- + key : Hashable + + Returns + ------- + is_cached : bool + """ + return key in self._cache + + def __getitem__(self, key): + """Get an existing item from the cache. + + Parameters + ---------- + key : Hashable + + Returns + ------- + value : Any + """ + value = self._cache[key] + self._cache.move_to_end(key) + return value + + def get(self, key, *, default=None): + """Get an item from the cache with fallback. + + Parameters + ---------- + key : Hashable + default : Any, optional + Value to fall back to if item is not cached. + + Returns + ------- + value : Any + """ + value = self._cache.get(key, self._MISS) + if value is self._MISS: + value = default + else: + self._cache.move_to_end(key) + return value + + def __setitem__(self, key, value): + """Add item to the cache. + + Parameters + ---------- + key : Hashable + value : Any + """ + if key in self._cache: + self._cache.pop(key) + self._cache[key] = value + if len(self._cache) > self.maxsize: + self._cache.popitem(last=False) + + def __len__(self): + """Return current size of the cache. + + Returns + ------- + current_size : int + """ + return len(self._cache) + + def __repr__(self) -> str: + return f"<{type(self).__name__}: {len(self)} (maxsize={self.maxsize}>" + + def __iter__(self): + """ + Returns + ------- + keys : Iterable[Any] + """ + return iter(self._cache) diff --git a/src/docstub/_concurrency.py b/src/docstub/_concurrency.py index 52a7ad0..6616ae5 100644 --- a/src/docstub/_concurrency.py +++ b/src/docstub/_concurrency.py @@ -151,6 +151,14 @@ def __exit__(self, exc_type, exc_val, exc_tb): logger.debug("Calling `cancel_join_thread` on queue") self._queue.cancel_join_thread() + elif exc_type: + logger.exception("Unexpected error in worker") + logger.debug("Terminating workers") + self._pool.terminate_workers() + # Ensure that the queue doesn't block (not sure if necessary) + logger.debug("Calling `cancel_join_thread` on queue") + self._queue.cancel_join_thread() + else: # We want to wait for any log record to reach the listener logger.debug("Shutting down pool") diff --git a/src/docstub/_config.py b/src/docstub/_config.py index e7fe44a..0348975 100644 --- a/src/docstub/_config.py +++ b/src/docstub/_config.py @@ -21,6 +21,7 @@ class Config: type_prefixes: dict[str, str] = dataclasses.field(default_factory=dict) type_nicknames: dict[str, str] = dataclasses.field(default_factory=dict) ignore_files: list[str] = dataclasses.field(default_factory=list) + plugin_paths: list[Path] = dataclasses.field(default_factory=list) config_paths: tuple[Path, ...] = () @@ -36,10 +37,25 @@ def from_toml(cls, path): ------- config : Self """ - path = Path(path) + path = Path(path).resolve() with open(path, "rb") as fp: raw = tomllib.load(fp) - config = cls(**raw.get("tool", {}).get("docstub", {}), config_paths=(path,)) + + docstub_table = raw.get("tool", {}).get("docstub", {}) + toml_dir = path.parent + + def normalize_plugin_paths(p): + p = Path(p) + if not p.is_absolute(): + p = toml_dir / str(p) + p = p.resolve() + return p + + plugin_paths = docstub_table.get("plugin_paths", []) + plugin_paths = [normalize_plugin_paths(p) for p in plugin_paths] + docstub_table["plugin_paths"] = plugin_paths + + config = cls(**docstub_table, config_paths=(path,)) logger.debug("Created `Config` from %s", path) return config @@ -56,11 +72,13 @@ def merge(self, other): """ if not isinstance(other, type(self)): return NotImplemented + new = Config( types=self.types | other.types, type_prefixes=self.type_prefixes | other.type_prefixes, type_nicknames=self.type_nicknames | other.type_nicknames, ignore_files=self.ignore_files + other.ignore_files, + plugin_paths=self.plugin_paths + other.plugin_paths, config_paths=self.config_paths + other.config_paths, ) logger.debug("Merged Config from %s", new.config_paths) @@ -107,3 +125,8 @@ def validate(mapping): for value in sequence: if not isinstance(value, str): raise TypeError(f"`{value}` in {field} must be a str") + + for path in mapping["plugin_paths"]: + if not path.is_file(): + raise ValueError(f"plugin file {path} does not exist") + assert path.is_absolute() diff --git a/src/docstub/_docstrings.py b/src/docstub/_docstrings.py index 7323191..9073df5 100644 --- a/src/docstub/_docstrings.py +++ b/src/docstub/_docstrings.py @@ -12,66 +12,17 @@ import numpydoc.docscrape as npds # TODO Uncouple docstrings & analysis module -# It should be possible to transform docstrings without matching to valid -# types and imports. I think that could very well be done at a higher level, -# e.g. in the stubs module. +# Matching types and imports in the _stubs.py would give the opportunity to +# cache extraction of annotations from docstrings. May be a good performance +# improvement but would require marking qualnames in annotations. from ._analysis import PyImport, TypeMatcher -from ._doctype import BlacklistedQualname, Term, TermKind, parse_doctype from ._report import ContextReporter, Stats from ._utils import escape_qualname +from .doctype import BlacklistedQualname, DoctypeTranspiler, TermKind logger: logging.Logger = logging.getLogger(__name__) -def _update_qualnames(expr, *, _parents=()): - """Yield and receive names in `expr`. - - This generator works as a coroutine. - - Parameters - ---------- - expr : ~.Expr - _parents : tuple of (~._doctype.Expr, ...) - - Yields - ------ - parents : tuple of (~._doctype.Expr, ...) - name : ~._doctype.Term - - Receives - -------- - new_name : str - - Examples - -------- - >>> from docstub._doctype import parse_doctype - >>> expr = parse_doctype("tuple of (tuple or str, ...)") - >>> updater = _update_qualnames(expr) - >>> for parents, name in updater: - ... if name == "tuple" and parents[-1].rule == "union": - ... updater.send("list") - ... if name == "str": - ... updater.send("bytes") - >>> expr.as_code() - 'tuple[list | bytes, ...]' - """ - _parents += (expr,) - children = expr.children.copy() - - for i, child in enumerate(children): - if hasattr(child, "children"): - yield from _update_qualnames(child, _parents=_parents) - - elif child.kind == TermKind.NAME: - new_name = yield _parents, child - if new_name is not None: - new_term = Term(new_name, kind=child.kind) - expr.children[i] = new_term - # `send` was called, yield `None` to return from `send`, - # otherwise send would return the next child - yield - - @dataclass(frozen=True, slots=True, kw_only=True) class Annotation: """Python-ready type annotation with attached import information.""" @@ -246,7 +197,9 @@ def _red_partial_underline(doctype, *, start, stop): return underlined -def doctype_to_annotation(doctype, *, matcher=None, reporter=None, stats=None): +def doctype_to_annotation( + doctype, *, matcher=None, reporter=None, stats=None, transpiler=None +): """Convert a type description to a Python-ready type. Parameters @@ -255,6 +208,7 @@ def doctype_to_annotation(doctype, *, matcher=None, reporter=None, stats=None): matcher : ~.TypeMatcher, optional reporter : ~.ContextReporter, optional stats : ~.Stats, optional + transpiler : ~.DoctypeTranspiler, optional Returns ------- @@ -265,38 +219,48 @@ def doctype_to_annotation(doctype, *, matcher=None, reporter=None, stats=None): matcher = matcher or TypeMatcher() reporter = reporter or ContextReporter(logger=logger) stats = Stats() if stats is None else stats + transpiler = transpiler or DoctypeTranspiler(reporter=reporter) try: - expression = parse_doctype(doctype, reporter=reporter) + expression = transpiler(doctype) stats.inc_counter("transformed_doctypes") reporter.debug( "Transformed doctype", details=(" %s\n-> %s", doctype, expression) ) + updated_terms = [] imports = set() unknown_qualnames = set() - updater = _update_qualnames(expression) - for _, name in updater: - search_name = str(name) + + for term in expression.terms: + if term.kind != TermKind.NAME: + # Syntax or literal, leave as is and continue to next term + updated_terms.append(term) + continue + + search_name = str(term) matched_name, py_import = matcher.match(search_name) + if matched_name is None: assert py_import is None - unknown_qualnames.add((search_name, *name.pos)) + unknown_qualnames.add((search_name, *term.pos)) matched_name = escape_qualname(search_name) - _ = updater.send(matched_name) - assert _ is None + updated_terms.append(matched_name) if py_import is None: + # Did not find match, alias to `_typeshed.Incomplete` incomplete_alias = PyImport( from_="_typeshed", import_="Incomplete", as_=matched_name, ) imports.add(incomplete_alias) + elif py_import.has_import: imports.add(py_import) - annotation = Annotation(value=str(expression), imports=frozenset(imports)) + updated_terms = "".join(updated_terms) + annotation = Annotation(value=updated_terms, imports=frozenset(imports)) except ( lark.exceptions.LexError, @@ -357,7 +321,15 @@ class DocstringAnnotations: dict_keys(['a', 'b', 'c']) """ - def __init__(self, docstring, *, matcher=None, reporter=None, stats=None): + def __init__( + self, + docstring, + *, + matcher=None, + reporter=None, + stats=None, + doctype_transpiler=None, + ): """ Parameters ---------- @@ -365,10 +337,12 @@ def __init__(self, docstring, *, matcher=None, reporter=None, stats=None): matcher : ~.TypeMatcher, optional reporter : ~.ContextReporter, optional stats : ~.Stats, optional + doctype_transpiler : ~.DoctypeTranspiler, optional """ self.docstring = docstring self.matcher = matcher or TypeMatcher() self.stats = Stats() if stats is None else stats + self.doctype_transpiler = doctype_transpiler if reporter is None: reporter = ContextReporter(logger=logger, line=0) @@ -554,6 +528,7 @@ def _section_annotations(self, name): matcher=self.matcher, reporter=self.reporter.copy_with(line_offset=ds_line), stats=self.stats, + transpiler=self.doctype_transpiler, ) annotated_params[param.name.strip()] = annotation diff --git a/src/docstub/_report.py b/src/docstub/_report.py index 0de0282..773a1bd 100644 --- a/src/docstub/_report.py +++ b/src/docstub/_report.py @@ -2,7 +2,6 @@ import dataclasses import logging -from collections.abc import Mapping from pathlib import Path from textwrap import indent @@ -62,7 +61,7 @@ def copy_with(self, *, logger=None, path=None, line=None, line_offset=None): Parameters ---------- - logger : logging.Logger + logger : logging.Logger, optional path : Path, optional line : int, optional line_offset : int, optional @@ -487,8 +486,8 @@ def update_with_add_values(*mappings, out=None): return out -class Stats(Mapping): - """Collect statistics +class Stats(dict): + """Collect statistics. Examples -------- @@ -517,31 +516,7 @@ def __init__(self, stats=None): ---------- stats : dict[str, list[Any] or str] """ - self._stats = {} if stats is None else stats - - def __getitem__(self, key): - """Retrieve a statistic. - - Parameters - ---------- - key : str - - Returns - ------- - value : list[Any] or int - """ - return self._stats[key] - - def __iter__(self): - """ - Returns - ------- - out : Iterator - """ - yield from self._stats - - def __len__(self) -> int: - return len(self._stats) + super().__init__(stats or {}) def inc_counter(self, key, *, inc=1): """Increase counter of a statistic. @@ -551,10 +526,9 @@ def inc_counter(self, key, *, inc=1): key : str inc : int, optional """ - if key not in self._stats: - self._stats[key] = 0 - assert isinstance(inc, int) - self._stats[key] += inc + if key not in self: + self[key] = 0 + self[key] += inc def append_to_list(self, key, value): """Append `value` to statistic. @@ -564,9 +538,9 @@ def append_to_list(self, key, value): key : str value : Any """ - if key not in self._stats: - self._stats[key] = [] - self._stats[key].append(value) + if key not in self: + self[key] = [] + self[key].append(value) @classmethod def merge(cls, *stats): @@ -585,27 +559,9 @@ def merge(cls, *stats): return out def __repr__(self) -> str: - keys = ", ".join(self._stats.keys()) + keys = ", ".join(self.keys()) return f"<{type(self).__name__}: {keys}>" - def pop(self, key, *, default=_UNSET): - """Return and remove a statistic from this container. - - Parameters - ---------- - key : str - default : Any, optional - If given, falls back to the given default value if `key` is not - found. - - Returns - ------- - value : list[Any] or int - """ - if key in self._stats or default is self._UNSET: - return self._stats.pop(key) - return default - def pop_all(self): """Return and remove all statistics from this container. @@ -613,6 +569,6 @@ def pop_all(self): ------- stats : dict[str, list[Any] or int] """ - out = self._stats - self._stats = {} + out = dict(self) + self.clear() return out diff --git a/src/docstub/_stubs.py b/src/docstub/_stubs.py index b8f786a..398b856 100644 --- a/src/docstub/_stubs.py +++ b/src/docstub/_stubs.py @@ -24,7 +24,19 @@ def try_format_stub(stub: str) -> str: - """Try to format a stub file with isort and black if available.""" + """Try to format a stub file with isort and black if available. + + Calls isort and black wiht their standard configuration if they are + available. + """ + + # TODO This seems faster even though we invoke a subprocess? + # But `ruff format` doesn't sort imports... + # output = subprocess.check_output( + # ["/usr/bin/ruff", "format", "-"], input=stub, text=True + # ) + # return output + try: import isort # noqa: PLC0415 @@ -324,13 +336,15 @@ def print_upper(x: Incomplete) -> None: ... ) _Annotation_None: ClassVar[cst.Annotation] = cst.Annotation(cst.Name("None")) - def __init__(self, *, matcher=None): + def __init__(self, *, matcher=None, doctype_transpiler=None): """ Parameters ---------- matcher : ~.TypeMatcher + doctype_transpiler : ~.DoctypeTranspiler, optional """ self.matcher = matcher + self.doctype_transpiler = doctype_transpiler self.reporter = ContextReporter(logger=logger) self.stats = Stats() # Relevant docstring for the current context @@ -917,6 +931,7 @@ def _annotations_from_node(self, node): matcher=self.matcher, reporter=reporter, stats=self.stats, + doctype_transpiler=self.doctype_transpiler, ) except (SystemExit, KeyboardInterrupt): raise diff --git a/src/docstub/config_template.toml b/src/docstub/config_template.toml index 192d677..b16626e 100644 --- a/src/docstub/config_template.toml +++ b/src/docstub/config_template.toml @@ -7,6 +7,10 @@ ignore_files = [] +# Register docstub plugins at these file paths. +plugin_paths = [] + + # Types and their external modules to use in docstrings. # Docstub can't yet automatically discover where to import types from other # packages from. Instead, you can provide this information explicitly. diff --git a/src/docstub/_doctype.py b/src/docstub/doctype.py similarity index 62% rename from src/docstub/_doctype.py rename to src/docstub/doctype.py index 82ed349..8c23f4f 100644 --- a/src/docstub/_doctype.py +++ b/src/docstub/doctype.py @@ -11,6 +11,7 @@ import logging from collections.abc import Iterable from dataclasses import dataclass +from functools import cached_property from pathlib import Path from textwrap import indent from typing import Final, Self @@ -18,20 +19,15 @@ import lark import lark.visitors -from ._report import ContextReporter +from ._cache import LRUMemoryCache +from ._report import ContextReporter, Stats from ._utils import DocstubError logger: Final[logging.Logger] = logging.getLogger(__name__) -grammar_path: Final[Path] = Path(__file__).parent / "doctype.lark" - -with grammar_path.open() as file: - _grammar: Final[str] = file.read() - -# TODO try passing `transformer=DoctypeTransformer()`, may be faster [1] -# [1] https://lark-parser.readthedocs.io/en/latest/classes.html#:~:text=after%20the%20parse%2C-,but%20faster,-) -_lark: Final[lark.Lark] = lark.Lark(_grammar, propagate_positions=True) +DOCTYPE_GRAMMAR_PATH: Final[Path] = Path(__file__).parent / "doctype.lark" +DOCTYPE_GRAMMAR: Final[str] = DOCTYPE_GRAMMAR_PATH.read_text() def flatten_recursive(iterable): @@ -98,7 +94,7 @@ class TermKind(enum.StrEnum): class Term(str): - """A terminal / symbol representing an atomic part of a doctype. + """A terminal / symbol representing an atomic part in an annotation expression. Attributes ---------- @@ -151,9 +147,9 @@ def __getnewargs_ex__(self): return (), kwargs -@dataclass(slots=True) +@dataclass(slots=True, frozen=True) class Expr: - """An expression that forms or is part of a doctype. + """An annotation expression. Parameters ---------- @@ -164,7 +160,12 @@ class Expr: """ rule: str - children: list[Self | Term] + children: tuple[Self | Term, ...] + + def __post_init__(self): + assert isinstance(self.rule, str) + assert self.rule + assert isinstance(self.children, tuple) @property def terms(self): @@ -184,7 +185,7 @@ def names(self): ------- names : list of Term """ - return [term for term in self.terms if term.kind == TermKind.NAME] + return self.filter(kind=TermKind.NAME) @property def sub_expressions(self): @@ -223,7 +224,7 @@ def format_tree(self): formatted_children = ",\n".join(formatted_children) formatted_children = indent(formatted_children, prefix=" ") return ( - f"{type(self).__name__}({self.rule!r}, children=[\n{formatted_children}])" + f"{type(self).__name__}({self.rule!r}, children=(\n{formatted_children}))" ) def print_tree(self): @@ -239,6 +240,55 @@ def __str__(self) -> str: def as_code(self) -> str: return str(self) + def replace(self, old, new): + """ + Parameters + ---------- + old : Expr or Term + new : Expr or Term + + Returns + ------- + replaced : Expr + """ + cls = type(self) + new_children = [] + for child in self.children: + if isinstance(child, cls): + child = child.replace(old, new) # noqa: PLW2901 + if child == old: + child = new # noqa: PLW2901 + new_children.append(child) + return cls(rule=self.rule, children=tuple(new_children)) + + def filter(self, func=None, /, **attributes): + """Filter children of the expression recursively. + + Parameters + ---------- + func : Callable, optional + Only return children for which ``func(child)`` is true. + attributes : Any + Only return children for which have an attribute with a matching + value. + + Yields + ------ + match : Expr or Term + """ + for child in self.children: + matches = True + if func: + matches = func(child) + for name, value in attributes.items(): + if not hasattr(child, name) or getattr(child, name) != value: + matches = False + break + if matches: + yield child + if hasattr(child, "filter"): + yield from child.filter(func, **attributes) + BLACKLISTED_QUALNAMES: Final[set[str]] = set(keyword.kwlist) - {"None", "True", "False"} @@ -253,7 +303,9 @@ class DoctypeTransformer(lark.visitors.Transformer): Examples -------- - >>> tree = _lark.parse("int or tuple of (int, ...)") + >>> import lark + >>> parser = lark.Lark(DOCTYPE_GRAMMAR, propagate_positions=True) + >>> tree = parser.parse("int or tuple of (int, ...)") >>> transformer = DoctypeTransformer() >>> str(transformer.transform(tree=tree)) 'int | tuple[int, ...]' @@ -278,7 +330,7 @@ def start(self, tree): ------- out : Expr """ - return Expr(rule="start", children=tree.children) + return Expr(rule="start", children=tuple(tree.children)) def qualname(self, tree): """ @@ -315,7 +367,7 @@ def rst_role(self, tree): """ # Drop rst_prefix children = [c for c in tree.children if isinstance(c, Term)] - expr = Expr(rule="rst_role", children=children) + expr = Expr(rule="rst_role", children=tuple(children)) return expr def ELLIPSES(self, token): @@ -341,7 +393,8 @@ def union(self, tree): out : Expr """ sep = Term(" | ", kind=TermKind.SYNTAX) - expr = Expr(rule="union", children=insert_between(tree.children, sep=sep)) + children = insert_between(tree.children, sep=sep) + expr = Expr(rule="union", children=tuple(children)) return expr def subscription(self, tree): @@ -373,7 +426,7 @@ def param_spec(self, tree): *insert_between(tree.children, sep=sep), Term("]", kind=TermKind.SYNTAX), ] - expr = Expr(rule="param_spec", children=children) + expr = Expr(rule="param_spec", children=tuple(children)) return expr def callable(self, tree): @@ -443,6 +496,7 @@ def literal_item(self, tree): assert not other kind = TermKind.LITERAL if isinstance(item, Term): + # Preserve kind=Name for qualnames in literal expressions kind = item.kind out = Term(item, kind=kind, pos=(tree.meta.start_pos, tree.meta.end_pos)) return out @@ -497,7 +551,7 @@ def dtype(self, tree): ------- out : Expr """ - return Expr(rule="dtype", children=tree.children) + return Expr(rule="dtype", children=tuple(tree.children)) def shape(self, tree): """ @@ -538,7 +592,8 @@ def extra_info(self, tree): logger.debug("dropping extra info") return lark.Discard - def _format_subscription(self, sequence, *, rule): + @staticmethod + def _format_subscription(sequence, *, rule): """Format a `name[...]` style expression. Parameters @@ -565,39 +620,137 @@ def _format_subscription(self, sequence, *, rule): else: children = [container] - expr = Expr(rule=rule, children=children) + expr = Expr(rule=rule, children=tuple(children)) return expr -def parse_doctype(doctype, *, reporter=None): - """Turn a type description in a docstring into a type annotation. +class DoctypeTranspiler: + """Transpile a :term:`doctype` into a structured expression. + + This is could be a simple function, but as a class dependency injection with + custom grammars or transformers is made easier. Parameters ---------- - doctype : str - The doctype to parse. + grammar : str, optional + Use custom grammar to parse doctypes. You can reuse parts of the + default grammar (see ``DOCTYPE_GRAMMAR_PATH`) with + ``%import doctype (start, ...)``. Defaults to ``DOCTYPE_GRAMMAR``. + grammar_import_paths : tuple of (Path, ...) + Additional paths to directories with grammars that are imported in + `grammar`. + transformer : lark.visitors.Transformer, optional + Use a custom transformer to transform the parsed tree. reporter : ~.ContextReporter, optional - - Returns - ------- - parsed : Expr - - Raises - ------ - lark.exceptions.VisitError - Raised when the transformation is interrupted by an exception. - See :class:`lark.exceptions.VisitError`. - BlacklistedQualname - Raised when a qualname is a forbidden keyword. + Pass file and line context for logging. + max_cache_size : int, optional + Size of the LRU cache used to cache transpiled doctypes. Setting it + to 0, disables the cache. Examples -------- - >>> parse_doctype("tuple of (int, ...)") + >>> transpile = DoctypeTranspiler() + >>> transpile("tuple of (int, ...)") - >>> parse_doctype("ndarray of dtype (float or int)") + >>> transpile("ndarray of dtype (float or int)") """ - tree = _lark.parse(doctype) - transformer = DoctypeTransformer(reporter=reporter) - expression = transformer.transform(tree=tree) - return expression + + def __init__( + self, + *, + grammar=None, + grammar_import_paths=(DOCTYPE_GRAMMAR_PATH.parent,), + transformer=None, + reporter=None, + max_cache_size=1000, + ): + if grammar is None: + grammar = DOCTYPE_GRAMMAR + + default_transformer = DoctypeTransformer(reporter=reporter) + if transformer is None: + transformer = default_transformer + + if grammar != DOCTYPE_GRAMMAR: + # A custom grammar may import the existing doctype grammar. In that + # case Lark will put imported rules into their own "doctype__" prefixed + # namespace. By providing `DoctypeTransformer` for that namespace we + # make it work. + transformer = lark.visitors.merge_transformers( + transformer, doctype=default_transformer + ) + + self.max_cache_size = max_cache_size + self.reporter = reporter or ContextReporter(logger=logger) + self._stats = Stats() + self._grammar = grammar + self._grammar_import_paths = grammar_import_paths + self._transformer = transformer + + # Can't easily use `lru_cache` here because it should not be applied to + # methods (B19, https://docs.astral.sh/ruff/rules/cached-instance-method). + # An awkward workaround, assigning the LRU wrapper to the instance + # instead of class, wouldn't be pickle-able. + self._lru_cache = LRUMemoryCache(maxsize=max_cache_size) + + def __call__(self, doctype): + """Transpile single doctype into a structured expression. + + Parameters + ---------- + doctype : str + + Returns + ------- + parsed : Expr + + Raises + ------ + lark.exceptions.VisitError + Raised when the transformation is interrupted by an exception. + See :cls:`lark.exceptions.VisitError`. + BlacklistedQualname + Raised when a qualname is a forbidden keyword. + """ + if doctype in self._lru_cache: + self._stats.inc_counter("doctype_transpiler_cache_hits") + return self._lru_cache[doctype] + + tree = self._parser.parse(doctype) + expression = self._transformer.transform(tree=tree) + assert isinstance(expression, Expr) + + self._lru_cache[doctype] = expression + self._stats.inc_counter("doctype_transpiler_cache_misses") + return expression + + @property + def stats(self): + self._stats.append_to_list( + "doctype_transpiler_curr_cache_size", len(self._lru_cache) + ) + return self._stats + + @cached_property + def _parser(self): + """Initialize Lark parser. + + Use a cached property to construct the lark parser as few times as + possible just before use. Doing that in the ``__init__`` doesn't work + because :class:`lark.Lark` is not pickle-able; and we want to pickle + this class for dependency injection before forking. + + Returns + ------- + parser : lark.Lark + """ + # TODO If we used an LARL conforming grammar we could pass the `transformer` + # directly here which should be faster [1] + # [1] https://lark-parser.readthedocs.io/en/latest/classes.html#lark:~:text=faster + parser = lark.Lark( + grammar=self._grammar, + propagate_positions=True, + import_paths=list(self._grammar_import_paths), + ) + return parser diff --git a/src/docstub/plugins.py b/src/docstub/plugins.py new file mode 100644 index 0000000..23a5f86 --- /dev/null +++ b/src/docstub/plugins.py @@ -0,0 +1,50 @@ +"""Docstub's plugin interface. + +The functions prefixed with ``docstub_`` declare hooks that can be implemented +by plugins. To create a hook, create a function with the same signature and +decorate it with ``hookimpl``: + +>>> from docstub.plugins import hookimpl +... +>>> @hookimpl +... def docstub_load_doctype_grammar(): +... return "my grammar..." + +Attributes +---------- +hookimpl : Final[pluggy.HookimplMarker] +""" + +import pluggy + +_hookspec: pluggy.HookspecMarker = pluggy.HookspecMarker("docstub") + +hookimpl = pluggy.HookimplMarker("docstub") + + +@_hookspec(firstresult=True) +def docstub_load_doctype_grammar(): + """Override the default doctype grammar. + + You can import from doctype's default grammar with:: + + %import doctype (start, WS_INLINE, ...) + %ignore WS_INLINE + + Returns + ------- + grammar : str + The custom grammar to use for parsing doctypes. + """ + + +@_hookspec(firstresult=True) +def docstub_load_doctype_transformer(): + """Override the default transformer for the doctype grammar. + + Returns + ------- + transformer : ~.DoctypeTransformer + Rules must return :class:`~.Expr`, and :class:`~.Term` objects. + Must be pickle-able. + """ diff --git a/stubtest_allow.txt b/stubtest_allow.txt index e9603e4..507985b 100644 --- a/stubtest_allow.txt +++ b/stubtest_allow.txt @@ -1,4 +1,4 @@ docstub\._version\..* docstub._cache.FuncSerializer.__type_params__ -docstub._doctype.TermKind..* +docstub.doctype.TermKind..* (docstub\..*__conditional_annotations__)? diff --git a/tests/test_doctype.py b/tests/test_doctype.py index 3d4b74f..3c84b7b 100644 --- a/tests/test_doctype.py +++ b/tests/test_doctype.py @@ -1,13 +1,21 @@ import logging +from itertools import product import lark import lark.exceptions import pytest -from docstub._doctype import BLACKLISTED_QUALNAMES, parse_doctype +from docstub.doctype import ( + BLACKLISTED_QUALNAMES, + DoctypeTransformer, + DoctypeTranspiler, + Expr, + Term, + insert_between, +) -class Test_parse_doctype: +class Test_DoctypeTranspiler: @pytest.mark.parametrize( "doctype", [ @@ -27,12 +35,12 @@ class Test_parse_doctype: ) def test_edge_case_errors(self, doctype): with pytest.raises(lark.exceptions.UnexpectedInput): - parse_doctype(doctype) + DoctypeTranspiler()(doctype) @pytest.mark.parametrize("doctype", BLACKLISTED_QUALNAMES) def test_reserved_keywords(self, doctype): with pytest.raises(lark.exceptions.VisitError): - parse_doctype(doctype) + DoctypeTranspiler()(doctype) @pytest.mark.parametrize( ("doctype", "expected"), @@ -42,7 +50,7 @@ def test_reserved_keywords(self, doctype): ], ) def test_natlang_union(self, doctype, expected): - expr = parse_doctype(doctype) + expr = DoctypeTranspiler()(doctype) assert expr.as_code() == expected @pytest.mark.parametrize( @@ -67,7 +75,7 @@ def test_natlang_union(self, doctype, expected): ], ) def test_subscription(self, doctype, expected): - expr = parse_doctype(doctype) + expr = DoctypeTranspiler()(doctype) assert expr.as_code() == expected @pytest.mark.parametrize( @@ -93,7 +101,7 @@ def test_subscription(self, doctype, expected): ], ) def test_natlang_container(self, doctype, expected): - expr = parse_doctype(doctype) + expr = DoctypeTranspiler()(doctype) assert expr.as_code() == expected @pytest.mark.parametrize( @@ -113,7 +121,7 @@ def test_natlang_container(self, doctype, expected): ) def test_subscription_error(self, doctype): with pytest.raises(lark.exceptions.UnexpectedInput): - parse_doctype(doctype) + DoctypeTranspiler()(doctype) @pytest.mark.parametrize( ("doctype"), @@ -136,7 +144,7 @@ def test_subscription_error(self, doctype): ], ) def test_literals(self, doctype): - expr = parse_doctype(doctype) + expr = DoctypeTranspiler()(doctype) assert expr.as_code() == doctype assert "literal" in [e.rule for e in expr.sub_expressions] @@ -167,12 +175,12 @@ def test_literals(self, doctype): ], ) def test_natlang_literals(self, doctype, expected): - expr = parse_doctype(doctype) + expr = DoctypeTranspiler()(doctype) assert expr.as_code() == expected assert "natlang_literal" in [e.rule for e in expr.sub_expressions] def test_single_natlang_literal_warning(self, caplog): - expr = parse_doctype("{True}") + expr = DoctypeTranspiler()("{True}") assert expr.as_code() == "Literal[True]" assert caplog.messages == ["Natural language literal with one item: `{True}`"] assert caplog.records[0].levelno == logging.WARNING @@ -205,7 +213,7 @@ def test_single_natlang_literal_warning(self, caplog): ) def test_optional_info(self, doctype, expected, optional_info): doctype_with_optional = doctype + optional_info - expr = parse_doctype(doctype_with_optional) + expr = DoctypeTranspiler()(doctype_with_optional) assert expr.as_code() == expected @pytest.mark.parametrize( @@ -220,7 +228,7 @@ def test_optional_info(self, doctype, expected, optional_info): ], ) def test_callable(self, doctype): - expr = parse_doctype(doctype) + expr = DoctypeTranspiler()(doctype) assert expr.as_code() == doctype assert "callable" in [e.rule for e in expr.sub_expressions] @@ -233,7 +241,7 @@ def test_callable(self, doctype): ], ) def test_callable_subscriptions_form(self, doctype): - expr = parse_doctype(doctype) + expr = DoctypeTranspiler()(doctype) assert expr.as_code() == doctype assert "callable" not in [e.rule for e in expr.sub_expressions] @@ -247,7 +255,7 @@ def test_callable_subscriptions_form(self, doctype): ) def test_callable_error(self, doctype): with pytest.raises(lark.exceptions.UnexpectedInput): - parse_doctype(doctype) + DoctypeTranspiler()(doctype) @pytest.mark.parametrize( ("doctype", "expected"), @@ -261,7 +269,7 @@ def test_callable_error(self, doctype): ], ) def test_rst_role(self, doctype, expected): - expr = parse_doctype(doctype) + expr = DoctypeTranspiler()(doctype) assert expr.as_code() == expected # fmt: off @@ -280,7 +288,7 @@ def test_rst_role(self, doctype, expected): def test_natlang_array(self, fmt, expected_fmt, name, dtype, shape): doctype = fmt.format(name=name, dtype=dtype, shape=shape) expected = expected_fmt.format(name=name, dtype=dtype, shape=shape) - expr = parse_doctype(doctype) + expr = DoctypeTranspiler()(doctype) assert expr.as_code() == expected assert "natlang_array" in [e.rule for e in expr.sub_expressions] # fmt: on @@ -293,7 +301,7 @@ def test_natlang_array(self, fmt, expected_fmt, name, dtype, shape): ], ) def test_natlang_array_specific(self, doctype, expected): - expr = parse_doctype(doctype) + expr = DoctypeTranspiler()(doctype) assert expr.as_code() == expected assert "natlang_array" in [e.rule for e in expr.sub_expressions] @@ -301,4 +309,201 @@ def test_natlang_array_specific(self, doctype, expected): def test_natlang_array_invalid_shape(self, shape): doctype = f"array of shape {shape}" with pytest.raises(lark.exceptions.UnexpectedInput): - _ = parse_doctype(doctype) + _ = DoctypeTranspiler()(doctype) + + +class Test_DoctypeTranspiler_override: + @pytest.mark.parametrize( + ("doctype", "expected"), + [ + ("ndarray of dtype uint8 or int", "ndarray[uint8] | int"), + ("ndarray of dtype in {uint8, uint16}", "ndarray[uint8 | uint16]"), + # ("ndarray of dtype (uint8, uint16)", "ndarray[uint8 | uint16]"), + ], + ) + def test_in_dtype(self, doctype, expected): + grammar = """ + %import doctype (start, dtype, qualname, WS_INLINE) + dtype_union : "in {" qualname ("," qualname)+ "}" -> union + %override dtype: qualname | dtype_union + %ignore WS_INLINE + """ + transpile = DoctypeTranspiler(grammar=grammar) + expr = transpile(doctype) + assert expr.as_code() == expected + assert "natlang_array" in [e.rule for e in expr.sub_expressions] + assert "dtype" in [e.rule for e in expr.sub_expressions] + + @pytest.mark.parametrize( + ("doctype", "expected"), + [ + ("my_array of dtype int", "my_array[int]"), + ("my_array of shape (4, 4)", "my_array"), + ], + ) + def test_doc_example_new_array_name(self, doctype, expected): + grammar = """ + %import doctype (start, WS_INLINE, ARRAY_NAME) + %extend ARRAY_NAME: "my_array" + %ignore WS_INLINE + """ + transpile = DoctypeTranspiler(grammar=grammar) + expr = transpile(doctype) + assert expr.as_code() == expected + assert len(list(expr.filter(rule="natlang_array"))) == 1 + + @pytest.mark.parametrize( + ("doctype", "expected"), + [ + ("int or float or complex", "complex | float | int"), + ("int or tuple[int or float, ...]", "int | tuple[float | int, ...]"), + ], + ) + def test_transformer_only(self, doctype, expected): + @lark.visitors.v_args(tree=True) + class Transformer(DoctypeTransformer): + def union(self, tree): + # Sort items in union alphabetically + tree.children = sorted(tree.children, key=str) + out = super().union(tree) + return out + + transpile = DoctypeTranspiler(transformer=Transformer()) + expr = transpile(doctype) + assert expr.as_code() == expected + assert "union" in [e.rule for e in expr.sub_expressions] + + @pytest.mark.parametrize( + ("doctype", "expected"), + [ + ("(int) -> bool", "Callable[[int], bool]"), + ("(str, bytes) -> int or bool", "Callable[[str, bytes], int | bool]"), + ], + ) + def test_natlang_callable(self, doctype, expected): + grammar = """ + %import doctype (start, _type, WS_INLINE) + params: _type ("," _type)* + natlang_callable : "(" params ") ->" _type + %extend _type: natlang_callable + %ignore WS_INLINE + """ + + class Transformer(lark.visitors.Transformer): + def start(self, children): + return Expr(rule="start", children=tuple(children)) + + def params(self, children): + sep = Term(", ", kind="syntax") + content = insert_between(children, sep=sep) + items = [ + Term("[", kind="syntax"), + *content, + Term("]", kind="syntax"), + ] + out = Expr(rule="params", children=tuple(items)) + return out + + def natlang_callable(self, children): + items = [Term("Callable", kind="name"), *children] + out = DoctypeTransformer._format_subscription( + items, rule="natlang_callable" + ) + return out + + transpile = DoctypeTranspiler(grammar=grammar, transformer=Transformer()) + expr = transpile(doctype) + assert expr.as_code() == expected + assert "natlang_callable" in [e.rule for e in expr.sub_expressions] + assert "params" in [e.rule for e in expr.sub_expressions] + + @pytest.mark.parametrize( + ("doctype", "expected"), + [ + ("{Iterable, Sequence} of int(s)", "Iterable[int] | Sequence[int]"), + ( + "{list, set} of {int, float}", + "list[int] | list[float] | set[int] | set[float]", + ), + ], + ) + def test_expansion_of_qualnames(self, doctype, expected): + grammar = """ + %import doctype (start, WS_INLINE, qualname) + %extend qualname: "{" qualname ("," qualname)+ "}" -> expansion + %ignore WS_INLINE + """ + + @lark.visitors.v_args(tree=True) + class ExpansionTransformer(DoctypeTransformer): + def start(self, tree): + """Transform "start" rule. + + Parameters + ---------- + tree : lark.Tree + + Returns + ------- + expr : Expr + """ + children = map(self._expand_expression, tree.children) + return Expr(rule="start", children=tuple(children)) + + def expansion(self, tree): + """Transform "expansion" rule. + + Actual expansion happens in `start` when the complete expression with + all potential expansions can be seen. + + Parameters + ---------- + tree : lark.Tree + + Returns + ------- + expr : Expr + """ + return Expr(rule="expansion", children=tuple(tree.children)) + + def _expand_expression(self, expr): + """Join combination of all expansions in a new union. + + Parameters + ---------- + expr : Expr + + Returns + ------- + expanded : Expr + """ + expansions = expr.filter(rule="expansion") + if not expansions: + return expr + + # Create Cartesian product of all expansions where each + # value is prefixed by the expansion it will replace + expansions = [ + [(expansion, value) for value in expansion.children] + for expansion in expansions + ] + combinations = list(product(*expansions)) + + # Create a new expression for each combination + interpolated = [] + for combination in combinations: + new_expr = expr + for expansion, value in combination: + new_expr = new_expr.replace(expansion, value) + interpolated.append(new_expr) + + # Join new expressions as union + expanded = self.union(lark.Tree("union", children=interpolated)) + return expanded + + transpile = DoctypeTranspiler( + grammar=grammar, transformer=ExpansionTransformer() + ) + expr = transpile(doctype) + assert expr.as_code() == expected + assert list(expr.filter(rule="natlang_array")) == []