fix-unions-for-return-annotations-too

Author: Darshan808

IPython/core/guarded_eval.py

@@ -793,9 +793,12 @@ def eval_node(node: Union[ast.AST, None], context: EvaluationContext):
 
         if is_property:
             if return_type is not None:
-                context.transient_locals[node.name] = _resolve_annotation(
-                    return_type, context
-                )
+                if is_type_annotation(return_type):
+                    context.transient_locals[node.name] = _resolve_annotation(
+                        return_type, context
+                    )
+                else:
+                    context.transient_locals[node.name] = return_type
             else:
                 return_value = _infer_return_value(node, func_context)
                 context.transient_locals[node.name] = return_value
@@ -806,7 +809,10 @@ def dummy_function(*args, **kwargs):
             pass
 
         if return_type is not None:
-            dummy_function.__annotations__["return"] = return_type
+            if is_type_annotation(return_type):
+                dummy_function.__annotations__["return"] = return_type
+            else:
+                dummy_function.__inferred_return__ = return_type
         else:
             inferred_return = _infer_return_value(node, func_context)
             if inferred_return is not None:

tests/test_completer.py

@@ -2536,6 +2536,7 @@ def _(expected):
             ),
             ["capitalize"],
         ],
+        # Test union types
         [
             "\n".join(
                 [
@@ -2545,6 +2546,27 @@ def _(expected):
             ),
             ["bit_length", "capitalize"],
         ],
+        [
+            "\n".join(
+                [
+                    "def func() -> int | str: pass",
+                    "func().",
+                ]
+            ),
+            ["bit_length", "capitalize"],
+        ],
+        [
+            "\n".join(
+                [
+                    "class T:",
+                    "   @property",
+                    "   def p(self) -> int | str: pass",
+                    "t = T()",
+                    "t.p.",
+                ]
+            ),
+            ["bit_length", "capitalize"],
+        ],
     ],
 )
 def test_undefined_variables(use_jedi, evaluation, code, insert_text):