PY-12745 Support 'six' library via typeshed

python/helpers/tools/packages

@@ -2049,8 +2049,7 @@ wpd wpd.mmxi.countdown
 tweet-tool Tweet-Command-Line-Tool 
 cmsplugin_bootstrap djangocms-bootstrap 
 ipaddress backport_ipaddress 
-six pi3d 
-multiselectbox TracMultiSelectBoxPlugin 
+multiselectbox TracMultiSelectBoxPlugin
 httplog django-httplog 
 genxmlif minixsv 
 tracking django-tracking-jl 

python/openapi/src/com/jetbrains/python/packaging/PyPackageManager.java

@@ -19,6 +19,7 @@ import com.intellij.execution.ExecutionException;
 import com.intellij.openapi.module.Module;
 import com.intellij.openapi.projectRoots.Sdk;
 import com.intellij.openapi.util.Key;
+import com.intellij.util.messages.Topic;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
@@ -32,6 +33,7 @@ public abstract class PyPackageManager {
   public static final Key<Boolean> RUNNING_PACKAGING_TASKS = Key.create("PyPackageRequirementsInspection.RunningPackagingTasks");
 
   public static final String USE_USER_SITE = "--user";
+  public static final Topic<Listener> PACKAGE_MANAGER_TOPIC = Topic.create("Python package manager", Listener.class);
 
   @NotNull
   public static PyPackageManager getInstance(@NotNull Sdk sdk) {
@@ -64,4 +66,8 @@ public abstract class PyPackageManager {
 
   @NotNull
   public abstract Set<PyPackage> getDependents(@NotNull PyPackage pkg) throws ExecutionException;
+
+  public interface Listener {
+    void packagesRefreshed(@NotNull Sdk sdk);
+  }
 }

python/python-community-configure/src/com/jetbrains/python/configuration/PythonPathEditor.java

@@ -37,6 +37,7 @@ import com.intellij.ui.ToolbarDecorator;
 import com.intellij.ui.components.JBList;
 import com.intellij.util.ArrayUtil;
 import com.jetbrains.python.PyBundle;
+import com.jetbrains.python.codeInsight.typing.PyTypeShed;
 import com.jetbrains.python.codeInsight.userSkeletons.PyUserSkeletonsUtil;
 import com.jetbrains.python.sdk.PythonSdkAdditionalData;
 import com.jetbrains.python.sdk.PythonSdkType;
@@ -291,6 +292,9 @@ public class PythonPathEditor extends SdkPathEditor {
       else if (file.equals(PyUserSkeletonsUtil.getUserSkeletonsDirectory())) {
         return true;
       }
+      else if (PyTypeShed.INSTANCE.isInside(file)) {
+        return true;
+      }
       else {
         return false;
       }

python/src/com/jetbrains/python/codeInsight/typing/PyTypeShed.kt

@@ -0,0 +1,115 @@
+/*
+ * Copyright 2000-2016 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+/*
+ * Copyright 2000-2016 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.jetbrains.python.codeInsight.typing
+
+import com.intellij.openapi.projectRoots.Sdk
+import com.intellij.openapi.vfs.StandardFileSystems
+import com.intellij.openapi.vfs.VfsUtilCore
+import com.intellij.openapi.vfs.VirtualFile
+import com.intellij.psi.util.QualifiedName
+import com.jetbrains.python.PythonHelpersLocator
+import com.jetbrains.python.packaging.PyPIPackageUtil
+import com.jetbrains.python.packaging.PyPackageManagers
+import com.jetbrains.python.packaging.PyPackageUtil
+import com.jetbrains.python.psi.LanguageLevel
+import com.jetbrains.python.sdk.PythonSdkType
+
+/**
+ * @author vlan
+ */
+object PyTypeShed {
+  private val ONLY_SUPPORTED_PY2_MINOR = 7
+  private val SUPPORTED_PY3_MINORS = 2..5
+  // TODO: Add `typing` to the white list and fix the tests
+  // TODO: Warn about unresolved `import typing` but still resolve it internally for type inference
+  private val WHITE_LIST = setOf("six")
+
+  /**
+   * Returns true if we allow to search typeshed for a stub for [name].
+   */
+  fun maySearchForStubInRoot(name: QualifiedName, root: VirtualFile, sdk : Sdk): Boolean {
+    val topLevelPackage = name.firstComponent ?: return false
+    if (topLevelPackage !in WHITE_LIST) {
+      return false
+    }
+    if (isInStandardLibrary(root)) {
+      return true
+    }
+    if (isInThirdPartyLibraries(root)) {
+      val pyPIPackage = PyPIPackageUtil.PACKAGES_TOPLEVEL[topLevelPackage] ?: topLevelPackage
+      val packages = PyPackageManagers.getInstance().forSdk(sdk).packages ?: return true
+      return PyPackageUtil.findPackage(packages, pyPIPackage) != null
+    }
+    return false
+  }
+
+  /**
+   * Returns list of roots in typeshed for Python language level of [sdk].
+   */
+  fun findRootsForSdk(sdk: Sdk): List<VirtualFile> {
+    val level = PythonSdkType.getLanguageLevelForSdk(sdk)
+    val minor = when (level.major) {
+      2 -> ONLY_SUPPORTED_PY2_MINOR
+      3 -> Math.min(Math.max(level.minor, SUPPORTED_PY3_MINORS.start), SUPPORTED_PY3_MINORS.endInclusive)
+      else -> return emptyList()
+    }
+    val dir = directory ?: return emptyList()
+    val paths = listOf("stdlib/${level.major}.${minor}",
+                       "stdlib/${level.major}",
+                       "stdlib/2and3",
+                       "third_party/${level.major}",
+                       "third_party/2and3")
+    return paths.asSequence()
+        .map { dir.findFileByRelativePath(it) }
+        .filterNotNull()
+        .toList()
+  }
+
+  fun isInside(file: VirtualFile): Boolean {
+    val dir = directory
+    return dir != null && VfsUtilCore.isAncestor(dir, file, true)
+  }
+
+  val directory: VirtualFile? by lazy {
+    val path = PythonHelpersLocator.getHelperPath("typeshed")
+    StandardFileSystems.local().findFileByPath(path)
+  }
+
+  fun isInThirdPartyLibraries(file: VirtualFile) = "third_party" in file.path
+
+  private fun isInStandardLibrary(file: VirtualFile) = "stdlib" in file.path
+
+  private val LanguageLevel.major: Int
+    get() = this.version / 10
+
+  private val LanguageLevel.minor: Int
+    get() = this.version % 10
+}

python/src/com/jetbrains/python/packaging/PyPackageManagerImpl.java

@@ -363,6 +363,7 @@ public class PyPackageManagerImpl extends PyPackageManager {
         final List<PyPackage> packages = collectPackages();
         LOG.debug("Packages installed in " + mySdk.getName() + ": " + packages);
         myPackagesCache = packages;
+        ApplicationManager.getApplication().getMessageBus().syncPublisher(PACKAGE_MANAGER_TOPIC).packagesRefreshed(mySdk);
         return Collections.unmodifiableList(packages);
       }
       catch (ExecutionException e) {

python/src/com/jetbrains/python/psi/impl/PyFileImpl.java

@@ -53,7 +53,6 @@ import com.jetbrains.python.psi.stubs.PyFileStub;
 import com.jetbrains.python.psi.types.PyModuleType;
 import com.jetbrains.python.psi.types.PyType;
 import com.jetbrains.python.psi.types.TypeEvalContext;
-import com.jetbrains.python.pyi.PyiUtil;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
@@ -271,12 +270,12 @@ public class PyFileImpl extends PsiFileBase implements PyFile, PyExpression {
       }
     }
   }
-
-  @Override
-  public PsiElement getNavigationElement() {
-    final PsiElement element = PyiUtil.getOriginalElement(this);
-    return element != null ? element : super.getNavigationElement();
-  }
+  //
+  //@Override
+  //public PsiElement getNavigationElement() {
+  //  final PsiElement element = PyiUtil.getOriginalElement(this);
+  //  return element != null ? element : super.getNavigationElement();
+  //}
 
   public boolean isAcceptedFor(@NotNull Class visitorClass) {
     for (Language lang : getViewProvider().getLanguages()) {

python/src/com/jetbrains/python/psi/resolve/PyResolveImportUtil.kt

@@ -34,6 +34,7 @@ import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiFileSystemItem
 import com.intellij.psi.PsiManager
 import com.intellij.psi.util.QualifiedName
+import com.jetbrains.python.codeInsight.typing.PyTypeShed
 import com.jetbrains.python.codeInsight.userSkeletons.PyUserSkeletonsUtil
 import com.jetbrains.python.facet.PythonPathContributingFacet
 import com.jetbrains.python.psi.LanguageLevel
@@ -125,7 +126,7 @@ fun resolveModuleAt(name: QualifiedName, directory: PsiDirectory?, context: PyQu
  * Creates a [PyQualifiedNameResolveContext] from a [foothold] element.
  */
 fun fromFoothold(foothold: PsiElement): PyQualifiedNameResolveContext {
-  val module = ModuleUtilCore.findModuleForPsiElement(foothold)
+  val module = ModuleUtilCore.findModuleForPsiElement(foothold.containingFile)
   val sdk = module?.let { ModuleRootManager.getInstance(module).sdk }
   return PyQualifiedNameResolveContextImpl(foothold.manager, module, foothold, sdk)
 }
@@ -181,7 +182,7 @@ private fun findFirstResults(results: List<PsiElement>) =
     if (results.all(::isNamespacePackage))
       results
     else {
-      val stubFile = results.firstOrNull { it is PyiFile }
+      val stubFile = results.firstOrNull { it is PyiFile || PyUtil.turnDirIntoInit(it) is PyiFile }
       if (stubFile != null)
         listOf(stubFile)
       else
@@ -214,9 +215,15 @@ private fun resultsFromRoots(name: QualifiedName, context: PyQualifiedNameResolv
   val moduleResults = mutableListOf<PsiElement>()
   val sdkResults = mutableListOf<PsiElement>()
 
+  val sdk = context.effectiveSdk
+  val module = context.module
+  val footholdFile = context.footholdFile
+
   val visitor = RootVisitor { root, module, sdk, isModuleSource ->
     val results = if (isModuleSource) moduleResults else sdkResults
-    if (!root.isValid || root == PyUserSkeletonsUtil.getUserSkeletonsDirectory()) {
+    if (!root.isValid ||
+        root == PyUserSkeletonsUtil.getUserSkeletonsDirectory() ||
+        sdk != null && PyTypeShed.isInside(root) && !PyTypeShed.maySearchForStubInRoot(name, root, sdk)) {
       return@RootVisitor true
     }
     val result = resolveInRoot(name, root, context)
@@ -232,10 +239,6 @@ private fun resultsFromRoots(name: QualifiedName, context: PyQualifiedNameResolv
     return@RootVisitor true
   }
 
-  val sdk = context.effectiveSdk
-  val module = context.module
-  val footholdFile = context.footholdFile
-
   when {
     context.visitAllModules -> {
       ModuleManager.getInstance(context.project).modules.forEach {

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

@@ -23,6 +23,8 @@ import com.intellij.openapi.projectRoots.Sdk;
 import com.intellij.openapi.roots.ModuleRootEvent;
 import com.intellij.openapi.roots.ModuleRootListener;
 import com.intellij.openapi.vfs.VirtualFileManager;
+import com.intellij.util.messages.MessageBusConnection;
+import com.jetbrains.python.packaging.PyPackageManager;
 import com.jetbrains.python.sdk.PythonSdkType;
 
 /**
@@ -35,12 +37,20 @@ public class PythonModulePathCache extends PythonPathCache implements Disposable
 
   @SuppressWarnings({"UnusedDeclaration"})
   public PythonModulePathCache(final Module module) {
-    module.getMessageBus().connect().subscribe(ProjectTopics.PROJECT_ROOTS, new ModuleRootListener() {
+    final MessageBusConnection connection = module.getMessageBus().connect();
+    connection.subscribe(ProjectTopics.PROJECT_ROOTS, new ModuleRootListener() {
       public void rootsChanged(ModuleRootEvent event) {
         updateCacheForSdk(module);
         clearCache();
       }
     });
+    connection.subscribe(PyPackageManager.PACKAGE_MANAGER_TOPIC, sdk -> {
+      final Sdk moduleSdk = PythonSdkType.findPythonSdk(module);
+      if (sdk == moduleSdk) {
+        updateCacheForSdk(module);
+        clearCache();
+      }
+    });
     VirtualFileManager.getInstance().addVirtualFileListener(new MyVirtualFileAdapter(), this);
     updateCacheForSdk(module);
   }

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

@@ -26,6 +26,8 @@ import com.intellij.openapi.util.Disposer;
 import com.intellij.openapi.util.Key;
 import com.intellij.openapi.vfs.VirtualFileManager;
 import com.intellij.util.containers.WeakHashMap;
+import com.intellij.util.messages.MessageBusConnection;
+import com.jetbrains.python.packaging.PyPackageManager;
 import com.jetbrains.python.psi.impl.PyBuiltinCache;
 import org.jetbrains.annotations.NotNull;
 
@@ -65,7 +67,8 @@ public class PythonSdkPathCache extends PythonPathCache implements Disposable {
       return;
     }
     Disposer.register(project, this);
-    project.getMessageBus().connect(this).subscribe(ProjectJdkTable.JDK_TABLE_TOPIC, new ProjectJdkTable.Adapter() {
+    final MessageBusConnection connection = project.getMessageBus().connect(this);
+    connection.subscribe(ProjectJdkTable.JDK_TABLE_TOPIC, new ProjectJdkTable.Adapter() {
       @Override
       public void jdkRemoved(Sdk jdk) {
         if (jdk == sdk) {
@@ -73,6 +76,11 @@ public class PythonSdkPathCache extends PythonPathCache implements Disposable {
         }
       }
     });
+    connection.subscribe(PyPackageManager.PACKAGE_MANAGER_TOPIC, eventSdk -> {
+      if (eventSdk == sdk) {
+        clearCache();
+      }
+    });
     sdk.getRootProvider().addRootSetChangedListener(new RootProvider.RootSetChangedListener() {
       @Override
       public void rootSetChanged(RootProvider wrapper) {

python/src/com/jetbrains/python/psi/search/PyProjectScopeBuilder.java

@@ -28,13 +28,16 @@ import com.intellij.openapi.util.io.FileUtil;
 import com.intellij.openapi.vfs.VirtualFile;
 import com.intellij.psi.PsiElement;
 import com.intellij.psi.search.*;
+import com.jetbrains.python.codeInsight.typing.PyTypeShed;
 import com.jetbrains.python.sdk.PythonSdkType;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
 import java.io.File;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.List;
+import java.util.stream.Collectors;
 
 /**
  *
@@ -115,18 +118,23 @@ public class PyProjectScopeBuilder extends ProjectScopeBuilderImpl {
   }
 
   @Nullable
-  public static GlobalSearchScope excludeSdkTestsScope(Project project, Sdk sdk) {
+  private static GlobalSearchScope excludeSdkTestsScope(Project project, Sdk sdk) {
     if (sdk != null && sdk.getSdkType() instanceof PythonSdkType) {
+      List<VirtualFile> excludedDirs = new ArrayList<>();
       VirtualFile libDir = findLibDir(sdk);
       if (libDir != null) {
         // superset of test dirs found in Python 2.5 to 3.1
-        List<VirtualFile> testDirs = findTestDirs(libDir, "test", "bsddb/test", "ctypes/test", "distutils/tests", "email/test",
-                                                  "importlib/test", "json/tests", "lib2to3/tests", "sqlite3/test", "tkinter/test",
-                                                  "idlelib/testcode.py");
-        if (!testDirs.isEmpty()) {
-          GlobalSearchScope scope = buildUnionScope(project, testDirs);
-          return GlobalSearchScope.notScope(scope);
-        }
+        excludedDirs.addAll(findTestDirs(libDir, "test", "bsddb/test", "ctypes/test", "distutils/tests", "email/test",
+                                         "importlib/test", "json/tests", "lib2to3/tests", "sqlite3/test", "tkinter/test",
+                                         "idlelib/testcode.py"));
+      }
+      // XXX: Disable resolving to any third-party libraries from typeshed in the same places where we don't want SDK tests
+      excludedDirs.addAll(Arrays.stream(sdk.getRootProvider().getFiles(OrderRootType.CLASSES))
+                            .filter(file -> PyTypeShed.INSTANCE.isInside(file) || PyTypeShed.INSTANCE.isInThirdPartyLibraries(file))
+                            .collect(Collectors.toList()));
+      if (!excludedDirs.isEmpty()) {
+        GlobalSearchScope scope = buildUnionScope(project, excludedDirs);
+        return GlobalSearchScope.notScope(scope);
       }
     }
     return null;

python/src/com/jetbrains/python/pyi/PyiRelatedItemLineMarkerProvider.java

@@ -26,6 +26,7 @@ import com.intellij.psi.PsiElement;
 import com.intellij.psi.SmartPointerManager;
 import com.intellij.psi.SmartPsiElementPointer;
 import com.intellij.util.PsiNavigateUtil;
+import com.jetbrains.python.psi.PyClass;
 import com.jetbrains.python.psi.PyElement;
 import com.jetbrains.python.psi.PyFunction;
 import com.jetbrains.python.psi.PyTargetExpression;
@@ -45,7 +46,7 @@ public class PyiRelatedItemLineMarkerProvider extends RelatedItemLineMarkerProvi
 
   @Override
   protected void collectNavigationMarkers(@NotNull PsiElement element, Collection<? super RelatedItemLineMarkerInfo> result) {
-    if (element instanceof PyFunction || element instanceof PyTargetExpression) {
+    if (element instanceof PyFunction || element instanceof PyTargetExpression || element instanceof PyClass) {
       final PsiElement pythonStub = PyiUtil.getPythonStub((PyElement)element);
       if (pythonStub != null) {
         result.add(createLineMarkerInfo(element, pythonStub, "Has stub item"));

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

@@ -43,6 +43,7 @@ import com.intellij.util.PathMappingSettings;
 import com.intellij.util.concurrency.BlockingSet;
 import com.intellij.util.concurrency.EdtExecutorService;
 import com.jetbrains.python.PyBundle;
+import com.jetbrains.python.codeInsight.typing.PyTypeShed;
 import com.jetbrains.python.codeInsight.userSkeletons.PyUserSkeletonsUtil;
 import com.jetbrains.python.packaging.PyPackageManager;
 import com.jetbrains.python.psi.PyUtil;
@@ -278,6 +279,7 @@ public class PythonSdkUpdater implements StartupActivity {
       .addAll(filterRootPaths(sdk, evaluateSysPath(sdk), project))
       .addAll(getSkeletonsPaths(sdk))
       .addAll(getUserAddedPaths(sdk))
+      .addAll(PyTypeShed.INSTANCE.findRootsForSdk(sdk))
       .build();
   }
 
@@ -292,6 +294,7 @@ public class PythonSdkUpdater implements StartupActivity {
       .addAll(getRemoteSdkMappedPaths(sdk, project))
       .addAll(getSkeletonsPaths(sdk))
       .addAll(getUserAddedPaths(sdk))
+      .addAll(PyTypeShed.INSTANCE.findRootsForSdk(sdk))
       .build();
   }
 

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

@@ -25,6 +25,7 @@ import com.intellij.openapi.vfs.LocalFileSystem;
 import com.intellij.openapi.vfs.VirtualFile;
 import com.intellij.psi.stubs.StubUpdatingIndex;
 import com.intellij.util.indexing.FileBasedIndex;
+import com.jetbrains.python.codeInsight.typing.PyTypeShed;
 import com.jetbrains.python.codeInsight.userSkeletons.PyUserSkeletonsUtil;
 import com.jetbrains.python.psi.stubs.PyModuleNameIndex;
 import com.jetbrains.python.sdk.PythonSdkType;
@@ -75,6 +76,7 @@ public class PythonMockSdk {
     }
 
     sdkModificator.addRoot(PyUserSkeletonsUtil.getUserSkeletonsDirectory(), OrderRootType.CLASSES);
+    PyTypeShed.INSTANCE.findRootsForSdk(sdk).forEach(file -> sdkModificator.addRoot(file, OrderRootType.CLASSES));
 
     String mock_stubs_path = mock_path + PythonSdkType.SKELETON_DIR_NAME;
     sdkModificator.addRoot(LocalFileSystem.getInstance().refreshAndFindFileByPath(mock_stubs_path), PythonSdkType.BUILTIN_ROOT_TYPE);