Suggest argument names for parametrized classes

GitOrigin-RevId: df1a269bea4e5b92e4f30e294b544b084fdba5a1

python/src/com/jetbrains/python/testing/pyTestParametrized/PyTestParamterized.kt

@@ -76,11 +76,13 @@ private fun getParametersFromDecorator(decorator: PyDecorator, function: PyFunct
   return parameterNames.mapIndexed { i, name -> PyTestParameter(name, parameterTypes[i]) }
 }
 
-private fun patchTypesWithIndirectFixture(paramNames: List<String>,
-                                          paramTypes: Array<PyType?>,
-                                          function: PyFunction,
-                                          evalContext: TypeEvalContext,
-                                          indirectKeyword: PyKeywordArgument) {
+private fun patchTypesWithIndirectFixture(
+  paramNames: List<String>,
+  paramTypes: Array<PyType?>,
+  function: PyFunction,
+  evalContext: TypeEvalContext,
+  indirectKeyword: PyKeywordArgument,
+) {
   if (indirectKeyword.keyword != "indirect") return
   val indirectParams = when (val expression = PyEvaluator().evaluate(indirectKeyword.valueExpression)) {
     is Boolean -> if (expression) paramNames else emptyList() // indirect=True
@@ -106,13 +108,23 @@ internal data class PyTestParameter(val name: String, val type: PyType? = null)
  * @return List<String> if test function decorated with parametrize -- return parameter names
  */
 internal fun PyFunction.getParametersOfParametrized(evalContext: TypeEvalContext): List<PyTestParameter> {
-  val decoratorList = decoratorList ?: return emptyList()
-  if (!isTestElement(this, ThreeState.NO, evalContext)) {
-    return emptyList()
-  }
-  return decoratorList.decorators
-    .filter { it.name == "parametrize" }
-    .flatMap { getParametersFromDecorator(it, this, evalContext) }
+  val result = mutableListOf<PyTestParameter>()
+  // function-level decorators
+  decoratorList?.decorators
+    ?.filter { it.name == "parametrize" }
+    ?.forEach { result.addAll(getParametersFromDecorator(it, this, evalContext)) }
 
+  // class-level decorators applied to containing test class
+  containingClass?.decoratorList?.decorators
+    ?.filter { it.name == "parametrize" }
+    ?.forEach { result.addAll(getParametersFromDecorator(it, this, evalContext)) }
+
+  // de-duplicate preserving order
+  val seen = HashSet<String>()
+  val unique = ArrayList<PyTestParameter>(result.size)
+  for (p in result) {
+    if (seen.add(p.name)) unique.add(p)
+  }
+  return unique
 }
 

python/testData/testCompletion/after_test_class_level_parametrize.txt

@@ -0,0 +1,7 @@
+import pytest
+
+
+@pytest.mark.parametrize("arg", [])
+class Test:
+    def test_(self, arg  ):
+        pass

python/testData/testCompletion/test_class_level_parametrize.py

@@ -0,0 +1,7 @@
+import pytest
+
+
+@pytest.mark.parametrize("arg", [])
+class Test:
+    def test_(self, ar<caret>):
+        pass

python/testSrc/com/jetbrains/python/testing/PyTestFixtureAndParametrizedTest.kt

@@ -82,6 +82,14 @@ class PyTestFixtureAndParametrizedTest : PyTestCase() {
     assertFalse(variants.contains("tmp_path"))
   }
 
+  fun testClassLevelParametrizeCompletion() {
+    myFixture.copyDirectoryToProject(".", ".")
+    myFixture.configureByFile("test_class_level_parametrize.py")
+    myFixture.completeBasic()
+    val text = myFixture.editor.document.text
+    assertTrue(text.contains("def test_(self, arg"))
+  }
+
   fun testRename() {
     myFixture.configureByFile("test_for_rename.py")
     myFixture.renameElementAtCaret("spam")