Add support of callable expressions

Author: lagru

src/docstub/_docstrings.py

@@ -384,6 +384,35 @@ def subscription(self, tree):
         out = f"{_container}[{_content}]"
         return out
 
+    def param_spec(self, tree):
+        """
+        Parameters
+        ----------
+        tree : lark.Tree
+
+        Returns
+        -------
+        out : str
+        """
+        _content = ", ".join(tree.children)
+        out = f"[{_content}]"
+        return out
+
+    def callable(self, tree):
+        """
+        Parameters
+        ----------
+        tree : lark.Tree
+
+        Returns
+        -------
+        out : str
+        """
+        _callable, *_content = tree.children
+        _content = ", ".join(_content)
+        out = f"{_callable}[{_content}]"
+        return out
+
     def natlang_literal(self, tree):
         """
         Parameters

src/docstub/doctype.lark

@@ -20,6 +20,7 @@
 ?type: qualname
     | union
     | subscription
+    | callable
     | natlang_literal
     | natlang_container
     | natlang_array
@@ -47,6 +48,17 @@ _OR: "or" | "|"
 subscription: qualname "[" type ("," type)* ("," ELLIPSES)? "]"
 
 
+// An expression describing a callable like "Callable[[int], str]"
+// [1] https://typing.python.org/en/latest/spec/callables.html#callable
+//
+callable: qualname "[" ELLIPSES ("," type)? "]"
+    | qualname "[" param_spec "," type "]"
+
+
+// The parameter specification inside a callable expression.
+param_spec: "[" type? ("," type)* ("," ELLIPSES)? "]"
+
+
 // Allow Python's ellipses object
 ELLIPSES: "..."
 

tests/test_docstrings.py

@@ -154,6 +154,40 @@ def test_subscription_error(self, doctype):
         with pytest.raises(lark.exceptions.UnexpectedInput):
             transformer.doctype_to_annotation(doctype)
 
+    @pytest.mark.parametrize(
+        "doctype",
+        [
+            "Callable[[int], str]",
+            "some_func[[int], str]",
+            "Callable[[int, float, byte], list[str]]",
+            "Callable[..., str]",
+            "Callable[[], str]",
+            "Callback[...]",
+            "Callable[Concatenate[int, float], str]",
+            "Callable[Concatenate[int, ...], str]",
+            "Callable[P, str]",
+        ],
+    )
+    def test_callable(self, doctype):
+        transformer = DoctypeTransformer()
+        annotation, _ = transformer.doctype_to_annotation(doctype)
+        assert annotation.value == doctype
+
+    @pytest.mark.parametrize(
+        "doctype",
+        [
+            "Callable[[...], int]",
+            "Callable[[..., str], int]",
+            "Callable[[float, str], int, byte]",
+            "Callable[Ellipsis]",
+            "Callable[Ellipsis, int]",
+        ],
+    )
+    def test_callable_error(self, doctype):
+        transformer = DoctypeTransformer()
+        with pytest.raises(lark.exceptions.UnexpectedInput):
+            transformer.doctype_to_annotation(doctype)
+
     @pytest.mark.parametrize(
         ("doctype", "expected"),
         [