Basic type inference for list literals

SonarSource · Apr 11, 2024 · f3c30a6 · f3c30a6
1 parent 31aef37
commit f3c30a6
Show file tree

Hide file tree

Showing 8 changed files with 65 additions and 5 deletions.
diff --git a/python-frontend/src/main/java/org/sonar/python/semantic/v2/ProjectLevelTypeTable.java b/python-frontend/src/main/java/org/sonar/python/semantic/v2/ProjectLevelTypeTable.java
@@ -2,15 +2,23 @@
 
 import java.util.Collections;
 import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
 import java.util.Set;
 import java.util.function.Function;
 import java.util.stream.Collectors;
+import org.sonar.plugins.python.api.symbols.AmbiguousSymbol;
+import org.sonar.plugins.python.api.symbols.ClassSymbol;
+import org.sonar.plugins.python.api.symbols.FunctionSymbol;
 import org.sonar.plugins.python.api.symbols.Symbol;
 import org.sonar.python.semantic.ProjectLevelSymbolTable;
 import org.sonar.python.types.TypeShed;
+import org.sonar.python.types.v2.ClassType;
+import org.sonar.python.types.v2.FunctionType;
+import org.sonar.python.types.v2.Member;
 import org.sonar.python.types.v2.ModuleType;
 import org.sonar.python.types.v2.PythonType;
+import org.sonar.python.types.v2.UnionType;
 
 public class ProjectLevelTypeTable {
 
@@ -48,6 +56,31 @@ public ModuleType getModule(String moduleName) {
   }
 
   private PythonType convertToType(Symbol symbol) {
+    return switch (symbol.kind()) {
+      case CLASS -> converToClassType((ClassSymbol) symbol);
+      case FUNCTION -> convertToFunctionType((FunctionSymbol) symbol);
+      case AMBIGUOUS -> convertToUnionType((AmbiguousSymbol) symbol);
+      case OTHER -> converToObjectType(symbol);
+    };
+  }
+
+  private PythonType converToObjectType(Symbol symbol) {
+    // What should we have here?
     return PythonType.UNKNOWN;
   }
+
+  private PythonType convertToFunctionType(FunctionSymbol symbol) {
+    return new FunctionType(symbol.name(), List.of(), List.of(), List.of(), List.of(), PythonType.UNKNOWN);
+  }
+
+  private PythonType converToClassType(ClassSymbol symbol) {
+    Set<Member> members = symbol.declaredMembers().stream().map(m -> new Member(m.name(), convertToType(m))).collect(Collectors.toSet());
+    List<PythonType> superClasses = symbol.superClasses().stream().map(this::convertToType).toList();
+    return new ClassType(symbol.name(), members, List.of(), superClasses, List.of());
+  }
+
+  private PythonType convertToUnionType(AmbiguousSymbol ambiguousSymbol) {
+    List<PythonType> pythonTypes = ambiguousSymbol.alternatives().stream().map(this::convertToType).toList();
+    return new UnionType(pythonTypes);
+  }
 }
diff --git a/python-frontend/src/main/java/org/sonar/python/semantic/v2/TypeInferenceV2.java b/python-frontend/src/main/java/org/sonar/python/semantic/v2/TypeInferenceV2.java
@@ -4,15 +4,20 @@
 import java.util.ArrayList;
 import java.util.Deque;
 import java.util.HashMap;
+import java.util.List;
+import org.sonar.plugins.python.api.tree.AssignmentStatement;
 import org.sonar.plugins.python.api.tree.BaseTreeVisitor;
 import org.sonar.plugins.python.api.tree.ClassDef;
+import org.sonar.plugins.python.api.tree.Expression;
 import org.sonar.plugins.python.api.tree.FileInput;
 import org.sonar.plugins.python.api.tree.FunctionDef;
 import org.sonar.plugins.python.api.tree.ImportName;
+import org.sonar.plugins.python.api.tree.ListLiteral;
 import org.sonar.plugins.python.api.tree.Name;
 import org.sonar.plugins.python.api.tree.NumericLiteral;
 import org.sonar.plugins.python.api.tree.StringLiteral;
 import org.sonar.plugins.python.api.types.InferredType;
+import org.sonar.python.tree.ListLiteralImpl;
 import org.sonar.python.tree.NameImpl;
 import org.sonar.python.tree.NumericLiteralImpl;
 import org.sonar.python.tree.StringLiteralImpl;
@@ -21,6 +26,7 @@
 import org.sonar.python.types.v2.FunctionType;
 import org.sonar.python.types.v2.Member;
 import org.sonar.python.types.v2.ModuleType;
