From 832e7dcc3fec4800d606a9905cd856df57f3715f Mon Sep 17 00:00:00 2001
From: Matthias Bussonnier <bussonniermatthias@gmail.com>
Date: Mon, 2 Oct 2023 14:46:53 +0200
Subject: [PATCH] add examples

---
 papyri/examples.py              |  7 ++++-
 papyri/signature.py             | 14 +++++++---
 papyri/tests/test_gen.py        |  1 +
 papyri/tests/test_signatures.py | 47 +++++++++++++++++++++++++++++++++
 4 files changed, 65 insertions(+), 4 deletions(-)

diff --git a/papyri/examples.py b/papyri/examples.py
index a596d788..f542ef17 100644
--- 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
index b9a7dba8..6fbfe6a3 100644
--- 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
index da90f58a..cf22474a 100644
--- 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
index 9e337cf6..8711a328 100644
--- 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,