Support nesting of natlang containers

src/docstub/doctype.lark

@@ -65,7 +65,7 @@ natlang_literal: "{" literal_item ("," literal_item)* "}"
 // These forms allow nesting with other expressions. But it's discouraged to do
 // so extensively to maintain readability.
 natlang_container: qualname "of" qualname _PLURAL_S?
-    | qualname "of" "(" union ")"
+    | qualname "of" "(" type ")"
     | _natlang_tuple
     | _natlang_mapping
 
@@ -83,7 +83,7 @@ _natlang_tuple: qualname "of" "(" type "," ELLIPSES ")"
 
 
 // Natural language container variant for mappings.
-_natlang_mapping: qualname "of" "{" type ":" (type | union) "}"
+_natlang_mapping: qualname "of" "{" type ":" type "}"
 
 
 // A natural language alternative to describe arrays with a dtype and shape.

tests/test_docstrings.py

@@ -34,6 +34,48 @@ def test_unexpected_value(self):
 
 
 class Test_DoctypeTransformer:
+    @pytest.mark.parametrize(
+        "doctype",
+        [
+            "((float))",
+            "(float,)",
+            "(, )",
+            "...",
+            "(..., ...)",
+            "{}",
+            "{:}",
+            "{a:}",
+            "{:b}",
+            "{'a',}",
+            "a or (b or c)",
+            ",, optional",
+        ],
+    )
+    def test_edge_case_errors(self, doctype):
+        transformer = DoctypeTransformer()
+        with pytest.raises(lark.exceptions.UnexpectedInput):
+            transformer.doctype_to_annotation(doctype)
+
+    @pytest.mark.parametrize("doctype", DoctypeTransformer.blacklisted_qualnames)
+    def test_reserved_keywords(self, doctype):
+        assert DoctypeTransformer.blacklisted_qualnames
+
+        transformer = DoctypeTransformer()
+        with pytest.raises(lark.exceptions.VisitError):
+            transformer.doctype_to_annotation(doctype)
+
+    @pytest.mark.parametrize(
+        ("doctype", "expected"),
+        [
+            ("int or float", "int | float"),
+            ("int or float or str", "int | float | str"),
+        ],
+    )
+    def test_natlang_union(self, doctype, expected):
+        transformer = DoctypeTransformer()
+        annotation, _ = transformer.doctype_to_annotation(doctype)
+        assert annotation.value == expected
+
     @pytest.mark.parametrize(
         ("doctype", "expected"),
         [
@@ -67,14 +109,19 @@ 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, ...]"),
             # 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]]"),
+            # Nesting is possible but probably rarely a good idea
+            ("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):
@@ -86,7 +133,7 @@ def test_natlang_container(self, doctype, expected):
         "doctype",
         [
             "list of int (s)",
-            "list of (float)",
+            "list of ((float))",
             "list of (float,)",
             "list of (, )",
             "list of ...",