PY-31051: pytest fixtures may use yield

In this case we must fetch real type from generator.
We also should open union if one is returned instead
of generator itself.

Any decorator leads to weak type (Union[Any, T]). To fix it, we use KnownDecorators (its better to use EP there)

python/src/com/jetbrains/python/psi/PyKnownDecoratorUtil.java

@@ -75,7 +75,8 @@ public class PyKnownDecoratorUtil {
     ATTR_S("attr.__init__.s"),
     ATTR_ATTRS("attr.__init__.attrs"),
     ATTR_ATTRIBUTES("attr.__init__.attributes"),
-    ATTR_DATACLASS("attr.__init__.dataclass");
+    ATTR_DATACLASS("attr.__init__.dataclass"),
+    PYTEST_FIXTURE("pytest.fixture");
 
     private final QualifiedName myQualifiedName;
 

python/src/com/jetbrains/python/testing/pyTestFixtures/PyTestFixtureReferenceContributor.kt

@@ -9,9 +9,7 @@ import com.jetbrains.python.BaseReference
 import com.jetbrains.python.psi.PyElementGenerator
 import com.jetbrains.python.psi.PyNamedParameter
 import com.jetbrains.python.psi.PyParameter
-import com.jetbrains.python.psi.types.PyType
-import com.jetbrains.python.psi.types.PyTypeProviderBase
-import com.jetbrains.python.psi.types.TypeEvalContext
+import com.jetbrains.python.psi.types.*
 
 class PyTestFixtureReference(namedParameter: PyNamedParameter, fixture: PyTestFixture) : BaseReference(namedParameter) {
   private val functionRef = fixture.function?.let { SmartPointerManager.createPointer(it) }
@@ -32,8 +30,23 @@ object PyTextFixtureTypeProvider : PyTypeProviderBase() {
   override fun getReferenceType(referenceTarget: PsiElement, context: TypeEvalContext, anchor: PsiElement?): Ref<PyType>? {
     val param = referenceTarget as? PyNamedParameter ?: return null
     val fixtureFunc = param.references.filterIsInstance(PyTestFixtureReference::class.java).firstOrNull()?.getFunction() ?: return null
-    return context.getReturnType(fixtureFunc)?.let { Ref(it) }
-
+    val returnType = context.getReturnType(fixtureFunc)
+    if (!fixtureFunc.isGenerator) {
+      return Ref(returnType)
+    }
+    else {
+      //If generator function returns collection this collection is generator
+      // which generates iteratedItemType.
+      // We also must open union (toStream)
+      val itemTypes = PyTypeUtil.toStream(returnType)
+                     .map {
+                       when (it) {
+                         is PyCollectionType -> it.iteratedItemType
+                         else -> it
+                       }
+                     }.toList() ?: return null
+      return Ref(PyUnionType.union(itemTypes))
+    }
   }
 }
 

python/testData/testCompletion/after_test_test.txt

@@ -3,16 +3,20 @@ import pytest
 #
 @pytest.fixture
 def my_fixture():
-    pass
+    yield "hello"
 
 @pytest.fixture
 def ham(my_fixture  ):
-    pass
+    return 42
 
 
 def test_sample_test(my_fixture ):
     pass
 
+def test_ham(ham, my_fixture):
+    ham.bit_length()    #
+    my_fixture.swapcase()   #
+
 @pytest.mark.parametrize(('spam,eggs'), [(1,1), (2,3), (3,3)])
 @pytest.mark.parametrize('first,second', [(1,1), (2,3), (3,3)])
 def test_sample(spam    ,eggs    ,first   ,second ):

python/testData/testCompletion/test_test.py

@@ -3,16 +3,20 @@ import pytest
 #
 @pytest.fixture
 def my_fixture():
-    pass
+    yield "hello"
 
 @pytest.fixture
 def ham(my_fixt<caret>):
-    pass
+    return 42
 
 
 def test_sample_test(my_fixt<caret>):
     pass
 
+def test_ham(ham, my_fixture):
+    ham.bit_len<caret>#
+    my_fixture.swapca<caret>#
+
 @pytest.mark.parametrize(('spam,eggs'), [(1,1), (2,3), (3,3)])
 @pytest.mark.parametrize('first,second', [(1,1), (2,3), (3,3)])
 def test_sample(spa<caret>,egg<caret>,firs<caret>,secon<caret>):