Basic support for user-defined skeletons

python/src/META-INF/python-plugin-common.xml

@@ -497,6 +497,9 @@
     <postStartupActivity implementation="com.jetbrains.python.packaging.PyPIPackagesUpdater"/>
     <postStartupActivity implementation="com.jetbrains.python.testing.PyTestRunnerUpdater"/>
 
+    <!-- User skeletons -->
+    <codeInsight.lineMarkerProvider language="Python" implementationClass="com.jetbrains.python.codeInsight.userSkeletons.PyUserSkeletonsLineMarkerProvider"/>
+
     <!-- Django -->
     <codeInsight.lineMarkerProvider language="Python" implementationClass="com.jetbrains.django.annotator.DjangoTemplateReferenceLineMarkerProvider"/>
     <configurationType implementation="com.jetbrains.django.run.DjangoServerRunConfigurationType"/>
@@ -856,6 +859,10 @@
     <documentationLinkProvider implementation="com.jetbrains.python.codeInsight.stdlib.PyStdlibDocumentationLinkProvider"/>
     <canonicalPathProvider implementation="com.jetbrains.python.codeInsight.stdlib.PyStdlibCanonicalPathProvider"/>
 
+    <!-- User skeletons -->
+    <pyModuleMembersProvider implementation="com.jetbrains.python.codeInsight.userSkeletons.PyUserSkeletonsModuleMembersProvider"/>
+    <typeProvider implementation="com.jetbrains.python.codeInsight.userSkeletons.PyUserSkeletonsTypeProvider"/>
+
     <typeProvider implementation="com.jetbrains.python.debugger.PyCallSignatureTypeProvider"/>
 
     <!-- Django -->

python/src/com/jetbrains/python/codeInsight/userSkeletons/PyUserSkeletonsLineMarkerProvider.java

@@ -0,0 +1,70 @@
+package com.jetbrains.python.codeInsight.userSkeletons;
+
+import com.intellij.codeHighlighting.Pass;
+import com.intellij.codeInsight.daemon.GutterIconNavigationHandler;
+import com.intellij.codeInsight.daemon.LineMarkerInfo;
+import com.intellij.codeInsight.daemon.LineMarkerProvider;
+import com.intellij.icons.AllIcons;
+import com.intellij.openapi.editor.markup.GutterIconRenderer;
+import com.intellij.psi.PsiElement;
+import com.intellij.util.Function;
+import com.intellij.util.PsiNavigateUtil;
+import com.jetbrains.python.psi.PyElement;
+import com.jetbrains.python.psi.PyFunction;
+import com.jetbrains.python.psi.PyTargetExpression;
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+
+import javax.swing.*;
+import java.awt.event.MouseEvent;
+import java.util.Collection;
+import java.util.List;
+
+/**
+ * @author vlan
+ */
+public class PyUserSkeletonsLineMarkerProvider implements LineMarkerProvider {
+  // TODO: Create an icon for a related user skeleton
+  public static final Icon ICON = AllIcons.Gutter.Unique;
+
+  @Nullable
+  @Override
+  public LineMarkerInfo getLineMarkerInfo(@NotNull PsiElement element) {
+    return null;
+  }
+
+  @Override
+  public void collectSlowLineMarkers(@NotNull List<PsiElement> elements, @NotNull Collection<LineMarkerInfo> result) {
+    for (PsiElement element : elements) {
+      final PyElement skeleton = getUserSkeleton(element);
+      if (skeleton != null) {
+        result.add(new LineMarkerInfo<PsiElement>(
+          element, element.getTextRange(), ICON, Pass.VISIBLE_LINE_MARKERS,
+          new Function<PsiElement, String>() {
+            @Override
+            public String fun(PsiElement e) {
+              return "Has user skeleton";
+            }
+          },
+          new GutterIconNavigationHandler<PsiElement>() {
+            @Override
+            public void navigate(MouseEvent e, PsiElement elt) {
+              final PyElement s = getUserSkeleton(elt);
+              if (s != null) {
+                PsiNavigateUtil.navigate(s);
+              }
+            }
+          },
+          GutterIconRenderer.Alignment.RIGHT));
+      }
+    }
+  }
+
+  @Nullable
+  private static PyElement getUserSkeleton(@NotNull PsiElement element) {
+    if (element instanceof PyFunction || element instanceof PyTargetExpression) {
+      return PyUserSkeletonsUtil.getUserSkeleton((PyElement)element);
+    }
+    return null;
+  }
+}

