diff --git a/docs/conf.py b/docs/conf.py
index 13c76a1..e45ffe0 100644
--- a/docs/conf.py
+++ b/docs/conf.py
@@ -44,6 +44,7 @@
intersphinx_mapping = {
"python": ("https://docs.python.org/3/", None),
+ "numpy": ("https://numpy.org/doc/stable/", None),
"typing": ("https://typing.python.org/en/latest/", None),
"lark": ("https://lark-parser.readthedocs.io/en/latest/", None),
"libcst": ("https://libcst.readthedocs.io/en/latest/", None),
diff --git a/docs/create_plugin.md b/docs/create_plugin.md
index ff18bd7..02982c3 100644
--- a/docs/create_plugin.md
+++ b/docs/create_plugin.md
@@ -26,9 +26,7 @@ from docstub.doctype import DoctypeTransformer, Expr, Term
my_grammar = """
%import doctype (start, WS_INLINE, ARRAY_NAME)
-
%extend ARRAY_NAME: "my_array"
-
%ignore WS_INLINE
"""
@@ -62,13 +60,11 @@ from itertools import product
import lark
from docstub.plugins import hookimpl
-from docstub.doctype import DoctypeTransformer, Expr, Term
+from docstub.doctype import DoctypeTransformer, Expr, Term, expand_expression
my_grammar = """
%import doctype (start, WS_INLINE, qualname)
-
%extend qualname: "{" qualname ("," qualname)+ "}" -> expansion
-
%ignore WS_INLINE
"""
@@ -86,59 +82,9 @@ class ExpansionTransformer(DoctypeTransformer):
-------
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
+ assert len(tree.children) == 1
+ child = expand_expression(tree.children[0], expand_rule="expansion")
+ return Expr(rule="start", children=(child,))
@hookimpl
diff --git a/docs/introduction.md b/docs/introduction.md
index 4ac49ba..2116e6e 100644
--- a/docs/introduction.md
+++ b/docs/introduction.md
@@ -19,7 +19,7 @@ def example_metric(image, *, mask=None, sigma=1.0, method='standard'):
----------
image : array-like
First image.
- mask : array of dtype uint8, optional
+ mask : array of dtype uint8 and shape (M, N), optional
Second image.
sigma : float or Iterable of float, default: 1.0
Sigma value for each dimension in `image`. A single value is broadcast
@@ -51,7 +51,7 @@ will create the [stub file](https://typing.python.org/en/latest/spec/distributin
# File generated with docstub
from collections.abc import Iterable
-from typing import Literal
+from typing import Any, Literal
import numpy as np
from numpy.typing import ArrayLike, NDArray
@@ -59,10 +59,10 @@ from numpy.typing import ArrayLike, NDArray
def example_metric(
image: ArrayLike,
*,
- mask: NDArray[np.uint8] | None = ...,
+ mask: NDArray[tuple[int, int], np.dtype[np.uint8]] | None = ...,
sigma: float | Iterable[float] = ...,
method: Literal["standard", "modified"] = ...
-) -> NDArray[float]: ...
+) -> NDArray[tuple[Any, ...], np.dtype[float]]: ...
```
diff --git a/docs/typing_syntax.md b/docs/typing_syntax.md
index 4f2fb64..ffbbcd7 100644
--- a/docs/typing_syntax.md
+++ b/docs/typing_syntax.md
@@ -1,9 +1,8 @@
-# Typing syntax in docstrings
+# Extended typing syntax
-Docstub defines its own [grammar](../src/docstub/doctype.lark) to parse and transform type information in docstrings into valid Python type expressions.
-This grammar fully supports [Python's conventional typing syntax](https://typing.python.org/en/latest/index.html).
-So any {term}`annotation expression` that is valid in Python, can be used in a docstrings as is.
-In addition, docstub extends this syntax with several "natural language" expressions that are commonly used in the scientific Python ecosystem.
+Any valid Python {term}`annotation expression` can be used in a docstrings as is.
+However, Docstub extends [Python's conventional typing syntax](https://typing.python.org/en/latest/index.html) with several "natural language" expressions that are commonly used in the scientific Python ecosystem.
+This extended syntax is defined in its own [grammar](#appendix-grammar).
Docstrings should follow a form that is inspired by the [NumPyDoc style](https://numpydoc.readthedocs.io/en/latest/format.html):
```none
@@ -28,10 +27,10 @@ Combining multiple names that share a doctype and description is supported.
In addition to Python's conventional shorthand `|` syntax for [union types](https://typing.python.org/en/latest/spec/concepts.html#union-types), you can use `or` to join types.
-| Docstring type | Python type annotation |
-|----------------|------------------------|
-| `X or Y` | `X \| Y` |
-| `int or float` | `int \| float` |
+| doctype | Python annotation |
+|----------------|-------------------|
+| `X or Y` | `X \| Y` |
+| `int or float` | `int \| float` |
## Containers
@@ -39,64 +38,68 @@ In addition to Python's conventional shorthand `|` syntax for [union types](http
The content of containers can be typed using a `CONTAINER of X` like form.
This extends the basic subscription syntax for [generics](https://typing.python.org/en/latest/spec/generics.html#generics).
-| Docstring type | Python type annotation |
-|-------------------------|------------------------|
-| `CONTAINER of X` | `CONTAINER[X]` |
-| `CONTAINER of (X or Y)` | `CONTAINER[X \| Y]` |
+| doctype | Python annotation |
+|-------------------------|---------------------|
+| `CONTAINER of X` | `CONTAINER[X]` |
+| `CONTAINER of (X or Y)` | `CONTAINER[X \| Y]` |
For the simple case `CONTAINER of X`, where `X` is a name, you can append `(s)` to indicate the plural form.
For example, `list of float(s)`.
Variants of for [**tuples**](https://typing.python.org/en/latest/spec/tuples.html)
-| Docstring type | Python type annotation |
-|---------------------|------------------------|
-| `tuple of (X, Y)` | `tuple[X, Y]` |
-| `tuple of (X, ...)` | `tuple[X, ...]` |
+| doctype | Python annotation |
+|---------------------|-------------------|
+| `tuple of (X, Y)` | `tuple[X, Y]` |
+| `tuple of (X, ...)` | `tuple[X, ...]` |
and **mappings** exist.
-| Docstring type | Python type annotation |
-|----------------------|------------------------|
-| `MAPPING of {X: Y}` | `MAPPING[X, Y]` |
-| `dict of {str: int}` | `dict[str, int]` |
+| doctype | Python annotation |
+|----------------------|-------------------|
+| `MAPPING of {X: Y}` | `MAPPING[X, Y]` |
+| `dict of {str: int}` | `dict[str, int]` |
:::{tip}
-While it is possible to nest these variants repeatedly, it decreases the readability.
+While it is possible to nest these variants, it decreases the readability.
For complex nested annotations with nested containers, consider using Python's conventional syntax.
In the future, docstub may warn against or disallow nesting these natural language variants.
:::
(natlang-array)=
-## Shape and dtype syntax for arrays
+## NDArrays with shape and dtype
-This expression allows adding shape and datatype information for data structures like [NumPy arrays](https://numpy.org/doc/stable/reference/generated/numpy.ndarray.html).
+This "natural language" expression describes shape and dtype information for [NumPy arrays](https://numpy.org/doc/stable/reference/generated/numpy.ndarray.html).
+The basic form is as follows:
-`array` and `ndarray`, and `array-like` and `array_like` can be used interchange-ably for the variable `ARRAY` below.
+| doctype | Python annotation |
+|----------------------------------------|-------------------------------------------|
+| `ARRAY of shape SHAPE` | `ARRAY[SHAPE]` |
+| `ARRAY of dtype DTYPE` | `ARRAY[tuple[Any, ...], np.dtype[DTYPE]]` |
+| `ARRAY of dtype DTYPE and shape SHAPE` | `ARRAY[SHAPE, np.dtype[DTYPE]]` |
+| `ARRAY of shape SHAPE and dtype DTYPE` | `ARRAY[SHAPE, np.dtype[DTYPE]]` |
-| Docstring type | Python type annotation |
-|----------------------------------------|------------------------|
-| `ARRAY of dtype DTYPE` | `ARRAY[DTYPE]` |
-| `ARRAY of dtype DTYPE and shape SHAPE` | `ARRAY[DTYPE]` |
-| `ARRAY of shape SHAPE` | `ARRAY[DTYPE]` |
-| `ARRAY of shape SHAPE and dtype DTYPE` | `ARRAY[DTYPE]` |
+Only `array`, `ndarray`, `array-like` and `array_like` can be used in place of `ARRAY` above.
+`array-like` and `array_like` map to {obj}`numpy.typing.ArrayLike` which does not support adding shape and dtype information.
+In this case, Docstub omits the dtype and shape automatically.
-For example
+`SHAPE` can be given in forms like `(n,)`, `(M, N)`, `(x, y, 4)`, or `(..., 2)`.
+Optional dimensions are supported by enclosing them in square brackets like `(x, y[, z])`.
-| Docstring type | Python type annotation |
-|------------------------------------------|------------------------|
-| `array of dtype int` | `ndarray[int]` |
-| `ndarray of dtype bool and shape (4, 4)` | `ndarray[bool]` |
-| `array-like of dtype float` | `ArrayLike[float]` |
-| `array_like of shape (M, 2)` | `ArrayLike` |
+`DTYPE` can be any valid name or a union of names.
+For example:
-:::{note}
-Noting the **shape** of an array in the docstring is supported.
-However, [support for including shapes in generated stubs](https://github.com/scientific-python/docstub/issues/76) is not yet included in docstub.
-:::
+| doctype | Python annotation |
+|------------------------------------------|--------------------------------------------|
+| `array of dtype int` | `NDArray[tuple[Any, ...], np.dtype[int]]` |
+| `array of shape (..., 2)` | `NDArray[tuple[Any, ...]]` |
+| `array of shape (x, y[, z])` | `NDArray[tuple[int, int]]`
`\| NDArray[tuple[int, int, int]]` |
+| `ndarray of dtype bool and shape (4, 4)` | `NDArray[tuple[int, int], np.dtype[bool]]` |
+| `array-like of dtype float` | `ArrayLike` |
+| `array_like of shape (M, 2)` | `ArrayLike` |
## Literals
@@ -104,7 +107,7 @@ However, [support for including shapes in generated stubs](https://github.com/sc
[Literals](https://typing.python.org/en/latest/spec/literal.html#literals) indicate a concrete value instead of type.
Instead of using [`typing.Literal`](https://docs.python.org/3/library/typing.html#typing.Literal), you can enclose literal values in `{...}` in docstrings.
-| Docstring type | Python type annotation |
+| doctype | Python annotation |
|---------------------------|----------------------------------|
| `{-1, 0, 3, True, False}` | `Literal[-1, 0, 3, True, False]` |
| `{"red", "blue", None}` | `Literal["red", "blue", None]` |
@@ -117,19 +120,27 @@ However, `Literal[X]` is more explicit.
:::{warning}
Python's `typing.Literal` only supports a restricted set of parameters.
For example, `float` literals are not yet supported by the type system but are allowed by docstub.
-Addressing this use case is on the roadmap.
-See [issue 47](https://github.com/scientific-python/docstub/issues/47) for more details.
+Use a broader inlined type such as `Literal[0, 1] | float` if you want to use something like `{0, .5, 1}` in the docstring.
:::
## reStructuredText role
Since docstrings are also used to generate documentation with Sphinx, you may want to use [restructuredText roles](https://docutils.sourceforge.io/docs/ref/rst/roles.html).
-This is supported anywhere in where a {term}`type name` is used in a {term}`doctype`.
-
-| Docstring type | Python type annotation |
-|-------------------------|------------------------|
-| `` `X` `` | `X` |
-| ``:ref:`X` `` | `X` |
-| ``:class:`X`[Y, ...] `` | `X[Y, ...]` |
-| ``:class:`Y.X` `` | `Y.X` |
-| ``:py:class:`Y.X` `` | `Y.X` |
+This is supported anywhere in where a {term}`type name` used in a {term}`doctype`.
+
+| doctype | Python annotation |
+|-------------------------|-------------------|
+| `` `X` `` | `X` |
+| ``:ref:`X` `` | `X` |
+| ``:class:`X`[Y, ...] `` | `X[Y, ...]` |
+| ``:class:`Y.X` `` | `Y.X` |
+| ``:py:class:`Y.X` `` | `Y.X` |
+
+
+---
+
+## Appendix: Grammar
+
+:::{literalinclude} ../src/docstub/doctype.lark
+:caption: doctype.lark
+:::
diff --git a/examples/example_pkg-stubs/_numpy.pyi b/examples/example_pkg-stubs/_numpy.pyi
index 590c0a7..8fed556 100644
--- a/examples/example_pkg-stubs/_numpy.pyi
+++ b/examples/example_pkg-stubs/_numpy.pyi
@@ -1,5 +1,7 @@
# File generated with docstub
+from typing import Any
+
import numpy
import numpy as np
from numpy.typing import ArrayLike, NDArray
@@ -8,8 +10,14 @@ def func_object_with_numpy_objects(
a1: numpy.int8, a2: np.int16, a3: numpy.typing.DTypeLike, a4: np.typing.DTypeLike
) -> None: ...
def func_ndarray(
- a1: NDArray, a2: np.NDArray, a3: NDArray[float], a4: NDArray[np.uint8] | None = ...
-) -> tuple[NDArray[np.uint8], NDArray[complex]]: ...
+ a1: NDArray,
+ a2: np.NDArray,
+ a3: NDArray[tuple[int, int], np.dtype[float]],
+ a4: NDArray[tuple[int], np.dtype[np.uint8]] | None = ...,
+) -> tuple[
+ NDArray[tuple[Any, ...], np.dtype[np.uint8]],
+ NDArray[tuple[Any, ...], np.dtype[complex]],
+]: ...
def func_array_like(
- a1: ArrayLike, a2: ArrayLike, a3: ArrayLike[float], a4: ArrayLike[np.uint8]
-) -> tuple[ArrayLike[np.uint8], ArrayLike[complex]]: ...
+ a1: ArrayLike, a2: ArrayLike, a3: ArrayLike, a4: ArrayLike
+) -> tuple[ArrayLike, ArrayLike]: ...
diff --git a/src/docstub-stubs/_docstrings.pyi b/src/docstub-stubs/_docstrings.pyi
index 937f3e3..95fd144 100644
--- a/src/docstub-stubs/_docstrings.pyi
+++ b/src/docstub-stubs/_docstrings.pyi
@@ -15,7 +15,7 @@ import numpydoc.docscrape as npds
from ._analysis import PyImport, TypeMatcher
from ._report import ContextReporter, Stats
from ._utils import escape_qualname
-from .doctype import BlacklistedQualname, DoctypeTranspiler, TermKind
+from .doctype import DoctypeTranspiler, QualnameIsForbidden, TermKind
logger: logging.Logger
diff --git a/src/docstub-stubs/doctype.pyi b/src/docstub-stubs/doctype.pyi
index cac18dd..e15111b 100644
--- a/src/docstub-stubs/doctype.pyi
+++ b/src/docstub-stubs/doctype.pyi
@@ -6,9 +6,10 @@ import logging
from collections.abc import Callable, Generator, Iterable, Sequence
from dataclasses import dataclass
from functools import cached_property
+from itertools import product
from pathlib import Path
from textwrap import indent
-from typing import Any, Final, Self
+from typing import Any, ClassVar, Final, Self
import lark
import lark.visitors
@@ -25,6 +26,7 @@ DOCTYPE_GRAMMAR: Final[str]
def flatten_recursive(iterable: Iterable[Iterable | str]) -> Generator[str]: ...
def insert_between(iterable: Iterable, *, sep: Any) -> list[Any]: ...
+def expand_expression(expr: Expr, *, expand_rule: str) -> Expr: ...
class TermKind(enum.StrEnum):
@@ -53,10 +55,6 @@ class Expr:
def __post_init__(self) -> None: ...
@property
def terms(self) -> list[Term]: ...
- @property
- def names(self) -> list[Term]: ...
- @property
- def sub_expressions(self) -> list[Self]: ...
def __iter__(self) -> Generator[Expr | Term]: ...
def format_tree(self) -> str: ...
def print_tree(self) -> None: ...
@@ -64,17 +62,30 @@ class Expr:
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
+ def _filter_iter(
+ self, func: Callable | None = ..., /, **attrs: Any
) -> Generator[Expr | Term]: ...
+ def find_all(
+ self, func: Callable | None = ..., /, **attrs: Any
+ ) -> list[Expr | Term]: ...
+ def try_find_one(
+ self, func: Callable | None = ..., /, **attrs: Any
+ ) -> Expr | Term | None: ...
+ def find_one(self, func: Callable | None = ..., /, **attrs: Any) -> Expr | Term: ...
-BLACKLISTED_QUALNAMES: Final[set[str]]
+FORBIDDEN_KEYWORDS: Final[set[str]]
-class BlacklistedQualname(DocstubError):
+class QualnameIsForbidden(DocstubError):
pass
class DoctypeTransformer(lark.visitors.Transformer):
+
+ any_shape_expr: ClassVar[Expr]
+
def __init__(self, *, reporter: ContextReporter | None = ...) -> None: ...
+ def __default__(
+ self, data: str, children: list[Expr | Term], meta: lark.tree.Meta
+ ) -> Expr: ...
def start(self, tree: lark.Tree) -> Expr: ...
def qualname(self, tree: lark.Tree) -> Term: ...
def rst_role(self, tree: lark.Tree) -> Expr: ...
@@ -90,11 +101,12 @@ class DoctypeTransformer(lark.visitors.Transformer):
def natlang_array(self, tree: lark.Tree) -> Expr: ...
def array_name(self, tree: lark.Tree) -> Term: ...
def dtype(self, tree: lark.Tree) -> Expr: ...
- def shape(self, tree: lark.Tree) -> lark.visitors._DiscardType: ...
+ def dtype_to_expand(self, tree: lark.Tree) -> Expr: ...
+ def shape(self, tree: lark.Tree) -> Expr: ...
def optional_info(self, tree: lark.Tree) -> lark.visitors._DiscardType: ...
def extra_info(self, tree: lark.Tree) -> lark.visitors._DiscardType: ...
@staticmethod
- def _format_subscription(sequence: Sequence[str], *, rule: str) -> Expr: ...
+ def _create_subscription_expr(children: Sequence[str], *, rule: str) -> Expr: ...
class DoctypeTranspiler:
def __init__(
diff --git a/src/docstub/_app_generate_stubs.py b/src/docstub/_app_generate_stubs.py
index 2355aff..541b5b1 100644
--- a/src/docstub/_app_generate_stubs.py
+++ b/src/docstub/_app_generate_stubs.py
@@ -473,7 +473,7 @@ def generate_stubs(
if total_errors:
logger.error("Total errors: %i", total_errors)
- assert len(stats) == 0
+ assert len(stats) == 0, stats
total_fails = total_errors
if fail_on_warning:
diff --git a/src/docstub/_docstrings.py b/src/docstub/_docstrings.py
index 9073df5..a57badb 100644
--- a/src/docstub/_docstrings.py
+++ b/src/docstub/_docstrings.py
@@ -18,7 +18,7 @@
from ._analysis import PyImport, TypeMatcher
from ._report import ContextReporter, Stats
from ._utils import escape_qualname
-from .doctype import BlacklistedQualname, DoctypeTranspiler, TermKind
+from .doctype import DoctypeTranspiler, QualnameIsForbidden, TermKind
logger: logging.Logger = logging.getLogger(__name__)
@@ -276,8 +276,8 @@ def doctype_to_annotation(
except lark.visitors.VisitError as error:
original_error = error.orig_exc
- if isinstance(original_error, BlacklistedQualname):
- msg = "Blacklisted keyword argument in doctype"
+ if isinstance(original_error, QualnameIsForbidden):
+ msg = "Forbidden keyword argument in doctype"
details = _red_partial_underline(
doctype,
start=error.obj.meta.start_pos,
diff --git a/src/docstub/_stubs.py b/src/docstub/_stubs.py
index 398b856..029b153 100644
--- a/src/docstub/_stubs.py
+++ b/src/docstub/_stubs.py
@@ -529,6 +529,8 @@ def leave_FunctionDef(self, original_node, updated_node):
annotation_value = ds_annotations.returns.value
if original_node.returns is None:
+ # TODO Supply cst config! Also on other places
+ # https://libcst.readthedocs.io/en/latest/best_practices.html#provide-a-config-when-generating-code-from-templates
annotation = cst.Annotation(cst.parse_expression(annotation_value))
node_changes["returns"] = annotation
# TODO: check imports
diff --git a/src/docstub/doctype.lark b/src/docstub/doctype.lark
index 4bd6826..65015ac 100644
--- a/src/docstub/doctype.lark
+++ b/src/docstub/doctype.lark
@@ -17,8 +17,12 @@ _annotation_with_meta: _type ("," optional_info)?
// A type annotation. Can range from a simple qualified name to a complex
// nested construct of types.
-_type: qualname
- | union
+_type: _type_without_union | union
+
+
+// Used in `union` to prevent directly nested unions. Without
+// `a | b | c` is parsed as `union(a, union(b, c))`.
+_type_without_union: qualname
| literal
| subscription
| callable
@@ -38,7 +42,7 @@ qualname: (/~/ ".")? (NAME ".")* NAME
// An union of different types, joined either by "or" or "|".
-union: _type (_OR _type)+
+union: _type_without_union (_OR _type_without_union)+
// Operator used in unions.
@@ -46,7 +50,9 @@ _OR: "or" | "|"
// An expression where an object is subscribed with "A[v, ...]".
-subscription: qualname "[" _type ("," _type)* ("," ELLIPSES)? "]"
+// Assign higher priority (`.1`) so it wins over the `literal` rule which has
+// overlap with this one. See description for `literal_item` for more.
+subscription.1: qualname "[" _type ("," _type)* ("," ELLIPSES)? "]"
// An expression describing a callable like "Callable[[int], str]"
@@ -70,10 +76,8 @@ natlang_literal: "{" literal_item ("," literal_item)* "}"
// A literal expression as supported by Python proper. The rule "subscription"
-// isn't allowed to contain "literal_items", so we need to define this.
-// Assign a higher priority so that things like `Literal[Some.ENUM]` is marked
-// as a literal expression.
-literal.1: qualname "[" literal_item ("," literal_item)* "]"
+// isn't allowed to contain "literal_items", so this rule covers that case.
+literal: qualname "[" literal_item ("," literal_item)* "]"
// An single item in a literal expression (or `optional`). We must also allow
@@ -111,7 +115,7 @@ natlang_array: array_name "of dtype" dtype ("and shape" shape)?
| array_name "of shape" shape ("and dtype" dtype)?
-// Currently a bit of a hack. Since the `array_expression` is ambiguous, we
+// Currently a bit of a hack. Since `natlang_array` is ambiguous, we
// want to make sure it only works for real arrays. For now, we are using a
// hack here, that only allows specific names in `array_name`. In the
// transformer we alias this to qualname.
@@ -122,30 +126,30 @@ ARRAY_NAME: "array" | "ndarray" | "array-like" | "array_like"
// The dtype used in an array expression.
-dtype: qualname | "(" union ")"
+dtype: qualname | dtype_to_expand
+
+
+// Dtypes may be specified as a union. But I'm not sure if `np.dtype` accepts
+// a union. So we use expansion here instead.
+dtype_to_expand: "(" qualname (_OR qualname)+ ")"
// The shape used in an array expression. Possibly to liberal right now in
// what it allows. Since there is currently no support to type the shape of
// NumPy arrays, this information is dropped during the transformation.
-shape: "(" dim ",)"
- | "(" leading_optional_dim? dim (("," dim | insert_optional_dim))* ")"
- | INT "-"? "D"
-
-
-// Optional dimensions in a `shape` expression placed at the start,
-// e.g., `([3 ,] N)`.
-?leading_optional_dim: "[" dim ("," dim)* ",]"
+shape: "()"
+ | "(" dim ",)"
+ | "(" "[" optional_dimensions "," "]" dim ")"
+ | "(" ("[" optional_dimensions "," "]")? dim ("," dim | "[" "," optional_dimensions "]")+ ")"
-// Optional dimensions in a `shape` expression placed anywhere but the start,
-// e.g., `(A[, B], C[, D])`.
-?insert_optional_dim: "[," dim ("," dim)* "]"
+// An optional dimension used inside a `shape` rule
+optional_dimensions: dim ("," dim)*
// Dimension can be a number, ellipses ('...') or a simple name. A simple name
// can be bound to a specific number, e.g. `N=3`.
-?dim: INT | ELLIPSES | NAME ("=" INT)?
+dim: INT | ELLIPSES | NAME ("=" INT)?
// Optional meta information added after the docstring annotation, e.g.
diff --git a/src/docstub/doctype.py b/src/docstub/doctype.py
index 8c23f4f..610c07a 100644
--- a/src/docstub/doctype.py
+++ b/src/docstub/doctype.py
@@ -12,9 +12,10 @@
from collections.abc import Iterable
from dataclasses import dataclass
from functools import cached_property
+from itertools import product
from pathlib import Path
from textwrap import indent
-from typing import Final, Self
+from typing import ClassVar, Final, Self
import lark
import lark.visitors
@@ -83,6 +84,45 @@ def insert_between(iterable, *, sep):
return out[:-1]
+def expand_expression(expr, *, expand_rule):
+ """Create union of `expr` form a combination of all matching children.
+
+ Parameters
+ ----------
+ expr : Expr
+ The expression to expand into multiple ones.
+ expand_rule : str
+
+ Returns
+ -------
+ expanded : Expr
+ """
+ expansions = expr.find_all(rule=expand_rule)
+ 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
+ expanded = []
+ for combination in combinations:
+ new_expr = expr
+ for expansion, value in combination:
+ new_expr = new_expr.replace(expansion, value)
+ expanded.append(new_expr)
+
+ # Join new expressions as union
+ sep = Term(" | ", kind=TermKind.SYNTAX)
+ expanded = insert_between(expanded, sep=sep)
+ expanded = Expr(rule="union", children=tuple(expanded))
+ return expanded
+
+
class TermKind(enum.StrEnum):
"""Encodes the different kinds of :class:`Term`."""
@@ -149,13 +189,13 @@ def __getnewargs_ex__(self):
@dataclass(slots=True, frozen=True)
class Expr:
- """An annotation expression.
+ """An expression that groups a logical part of an annotation expression.
Parameters
----------
- rule :
+ rule : str
The name of the (grammar) rule corresponding to this expression.
- children : list of (Expr or Term)
+ children : Iterable of (Expr or Term)
Sub-expressions or terms that make up this expression.
"""
@@ -177,30 +217,6 @@ def terms(self):
"""
return list(flatten_recursive(self))
- @property
- def names(self):
- """Name terms in the expression.
-
- Returns
- -------
- names : list of Term
- """
- return self.filter(kind=TermKind.NAME)
-
- @property
- def sub_expressions(self):
- """Iterate expressions inside the current one.
-
- Returns
- -------
- names : list of Self
- """
- cls = type(self)
- for child in self.children:
- if isinstance(child, cls):
- yield child
- yield from child.sub_expressions
-
def __iter__(self):
"""Iterate over children of this expression.
@@ -240,8 +256,11 @@ def __str__(self) -> str:
def as_code(self) -> str:
return str(self)
+ # TODO Candidate for optimization, benchmark
+ # Introduce `count` parameter, break iteration earlier?
def replace(self, old, new):
- """
+ """Replace children in current expression and return new one.
+
Parameters
----------
old : Expr or Term
@@ -254,21 +273,24 @@ def replace(self, old, new):
cls = type(self)
new_children = []
for child in self.children:
- if isinstance(child, cls):
- child = child.replace(old, new) # noqa: PLW2901
+ new_child = None
if child == old:
- child = new # noqa: PLW2901
- new_children.append(child)
- return cls(rule=self.rule, children=tuple(new_children))
+ new_child = new
+ elif isinstance(child, cls):
+ # Potentially replace in child's children
+ new_child = child.replace(old, new)
+ new_children.append(new_child or child)
+ replaced = cls(rule=self.rule, children=tuple(new_children))
+ return replaced
- def filter(self, func=None, /, **attributes):
- """Filter children of the expression recursively.
+ def _filter_iter(self, func=None, /, **attrs):
+ """Find matching children in the expression recursively.
Parameters
----------
func : Callable, optional
Only return children for which ``func(child)`` is true.
- attributes : Any
+ attrs : Any, optional
Only return children for which have an attribute with a matching
value.
@@ -280,20 +302,79 @@ def filter(self, func=None, /, **attributes):
matches = True
if func:
matches = func(child)
- for name, value in attributes.items():
+ for name, value in attrs.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)
+ if hasattr(child, "_filter_iter"):
+ yield from child._filter_iter(func, **attrs)
+
+ def find_all(self, func=None, /, **attrs):
+ """Find all matching children in the expression recursively.
+
+ Parameters
+ ----------
+ func : Callable, optional
+ Only return children for which ``func(child)`` is true.
+ attrs : Any, optional
+ Only return children for which have an attribute with a matching
+ value.
+
+ Returns
+ -------
+ matches : list of (Expr or Term)
+ """
+ return list(self._filter_iter(func, **attrs))
+
+ def try_find_one(self, func=None, /, **attrs):
+ """Try to find one matching child in expression or return None
+
+ Parameters
+ ----------
+ func : Callable, optional
+ Only return children for which ``func(child)`` is true.
+ attrs : Any, optional
+ Only return children for which have an attribute with a matching
+ value.
+
+ Returns
+ -------
+ match : Expr or Term or None
+ """
+ matches = self.find_all(func, **attrs)
+ if len(matches) == 1:
+ return matches[0]
+ if len(matches) > 1:
+ raise RuntimeError(f"found multiple matches: {matches}")
+ return None
+
+ def find_one(self, func=None, /, **attrs):
+ """Find exactly one matching child in expression.
+
+ Parameters
+ ----------
+ func : Callable, optional
+ Only return children for which ``func(child)`` is true.
+ attrs : Any, optional
+ Only return children for which have an attribute with a matching
+ value.
+
+ Returns
+ -------
+ match : Expr or Term
+ """
+ one = self.try_find_one(func, **attrs)
+ if one is None:
+ raise RuntimeError("found no match")
+ return one
-BLACKLISTED_QUALNAMES: Final[set[str]] = set(keyword.kwlist) - {"None", "True", "False"}
+FORBIDDEN_KEYWORDS: Final[set[str]] = set(keyword.kwlist) - {"None", "True", "False"}
-class BlacklistedQualname(DocstubError):
+class QualnameIsForbidden(DocstubError):
"""Raised when a qualname is a forbidden keyword."""
@@ -311,6 +392,23 @@ class DoctypeTransformer(lark.visitors.Transformer):
'int | tuple[int, ...]'
"""
+ any_shape_expr: ClassVar[Expr] = Expr(
+ rule="shape",
+ children=(
+ Expr(
+ "subscription",
+ children=(
+ Term("tuple", kind="name"),
+ Term("[", kind="syntax"),
+ Term("Any", kind="name"),
+ Term(", ", kind="syntax"),
+ Term("...", kind="literal"),
+ Term("]", kind="syntax"),
+ ),
+ ),
+ ),
+ )
+
def __init__(self, *, reporter=None):
"""
Parameters
@@ -320,6 +418,26 @@ def __init__(self, *, reporter=None):
reporter = reporter or ContextReporter(logger=logger)
self.reporter = reporter.copy_with(logger=logger)
+ def __default__(self, data, children, meta):
+ """Create expression from rule which has no explicit method.
+
+ Parameters
+ ----------
+ data : str
+ Name of the rule/expression.
+ children : list[Expr | Term]
+ Children of the rule/expression.
+ meta : lark.tree.Meta
+ Meta data like position.
+
+ Returns
+ -------
+ expr : Expr
+ """
+ if "__" in data:
+ _prefix, data = data.split("__")
+ return Expr(rule=data, children=tuple(children))
+
def start(self, tree):
"""
Parameters
@@ -345,8 +463,8 @@ def qualname(self, tree):
children = tree.children
_qualname = ".".join(children)
- if _qualname in BLACKLISTED_QUALNAMES:
- raise BlacklistedQualname(_qualname)
+ if _qualname in FORBIDDEN_KEYWORDS:
+ raise QualnameIsForbidden(_qualname)
_qualname = Term(
_qualname,
@@ -408,7 +526,7 @@ def subscription(self, tree):
out : Expr
"""
assert len(tree.children) > 1
- return self._format_subscription(tree.children, rule="subscription")
+ return self._create_subscription_expr(tree.children, rule="subscription")
def param_spec(self, tree):
"""
@@ -440,7 +558,7 @@ def callable(self, tree):
out : Expr
"""
assert len(tree.children) > 1
- return self._format_subscription(tree.children, rule="callable")
+ return self._create_subscription_expr(tree.children, rule="callable")
def literal(self, tree):
"""
@@ -453,7 +571,7 @@ def literal(self, tree):
out : Expr
"""
assert len(tree.children) > 1
- out = self._format_subscription(tree.children, rule="literal")
+ out = self._create_subscription_expr(tree.children, rule="literal")
return out
def natlang_literal(self, tree):
@@ -466,11 +584,8 @@ def natlang_literal(self, tree):
-------
out : Expr
"""
- items = [
- Term("Literal", kind=TermKind.NAME),
- *tree.children,
- ]
- out = self._format_subscription(items, rule="natlang_literal")
+ items = [Term("Literal", kind=TermKind.NAME), *tree.children]
+ out = self._create_subscription_expr(items, rule="natlang_literal")
assert len(tree.children) >= 1
if len(tree.children) == 1:
@@ -512,7 +627,7 @@ def natlang_container(self, tree):
out : Expr
"""
assert len(tree.children) >= 1
- return self._format_subscription(tree.children, rule="natlang_container")
+ return self._create_subscription_expr(tree.children, rule="natlang_container")
def natlang_array(self, tree):
"""
@@ -524,7 +639,45 @@ def natlang_array(self, tree):
-------
out : Expr
"""
- return self._format_subscription(tree.children, rule="natlang_array")
+ match tree.children:
+ # Array without dtype or shape
+ case [Expr(rule="array_name") as array_name]:
+ children = (array_name,)
+
+ # Only dtype
+ case [Expr(rule="array_name") as array_name, Expr(rule="dtype") as dtype]:
+ # Insert missing shape before dtype
+ children = (array_name, self.any_shape_expr, dtype)
+
+ # Only shape
+ case [Expr(rule="array_name") as array_name, Expr(rule="shape") as shape]:
+ children = (array_name, shape)
+
+ # Shape and dtype in any order
+ case [
+ Expr(rule="array_name") as array_name,
+ Expr(rule="shape") as shape,
+ Expr(rule="dtype") as dtype,
+ ] | [
+ Expr(rule="array_name") as array_name,
+ Expr(rule="dtype") as dtype,
+ Expr(rule="shape") as shape,
+ ]:
+ children = (array_name, shape, dtype)
+
+ case _:
+ raise ValueError(f"{tree.children=} did not match a case")
+
+ if str(array_name) in ["array-like", "array_like"]:
+ # `numpy.ArrayLike` is a union and doesn't support subscription.
+ # So return expression without shape and dtype
+ return Expr(rule="natlang_array", children=(array_name,))
+
+ expr = self._create_subscription_expr(children, rule="natlang_array")
+
+ # Expand shapes with optional dimensions and dtype unions
+ expr = expand_expression(expr, expand_rule="to_expand")
+ return expr
def array_name(self, tree):
"""
@@ -537,9 +690,10 @@ def array_name(self, tree):
out : Term
"""
# This currently relies on a hack that only allows specific names
- # in `array_expression` (see `ARRAY_NAME` terminal in gramar)
+ # in `array_expression` (see `ARRAY_NAME` terminal in grammar)
qualname = self.qualname(tree)
- return qualname
+ array_name = Expr(rule="array_name", children=(qualname,))
+ return array_name
def dtype(self, tree):
"""
@@ -551,7 +705,22 @@ def dtype(self, tree):
-------
out : Expr
"""
- return Expr(rule="dtype", children=tuple(tree.children))
+ children = [Term("np.dtype", kind=TermKind.NAME), *tree.children]
+ expr = self._create_subscription_expr(children, rule="dtype")
+ return expr
+
+ def dtype_to_expand(self, tree):
+ """
+ Parameters
+ ----------
+ tree : lark.Tree
+
+ Returns
+ -------
+ out : Expr
+ """
+ # Remap to rule=to_expand
+ return Expr(rule="to_expand", children=tuple(tree.children))
def shape(self, tree):
"""
@@ -561,10 +730,52 @@ def shape(self, tree):
Returns
-------
- out : lark.visitors._DiscardType
+ out : Expr
"""
- logger.debug("dropping shape information")
- return lark.Discard
+ if len(tree.children) == 0:
+ # 0-D shape
+ children = self._create_subscription_expr(
+ [Term("tuple", kind=TermKind.NAME), Term("()", kind=TermKind.SYNTAX)],
+ rule="subscription",
+ )
+ shape = Expr(rule="shape", children=(children,))
+ return shape
+
+ expr = Expr(rule="shape", children=tuple(tree.children))
+
+ if expr.find_all(lambda x: x == "...", kind=TermKind.LITERAL):
+ # "..." effectively allows any shape
+ return self.any_shape_expr
+
+ # Count required and optional dimensions
+ required_count = len(expr.find_all(rule="dim"))
+ optional_count = 0
+ for optional in expr.find_all(rule="optional_dimensions"):
+ optional_count += len(optional.find_all(rule="dim"))
+ required_count -= optional_count
+
+ tuple_ = Term("tuple", kind=TermKind.NAME)
+ dim = Term("int", kind=TermKind.NAME)
+
+ # Construct a "shape tuple" for each of the possible dimension lengths
+ shape_tuples = (
+ (tuple_,) + (dim,) * length
+ for length in range(required_count, required_count + optional_count + 1)
+ )
+ shape_tuples = tuple(
+ self._create_subscription_expr(tuple_, rule="subscription")
+ for tuple_ in shape_tuples
+ )
+
+ if len(shape_tuples) > 1:
+ # Mark for expansion later
+ shape_tuples = Expr(rule="to_expand", children=shape_tuples)
+ elif len(shape_tuples) == 1:
+ # Otherwise unpack single tuple
+ shape_tuples = shape_tuples[0]
+
+ shape = Expr(rule="shape", children=(shape_tuples,))
+ return shape
def optional_info(self, tree):
"""
@@ -593,20 +804,23 @@ def extra_info(self, tree):
return lark.Discard
@staticmethod
- def _format_subscription(sequence, *, rule):
- """Format a `name[...]` style expression.
+ def _create_subscription_expr(children, *, rule):
+ """Create a `name[...]` style expression.
Parameters
----------
- sequence : Sequence[str]
+ children : Sequence[str]
+ Children for the created expression, including the name before the square
+ brackets.
rule : str, optional
+ The `rule` for the created expression.
Returns
-------
out : Expr
"""
sep = Term(", ", kind=TermKind.SYNTAX)
- container, *content = sequence
+ container, *content = children
assert container
if content:
@@ -652,8 +866,8 @@ class DoctypeTranspiler:
>>> transpile = DoctypeTranspiler()
>>> transpile("tuple of (int, ...)")
- >>> transpile("ndarray of dtype (float or int)")
-
+ >>> transpile("ndarray of dtype int and shape (M, N)")
+
"""
def __init__(
@@ -715,10 +929,15 @@ def __call__(self, doctype):
"""
if doctype in self._lru_cache:
self._stats.inc_counter("doctype_transpiler_cache_hits")
- return self._lru_cache[doctype]
+ expression = self._lru_cache[doctype]
+ return expression
+ # TODO Cache exceptions raised in this block?
tree = self._parser.parse(doctype)
- expression = self._transformer.transform(tree=tree)
+ try:
+ expression = self._transformer.transform(tree=tree)
+ except lark.visitors.VisitError as e:
+ raise e.orig_exc from None # Unwrap original error
assert isinstance(expression, Expr)
self._lru_cache[doctype] = expression
diff --git a/tests/test_doctype.py b/tests/test_doctype.py
index 3c84b7b..dd794da 100644
--- a/tests/test_doctype.py
+++ b/tests/test_doctype.py
@@ -1,16 +1,17 @@
import logging
-from itertools import product
import lark
import lark.exceptions
import pytest
from docstub.doctype import (
- BLACKLISTED_QUALNAMES,
+ FORBIDDEN_KEYWORDS,
DoctypeTransformer,
DoctypeTranspiler,
Expr,
+ QualnameIsForbidden,
Term,
+ expand_expression,
insert_between,
)
@@ -37,9 +38,11 @@ def test_edge_case_errors(self, doctype):
with pytest.raises(lark.exceptions.UnexpectedInput):
DoctypeTranspiler()(doctype)
- @pytest.mark.parametrize("doctype", BLACKLISTED_QUALNAMES)
- def test_reserved_keywords(self, doctype):
- with pytest.raises(lark.exceptions.VisitError):
+ @pytest.mark.parametrize("keyword", FORBIDDEN_KEYWORDS)
+ @pytest.mark.parametrize("template", ["{}", "Literal[{}]"])
+ def test_reserved_keywords(self, keyword, template):
+ doctype = template.format(keyword)
+ with pytest.raises(QualnameIsForbidden, match=keyword):
DoctypeTranspiler()(doctype)
@pytest.mark.parametrize(
@@ -52,31 +55,21 @@ def test_reserved_keywords(self, doctype):
def test_natlang_union(self, doctype, expected):
expr = DoctypeTranspiler()(doctype)
assert expr.as_code() == expected
+ assert expr.find_one(rule="union")
@pytest.mark.parametrize(
("doctype", "expected"),
[
- # Conventional
("list[float]", "list[float]"),
("dict[str, Union[int, str]]", "dict[str, Union[int, str]]"),
("tuple[int, ...]", "tuple[int, ...]"),
("Sequence[int | float]", "Sequence[int | float]"),
- # Natural language variant with "of" and optional plural "(s)"
- ("list of int", "list[int]"),
- ("list of int(s)", "list[int]"),
- # Natural tuple variant
- ("tuple of (float, int, str)", "tuple[float, int, str]"),
- ("tuple of (float, ...)", "tuple[float, ...]"),
- # Natural dict variant
- ("dict of {str: int}", "dict[str, int]"),
- ("dict of {str: int | float}", "dict[str, int | float]"),
- ("dict of {str: int or float}", "dict[str, int | float]"),
- ("dict[list of str]", "dict[list[str]]"),
],
)
def test_subscription(self, doctype, expected):
expr = DoctypeTranspiler()(doctype)
assert expr.as_code() == expected
+ assert expr.find_all(rule="subscription")
@pytest.mark.parametrize(
("doctype", "expected"),
@@ -85,7 +78,6 @@ def test_subscription(self, doctype, expected):
("list of int", "list[int]"),
("list of int(s)", "list[int]"),
("list of (int or float)", "list[int | float]"),
- ("list of (list of int)", "list[list[int]]"),
# Natural tuple variant
("tuple of (float, int, str)", "tuple[float, int, str]"),
("tuple of (float, ...)", "tuple[float, ...]"),
@@ -93,16 +85,27 @@ def test_subscription(self, doctype, expected):
("dict of {str: int}", "dict[str, int]"),
("dict of {str: int | float}", "dict[str, int | float]"),
("dict of {str: int or float}", "dict[str, int | float]"),
- # Nesting is possible but probably rarely a good idea
+ ],
+ )
+ def test_natlang_container(self, doctype, expected):
+ expr = DoctypeTranspiler()(doctype)
+ assert expr.as_code() == expected
+ assert expr.find_one(rule="natlang_container"), expr.format_tree()
+
+ @pytest.mark.parametrize(
+ ("doctype", "expected"),
+ [
+ ("list of (list of int)", "list[list[int]]"),
("list of (list of int(s))", "list[list[int]]"),
("tuple of (tuple of (float, ...), ...)", "tuple[tuple[float, ...], ...]"),
("dict of {str: dict of {str: float}}", "dict[str, dict[str, float]]"),
("dict of {str: list of (list of int(s))}", "dict[str, list[list[int]]]"),
],
)
- def test_natlang_container(self, doctype, expected):
+ def test_nested_natlang_container(self, doctype, expected):
expr = DoctypeTranspiler()(doctype)
assert expr.as_code() == expected
+ assert len(expr.find_all(rule="natlang_container")) > 1
@pytest.mark.parametrize(
"doctype",
@@ -124,29 +127,40 @@ def test_subscription_error(self, doctype):
DoctypeTranspiler()(doctype)
@pytest.mark.parametrize(
- ("doctype"),
+ "doctype",
[
"Literal[0]",
"Literal[-1, 1]",
- "Literal[None]",
- "Literal[True, False]",
+ "Literal[-2., 1.5]",
"""Literal['a', "bar"]""",
# Enum
- "Literal[SomeEnum.FIRST]",
"Literal[SomeEnum.FIRST, 1]",
"Literal[SomeEnum.FIRST, 2]",
"Literal[SomeEnum.FIRST, 3]",
# Nesting
"dict[Literal['a', 'b'], int]",
- # Custom qualname for literal
- "MyLiteral[0]",
- "MyLiteral[SomeEnum.FIRST]",
],
)
def test_literals(self, doctype):
expr = DoctypeTranspiler()(doctype)
assert expr.as_code() == doctype
- assert "literal" in [e.rule for e in expr.sub_expressions]
+ assert expr.find_one(rule="literal")
+
+ @pytest.mark.parametrize(
+ "doctype",
+ [
+ "Literal[None]",
+ "Literal[True, False]",
+ # Enum
+ "Literal[SomeEnum.FIRST]",
+ ],
+ )
+ def test_literal_only_qualnames(self, doctype):
+ """Parsed as a subscription because it only contains qualnames."""
+ expr = DoctypeTranspiler()(doctype)
+ assert expr.as_code() == doctype
+ assert expr.find_one(rule="subscription")
+ assert not expr.find_all(rule="literal")
@pytest.mark.parametrize(
("doctype", "expected"),
@@ -166,18 +180,18 @@ def test_literals(self, doctype):
# These aren't officially valid as an argument to `Literal` (yet)
# https://typing.python.org/en/latest/spec/literal.html
# TODO figure out how docstub should deal with these
- ("{-2., 1.}", "Literal[-2., 1.]"),
+ ("{-2., 1.5}", "Literal[-2., 1.5]"),
pytest.param(
"{-inf, inf, nan}",
- "Literal[, 1.]",
- marks=pytest.mark.xfail(reason="unsure how to support"),
+ "Literal[-inf, inf, nan]",
+ marks=pytest.mark.xfail(reason="unsure if to support"),
),
],
)
def test_natlang_literals(self, doctype, expected):
expr = DoctypeTranspiler()(doctype)
assert expr.as_code() == expected
- assert "natlang_literal" in [e.rule for e in expr.sub_expressions]
+ assert expr.find_one(rule="natlang_literal")
def test_single_natlang_literal_warning(self, caplog):
expr = DoctypeTranspiler()("{True}")
@@ -193,7 +207,6 @@ def test_single_natlang_literal_warning(self, caplog):
("doctype", "expected"),
[
("int", "int"),
- ("int | None", "int | None"),
("tuple of (int, float)", "tuple[int, float]"),
("{'a', 'b'}", "Literal['a', 'b']"),
],
@@ -215,6 +228,7 @@ def test_optional_info(self, doctype, expected, optional_info):
doctype_with_optional = doctype + optional_info
expr = DoctypeTranspiler()(doctype_with_optional)
assert expr.as_code() == expected
+ assert not expr.find_all(rule="optional_info")
@pytest.mark.parametrize(
"doctype",
@@ -230,7 +244,7 @@ def test_optional_info(self, doctype, expected, optional_info):
def test_callable(self, doctype):
expr = DoctypeTranspiler()(doctype)
assert expr.as_code() == doctype
- assert "callable" in [e.rule for e in expr.sub_expressions]
+ assert expr.find_one(rule="callable")
@pytest.mark.parametrize(
"doctype",
@@ -243,7 +257,7 @@ def test_callable(self, doctype):
def test_callable_subscriptions_form(self, doctype):
expr = DoctypeTranspiler()(doctype)
assert expr.as_code() == doctype
- assert "callable" not in [e.rule for e in expr.sub_expressions]
+ assert not expr.find_all(rule="callable")
@pytest.mark.parametrize(
"doctype",
@@ -271,74 +285,136 @@ def test_callable_error(self, doctype):
def test_rst_role(self, doctype, expected):
expr = DoctypeTranspiler()(doctype)
assert expr.as_code() == expected
+ assert expr.find_one(rule="rst_role")
- # fmt: off
+
+class Test_DoctypeTranspiler_arrays:
@pytest.mark.parametrize(
- ("fmt", "expected_fmt"),
+ ("in_shape", "out_shape"),
[
- ("{name} of shape {shape} and dtype {dtype}", "{name}[{dtype}]"),
- ("{name} of dtype {dtype} and shape {shape}", "{name}[{dtype}]"),
+ # Exact dimensionality
+ ("()", "tuple[()]"), # Special case: 0-D
+ ("(M,)", "tuple[int]"),
+ ("(M=3,)", "tuple[int]"),
+ ("(2,)", "tuple[int]"),
+ ("(M, N)", "tuple[int, int]"),
+ ("(M, N, O)", "tuple[int, int, int]"),
+ ("(0, 1, 2)", "tuple[int, int, int]"),
+ # Ellipsis expands to unconstraint shape
+ ("(M, ...)", "tuple[Any, ...]"),
+ ("(0, ...)", "tuple[Any, ...]"),
+ ("(..., M)", "tuple[Any, ...]"),
+ ("(M, ..., N)", "tuple[Any, ...]"),
+ ("(M[, ...], N)", "tuple[Any, ...]"),
],
)
- @pytest.mark.parametrize("name", ["array", "ndarray", "array-like", "array_like"])
- @pytest.mark.parametrize("dtype", ["int", "np.int8"])
- @pytest.mark.parametrize("shape",
- ["(2, 3)", "(N, m)", "3D", "2-D", "(N, ...)", "([P,] M, N)"]
- )
- 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)
+ def test_explicit_shape(self, in_shape, out_shape):
+ """Cases where the shape describes the exact dimensionality"""
+ doctype = f"ndarray of shape {in_shape}"
+ expected = f"ndarray[{out_shape}]"
expr = DoctypeTranspiler()(doctype)
assert expr.as_code() == expected
- assert "natlang_array" in [e.rule for e in expr.sub_expressions]
- # fmt: on
@pytest.mark.parametrize(
("doctype", "expected"),
[
- ("ndarray of dtype (int or float)", "ndarray[int | float]"),
- ("ndarray of shape (M, N)", "ndarray"),
+ (
+ "ndarray of shape (M[, N])",
+ "ndarray[tuple[int]] | ndarray[tuple[int, int]]",
+ ),
+ (
+ "ndarray of shape ([M, ]N)",
+ "ndarray[tuple[int]] | ndarray[tuple[int, int]]",
+ ),
+ (
+ "ndarray of shape ([N,] M [,O])",
+ "ndarray[tuple[int]]"
+ " | ndarray[tuple[int, int]]"
+ " | ndarray[tuple[int, int, int]]",
+ ),
+ (
+ "ndarray of shape (M[, N, O], P)",
+ "ndarray[tuple[int, int]]"
+ " | ndarray[tuple[int, int, int]]"
+ " | ndarray[tuple[int, int, int, int]]",
+ ),
],
)
- def test_natlang_array_specific(self, doctype, expected):
+ def test_optional_shape(self, doctype, expected):
+ """Optional dimensions are expanded into union of matching arrays."""
expr = DoctypeTranspiler()(doctype)
assert expr.as_code() == expected
- assert "natlang_array" in [e.rule for e in expr.sub_expressions]
- @pytest.mark.parametrize("shape", ["(-1, 3)", "(1.0, 2)", "-3D", "-2-D"])
- def test_natlang_array_invalid_shape(self, shape):
- doctype = f"array of shape {shape}"
- with pytest.raises(lark.exceptions.UnexpectedInput):
- _ = DoctypeTranspiler()(doctype)
+ @pytest.mark.parametrize("shape", ["(N)", "(2)", "(-1, 3)", "(1.0, 2)"])
+ def test_malformed_shape(self, shape):
+ doctype = f"ndarray of shape {shape}"
+ with pytest.raises(lark.LexError, match="No terminal matches"):
+ DoctypeTranspiler()(doctype)
+ @pytest.mark.parametrize(
+ ("doctype", "expected"),
+ [
+ ("ndarray of dtype int", "ndarray[tuple[Any, ...], np.dtype[int]]"),
+ (
+ "ndarray of dtype (int or float)",
+ "ndarray[tuple[Any, ...], np.dtype[int]]"
+ " | ndarray[tuple[Any, ...], np.dtype[float]]",
+ ),
+ (
+ "ndarray of dtype (int | float | complex)",
+ "ndarray[tuple[Any, ...], np.dtype[int]]"
+ " | ndarray[tuple[Any, ...], np.dtype[float]]"
+ " | ndarray[tuple[Any, ...], np.dtype[complex]]",
+ ),
+ ],
+ )
+ def test_dtype(self, doctype, expected):
+ expr = DoctypeTranspiler()(doctype)
+ assert expr.as_code() == expected
-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]"),
+ (
+ "ndarray of dtype int and shape (M, N)",
+ "ndarray[tuple[int, int], np.dtype[int]]",
+ ),
+ (
+ "ndarray of dtype (int or float) and shape ([N,] M)",
+ "ndarray[tuple[int], np.dtype[int]]"
+ " | ndarray[tuple[int], np.dtype[float]]"
+ " | ndarray[tuple[int, int], np.dtype[int]]"
+ " | ndarray[tuple[int, int], np.dtype[float]]",
+ ),
],
)
- 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)
+ def test_shape_and_dtype(self, doctype, expected):
+ expr = DoctypeTranspiler()(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(
+ "fmt",
+ [
+ "{name} of shape {shape} and dtype {dtype}",
+ "{name} of dtype {dtype} and shape {shape}",
+ ],
+ )
+ @pytest.mark.parametrize("name", ["array-like", "array_like"])
+ @pytest.mark.parametrize("dtype", ["int", "np.int8"])
+ @pytest.mark.parametrize("shape", ["(2, 3)", "(N, m)", "(N, ...)", "([P,] M, N)"])
+ def test_natlang_array_like(self, fmt, name, dtype, shape):
+ doctype = fmt.format(name=name, dtype=dtype, shape=shape)
+ expr = DoctypeTranspiler()(doctype)
+ assert expr.as_code() == name
+ assert expr.find_one(rule="natlang_array")
+
+
+class Test_DoctypeTranspiler_override:
@pytest.mark.parametrize(
("doctype", "expected"),
[
- ("my_array of dtype int", "my_array[int]"),
- ("my_array of shape (4, 4)", "my_array"),
+ ("my_array of dtype int", "my_array[tuple[Any, ...], np.dtype[int]]"),
+ ("my_array of shape (4,)", "my_array[tuple[int]]"),
],
)
def test_doc_example_new_array_name(self, doctype, expected):
@@ -350,7 +426,7 @@ def test_doc_example_new_array_name(self, doctype, expected):
transpile = DoctypeTranspiler(grammar=grammar)
expr = transpile(doctype)
assert expr.as_code() == expected
- assert len(list(expr.filter(rule="natlang_array"))) == 1
+ assert expr.find_one(rule="natlang_array")
@pytest.mark.parametrize(
("doctype", "expected"),
@@ -371,20 +447,20 @@ def union(self, tree):
transpile = DoctypeTranspiler(transformer=Transformer())
expr = transpile(doctype)
assert expr.as_code() == expected
- assert "union" in [e.rule for e in expr.sub_expressions]
+ assert expr.find_all(rule="union")
@pytest.mark.parametrize(
("doctype", "expected"),
[
- ("(int) -> bool", "Callable[[int], bool]"),
- ("(str, bytes) -> int or bool", "Callable[[str, bytes], int | bool]"),
+ ("f(int) -> bool", "Callable[[int], bool]"),
+ ("f(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
+ natlang_callable : "f(" params ") ->" _type
%extend _type: natlang_callable
%ignore WS_INLINE
"""
@@ -406,7 +482,7 @@ def params(self, children):
def natlang_callable(self, children):
items = [Term("Callable", kind="name"), *children]
- out = DoctypeTransformer._format_subscription(
+ out = DoctypeTransformer._create_subscription_expr(
items, rule="natlang_callable"
)
return out
@@ -414,8 +490,8 @@ def natlang_callable(self, children):
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]
+ assert expr.find_one(rule="natlang_callable")
+ assert expr.find_all(rule="params")
@pytest.mark.parametrize(
("doctype", "expected"),
@@ -447,63 +523,14 @@ def start(self, tree):
-------
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
+ assert len(tree.children) == 1
+ assert isinstance(tree.children[0], Expr)
+ child = expand_expression(tree.children[0], expand_rule="expansion")
+ return Expr(rule="start", children=(child,))
transpile = DoctypeTranspiler(
grammar=grammar, transformer=ExpansionTransformer()
)
expr = transpile(doctype)
assert expr.as_code() == expected
- assert list(expr.filter(rule="natlang_array")) == []
+ assert not expr.find_all(rule="expansion")