[python] Do not store unused decorator arguments in stubs

GitOrigin-RevId: d1a5f63dbed11a4f373d1d7936700af3bc74f70f

python/python-psi-api/src/com/jetbrains/python/psi/stubs/PyDecoratorStub.java

@@ -19,7 +19,6 @@ import com.intellij.psi.stubs.StubElement;
 import com.intellij.psi.util.QualifiedName;
 import com.jetbrains.python.psi.PyDecorator;
 import org.jetbrains.annotations.ApiStatus;
-import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
 public interface PyDecoratorStub extends StubElement<PyDecorator> {
@@ -34,11 +33,7 @@ public interface PyDecoratorStub extends StubElement<PyDecorator> {
    */
   boolean hasArgumentList();
 
-  @ApiStatus.Experimental
+  @ApiStatus.Internal
   @Nullable
-  String getPositionalArgumentLiteralText(int position);
-
-  @ApiStatus.Experimental
-  @Nullable
-  String getNamedArgumentLiteralText(@NotNull String name);
+  <T> T getCustomStub(Class<T> stubClass);
 }

python/python-psi-impl/resources/META-INF/PythonPsiImpl.xml

@@ -477,9 +477,8 @@
     <extensionPoint qualifiedName="Pythonid.customClassStubType"
                     interface="com.jetbrains.python.psi.impl.stubs.PyCustomClassStubType"
                     dynamic="true"/>
-
-    <extensionPoint qualifiedName="Pythonid.decoratorIndexer"
-                    interface="com.jetbrains.python.PyCustomDecoratorIndexer"
+    <extensionPoint qualifiedName="Pythonid.customDecoratorStubType"
+                    interface="com.jetbrains.python.psi.impl.stubs.PyCustomDecoratorStubType"
                     dynamic="true"/>
 
     <extensionPoint qualifiedName="Pythonid.visitorFilter"
@@ -520,6 +519,7 @@
     <customTargetExpressionStubType implementation="com.jetbrains.python.psi.impl.stubs.PyTypingNewTypeStubType"/>
 
     <customClassStubType implementation="com.jetbrains.python.psi.impl.stubs.PyDataclassStubType"/>
+    <customDecoratorStubType implementation="com.jetbrains.python.psi.stubs.PyTestFixtureDecoratorStubType"/>
 
     <typeProvider implementation="com.jetbrains.python.psi.types.PyCollectionTypeByModificationsProvider" order="last"/>
     <typeProvider implementation="com.jetbrains.python.codeInsight.decorator.PyDecoratedFunctionTypeProvider"/>

python/python-psi-impl/src/com/jetbrains/python/PyCustomDecoratorIndexer.kt

@@ -1,18 +0,0 @@
-package com.jetbrains.python
-
-import com.intellij.openapi.extensions.ExtensionPointName
-import com.intellij.psi.stubs.StubIndexKey
-import com.jetbrains.python.psi.PyDecorator
-import com.jetbrains.python.psi.stubs.PyDecoratorStub
-import org.jetbrains.annotations.ApiStatus
-
-
-@ApiStatus.Experimental
-interface PyCustomDecoratorIndexer {
-  fun getKey(): StubIndexKey<String, PyDecorator>
-  fun getKeyForStub(decorator: PyDecoratorStub): String?
-  companion object {
-    @JvmField val EP_NAME = ExtensionPointName.create<PyCustomDecoratorIndexer>("Pythonid.decoratorIndexer")
-  }
-}
-

python/python-psi-impl/src/com/jetbrains/python/psi/PyFileElementType.java

@@ -60,7 +60,7 @@ public class PyFileElementType extends IStubFileElementType<PyFileStub> {
   @Override
   public int getStubVersion() {
     // Don't forget to update versions of indexes that use the updated stub-based elements
-    return 88;
+    return 89;
   }
 
   @Nullable

python/python-psi-impl/src/com/jetbrains/python/psi/impl/pyDecoratorExt.kt

@@ -1,15 +0,0 @@
-package com.jetbrains.python.psi.impl
-
-import com.intellij.extapi.psi.StubBasedPsiElementBase
-import com.jetbrains.python.psi.PyDecorator
-import com.jetbrains.python.psi.impl.stubs.PyDecoratorCallElementType.evaluateArgumentValue
-
-import com.jetbrains.python.psi.stubs.PyDecoratorStub
-
-fun PyDecorator.getNamedArgument(name: String): String? {
-  return when(val stub = (this as StubBasedPsiElementBase<PyDecoratorStub>).greenStub) {
-    null -> getKeywordArgument(name)?.let(::evaluateArgumentValue)
-    else -> stub.getNamedArgumentLiteralText(name)
-  }
-}
-

python/python-psi-impl/src/com/jetbrains/python/psi/impl/stubs/PyCustomDecoratorStub.kt

@@ -0,0 +1,3 @@
+package com.jetbrains.python.psi.impl.stubs
+
+interface PyCustomDecoratorStub: PyCustomStub<PyCustomDecoratorStubType<*>>

python/python-psi-impl/src/com/jetbrains/python/psi/impl/stubs/PyCustomDecoratorStubType.kt

@@ -0,0 +1,13 @@
+package com.jetbrains.python.psi.impl.stubs
+
+import com.intellij.openapi.extensions.ExtensionPointName
+import com.jetbrains.python.psi.PyDecorator
+import org.jetbrains.annotations.ApiStatus
+
+@ApiStatus.Internal
+interface PyCustomDecoratorStubType<T : PyCustomDecoratorStub> : PyCustomStubType<PyDecorator, T> {
+  companion object {
+    @JvmField
+    val EP_NAME = ExtensionPointName.create<PyCustomDecoratorStubType<out PyCustomDecoratorStub>>("Pythonid.customDecoratorStubType")
+  }
+}

python/python-psi-impl/src/com/jetbrains/python/psi/impl/stubs/PyDecoratorCallElementType.java

@@ -8,27 +8,21 @@ import com.intellij.psi.stubs.StubElement;
 import com.intellij.psi.stubs.StubInputStream;
 import com.intellij.psi.stubs.StubOutputStream;
 import com.intellij.psi.util.QualifiedName;
-import com.intellij.util.containers.ContainerUtil;
-import com.jetbrains.python.PyCustomDecoratorIndexer;
-import com.jetbrains.python.psi.*;
+import com.jetbrains.python.psi.PyDecorator;
+import com.jetbrains.python.psi.PyStubElementType;
 import com.jetbrains.python.psi.impl.PyDecoratorImpl;
-import com.jetbrains.python.psi.impl.PyEvaluator;
 import com.jetbrains.python.psi.stubs.PyDecoratorStub;
 import com.jetbrains.python.psi.stubs.PyDecoratorStubIndex;
-import org.jetbrains.annotations.ApiStatus;
 import org.jetbrains.annotations.NotNull;
-import org.jetbrains.annotations.Nullable;
 
 import java.io.IOException;
-import java.util.ArrayList;
-import java.util.HashMap;
 import java.util.List;
-import java.util.Map;
 
 /**
  * Actual serialized data of a decorator call.
  */
-public class PyDecoratorCallElementType extends PyStubElementType<PyDecoratorStub, PyDecorator> {
+public class PyDecoratorCallElementType extends PyStubElementType<PyDecoratorStub, PyDecorator>
+  implements PyCustomizableStubElementType<PyDecorator, PyCustomDecoratorStub, PyCustomDecoratorStubType<? extends PyCustomDecoratorStub>> {
   public PyDecoratorCallElementType() {
     super("DECORATOR_CALL");
   }
@@ -47,43 +41,20 @@ public class PyDecoratorCallElementType extends PyStubElementType<PyDecoratorStu
   @Override
   @NotNull
   public PyDecoratorStub createStub(@NotNull PyDecorator psi, StubElement parentStub) {
-    PyExpression[] arguments = psi.getArguments();
-    List<String> positionalArguments = new ArrayList<>();
-    Map<String, String> namedArguments = new HashMap<>();
-    for (PyExpression argument : arguments) {
-      if (argument instanceof PyKeywordArgument keywordArgument) {
-        String keyword = keywordArgument.getKeyword();
-        String value = evaluateArgumentValue(keywordArgument);
-        if (keyword != null && value != null) {
-          namedArguments.put(keyword, value);
-        }
-      }
-      else {
-        String value = evaluateArgumentValue(argument);
-        positionalArguments.add(value);
-      }
-    }
-    return new PyDecoratorStubImpl(psi.getQualifiedName(), psi.hasArgumentList(), parentStub, ContainerUtil.unmodifiableOrEmptyList(
-      positionalArguments), ContainerUtil.unmodifiableOrEmptyMap(namedArguments));
+    return new PyDecoratorStubImpl(psi.getQualifiedName(), psi.hasArgumentList(), parentStub, createCustomStub(psi));
   }
 
-  @Nullable
-  @ApiStatus.Internal
-  public static String evaluateArgumentValue(@Nullable PyExpression expression) {
-    return PyEvaluator.evaluateNoResolve(expression, String.class);
+  @Override
+  public @NotNull List<PyCustomDecoratorStubType<? extends PyCustomDecoratorStub>> getExtensions() {
+    return PyCustomDecoratorStubType.EP_NAME.getExtensionList();
   }
 
   @Override
   public void serialize(@NotNull PyDecoratorStub stub, @NotNull StubOutputStream dataStream) throws IOException {
     QualifiedName.serialize(stub.getQualifiedName(), dataStream);
     dataStream.writeBoolean(stub.hasArgumentList());
-    PyDecoratorStubImpl decoratorStub = (PyDecoratorStubImpl)stub;
-    PyFileElementType.writeNullableList(dataStream, decoratorStub.getPositionalArguments());
-    dataStream.writeInt(decoratorStub.getKeywordArguments().size());
-    for (Map.Entry<String, String> entry : decoratorStub.getKeywordArguments().entrySet()) {
-      dataStream.writeName(entry.getKey());
-      dataStream.writeName(entry.getValue());
-    }
+    PyCustomDecoratorStub customStub = stub.getCustomStub(PyCustomDecoratorStub.class);
+    serializeCustomStub(customStub, dataStream);
   }
 
   @Override
@@ -92,12 +63,6 @@ public class PyDecoratorCallElementType extends PyStubElementType<PyDecoratorStu
     final QualifiedName qualifiedName = stub.getQualifiedName();
     if (qualifiedName != null) {
       sink.occurrence(PyDecoratorStubIndex.KEY, qualifiedName.toString());
-      PyCustomDecoratorIndexer.EP_NAME.getExtensionList().forEach(extension -> {
-        String keyForStub = extension.getKeyForStub(stub);
-        if (keyForStub != null) {
-          sink.occurrence(extension.getKey(), keyForStub);
-        }
-      });
     }
   }
 
@@ -106,18 +71,7 @@ public class PyDecoratorCallElementType extends PyStubElementType<PyDecoratorStu
   public PyDecoratorStub deserialize(@NotNull StubInputStream dataStream, StubElement parentStub) throws IOException {
     QualifiedName q_name = QualifiedName.deserialize(dataStream);
     boolean hasArgumentList = dataStream.readBoolean();
-    List<String> positionalArguments = PyFileElementType.readNullableList((dataStream));
-    int namedSize = dataStream.readInt();
-    Map<String, String> namedArguments = null;
-    if (namedSize > 0) {
-      namedArguments = new HashMap<>(namedSize);
-      for (int i = 0; i < namedSize; i++) {
-        String key = dataStream.readNameString();
-        String value = dataStream.readNameString();
-        namedArguments.put(key, value);
-      }
-    }
-    return new PyDecoratorStubImpl(q_name, hasArgumentList, parentStub, ContainerUtil.notNullize(positionalArguments),
-                                   ContainerUtil.notNullize(namedArguments));
+    PyCustomDecoratorStub customStub = deserializeCustomStub(dataStream);
+    return new PyDecoratorStubImpl(q_name, hasArgumentList, parentStub, customStub);
   }
 }

python/python-psi-impl/src/com/jetbrains/python/psi/impl/stubs/PyDecoratorStubImpl.java

@@ -4,35 +4,28 @@ package com.jetbrains.python.psi.impl.stubs;
 import com.intellij.psi.stubs.StubBase;
 import com.intellij.psi.stubs.StubElement;
 import com.intellij.psi.util.QualifiedName;
-import com.jetbrains.python.PyElementTypes;
+import com.intellij.util.ObjectUtils;
 import com.jetbrains.python.PyStubElementTypes;
 import com.jetbrains.python.psi.PyDecorator;
 import com.jetbrains.python.psi.stubs.PyDecoratorStub;
-import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
-import java.util.List;
-import java.util.Map;
-
 /**
  * Decorator stub storage.
  */
 public class PyDecoratorStubImpl extends StubBase<PyDecorator> implements PyDecoratorStub {
   private final QualifiedName myQualifiedName;
   private final boolean myHasArgumentList;
-  private final List<String> myPositionalArguments;
-  private final Map<String, String> myKeywordArguments;
+  private final @Nullable PyCustomDecoratorStub myCustomStub;
 
   protected PyDecoratorStubImpl(final QualifiedName qualname,
                                 final boolean hasArgumentList,
                                 final StubElement parent,
-                                final @NotNull List<String> positionalArguments,
-                                final @NotNull Map<String, String> namedArguments) {
+                                final @Nullable PyCustomDecoratorStub customStub) {
     super(parent, PyStubElementTypes.DECORATOR_CALL);
     myQualifiedName = qualname;
     myHasArgumentList = hasArgumentList;
-    myPositionalArguments = positionalArguments;
-    myKeywordArguments = namedArguments;
+    myCustomStub = customStub;
   }
 
   @Override
@@ -46,21 +39,7 @@ public class PyDecoratorStubImpl extends StubBase<PyDecorator> implements PyDeco
   }
 
   @Override
-  public @Nullable String getPositionalArgumentLiteralText(int position) {
-    if (position >= myPositionalArguments.size()) return null;
-    return myPositionalArguments.get(position);
-  }
-
-  @Override
-  public @Nullable String getNamedArgumentLiteralText(@NotNull String name) {
-    return myKeywordArguments.get(name);
-  }
-
-  List<String> getPositionalArguments() {
-    return myPositionalArguments;
-  }
-
-  Map<String, String> getKeywordArguments() {
-    return myKeywordArguments;
+  public <T> @Nullable T getCustomStub(Class<T> stubClass) {
+    return ObjectUtils.tryCast(myCustomStub, stubClass);
   }
 }

python/python-psi-impl/src/com/jetbrains/python/psi/impl/stubs/PyTargetExpressionStubImpl.java

@@ -19,7 +19,7 @@ import com.intellij.openapi.util.text.StringUtil;
 import com.intellij.psi.stubs.StubBase;
 import com.intellij.psi.stubs.StubElement;
 import com.intellij.psi.util.QualifiedName;
-import com.jetbrains.python.PyElementTypes;
+import com.intellij.util.ObjectUtils;
 import com.jetbrains.python.PyStubElementTypes;
 import com.jetbrains.python.psi.PyTargetExpression;
 import com.jetbrains.python.psi.stubs.PyTargetExpressionStub;
@@ -102,10 +102,7 @@ public class PyTargetExpressionStubImpl extends StubBase<PyTargetExpression> imp
   @Nullable
   @Override
   public <T> T getCustomStub(Class<T> stubClass) {
-    if (stubClass.isInstance(myCustomStub)) {
-      return stubClass.cast(myCustomStub);
-    }
-    return null;
+    return ObjectUtils.tryCast(myCustomStub, stubClass);
   }
 
   @Nullable

python/python-psi-impl/src/com/jetbrains/python/psi/stubs/PyTestFixtureDecoratorStub.kt

@@ -0,0 +1,7 @@
+package com.jetbrains.python.psi.stubs
+
+import com.jetbrains.python.psi.impl.stubs.PyCustomDecoratorStub
+
+interface PyTestFixtureDecoratorStub : PyCustomDecoratorStub {
+  val name: String
+}

python/python-psi-impl/src/com/jetbrains/python/psi/stubs/PyTestFixtureDecoratorStubType.kt

@@ -0,0 +1,32 @@
+package com.jetbrains.python.psi.stubs
+
+import com.intellij.psi.stubs.StubInputStream
+import com.intellij.psi.stubs.StubOutputStream
+import com.jetbrains.python.psi.PyDecorator
+import com.jetbrains.python.psi.impl.stubs.PyCustomDecoratorStubType
+import com.jetbrains.python.testing.pyTestFixtures.TEST_FIXTURE_DECORATOR_NAMES
+import com.jetbrains.python.testing.pyTestFixtures.getTestFixtureName
+
+class PyTestFixtureDecoratorStubType : PyCustomDecoratorStubType<PyTestFixtureDecoratorStub> {
+  override fun createStub(psi: PyDecorator): PyTestFixtureDecoratorStub? {
+    val qName = psi.getQualifiedName()
+    if (qName == null || qName.toString() !in TEST_FIXTURE_DECORATOR_NAMES) {
+      return null
+    }
+    val testFixtureName = getTestFixtureName(psi) ?: return null
+    return PyTestFixtureDecoratorStubImpl(testFixtureName)
+  }
+
+  override fun deserializeStub(stream: StubInputStream): PyTestFixtureDecoratorStub? {
+    val name = stream.readNameString() ?: return null
+    return PyTestFixtureDecoratorStubImpl(name)
+  }
+}
+
+private class PyTestFixtureDecoratorStubImpl(override val name: String) : PyTestFixtureDecoratorStub {
+  override fun getTypeClass(): Class<PyTestFixtureDecoratorStubType> = PyTestFixtureDecoratorStubType::class.java
+
+  override fun serialize(stream: StubOutputStream) {
+    stream.writeName(name)
+  }
+}

python/python-psi-impl/src/com/jetbrains/python/testing/pyTestFixtures/PyTestFixtureUtils.kt

@@ -1,6 +1,10 @@
 // Copyright 2000-2023 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
 package com.jetbrains.python.testing.pyTestFixtures
 
+import com.jetbrains.python.psi.PyDecorator
+import com.jetbrains.python.psi.impl.PyEvaluator
+import org.jetbrains.annotations.ApiStatus
+
 
 const val CONFTEST_PY = "conftest.py"
 const val REQUEST_FIXTURE = "request"
@@ -30,4 +34,13 @@ val reservedFixtureClassSet = setOf(
   "testdir",
   "tmpdir",
   "tmpdir_factory"
-)
+)
+
+@ApiStatus.Internal
+val TEST_FIXTURE_DECORATOR_NAMES = listOf("pytest.fixture", "fixture", "pytest_asyncio.fixture")
+
+@ApiStatus.Internal
+fun getTestFixtureName(decorator: PyDecorator): String? {
+  val nameArgument = decorator.getKeywordArgument("name") ?: return null
+  return PyEvaluator.evaluateNoResolve(nameArgument, String::class.java)
+}

python/src/com/jetbrains/python/testing/pyTestFixtures/PyTestFixture.kt

@@ -1,6 +1,7 @@
 // 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.
 package com.jetbrains.python.testing.pyTestFixtures
 
+import com.intellij.extapi.psi.StubBasedPsiElementBase
 import com.intellij.openapi.module.Module
 import com.intellij.openapi.module.ModuleUtilCore
 import com.intellij.psi.PsiDirectory
@@ -13,8 +14,9 @@ import com.intellij.util.Processor
 import com.intellij.util.ThreeState
 import com.jetbrains.python.extensions.getSdk
 import com.jetbrains.python.psi.*
-import com.jetbrains.python.psi.impl.getNamedArgument
+import com.jetbrains.python.psi.stubs.PyDecoratorStub
 import com.jetbrains.python.psi.stubs.PyDecoratorStubIndex
+import com.jetbrains.python.psi.stubs.PyTestFixtureDecoratorStub
 import com.jetbrains.python.psi.types.TypeEvalContext
 import com.jetbrains.python.sdk.basePath
 import com.jetbrains.python.testing.PyTestFactory
@@ -22,12 +24,10 @@ import com.jetbrains.python.testing.TestRunnerService
 import com.jetbrains.python.testing.autoDetectTests.PyAutoDetectionConfigurationFactory
 import com.jetbrains.python.testing.isTestElement
 
-private val decoratorNames = arrayOf("pytest.fixture", "fixture", "pytest_asyncio.fixture")
-
 private val PyFunction.asFixture: PyTestFixture?
-  get() = decoratorList?.decorators?.firstOrNull { it.name in decoratorNames }?.let { createFixture(it) }
+  get() = decoratorList?.decorators?.firstOrNull { it.name in TEST_FIXTURE_DECORATOR_NAMES }?.let { createFixture(it) }
 
-private fun PyDecoratorList.hasDecorator(vararg names: String) = names.any { findDecorator(it) != null }
+private fun PyDecoratorList.hasDecorator(names: Iterable<String>) = names.any { findDecorator(it) != null }
 
 private fun PyElement.getFixtureName() = name ?: (this as? PyStringLiteralExpression)?.stringValue
 
@@ -301,7 +301,7 @@ fun PyNamedParameter.isFixture(typeEvalContext: TypeEvalContext) = getFixtureLin
 /**
  * @return Boolean is function decorated as fixture or marked so by EP
  */
-internal fun PyFunction.isFixture() = decoratorList?.hasDecorator(*decoratorNames) ?: false || isCustomFixture()
+internal fun PyFunction.isFixture() = decoratorList?.hasDecorator(TEST_FIXTURE_DECORATOR_NAMES) ?: false || isCustomFixture()
 
 
 /**
@@ -312,7 +312,7 @@ internal fun PyFunction.isFixture() = decoratorList?.hasDecorator(*decoratorName
 data class PyTestFixture(val function: PyFunction? = null, val resolveTarget: PyElement? = function, val name: String)
 
 
-fun findDecoratorsByName(module: Module, vararg names: String): Iterable<PyDecorator> =
+fun findDecoratorsByName(module: Module, names: Iterable<String>): Iterable<PyDecorator> =
   names.flatMap { name ->
     StubIndex.getElements(PyDecoratorStubIndex.KEY, name, module.project,
                           GlobalSearchScope.union(
@@ -323,11 +323,18 @@ fun findDecoratorsByName(module: Module, vararg names: String): Iterable<PyDecor
 
 private fun createFixture(decorator: PyDecorator): PyTestFixture? {
   val target = decorator.target ?: return null
-  return (decorator.getNamedArgument("name") ?: target.name)?.let { name ->
+  return (decorator.getTestFixtureName() ?: target.name)?.let { name ->
     PyTestFixture(target, target, name)
   }
 }
 
+private fun PyDecorator.getTestFixtureName(): String? {
+  return when (@Suppress("UNCHECKED_CAST") val stub = (this as StubBasedPsiElementBase<PyDecoratorStub>).greenStub) {
+    null -> getTestFixtureName(this)
+    else -> stub.getCustomStub(PyTestFixtureDecoratorStub::class.java)?.name
+  }
+}
+
 /**
  * Gets list of fixtures suitable for certain function.
  *
@@ -345,7 +352,7 @@ internal fun getFixtures(module: Module, forWhat: PyFunction, typeEvalContext: T
     forWhat.isSubjectForFixture()
   ) {
     //Fixtures
-    (findDecoratorsByName(module, *decoratorNames)
+    (findDecoratorsByName(module, TEST_FIXTURE_DECORATOR_NAMES)
        .filter { it.target?.containingClass == null } //We need only top-level functions, class-based fixtures processed above
        .mapNotNull { createFixture(it) }
      + getCustomFixtures(typeEvalContext, forWhat))
@@ -367,7 +374,7 @@ internal fun getFixtures(module: Module, forWhat: PyFunction, typeEvalContext: T
 
 private fun getModuleFixtures(module: Module): List<PyTestFixture> {
   return if (isPyTestEnabled(module)) {
-    findDecoratorsByName(module, *decoratorNames).mapNotNull { createFixture(it) }
+    findDecoratorsByName(module, TEST_FIXTURE_DECORATOR_NAMES).mapNotNull { createFixture(it) }
   }
   else emptyList()
 }
@@ -381,6 +388,3 @@ internal fun isPyTestEnabled(module: Module): Boolean {
   val factoryId = (if (factory is PyAutoDetectionConfigurationFactory) factory.getFactory(sdk) else factory).id
   return factoryId == PyTestFactory.id
 }
-
-
-