diff --git a/python/src/com/jetbrains/python/psi/impl/PyClassImpl.java b/python/src/com/jetbrains/python/psi/impl/PyClassImpl.java index b06f3997a29c..9e3ddd802bc0 100644 --- a/python/src/com/jetbrains/python/psi/impl/PyClassImpl.java +++ b/python/src/com/jetbrains/python/psi/impl/PyClassImpl.java @@ -190,38 +190,11 @@ public class PyClassImpl extends PyBaseElementImpl implements PyCla public static boolean isSixWithMetaclassCall(@NotNull PyExpression expression) { if (expression instanceof PyCallExpression){ - final PyCallExpression call = (PyCallExpression)expression; - final PyExpression callee = call.getCallee(); - if (callee != null && "with_metaclass".equals(callee.getName())) { - // SUPPORTED CASES: - - // import six - // six.with_metaclass(...) - - // from six import metaclass - // with_metaclass(...) - return true; - } + final PyExpression callee = ((PyCallExpression)expression).getCallee(); if (callee instanceof PyReferenceExpression) { - // SUPPORTED CASES: - - // from six import with_metaclass as w_m - // w_m(...) - - final boolean importedWithMetaclass = StreamEx - .of(PyResolveUtil.resolveLocally((PyReferenceExpression)callee)) - .select(PyImportElement.class) - .map(PyImportElement::getImportedQName) - .nonNull() - .map(QualifiedName::getLastComponent) - .nonNull() - .findAny("with_metaclass"::equals) - .isPresent(); - - if (importedWithMetaclass) { - return true; - } + final QualifiedName sixWithMetaclass = QualifiedName.fromComponents("six", "with_metaclass"); + return PyResolveUtil.resolveImportedElementQNameLocally((PyReferenceExpression)callee).contains(sixWithMetaclass); } } @@ -1638,36 +1611,10 @@ public class PyClassImpl extends PyBaseElementImpl implements PyCla private static boolean isSixAddMetaclass(@NotNull PyDecorator decorator) { final PyExpression callee = decorator.getCallee(); - if (callee != null && "add_metaclass".equals(callee.getName())) { - // SUPPORTED CASES: - - // import six - // six.add_metaclass(...) - - // from six import add_metaclass - // add_metaclass(...) - return true; - } if (callee instanceof PyReferenceExpression) { - // SUPPORTED CASES: - - // from six import add_metaclass as a_m - // a_m(...) - - final boolean importedAddMetaclass = StreamEx - .of(PyResolveUtil.resolveLocally((PyReferenceExpression)callee)) - .select(PyImportElement.class) - .map(PyImportElement::getImportedQName) - .nonNull() - .map(QualifiedName::getLastComponent) - .nonNull() - .findAny("add_metaclass"::equals) - .isPresent(); - - if (importedAddMetaclass) { - return true; - } + final QualifiedName sixAddMetaclass = QualifiedName.fromComponents("six", "add_metaclass"); + return PyResolveUtil.resolveImportedElementQNameLocally((PyReferenceExpression)callee).contains(sixAddMetaclass); } return false; diff --git a/python/src/com/jetbrains/python/psi/impl/stubs/PyDataclassFieldStubImpl.kt b/python/src/com/jetbrains/python/psi/impl/stubs/PyDataclassFieldStubImpl.kt index 327009ced15a..8de3bd81762b 100644 --- a/python/src/com/jetbrains/python/psi/impl/stubs/PyDataclassFieldStubImpl.kt +++ b/python/src/com/jetbrains/python/psi/impl/stubs/PyDataclassFieldStubImpl.kt @@ -6,8 +6,11 @@ package com.jetbrains.python.psi.impl.stubs import com.intellij.psi.stubs.StubInputStream import com.intellij.psi.stubs.StubOutputStream import com.intellij.psi.util.QualifiedName -import com.jetbrains.python.psi.* +import com.jetbrains.python.psi.PyCallExpression +import com.jetbrains.python.psi.PyReferenceExpression +import com.jetbrains.python.psi.PyTargetExpression import com.jetbrains.python.psi.impl.PyEvaluator +import com.jetbrains.python.psi.impl.PyPsiUtils import com.jetbrains.python.psi.resolve.PyResolveUtil import com.jetbrains.python.psi.stubs.PyDataclassFieldStub import java.io.IOException @@ -21,7 +24,7 @@ class PyDataclassFieldStubImpl private constructor(private val calleeName: Quali val value = expression.findAssignedValue() as? PyCallExpression ?: return null val callee = value.callee as? PyReferenceExpression ?: return null - val calleeName = calculateFullyQCalleeName(callee) ?: calculateImportedCalleeName(callee) ?: return null + val calleeName = calculateCalleeName(callee) ?: return null val arguments = analyzeArguments(value) return PyDataclassFieldStubImpl(calleeName, arguments.first, arguments.second, arguments.third) @@ -37,41 +40,9 @@ class PyDataclassFieldStubImpl private constructor(private val calleeName: Quali return PyDataclassFieldStubImpl(QualifiedName.fromDottedString(calleeName), hasDefault, hasDefaultFactory, initValue) } - private fun calculateFullyQCalleeName(callee: PyReferenceExpression): QualifiedName? { - // SUPPORTED CASES: - - // import dataclasses - // ... = dataclasses.field(...) - - // import dataclasses as dc - // ... = dc.field(...) - - val calleeName = callee.name - val qualifier = callee.qualifier - - if (calleeName == "field" && qualifier is PyReferenceExpression && !qualifier.isQualified && resolvesToDataclassesModule(qualifier)) { - return QualifiedName.fromComponents(qualifier.name, calleeName) - } - - return null - } - - private fun calculateImportedCalleeName(callee: PyReferenceExpression): QualifiedName? { - // SUPPORTED CASES: - - // from dataclasses import field - // ... = field(...) - - // from dataclasses import field as F - // ... = F(...) - - for (element in PyResolveUtil.resolveLocally(callee)) { - if (element is PyImportElement && element.importedQName.toString() == "field") { - val importStatement = element.containingImportStatement - if (importStatement is PyFromImportStatement && importStatement.importSourceQName.toString() == "dataclasses") { - return QualifiedName.fromComponents(callee.name) - } - } + private fun calculateCalleeName(callee: PyReferenceExpression): QualifiedName? { + if (QualifiedName.fromComponents("dataclasses", "field") in PyResolveUtil.resolveImportedElementQNameLocally(callee)) { + return PyPsiUtils.asQualifiedName(callee) } return null @@ -84,10 +55,6 @@ class PyDataclassFieldStubImpl private constructor(private val calleeName: Quali return Triple(hasDefault, hasDefaultFactory, initValue) } - - private fun resolvesToDataclassesModule(referenceExpression: PyReferenceExpression): Boolean { - return PyResolveUtil.resolveLocally(referenceExpression).any { it is PyImportElement && it.importedQName.toString() == "dataclasses" } - } } override fun getTypeClass(): Class> { diff --git a/python/src/com/jetbrains/python/psi/impl/stubs/PyNamedTupleStubImpl.java b/python/src/com/jetbrains/python/psi/impl/stubs/PyNamedTupleStubImpl.java index ca0f58d8ccb6..e9a97ad80586 100644 --- a/python/src/com/jetbrains/python/psi/impl/stubs/PyNamedTupleStubImpl.java +++ b/python/src/com/jetbrains/python/psi/impl/stubs/PyNamedTupleStubImpl.java @@ -22,6 +22,7 @@ import com.intellij.psi.stubs.StubInputStream; import com.intellij.psi.stubs.StubOutputStream; import com.intellij.psi.util.QualifiedName; import com.intellij.util.ArrayUtil; +import com.intellij.util.containers.ContainerUtil; import com.jetbrains.python.PyNames; import com.jetbrains.python.codeInsight.typing.PyTypingTypeProvider; import com.jetbrains.python.psi.*; @@ -97,19 +98,15 @@ public class PyNamedTupleStubImpl implements PyNamedTupleStub { @Nullable public static PyNamedTupleStub deserialize(@NotNull StubInputStream stream) throws IOException { - String calleeName = stream.readNameString(); - String name = stream.readNameString(); + final String calleeName = stream.readNameString(); + final String name = stream.readNameString(); final LinkedHashMap> fields = deserializeFields(stream, stream.readVarInt()); if (calleeName == null || name == null) { return null; } - return new PyNamedTupleStubImpl( - QualifiedName.fromDottedString(calleeName), - name, - fields - ); + return new PyNamedTupleStubImpl(QualifiedName.fromDottedString(calleeName), name, fields); } @NotNull @@ -150,13 +147,19 @@ public class PyNamedTupleStubImpl implements PyNamedTupleStub { @Nullable private static Pair getCalleeNameAndNTModule(@NotNull PyReferenceExpression referenceExpression) { - final Pair name = getFullyQCalleeNameAndNTModule(referenceExpression); + final QualifiedName calleeName = PyPsiUtils.asQualifiedName(referenceExpression); + if (calleeName == null) return null; - if (name != null) { - return name; + for (String name : ContainerUtil.map(PyResolveUtil.resolveImportedElementQNameLocally(referenceExpression), QualifiedName::toString)) { + if (name.equals(PyNames.COLLECTIONS_NAMEDTUPLE_PY2)) { + return Pair.createNonNull(calleeName, NamedTupleModule.COLLECTIONS); + } + else if (name.equals(PyTypingTypeProvider.NAMEDTUPLE)) { + return Pair.createNonNull(calleeName, NamedTupleModule.TYPING); + } } - return getImportedCalleeNameAndNTModule(referenceExpression); + return null; } @Nullable @@ -187,106 +190,6 @@ public class PyNamedTupleStubImpl implements PyNamedTupleStub { return fields; } - @Nullable - private static Pair getFullyQCalleeNameAndNTModule(@NotNull PyReferenceExpression referenceExpression) { - // SUPPORTED CASES: - - // import collections - // Point = collections.namedtuple(...) - - // import collections as c - // Point = c.namedtuple(...) - - // import typing - // ... = typing.NamedTuple(...) - - // import typing as t - // ... = t.NamedTuple(...) - - final String referenceName = referenceExpression.getName(); - final NamedTupleModule module = PyNames.NAMEDTUPLE.equals(referenceName) - ? NamedTupleModule.COLLECTIONS - : PyTypingTypeProvider.NAMEDTUPLE_SIMPLE.equals(referenceName) - ? NamedTupleModule.TYPING - : null; - - if (module != null) { - final PyExpression qualifier = referenceExpression.getQualifier(); - - if (qualifier instanceof PyReferenceExpression) { - final PyReferenceExpression qualifierReference = (PyReferenceExpression)qualifier; - - if (!qualifierReference.isQualified() && resolvesToModule(qualifierReference, module)) { - return Pair.createNonNull(QualifiedName.fromComponents(qualifierReference.getName(), referenceName), module); - } - } - } - - return null; - } - - @Nullable - private static Pair getImportedCalleeNameAndNTModule(@NotNull PyReferenceExpression referenceExpression) { - // SUPPORTED CASES: - - // from collections import namedtuple - // Point = namedtuple(...) - - // from collections import namedtuple as NT - // Point = NT(...) - - // from typing import NamedTuple - // Point = NamedTuple(...) - - // from typing import NamedTuple as NT - // Point = NT(...) - - for (PsiElement element : PyResolveUtil.resolveLocally(referenceExpression)) { - if (element instanceof PyImportElement) { - final PyImportElement importElement = (PyImportElement)element; - final QualifiedName importedQName = importElement.getImportedQName(); - - final NamedTupleModule module = equals(importedQName, PyNames.NAMEDTUPLE) - ? NamedTupleModule.COLLECTIONS - : equals(importedQName, PyTypingTypeProvider.NAMEDTUPLE_SIMPLE) - ? NamedTupleModule.TYPING - : null; - - if (module != null) { - final PyStatement importStatement = importElement.getContainingImportStatement(); - - if (importStatement instanceof PyFromImportStatement) { - final PyFromImportStatement fromImportStatement = (PyFromImportStatement)importStatement; - - if (equals(fromImportStatement.getImportSourceQName(), module.getModuleName())) { - return Pair.createNonNull(QualifiedName.fromComponents(referenceExpression.getName()), module); - } - } - } - } - } - - return null; - } - - private static boolean resolvesToModule(@NotNull PyReferenceExpression referenceExpression, @NotNull NamedTupleModule module) { - for (PsiElement element : PyResolveUtil.resolveLocally(referenceExpression)) { - if (element instanceof PyImportElement) { - final PyImportElement importElement = (PyImportElement)element; - - if (equals(importElement.getImportedQName(), module.getModuleName())) { - return true; - } - } - } - - return false; - } - - private static boolean equals(@Nullable QualifiedName qualifiedName, @NotNull String name) { - return qualifiedName != null && name.equals(qualifiedName.toString()); - } - @Nullable private static LinkedHashMap> resolveCollectionsNTFields(@NotNull PyCallExpression callExpression) { // SUPPORTED CASES: diff --git a/python/src/com/jetbrains/python/psi/resolve/PyResolveUtil.java b/python/src/com/jetbrains/python/psi/resolve/PyResolveUtil.java index 7d300cccbf1e..230f8fca4f8b 100644 --- a/python/src/com/jetbrains/python/psi/resolve/PyResolveUtil.java +++ b/python/src/com/jetbrains/python/psi/resolve/PyResolveUtil.java @@ -22,6 +22,7 @@ import com.intellij.psi.ResolveState; import com.intellij.psi.scope.PsiScopeProcessor; import com.intellij.psi.util.PsiTreeUtil; import com.intellij.psi.util.QualifiedName; +import com.intellij.util.containers.ContainerUtil; import com.jetbrains.python.codeInsight.controlflow.ControlFlowCache; import com.jetbrains.python.codeInsight.controlflow.ScopeOwner; import com.jetbrains.python.codeInsight.dataflow.scope.Scope; @@ -141,6 +142,52 @@ public class PyResolveUtil { return processor.getElements(); } + @NotNull + public static List resolveImportedElementQNameLocally(@NotNull PyReferenceExpression expression) { + // SUPPORTED CASES: + + // import six + // six.with_metaclass(...) + + // from six import metaclass + // with_metaclass(...) + + // from six import with_metaclass as w_m + // w_m(...) + + final PyExpression qualifier = expression.getQualifier(); + if (qualifier instanceof PyReferenceExpression) { + final String name = expression.getName(); + + return name == null + ? Collections.emptyList() + : ContainerUtil.map(resolveImportedElementQNameLocally((PyReferenceExpression)qualifier), qn -> qn.append(name)); + } + else { + return StreamEx + .of(resolveLocally(expression)) + .select(PyImportElement.class) + .map( + element -> { + final PyStatement importStatement = element.getContainingImportStatement(); + + if (importStatement instanceof PyFromImportStatement) { + final QualifiedName importSourceQName = ((PyFromImportStatement)importStatement).getImportSourceQName(); + final QualifiedName importedQName = element.getImportedQName(); + + if (importSourceQName != null && importedQName != null) { + return importSourceQName.append(importedQName); + } + } + + return element.getImportedQName(); + } + ) + .nonNull() + .toList(); + } + } + /** * Resolve a symbol by its qualified name, starting from the specified file and then following the chain of type members. * This type of resolve is stub-safe, i.e. it's not supposed to cause any un-stubbing of external files unless it explicitly diff --git a/python/src/com/jetbrains/python/psi/types/PyTypingNewType.kt b/python/src/com/jetbrains/python/psi/types/PyTypingNewType.kt index 0738e2ef9a11..2d8cc1e85cda 100644 --- a/python/src/com/jetbrains/python/psi/types/PyTypingNewType.kt +++ b/python/src/com/jetbrains/python/psi/types/PyTypingNewType.kt @@ -3,7 +3,9 @@ package com.jetbrains.python.psi.types import com.intellij.psi.util.QualifiedName import com.jetbrains.python.codeInsight.typing.PyTypingTypeProvider -import com.jetbrains.python.psi.* +import com.jetbrains.python.psi.PyCallExpression +import com.jetbrains.python.psi.PyCallSiteExpression +import com.jetbrains.python.psi.PyReferenceExpression import com.jetbrains.python.psi.resolve.PyResolveUtil data class PyTypingNewType(internal val classType: PyClassType, internal val isDefinition: Boolean, internal val myName: String?) : PyClassTypeImpl( @@ -27,7 +29,7 @@ data class PyTypingNewType(internal val classType: PyClassType, internal val isD override fun isCallable() = classType.isCallable || isDefinition - override fun toString() = "TypingNewType: " + myName + override fun toString() = "TypingNewType: $myName" override fun getParameters(context: TypeEvalContext): List? { return if (isCallable) { @@ -45,27 +47,9 @@ data class PyTypingNewType(internal val classType: PyClassType, internal val isD } companion object { - private fun getImportedQualifiedName(referenceExpression: PyReferenceExpression): QualifiedName? { - val qualifier = referenceExpression.qualifier - if (qualifier is PyReferenceExpression) { - PyResolveUtil.resolveLocally(qualifier) - .filterIsInstance() - .firstOrNull { return it.importedQName?.append(referenceExpression.name) } - } - for (element in PyResolveUtil.resolveLocally(referenceExpression)) { - if (element is PyImportElement) { - val importStatement = element.containingImportStatement - if (importStatement is PyFromImportStatement) { - return importStatement.importSourceQName?.append(element.importedQName) - } - } - } - return null - } - fun isTypingNewType(callExpression: PyCallExpression): Boolean { - val calleeReference = callExpression.callee as? PyReferenceExpression ?: return false - return getImportedQualifiedName(calleeReference) == QualifiedName.fromDottedString(PyTypingTypeProvider.NEW_TYPE) + val callee = callExpression.callee as? PyReferenceExpression ?: return false + return QualifiedName.fromDottedString(PyTypingTypeProvider.NEW_TYPE) in PyResolveUtil.resolveImportedElementQNameLocally(callee) } } }