+import org.sonar.python.types.v2.ObjectType;
 import org.sonar.python.types.v2.PythonType;
 
 public class TypeInferenceV2 extends BaseTreeVisitor {
@@ -43,7 +49,7 @@ public void visitFileInput(FileInput fileInput) {
   public void visitStringLiteral(StringLiteral stringLiteral) {
     ModuleType builtins = this.projectLevelTypeTable.getModule("builtins");
     // TODO: multiple object types to represent str instance?
-    ((StringLiteralImpl) stringLiteral).typeV2(builtins.resolveMember("str"));
+    ((StringLiteralImpl) stringLiteral).typeV2(new ObjectType(builtins.resolveMember("str"), List.of(), List.of()));
   }
 
   @Override
@@ -52,10 +58,19 @@ public void visitNumericLiteral(NumericLiteral numericLiteral) {
     InferredType type = numericLiteral.type();
     String memberName = ((RuntimeType) type).getTypeClass().fullyQualifiedName();
     if (memberName != null) {
-      ((NumericLiteralImpl) numericLiteral).typeV2(builtins.resolveMember(memberName));
+      ((NumericLiteralImpl) numericLiteral).typeV2(new ObjectType(builtins.resolveMember(memberName), List.of(), List.of()));
     }
   }
 
+  @Override
+  public void visitListLiteral(ListLiteral listLiteral) {
+    ModuleType builtins = this.projectLevelTypeTable.getModule("builtins");
+    scan(listLiteral.elements());
+    List<PythonType> pythonTypes = listLiteral.elements().expressions().stream().map(Expression::typeV2).distinct().toList();
+    // TODO: cleanly reduce attributes
+    ((ListLiteralImpl) listLiteral).typeV2(new ObjectType(builtins.resolveMember("list"), pythonTypes, List.of()));
+  }
+
   @Override
   public void visitClassDef(ClassDef classDef) {
     scan(classDef.args());

diff --git a/python-frontend/src/main/java/org/sonar/python/tree/ListLiteralImpl.java b/python-frontend/src/main/java/org/sonar/python/tree/ListLiteralImpl.java
@@ -28,13 +28,16 @@
 import org.sonar.plugins.python.api.tree.TreeVisitor;
 import org.sonar.plugins.python.api.types.InferredType;
 import org.sonar.python.types.InferredTypes;
+import org.sonar.python.types.v2.PythonType;
 
 public class ListLiteralImpl extends PyTree implements ListLiteral {
 
   private final Token leftBracket;
   private final ExpressionList elements;
   private final Token rightBracket;
 
+  private PythonType typeV2;
+
   public ListLiteralImpl(Token leftBracket, ExpressionList elements, Token rightBracket) {
     this.leftBracket = leftBracket;
     this.elements = elements;
@@ -75,4 +78,12 @@ public List<Tree> computeChildren() {
   public InferredType type() {
     return InferredTypes.LIST;
   }
+
+  public PythonType typeV2() {
+    return this.typeV2;
+  }
+
+  public void typeV2(PythonType type) {
+    this.typeV2 = type;
+  }
 }
diff --git a/python-frontend/src/main/java/org/sonar/python/tree/StringLiteralImpl.java b/python-frontend/src/main/java/org/sonar/python/tree/StringLiteralImpl.java
@@ -31,7 +31,6 @@
 import org.sonar.plugins.python.api.tree.TreeVisitor;
 import org.sonar.plugins.python.api.types.InferredType;
 import org.sonar.python.types.InferredTypes;
-import org.sonar.python.types.v2.ObjectType;
 import org.sonar.python.types.v2.PythonType;
 
 public class StringLiteralImpl extends PyTree implements StringLiteral {

diff --git a/python-frontend/src/main/java/org/sonar/python/types/v2/FunctionType.java b/python-frontend/src/main/java/org/sonar/python/types/v2/FunctionType.java
@@ -31,4 +31,5 @@ public record FunctionType(
   List<PythonType> typeVars,
   List<Member> parameters,
   PythonType returnType) implements PythonType {
+  // TODO: Decorators? Parameters should express keyword/positional-only information - Difference between attributes and typeVars?
 }
diff --git a/python-frontend/src/main/java/org/sonar/python/types/v2/ObjectType.java b/python-frontend/src/main/java/org/sonar/python/types/v2/ObjectType.java
@@ -20,7 +20,6 @@
 package org.sonar.python.types.v2;
 
 import java.util.List;
-import org.sonar.plugins.python.api.symbols.ClassSymbol;
 
 public record ObjectType(PythonType type, List<PythonType> attributes, List<Member> members) implements PythonType {
 

diff --git a/python-frontend/src/main/java/org/sonar/python/types/v2/UnionType.java b/python-frontend/src/main/java/org/sonar/python/types/v2/UnionType.java
@@ -28,7 +28,7 @@ public UnionType() {
 
   @Override
   public String displayName() {
-    var candidatesName = candidates.stream().map(c -> c.displayName()).toList();
+    var candidatesName = candidates.stream().map(PythonType::displayName).toList();
     return "Union[%s]".formatted(String.join(", ", candidatesName));
   }
 

diff --git a/python-frontend/src/test/resources/semantic/v2/script.py b/python-frontend/src/test/resources/semantic/v2/script.py
@@ -15,3 +15,5 @@ def script_do_something(param):
 
 
 z = "hello"
+
+my_list = ["a", "b", "c"]