WIP

Author: lagru

src/docstub-stubs/_app_generate_stubs.pyi

@@ -6,6 +6,7 @@ from collections import Counter
 from collections.abc import Iterable, Sequence
 from contextlib import contextmanager
 from pathlib import Path
+from typing import Literal
 
 from ._analysis import PyImport, TypeCollector, TypeMatcher, common_known_types
 from ._cache import CACHE_DIR_NAME, FileCache

src/docstub-stubs/_cli.pyi

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

src/docstub-stubs/_docstrings.pyi

@@ -10,6 +10,7 @@ import click
 import lark
 import lark.visitors
 import numpydoc.docscrape as npds
+from _typeshed import Incomplete as Expr
 
 from ._analysis import PyImport, TypeMatcher
 from ._doctype import BlacklistedQualname, Expr, Term, TermKind, parse_doctype

src/docstub-stubs/_doctype.pyi

@@ -7,11 +7,12 @@ from collections.abc import Generator, Iterable, Sequence
 from dataclasses import dataclass
 from pathlib import Path
 from textwrap import indent
-from typing import Any, Final
+from typing import Any, Final, Literal, Self
 
 import lark
 import lark.visitors
 from _typeshed import Incomplete
+from _typeshed import Incomplete as Expression
 
 from ._utils import DocstubError
 
@@ -47,14 +48,17 @@ class Term(str):
 class Expr:
 
     rule: str
-    children: list[Expr | Term]
+    children: list[Self | Term]
 
     @property
     def terms(self) -> list[Term]: ...
     @property
     def names(self) -> list[Term]: ...
+    @property
+    def sub_expressions(self) -> list[Expression] | Literal[1]: ...
     def __iter__(self) -> Generator[Expr | Term]: ...
     def format_tree(self) -> str: ...
+    def print_tree(self) -> None: ...
     def __repr__(self) -> str: ...
     def __str__(self) -> str: ...
     def as_code(self) -> str: ...
@@ -67,9 +71,14 @@ class BlacklistedQualname(DocstubError):
 class DoctypeTransformer(lark.visitors.Transformer):
     def start(self, tree: lark.Tree) -> Expr: ...
     def qualname(self, tree: lark.Tree) -> Term: ...
+    def qualname(self, tree: lark.Tree) -> Term: ...
+    def rst_role(self, tree: lark.Tree) -> Expr: ...
     def ELLIPSES(self, token: lark.Token) -> Term: ...
     def union(self, tree: lark.Tree) -> Expr: ...
     def subscription(self, tree: lark.Tree) -> Expr: ...
+    def param_spec(self, tree: lark.Tree) -> Expr: ...
+    def callable(self, tree: lark.Tree) -> Expr: ...
+    def literal(self, tree: lark.Tree) -> Expr: ...
     def natlang_literal(self, tree: lark.Tree) -> Expr: ...
     def literal_item(self, tree: lark.Tree) -> Term: ...
     def natlang_container(self, tree: lark.Tree) -> Expr: ...
@@ -79,9 +88,7 @@ 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 _format_subscription(self, sequence: Sequence[str], *, rule: str) -> Expr: ...
 
 _transformer: Final
 

src/docstub-stubs/_report.pyi

@@ -5,7 +5,7 @@ import logging
 from collections.abc import Hashable, Iterator, Mapping, Sequence
 from pathlib import Path
 from textwrap import indent
-from typing import Any, ClassVar, Self, TextIO
+from typing import Any, ClassVar, Literal, Self, TextIO
 
 import click
 from pre_commit.envcontext import UNSET

src/docstub/_doctype.py

@@ -146,6 +146,20 @@ def names(self):
         """
         return [term for term in self.terms if term.kind == TermKind.NAME]
 
+    @property
+    def sub_expressions(self):
+        """Iterate expressions inside the current one.
+
+        Returns
+        -------
+        names : list of Expr or {1}
+        """
+        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.
 
@@ -253,6 +267,21 @@ def qualname(self, tree):
         )
         return _qualname
 
+    def rst_role(self, tree):
+        """
+        Parameters
+        ----------
+        tree : lark.Tree
+
+        Returns
+        -------
+        out : Expr
+        """
+        # Drop rst_prefix
+        children = [c for c in tree.children if isinstance(c, Term)]
+        expr = Expr(rule="rst_role", children=children)
+        return expr
+
     def ELLIPSES(self, token):
         """
         Parameters
