2 * Copyright 2000-2014 JetBrains s.r.o.
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
8 * http://www.apache.org/licenses/LICENSE-2.0
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
16 package com.jetbrains.python.psi.impl;
18 import com.intellij.lang.ASTNode;
19 import com.intellij.navigation.ItemPresentation;
20 import com.intellij.openapi.extensions.Extensions;
21 import com.intellij.openapi.util.Pair;
22 import com.intellij.psi.PsiElement;
23 import com.intellij.psi.PsiPolyVariantReference;
24 import com.intellij.psi.PsiReference;
25 import com.intellij.psi.search.GlobalSearchScope;
26 import com.intellij.psi.search.LocalSearchScope;
27 import com.intellij.psi.search.SearchScope;
28 import com.intellij.psi.stubs.IStubElementType;
29 import com.intellij.psi.stubs.StubElement;
30 import com.intellij.psi.util.PsiTreeUtil;
31 import com.intellij.psi.util.QualifiedName;
32 import com.intellij.util.IncorrectOperationException;
33 import com.intellij.util.PlatformIcons;
34 import com.jetbrains.python.PyElementTypes;
35 import com.jetbrains.python.PyNames;
36 import com.jetbrains.python.PyTokenTypes;
37 import com.jetbrains.python.PythonDialectsTokenSetProvider;
38 import com.jetbrains.python.codeInsight.controlflow.ControlFlowCache;
39 import com.jetbrains.python.codeInsight.controlflow.ScopeOwner;
40 import com.jetbrains.python.codeInsight.dataflow.scope.Scope;
41 import com.jetbrains.python.codeInsight.dataflow.scope.ScopeUtil;
42 import com.jetbrains.python.documentation.docstrings.DocStringUtil;
43 import com.jetbrains.python.psi.*;
44 import com.jetbrains.python.psi.impl.references.PyQualifiedReference;
45 import com.jetbrains.python.psi.impl.references.PyTargetReference;
46 import com.jetbrains.python.psi.impl.stubs.CustomTargetExpressionStub;
47 import com.jetbrains.python.psi.resolve.PyResolveContext;
48 import com.jetbrains.python.psi.resolve.QualifiedNameFinder;
49 import com.jetbrains.python.psi.resolve.RatedResolveResult;
50 import com.jetbrains.python.psi.stubs.PyClassStub;
51 import com.jetbrains.python.psi.stubs.PyFunctionStub;
52 import com.jetbrains.python.psi.stubs.PyTargetExpressionStub;
53 import com.jetbrains.python.psi.types.*;
54 import org.jetbrains.annotations.NotNull;
55 import org.jetbrains.annotations.Nullable;
58 import java.util.ArrayList;
59 import java.util.Collection;
60 import java.util.Collections;
61 import java.util.List;
66 public class PyTargetExpressionImpl extends PyBaseElementImpl<PyTargetExpressionStub> implements PyTargetExpression {
67 QualifiedName myQualifiedName;
69 public PyTargetExpressionImpl(ASTNode astNode) {
73 public PyTargetExpressionImpl(final PyTargetExpressionStub stub) {
74 super(stub, PyElementTypes.TARGET_EXPRESSION);
77 public PyTargetExpressionImpl(final PyTargetExpressionStub stub, IStubElementType nodeType) {
78 super(stub, nodeType);
82 protected void acceptPyVisitor(PyElementVisitor pyVisitor) {
83 pyVisitor.visitPyTargetExpression(this);
88 public String getName() {
89 final PyTargetExpressionStub stub = getStub();
91 return stub.getName();
93 ASTNode node = getNameElement();
94 return node != null ? node.getText() : null;
98 public int getTextOffset() {
99 final ASTNode nameElement = getNameElement();
100 return nameElement != null ? nameElement.getStartOffset() : getTextRange().getStartOffset();
104 public ASTNode getNameElement() {
105 return getNode().findChildByType(PyTokenTypes.IDENTIFIER);
108 public PsiElement getNameIdentifier() {
109 final ASTNode nameElement = getNameElement();
110 return nameElement == null ? null : nameElement.getPsi();
113 public String getReferencedName() {
117 public PsiElement setName(@NotNull String name) throws IncorrectOperationException {
118 final ASTNode oldNameElement = getNameElement();
119 if (oldNameElement != null) {
120 final ASTNode nameElement = PyUtil.createNewName(this, name);
121 getNode().replaceChild(oldNameElement, nameElement);
126 public PyType getType(@NotNull TypeEvalContext context, @NotNull TypeEvalContext.Key key) {
127 if (!TypeEvalStack.mayEvaluate(this)) {
131 if (PyNames.ALL.equals(getName())) {
132 // no type for __all__, to avoid unresolved reference errors for expressions where a qualifier is a name
133 // imported via __all__
136 final PyType pyType = PyReferenceExpressionImpl.getReferenceTypeFromProviders(this, context, null);
137 if (pyType != null) {
140 PyType type = getTypeFromDocString();
144 if (!context.maySwitchToAST(this)) {
145 final PsiElement value = resolveAssignedValue(PyResolveContext.noImplicits().withTypeEvalContext(context));
146 if (value instanceof PyTypedElement) {
147 type = context.getType((PyTypedElement)value);
148 if (type instanceof PyNoneType) {
155 type = getTypeFromComment(this);
159 final PsiElement parent = getParent();
160 if (parent instanceof PyAssignmentStatement) {
161 final PyAssignmentStatement assignmentStatement = (PyAssignmentStatement)parent;
162 PyExpression assignedValue = assignmentStatement.getAssignedValue();
163 if (assignedValue instanceof PyParenthesizedExpression) {
164 assignedValue = ((PyParenthesizedExpression)assignedValue).getContainedExpression();
166 if (assignedValue != null) {
167 if (assignedValue instanceof PyYieldExpression) {
170 return context.getType(assignedValue);
173 if (parent instanceof PyTupleExpression) {
174 PsiElement nextParent = parent.getParent();
175 while (nextParent instanceof PyParenthesizedExpression) {
176 nextParent = nextParent.getParent();
178 if (nextParent instanceof PyAssignmentStatement) {
179 final PyAssignmentStatement assignment = (PyAssignmentStatement)nextParent;
180 final PyExpression value = assignment.getAssignedValue();
182 final PyType assignedType = PyTypeChecker.toNonWeakType(context.getType(value), context);
183 if (assignedType instanceof PyTupleType) {
184 final PyType t = PyTypeChecker.getTargetTypeFromTupleAssignment(this, (PyTupleExpression)parent, (PyTupleType)assignedType);
192 if (parent instanceof PyWithItem) {
193 return getWithItemVariableType(context, (PyWithItem)parent);
195 PyType iterType = getTypeFromIteration(context);
196 if (iterType != null) {
199 PyType excType = getTypeFromExcept();
200 if (excType != null) {
206 TypeEvalStack.evaluated(this);
211 private static PyType getWithItemVariableType(TypeEvalContext context, PyWithItem item) {
212 final PyExpression expression = item.getExpression();
213 if (expression != null) {
214 final PyType exprType = context.getType(expression);
215 if (exprType instanceof PyClassType) {
216 final PyClass cls = ((PyClassType)exprType).getPyClass();
217 final PyFunction enter = cls.findMethodByName(PyNames.ENTER, true, null);
219 final PyType enterType = enter.getCallType(expression, Collections.<PyExpression, PyNamedParameter>emptyMap(), context);
220 if (enterType != null) {
223 for (PyTypeProvider provider: Extensions.getExtensions(PyTypeProvider.EP_NAME)) {
224 PyType typeFromProvider = provider.getContextManagerVariableType(cls, expression, context);
225 if (typeFromProvider != null) {
226 return typeFromProvider;
229 // Guess the return type of __enter__
230 return PyUnionType.createWeakType(exprType);
238 public PyType getTypeFromDocString() {
239 String typeName = null;
240 final String name = getName();
241 final StructuredDocString targetDocString = getStructuredDocString();
242 if (targetDocString != null) {
243 typeName = targetDocString.getParamType(null);
244 if (typeName == null) {
245 typeName = targetDocString.getParamType(name);
248 if (typeName == null && PyUtil.isAttribute(this)) {
249 final PyClass cls = getContainingClass();
251 final StructuredDocString classDocString = cls.getStructuredDocString();
252 if (classDocString != null) {
253 typeName = classDocString.getParamType(name);
257 if (typeName != null) {
258 return PyTypeParser.getTypeByName(this, typeName);
264 public static PyType getTypeFromComment(PyTargetExpressionImpl targetExpression) {
265 String docComment = DocStringUtil.getAttributeDocComment(targetExpression);
266 if (docComment != null) {
267 StructuredDocString structuredDocString = DocStringUtil.parse(docComment, targetExpression);
268 String typeName = structuredDocString.getParamType(null);
269 if (typeName == null) {
270 typeName = structuredDocString.getParamType(targetExpression.getName());
272 if (typeName != null) {
273 return PyTypeParser.getTypeByName(targetExpression, typeName);
280 private PyType getTypeFromIteration(@NotNull TypeEvalContext context) {
281 PyExpression target = null;
282 PyExpression source = null;
283 final PyForPart forPart = PsiTreeUtil.getParentOfType(this, PyForPart.class);
284 if (forPart != null) {
285 final PyExpression expr = forPart.getTarget();
286 if (PsiTreeUtil.isAncestor(expr, this, false)) {
288 source = forPart.getSource();
291 final PyComprehensionElement comprh = PsiTreeUtil.getParentOfType(this, PyComprehensionElement.class);
292 if (comprh != null) {
293 for (ComprhForComponent c : comprh.getForComponents()) {
294 final PyExpression expr = c.getIteratorVariable();
295 if (PsiTreeUtil.isAncestor(expr, this, false)) {
297 source = c.getIteratedList();
301 if (source != null) {
302 final PyType sourceType = context.getType(source);
303 final PyType type = getIterationType(sourceType, source, this, context);
304 if (type instanceof PyTupleType && target instanceof PyTupleExpression) {
305 return PyTypeChecker.getTargetTypeFromTupleAssignment(this, (PyTupleExpression)target, (PyTupleType)type);
307 if (target == this && type != null) {
315 private static PyType getIterationType(@Nullable PyType iterableType, @Nullable PyExpression source, @NotNull PsiElement anchor,
316 @NotNull TypeEvalContext context) {
317 if (iterableType instanceof PyTupleType) {
318 final PyTupleType tupleType = (PyTupleType)iterableType;
319 final List<PyType> memberTypes = new ArrayList<PyType>();
320 for (int i = 0; i < tupleType.getElementCount(); i++) {
321 memberTypes.add(tupleType.getElementType(i));
323 return PyUnionType.union(memberTypes);
325 else if (iterableType instanceof PyUnionType) {
326 final Collection<PyType> members = ((PyUnionType)iterableType).getMembers();
327 final List<PyType> iterationTypes = new ArrayList<PyType>();
328 for (PyType member : members) {
329 iterationTypes.add(getIterationType(member, source, anchor, context));
331 return PyUnionType.union(iterationTypes);
333 else if (iterableType != null && PyABCUtil.isSubtype(iterableType, PyNames.ITERABLE, context)) {
334 final PyFunction iterateMethod = findMethodByName(iterableType, PyNames.ITER, context);
335 if (iterateMethod != null) {
336 final PyType iterateReturnType = getContextSensitiveType(iterateMethod, context, source);
337 final PyType type = getCollectionElementType(iterateReturnType, context);
338 if (!isTrivialType(type)) {
342 final String nextMethodName = LanguageLevel.forElement(anchor).isAtLeast(LanguageLevel.PYTHON30) ?
343 PyNames.DUNDER_NEXT : PyNames.NEXT;
344 final PyFunction next = findMethodByName(iterableType, nextMethodName, context);
346 final PyType type = getContextSensitiveType(next, context, source);
347 if (!isTrivialType(type)) {
351 final PyFunction getItem = findMethodByName(iterableType, PyNames.GETITEM, context);
352 if (getItem != null) {
353 final PyType type = getContextSensitiveType(getItem, context, source);
354 if (!isTrivialType(type)) {
362 private static boolean isTrivialType(@Nullable PyType type) {
363 return type == null || type instanceof PyNoneType;
367 private static PyType getCollectionElementType(@Nullable PyType type, @NotNull TypeEvalContext context) {
368 if (type instanceof PyCollectionType) {
369 final List<PyType> elementTypes = ((PyCollectionType)type).getElementTypes(context);
370 // TODO: Select the parameter type that matches T in Iterable[T]
371 return elementTypes.isEmpty() ? null : elementTypes.get(0);
377 private static PyFunction findMethodByName(@NotNull PyType type, @NotNull String name, @NotNull TypeEvalContext context) {
378 final PyResolveContext resolveContext = PyResolveContext.defaultContext().withTypeEvalContext(context);
379 final List<? extends RatedResolveResult> results = type.resolveMember(name, null, AccessDirection.READ, resolveContext);
380 if (results != null && !results.isEmpty()) {
381 final RatedResolveResult result = results.get(0);
382 final PsiElement element = result.getElement();
383 if (element instanceof PyFunction) {
384 return (PyFunction)element;
391 private static PyType getContextSensitiveType(@NotNull PyFunction function, @NotNull TypeEvalContext context,
392 @Nullable PyExpression source) {
393 return function.getCallType(source, Collections.<PyExpression, PyNamedParameter>emptyMap(), context);
397 private PyType getTypeFromExcept() {
398 PyExceptPart exceptPart = PsiTreeUtil.getParentOfType(this, PyExceptPart.class);
399 if (exceptPart == null || exceptPart.getTarget() != this) {
402 final PyExpression exceptClass = exceptPart.getExceptClass();
403 if (exceptClass instanceof PyReferenceExpression) {
404 final PsiElement element = ((PyReferenceExpression)exceptClass).getReference().resolve();
405 if (element instanceof PyClass) {
406 return new PyClassTypeImpl((PyClass) element, false);
412 public PyExpression getQualifier() {
413 ASTNode qualifier = getNode().findChildByType(PythonDialectsTokenSetProvider.INSTANCE.getExpressionTokens());
414 return qualifier != null ? (PyExpression) qualifier.getPsi() : null;
419 public QualifiedName asQualifiedName() {
420 if (myQualifiedName == null) {
421 myQualifiedName = PyPsiUtils.asQualifiedName(this);
423 return myQualifiedName;
426 public String toString() {
427 return super.toString() + ": " + getName();
430 public Icon getIcon(final int flags) {
431 if (isQualified() || PsiTreeUtil.getStubOrPsiParentOfType(this, PyDocStringOwner.class) instanceof PyClass) {
432 return PlatformIcons.FIELD_ICON;
434 return PlatformIcons.VARIABLE_ICON;
437 public boolean isQualified() {
438 PyTargetExpressionStub stub = getStub();
440 return stub.isQualified();
442 return getQualifier() != null;
447 public PsiElement resolveAssignedValue(@NotNull PyResolveContext resolveContext) {
448 final TypeEvalContext context = resolveContext.getTypeEvalContext();
449 if (context.maySwitchToAST(this)) {
450 final PyExpression value = findAssignedValue();
452 final List<PsiElement> results = PyUtil.multiResolveTopPriority(value, resolveContext);
453 return !results.isEmpty() ? results.get(0) : null;
458 final QualifiedName qName = getAssignedQName();
460 final ScopeOwner owner = ScopeUtil.getScopeOwner(this);
461 if (owner instanceof PyTypedElement) {
462 final List<String> components = qName.getComponents();
463 if (!components.isEmpty()) {
464 PsiElement resolved = owner;
465 for (String component : components) {
466 if (!(resolved instanceof PyTypedElement)) {
469 final PyType qualifierType = context.getType((PyTypedElement)resolved);
470 if (qualifierType == null) {
473 final List<? extends RatedResolveResult> results = qualifierType.resolveMember(component, null, AccessDirection.READ,
475 if (results == null || results.isEmpty()) {
478 resolved = results.get(0).getElement();
490 public PyExpression findAssignedValue() {
492 PyAssignmentStatement assignment = PsiTreeUtil.getParentOfType(this, PyAssignmentStatement.class);
493 if (assignment != null) {
494 List<Pair<PyExpression, PyExpression>> mapping = assignment.getTargetsToValuesMapping();
495 for (Pair<PyExpression, PyExpression> pair : mapping) {
496 PyExpression assigned_to = pair.getFirst();
497 if (assigned_to == this) return pair.getSecond();
506 public QualifiedName getAssignedQName() {
507 final PyTargetExpressionStub stub = getStub();
509 if (stub.getInitializerType() == PyTargetExpressionStub.InitializerType.ReferenceExpression) {
510 return stub.getInitializer();
514 return PyPsiUtils.asQualifiedName(findAssignedValue());
518 public QualifiedName getCalleeName() {
519 final PyTargetExpressionStub stub = getStub();
521 final PyTargetExpressionStub.InitializerType initializerType = stub.getInitializerType();
522 if (initializerType == PyTargetExpressionStub.InitializerType.CallExpression) {
523 return stub.getInitializer();
525 else if (initializerType == PyTargetExpressionStub.InitializerType.Custom) {
526 final CustomTargetExpressionStub customStub = stub.getCustomStub(CustomTargetExpressionStub.class);
527 if (customStub != null) {
528 return customStub.getCalleeName();
533 final PyExpression value = findAssignedValue();
534 if (value instanceof PyCallExpression) {
535 final PyExpression callee = ((PyCallExpression)value).getCallee();
536 return PyPsiUtils.asQualifiedName(callee);
543 public PsiReference getReference() {
544 return getReference(PyResolveContext.defaultContext());
549 public PsiPolyVariantReference getReference(final PyResolveContext resolveContext) {
551 return new PyQualifiedReference(this, resolveContext);
553 return new PyTargetReference(this, resolveContext);
558 public SearchScope getUseScope() {
560 return super.getUseScope();
562 final ScopeOwner owner = ScopeUtil.getScopeOwner(this);
564 final Scope scope = ControlFlowCache.getScope(owner);
565 if (scope.isGlobal(getName())) {
566 return GlobalSearchScope.projectScope(getProject());
568 if (scope.isNonlocal(getName())) {
569 return new LocalSearchScope(getContainingFile());
573 // find highest level function containing our var
574 PyElement container = this;
576 PyElement parentContainer = PsiTreeUtil.getParentOfType(container, PyFunction.class, PyClass.class);
577 if (parentContainer instanceof PyClass) {
579 return super.getUseScope();
583 if (parentContainer == null) {
586 container = parentContainer;
588 if (container instanceof PyFunction) {
589 return new LocalSearchScope(container);
591 return super.getUseScope();
595 public PyClass getContainingClass() {
596 final PyTargetExpressionStub stub = getStub();
598 final StubElement parentStub = stub.getParentStub();
599 if (parentStub instanceof PyClassStub) {
600 return ((PyClassStub)parentStub).getPsi();
602 if (parentStub instanceof PyFunctionStub) {
603 final StubElement functionParent = parentStub.getParentStub();
604 if (functionParent instanceof PyClassStub) {
605 return ((PyClassStub) functionParent).getPsi();
612 final PsiElement parent = PsiTreeUtil.getParentOfType(this, PyFunction.class, PyClass.class);
613 if (parent instanceof PyClass) {
614 return (PyClass)parent;
616 if (parent instanceof PyFunction) {
617 return ((PyFunction)parent).getContainingClass();
623 public ItemPresentation getPresentation() {
624 return new PyElementPresentation(this) {
627 public String getLocationString() {
628 final PyClass containingClass = getContainingClass();
629 if (containingClass != null) {
630 return "(" + containingClass.getName() + " in " + getPackageForFile(getContainingFile()) + ")";
632 return super.getLocationString();
639 public String getDocStringValue() {
640 final PyTargetExpressionStub stub = getStub();
642 return stub.getDocString();
644 return DocStringUtil.getDocStringValue(this);
649 public StructuredDocString getStructuredDocString() {
650 return DocStringUtil.getStructuredDocString(this);
655 public PyStringLiteralExpression getDocStringExpression() {
656 final PsiElement parent = getParent();
657 if (parent instanceof PyAssignmentStatement) {
658 final PsiElement nextSibling = PyPsiUtils.getNextNonCommentSibling(parent, true);
659 if (nextSibling instanceof PyExpressionStatement) {
660 final PyExpression expression = ((PyExpressionStatement)nextSibling).getExpression();
661 if (expression instanceof PyStringLiteralExpression) {
662 return (PyStringLiteralExpression)expression;
670 public void subtreeChanged() {
671 super.subtreeChanged();
672 myQualifiedName = null;
677 public String getQualifiedName() {
678 return QualifiedNameFinder.getQualifiedName(this);