PY-48144: Fixed type inference for pytest.param()

Merge-request: IJ-MR-155371
Merged-by: Ilya Kazakevich <ilya.kazakevich@jetbrains.com>

GitOrigin-RevId: e8b74efe6cf28f953dd2f4709af26071a0af3f72

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<PyType>? {
+    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
+  }
 }

python/testData/testPytestFixtureResolving/_pytest/mark.py

@@ -0,0 +1,5 @@
+def param(*args, id=None):
+    pass
+
+def parametrize(names, *args):
+    pass

python/testData/testPytestFixtureResolving/pytest/__init__.py

@@ -1 +1,2 @@
-from _pytest.fixtures import fixture
+from _pytest.fixtures import fixture
+from _pytest.mark import param, parametrize

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, para<caret>m2):
+    pass

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)
+  }
 }