From 3a4b52660e5da1a208c2115680dc2efe1b3eeb94 Mon Sep 17 00:00:00 2001
From: "Ilya.Kazakevich" <Ilya.Kazakevich@jetbrains.com>
Date: Tue, 21 Aug 2018 22:24:58 +0300
Subject: [PATCH] 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/psi/PyKnownDecoratorUtil.java      |  3 ++-
 .../PyTestFixtureReferenceContributor.kt      | 23 +++++++++++++++----
 .../testCompletion/after_test_test.txt        |  8 +++++--
 python/testData/testCompletion/test_test.py   |  8 +++++--
 4 files changed, 32 insertions(+), 10 deletions(-)

diff --git a/python/src/com/jetbrains/python/psi/PyKnownDecoratorUtil.java b/python/src/com/jetbrains/python/psi/PyKnownDecoratorUtil.java
index d152af33e4c4..379fdae8c859 100644
--- a/python/src/com/jetbrains/python/psi/PyKnownDecoratorUtil.java
+++ b/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;
 
diff --git a/python/src/com/jetbrains/python/testing/pyTestFixtures/PyTestFixtureReferenceContributor.kt b/python/src/com/jetbrains/python/testing/pyTestFixtures/PyTestFixtureReferenceContributor.kt
index 11d8783bf174..274e7aa0cac8 100644
--- a/python/src/com/jetbrains/python/testing/pyTestFixtures/PyTestFixtureReferenceContributor.kt
+++ b/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))
+    }
   }
 }
 
diff --git a/python/testData/testCompletion/after_test_test.txt b/python/testData/testCompletion/after_test_test.txt
index 1a6e080281a6..c46d14aa237e 100644
--- a/python/testData/testCompletion/after_test_test.txt
+++ b/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 ):
diff --git a/python/testData/testCompletion/test_test.py b/python/testData/testCompletion/test_test.py
index f1d1e281d3ca..171af07d9458 100644
--- a/python/testData/testCompletion/test_test.py
+++ b/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>):