From 2123dbfa15ec2a8c5758b94ffa01a09a378ce460 Mon Sep 17 00:00:00 2001
From: "evgeny.bovykin" <evgeny.bovykin@jetbrains.com>
Date: Tue, 14 Oct 2025 14:54:42 +0200
Subject: [PATCH] Enable `TypeEvalContext#assumeType` only if 
 `python.use.better.control.flow.type.inference` is enabled

GitOrigin-RevId: da7e3452bfbdbfd618dbe166e5a6c199f28cae6f
---
 .../src/com/jetbrains/python/psi/types/TypeEvalContext.java    | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/python/python-psi-api/src/com/jetbrains/python/psi/types/TypeEvalContext.java b/python/python-psi-api/src/com/jetbrains/python/psi/types/TypeEvalContext.java
index aa300b7210ed..0f95a6641912 100644
--- a/python/python-psi-api/src/com/jetbrains/python/psi/types/TypeEvalContext.java
+++ b/python/python-psi-api/src/com/jetbrains/python/psi/types/TypeEvalContext.java
@@ -193,6 +193,9 @@ public sealed class TypeEvalContext {
 
   @ApiStatus.Internal
   public <R> @Nullable R assumeType(@NotNull PyTypedElement element, @Nullable PyType type, @NotNull Function<TypeEvalContext, R> func) {
+    if (!Registry.is("python.use.better.control.flow.type.inference")) {
+      return func.apply(this);
+    }
     if (getKnownType(element) != null) {
       // Temporary solution, as overwriting known type might introduce inconsistencies with its dependencies.
       return null;