From 7bba13e04cb263dd25902a5996aa6696a180b963 Mon Sep 17 00:00:00 2001
From: Maksim Grebeniuk <maksim.grebeniuk@sonarsource.com>
Date: Fri, 25 Oct 2024 10:40:38 +0200
Subject: [PATCH] CR fixes

---
 .../semantic/v2/types/TrivialTypeInferenceVisitor.java   | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/python-frontend/src/main/java/org/sonar/python/semantic/v2/types/TrivialTypeInferenceVisitor.java b/python-frontend/src/main/java/org/sonar/python/semantic/v2/types/TrivialTypeInferenceVisitor.java
index bfb615d362..ec6783a358 100644
--- a/python-frontend/src/main/java/org/sonar/python/semantic/v2/types/TrivialTypeInferenceVisitor.java
+++ b/python-frontend/src/main/java/org/sonar/python/semantic/v2/types/TrivialTypeInferenceVisitor.java
@@ -471,10 +471,17 @@ public void visitName(Name name) {
       .map(Expression.class::cast)
       .map(Expression::typeV2)
       // TODO: classes (SONARPY-1829) and functions should be propagated like other types
-      .filter(t -> (t instanceof ClassType) || (t instanceof FunctionType) || (t instanceof ModuleType) || (t instanceof UnknownType.UnresolvedImportType))
+      .filter(TrivialTypeInferenceVisitor::shouldTypeBeEagerlyPropagated)
       .ifPresent(type -> setTypeToName(name, type));
   }
 
+  private static boolean shouldTypeBeEagerlyPropagated(PythonType t) {
+    return (t instanceof ClassType)
+           || (t instanceof FunctionType)
+           || (t instanceof ModuleType)
+           || (t instanceof UnknownType.UnresolvedImportType);
+  }
+
   private PythonType currentType() {
     return typeStack.peek();
   }