diff --git a/crates/red_knot_python_semantic/resources/mdtest/call/str_startswith.md b/crates/red_knot_python_semantic/resources/mdtest/call/str_startswith.md new file mode 100644 index 00000000000000..643f1747635044 --- /dev/null +++ b/crates/red_knot_python_semantic/resources/mdtest/call/str_startswith.md @@ -0,0 +1,50 @@ +# `str.startswith` + +We special-case `str.startswith` to allow inference of precise Boolean literal types, because those +are used in [`sys.platform` checks]. + +```py +reveal_type("abc".startswith("")) # revealed: Literal[True] +reveal_type("abc".startswith("a")) # revealed: Literal[True] +reveal_type("abc".startswith("ab")) # revealed: Literal[True] +reveal_type("abc".startswith("abc")) # revealed: Literal[True] + +reveal_type("abc".startswith("abcd")) # revealed: Literal[False] +reveal_type("abc".startswith("bc")) # revealed: Literal[False] + +reveal_type("AbC".startswith("")) # revealed: Literal[True] +reveal_type("AbC".startswith("A")) # revealed: Literal[True] +reveal_type("AbC".startswith("Ab")) # revealed: Literal[True] +reveal_type("AbC".startswith("AbC")) # revealed: Literal[True] + +reveal_type("AbC".startswith("a")) # revealed: Literal[False] +reveal_type("AbC".startswith("aB")) # revealed: Literal[False] + +reveal_type("".startswith("")) # revealed: Literal[True] + +reveal_type("".startswith(" ")) # revealed: Literal[False] +``` + +Make sure that we fall back to `bool` for more complex cases: + +```py +reveal_type("abc".startswith("b", 1)) # revealed: bool +reveal_type("abc".startswith("bc", 1, 3)) # revealed: bool + +reveal_type("abc".startswith(("a", "x"))) # revealed: bool +``` + +And similiarly, we should still infer `bool` if the instance or the prefix are not string literals: + +```py +from typing_extensions import LiteralString + +def _(string_instance: str, literalstring: LiteralString): + reveal_type(string_instance.startswith("a")) # revealed: bool + reveal_type(literalstring.startswith("a")) # revealed: bool + + reveal_type("a".startswith(string_instance)) # revealed: bool + reveal_type("a".startswith(literalstring)) # revealed: bool +``` + +[`sys.platform` checks]: https://docs.python.org/3/library/sys.html#sys.platform diff --git a/crates/red_knot_python_semantic/resources/mdtest/sys_platform.md b/crates/red_knot_python_semantic/resources/mdtest/sys_platform.md index 1cf9410456f76f..3997cbb28bd794 100644 --- a/crates/red_knot_python_semantic/resources/mdtest/sys_platform.md +++ b/crates/red_knot_python_semantic/resources/mdtest/sys_platform.md @@ -31,13 +31,13 @@ reveal_type(sys.platform) # revealed: Literal["linux"] ## Testing for a specific platform -### Exact comparison - ```toml [environment] python-platform = "freebsd8" ``` +### Exact comparison + ```py import sys @@ -48,11 +48,11 @@ reveal_type(sys.platform == "linux") # revealed: Literal[False] ### Substring comparison It is [recommended](https://docs.python.org/3/library/sys.html#sys.platform) to use -`sys.platform.startswith(...)` for platform checks. This is not yet supported in type inference: +`sys.platform.startswith(...)` for platform checks: ```py import sys -reveal_type(sys.platform.startswith("freebsd")) # revealed: bool -reveal_type(sys.platform.startswith("linux")) # revealed: bool +reveal_type(sys.platform.startswith("freebsd")) # revealed: Literal[True] +reveal_type(sys.platform.startswith("linux")) # revealed: Literal[False] ``` diff --git a/crates/red_knot_python_semantic/src/types.rs b/crates/red_knot_python_semantic/src/types.rs index 4d3b05f28f9899..c3da9b99d27e65 100644 --- a/crates/red_knot_python_semantic/src/types.rs +++ b/crates/red_knot_python_semantic/src/types.rs @@ -2508,6 +2508,10 @@ impl<'db> Type<'db> { Type::MethodWrapper(MethodWrapperKind::PropertyDunderSet(property)), ) .into(), + Type::StringLiteral(literal) if name == "startswith" => Symbol::bound( + Type::MethodWrapper(MethodWrapperKind::StrStartswith(literal)), + ) + .into(), Type::ClassLiteral(class) if name == "__get__" && class.is_known(db, KnownClass::FunctionType) => @@ -3112,6 +3116,34 @@ impl<'db> Type<'db> { )) } + Type::MethodWrapper(MethodWrapperKind::StrStartswith(_)) => { + Signatures::single(CallableSignature::single( + self, + Signature::new( + Parameters::new([ + Parameter::positional_only(Some(Name::new_static("prefix"))) + .with_annotated_type(UnionType::from_elements( + db, + [ + KnownClass::Str.to_instance(db), + // TODO: tuple[str, ...] + KnownClass::Tuple.to_instance(db), + ], + )), + Parameter::positional_only(Some(Name::new_static("start"))) + // TODO: SupportsIndex | None + .with_annotated_type(Type::object(db)) + .with_default_type(Type::none(db)), + Parameter::positional_only(Some(Name::new_static("end"))) + // TODO: SupportsIndex | None + .with_annotated_type(Type::object(db)) + .with_default_type(Type::none(db)), + ]), + Some(KnownClass::Bool.to_instance(db)), + ), + )) + } + Type::FunctionLiteral(function_type) => match function_type.known(db) { Some( KnownFunction::IsEquivalentTo @@ -4238,6 +4270,7 @@ impl<'db> Type<'db> { | Type::AlwaysTruthy | Type::AlwaysFalsy | Type::WrapperDescriptor(_) + | Type::MethodWrapper(MethodWrapperKind::StrStartswith(_)) | Type::ModuleLiteral(_) // A non-generic class never needs to be specialized. A generic class is specialized // explicitly (via a subscript expression) or implicitly (via a call), and not because @@ -6155,6 +6188,8 @@ pub enum MethodWrapperKind<'db> { PropertyDunderGet(PropertyInstanceType<'db>), /// Method wrapper for `some_property.__set__` PropertyDunderSet(PropertyInstanceType<'db>), + /// Method wrapper for `str.startswith` + StrStartswith(StringLiteralType<'db>), } /// Represents a specific instance of `types.WrapperDescriptorType` diff --git a/crates/red_knot_python_semantic/src/types/call/bind.rs b/crates/red_knot_python_semantic/src/types/call/bind.rs index d4c88d013570c2..8beb6667be4f76 100644 --- a/crates/red_knot_python_semantic/src/types/call/bind.rs +++ b/crates/red_knot_python_semantic/src/types/call/bind.rs @@ -395,6 +395,16 @@ impl<'db> Bindings<'db> { } } + Type::MethodWrapper(MethodWrapperKind::StrStartswith(literal)) => { + if let [Some(Type::StringLiteral(prefix)), None, None] = + overload.parameter_types() + { + overload.set_return_type(Type::BooleanLiteral( + literal.value(db).starts_with(&**prefix.value(db)), + )); + } + } + Type::BoundMethod(bound_method) if bound_method.self_instance(db).is_property_instance() => { diff --git a/crates/red_knot_python_semantic/src/types/display.rs b/crates/red_knot_python_semantic/src/types/display.rs index 87825bc5a684ca..84747c01571674 100644 --- a/crates/red_knot_python_semantic/src/types/display.rs +++ b/crates/red_knot_python_semantic/src/types/display.rs @@ -145,6 +145,9 @@ impl Display for DisplayRepresentation<'_> { Type::MethodWrapper(MethodWrapperKind::PropertyDunderSet(_)) => { write!(f, "",) } + Type::MethodWrapper(MethodWrapperKind::StrStartswith(_)) => { + write!(f, "",) + } Type::WrapperDescriptor(kind) => { let (method, object) = match kind { WrapperDescriptorKind::FunctionTypeDunderGet => ("__get__", "function"),