python/src/com/jetbrains/python/codeInsight/userSkeletons/PyUserSkeletonsModuleMembersProvider.java

@@ -0,0 +1,44 @@
+package com.jetbrains.python.codeInsight.userSkeletons;
+
+import com.intellij.psi.PsiElement;
+import com.jetbrains.python.codeInsight.PyDynamicMember;
+import com.jetbrains.python.psi.PyElement;
+import com.jetbrains.python.psi.PyFile;
+import com.jetbrains.python.psi.types.PyModuleMembersProvider;
+import org.jetbrains.annotations.Nullable;
+
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.List;
+
+/**
+ * @author vlan
+ */
+public class PyUserSkeletonsModuleMembersProvider extends PyModuleMembersProvider {
+  @Nullable
+  @Override
+  public PsiElement resolveMember(PyFile module, String name) {
+    final PyFile moduleSkeleton = PyUserSkeletonsUtil.getUserSkeleton(module);
+    if (moduleSkeleton != null) {
+      return moduleSkeleton.getElementNamed(name);
+    }
+    return null;
+  }
+
+  @Override
+  protected Collection<PyDynamicMember> getMembersByQName(PyFile module, String qName) {
+   final PyFile moduleSkeleton = PyUserSkeletonsUtil.getUserSkeletonForModuleQName(qName, module);
+    if (moduleSkeleton != null) {
+      final List<PyDynamicMember> results = new ArrayList<PyDynamicMember>();
+      for (PyElement element : moduleSkeleton.iterateNames()) {
+        final String name = element.getName();
+        if (name != null) {
+          results.add(new PyDynamicMember(name, element));
+        }
+      }
+      return results;
+    }
+    return Collections.emptyList();
+  }
+}

python/src/com/jetbrains/python/codeInsight/userSkeletons/PyUserSkeletonsTypeProvider.java

@@ -0,0 +1,49 @@
+package com.jetbrains.python.codeInsight.userSkeletons;
+
+import com.intellij.psi.PsiElement;
+import com.jetbrains.python.psi.*;
+import com.jetbrains.python.psi.types.PyType;
+import com.jetbrains.python.psi.types.PyTypeProviderBase;
+import com.jetbrains.python.psi.types.TypeEvalContext;
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+
+/**
+ * @author vlan
+ */
+public class PyUserSkeletonsTypeProvider extends PyTypeProviderBase {
+  @Override
+  public PyType getParameterType(@NotNull PyNamedParameter param, @NotNull PyFunction func, @NotNull TypeEvalContext context) {
+    final String name = param.getName();
+    if (name != null) {
+      final PyFunction functionSkeleton = PyUserSkeletonsUtil.getUserSkeleton(func);
+      if (functionSkeleton != null) {
+        final PyNamedParameter paramSkeleton = functionSkeleton.getParameterList().findParameterByName(name);
+        if (paramSkeleton != null) {
+          return context.getType(paramSkeleton);
+        }
+      }
+    }
+    return null;
+  }
+
+  @Override
+  public PyType getReturnType(@NotNull PyFunction function, @Nullable PyQualifiedExpression callSite, @NotNull TypeEvalContext context) {
+    final PyFunction functionSkeleton = PyUserSkeletonsUtil.getUserSkeleton(function);
+    if (functionSkeleton != null) {
+      return functionSkeleton.getReturnType(context, callSite);
+    }
+    return null;
+  }
+
+  @Override
+  public PyType getReferenceType(@NotNull PsiElement target, TypeEvalContext context, @Nullable PsiElement anchor) {
+    if (target instanceof PyTargetExpression) {
+      final PyTargetExpression targetSkeleton = PyUserSkeletonsUtil.getUserSkeleton((PyTargetExpression)target);
+      if (targetSkeleton != null) {
+        return context.getType(targetSkeleton);
+      }
+    }
+    return null;
+  }
+}

