WIP

Author: lagru

src/docstub/_doctype.py

@@ -7,7 +7,7 @@
 from dataclasses import dataclass
 from pathlib import Path
 from textwrap import indent
-from typing import Final
+from typing import Final, Self
 
 import lark
 import lark.visitors
@@ -124,7 +124,7 @@ class Expr:
     """
 
     rule: str
-    children: list[Expr | Term]
+    children: list[Self | Term]
 
     @property
     def terms(self):
@@ -226,6 +226,29 @@ def qualname(self, tree):
         )
         return _qualname
 
+    def qualname(self, tree):
+        """
+        Parameters
+        ----------
+        tree : lark.Tree
+
+        Returns
+        -------
+        out : Term
+        """
+        children = tree.children
+        _qualname = ".".join(children)
+
+        if _qualname in BLACKLISTED_QUALNAMES:
+            raise BlacklistedQualname(_qualname)
+
+        _qualname = Term(
+            _qualname,
+            kind=TermKind.NAME,
+            pos=(tree.meta.start_pos, tree.meta.end_pos),
+        )
+        return _qualname
+
     def ELLIPSES(self, token):
         """
         Parameters
@@ -262,7 +285,7 @@ def subscription(self, tree):
         -------
         out : Expr
         """
-        return self._format_subscription(tree.children)
+        return self._format_subscription(tree.children, rule="subscription")
 
     def param_spec(self, tree):
         """
@@ -295,6 +318,23 @@ def callable(self, tree):
         """
         return self._format_subscription(tree.children, rule="callable")
 
+    def literal(self, tree):
+        """
+        Parameters
+        ----------
+        tree : lark.Tree
+
+        Returns
+        -------
+        out : Expr
+        """
+        items = [
+            Term("Literal", kind=TermKind.NAME),
+            *tree.children,
+        ]
+        out = self._format_subscription(items, rule="literal")
+        return out
+
     def natlang_literal(self, tree):
         """
         Parameters
@@ -306,7 +346,7 @@ def natlang_literal(self, tree):
         out : Expr
         """
         items = [
-            Term("Literal", kind=TermKind.SYNTAX),
+            Term("Literal", kind=TermKind.NAME),
             *tree.children,
         ]
         out = self._format_subscription(items, rule="natlang_literal")
@@ -335,7 +375,8 @@ def literal_item(self, tree):
         kind = TermKind.LITERAL
         if isinstance(item, Term):
             kind = item.kind
-        return Term(item, kind=kind, pos=(tree.meta.start_pos, tree.meta.end_pos))
+        out = Term(item, kind=kind, pos=(tree.meta.start_pos, tree.meta.end_pos))
+        return out
 
     def natlang_container(self, tree):
         """
@@ -427,7 +468,7 @@ def extra_info(self, tree):
         logger.debug("dropping extra info")
         return lark.Discard
 
-    def _format_subscription(self, sequence, rule="subscription"):
+    def _format_subscription(self, sequence, *, rule):
         """
         Parameters
         ----------

src/docstub/doctype.lark

@@ -19,6 +19,7 @@ _annotation_with_meta: _type ("," optional_info)?
 // nested construct of types.
 _type: qualname
     | union
+    | literal
     | subscription
     | callable
     | natlang_literal
@@ -68,6 +69,13 @@ ELLIPSES: "..."
 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: "Literal[" literal_item ("," literal_item)* "]"
+
+
 // An single item in a literal expression (or `optional`). We must also allow
 // for qualified names, since a "class" or enum can be used as a literal too.
 literal_item: ELLIPSES | STRING | SIGNED_NUMBER | qualname

tests/test_doctype.py

@@ -115,6 +115,27 @@ def test_subscription_error(self, doctype):
         with pytest.raises(lark.exceptions.UnexpectedInput):
             parse_doctype(doctype)
 
+    @pytest.mark.parametrize(
+        ("doctype"),
+        [
+            "Literal[0]",
+            "Literal[-1, 1]",
+            "Literal[None]",
+            "Literal[True, False]",
+            """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]",
+        ],
+    )
+    def test_literals(self, doctype):
+        expr = parse_doctype(doctype)
+        assert expr.as_code() == doctype
+
     @pytest.mark.parametrize(
         ("doctype", "expected"),
         [
@@ -124,10 +145,10 @@ def test_subscription_error(self, doctype):
             ("{True, False}", "Literal[True, False]"),
             ("""{'a', "bar"}""", """Literal['a', "bar"]"""),
             # Enum
-            ("{SomeEnum.FIRST}", "Literal[SomeEnum_FIRST]"),
-            ("{`SomeEnum.FIRST`, 1}", "Literal[SomeEnum_FIRST, 1]"),
-            ("{:ref:`SomeEnum.FIRST`, 2}", "Literal[SomeEnum_FIRST, 2]"),
-            ("{:py:ref:`SomeEnum.FIRST`, 3}", "Literal[SomeEnum_FIRST, 3]"),
+            ("{SomeEnum.FIRST}", "Literal[SomeEnum.FIRST]"),
+            ("{`SomeEnum.FIRST`, 1}", "Literal[SomeEnum.FIRST, 1]"),
+            ("{:ref:`SomeEnum.FIRST`, 2}", "Literal[SomeEnum.FIRST, 2]"),
+            ("{:py:ref:`SomeEnum.FIRST`, 3}", "Literal[SomeEnum.FIRST, 3]"),
             # Nesting
             ("dict[{'a', 'b'}, int]", "dict[Literal['a', 'b'], int]"),
             # These aren't officially valid as an argument to `Literal` (yet)
@@ -141,7 +162,7 @@ def test_subscription_error(self, doctype):
             ),
         ],
     )
-    def test_literals(self, doctype, expected):
+    def test_natlang_literals(self, doctype, expected):
         expr = parse_doctype(doctype)
         assert expr.as_code() == expected