Skip to content

Commit f2e7cf5

Browse files
andrea-guarino-sonarsourceandrea-guarino-sonarsource
authored
Type inference: consider parameter type annotation (SonarSource#830)
1 parent a8e735d commit f2e7cf5

3 files changed

Lines changed: 65 additions & 15 deletions

File tree

‎python-frontend/src/main/java/org/sonar/python/types/FlowSensitiveTypeInference.java‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -45,19 +45,22 @@ class FlowSensitiveTypeInference extends ForwardAnalysis {
4545
private final Set<Symbol> trackedVars;
4646
private final Map<QualifiedExpression, MemberAccess> memberAccessesByQualifiedExpr;
4747
private final Map<AssignmentStatement, Assignment> assignmentsByAssignmentStatement;
48+
private final Map<String, InferredType> parameterTypesByName;
4849

4950
public FlowSensitiveTypeInference(Set<Symbol> trackedVars, Map<QualifiedExpression, MemberAccess> memberAccessesByQualifiedExpr,
50-
Map<AssignmentStatement, Assignment> assignmentsByAssignmentStatement) {
51+
Map<AssignmentStatement, Assignment> assignmentsByAssignmentStatement, Map<String, InferredType> parameterTypesByName) {
5152
this.trackedVars = trackedVars;
5253
this.memberAccessesByQualifiedExpr = memberAccessesByQualifiedExpr;
5354
this.assignmentsByAssignmentStatement = assignmentsByAssignmentStatement;
55+
this.parameterTypesByName = parameterTypesByName;
5456
}
5557

5658
@Override
5759
public ProgramState initialState() {
5860
TypeInferenceProgramState initialState = new TypeInferenceProgramState();
5961
for (Symbol variable : trackedVars) {
60-
initialState.setTypes(variable, Collections.emptySet());
62+
InferredType inferredType = parameterTypesByName.get(variable.name());
63+
initialState.setTypes(variable, inferredType != null ? Collections.singleton(inferredType) : Collections.emptySet());
6164
}
6265
return initialState;
6366
}

‎python-frontend/src/main/java/org/sonar/python/types/TypeInference.java‎

Lines changed: 32 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@
3232
import java.util.stream.Collectors;
3333
import org.sonar.plugins.python.api.PythonFile;
3434
import org.sonar.plugins.python.api.cfg.ControlFlowGraph;
35+
import org.sonar.plugins.python.api.symbols.FunctionSymbol;
3536
import org.sonar.plugins.python.api.symbols.Symbol;
3637
import org.sonar.plugins.python.api.symbols.Usage;
3738
import org.sonar.plugins.python.api.tree.AssignmentStatement;
@@ -41,12 +42,15 @@
4142
import org.sonar.plugins.python.api.tree.FileInput;
4243
import org.sonar.plugins.python.api.tree.FunctionDef;
4344
import org.sonar.plugins.python.api.tree.Name;
45+
import org.sonar.plugins.python.api.tree.Parameter;
4446
import org.sonar.plugins.python.api.tree.QualifiedExpression;
4547
import org.sonar.plugins.python.api.tree.Tree;
4648
import org.sonar.plugins.python.api.tree.TryStatement;
4749
import org.sonar.plugins.python.api.types.InferredType;
4850
import org.sonar.python.semantic.SymbolImpl;
51+
import org.sonar.python.tree.FunctionDefImpl;
4952
import org.sonar.python.tree.NameImpl;
53+
import org.sonar.python.tree.TreeUtils;
5054

5155
public class TypeInference extends BaseTreeVisitor {
5256

@@ -58,6 +62,7 @@ public class TypeInference extends BaseTreeVisitor {
5862
private final Map<Symbol, Set<Assignment>> assignmentsByLhs = new HashMap<>();
5963
private final Map<QualifiedExpression, MemberAccess> memberAccessesByQualifiedExpr = new HashMap<>();
6064
private final Map<AssignmentStatement, Assignment> assignmentsByAssignmentStatement = new HashMap<>();
65+
private Map<String, InferredType> parameterTypesByName = new HashMap<>();
6166

6267
public static void inferTypes(FileInput fileInput, PythonFile pythonFile) {
6368
fileInput.accept(new BaseTreeVisitor() {
@@ -67,7 +72,6 @@ public void visitFunctionDef(FunctionDef funcDef) {
6772
inferTypesAndMemberAccessSymbols(funcDef, pythonFile);
6873
}
6974
});
70-
7175
fileInput.accept(new BaseTreeVisitor() {
7276
@Override
7377
public void visitQualifiedExpression(QualifiedExpression qualifiedExpression) {
@@ -82,28 +86,32 @@ public void visitQualifiedExpression(QualifiedExpression qualifiedExpression) {
8286
});
8387
}
8488

85-
private static void inferTypesAndMemberAccessSymbols(FunctionDef functionDef, PythonFile pythonFile) {
86-
TypeInference visitor = new TypeInference();
87-
functionDef.accept(visitor);
89+
private static Set<Symbol> getTrackedVars(Set<Symbol> localVariables, Set<Name> assignedNames) {
8890
Set<Symbol> trackedVars = new HashSet<>();
89-
Set<Name> assignedNames = visitor.assignmentsByLhs.values().stream()
90-
.flatMap(Collection::stream)
91-
.map(a -> a.lhsName)
92-
.collect(Collectors.toSet());
93-
for (Symbol variable : functionDef.localVariables()) {
91+
for (Symbol variable : localVariables) {
9492
boolean hasMissingBindingUsage = variable.usages().stream()
9593
.filter(Usage::isBindingUsage)
9694
.anyMatch(u -> !assignedNames.contains(u.tree()));
9795
if (!hasMissingBindingUsage) {
9896
trackedVars.add(variable);
9997
}
10098
}
99+
return trackedVars;
100+
}
101+
102+
private static void inferTypesAndMemberAccessSymbols(FunctionDef functionDef, PythonFile pythonFile) {
103+
TypeInference visitor = new TypeInference();
104+
functionDef.accept(visitor);
105+
Set<Name> assignedNames = visitor.assignmentsByLhs.values().stream()
106+
.flatMap(Collection::stream)
107+
.map(a -> a.lhsName)
108+
.collect(Collectors.toSet());
101109

102110
TryStatementVisitor tryStatementVisitor = new TryStatementVisitor(functionDef);
103111
functionDef.body().accept(tryStatementVisitor);
104112
if (tryStatementVisitor.hasTryStatement) {
105113
// CFG doesn't model precisely try-except statements. Hence we fallback to AST based type inference
106-
visitor.processPropagations(trackedVars);
114+
visitor.processPropagations(getTrackedVars(functionDef.localVariables(), assignedNames));
107115
functionDef.body().accept(new BaseTreeVisitor() {
108116
@Override
109117
public void visitFunctionDef(FunctionDef visited) {
@@ -122,7 +130,12 @@ public void visitName(Name name) {
122130
if (cfg == null) {
123131
return;
124132
}
125-
visitor.flowSensitiveTypeInference(cfg, trackedVars);
133+
Set<Name> annotatedParamNames = TreeUtils.nonTupleParameters(functionDef).stream()
134+
.filter(parameter -> parameter.typeAnnotation() != null)
135+
.map(Parameter::name)
136+
.collect(Collectors.toSet());
137+
assignedNames.addAll(annotatedParamNames);
138+
visitor.flowSensitiveTypeInference(cfg, getTrackedVars(functionDef.localVariables(), assignedNames), functionDef);
126139
}
127140
}
128141

@@ -198,9 +211,15 @@ public void visitQualifiedExpression(QualifiedExpression qualifiedExpression) {
198211
* if i > 0: b = a.capitalize() # at the end of second execution, type of b is inferred to be "STR"
199212
* else: a = 'abc'
200213
*/
201-
private void flowSensitiveTypeInference(ControlFlowGraph cfg, Set<Symbol> trackedVars) {
214+
private void flowSensitiveTypeInference(ControlFlowGraph cfg, Set<Symbol> trackedVars, FunctionDef functionDef) {
215+
Optional.ofNullable(((FunctionDefImpl) functionDef).functionSymbol()).ifPresent(functionSymbol ->
216+
parameterTypesByName = functionSymbol.parameters()
217+
.stream()
218+
.filter(parameter -> parameter.name() != null)
219+
.collect(Collectors.toMap(FunctionSymbol.Parameter::name, FunctionSymbol.Parameter::declaredType)));
220+
202221
FlowSensitiveTypeInference flowSensitiveTypeInference =
203-
new FlowSensitiveTypeInference(trackedVars, memberAccessesByQualifiedExpr, assignmentsByAssignmentStatement);
222+
new FlowSensitiveTypeInference(trackedVars, memberAccessesByQualifiedExpr, assignmentsByAssignmentStatement, parameterTypesByName);
204223

205224
flowSensitiveTypeInference.compute(cfg);
206225
flowSensitiveTypeInference.compute(cfg);

‎python-frontend/src/test/java/org/sonar/python/types/TypeInferenceTest.java‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,8 @@
3030
import org.sonar.plugins.python.api.tree.FileInput;
3131
import org.sonar.plugins.python.api.tree.RegularArgument;
3232
import org.sonar.plugins.python.api.tree.Tree;
33+
import org.sonar.plugins.python.api.types.BuiltinTypes;
34+
import org.sonar.plugins.python.api.types.InferredType;
3335
import org.sonar.python.PythonTestUtils;
3436

3537
import static org.assertj.core.api.Assertions.assertThat;
@@ -50,6 +52,7 @@
5052
import static org.sonar.python.types.InferredTypes.anyType;
5153
import static org.sonar.python.types.InferredTypes.or;
5254
import static org.sonar.python.types.InferredTypes.runtimeType;
55+
import static org.sonar.python.types.InferredTypes.typeName;
5356

5457
//import static org.sonar.python.types.InferredTypes.BYTES;
5558

@@ -104,6 +107,26 @@ public void parameter() {
104107
" a").type()).isEqualTo(anyType());
105108
}
106109

110+
@Test
111+
public void parameter_with_annotation() {
112+
assertDeclaredType(lastExpression("def f(p: int): p").type(), BuiltinTypes.INT);
113+
assertDeclaredType(lastExpression("def f(p: str): p").type(), BuiltinTypes.STR);
114+
assertDeclaredType(lastExpression("class A: ...\ndef f(p: A): p").type(), "A");
115+
assertThat(lastExpression("def f(p: unknown): p").type()).isEqualTo(InferredTypes.anyType());
116+
assertDeclaredType(lastExpression("def f(p1: int, *, p2: str): p2").type(), BuiltinTypes.STR);
117+
118+
assertThat(lastExpression(
119+
"def f(p: int):",
120+
" p = 'str'",
121+
" p").type()).isEqualTo(STR);
122+
123+
assertThat(lastExpression(
124+
"def f(p: int):",
125+
" try: ...",
126+
" except: ...",
127+
" p").type()).isEqualTo(anyType());
128+
}
129+
107130
@Test
108131
public void local_variable() {
109132
assertThat(lastExpressionInFunction(
@@ -500,4 +523,9 @@ public void execution_order_assignment_statement() {
500523
Expression xLhs = assignment.lhsExpressions().get(0).expressions().get(0);
501524
assertThat(xLhs.type()).isEqualTo(STR);
502525
}
526+
527+
private static void assertDeclaredType(InferredType type, String typeName) {
528+
assertThat(type).isInstanceOf(DeclaredType.class);
529+
assertThat(typeName(type)).isEqualTo(typeName);
530+
}
503531
}

0 commit comments

Comments
 (0)