Fix AttributeError when using union of dtypes in array expression (#52)

Author: lagru

src/docstub/_docstrings.py

@@ -32,8 +32,23 @@
 
 
 def _find_one_token(tree: lark.Tree, *, name: str) -> lark.Token:
-    """Find token with a specific type name in tree."""
-    tokens = [child for child in tree.children if child.type == name]
+    """Find token with a specific type name in tree.
+
+    Parameters
+    ----------
+    tree : lark.Tree
+    name : str
+        Name of the token to find in the children of `tree`.
+
+    Returns
+    -------
+    token : lark.Token
+    """
+    tokens = [
+        child
+        for child in tree.children
+        if hasattr(child, "type") and child.type == name
+    ]
     if len(tokens) != 1:
         msg = f"expected exactly one Token of type {name}, found {len(tokens)}"
         raise ValueError(msg)

tests/test_docstrings.py

@@ -185,7 +185,9 @@ def test_rst_role(self, doctype, expected):
     )
     @pytest.mark.parametrize("name", ["array", "ndarray", "array-like", "array_like"])
     @pytest.mark.parametrize("dtype", ["int", "np.int8"])
-    @pytest.mark.parametrize("shape", ["(2, 3)", "(N, m)", "3D", "2-D", "(N, ...)"])
+    @pytest.mark.parametrize("shape",
+        ["(2, 3)", "(N, m)", "3D", "2-D", "(N, ...)", "([P,] M, N)"]
+     )
     def test_natlang_array(self, fmt, expected_fmt, name, dtype, shape):
 
         def escape(name: str) -> str:
@@ -202,6 +204,18 @@ def escape(name: str) -> str:
         assert annotation.value == expected
     # fmt: on
 
+    @pytest.mark.parametrize(
+        ("doctype", "expected"),
+        [
+            ("ndarray of dtype (int or float)", "ndarray[int | float]"),
+            ("([P,] M, N) (int or float) array", "array[int | float]"),
+        ],
+    )
+    def test_natlang_array_specific(self, doctype, expected):
+        transformer = DoctypeTransformer()
+        annotation, _ = transformer.doctype_to_annotation(doctype)
+        assert annotation.value == expected
+
     @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}"