Skip to content

Commit 7ead3ec

Browse files
marc-jasper-sonarsourcesonartech
authored andcommitted
SONARPY-3235 Resolve parameters type hinted with typing.Self (#688)
GitOrigin-RevId: 9f30af59aab6424963762ea5ec0043b5f229214a
1 parent 450ec50 commit 7ead3ec

File tree

2 files changed

+260
-12
lines changed

2 files changed

+260
-12
lines changed

python-frontend/src/main/java/org/sonar/python/semantic/v2/types/TrivialTypeInferenceVisitor.java

Lines changed: 53 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@
2929
import java.util.Set;
3030
import java.util.function.Predicate;
3131
import java.util.stream.Stream;
32+
import javax.annotation.CheckForNull;
3233
import javax.annotation.Nullable;
3334
import org.sonar.plugins.python.api.PythonFile;
3435
import org.sonar.plugins.python.api.TriBool;
@@ -49,6 +50,7 @@
4950
import org.sonar.plugins.python.api.tree.ExpressionList;
5051
import org.sonar.plugins.python.api.tree.FileInput;
5152
import org.sonar.plugins.python.api.tree.FunctionDef;
53+
import org.sonar.plugins.python.api.tree.FunctionLike;
5254
import org.sonar.plugins.python.api.tree.ImportFrom;
5355
import org.sonar.plugins.python.api.tree.ImportName;
5456
import org.sonar.plugins.python.api.tree.ListLiteral;
@@ -113,6 +115,11 @@ public class TrivialTypeInferenceVisitor extends BaseTreeVisitor {
113115
TypeInferenceMatchers.isObjectOfType("bytes"),
114116
TypeInferenceMatchers.isObjectOfType("complex")));
115117

