diff --git a/python/psi-api/src/com/jetbrains/python/psi/impl/stubs/PyCustomClassStub.java b/python/psi-api/src/com/jetbrains/python/psi/impl/stubs/PyCustomClassStub.java new file mode 100644 index 000000000000..ef1cb6c57e42 --- /dev/null +++ b/python/psi-api/src/com/jetbrains/python/psi/impl/stubs/PyCustomClassStub.java @@ -0,0 +1,24 @@ +// 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.psi.impl.stubs; + +import com.intellij.psi.stubs.StubInputStream; +import com.intellij.psi.stubs.StubOutputStream; +import org.jetbrains.annotations.NotNull; + +import java.io.IOException; + +public interface PyCustomClassStub { + + /** + * @return type class to distinguish one custom stub from another. + */ + @NotNull + Class getTypeClass(); + + /** + * @param stream stream to serialize {@code this} stub + * @throws IOException + * @see PyCustomClassStubType#deserializeStub(StubInputStream) + */ + void serialize(@NotNull StubOutputStream stream) throws IOException; +} diff --git a/python/psi-api/src/com/jetbrains/python/psi/impl/stubs/PyCustomClassStubType.java b/python/psi-api/src/com/jetbrains/python/psi/impl/stubs/PyCustomClassStubType.java new file mode 100644 index 000000000000..7651e3dc9225 --- /dev/null +++ b/python/psi-api/src/com/jetbrains/python/psi/impl/stubs/PyCustomClassStubType.java @@ -0,0 +1,32 @@ +// 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.psi.impl.stubs; + +import com.intellij.openapi.extensions.ExtensionPointName; +import com.intellij.psi.stubs.StubInputStream; +import com.intellij.psi.stubs.StubOutputStream; +import com.jetbrains.python.psi.PyClass; +import org.jetbrains.annotations.NotNull; +import org.jetbrains.annotations.Nullable; + +import java.io.IOException; + +public abstract class PyCustomClassStubType { + + public static final ExtensionPointName EP_NAME = ExtensionPointName.create("Pythonid.customClassStubType"); + + /** + * @param psi class to create stub for + * @return custom stub for given class or null. + */ + @Nullable + public abstract T createStub(@NotNull PyClass psi); + + /** + * @param stream stream containing serialized stub + * @return a custom stub instance or null if it could not be read. + * @throws IOException + * @see PyCustomClassStub#serialize(StubOutputStream) + */ + @Nullable + public abstract T deserializeStub(@NotNull StubInputStream stream) throws IOException; +} diff --git a/python/psi-api/src/com/jetbrains/python/psi/stubs/PyClassStub.java b/python/psi-api/src/com/jetbrains/python/psi/stubs/PyClassStub.java index 324a27d8960b..5f90dad8e8fe 100644 --- a/python/psi-api/src/com/jetbrains/python/psi/stubs/PyClassStub.java +++ b/python/psi-api/src/com/jetbrains/python/psi/stubs/PyClassStub.java @@ -22,6 +22,7 @@ package com.jetbrains.python.psi.stubs; import com.intellij.psi.stubs.NamedStub; import com.intellij.psi.util.QualifiedName; import com.jetbrains.python.psi.PyClass; +import com.jetbrains.python.psi.impl.stubs.PyCustomClassStub; import org.jetbrains.annotations.NotNull; import org.jetbrains.annotations.Nullable; @@ -66,4 +67,17 @@ public interface PyClassStub extends NamedStub { default List getSuperClassesText() { return Collections.emptyList(); } + + /** + * @param stubClass custom stub class + * @param custom stub class + * @return saved stub if it is an instance of passed class or null. + * @apiNote This method will be marked as abstract in 2019.2. + * @see PyCustomClassStub + * @see com.jetbrains.python.psi.impl.stubs.PyCustomClassStubType + */ + @Nullable + default T getCustomStub(@NotNull Class stubClass) { + return null; + } } \ No newline at end of file diff --git a/python/src/META-INF/python-core-common.xml b/python/src/META-INF/python-core-common.xml index 042ea9f9cab9..61acc5c4e22a 100644 --- a/python/src/META-INF/python-core-common.xml +++ b/python/src/META-INF/python-core-common.xml @@ -692,6 +692,8 @@ + diff --git a/python/src/com/jetbrains/python/psi/PyFileElementType.java b/python/src/com/jetbrains/python/psi/PyFileElementType.java index 20d43146441e..7f4fef3c47f3 100644 --- a/python/src/com/jetbrains/python/psi/PyFileElementType.java +++ b/python/src/com/jetbrains/python/psi/PyFileElementType.java @@ -62,7 +62,7 @@ public class PyFileElementType extends IStubFileElementType { @Override public int getStubVersion() { // Don't forget to update versions of indexes that use the updated stub-based elements - return 69; + return 70; } @Nullable diff --git a/python/src/com/jetbrains/python/psi/impl/stubs/PyClassElementType.java b/python/src/com/jetbrains/python/psi/impl/stubs/PyClassElementType.java index bf1fd04f4d5a..fcb14acf4a97 100644 --- a/python/src/com/jetbrains/python/psi/impl/stubs/PyClassElementType.java +++ b/python/src/com/jetbrains/python/psi/impl/stubs/PyClassElementType.java @@ -12,6 +12,7 @@ import com.jetbrains.python.psi.impl.PyClassImpl; import com.jetbrains.python.psi.impl.PyPsiUtils; import com.jetbrains.python.psi.resolve.PyResolveUtil; import com.jetbrains.python.psi.stubs.*; +import one.util.streamex.StreamEx; import org.jetbrains.annotations.NonNls; import org.jetbrains.annotations.NotNull; import org.jetbrains.annotations.Nullable; @@ -26,6 +27,9 @@ import static com.jetbrains.python.psi.PyUtil.as; */ public class PyClassElementType extends PyStubElementType { + @Nullable + private List myCustomStubTypes; + public PyClassElementType() { this("CLASS_DECLARATION"); } @@ -48,6 +52,12 @@ public class PyClassElementType extends PyStubElementType @Override @NotNull public PyClassStub createStub(@NotNull final PyClass psi, final StubElement parentStub) { + final PyCustomClassStub customStub = StreamEx + .of(getCustomStubTypes()) + .map(type -> type.createStub(psi)) + .findFirst(Objects::nonNull) + .orElse(null); + return new PyClassStubImpl(psi.getName(), parentStub, getSuperClassQNames(psi), @@ -56,7 +66,8 @@ public class PyClassElementType extends PyStubElementType PyPsiUtils.asQualifiedName(psi.getMetaClassExpression()), psi.getOwnSlots(), PyPsiUtils.strValue(psi.getDocStringExpression()), - getStubElementType()); + getStubElementType(), + customStub); } @NotNull @@ -148,6 +159,20 @@ public class PyClassElementType extends PyStubElementType final String docString = pyClassStub.getDocString(); dataStream.writeUTFFast(docString != null ? docString : ""); + + serializeCustomStub(pyClassStub, dataStream); + } + + private static void serializeCustomStub(@NotNull PyClassStub stub, @NotNull StubOutputStream stream) throws IOException { + final PyCustomClassStub customStub = stub.getCustomStub(PyCustomClassStub.class); + + final boolean hasCustomStub = customStub != null; + stream.writeBoolean(hasCustomStub); + + if (hasCustomStub) { + stream.writeName(customStub.getTypeClass().getCanonicalName()); + customStub.serialize(stream); + } } @Override @@ -181,8 +206,25 @@ public class PyClassElementType extends PyStubElementType final String docStringInStub = dataStream.readUTFFast(); final String docString = docStringInStub.length() > 0 ? docStringInStub : null; + final PyCustomClassStub customStub = deserializeCustomStub(dataStream); + return new PyClassStubImpl(name, parentStub, superClasses, parametrizedBaseClasses, baseClassesText, metaClass, slots, docString, - getStubElementType()); + getStubElementType(), customStub); + } + + @Nullable + private PyCustomClassStub deserializeCustomStub(@NotNull StubInputStream stream) throws IOException { + if (stream.readBoolean()) { + final String typeName = stream.readNameString(); + for (PyCustomClassStubType type : getCustomStubTypes()) { + if (type.getClass().getCanonicalName().equals(typeName)) { + return type.deserializeStub(stream); + } + } + throw new IOException("Unknown custom class stub type " + typeName); + } + + return null; } @Override @@ -209,4 +251,12 @@ public class PyClassElementType extends PyStubElementType protected IStubElementType getStubElementType() { return PyElementTypes.CLASS_DECLARATION; } + + @NotNull + private List getCustomStubTypes() { + if (myCustomStubTypes == null) { + myCustomStubTypes = PyCustomClassStubType.EP_NAME.getExtensionList(); + } + return myCustomStubTypes; + } } \ No newline at end of file diff --git a/python/src/com/jetbrains/python/psi/impl/stubs/PyClassStubImpl.java b/python/src/com/jetbrains/python/psi/impl/stubs/PyClassStubImpl.java index 30e8e72f2aed..0109bf9e6f4e 100644 --- a/python/src/com/jetbrains/python/psi/impl/stubs/PyClassStubImpl.java +++ b/python/src/com/jetbrains/python/psi/impl/stubs/PyClassStubImpl.java @@ -5,6 +5,7 @@ import com.intellij.psi.stubs.IStubElementType; import com.intellij.psi.stubs.StubBase; import com.intellij.psi.stubs.StubElement; import com.intellij.psi.util.QualifiedName; +import com.intellij.util.ObjectUtils; import com.jetbrains.python.psi.PyClass; import com.jetbrains.python.psi.stubs.PyClassStub; import org.jetbrains.annotations.NotNull; @@ -38,8 +39,11 @@ public class PyClassStubImpl extends StubBase implements PyClassStub { @NotNull private final List mySuperClassesText; + @Nullable + private final PyCustomClassStub myCustomStub; + /** - * @deprecated Use {@link PyClassStubImpl#PyClassStubImpl(String, StubElement, Map, List, List, QualifiedName, List, String, IStubElementType)} instead. + * @deprecated Use {@link PyClassStubImpl#PyClassStubImpl(String, StubElement, Map, List, QualifiedName, List, String, IStubElementType, PyCustomClassStub)} instead. * This constructor will be removed in 2019.2. */ @Deprecated @@ -51,9 +55,14 @@ public class PyClassStubImpl extends StubBase implements PyClassStub { @Nullable List slots, @Nullable String docString, @NotNull IStubElementType stubElementType) { - this(name, parentStub, superClasses, subscriptedSuperClassesText, Collections.emptyList(), metaClass, slots, docString, stubElementType); + this(name, parentStub, superClasses, subscriptedSuperClassesText, Collections.emptyList(), metaClass, slots, docString, stubElementType, null); } + /** + * @deprecated Use {@link PyClassStubImpl#PyClassStubImpl(String, StubElement, Map, List, QualifiedName, List, String, IStubElementType, PyCustomClassStub)} instead. + * This constructor will be removed in 2019.2. + */ + @Deprecated public PyClassStubImpl(@Nullable String name, @Nullable StubElement parentStub, @NotNull Map superClasses, @@ -62,7 +71,8 @@ public class PyClassStubImpl extends StubBase implements PyClassStub { @Nullable QualifiedName metaClass, @Nullable List slots, @Nullable String docString, - @NotNull IStubElementType stubElementType) { + @NotNull IStubElementType stubElementType, + @Nullable PyCustomClassStub customStub) { super(parentStub, stubElementType); myName = name; mySuperClasses = superClasses; @@ -71,6 +81,19 @@ public class PyClassStubImpl extends StubBase implements PyClassStub { myMetaClass = metaClass; mySlots = slots; myDocString = docString; + myCustomStub = customStub; + } + + public PyClassStubImpl(@Nullable String name, + @Nullable StubElement parentStub, + @NotNull Map superClasses, + @NotNull List superClassesText, + @Nullable QualifiedName metaClass, + @Nullable List slots, + @Nullable String docString, + @NotNull IStubElementType stubElementType, + @Nullable PyCustomClassStub customStub) { + this(name, parentStub, superClasses, Collections.emptyList(), superClassesText, metaClass, slots, docString, stubElementType, customStub); } @Override @@ -115,6 +138,12 @@ public class PyClassStubImpl extends StubBase implements PyClassStub { return mySuperClassesText; } + @Nullable + @Override + public T getCustomStub(@NotNull Class stubClass) { + return ObjectUtils.tryCast(myCustomStub, stubClass); + } + @Override public String toString() { return "PyClassStub(" + myName + ")";