PY-75831 Split cache(resolve/type) into library and user part

GitOrigin-RevId: 8dfd0120379c9a34051d66e147ffdc2c69f0db66

python/python-psi-api/src/com/jetbrains/python/psi/types/PyLibraryModificationTracker.kt

@@ -0,0 +1,80 @@
+// Copyright 2000-2025 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package com.jetbrains.python.psi.types
+
+import com.intellij.openapi.Disposable
+import com.intellij.openapi.components.Service
+import com.intellij.openapi.components.service
+import com.intellij.openapi.diagnostic.Logger
+import com.intellij.openapi.diagnostic.ThrottledLogger
+import com.intellij.openapi.fileEditor.FileDocumentManagerListener
+import com.intellij.openapi.project.DumbService
+import com.intellij.openapi.project.Project
+import com.intellij.openapi.roots.ProjectRootManager
+import com.intellij.openapi.util.ModificationTracker
+import com.intellij.openapi.util.SimpleModificationTracker
+import com.intellij.openapi.vfs.VirtualFile
+import com.intellij.psi.PsiManager
+import com.intellij.psi.PsiTreeChangeAdapter
+import com.intellij.psi.PsiTreeChangeEvent
+import com.intellij.psi.search.GlobalSearchScope
+import com.intellij.psi.search.ProjectScope
+import com.intellij.util.ForcefulReparseModificationTracker
+import java.util.concurrent.TimeUnit
+
+@Service(Service.Level.PROJECT)
+class PyLibraryModificationTracker(project: Project) : ModificationTracker, Disposable {
+  private val myProjectRootManager: ModificationTracker = ProjectRootManager.getInstance(project)
+  private val myDumbServiceModificationTracker: ModificationTracker = DumbService.getInstance(project).modificationTracker
+  private val myForcefulReparseModificationTracker: ModificationTracker = ForcefulReparseModificationTracker.getInstance() // PsiClass from libraries may become invalid on reparse
+  private val myOnContentReloadModificationTracker: SimpleModificationTracker = SimpleModificationTracker()
+  private val creationStack = Throwable()
+  val projectLibraryScope: GlobalSearchScope = ProjectScope.getLibrariesScope(project)
+
+  init {
+    val connection = project.getMessageBus().connect(this)
+    PsiManager.getInstance(project).addPsiTreeChangeListener(object : PsiTreeChangeAdapter() {
+      override fun childrenChanged(event: PsiTreeChangeEvent) {
+        val file = event.file ?: return
+        val virtualFile = file.virtualFile ?: return
+        if (isLibraryFile(virtualFile)) {
+          myOnContentReloadModificationTracker.incModificationCount()
+        }
+      }
+    }, this)
+
+    connection.subscribe<FileDocumentManagerListener>(FileDocumentManagerListener.TOPIC, object : FileDocumentManagerListener {
+      override fun fileWithNoDocumentChanged(file: VirtualFile) {
+        if (!project.isInitialized()) {
+          THROTTLED_LOG.warn("SearchScope.contains(file) would log an error because WorkspaceFileIndex is not yet initialized. " +
+                             "Probably LibraryModificationTracker was created too early. " +
+                             "See LibraryModificationTracker creation stacktrace: ", creationStack)
+          return
+        }
+        if (isLibraryFile(file)) {
+          myOnContentReloadModificationTracker.incModificationCount()
+        }
+      }
+    })
+  }
+
+  private fun isLibraryFile(file: VirtualFile): Boolean {
+    return "pyi" == file.extension || projectLibraryScope.contains(file)
+  }
+
+  override fun getModificationCount(): Long {
+    return (myProjectRootManager.getModificationCount()
+            + myDumbServiceModificationTracker.getModificationCount()
+            + myForcefulReparseModificationTracker.getModificationCount()
+            + myOnContentReloadModificationTracker.getModificationCount())
+  }
+
+
+  override fun dispose() {
+  }
+
+  companion object {
+    private val THROTTLED_LOG = ThrottledLogger(Logger.getInstance(PyLibraryModificationTracker::class.java), TimeUnit.SECONDS.toMillis(30))
+
+    fun getInstance(project: Project): PyLibraryModificationTracker = project.service()
+  }
+}

python/python-psi-api/src/com/jetbrains/python/psi/types/TypeEvalContext.java

@@ -2,10 +2,12 @@
 package com.jetbrains.python.psi.types;
 
 import com.intellij.openapi.project.Project;
+import com.intellij.openapi.roots.ProjectFileIndex;
 import com.intellij.openapi.util.Pair;
 import com.intellij.openapi.util.RecursionManager;
 import com.intellij.openapi.util.registry.Registry;
 import com.intellij.openapi.util.text.StringUtil;
+import com.intellij.openapi.vfs.VirtualFile;
 import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiFile;
 import com.intellij.util.ArrayUtil;
@@ -233,7 +235,28 @@ public sealed class TypeEvalContext {
     return null;
   }
 
+  private static boolean isLibraryElement(@NotNull PsiElement element) {
+    VirtualFile vFile = element.getContainingFile().getOriginalFile().getVirtualFile();
+    return vFile != null && ("pyi".equals(vFile.getExtension()) || ProjectFileIndex.getInstance(element.getProject()).isInLibrary(vFile));
+  }
+
+  private @NotNull TypeEvalContext getLibraryContext(@NotNull Project project) {
+    return project.getService(TypeEvalContextCache.class).getLibraryContext(new LibraryTypeEvalContext(getConstraints()));
+  }
+
+  /**
+   * If true the element's type will be calculated and stored in the long-life context bounded to the PyLibraryModificationTracker.
+   */
+  protected boolean canDelegateToLibraryContext(PyTypedElement element) {
+    return Registry.is("python.use.separated.libraries.type.cache") && isLibraryElement(element);
+  }
+
   public @Nullable PyType getType(final @NotNull PyTypedElement element) {
+    if (canDelegateToLibraryContext(element)) {
+      var context = getLibraryContext(element.getProject());
+      return context.getType(element);
+    }
+
     final PyType knownType = getKnownType(element);
     if (knownType != null) {
       return knownType == PyNullType.INSTANCE ? null : knownType;
@@ -251,6 +274,11 @@ public sealed class TypeEvalContext {
   }
 
   public @Nullable PyType getReturnType(final @NotNull PyCallable callable) {
+    if (canDelegateToLibraryContext(callable)) {
+      var context = getLibraryContext(callable.getProject());
+      return context.getReturnType(callable);
+    }
+
     final PyType knownReturnType = getKnownReturnType(callable);
     if (knownReturnType != null) {
       return knownReturnType == PyNullType.INSTANCE ? null : knownReturnType;
@@ -416,4 +444,16 @@ public sealed class TypeEvalContext {
       return this == o;
     }
   }
+
+  final static class LibraryTypeEvalContext extends TypeEvalContext {
+    private LibraryTypeEvalContext(@NotNull TypeEvalConstraints constraints) {
+      super(constraints);
+    }
+
+    @Override
+    protected boolean canDelegateToLibraryContext(PyTypedElement element) {
+      // It's already the library-context.
+      return false;
+    }
+  }
 }

python/python-psi-api/src/com/jetbrains/python/psi/types/TypeEvalContextCache.java

@@ -36,4 +36,7 @@ public interface TypeEvalContextCache {
    */
   @NotNull
   TypeEvalContext getContext(@NotNull TypeEvalContext standard);
+
+  @NotNull
+  TypeEvalContext getLibraryContext(@NotNull TypeEvalContext standard);
 }