[red-knot] Type inference for str.startswith

Author: sharkdp

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

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]
 ```

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`

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() =>
                 {

crates/red_knot_python_semantic/src/types/display.rs

@@ -145,6 +145,9 @@ impl Display for DisplayRepresentation<'_> {
             Type::MethodWrapper(MethodWrapperKind::PropertyDunderSet(_)) => {
                 write!(f, "<method-wrapper `__set__` of `property` object>",)
             }
+            Type::MethodWrapper(MethodWrapperKind::StrStartswith(_)) => {
+                write!(f, "<method-wrapper `startswith` of `str` object>",)
+            }
             Type::WrapperDescriptor(kind) => {
                 let (method, object) = match kind {
                     WrapperDescriptorKind::FunctionTypeDunderGet => ("__get__", "function"),