add examples

jupyter · Oct 2, 2023 · 832e7dc · 832e7dc
1 parent 8516262
commit 832e7dc
Show file tree

Hide file tree

Showing 4 changed files with 65 additions and 4 deletions.
diff --git a/papyri/examples.py b/papyri/examples.py
@@ -39,8 +39,12 @@
 
 """
 
+from typing import Optional, Union
 
-def example1(pos, only, /, var, args, *, kwargs, also=None):
+
+def example1(
+    pos: int, only: None, /, var: Union[float, bool], args=1, *, kwargs, also=None
+) -> Optional[str]:
     """
     first example.
 
@@ -56,6 +60,7 @@ def example1(pos, only, /, var, args, *, kwargs, also=None):
     ... plt.show()
 
     """
+    return "ok"
 
 
 def example2():

diff --git a/papyri/signature.py b/papyri/signature.py
@@ -46,6 +46,7 @@ def to_parameter(self) -> inspect.Parameter:
 class SignatureNode(Node):
     kind: str  # maybe enum, is it a function, async generator, generator, etc.
     parameters: List[ParameterNode]  # of pairs, we don't use dict because of ordering
+    return_annotation: Union[Empty, str]
 
     def to_signature(self):
         return inspect.Signature([p.to_parameter() for p in self.parameters])
@@ -101,6 +102,7 @@ def to_node(self) -> SignatureNode:
 
         parameters = []
         for param in self.parameters.values():
+            annotation: Union[Empty, str]
             if param.annotation is inspect._empty:
                 annotation = _empty
             elif isinstance(param.annotation, str):
@@ -119,7 +121,13 @@ def to_node(self) -> SignatureNode:
                 )
             )
         assert isinstance(kind, str)
-        return SignatureNode(kind=kind, parameters=parameters)
+        return SignatureNode(
+            kind=kind,
+            parameters=parameters,
+            return_annotation=_empty
+            if self._sig.return_annotation is inspect._empty
+            else str(self._sig.return_annotation),
+        )
 
     @property
     def parameters(self):
@@ -145,10 +153,10 @@ def annotations(self) -> bool:
         return self.target_item.__annotations__
 
     @property
-    def return_annotation(self) -> Optional[str]:
+    def return_annotation(self) -> Union[Empty, str]:
         return_annotation = self._sig.return_annotation
         return (
-            None
+            _empty
             if return_annotation is inspect._empty
             else inspect.formatannotation(return_annotation)
         )

diff --git a/papyri/tests/test_gen.py b/papyri/tests/test_gen.py
@@ -151,5 +151,6 @@ def test_self():
                 "default": {"data": "None", "type": "str"},
             },
         ],
+        "return_annotation": {"type": "Empty"},
     }
     assert g.data["papyri"].to_dict()["signature"] is None
diff --git a/papyri/tests/test_signatures.py b/papyri/tests/test_signatures.py
@@ -4,6 +4,7 @@
 """
 from papyri.signature import Signature as SignatureObject, SignatureNode
 import json
+from typing import Union, Optional
 
 import pytest
 
@@ -78,6 +79,7 @@ def function_1(posonly, /, pos_or_k, pos_ok_k_d=1, *varargs, **varkwargs):
           "type": "ParameterNode"
         }
       ],
+      "return_annotation": {"type": "Empty"},
       "type": "SignatureNode"
     }"""
     pass
@@ -145,6 +147,7 @@ def async_function_2(posonly, /, pos_or_k, pos_ok_k_d=1, *varargs, **varkwargs):
           "type": "ParameterNode"
         }
       ],
+      "return_annotation": {"type": "Empty"},
       "type": "SignatureNode"
     }"""
     pass
@@ -212,6 +215,7 @@ def generator_function_3(posonly, /, pos_or_k, pos_ok_k_d=1, *varargs, **varkwar
           "type": "ParameterNode"
         }
       ],
+      "return_annotation": {"type": "Empty"},
       "type": "SignatureNode"
     }"""
     yield
@@ -281,11 +285,54 @@ async def async_generator_function_4(
           "type": "ParameterNode"
         }
       ],
+      "return_annotation": {"type": "Empty"},
       "type": "SignatureNode"
     }"""
     yield
 
 
+@add
+def function_with_annotation5(a: int, b: Union[int, float]) -> Optional[bool]:
+    """
+    {
+      "kind": "function",
+      "parameters": [
+        {
+          "annotation": {
+            "data": "int",
+            "type": "str"
+          },
+          "default": {
+            "type": "Empty"
+          },
+          "kind": "POSITIONAL_OR_KEYWORD",
+          "name": "a",
+          "type": "ParameterNode"
+        },
+        {
+          "annotation": {
+            "data": "Union[int, float]",
+            "type": "str"
+          },
+          "default": {
+            "type": "Empty"
+          },
+          "kind": "POSITIONAL_OR_KEYWORD",
+          "name": "b",
+          "type": "ParameterNode"
+        }
+      ],
+      "return_annotation": {
+          "data": "typing.Optional[bool]",
+          "type": "str"
+      },
+      "type": "SignatureNode"
+    }
+
+    """
+    pass
+
+
 @pytest.mark.parametrize(
     "func",
     all_funcs,