[red-knot] Special case @final, @staticmethod, @override

Author: dhruvmanila

crates/red_knot_python_semantic/src/types.rs

@@ -5882,6 +5882,12 @@ bitflags! {
         const OVERLOAD = 1 << 2;
         /// `@abc.abstractmethod`
         const ABSTRACT_METHOD = 1 << 3;
+        /// `@typing.final`
+        const FINAL = 1 << 4;
+        /// `@staticmethod`
+        const STATICMETHOD = 1 << 5;
+        /// `@typing.override`
+        const OVERRIDE = 1 << 6;
     }
 }
 
@@ -6265,6 +6271,8 @@ pub enum KnownFunction {
     Cast,
     /// `typing(_extensions).overload`
     Overload,
+    /// `typing(_extensions).override`
+    Override,
     /// `typing(_extensions).is_protocol`
     IsProtocol,
     /// `typing(_extensions).get_protocol_members`
@@ -6332,6 +6340,7 @@ impl KnownFunction {
             | Self::AssertNever
             | Self::Cast
             | Self::Overload
+            | Self::Override
             | Self::RevealType
             | Self::Final
             | Self::IsProtocol
@@ -7709,6 +7718,7 @@ pub(crate) mod tests {
                 KnownFunction::Cast
                 | KnownFunction::Final
                 | KnownFunction::Overload
+                | KnownFunction::Override
                 | KnownFunction::RevealType
                 | KnownFunction::AssertType
                 | KnownFunction::AssertNever

crates/red_knot_python_semantic/src/types/call/bind.rs

@@ -585,6 +585,14 @@ impl<'db> Bindings<'db> {
                         }
                     }
 
+                    Some(KnownFunction::Override) => {
+                        // TODO: This can be removed once we understand legacy generics because the
+                        // typeshed definition for `typing.overload` is an identity function.
+                        if let [Some(ty)] = overload.parameter_types() {
+                            overload.set_return_type(*ty);
+                        }
+                    }
+
                     Some(KnownFunction::AbstractMethod) => {
                         // TODO: This can be removed once we understand legacy generics because the
                         // typeshed definition for `abc.abstractmethod` is an identity function.
@@ -593,6 +601,14 @@ impl<'db> Bindings<'db> {
                         }
                     }
 
+                    Some(KnownFunction::Final) => {
+                        // TODO: This can be removed once we understand legacy generics because the
+                        // typeshed definition for `abc.abstractmethod` is an identity function.
+                        if let [Some(ty)] = overload.parameter_types() {
+                            overload.set_return_type(*ty);
+                        }
+                    }
+
                     Some(KnownFunction::GetattrStatic) => {
                         let [Some(instance_ty), Some(attr_name), default] =
                             overload.parameter_types()

crates/red_knot_python_semantic/src/types/class.rs

@@ -1933,6 +1933,7 @@ pub(crate) enum KnownClass {
     BaseException,
     BaseExceptionGroup,
     Classmethod,
+    Staticmethod,
     Super,
     // enum
     Enum,
@@ -2052,6 +2053,7 @@ impl<'db> KnownClass {
             // and raises a `TypeError` in Python >=3.14
             // (see https://docs.python.org/3/library/constants.html#NotImplemented)
             | Self::NotImplementedType
+            | Self::Staticmethod
             | Self::Classmethod => Truthiness::Ambiguous,
         }
     }
@@ -2079,6 +2081,7 @@ impl<'db> KnownClass {
             Self::BaseException => "BaseException",
             Self::BaseExceptionGroup => "BaseExceptionGroup",
             Self::Classmethod => "classmethod",
+            Self::Staticmethod => "staticmethod",
             Self::GenericAlias => "GenericAlias",
             Self::ModuleType => "ModuleType",
             Self::FunctionType => "FunctionType",
@@ -2256,6 +2259,7 @@ impl<'db> KnownClass {
             | Self::BaseException
             | Self::BaseExceptionGroup
             | Self::Classmethod
+            | Self::Staticmethod
             | Self::Slice
             | Self::Range
             | Self::Super
@@ -2345,6 +2349,7 @@ impl<'db> KnownClass {
             | Self::BaseException
             | Self::BaseExceptionGroup
             | Self::Classmethod
+            | Self::Staticmethod
             | Self::GenericAlias
             | Self::ModuleType
             | Self::FunctionType
@@ -2419,6 +2424,7 @@ impl<'db> KnownClass {
             | Self::BaseException
             | Self::BaseExceptionGroup
             | Self::Classmethod
+            | Self::Staticmethod
             | Self::TypeVar
             | Self::ParamSpec
             | Self::ParamSpecArgs
@@ -2462,6 +2468,7 @@ impl<'db> KnownClass {
             "BaseException" => Self::BaseException,
             "BaseExceptionGroup" => Self::BaseExceptionGroup,
             "classmethod" => Self::Classmethod,
+            "staticmethod" => Self::Staticmethod,
             "GenericAlias" => Self::GenericAlias,
             "NoneType" => Self::NoneType,
             "ModuleType" => Self::ModuleType,
@@ -2540,6 +2547,7 @@ impl<'db> KnownClass {
             | Self::EllipsisType
             | Self::BaseExceptionGroup
             | Self::Classmethod
+            | Self::Staticmethod
             | Self::FunctionType
             | Self::MethodType
             | Self::MethodWrapperType

crates/red_knot_python_semantic/src/types/infer.rs

@@ -1500,15 +1500,28 @@ impl<'db> TypeInferenceBuilder<'db> {
                             function_decorators |= FunctionDecorators::ABSTRACT_METHOD;
                             continue;
                         }
+                        Some(KnownFunction::Final) => {
+                            function_decorators |= FunctionDecorators::FINAL;
+                            continue;
+                        }
+                        Some(KnownFunction::Override) => {
+                            function_decorators |= FunctionDecorators::OVERRIDE;
+                            continue;
+                        }
                         _ => {}
                     }
                 }
-                Type::ClassLiteral(class) => {
-                    if class.is_known(self.db(), KnownClass::Classmethod) {
+                Type::ClassLiteral(class) => match class.known(self.db()) {
+                    Some(KnownClass::Classmethod) => {
                         function_decorators |= FunctionDecorators::CLASSMETHOD;
                         continue;
                     }
-                }
+                    Some(KnownClass::Staticmethod) => {
+                        function_decorators |= FunctionDecorators::STATICMETHOD;
+                        continue;
+                    }
+                    _ => {}
+                },
                 Type::DataclassTransformer(params) => {
                     dataclass_transformer_params = Some(params);
                 }