python/src/com/jetbrains/python/codeInsight/userSkeletons/PyUserSkeletonsUtil.java

@@ -0,0 +1,143 @@
+package com.jetbrains.python.codeInsight.userSkeletons;
+
+import com.intellij.openapi.application.ApplicationManager;
+import com.intellij.openapi.application.PathManager;
+import com.intellij.openapi.project.Project;
+import com.intellij.openapi.projectRoots.Sdk;
+import com.intellij.openapi.projectRoots.SdkModificator;
+import com.intellij.openapi.roots.OrderRootType;
+import com.intellij.openapi.util.text.StringUtil;
+import com.intellij.openapi.vfs.LocalFileSystem;
+import com.intellij.openapi.vfs.VirtualFile;
+import com.intellij.psi.PsiDirectory;
+import com.intellij.psi.PsiElement;
+import com.intellij.psi.PsiFile;
+import com.intellij.psi.PsiManager;
+import com.jetbrains.python.codeInsight.controlflow.ScopeOwner;
+import com.jetbrains.python.codeInsight.dataflow.scope.ScopeUtil;
+import com.jetbrains.python.psi.*;
+import com.jetbrains.python.psi.impl.PyQualifiedName;
+import com.jetbrains.python.psi.resolve.*;
+import com.jetbrains.python.psi.types.PyClassLikeType;
+import com.jetbrains.python.psi.types.PyType;
+import com.jetbrains.python.psi.types.TypeEvalContext;
+import com.jetbrains.python.sdk.PythonSdkType;
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+
+import java.io.File;
+import java.util.Collections;
+import java.util.List;
+
+/**
+ * @author vlan
+ */
+public class PyUserSkeletonsUtil {
+  @Nullable private static VirtualFile ourUserSkeletonsDirectory;
+
+  @NotNull
+  public static String getUserSkeletonsPath() {
+    if (ApplicationManager.getApplication().isInternal()) {
+      return StringUtil.join(new String[] {PathManager.getHomePath(), "python", "skeletons", "user-skeletons"}, File.separator);
+    }
+    // TODO: Copy user skeletons to this path
+    return PathManager.getSystemPath() + File.separator + "user-skeletons";
+  }
+
+  @Nullable
+  public static VirtualFile getUserSkeletonsDirectory() {
+    if (ourUserSkeletonsDirectory == null) {
+      ourUserSkeletonsDirectory = LocalFileSystem.getInstance().findFileByPath(getUserSkeletonsPath());
+    }
+    return ourUserSkeletonsDirectory;
+  }
+
+  @Nullable
+  public static <T extends PyElement> T getUserSkeleton(@NotNull T element) {
+    final PsiFile file = element.getContainingFile();
+    if (file instanceof PyFile) {
+      final PyFile skeletonFile = getUserSkeletonForFile((PyFile)file);
+      if (skeletonFile != null && skeletonFile != file) {
+        final PsiElement skeletonElement = getUserSkeleton(element, skeletonFile);
+        if (element.getClass().isInstance(skeletonElement) && skeletonElement != element) {
+          //noinspection unchecked
+          return (T)skeletonElement;
+        }
+      }
+    }
+    return null;
+  }
+
+  @Nullable
+  public static PyFile getUserSkeletonForModuleQName(@NotNull String qName, @NotNull PsiElement foothold) {
+    final Sdk sdk = PythonSdkType.getSdk(foothold);
+    if (sdk != null) {
+      final Project project = foothold.getProject();
+      final PythonSdkPathCache cache = PythonSdkPathCache.getInstance(project, sdk);
+      final PyQualifiedName cacheQName = PyQualifiedName.fromDottedString("user-skeletons." + qName);
+      final List<PsiElement> results = cache.get(cacheQName);
+      if (results != null) {
+        final PsiElement element = results.isEmpty() ? null : results.get(0);
+        if (element instanceof PyFile) {
+          return (PyFile)element;
+        }
+      }
+      final VirtualFile directory = getUserSkeletonsDirectory();
+      if (directory != null) {
+        final PsiDirectory psiDirectory = PsiManager.getInstance(project).findDirectory(directory);
+        final PsiElement fileSkeleton = new QualifiedNameResolverImpl(qName).resolveModuleAt(psiDirectory);
+        if (fileSkeleton instanceof PyFile) {
+          cache.put(cacheQName, Collections.singletonList(fileSkeleton));
+          return (PyFile)fileSkeleton;
+        }
+      }
+      cache.put(cacheQName, Collections.<PsiElement>emptyList());
+    }
+    return null;
+  }
+
+  public static void addUserSkeletonsRoot(@NotNull SdkModificator sdkModificator) {
+    final VirtualFile root = LocalFileSystem.getInstance().refreshAndFindFileByPath(getUserSkeletonsPath());
+    sdkModificator.addRoot(root, OrderRootType.CLASSES);
+  }
+
+  @Nullable
+  private static PsiElement getUserSkeleton(@NotNull PyElement element, @NotNull PyFile skeletonFile) {
+    if (element instanceof PyFile) {
+      return skeletonFile;
+    }
+    final ScopeOwner owner = ScopeUtil.getScopeOwner(element);
+    final String name = element.getName();
+    if (owner != null && name != null) {
+      assert owner != element;
+      final PsiElement originalOwner = getUserSkeleton(owner, skeletonFile);
+      if (originalOwner instanceof PyClass) {
+        final PyType type = TypeEvalContext.codeInsightFallback().getType((PyClass)originalOwner);
+        if (type instanceof PyClassLikeType) {
+          final PyClassLikeType classType = (PyClassLikeType)type;
+          final List<? extends RatedResolveResult> resolveResults = classType.resolveMember(name, null, AccessDirection.READ,
+                                                                                            PyResolveContext.noImplicits(), false);
+          if (resolveResults != null && !resolveResults.isEmpty()) {
+            return resolveResults.get(0).getElement();
+          }
+        }
+      }
+      else if (originalOwner instanceof NameDefiner) {
+        return ((NameDefiner)originalOwner).getElementNamed(name);
+      }
+    }
+    return null;
+  }
+
+  @Nullable
+  private static PyFile getUserSkeletonForFile(@NotNull PyFile file) {
+    final VirtualFile moduleVirtualFile = file.getVirtualFile();
+    if (moduleVirtualFile != null) {
+      final String moduleName = QualifiedNameFinder.findShortestImportableName(file, moduleVirtualFile);
+      if (moduleName != null) {
+        return getUserSkeletonForModuleQName(moduleName, file);
+      }
+    }
+    return null;
+  }
+}

