[red-knot] Type inference for str.startswith

Author: sharkdp

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

@@ -2000,6 +2000,12 @@ impl<'db> Type<'db> {
                         ))
                         .into(),
                     ),
+                    (Some(KnownClass::Str), "startswith") => Some(
+                        Symbol::bound(Type::WrapperDescriptor(
+                            WrapperDescriptorKind::StrStartswith,
+                        ))
+                        .into(),
+                    ),
                     // TODO:
                     // We currently hard-code the knowledge that the following known classes are not
                     // descriptors, i.e. that they have no `__get__` method. This is not wrong and
@@ -2508,6 +2514,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) =>
@@ -2533,6 +2543,14 @@ impl<'db> Type<'db> {
                 ))
                 .into()
             }
+            Type::ClassLiteral(class)
+                if name == "startswith" && class.is_known(db, KnownClass::Str) =>
+            {
+                Symbol::bound(Type::WrapperDescriptor(
+                    WrapperDescriptorKind::StrStartswith,
+                ))
+                .into()
+            }
             Type::BoundMethod(bound_method) => match name_str {
                 "__self__" => Symbol::bound(bound_method.self_instance(db)).into(),
                 "__func__" => {
@@ -3043,7 +3061,7 @@ impl<'db> Type<'db> {
                     WrapperDescriptorKind::PropertyDunderGet => {
                         KnownClass::Property.to_instance(db)
                     }
-                    WrapperDescriptorKind::PropertyDunderSet => {
+                    _ => {
                         unreachable!("Not part of outer match pattern")
                     }
                 };
@@ -3112,6 +3130,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 +4284,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 +6202,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`
@@ -6166,6 +6215,8 @@ pub enum WrapperDescriptorKind {
     PropertyDunderGet,
     /// `property.__set__`
     PropertyDunderSet,
+    /// `str.startswith`
+    StrStartswith,
 }
 
 #[salsa::interned(debug)]

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() =>
                 {
@@ -1088,6 +1098,7 @@ impl<'db> CallableDescription<'db> {
                     WrapperDescriptorKind::FunctionTypeDunderGet => "FunctionType.__get__",
                     WrapperDescriptorKind::PropertyDunderGet => "property.__get__",
                     WrapperDescriptorKind::PropertyDunderSet => "property.__set__",
+                    WrapperDescriptorKind::StrStartswith => "str.startswith",
                 },
             }),
             _ => None,

crates/red_knot_python_semantic/src/types/display.rs

@@ -145,11 +145,15 @@ 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"),
                     WrapperDescriptorKind::PropertyDunderGet => ("__get__", "property"),
                     WrapperDescriptorKind::PropertyDunderSet => ("__set__", "property"),
+                    WrapperDescriptorKind::StrStartswith => ("startswith", "str"),
                 };
                 write!(f, "<wrapper-descriptor `{method}` of `{object}` objects>")
             }