qualified expression and call expression type inference

SonarSource · Apr 15, 2024 · 23fb013 · 23fb013
1 parent 82cc5cc
commit 23fb013
Show file tree

Hide file tree

Showing 3 changed files with 77 additions and 2 deletions.
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
@@ -21,19 +21,24 @@
 
 import java.util.ArrayDeque;
 import java.util.ArrayList;
+import java.util.Collection;
 import java.util.Deque;
 import java.util.HashMap;
 import java.util.List;
+import java.util.Optional;
 import org.sonar.plugins.python.api.tree.AssignmentStatement;
 import org.sonar.plugins.python.api.tree.BaseTreeVisitor;
+import org.sonar.plugins.python.api.tree.CallExpression;
 import org.sonar.plugins.python.api.tree.ClassDef;
 import org.sonar.plugins.python.api.tree.Expression;
+import org.sonar.plugins.python.api.tree.ExpressionList;
 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.QualifiedExpression;
 import org.sonar.plugins.python.api.tree.StringLiteral;
 import org.sonar.plugins.python.api.types.InferredType;
 import org.sonar.python.tree.ListLiteralImpl;
@@ -46,6 +51,7 @@
 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.ObjectTypeBuilder;
 import org.sonar.python.types.v2.PythonType;
 
 public class TypeInferenceV2 extends BaseTreeVisitor {
@@ -130,9 +136,50 @@ public void visitImportName(ImportName importName) {
 
   @Override
   public void visitAssignmentStatement(AssignmentStatement assignmentStatement) {
+    scan(assignmentStatement.assignedValue());
+    Optional.of(assignmentStatement)
+      .map(AssignmentStatement::lhsExpressions)
+      .filter(lhs -> lhs.size() == 1)
+      .map(lhs -> lhs.get(0))
+      .map(ExpressionList::expressions)
+      .filter(lhs -> lhs.size() == 1)
+      .map(lhs -> lhs.get(0))
+      .filter(NameImpl.class::isInstance)
+      .map(NameImpl.class::cast)
+      .ifPresent(lhsName -> {
+        var assignedValueType = assignmentStatement.assignedValue().typeV2();
+        lhsName.typeV2(assignedValueType);
+      });
+  }
+
+  @Override
+  public void visitQualifiedExpression(QualifiedExpression qualifiedExpression) {
+    scan(qualifiedExpression.qualifier());
+    if (qualifiedExpression.name() instanceof NameImpl name) {
+      var nameType = qualifiedExpression.qualifier().typeV2().resolveMember(qualifiedExpression.name().name());
+      name.typeV2(nameType);
+    }
+  }
 
+  @Override
+  public void visitName(Name name) {
+    var type = Optional.of(name)
+      .map(Name::symbolV2)
+      .map(SymbolV2::usages)
+      .stream()
+      .flatMap(Collection::stream)
+      .filter(UsageV2::isBindingUsage)
+      .map(UsageV2::tree)
+      .filter(Expression.class::isInstance)
+      .map(Expression.class::cast)
+      .map(Expression::typeV2)
+      .toList();
+
+    if (type.size() == 1 && name instanceof NameImpl nameImpl) {
+      nameImpl.typeV2(type.get(0));
+    }
   }
-  
+
   private PythonType currentType() {
     return typeStack.peek();
   }

diff --git a/python-frontend/src/main/java/org/sonar/python/tree/CallExpressionImpl.java b/python-frontend/src/main/java/org/sonar/python/tree/CallExpressionImpl.java
@@ -46,6 +46,10 @@
 import org.sonar.python.types.DeclaredType;
 import org.sonar.python.types.HasTypeDependencies;
 import org.sonar.python.types.InferredTypes;
+import org.sonar.python.types.v2.ClassType;
+import org.sonar.python.types.v2.FunctionType;
+import org.sonar.python.types.v2.ObjectType;
+import org.sonar.python.types.v2.PythonType;
 
 import static org.sonar.plugins.python.api.symbols.Symbol.Kind.CLASS;
 import static org.sonar.plugins.python.api.tree.Tree.Kind.SUBSCRIPTION;
@@ -179,4 +183,15 @@ private static InferredType getType(Symbol symbol) {
   public List<Expression> typeDependencies() {
     return Collections.singletonList(callee);
   }
+
+  @Override
+  public PythonType typeV2() {
+    if (callee().typeV2() instanceof ClassType classType) {
+      return new ObjectType(classType);
+    }
+    if (callee().typeV2() instanceof FunctionType functionType) {
+      return new ObjectType(functionType.returnType());
+    }
+    return PythonType.UNKNOWN;
+  }
 }
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
@@ -19,10 +19,16 @@
  */
 package org.sonar.python.types.v2;
 
+import java.util.ArrayList;
 import java.util.List;
+import java.util.Objects;
 
 public record ObjectType(PythonType type, List<PythonType> attributes, List<Member> members) implements PythonType {
 
+  public ObjectType(PythonType type) {
+    this(type, new ArrayList<>(), new ArrayList<>());
+  }
+
   @Override
   public String displayName() {
     return type.displayName();
@@ -33,5 +39,12 @@ public boolean isCompatibleWith(PythonType another) {
     return this.type.isCompatibleWith(another);
   }
 
-
+  @Override
+  public PythonType resolveMember(String memberName) {
+    return members().stream()
+      .filter(member -> Objects.equals(member.name(), memberName))
+      .map(Member::type)
+      .findFirst()
+      .orElseGet(() -> type.resolveMember(memberName));
+  }
 }