From 148fffae2877509a95607d2113c4b2df8cf2df79 Mon Sep 17 00:00:00 2001
From: Dave Halter <davidhalter88@gmail.com>
Date: Fri, 27 Dec 2019 01:50:17 +0100
Subject: [PATCH] Make yield pytest fixtures work

---
 jedi/inference/value/function.py |  8 +++++---
 jedi/plugins/pytest.py           | 15 +++++++++++++--
 test/completion/pytest.py        | 21 +++++++++++++++++++--
 3 files changed, 37 insertions(+), 7 deletions(-)

diff --git a/jedi/inference/value/function.py b/jedi/inference/value/function.py
index dfe190f9..18963f64 100644
--- a/jedi/inference/value/function.py
+++ b/jedi/inference/value/function.py
@@ -277,17 +277,19 @@ class BaseFunctionExecutionContext(ValueContext, TreeContextMixin):
             for lazy_value in self.get_yield_lazy_values()
         )
 
+    def is_generator(self):
+        return bool(get_yield_exprs(self.inference_state, self.tree_node))
+
     def infer(self):
         """
         Created to be used by inheritance.
         """
         inference_state = self.inference_state
         is_coroutine = self.tree_node.parent.type in ('async_stmt', 'async_funcdef')
-        is_generator = bool(get_yield_exprs(inference_state, self.tree_node))
         from jedi.inference.gradual.base import GenericClass
 
         if is_coroutine:
-            if is_generator:
+            if self.is_generator():
                 if inference_state.environment.version_info < (3, 6):
                     return NO_VALUES
                 async_generator_classes = inference_state.typing_module \
@@ -312,7 +314,7 @@ class BaseFunctionExecutionContext(ValueContext, TreeContextMixin):
                     GenericClass(c, TupleGenericManager(generics)) for c in async_classes
                 ).execute_annotation()
         else:
-            if is_generator:
+            if self.is_generator():
                 return ValueSet([iterable.Generator(inference_state, self)])
             else:
                 return self.get_return_values()
diff --git a/jedi/plugins/pytest.py b/jedi/plugins/pytest.py
index 3629b49e..05245737 100644
--- a/jedi/plugins/pytest.py
+++ b/jedi/plugins/pytest.py
@@ -17,8 +17,17 @@ def execute(callback):
 
 def infer_anonymous_param(func):
     def get_returns(value):
+        if value.tree_node.annotation is not None:
+            return value.execute_with_values()
+
+        # In pytest we need to differentiate between generators and normal
+        # returns.
+        # Parameters still need to be anonymous, .as_context() ensures that.
         function_context = value.as_context()
-        return function_context.get_return_values()
+        if function_context.is_generator():
+            return function_context.merge_yield_values()
+        else:
+            return function_context.get_return_values()
 
     def wrapper(param):
         module = param.get_root_context()
@@ -59,4 +68,6 @@ def _iter_pytest_modules(module_context):
 class FixtureFilter(ParserTreeFilter):
     def _filter(self, names):
         for name in super(FixtureFilter, self)._filter(names):
-            yield name
+            if name.parent.type == 'funcdef':
+                # Class fixtures are not supported
+                yield name
diff --git a/test/completion/pytest.py b/test/completion/pytest.py
index bf9bb20d..b8423c16 100644
--- a/test/completion/pytest.py
+++ b/test/completion/pytest.py
@@ -13,14 +13,25 @@ def my_simple_fixture():
     return 1
 
 
+@fixture
+def my_yield_fixture():
+    yield 1
+
+
+@fixture
+class MyClassFixture():
+    pass
+
 # -----------------
 # goto/infer
 # -----------------
 
-#! 18 'def my_conftest_fixture'
-def test_x(my_conftest_fixture, my_fixture, my_not_existing_fixture):
+#! 18 ['def my_conftest_fixture']
+def test_x(my_conftest_fixture, my_fixture, my_not_existing_fixture, my_yield_fixture):
     #? str()
     my_fixture
+    #? int()
+    my_yield_fixture
     #?
     my_not_existing_fixture
     #? float()
@@ -30,6 +41,12 @@ def test_x(my_conftest_fixture, my_fixture, my_not_existing_fixture):
 def test_x(my_conftest_fixture, my_fixture):
     pass
 
+
+#! 18 ['param MyClassFixture']
+def test_x(MyClassFixture):
+    #?
+    MyClassFixture
+
 # -----------------
 # completion
 # -----------------