543d635050d5073dc618af3930b136868f6df719
[idea/community.git] / python / psi-api / src / com / jetbrains / python / codeInsight / PyCustomMember.java
1 /*
2  * Copyright 2000-2015 JetBrains s.r.o.
3  *
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
7  *
8  * http://www.apache.org/licenses/LICENSE-2.0
9  *
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.
15  */
16 package com.jetbrains.python.codeInsight;
17
18 import com.intellij.extapi.psi.ASTWrapperPsiElement;
19 import com.intellij.icons.AllIcons;
20 import com.intellij.openapi.util.Key;
21 import com.intellij.openapi.util.UserDataHolderBase;
22 import com.intellij.psi.PsiElement;
23 import com.intellij.psi.PsiReference;
24 import com.intellij.psi.util.*;
25 import com.intellij.util.Function;
26 import com.jetbrains.python.psi.PyClass;
27 import com.jetbrains.python.psi.PyFunction;
28 import com.jetbrains.python.psi.PyPsiFacade;
29 import com.jetbrains.python.psi.PyTypedElement;
30 import com.jetbrains.python.psi.types.PyType;
31 import com.jetbrains.python.psi.types.TypeEvalContext;
32 import org.jetbrains.annotations.NotNull;
33 import org.jetbrains.annotations.Nullable;
34
35 import javax.swing.*;
36
37 /**
38  * @author Dennis.Ushakov
39  */
40 public class PyCustomMember extends UserDataHolderBase {
41   private static final Key<ParameterizedCachedValue<PyClass, PsiElement>>
42     RESOLVE = Key.create("resolve");
43   private final String myName;
44   private final boolean myResolveToInstance;
45   private final Function<PsiElement, PyType> myTypeCallback;
46   @Nullable
47   private final String myTypeName;
48
49   private final PsiElement myTarget;
50   private PyPsiPath myPsiPath;
51
52   boolean myFunction = false;
53
54   /**
55    * Force resolving to {@link MyInstanceElement} even if element is function
56    */
57   private boolean myAlwaysResolveToCustomElement;
58
59   public PyCustomMember(@NotNull final String name, @Nullable final String type, final boolean resolveToInstance) {
60     myName = name;
61     myResolveToInstance = resolveToInstance;
62     myTypeName = type;
63
64     myTarget = null;
65     myTypeCallback = null;
66   }
67
68   public PyCustomMember(@NotNull final String name) {
69     myName = name;
70     myResolveToInstance = false;
71     myTypeName = null;
72
73     myTarget = null;
74     myTypeCallback = null;
75   }
76
77   public PyCustomMember(@NotNull final String name,
78                         @Nullable final String type,
79                         final Function<PsiElement, PyType> typeCallback) {
80     myName = name;
81
82     myResolveToInstance = false;
83     myTypeName = type;
84
85     myTarget = null;
86     myTypeCallback = typeCallback;
87   }
88
89   public PyCustomMember(@NotNull final String name, @Nullable final PsiElement target, @Nullable String typeName) {
90     myName = name;
91     myTarget = target;
92     myResolveToInstance = false;
93     myTypeName = typeName;
94     myTypeCallback = null;
95   }
96
97   public PyCustomMember(@NotNull final String name, @Nullable final PsiElement target) {
98     this(name, target, null);
99   }
100
101   public PyCustomMember resolvesTo(String moduleQName) {
102     myPsiPath = new PyPsiPath.ToFile(moduleQName);
103     return this;
104   }
105
106   public PyCustomMember resolvesToClass(String classQName) {
107     myPsiPath = new PyPsiPath.ToClassQName(classQName);
108     return this;
109   }
110
111   /**
112    * Force resolving to {@link MyInstanceElement} even if element is function
113    */
114   @NotNull
115   public final PyCustomMember alwaysResolveToCustomElement() {
116     myAlwaysResolveToCustomElement = true;
117     return this;
118   }
119
120   public PyCustomMember toClass(String name) {
121     myPsiPath = new PyPsiPath.ToClass(myPsiPath, name);
122     return this;
123   }
124
125   public PyCustomMember toFunction(String name) {
126     myPsiPath = new PyPsiPath.ToFunction(myPsiPath, name);
127     return this;
128   }
129
130   public PyCustomMember toFunctionRecursive(String name) {
131     myPsiPath = new PyPsiPath.ToFunctionRecursive(myPsiPath, name);
132     return this;
133   }
134
135   public PyCustomMember toClassAttribute(String name) {
136     myPsiPath = new PyPsiPath.ToClassAttribute(myPsiPath, name);
137     return this;
138   }
139
140   public PyCustomMember toCall(String name, String... args) {
141     myPsiPath = new PyPsiPath.ToCall(myPsiPath, name, args);
142     return this;
143   }
144
145   public PyCustomMember toAssignment(String assignee) {
146     myPsiPath = new PyPsiPath.ToAssignment(myPsiPath, assignee);
147     return this;
148   }
149
150   public PyCustomMember toPsiElement(final PsiElement psiElement) {
151     myPsiPath = new PyPsiPath() {
152
153       @Override
154       public PsiElement resolve(PsiElement module) {
155         return psiElement;
156       }
157     };
158     return this;
159   }
160
161   public String getName() {
162     return myName;
163   }
164
165   public Icon getIcon() {
166     if (myTarget != null) {
167       return myTarget.getIcon(0);
168     }
169     return AllIcons.Nodes.Method;
170   }
171
172   @Nullable
173   public PsiElement resolve(@NotNull final PsiElement context) {
174
175     if (myTarget != null) {
176       return myTarget;
177     }
178
179     PyClass targetClass = null;
180     if (myTypeName != null) {
181
182       final ParameterizedCachedValueProvider<PyClass, PsiElement> provider = new ParameterizedCachedValueProvider<PyClass, PsiElement>() {
183         @Nullable
184         @Override
185         public CachedValueProvider.Result<PyClass> compute(
186           final PsiElement param) {
187           final PyClass result = PyPsiFacade.getInstance(param.getProject()).createClassByQName(myTypeName, param);
188           return CachedValueProvider.Result.create(result, PsiModificationTracker.MODIFICATION_COUNT);
189         }
190       };
191       targetClass = CachedValuesManager.getManager(context.getProject()).getParameterizedCachedValue(this, RESOLVE,
192                                                                                                      provider, false, context);
193     }
194     final PsiElement resolveTarget = findResolveTarget(context);
195     if (resolveTarget instanceof PyFunction && !myAlwaysResolveToCustomElement) {
196       return resolveTarget;
197     }
198     if (resolveTarget != null || targetClass != null) {
199       return new MyInstanceElement(targetClass, context, resolveTarget);
200     }
201     return null;
202   }
203
204   @Nullable
205   private PsiElement findResolveTarget(@NotNull PsiElement context) {
206     if (myPsiPath != null) {
207       return myPsiPath.resolve(context);
208     }
209     return null;
210   }
211
212   @Nullable
213   public String getShortType() {
214     if (myTypeName == null) {
215       return null;
216     }
217     int pos = myTypeName.lastIndexOf('.');
218     return myTypeName.substring(pos + 1);
219   }
220
221   public PyCustomMember asFunction() {
222     myFunction = true;
223     return this;
224   }
225
226   public boolean isFunction() {
227     return myFunction;
228   }
229
230   /**
231    * Checks if some reference points to this element
232    *
233    * @param reference reference to check
234    * @return true if reference points to it
235    */
236   public final boolean isReferenceToMe(@NotNull final PsiReference reference) {
237     final PsiElement element = reference.resolve();
238     if (!(element instanceof MyInstanceElement)) {
239       return false;
240     }
241     return ((MyInstanceElement)element).getThis().equals(this);
242   }
243
244   private class MyInstanceElement extends ASTWrapperPsiElement implements PyTypedElement {
245     private final PyClass myClass;
246     private final PsiElement myContext;
247
248     public MyInstanceElement(PyClass clazz, PsiElement context, PsiElement resolveTarget) {
249       super(resolveTarget != null ? resolveTarget.getNode() : clazz.getNode());
250       myClass = clazz;
251       myContext = context;
252     }
253
254     private PyCustomMember getThis() {
255       return PyCustomMember.this;
256     }
257
258     public PyType getType(@NotNull TypeEvalContext context, @NotNull TypeEvalContext.Key key) {
259       if (myTypeCallback != null) {
260         return myTypeCallback.fun(myContext);
261       }
262       else if (myClass != null) {
263         return PyPsiFacade.getInstance(getProject()).createClassType(myClass, !myResolveToInstance);
264       }
265       return null;
266     }
267   }
268 }