github653224
diff --git a/‎python-frontend/src/main/java/org/sonar/python/types/InferredTypes.java‎
Lines changed: 15 additions & 0 deletions b/‎python-frontend/src/main/java/org/sonar/python/types/InferredTypes.java‎
Lines changed: 15 additions & 0 deletions
diff --git a/‎python-frontend/src/main/java/org/sonar/python/types/RuntimeType.java‎
Lines changed: 4 additions & 0 deletions b/‎python-frontend/src/main/java/org/sonar/python/types/RuntimeType.java‎
Lines changed: 4 additions & 0 deletions
diff --git a/‎python-frontend/src/test/java/org/sonar/python/types/InferredTypesTest.java‎
Lines changed: 14 additions & 0 deletions b/‎python-frontend/src/test/java/org/sonar/python/types/InferredTypesTest.java‎
Lines changed: 14 additions & 0 deletions
@@ -19,10 +19,13 @@
  */
 package org.sonar.python.types;
 
+import java.util.Collection;
 import java.util.Collections;
 import java.util.HashMap;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 import java.util.stream.Stream;
 import javax.annotation.Nullable;
 import org.sonar.plugins.python.api.symbols.AmbiguousSymbol;
@@ -149,4 +152,16 @@ private static InferredType optionalDeclaredType(List<Expression> subscripts, Ma
     InferredType noneType = InferredTypes.runtimeType(builtinSymbols.get(BuiltinTypes.NONE_TYPE));
     return InferredTypes.or(declaredType(subscripts.get(0), builtinSymbols), noneType);
   }
+
+  public static Collection<ClassSymbol> typeSymbols(InferredType inferredType) {
+    if (inferredType instanceof RuntimeType) {
+      return Collections.singleton(((RuntimeType) inferredType).getTypeClass());
+    }
+    if (inferredType instanceof UnionType) {
+      Set<ClassSymbol> typeClasses = new HashSet<>();
+      ((UnionType) inferredType).types().forEach(type -> typeClasses.addAll(typeSymbols(type)));
+      return typeClasses;
+    }
+    return Collections.emptySet();
+  }
 }
@@ -131,4 +131,8 @@ public int hashCode() {
   public String toString() {
     return "RuntimeType(" + typeClass.fullyQualifiedName() + ')';
   }
+
+  public ClassSymbol getTypeClass() {
+    return typeClass;
+  }
 }
@@ -32,9 +32,12 @@
 import org.sonar.python.semantic.SymbolImpl;
 
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.sonar.python.types.InferredTypes.INT;
+import static org.sonar.python.types.InferredTypes.STR;
 import static org.sonar.python.types.InferredTypes.anyType;
 import static org.sonar.python.types.InferredTypes.or;
 import static org.sonar.python.types.InferredTypes.runtimeType;
+import static org.sonar.python.types.TypeShed.typeShedClass;
 
 public class InferredTypesTest {
 
@@ -123,6 +126,17 @@ public void test_optional_type_annotations() {
     assertThat(InferredTypes.declaredType(typeAnnotation)).isEqualTo(InferredTypes.anyType());
   }
 
+  @Test
+  public void test_typeSymbol() {
+    assertThat(InferredTypes.typeSymbols(STR)).containsExactly(typeShedClass("str"));
+
+    ClassSymbol a = new ClassSymbolImpl("A", "mod.A");
+    assertThat(InferredTypes.typeSymbols(new RuntimeType(a))).containsExactly(a);
+
+    assertThat(InferredTypes.typeSymbols(or(STR, INT))).containsExactlyInAnyOrder(typeShedClass("str"), typeShedClass("int"));
+    assertThat(InferredTypes.typeSymbols(InferredTypes.anyType())).isEmpty();
+  }
+
   private TypeAnnotation typeAnnotation(String... code) {
     return PythonTestUtils.getLastDescendant(PythonTestUtils.parse(code), tree -> tree.is(Tree.Kind.VARIABLE_TYPE_ANNOTATION));
   }
Original file line number	Diff line number	Diff line change
`@@ -131,4 +131,8 @@ public int hashCode() {`
`131`	`131`	`public String toString() {`
`132`	`132`	`return "RuntimeType(" + typeClass.fullyQualifiedName() + ')';`
`133`	`133`	`}`
	`134`	`+`
	`135`	`+ public ClassSymbol getTypeClass() {`
	`136`	`+ return typeClass;`
	`137`	`+ }`
`134`	`138`	`}`