PY-81269 Extremely High RAM Usage

GitOrigin-RevId: bb534d62931c176d2395e2cceb0afabc57464669

python/python-psi-api/src/com/jetbrains/python/psi/PyInstantTypeProvider.kt

@@ -0,0 +1,7 @@
+// 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
+
+import org.jetbrains.annotations.ApiStatus
+
+@ApiStatus.Internal
+interface PyInstantTypeProvider : PyTypedElement

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

@@ -118,7 +118,7 @@ public sealed class TypeEvalContext {
    * Inspections should not create a new type evaluation context. They should re-use the context of the inspection session.
    */
   public static TypeEvalContext codeAnalysis(final @NotNull Project project, final @Nullable PsiFile origin) {
-    return getContextFromCache(project, new TypeEvalContext(false, false, false, origin));
+    return getContextFromCache(project, buildCodeAnalysisContext(origin));
   }
 
   /**
@@ -145,6 +145,13 @@ public sealed class TypeEvalContext {
     return getContextFromCache(project, new TypeEvalContext(false, true, false, null));
   }
 
+  private static TypeEvalContext buildCodeAnalysisContext(@Nullable PsiFile origin) {
+    if (Registry.is("python.optimized.type.eval.context")) {
+      return new OptimizedTypeEvalContext(false, false, false, origin);
+    }
+    return new TypeEvalContext(false, false, false, origin);
+  }
+
   /**
    * Moves context through cache returning one from cache (if exists).
    *
@@ -213,7 +220,7 @@ public sealed class TypeEvalContext {
   }
 
   protected @Nullable PyType getKnownType(final @NotNull PyTypedElement element) {
-    if (element instanceof PyNoneLiteralExpression) {
+    if (element instanceof PyInstantTypeProvider) {
       return element.getType(this, Key.INSTANCE);
     }
     final PyType cachedType = myEvaluated.get(element);
@@ -416,4 +423,64 @@ public sealed class TypeEvalContext {
       return this == o;
     }
   }
+
+  final static class OptimizedTypeEvalContext extends TypeEvalContext {
+    private volatile TypeEvalContext codeInsightFallback;
+
+    OptimizedTypeEvalContext(boolean allowDataFlow, boolean allowStubToAST, boolean allowCallContext, @Nullable PsiFile origin) {
+      super(allowDataFlow, allowStubToAST, allowCallContext, origin);
+    }
+
+    private boolean shouldSwitchToFallbackContext(PsiElement element) {
+      PsiFile file = element.getContainingFile();
+      if (file instanceof PyExpressionCodeFragment codeFragment) {
+        PsiElement context = codeFragment.getContext();
+        if (context != null) {
+          file = context.getContainingFile();
+        }
+      }
+      TypeEvalConstraints constraints = getConstraints();
+      return constraints.myOrigin != null && file != constraints.myOrigin && (file instanceof PyFile) &&
+             !constraints.myAllowDataFlow && !constraints.myAllowStubToAST && !constraints.myAllowCallContext;
+    }
+
+    private TypeEvalContext getFallbackContext(Project project) {
+      if (codeInsightFallback == null) {
+        codeInsightFallback = codeInsightFallback(project);
+      }
+      return codeInsightFallback;
+    }
+
+    @Override
+    protected @Nullable PyType getKnownType(@NotNull PyTypedElement element) {
+      if (shouldSwitchToFallbackContext(element)) {
+        return getFallbackContext(element.getProject()).getKnownType(element);
+      }
+      return super.getKnownType(element);
+    }
+
+    @Override
+    protected @Nullable PyType getKnownReturnType(@NotNull PyCallable callable) {
+      if (shouldSwitchToFallbackContext(callable)) {
+        return getFallbackContext(callable.getProject()).getKnownReturnType(callable);
+      }
+      return super.getKnownReturnType(callable);
+    }
+
+    @Override
+    public @Nullable PyType getType(@NotNull PyTypedElement element) {
+      if (shouldSwitchToFallbackContext(element)) {
+        return getFallbackContext(element.getProject()).getType(element);
+      }
+      return super.getType(element);
+    }
+
+    @Override
+    public @Nullable PyType getReturnType(@NotNull PyCallable callable) {
+      if (shouldSwitchToFallbackContext(callable)) {
+        return getFallbackContext(callable.getProject()).getReturnType(callable);
+      }
+      return super.getReturnType(callable);
+    }
+  }
 }

python/python-psi-impl/resources/intellij.python.psi.impl.xml

@@ -492,6 +492,8 @@
                  description="Require marking namespace packages explicitly, treat regular directories as implicit source roots"/>
     <registryKey key="python.type.hints.literal.string" defaultValue="true"
                  description="When enabled, activates LiteralString inference for Python string literals" />
+    <registryKey key="python.optimized.type.eval.context" defaultValue="true"
+                 description="It optimises cache usage for storing calculated types."/>
     <registryKey key="python.statement.lists.incremental.reparse" defaultValue="false"
                  description="Enables incremental reparse for statement lists"/>
   </extensions>

python/python-psi-impl/src/com/jetbrains/python/psi/impl/PyBoolLiteralExpressionImpl.java

@@ -4,12 +4,13 @@ package com.jetbrains.python.psi.impl;
 import com.intellij.lang.ASTNode;
 import com.jetbrains.python.psi.PyBoolLiteralExpression;
 import com.jetbrains.python.psi.PyElementVisitor;
+import com.jetbrains.python.psi.PyInstantTypeProvider;
 import com.jetbrains.python.psi.types.PyType;
 import com.jetbrains.python.psi.types.TypeEvalContext;
 import org.jetbrains.annotations.NotNull;
 
 
-public class PyBoolLiteralExpressionImpl extends PyElementImpl implements PyBoolLiteralExpression {
+public class PyBoolLiteralExpressionImpl extends PyElementImpl implements PyBoolLiteralExpression, PyInstantTypeProvider {
   public PyBoolLiteralExpressionImpl(ASTNode astNode) {
     super(astNode);
   }

python/python-psi-impl/src/com/jetbrains/python/psi/impl/PyDoubleStarExpressionImpl.java

@@ -18,12 +18,13 @@ package com.jetbrains.python.psi.impl;
 import com.intellij.lang.ASTNode;
 import com.jetbrains.python.psi.PyDoubleStarExpression;
 import com.jetbrains.python.psi.PyElementVisitor;
+import com.jetbrains.python.psi.PyInstantTypeProvider;
 import com.jetbrains.python.psi.types.PyType;
 import com.jetbrains.python.psi.types.TypeEvalContext;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
-public class PyDoubleStarExpressionImpl extends PyElementImpl implements PyDoubleStarExpression {
+public class PyDoubleStarExpressionImpl extends PyElementImpl implements PyDoubleStarExpression, PyInstantTypeProvider {
   public PyDoubleStarExpressionImpl(ASTNode astNode) {
     super(astNode);
   }

python/python-psi-impl/src/com/jetbrains/python/psi/impl/PyDoubleStarPatternImpl.java

@@ -3,12 +3,13 @@ package com.jetbrains.python.psi.impl;
 import com.intellij.lang.ASTNode;
 import com.jetbrains.python.psi.PyDoubleStarPattern;
 import com.jetbrains.python.psi.PyElementVisitor;
+import com.jetbrains.python.psi.PyInstantTypeProvider;
 import com.jetbrains.python.psi.types.PyType;
 import com.jetbrains.python.psi.types.TypeEvalContext;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
-public class PyDoubleStarPatternImpl extends PyElementImpl implements PyDoubleStarPattern {
+public class PyDoubleStarPatternImpl extends PyElementImpl implements PyDoubleStarPattern, PyInstantTypeProvider {
   public PyDoubleStarPatternImpl(ASTNode astNode) {
     super(astNode);
   }

python/python-psi-impl/src/com/jetbrains/python/psi/impl/PyEllipsisLiteralExpressionImpl.kt

@@ -18,10 +18,11 @@ package com.jetbrains.python.psi.impl
 import com.intellij.lang.ASTNode
 import com.jetbrains.python.psi.PyElementVisitor
 import com.jetbrains.python.psi.PyEllipsisLiteralExpression
+import com.jetbrains.python.psi.PyInstantTypeProvider
 import com.jetbrains.python.psi.types.PyClassType
 import com.jetbrains.python.psi.types.TypeEvalContext
 
-class PyEllipsisLiteralExpressionImpl(astNode: ASTNode?) : PyElementImpl(astNode), PyEllipsisLiteralExpression {
+class PyEllipsisLiteralExpressionImpl(astNode: ASTNode?) : PyElementImpl(astNode), PyEllipsisLiteralExpression, PyInstantTypeProvider {
   override fun getType(context: TypeEvalContext, key: TypeEvalContext.Key): PyClassType? {
     return PyBuiltinCache.getInstance(this).ellipsisType
   }

python/python-psi-impl/src/com/jetbrains/python/psi/impl/PyEmptyExpressionImpl.java

@@ -3,12 +3,13 @@ package com.jetbrains.python.psi.impl;
 
 import com.intellij.lang.ASTNode;
 import com.jetbrains.python.psi.PyEmptyExpression;
+import com.jetbrains.python.psi.PyInstantTypeProvider;
 import com.jetbrains.python.psi.types.PyType;
 import com.jetbrains.python.psi.types.TypeEvalContext;
 import org.jetbrains.annotations.NotNull;
 
 
-public class PyEmptyExpressionImpl extends PyElementImpl implements PyEmptyExpression {
+public class PyEmptyExpressionImpl extends PyElementImpl implements PyEmptyExpression, PyInstantTypeProvider {
   public PyEmptyExpressionImpl(ASTNode astNode) {
     super(astNode);
   }

python/python-psi-impl/src/com/jetbrains/python/psi/impl/PyKeyValuePatternImpl.java

@@ -2,12 +2,13 @@ package com.jetbrains.python.psi.impl;
 
 import com.intellij.lang.ASTNode;
 import com.jetbrains.python.psi.PyElementVisitor;
+import com.jetbrains.python.psi.PyInstantTypeProvider;
 import com.jetbrains.python.psi.PyKeyValuePattern;
 import com.jetbrains.python.psi.types.PyType;
 import com.jetbrains.python.psi.types.TypeEvalContext;
 import org.jetbrains.annotations.NotNull;
 
-public class PyKeyValuePatternImpl extends PyElementImpl implements PyKeyValuePattern {
+public class PyKeyValuePatternImpl extends PyElementImpl implements PyKeyValuePattern, PyInstantTypeProvider {
   public PyKeyValuePatternImpl(ASTNode astNode) {
     super(astNode);
   }

python/python-psi-impl/src/com/jetbrains/python/psi/impl/PyNoneLiteralExpressionImpl.kt

@@ -17,11 +17,12 @@ package com.jetbrains.python.psi.impl
 
 import com.intellij.lang.ASTNode
 import com.jetbrains.python.psi.PyElementVisitor
+import com.jetbrains.python.psi.PyInstantTypeProvider
 import com.jetbrains.python.psi.PyNoneLiteralExpression
 import com.jetbrains.python.psi.types.PyType
 import com.jetbrains.python.psi.types.TypeEvalContext
 
-class PyNoneLiteralExpressionImpl(astNode: ASTNode) : PyElementImpl(astNode), PyNoneLiteralExpression {
+class PyNoneLiteralExpressionImpl(astNode: ASTNode) : PyElementImpl(astNode), PyNoneLiteralExpression, PyInstantTypeProvider {
   override fun getType(context: TypeEvalContext, key: TypeEvalContext.Key): PyType? {
     return PyBuiltinCache.getInstance(this).noneType
   }

python/python-psi-impl/src/com/jetbrains/python/psi/impl/PyNumericLiteralExpressionImpl.java

@@ -19,6 +19,7 @@ import com.intellij.lang.ASTNode;
 import com.intellij.psi.tree.IElementType;
 import com.jetbrains.python.PyElementTypes;
 import com.jetbrains.python.psi.PyElementVisitor;
+import com.jetbrains.python.psi.PyInstantTypeProvider;
 import com.jetbrains.python.psi.PyNumericLiteralExpression;
 import com.jetbrains.python.psi.types.PyType;
 import com.jetbrains.python.psi.types.TypeEvalContext;
@@ -26,7 +27,7 @@ import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
 
-public class PyNumericLiteralExpressionImpl extends PyElementImpl implements PyNumericLiteralExpression {
+public class PyNumericLiteralExpressionImpl extends PyElementImpl implements PyNumericLiteralExpression, PyInstantTypeProvider {
 
   public PyNumericLiteralExpressionImpl(@NotNull ASTNode astNode) {
     super(astNode);

python/python-psi-impl/src/com/jetbrains/python/psi/impl/PyReprExpressionImpl.java

@@ -3,13 +3,14 @@ package com.jetbrains.python.psi.impl;
 
 import com.intellij.lang.ASTNode;
 import com.jetbrains.python.psi.PyElementVisitor;
+import com.jetbrains.python.psi.PyInstantTypeProvider;
 import com.jetbrains.python.psi.PyReprExpression;
 import com.jetbrains.python.psi.types.PyType;
 import com.jetbrains.python.psi.types.TypeEvalContext;
 import org.jetbrains.annotations.NotNull;
 
 
-public class PyReprExpressionImpl extends PyElementImpl implements PyReprExpression {
+public class PyReprExpressionImpl extends PyElementImpl implements PyReprExpression, PyInstantTypeProvider {
   public PyReprExpressionImpl(ASTNode astNode) {
     super(astNode);
   }

python/python-psi-impl/src/com/jetbrains/python/psi/impl/PySingleStarPatternImpl.java

@@ -9,7 +9,7 @@ import org.jetbrains.annotations.Nullable;
 import java.util.Collections;
 import java.util.List;
 
-public class PySingleStarPatternImpl extends PyElementImpl implements PySingleStarPattern {
+public class PySingleStarPatternImpl extends PyElementImpl implements PySingleStarPattern, PyInstantTypeProvider {
   public PySingleStarPatternImpl(ASTNode astNode) {
     super(astNode);
   }

python/python-psi-impl/src/com/jetbrains/python/psi/impl/PySliceItemImpl.java

@@ -2,6 +2,7 @@
 package com.jetbrains.python.psi.impl;
 
 import com.intellij.lang.ASTNode;
+import com.jetbrains.python.psi.PyInstantTypeProvider;
 import com.jetbrains.python.psi.PySliceItem;
 import com.jetbrains.python.psi.types.PyType;
 import com.jetbrains.python.psi.types.TypeEvalContext;
@@ -9,7 +10,7 @@ import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
 
-public class PySliceItemImpl extends PyElementImpl implements PySliceItem {
+public class PySliceItemImpl extends PyElementImpl implements PySliceItem, PyInstantTypeProvider {
   public PySliceItemImpl(ASTNode astNode) {
     super(astNode);
   }

python/python-psi-impl/src/com/jetbrains/python/psi/impl/PyStarArgumentImpl.java

@@ -2,13 +2,14 @@
 package com.jetbrains.python.psi.impl;
 
 import com.intellij.lang.ASTNode;
+import com.jetbrains.python.psi.PyInstantTypeProvider;
 import com.jetbrains.python.psi.PyStarArgument;
 import com.jetbrains.python.psi.types.PyType;
 import com.jetbrains.python.psi.types.TypeEvalContext;
 import org.jetbrains.annotations.NotNull;
 
 
-public class PyStarArgumentImpl extends PyElementImpl implements PyStarArgument {
+public class PyStarArgumentImpl extends PyElementImpl implements PyStarArgument, PyInstantTypeProvider {
   public PyStarArgumentImpl(ASTNode astNode) {
     super(astNode);
   }

python/python-psi-impl/src/com/jetbrains/python/psi/impl/PyStarExpressionImpl.java

@@ -3,12 +3,13 @@ package com.jetbrains.python.psi.impl;
 
 import com.intellij.lang.ASTNode;
 import com.jetbrains.python.psi.PyElementVisitor;
+import com.jetbrains.python.psi.PyInstantTypeProvider;
 import com.jetbrains.python.psi.PyStarExpression;
 import com.jetbrains.python.psi.types.PyType;
 import com.jetbrains.python.psi.types.TypeEvalContext;
 import org.jetbrains.annotations.NotNull;
 
-public class PyStarExpressionImpl extends PyElementImpl implements PyStarExpression {
+public class PyStarExpressionImpl extends PyElementImpl implements PyStarExpression, PyInstantTypeProvider {
   public PyStarExpressionImpl(ASTNode astNode) {
     super(astNode);
   }