Allow signed numbers in literals

src/docstub/doctype.lark

@@ -58,12 +58,12 @@ natlang_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 | NUMBER | qualname
+?literal_item: ELLIPSES | STRING | SIGNED_NUMBER | qualname
 
 
 // Natural language forms of the subscription expression for containers.
-// These forms allow nesting allow nesting in and with other expressions. But
-// it's discouraged to do so extensively to maintain readability.
+// 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 ")"
     | _natlang_tuple
@@ -114,7 +114,7 @@ ARRAY_NAME: "array" | "ndarray" | "array-like" | "array_like"
 // NumPy arrays, this information is dropped during the transformation.
 shape: "(" dim ",)"
     | "(" leading_optional_dim? dim (("," dim | insert_optional_dim))* ")"
-    | NUMBER "-"? "D"
+    | INT "-"? "D"
 
 
 // Optional dimensions in a `shape` expression placed at the start,
@@ -129,7 +129,7 @@ shape: "(" 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: NUMBER | ELLIPSES | NAME ("=" NUMBER)?
+?dim: INT | ELLIPSES | NAME ("=" INT)?
 
 
 // Optional information about a parameter has a default value, added after the
@@ -147,6 +147,6 @@ NAME: /[^\W\d][\w-]*/
 
 
 %import python (STRING)
-%import common (NUMBER, WS_INLINE)
+%import common (SIGNED_NUMBER, INT, WS_INLINE)
 
 %ignore WS_INLINE

tests/test_docstrings.py

@@ -108,12 +108,26 @@ def test_subscription_error(self, doctype):
         ("doctype", "expected"),
         [
             ("{0}", "Literal[0]"),
-            ("{'a', 1, None, False}", "Literal['a', 1, None, False]"),
-            ("dict[{'a', 'b'}, int]", "dict[Literal['a', 'b'], int]"),
+            ("{-1, 1}", "Literal[-1, 1]"),
+            ("{None}", "Literal[None]"),
+            ("{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]"),
+            # Nesting
+            ("dict[{'a', 'b'}, int]", "dict[Literal['a', 'b'], int]"),
+            # 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.]"),
+            pytest.param(
+                "{-inf, inf, nan}",
+                "Literal[, 1.]",
+                marks=pytest.mark.xfail(reason="unsure how to support"),
+            ),
         ],
     )
     def test_literals(self, doctype, expected):
@@ -188,6 +202,13 @@ def escape(name: str) -> str:
         assert annotation.value == expected
     # fmt: on
 
+    @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}"
+        transformer = DoctypeTransformer()
+        with pytest.raises(lark.exceptions.UnexpectedInput):
+            transformer.doctype_to_annotation(doctype)
+
     def test_unknown_name(self):
         # Simple unknown name is aliased to typing.Any
         transformer = DoctypeTransformer()