@@ -332,11 +361,7 @@ def literal(self, tree):
         -------
         out : Expr
         """
-        items = [
-            Term("Literal", kind=TermKind.NAME),
-            *tree.children,
-        ]
-        out = self._format_subscription(items, rule="literal")
+        out = self._format_subscription(tree.children, rule="literal")
         return out
 
     def natlang_literal(self, tree):

src/docstub/doctype.lark

@@ -34,7 +34,7 @@ _type: qualname
 // [1] https://docutils.sourceforge.io/docs/ref/rst/roles.html
 //
 qualname: (/~/ ".")? (NAME ".")* NAME
-    | (":" (NAME ":")? NAME ":")? "`" qualname "`"
+    | (":" (NAME ":")? NAME ":")? "`" qualname "`" -> rst_role
 
 
 // An union of different types, joined either by "or" or "|".
@@ -73,7 +73,7 @@ natlang_literal: "{" literal_item ("," literal_item)* "}"
 // 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: "Literal[" literal_item ("," literal_item)* "]"
+literal.1: qualname "[" literal_item ("," literal_item)* "]"
 
 
 // An single item in a literal expression (or `optional`). We must also allow

tests/test_docstrings.py

@@ -5,6 +5,7 @@
 from docstub._analysis import PyImport
 from docstub._docstrings import (
     Annotation,
+    doctype_to_annotation,
     DocstringAnnotations,
 )
 
@@ -35,6 +36,37 @@ def test_unexpected_value(self):
             Annotation(value="~.foo")
 
 
+class Test_doctype_to_annotation:
+
+    def test_unknown_name(self):
+        # Simple unknown name is aliased to typing.Any
+        annotation = doctype_to_annotation("a")
+        assert annotation.value == "a"
+        assert annotation.imports == {
+            PyImport(import_="Incomplete", from_="_typeshed", as_="a")
+        }
+        assert unknown_names == [("a", 0, 1)]
+
+    def test_unknown_qualname(self):
+        # Unknown qualified name is escaped and aliased to typing.Any as well
+        annotation = doctype_to_annotation("a.b")
+        assert annotation.value == "a_b"
+        assert annotation.imports == {
+            PyImport(import_="Incomplete", from_="_typeshed", as_="a_b")
+        }
+        assert unknown_names == [("a.b", 0, 3)]
+
+    def test_multiple_unknown_names(self):
+        # Multiple names are aliased to typing.Any
+        annotation = doctype_to_annotation("a.b of c")
+        assert annotation.value == "a_b[c]"
+        assert annotation.imports == {
+            PyImport(import_="Incomplete", from_="_typeshed", as_="a_b"),
+            PyImport(import_="Incomplete", from_="_typeshed", as_="c"),
+        }
+        assert unknown_names == [("a.b", 0, 3), ("c", 7, 8)]
+
+
 class Test_DocstringAnnotations:
     def test_empty_docstring(self):
         docstring = dedent("""No sections in this docstring.""")
@@ -247,4 +279,4 @@ def test_combined_numpydoc_params(self):
         assert annotations.parameters["c"].value == "bool"
 
         assert "d" not in annotations.parameters
-        assert "e" not in annotations.parameters
+        assert "e" not in annotations.parameters

tests/test_doctype.py

@@ -130,11 +130,15 @@ def test_subscription_error(self, doctype):
             "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 = parse_doctype(doctype)
         assert expr.as_code() == doctype
+        assert "literal" in [e.rule for e in expr.sub_expressions]
 
     @pytest.mark.parametrize(
         ("doctype", "expected"),
@@ -165,6 +169,7 @@ def test_literals(self, doctype):
     def test_natlang_literals(self, doctype, expected):
         expr = parse_doctype(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}")
@@ -220,6 +225,7 @@ def test_optional_info(self, doctype, expected, optional_info):
     def test_callable(self, doctype):
         expr = parse_doctype(doctype)
         assert expr.as_code() == doctype
+        assert "callable" in [e.rule for e in expr.sub_expressions]
 
     @pytest.mark.parametrize(
         "doctype",
@@ -240,7 +246,7 @@ def test_callable_error(self, doctype):
             (":class:`Generator`", "Generator"),
             (":py:class:`Generator`", "Generator"),
             (":py:class:`Generator`[int]", "Generator[int]"),
-            (":py:ref:`~.Foo`[int]", "_Foo[int]"),
+            (":py:ref:`~.Foo`[int]", "~.Foo[int]"),
             ("list[:py:class:`Generator`]", "list[Generator]"),
         ],
     )
@@ -284,31 +290,3 @@ def test_natlang_array_invalid_shape(self, shape):
         doctype = f"array of shape {shape}"
         with pytest.raises(lark.exceptions.UnexpectedInput):
             _ = parse_doctype(doctype)
-
-    def test_unknown_name(self):
-        # Simple unknown name is aliased to typing.Any
-        annotation, unknown_names = parse_doctype("a")
-        assert annotation.value == "a"
-        assert annotation.imports == {
-            PyImport(import_="Incomplete", from_="_typeshed", as_="a")
-        }
-        assert unknown_names == [("a", 0, 1)]
-
-    def test_unknown_qualname(self):
-        # Unknown qualified name is escaped and aliased to typing.Any as well
-        annotation, unknown_names = parse_doctype("a.b")
-        assert annotation.value == "a_b"
-        assert annotation.imports == {
-            PyImport(import_="Incomplete", from_="_typeshed", as_="a_b")
-        }
-        assert unknown_names == [("a.b", 0, 3)]
-
-    def test_multiple_unknown_names(self):
-        # Multiple names are aliased to typing.Any
-        annotation, unknown_names = parse_doctype("a.b of c")
-        assert annotation.value == "a_b[c]"
-        assert annotation.imports == {
-            PyImport(import_="Incomplete", from_="_typeshed", as_="a_b"),
-            PyImport(import_="Incomplete", from_="_typeshed", as_="c"),
-        }
-        assert unknown_names == [("a.b", 0, 3), ("c", 7, 8)]