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? { 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): - pass + return 42 def test_sample_test(my_fixt): pass +def test_ham(ham, my_fixture): + ham.bit_len# + my_fixture.swapca# + @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,egg,firs,secon):