118+
private static final TypeInferenceMatcher IS_TYPING_SELF = TypeInferenceMatcher.of(
119+
TypeInferenceMatchers.any(
120+
TypeInferenceMatchers.isType("typing.Self"),
121+
TypeInferenceMatchers.isType("typing_extensions.Self")));
122+
116123
private final TypeTable projectLevelTypeTable;
117124
private final String fileId;
118125
private final String fullyQualifiedModuleName;
@@ -347,17 +354,15 @@ private FunctionType buildFunctionType(FunctionDef functionDef) {
347354
FunctionTypeBuilder functionTypeBuilder = new FunctionTypeBuilder()
348355
.fromFunctionDef(functionDef, fullyQualifiedName, fileId, projectLevelTypeTable)
349356
.withDefinitionLocation(locationInFile(functionDef.name(), fileId));
350-
ClassType owner = null;
351-
if (currentType() instanceof ClassType classType) {
352-
owner = classType;
353-
}
357+
ClassType owner = computeDirectClassOwner();
354358
if (owner != null) {
355359
functionTypeBuilder.withOwner(owner);
356360
}
357361
TypeAnnotation typeAnnotation = functionDef.returnTypeAnnotation();
358362
if (typeAnnotation != null) {
359-
PythonType returnType = typeAnnotation.expression().typeV2();
360-
functionTypeBuilder.withReturnType(returnType instanceof UnknownType ? returnType : ObjectType.Builder.fromType(returnType).withTypeSource(TypeSource.TYPE_HINT).build());
363+
PythonType returnType = resolveTypeAnnotationExpressionType(typeAnnotation.expression(), owner);
364+
functionTypeBuilder.withReturnType(returnType instanceof UnknownType ? returnType :
365+
ObjectType.Builder.fromType(returnType).withTypeSource(TypeSource.TYPE_HINT).build());
361366
functionTypeBuilder.withTypeOrigin(TypeOrigin.LOCAL);
362367
}
363368
FunctionType functionType = functionTypeBuilder.build();
@@ -618,16 +623,37 @@ private void addStaticFieldToClass(Name name, PythonType type) {
618623
public void visitParameter(Parameter parameter) {
619624
scan(parameter.typeAnnotation());
620625
scan(parameter.defaultValue());
626+
ClassType owner = computeOwnerForParameter(parameter);
621627
Optional.ofNullable(parameter.typeAnnotation())
622628
.map(TypeAnnotation::expression)
623-
.map(TrivialTypeInferenceVisitor::resolveTypeAnnotationExpressionType)
629+
.map(expr -> resolveTypeAnnotationExpressionType(expr, owner))
624630
.ifPresent(type -> setTypeToName(parameter.name(), type));
625631
scan(parameter.name());
626632
}
627633

628-
private static PythonType resolveTypeAnnotationExpressionType(Expression expression) {
634+
@CheckForNull
635+
private ClassType computeOwnerForParameter(Parameter parameter) {
636+
FunctionLike enclosingFunction = TreeUtils.firstAncestorOfClass(parameter, FunctionLike.class);
637+
if (enclosingFunction instanceof FunctionDef functionDef) {
638+
// During the second scan of the function, the FunctionType has already been built
639+
// and we can retrieve the owner directly from it.
640+
PythonType funcType = functionDef.name().typeV2();
641+
if (funcType instanceof FunctionType ft) {
642+
return ft.owner() instanceof ClassType ct ? ct : null;
643+
}
644+
}
645+
return computeDirectClassOwner();
646+
}
647+
648+
@CheckForNull
649+
private ClassType computeDirectClassOwner() {
650+
return currentType() instanceof ClassType classType ? classType : null;
651+
}
652+
653+
private PythonType resolveTypeAnnotationExpressionType(Expression expression, @Nullable ClassType enclosingClassType) {
629654
if (expression instanceof Name name && name.typeV2() != PythonType.UNKNOWN) {
630-
return ObjectType.Builder.fromType(name.typeV2()).withTypeSource(TypeSource.TYPE_HINT).build();
655+
PythonType resolvedType = resolveSelfType(name.typeV2(), enclosingClassType);
656+
return ObjectType.Builder.fromType(resolvedType).withTypeSource(TypeSource.TYPE_HINT).build();
631657
} else if (expression instanceof SubscriptionExpression subscriptionExpression && subscriptionExpression.object().typeV2() != PythonType.UNKNOWN) {
632658
var candidateTypes = subscriptionExpression.subscripts()
633659
.expressions()
@@ -645,17 +671,32 @@ private static PythonType resolveTypeAnnotationExpressionType(Expression express
645671
.withTypeSource(TypeSource.TYPE_HINT)
646672
.build();
647673
} else if (expression instanceof BinaryExpression binaryExpression) {
648-
var left = resolveTypeAnnotationExpressionType(binaryExpression.leftOperand());
649-
var right = resolveTypeAnnotationExpressionType(binaryExpression.rightOperand());
674+
var left = resolveTypeAnnotationExpressionType(binaryExpression.leftOperand(), enclosingClassType);
675+
var right = resolveTypeAnnotationExpressionType(binaryExpression.rightOperand(), enclosingClassType);
650676
// TODO: we need to make a decision on should here be a union type of object types or an object type of a union type.
651677
// ATM it is blocked by the generic types resolution redesign
652678
return UnionType.or(left, right);
653679
} else if (expression.typeV2() instanceof ClassType classType) {
654680
return ObjectType.Builder.fromType(classType).withTypeSource(TypeSource.TYPE_HINT).build();
681+
} else if (expression instanceof NoneExpression noneExpression) {
682+
return noneExpression.typeV2();
655683
}
656684
return PythonType.UNKNOWN;
657685
}
658686

687+
private PythonType resolveSelfType(PythonType type, @Nullable ClassType enclosingClassType) {
688+
if (isTypingSelf(type)) {
689+
return Optional.ofNullable(enclosingClassType)
690+
.map(SelfType::of)
691+
.orElse(type);
692+
}
693+
return type;
694+
}
695+
696+
private boolean isTypingSelf(PythonType type) {
697+
return IS_TYPING_SELF.evaluate(type, typePredicateContext).isTrue();
698+
}
699+
659700
@Override
660701
public void visitQualifiedExpression(QualifiedExpression qualifiedExpression) {
661702
scan(qualifiedExpression.qualifier());
@@ -767,4 +808,4 @@ private static void setTypeToName(@Nullable Tree tree, @Nullable PythonType type
767808
name.typeV2(type);
768809
}
769810
}
770-
}
811+
}

python-frontend/src/test/java/org/sonar/python/semantic/v2/TypeInferenceV2Test.java

Lines changed: 207 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4105,6 +4105,213 @@ def foo((a, b)):
41054105
}
41064106

41074107
@Test
4108+
void inferParameterTypeHintedWithTypingSelf() {
4109+
FileInput root = inferTypes("""
4110+
from typing import Self
4111+
class A:
4112+
def foo(self, other: Self) -> Self:
4113+
...
4114+
""");
4115+
4116+
var classDef = (ClassDef) root.statements().statements().get(1);
4117+
var classType = (ClassType) classDef.name().typeV2();
4118+
var functionDef = (FunctionDef) classDef.body().statements().get(0);
4119+
var functionType = (FunctionType) functionDef.name().typeV2();
4120+
4121+
var otherParamType = functionType.parameters().get(1).declaredType().type();
4122+
assertThat(otherParamType).isInstanceOf(ObjectType.class);
4123+
var otherParamUnwrapped = otherParamType.unwrappedType();
4124+
assertThat(otherParamUnwrapped).isInstanceOf(SelfType.class);
4125+
assertThat(((SelfType) otherParamUnwrapped).innerType()).isEqualTo(classType);
4126+
4127+
var returnType = functionType.returnType();
4128+
assertThat(returnType).isInstanceOf(ObjectType.class);
4129+
var returnTypeUnwrapped = returnType.unwrappedType();
4130+
assertThat(returnTypeUnwrapped).isInstanceOf(SelfType.class);
4131+
assertThat(((SelfType) returnTypeUnwrapped).innerType()).isEqualTo(classType);
4132+
}
4133+
4134+
@Test
4135+
void parameterTypeHintedWithNonSelfUnresolvedImport() {
4136+
FileInput root = inferTypes("""
4137+
from unknown_module import SomeType
4138+
class A:
4139+
def foo(self, x: SomeType) -> None:
4140+
...
4141+
""");
4142+
4143+
var classDef = (ClassDef) root.statements().statements().get(1);
4144+
var functionDef = (FunctionDef) classDef.body().statements().get(0);
4145+
var functionType = (FunctionType) functionDef.name().typeV2();
4146+
4147+
var xParamType = functionType.parameters().get(1).declaredType().type();
4148+
assertThat(xParamType).isInstanceOf(ObjectType.class);
4149+
assertThat(xParamType.unwrappedType()).isInstanceOf(UnknownType.UnresolvedImportType.class);
4150+
assertThat(xParamType.unwrappedType()).isNotInstanceOf(SelfType.class);
4151+
}
4152+
4153+
@Test
4154+
void parameterTypeHintedWithNonSelfSpecialForm() {
4155+
FileInput root = inferTypes("""
4156+
from typing import Final
4157+
class A:
4158+
def foo(self, x: Final) -> None:
4159+
...
4160+
""");
4161+
4162+
var classDef = (ClassDef) root.statements().statements().get(1);
4163+
var functionDef = (FunctionDef) classDef.body().statements().get(0);
4164+
var functionType = (FunctionType) functionDef.name().typeV2();
4165+
4166+
var xParamType = functionType.parameters().get(1).declaredType().type();
4167+
assertThat(xParamType).isInstanceOf(ObjectType.class);
4168+
assertThat(xParamType.unwrappedType()).isNotInstanceOf(SelfType.class);
4169+
}
4170+
4171+
@Test
4172+
void selfReturnTypeWithoutEnclosingClass() {
4173+
FileInput root = inferTypes("""
4174+
from typing import Self
4175+
def foo() -> Self:
4176+
...
4177+
""");
4178+
4179+
var functionDef = (FunctionDef) root.statements().statements().get(1);
4180+
var functionType = (FunctionType) functionDef.name().typeV2();
4181+
4182+
var returnType = functionType.returnType();
4183+
assertThat(returnType).isNotInstanceOf(SelfType.class);
4184+
assertThat(returnType.unwrappedType()).isNotInstanceOf(SelfType.class);
4185+
}
4186+
4187+
@Test
4188+
void selfReturnTypeInNestedFunctionWithoutEnclosingClass() {
4189+
FileInput root = inferTypes("""
4190+
from typing import Self
4191+
def foo():
4192+
def bar() -> Self:
4193+
pass
4194+
""");
4195+
4196+
var outerFunctionDef = (FunctionDef) root.statements().statements().get(1);
4197+
var innerFunctionDef = (FunctionDef) outerFunctionDef.body().statements().get(0);
4198+
var innerFunctionType = (FunctionType) innerFunctionDef.name().typeV2();
4199+
4200+
var returnType = innerFunctionType.returnType();
4201+
assertThat(returnType).isNotInstanceOf(SelfType.class);
4202+
assertThat(returnType.unwrappedType()).isNotInstanceOf(SelfType.class);
4203+
}
4204+
4205+
@Test
4206+
void selfReturnTypeInNestedFunctionInsideMethod() {
4207+
FileInput root = inferTypes("""
4208+
from typing import Self
4209+
class A:
4210+
def foo(self):
4211+
def bar() -> Self:
4212+
pass
4213+
""");
4214+
4215+
var classDef = (ClassDef) root.statements().statements().get(1);
4216+
var methodDef = (FunctionDef) classDef.body().statements().get(0);
4217+
var innerFunctionDef = (FunctionDef) methodDef.body().statements().get(0);
4218+
var innerFunctionType = (FunctionType) innerFunctionDef.name().typeV2();
4219+
4220+
// bar's return type should NOT be resolved to SelfType(A) because bar is not directly owned by class A
4221+
var returnType = innerFunctionType.returnType();
4222+
assertThat(returnType).isNotInstanceOf(SelfType.class);
4223+
assertThat(returnType.unwrappedType()).isNotInstanceOf(SelfType.class);
4224+
}
4225+
4226+
@Test
4227+
void stringLiteralTypeInference() {
4228+
Expression root = lastExpression("""
4229+
def bar(param: "MyParameterType") -> "MyReturnType":
4230+
return param
4231+
bar
4232+
""");
4233+
4234+
assertThat(root.typeV2()).isInstanceOf(FunctionType.class);
4235+
FunctionType functionType = (FunctionType) root.typeV2();
4236+
assertThat(functionType.parameters().get(0).declaredType().type().unwrappedType()).isEqualTo(PythonType.UNKNOWN);
4237+
assertThat(functionType.returnType().unwrappedType()).isEqualTo(PythonType.UNKNOWN);
4238+
}
4239+
4240+
@Test
4241+
void selfInUnionReturnType() {
4242+
FileInput root = inferTypes("""
4243+
from typing import Self
4244+
class A:
4245+
def foo(self) -> Self | None:
4246+
...
4247+
""");
4248+
4249+
var classDef = (ClassDef) root.statements().statements().get(1);
4250+
var classType = (ClassType) classDef.name().typeV2();
4251+
var functionDef = (FunctionDef) classDef.body().statements().get(0);
4252+
var functionType = (FunctionType) functionDef.name().typeV2();
4253+
4254+
var returnType = functionType.returnType();
4255+
assertThat(returnType).isInstanceOf(ObjectType.class);
4256+
var unwrappedReturnType = returnType.unwrappedType();
4257+
assertThat(unwrappedReturnType).isInstanceOf(UnionType.class);
4258+
4259+
var unionType = (UnionType) unwrappedReturnType;
4260+
var candidates = unionType.candidates();
4261+
assertThat(candidates).hasSize(2);
4262+
4263+
var selfCandidate = candidates.stream()
4264+
.filter(c -> c.unwrappedType() instanceof SelfType)
4265+
.findFirst();
4266+
assertThat(selfCandidate).isPresent();
4267+
var selfType = (SelfType) selfCandidate.get().unwrappedType();
4268+
assertThat(selfType.innerType()).isEqualTo(classType);
4269+
4270+
var noneCandidate = candidates.stream()
4271+
.filter(c -> !(c.unwrappedType() instanceof SelfType))
4272+
.findFirst();
4273+
assertThat(noneCandidate).isPresent();
4274+
}
4275+
4276+
@Test
4277+
void selfInUnionReturnTypeWithoutEnclosingClass() {
4278+
FileInput root = inferTypes("""
4279+
from typing import Self
4280+
def foo() -> Self | None:
4281+
...
4282+
""");
4283+
4284+
var functionDef = (FunctionDef) root.statements().statements().get(1);
4285+
var functionType = (FunctionType) functionDef.name().typeV2();
4286+
4287+
var returnType = functionType.returnType();
4288+
assertThat(returnType).isInstanceOf(ObjectType.class);
4289+
var unwrappedReturnType = returnType.unwrappedType();
4290+
assertThat(unwrappedReturnType).isInstanceOf(UnionType.class);
4291+
4292+
var unionType = (UnionType) unwrappedReturnType;
4293+
assertThat(unionType.candidates()).noneMatch(c -> c.unwrappedType() instanceof SelfType);
4294+
}
4295+
4296+
@Test
4297+
void selfOrNoneTypeAlias() {
4298+
FileInput root = inferTypes("""
4299+
from typing import Self
4300+
SelfOrNone = Self | None
4301+
class A:
4302+
def foo(self) -> SelfOrNone:
4303+
...
4304+
""");
4305+
4306+
var classDef = (ClassDef) root.statements().statements().get(2);
4307+
var functionDef = (FunctionDef) classDef.body().statements().get(0);
4308+
var functionType = (FunctionType) functionDef.name().typeV2();
4309+
4310+
var returnType = functionType.returnType();
4311+
assertThat(returnType).isSameAs(PythonType.UNKNOWN); // SONARPY-3559 should change this to correctly resolve to a union type of Self and None
4312+
}
4313+
4314+
@Test
41084315
void anyResolvesToUnknown() {
41094316
Expression anyExpression = lastExpression("""
41104317
from typing import Any

0 commit comments

Comments
 (0)