intellij: Intellij doesn't resolve python imports from pip_install()

Python rules used from https://github.com/bazelbuild/rules_python

A pip dependency added to my target in deps. E.g.: requirement("tensorflow")

import tensorflow OR from pypi__tensorflow as tensorflow

Both of them work with bazel commandline; but Intellij IDE only resolves the second one with Bazel plugin. Looking at the code, the behaviour makes sense.

Following patch reverses the IDE behavior (i.e. first import is resolved and second one isn’t- which seems correct behavior): In https://github.com/bazelbuild/intellij/blob/master/python/src/com/google/idea/blaze/python/resolve/provider/BazelPyImportResolverStrategy.java#L50:

diff --git a/python/src/com/google/idea/blaze/python/resolve/provider/BazelPyImportResolverStrategy.java b/python/src/com/google/idea/blaze/python/resolve/provider/BazelPyImportResolverStrategy.java
index ee0acbda..56d1f92c 100644
--- a/python/src/com/google/idea/blaze/python/resolve/provider/BazelPyImportResolverStrategy.java
+++ b/python/src/com/google/idea/blaze/python/resolve/provider/BazelPyImportResolverStrategy.java
@@ -22,6 +22,8 @@ import com.intellij.psi.PsiElement;
 import com.intellij.psi.util.QualifiedName;
 import com.jetbrains.python.psi.resolve.PyQualifiedNameResolveContext;
 import java.io.File;
+import java.nio.file.Path;
+import java.nio.file.Paths;
 import javax.annotation.Nullable;
 
 /**
@@ -32,6 +34,10 @@ import javax.annotation.Nullable;
  */
 public class BazelPyImportResolverStrategy extends AbstractPyImportResolverStrategy {
 
+  // Packages installed with pip_install in python_rules are installed in a subdir that starts
+  // with this prefix.
+  private static final String PIP_INSTALL_DIR_PREFIX = "pypi__";
+
   @Override
   public boolean appliesToBuildSystem(BuildSystem buildSystem) {
     return buildSystem == BuildSystem.Bazel;
@@ -51,6 +57,11 @@ public class BazelPyImportResolverStrategy extends AbstractPyImportResolverStrat
     if (source.isGenerated() || !source.getRelativePath().endsWith(".py")) {
       return null;
     }
-    return fromRelativePath(source.getRelativePath());
+    String relativePath = source.getRelativePath();
+    if (relativePath.startsWith(PIP_INSTALL_DIR_PREFIX)) {
+      Path path = Paths.get(relativePath);
+      relativePath = path.subpath(1, path.getNameCount()).toString();
+    }
+    return fromRelativePath(relativePath);
   }
 }

If it makes sense, I can send a pull request.