Split docname into KnownImport and replace

Reduce responsibility of the former DocName class. Replacing docstring
specific type description should be handled separately.

Author: lagru

src/docstub/_docstrings.py

@@ -266,19 +266,21 @@ def contains(self, tree):
         return out
 
     def literals(self, tree):
-        out = " , ".join(tree.children)
+        out = ", ".join(tree.children)
         out = f"Literal[{out}]"
         _, known_import = self.inspector.query("Literal")
-        self._collected_imports.add(known_import)
+        if known_import:
+            self._collected_imports.add(known_import)
         return out
 
     def _find_import(self, qualname):
         """Match type names to known imports."""
         try:
-            qualname, known_import = self.inspector.query(qualname)
-            if known_import:
-                if known_import.has_import:
-                    self._collected_imports.add(known_import)
+            annotation_name, known_import = self.inspector.query(qualname)
+            if known_import and known_import.has_import:
+                self._collected_imports.add(known_import)
+            if annotation_name:
+                qualname = annotation_name
             else:
                 logger.warning(
                     "unknown import for %r in %s",

src/docstub/doctype.lark

@@ -1,15 +1,14 @@
 ?start : annotation
 
-annotation : (literals | types_or) ("," optional)? ("," extra_info)?
+annotation : types_or ("," optional)? ("," extra_info)?
 
-literals : "{" literal ("," literal)* "}"
-
-types_or : type (("or" | "|") type)*
+?types_or : type (("or" | "|") type)*
 
 ?type : qualname
       | sphinx_ref
       | container
       | shape_n_dtype
+      | literals
 
 optional : "optional"
          | "default" ("=" | ":")? literal
@@ -18,10 +17,9 @@ extra_info : /[^\r\n]+/
 
 sphinx_ref : ":" (NAME ":")? NAME ":`" qualname "`"
 
-container: qualname "[" types_or ("," types_or)* "]"
-         | qualname "[" types_or "," PY_ELLIPSES "]"
+container: qualname "[" types_or ("," types_or)* ("," PY_ELLIPSES)? "]"
          | qualname "of" type
-         | qualname "of" "(" types_or ("," types_or)* ")"
+         | qualname "of" "(" types_or ("," types_or)* ("," PY_ELLIPSES)? ")"
          | qualname "of" "{" types_or ":" types_or "}"
 
 // Name with leading dot separated path
@@ -42,14 +40,16 @@ dtype : qualname
 shape : "(" dim ",)"
       | "(" leading_optional? dim (("," dim | insert_optional))* ")"
       | NUMBER "-"? "D"
-leading_optional : "[" dim ("," dim)* ",]" -> optional
-insert_optional : "[," dim ("," dim)* "]" -> optional
+leading_optional : "[" dim ("," dim)* ",]"
+insert_optional : "[," dim ("," dim)* "]"
 ?dim : NUMBER
      | PY_ELLIPSES
      | NAME
 
 
 // Python
+literals : "{" literal ("," literal)* "}"
+
 literal : PY_ELLIPSES
         | STRING
         | NUMBER

tests/test_dev.py

@@ -0,0 +1,99 @@
+import pytest
+
+from docstub._analysis import KnownImport, StaticInspector, common_known_imports
+from docstub._docstrings import DoctypeTransformer, _lark
+
+
+@pytest.fixture()
+def transformer():
+    inspector = StaticInspector(known_imports=common_known_imports())
+    transformer = DoctypeTransformer(inspector=inspector, replace_doctypes={})
+    return transformer
+
+
+class Test_DoctypeTransformer:
+    # fmt: off
+    @pytest.mark.parametrize(
+        ("raw", "expected"),
+        [
+            ("list[float]", "list[float]"),
+            ("dict[str, Union[int, str]]", "dict[str, Union[int, str]]"),
+            ("tuple[int, ...]", "tuple[int, ...]"),
+
+            ("list of int", "list[int]"),
+            ("tuple of float", "tuple[float]"),
+            ("tuple of (float, ...)", "tuple[float, ...]"),
+
+            ("Sequence[int | float]", "Sequence[int | float]"),
+
+            ("dict of {str: int}", "dict[str, int]"),
+        ],
+    )
+    def test_container(self, raw, expected, transformer):
+        tree = _lark.parse(raw)
+        annotation = transformer.transform(tree)
+
+        assert annotation.value == expected
+    # fmt: on
+
+    @pytest.mark.parametrize(
+        ("raw", "expected"),
+        [
+            ("{'a', 1, None, False}", "Literal['a', 1, None, False]"),
+            ("dict[{'a', 'b'}, int]", "dict[Literal['a', 'b'], int]"),
+        ],
+    )
+    def test_literals(self, raw, expected, transformer):
+        tree = _lark.parse(raw)
+        annotation = transformer.transform(tree)
+
+        assert annotation.value == expected
+        assert annotation.imports == frozenset(
+            {KnownImport(import_path="typing", import_name="Literal")}
+        )
+
+    @pytest.mark.parametrize(
+        ("raw", "expected"),
+        [
+            ("int, optional", "int | None"),
+            # None isn't appended, since the type should cover the default
+            ("int, default 1", "int"),
+            ("int, default = 1", "int"),
+            ("int, default: 1", "int"),
+        ],
+    )
+    @pytest.mark.parametrize("extra_info", [None, "int", ", extra, info"])
+    def test_optional_extra_info(self, raw, expected, extra_info, transformer):
+        doctype = raw
+        if extra_info:
+            doctype = f"{doctype}, {extra_info}"
+
+        tree = _lark.parse(doctype)
+        annotation = transformer.transform(tree)
+
+        assert annotation.value == expected
+
+    # fmt: off
+    @pytest.mark.parametrize(
+        ("fmt", "expected_fmt"),
+        [
+            ("{shape} {name}",                            "{name}"),
+            ("{shape} {name} of {dtype}",                 "{name}[{dtype}]"),
+            ("{shape} {dtype} {name}",                    "{name}[{dtype}]"),
+            ("{dtype} {name}",                            "{name}[{dtype}]"),
+            ("{name} of shape {shape} and dtype {dtype}", "{name}[{dtype}]"),
+            ("{name} of dtype {dtype} and shape {shape}", "{name}[{dtype}]"),
+        ],
+    )
+    @pytest.mark.parametrize("name", ["array", "ndarray", "array-like", "array_like"])
+    @pytest.mark.parametrize("dtype", ["int", "np.int8", "~.foo"])
+    @pytest.mark.parametrize("shape", ["(2, 3)", "(N, m)", "3D", "2-D", "(N, ...)"])
+    def test_shape_n_dtype(self, fmt, expected_fmt, name, dtype, shape, transformer):
+        doctype = fmt.format(name=name, dtype=dtype, shape=shape)
+        expected = expected_fmt.format(name=name, dtype=dtype, shape=shape)
+
+        tree = _lark.parse(doctype)
+        annotation = transformer.transform(tree)
+
+        assert annotation.value == expected
+    # fmt: on