python/src/com/jetbrains/python/psi/resolve/QualifiedNameResolverImpl.java

@@ -13,6 +13,7 @@ import com.intellij.psi.PsiDirectory;
 import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiFile;
 import com.jetbrains.django.facet.DjangoFacetType;
+import com.jetbrains.python.codeInsight.userSkeletons.PyUserSkeletonsUtil;
 import com.jetbrains.python.console.PydevConsoleRunner;
 import com.jetbrains.python.psi.impl.PyBuiltinCache;
 import com.jetbrains.python.psi.impl.PyImportResolver;
@@ -157,6 +158,9 @@ public class QualifiedNameResolverImpl implements RootVisitor, QualifiedNameReso
     if (!root.isValid()) {
       return true;
     }
+    if (root.equals(PyUserSkeletonsUtil.getUserSkeletonsDirectory())) {
+      return true;
+    }
     PsiElement resolveResult = resolveInRoot(root);
     if (resolveResult != null) {
       addRoot(resolveResult, isModuleSource);
@@ -331,7 +335,7 @@ public class QualifiedNameResolverImpl implements RootVisitor, QualifiedNameReso
    * @param directory where to start from; top qualifier will be searched for here.
    */
   @Nullable
-  private PsiElement resolveModuleAt(@Nullable PsiDirectory directory) {
+  public PsiElement resolveModuleAt(@Nullable PsiDirectory directory) {
     // prerequisites
     if (directory == null || !directory.isValid()) return null;
 

python/src/com/jetbrains/python/sdk/PythonSdkType.java

@@ -46,6 +46,7 @@ import com.jetbrains.python.PyBundle;
 import com.jetbrains.python.PyNames;
 import com.jetbrains.python.PythonFileType;
 import com.jetbrains.python.PythonHelpersLocator;
+import com.jetbrains.python.codeInsight.userSkeletons.PyUserSkeletonsUtil;
 import com.jetbrains.python.facet.PythonFacetSettings;
 import com.jetbrains.python.psi.LanguageLevel;
 import com.jetbrains.python.psi.impl.PyBuiltinCache;
@@ -573,6 +574,7 @@ public class PythonSdkType extends SdkType {
       }
     }
 
+    PyUserSkeletonsUtil.addUserSkeletonsRoot(sdkModificator);
     addSkeletonsRoot(sdkModificator, bin_path);
 
     if (not_in_unit_test_mode) {

python/src/com/jetbrains/python/sdk/PythonSdkUpdater.java

@@ -19,6 +19,7 @@ import com.intellij.openapi.startup.StartupActivity;
 import com.intellij.openapi.util.io.FileUtilRt;
 import com.intellij.openapi.vfs.VirtualFile;
 import com.jetbrains.python.PyBundle;
+import com.jetbrains.python.codeInsight.userSkeletons.PyUserSkeletonsUtil;
 import com.jetbrains.python.sdk.skeletons.PySkeletonRefresher;
 import org.jetbrains.annotations.NotNull;
 
@@ -125,6 +126,7 @@ public class PythonSdkUpdater implements StartupActivity {
   private static void updateSysPath(final Sdk sdk) throws InvalidSdkException {
     long start_time = System.currentTimeMillis();
     final List<String> sysPath = PythonSdkType.getSysPath(sdk.getHomePath());
+    sysPath.add(PyUserSkeletonsUtil.getUserSkeletonsPath());
     ApplicationManager.getApplication().invokeLater(new Runnable() {
       @Override
       public void run() {

python/testSrc/com/jetbrains/python/PythonMockSdk.java

@@ -7,6 +7,7 @@ import com.intellij.openapi.projectRoots.SdkType;
 import com.intellij.openapi.projectRoots.impl.ProjectJdkImpl;
 import com.intellij.openapi.roots.OrderRootType;
 import com.intellij.openapi.vfs.LocalFileSystem;
+import com.jetbrains.python.codeInsight.userSkeletons.PyUserSkeletonsUtil;
 import com.jetbrains.python.sdk.PythonSdkType;
 import org.jetbrains.annotations.NonNls;
 
@@ -53,6 +54,8 @@ public class PythonMockSdk {
       sdkModificator.addRoot(LocalFileSystem.getInstance().refreshAndFindFileByIoFile(libPath), OrderRootType.CLASSES);
     }
 
+    PyUserSkeletonsUtil.addUserSkeletonsRoot(sdkModificator);
+
     String mock_stubs_path = mock_path + PythonSdkType.SKELETON_DIR_NAME;
     sdkModificator.addRoot(LocalFileSystem.getInstance().refreshAndFindFileByPath(mock_stubs_path), PythonSdkType.BUILTIN_ROOT_TYPE);