feat: Query Sync for Python

examples/python/with_numpy/lib/BUILD.bazel

@@ -0,0 +1,5 @@
+py_library(
+    name = "symbols",
+    srcs = glob(["*.py"]),
+    imports = ["."],
+)

examples/python/with_numpy/lib/symbols.py

@@ -0,0 +1 @@
+ALPHA="alpha"

python/BUILD

@@ -29,7 +29,9 @@ java_library(
         "//intellij_platform_sdk:jsr305",
         "//intellij_platform_sdk:plugin_api",
         "//proto:proto_deps",
+        "//querysync",
         "//sdkcompat",
+        "//shared",
         "//third_party/python",
     ],
 )
@@ -85,6 +87,7 @@ intellij_integration_test_suite(
         "//intellij_platform_sdk:jsr305",
         "//intellij_platform_sdk:plugin_api_for_tests",
         "//intellij_platform_sdk:test_libs",
+        "//querysync",
         "//third_party/python:python_for_tests",
         "//third_party/toml:toml_for_tests",
         "@com_google_guava_guava//jar",

python/src/com/google/idea/blaze/python/resolve/BlazePyResolverUtils.java

@@ -20,6 +20,8 @@
 import com.google.idea.blaze.base.io.VirtualFileSystemProvider;
 import com.google.idea.blaze.base.model.BlazeProjectData;
 import com.google.idea.blaze.base.model.RemoteOutputArtifacts;
+import com.google.idea.blaze.base.settings.Blaze;
+import com.google.idea.blaze.base.settings.BlazeImportSettings;
 import com.google.idea.blaze.base.sync.data.BlazeProjectDataManager;
 import com.google.idea.blaze.base.sync.workspace.WorkspacePathResolver;
 import com.intellij.openapi.project.Project;
@@ -51,6 +53,9 @@ public static PsiElement resolveGenfilesPath(
   private static Optional<File> resolveGenfilesPath(Project project, String relativePath) {
     BlazeProjectData projectData =
         BlazeProjectDataManager.getInstance(project).getBlazeProjectData();
+    if(Blaze.getProjectType(project) == BlazeImportSettings.ProjectType.QUERY_SYNC) {
+      return Optional.empty();
+    }
     if (projectData == null) {
       return Optional.empty();
     }

python/src/com/google/idea/blaze/python/resolve/provider/AbstractPyImportResolverStrategy.java

@@ -26,12 +26,17 @@
 import com.google.idea.blaze.base.ideinfo.TargetIdeInfo;
 import com.google.idea.blaze.base.model.BlazeProjectData;
 import com.google.idea.blaze.base.model.primitives.LanguageClass;
+import com.google.idea.blaze.base.model.primitives.WorkspaceRoot;
+import com.google.idea.blaze.base.qsync.QuerySyncManager;
 import com.google.idea.blaze.base.settings.Blaze;
 import com.google.idea.blaze.base.settings.BlazeImportSettings.ProjectType;
 import com.google.idea.blaze.base.sync.SyncCache;
 import com.google.idea.blaze.base.sync.workspace.ArtifactLocationDecoder;
 import com.google.idea.blaze.base.sync.workspace.WorkspacePathResolver;
+import com.google.idea.blaze.common.Label;
 import com.google.idea.blaze.python.resolve.BlazePyResolverUtils;
+import com.google.idea.blaze.qsync.project.ProjectTarget;
+import com.google.idea.blaze.qsync.query.Query;
 import com.intellij.openapi.project.Project;
 import com.intellij.openapi.util.text.StringUtil;
 import com.intellij.psi.PsiElement;
@@ -46,11 +51,7 @@
 import com.jetbrains.python.psi.resolve.PyQualifiedNameResolveContext;
 import java.io.File;
 import java.nio.file.Path;
-import java.util.Collection;
-import java.util.HashMap;
-import java.util.List;
-import java.util.Map;
-import java.util.Set;
+import java.util.*;
 import javax.annotation.Nullable;
 
 /**
@@ -129,19 +130,21 @@ public final void addImportCandidates(
   @Nullable
   private PySourcesIndex getSourcesIndex(Project project) {
     if (Blaze.getProjectType(project) == ProjectType.QUERY_SYNC) {
-      return null;
+      return SyncCache.getInstance(project).get(getClass(), this::buildSourcesIndexQuerySync);
     }
     return SyncCache.getInstance(project).get(getClass(), this::buildSourcesIndex);
   }
 
   // exposed package-private for testing
   @SuppressWarnings("unused")
   PySourcesIndex buildSourcesIndex(Project project, BlazeProjectData projectData) {
+
     ImmutableSetMultimap.Builder<String, QualifiedName> shortNames = ImmutableSetMultimap.builder();
     Map<QualifiedName, PsiElementProvider> map = new HashMap<>();
     ArtifactLocationDecoder decoder = projectData.getArtifactLocationDecoder();
     for (TargetIdeInfo target : projectData.getTargetMap().targets()) {
       List<QualifiedName> importRoots = assembleImportRoots(target);
+
       for (ArtifactLocation source : getPySources(projectData.getWorkspacePathResolver(), target)) {
         List<QualifiedName> sourceImports = assembleSourceImportsFromImportRoots(importRoots,
             toImportString(source));
@@ -161,6 +164,46 @@ PySourcesIndex buildSourcesIndex(Project project, BlazeProjectData projectData)
     return new PySourcesIndex(shortNames.build(), ImmutableMap.copyOf(map));
   }
 
+  // exposed package-private for testing
+  @SuppressWarnings("unused")
+  PySourcesIndex buildSourcesIndexQuerySync(Project project, BlazeProjectData projectData) {
+
+    ImmutableSetMultimap.Builder<String, QualifiedName> shortNames = ImmutableSetMultimap.builder();
+    Map<QualifiedName, PsiElementProvider> map = new HashMap<>();
+
+    var currentSnapshot =QuerySyncManager.getInstance(project).getLoadedProject().get().getSnapshotHolder().getCurrent()
+            .orElse(null);
+    if(currentSnapshot == null) {
+      return null;
+    }
+
+    var workspaceRoot = WorkspaceRoot.fromProjectSafe(project);
+    if (workspaceRoot == null) {
+      return null;
+    }
+
+    for(var t: currentSnapshot.getTargetMap().entrySet()){
+      List<QualifiedName> importRoots = assembleImportRootsQuerySync(t.getKey(), t.getValue(), project);
+      for (var source : t.getValue().sourceLabels().get(ProjectTarget.SourceType.REGULAR)) {
+        List<QualifiedName> sourceImports = assembleSourceImportsFromImportRoots(importRoots,
+                fromRelativePath(source.toFilePath().toString()));
+        var file = workspaceRoot.path().resolve(source.toFilePath());
+        for (QualifiedName sourceImport : sourceImports) {
+          if (null != sourceImport.getLastComponent()) {
+            shortNames.put(sourceImport.getLastComponent(), sourceImport);
+            PsiElementProvider psiProvider = mapQuerySync(file.toFile());
+            map.put(sourceImport, psiProvider);
+            if (includeParentDirectory(file.toFile())) {
+              map.put(sourceImport.removeTail(1), PsiElementProvider.getParent(psiProvider));
+            }
+          }
+        }
+      }
+    }
+
+    return new PySourcesIndex(shortNames.build(), ImmutableMap.copyOf(map));
+  }
+
   /**
    * This method will extract sources from the supplied target. If any of the sources
    * are a directory rather than a file then it will descend through the directory
@@ -280,6 +323,10 @@ private static boolean includeParentDirectory(ArtifactLocation source) {
     return source.getRelativePath().endsWith(".py");
   }
 
+  private static boolean includeParentDirectory(File source) {
+    return source.getName().endsWith(".py");
+  }
+
   static QualifiedName fromRelativePath(String relativePath) {
     relativePath = StringUtil.trimEnd(relativePath, File.separator + PyNames.INIT_DOT_PY);
     relativePath = StringUtil.trimExtensions(relativePath);
@@ -361,6 +408,49 @@ private static List<QualifiedName> assembleImportRoots(TargetIdeInfo target) {
     return resultBuilder.build();
   }
 
+  private static List<QualifiedName> assembleImportRootsQuerySync(Label label, ProjectTarget target, Project project) {
+
+    ImmutableList.Builder<QualifiedName> resultBuilder = ImmutableList.builder();
+//      var newLabel = com.google.idea.blaze.common.Label.of(target.getKey().getLabel().toString());
+
+      // In query sync, when analysis is disabled, we rely on  the `imports` attribute
+      // of the rules. Once it's anabled we should switch to what we got from providers
+      var imports = QuerySyncManager.getInstance(project).getLoadedProject()
+              .flatMap(it -> it.getSnapshotHolder().getCurrent())
+              .map(it -> it.queryData().querySummary().getRulesMap())
+              .flatMap(it -> Optional.ofNullable(it.get(label)))
+              .map(Query.Rule::getImportsList)
+              .map(it -> it.stream().toList())
+              .orElse(ImmutableList.of());
+
+      var buildParentPath = label.getPackage();
+      for (String imp : imports) {
+        Path impPath = buildParentPath.resolve(imp).normalize();
+        String[] impPathParts = new String[impPath.getNameCount()];
+
+        for (int i = impPath.getNameCount() - 1; i >= 0; i--) {
+          impPathParts[i] = impPath.getName(i).toString();
+        }
+
+        resultBuilder.add(QualifiedName.fromComponents(impPathParts));
+      }
+
+
+    return resultBuilder.build();
+  }
+
+
+  public PsiElementProvider mapQuerySync(File file0) {
+    return (manager) -> {
+      var file = file0;
+      if (PyNames.INIT_DOT_PY.equals(file.getName())) {
+        file = file.getParentFile();
+      }
+      return BlazePyResolverUtils.resolveFile(manager, file);
+    };
+  }
+
+
   /**
    * For each of the <code>importRoots</code>, see if it matches as a prefix on the
    * <code>sourceImport</code> and then trim off the prefix; yielding the true module name.

querysync/BUILD

@@ -9,6 +9,7 @@ package(default_visibility = [
     "//java/com/google/devtools/intellij/blaze/plugin/base:__subpackages__",
     "//java/com/google/devtools/intellij/blaze/plugin/cpp:__subpackages__",
     "//java/com/google/devtools/intellij/blaze/plugin/querysync:__subpackages__",
+    "//python:__subpackages__",
     "//querysync/javatests:__subpackages__",
 ])
 

querysync/java/com/google/idea/blaze/qsync/AffectedPackagesCalculator.java

@@ -235,7 +235,7 @@ public AffectedPackages getAffectedPackages() {
 
   private boolean isIncludedInProject(Path file) {
     for (Path includePath : projectIncludes()) {
-      if (file.startsWith(includePath)) {
+      if (file.startsWith(includePath) || includePath.toString().isEmpty()) {
         for (Path excludePath : projectExcludes()) {
           if (file.startsWith(excludePath)) {
             return false;

querysync/java/com/google/idea/blaze/qsync/BlazeQueryParser.java

@@ -137,6 +137,9 @@ public BuildGraphData parse() {
       if (RuleKinds.isProtoSource(ruleClass)) {
         visitProtoRule(ruleEntry.getValue(), targetBuilder);
       }
+      if (RuleKinds.isPythonSource(ruleClass)) {
+        visitPythonRule(ruleEntry.getKey(), ruleEntry.getValue(), targetBuilder);
+      }
       if (alwaysBuildRuleKinds.contains(ruleClass)) {
         projectTargetsToBuild.add(ruleEntry.getKey());
       }
@@ -174,6 +177,21 @@ public BuildGraphData parse() {
     return graph;
   }
 
+  private void visitPythonRule(Label label, Rule rule, ProjectTarget.Builder targetBuilder) {
+    graphBuilder.allTargetsBuilder().add(label);
+    targetBuilder.languagesBuilder().add(QuerySyncLanguage.PYTHON);
+    targetBuilder
+            .sourceLabelsBuilder()
+            .putAll(SourceType.REGULAR, expandFileGroupValues(rule.getSourcesList()));
+
+
+    Set<Label> thisDeps = Sets.newHashSet(toLabelList(rule.getDepsList()));
+    targetBuilder.depsBuilder().addAll(thisDeps);
+
+    targetBuilder.runtimeDepsBuilder().addAll(toLabelList(rule.getRuntimeDepsList()));
+    javaDeps.addAll(thisDeps);
+  }
+
   private void visitProtoRule(Query.Rule rule, ProjectTarget.Builder targetBuilder) {
     targetBuilder
         .sourceLabelsBuilder()
@@ -182,7 +200,7 @@ private void visitProtoRule(Query.Rule rule, ProjectTarget.Builder targetBuilder
 
   private void visitJavaRule(Label label, Query.Rule rule, ProjectTarget.Builder targetBuilder) {
     graphBuilder.allTargetsBuilder().add(label);
-    targetBuilder.languagesBuilder().add(QuerySyncLanguage.JAVA);
+    targetBuilder.languagesBuilder().add(QuerySyncLanguage.PYTHON);
     targetBuilder
         .sourceLabelsBuilder()
         .putAll(SourceType.REGULAR, expandFileGroupValues(rule.getSourcesList()))

querysync/java/com/google/idea/blaze/qsync/project/QuerySyncLanguage.java

@@ -26,7 +26,8 @@
 public enum QuerySyncLanguage {
   JAVA(LanguageClass.LANGUAGE_CLASS_JAVA, DependencyTrackingBehavior.EXTERNAL_DEPENDENCIES),
   KOTLIN(LanguageClass.LANGUAGE_CLASS_KOTLIN, DependencyTrackingBehavior.EXTERNAL_DEPENDENCIES),
-  CC(LanguageClass.LANGUAGE_CLASS_CC, DependencyTrackingBehavior.SELF);
+  CC(LanguageClass.LANGUAGE_CLASS_CC, DependencyTrackingBehavior.SELF),
+  PYTHON(LanguageClass.LANGUAGE_CLASS_PYTHON, DependencyTrackingBehavior.EXTERNAL_DEPENDENCIES);
 
   QuerySyncLanguage(
       LanguageClassProto.LanguageClass protoValue,

querysync/java/com/google/idea/blaze/qsync/project/language_class.proto

@@ -11,5 +11,6 @@ enum LanguageClass {
   LANGUAGE_CLASS_JAVA = 1;
   LANGUAGE_CLASS_KOTLIN = 2;
   LANGUAGE_CLASS_CC = 3;
+  LANGUAGE_CLASS_PYTHON = 4;
 }
 

querysync/java/com/google/idea/blaze/qsync/query/QuerySummary.java

@@ -190,8 +190,10 @@ public static QuerySummary create(InputStream protoInputStream) throws IOExcepti
 
             if (a.getName().equals("test_app")) {
               rule.setTestApp(a.getStringValue());
-            } else if (a.getName().equals("instruments")) {
+            } if (a.getName().equals("instruments")) {
               rule.setInstruments(a.getStringValue());
+            } else if (a.getName().equals("imports")) {
+              rule.addAllImports(a.getStringListValueList());
             }
           }
           ruleMap.put(Label.of(target.getRule().getName()).toString(), rule.build());

querysync/java/com/google/idea/blaze/qsync/query/querysummary.proto

@@ -38,4 +38,5 @@ message Rule {
   repeated string copts = 13;
   repeated string tags = 14;
   string main_class = 15;
+  repeated string imports = 16;
 }

shared/java/com/google/idea/blaze/common/RuleKinds.java

@@ -53,6 +53,9 @@ private RuleKinds() {}
   private static final ImmutableSet<String> PROTO_SOURCE_RULE_KINDS =
       ImmutableSet.of("proto_library");
 
+  private static final ImmutableSet<String> PYTHON_RULE_KINDS =
+          ImmutableSet.of("py_library", "py_binary", "py_test");
+
   public static boolean isJava(String ruleClass) {
     return JAVA_RULE_KINDS.contains(ruleClass) || ANDROID_RULE_KINDS.contains(ruleClass);
   }
@@ -68,4 +71,8 @@ public static boolean isCc(String ruleClass) {
   public static boolean isProtoSource(String ruleClass) {
     return PROTO_SOURCE_RULE_KINDS.contains(ruleClass);
   }
+
+  public static boolean isPythonSource(String ruleClass) {
+    return PYTHON_RULE_KINDS.contains(ruleClass);
+  }
 }