SONARPY-2266 TypeAnnotationToPythonTypeConverter creates nested LazyU…

…nionType (#2222)
SonarSource · Dec 9, 2024 · ac3982b · ac3982b
1 parent fceee63
commit ac3982b
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 2 deletions.
diff --git a/python-frontend/src/main/java/org/sonar/python/types/v2/LazyUnionType.java b/python-frontend/src/main/java/org/sonar/python/types/v2/LazyUnionType.java
@@ -16,15 +16,19 @@
  */
 package org.sonar.python.types.v2;
 
+import com.google.common.annotations.VisibleForTesting;
+import java.util.Collections;
 import java.util.HashSet;
 import java.util.Set;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
 
 public class LazyUnionType implements PythonType, ResolvableType {
 
-  private final Set<PythonType> candidates = new HashSet<>();
+  private final Set<PythonType> candidates;
 
   public LazyUnionType(Set<PythonType> candidates) {
-    this.candidates.addAll(candidates);
+    this.candidates = candidates.stream().flatMap(LazyUnionType::flattenLazyUnionTypes).collect(Collectors.toCollection(HashSet::new));
   }
 
   public PythonType resolve() {
@@ -37,4 +41,16 @@ public PythonType resolve() {
     }
     return UnionType.or(resolvedCandidates);
   }
+
+  private static Stream<PythonType> flattenLazyUnionTypes(PythonType type) {
+    if (type instanceof LazyUnionType lazyUnionType) {
+      return lazyUnionType.candidates.stream();
+    }
+    return Stream.of(type);
+  }
+
+  @VisibleForTesting
+  protected Set<PythonType> candidates() {
+    return Collections.unmodifiableSet(candidates);
+  }
 }
diff --git a/python-frontend/src/test/java/org/sonar/python/types/v2/LazyUnionTypeTest.java b/python-frontend/src/test/java/org/sonar/python/types/v2/LazyUnionTypeTest.java
@@ -19,7 +19,9 @@
 import java.util.Set;
 import org.junit.jupiter.api.Test;
 import org.mockito.Mockito;
+import org.sonar.python.semantic.ProjectLevelSymbolTable;
 import org.sonar.python.semantic.v2.LazyTypesContext;
+import org.sonar.python.semantic.v2.ProjectLevelTypeTable;
 
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.mockito.Mockito.when;
@@ -37,4 +39,23 @@ void lazyUnionTypeResolvesNestedLazyTypesWhenAccessed() {
     UnionType unionType = (UnionType) lazyUnionType.resolve();
     assertThat(unionType.candidates()).containsExactlyInAnyOrder(INT_TYPE, FLOAT_TYPE);
   }
+
+  @Test
+  void flattened() {
+    LazyTypesContext lazyTypesContext = Mockito.spy(new LazyTypesContext(new ProjectLevelTypeTable(ProjectLevelSymbolTable.empty())));
+
+    var lazyType1 = lazyTypeUnresolved("lazy1", lazyTypesContext);
+    var lazyType2 = lazyTypeUnresolved("lazy2", lazyTypesContext);
+    var lazyType3 = lazyTypeUnresolved("lazy3", lazyTypesContext);
+
+    var lazyUnionType1 = new LazyUnionType(Set.of(lazyType1, lazyType2, lazyType3));
+    var lazyUnionType2 = new LazyUnionType(Set.of(lazyType1, lazyType2, lazyType3, lazyUnionType1));
+    assertThat(lazyUnionType2.candidates()).containsExactlyInAnyOrder(lazyType1, lazyType2, lazyType3);
+  }
+
+  LazyType lazyTypeUnresolved(String name, LazyTypesContext lazyTypesContext) {
+    var lazyType = lazyTypesContext.getOrCreateLazyType(name);
+    when(lazyTypesContext.resolveLazyType(lazyType)).thenReturn(new UnknownType.UnresolvedImportType(name));
+    return lazyType;
+  }
 }