From a38f446b5ed90b088cc3cd588ada0b4e8749c238 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Alex=20Gr=C3=B6nholm?= Date: Wed, 19 Feb 2025 18:10:46 +0000 Subject: [PATCH] PY-48144: Fixed type inference for pytest.param() Merge-request: IJ-MR-155371 Merged-by: Ilya Kazakevich GitOrigin-RevId: e8b74efe6cf28f953dd2f4709af26071a0af3f72 --- .../PyTestParametrizedTypeProvider.kt | 21 ++++++++++++++++++- .../_pytest/mark.py | 5 +++++ .../pytest/__init__.py | 3 ++- .../testParameters/test_parameter_types.py | 8 +++++++ .../testing/PyTestFixtureResolvingTest.kt | 8 +++++++ 5 files changed, 43 insertions(+), 2 deletions(-) create mode 100644 python/testData/testPytestFixtureResolving/_pytest/mark.py create mode 100644 python/testData/testPytestFixtureResolving/testParameters/test_parameter_types.py diff --git a/python/src/com/jetbrains/python/testing/pyTestParametrized/PyTestParametrizedTypeProvider.kt b/python/src/com/jetbrains/python/testing/pyTestParametrized/PyTestParametrizedTypeProvider.kt index 569c412b63de..512db62f1a7e 100644 --- a/python/src/com/jetbrains/python/testing/pyTestParametrized/PyTestParametrizedTypeProvider.kt +++ b/python/src/com/jetbrains/python/testing/pyTestParametrized/PyTestParametrizedTypeProvider.kt @@ -1,9 +1,15 @@ -// Copyright 2000-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file. +// Copyright 2000-2025 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file. package com.jetbrains.python.testing.pyTestParametrized import com.intellij.openapi.util.Ref +import com.jetbrains.python.PyNames +import com.jetbrains.python.psi.PyCallSiteExpression import com.jetbrains.python.psi.PyFunction +import com.jetbrains.python.psi.PyKeywordArgument import com.jetbrains.python.psi.PyNamedParameter +import com.jetbrains.python.psi.impl.PyBuiltinCache +import com.jetbrains.python.psi.types.PyTupleType +import com.jetbrains.python.psi.types.PyType import com.jetbrains.python.psi.types.PyTypeProviderBase import com.jetbrains.python.psi.types.TypeEvalContext @@ -13,4 +19,17 @@ import com.jetbrains.python.psi.types.TypeEvalContext class PyTestParametrizedTypeProvider : PyTypeProviderBase() { override fun getParameterType(param: PyNamedParameter, func: PyFunction, context: TypeEvalContext) = param.asParametrized(context)?.type?.let { Ref(it) } + + override fun getCallType(function: PyFunction, callSite: PyCallSiteExpression, context: TypeEvalContext): Ref? { + var retval = super.getCallType(function, callSite, context) + if (retval == null && function.qualifiedName == "_pytest.mark.param") { + // Infer the types of positional arguments from the param() call and return them as a tuple + val tupleClass = PyBuiltinCache.getInstance(callSite).getClass(PyNames.TUPLE)?: return null + val args = callSite.getArguments(function) + val argTypes = args.filter { it !is PyKeywordArgument }.map { context.getType(it) } + val tupleType = PyTupleType(tupleClass, argTypes, false) + return Ref(tupleType) + } + return retval + } } diff --git a/python/testData/testPytestFixtureResolving/_pytest/mark.py b/python/testData/testPytestFixtureResolving/_pytest/mark.py new file mode 100644 index 000000000000..12d0f92cfb9e --- /dev/null +++ b/python/testData/testPytestFixtureResolving/_pytest/mark.py @@ -0,0 +1,5 @@ +def param(*args, id=None): + pass + +def parametrize(names, *args): + pass diff --git a/python/testData/testPytestFixtureResolving/pytest/__init__.py b/python/testData/testPytestFixtureResolving/pytest/__init__.py index a6208f4fbb28..930e2ca4c14a 100644 --- a/python/testData/testPytestFixtureResolving/pytest/__init__.py +++ b/python/testData/testPytestFixtureResolving/pytest/__init__.py @@ -1 +1,2 @@ -from _pytest.fixtures import fixture \ No newline at end of file +from _pytest.fixtures import fixture +from _pytest.mark import param, parametrize diff --git a/python/testData/testPytestFixtureResolving/testParameters/test_parameter_types.py b/python/testData/testPytestFixtureResolving/testParameters/test_parameter_types.py new file mode 100644 index 000000000000..a244c4dbe9f2 --- /dev/null +++ b/python/testData/testPytestFixtureResolving/testParameters/test_parameter_types.py @@ -0,0 +1,8 @@ +import pytest + +@pytest.mark.parametrize("param1, param2", [ + (5, 9) + pytest.param(1, "foo", id="namedparam") +]) +def test_named_parameter(param1, param2): + pass diff --git a/python/testSrc/com/jetbrains/python/testing/PyTestFixtureResolvingTest.kt b/python/testSrc/com/jetbrains/python/testing/PyTestFixtureResolvingTest.kt index b6f5af0aba85..3eb02b2a2f33 100644 --- a/python/testSrc/com/jetbrains/python/testing/PyTestFixtureResolvingTest.kt +++ b/python/testSrc/com/jetbrains/python/testing/PyTestFixtureResolvingTest.kt @@ -16,6 +16,7 @@ class PyTestFixtureResolvingTest : PyTestCase() { const val TESTS_SUBDIR = "/testPytestFixtureResolving" const val STR_TYPE_NAME = "str" const val STR_TYPE_DICT = "dict[$STR_TYPE_NAME, $STR_TYPE_NAME]" + const val INT_STR_UNION = "int | str" const val SIMPLE_TEST_DIR = "/testSimple" const val SIMPLE_TEST_CONFTEST_FIXTURE = "/test_conftest_fixture.py" @@ -99,6 +100,9 @@ class PyTestFixtureResolvingTest : PyTestCase() { const val TEST_ASYNC_ITERABLE = "/test_async_iterable.py" const val TEST_ASYNC_ITERATOR = "/test_async_iterator.py" + const val PARAMETRIZED_DIR = "/testParameters" + const val TEST_PARAMETER_TYPES = "/test_parameter_types.py" + } override fun getTestDataPath() = super.getTestDataPath() + TESTS_SUBDIR @@ -337,4 +341,8 @@ class PyTestFixtureResolvingTest : PyTestCase() { fun testAsyncFunction() { assertCorrectType(ASYNC_DIR, TEST_ASYNC_FUNCTION, STR_TYPE_DICT) } + + fun testNamedParameterTypes() { + assertCorrectType(PARAMETRIZED_DIR, TEST_PARAMETER_TYPES, INT_STR_UNION) + } } \ No newline at end of file