[ty] Infer type of self for decorated methods and properties

Author: sharkdp

crates/ruff_benchmark/benches/ty_walltime.rs

@@ -226,7 +226,7 @@ static STATIC_FRAME: Benchmark = Benchmark::new(
         max_dep_date: "2025-08-09",
         python_version: PythonVersion::PY311,
     },
-    750,
+    800,
 );
 
 #[track_caller]

crates/ty_ide/src/semantic_tokens.rs

@@ -1413,7 +1413,7 @@ u = List.__name__        # __name__ should be variable<CURSOR>
         "property" @ 168..176: Decorator
         "prop" @ 185..189: Method [definition]
         "self" @ 190..194: SelfParameter
-        "self" @ 212..216: Variable
+        "self" @ 212..216: TypeParameter
         "CONSTANT" @ 217..225: Variable [readonly]
         "obj" @ 227..230: Variable
         "MyClass" @ 233..240: Class

crates/ty_python_semantic/resources/mdtest/annotations/self.md

@@ -116,7 +116,7 @@ A.implicit_self(1)
 Passing `self` implicitly also verifies the type:
 
 ```py
-from typing import Never
+from typing import Never, Callable
 
 class Strange:
     def can_not_be_called(self: Never) -> None: ...
@@ -139,6 +139,9 @@ The first parameter of instance methods always has type `Self`, if it is not exp
 The name `self` is not special in any way.
 
 ```py
+def some_decorator(f: Callable) -> Callable:
+    return f
+
 class B:
     def name_does_not_matter(this) -> Self:
         reveal_type(this)  # revealed: Self@name_does_not_matter
@@ -153,18 +156,45 @@ class B:
         reveal_type(self)  # revealed: Self@keyword_only
         return self
 
+    @some_decorator
+    def decorated_method(self) -> Self:
+        reveal_type(self)  # revealed: Self@decorated_method
+        return self
+
     @property
     def a_property(self) -> Self:
-        # TODO: Should reveal Self@a_property
-        reveal_type(self)  # revealed: Unknown
+        reveal_type(self)  # revealed: Self@a_property
         return self
 
+    async def async_method(self) -> Self:
+        reveal_type(self)  # revealed: Self@async_method
+        return self
+
+    @staticmethod
+    def static_method(self):
+        # The parameter can be called `self`, but it is not treated as `Self`
+        reveal_type(self)  # revealed: Unknown
+
+    @staticmethod
+    @some_decorator
+    def decorated_static_method(self):
+        reveal_type(self)  # revealed: Unknown
+
+    @some_decorator
+    @staticmethod
+    def decorated_static_method_2(self):
+        reveal_type(self)  # revealed: Unknown
+
 reveal_type(B().name_does_not_matter())  # revealed: B
 reveal_type(B().positional_only(1))  # revealed: B
 reveal_type(B().keyword_only(x=1))  # revealed: B
+reveal_type(B().decorated_method())  # revealed: Unknown
 
 # TODO: this should be B
 reveal_type(B().a_property)  # revealed: Unknown
+
+async def _():
+    reveal_type(await B().async_method())  # revealed: B
 ```
 
 This also works for generic classes:

crates/ty_python_semantic/resources/mdtest/overloads.md

@@ -598,6 +598,7 @@ class CheckClassMethod:
     # error: [invalid-overload]
     def try_from3(cls, x: int | str) -> CheckClassMethod | None:
         if isinstance(x, int):
+            # error: [call-non-callable] "Object of type `CheckClassMethod` is not callable"
             return cls(x)
         return None
 

crates/ty_python_semantic/resources/mdtest/snapshots/overloads.md_-_Overloads_-_Invalid_-_Inconsistent_decorat…_-_`@classmethod`_(aaa04d4cfa3adaba).snap

@@ -53,20 +53,21 @@ mdtest path: crates/ty_python_semantic/resources/mdtest/overloads.md
 39 |     # error: [invalid-overload]
 40 |     def try_from3(cls, x: int | str) -> CheckClassMethod | None:
 41 |         if isinstance(x, int):
-42 |             return cls(x)
-43 |         return None
-44 | 
-45 |     @overload
-46 |     @classmethod
-47 |     def try_from4(cls, x: int) -> CheckClassMethod: ...
-48 |     @overload
-49 |     @classmethod
-50 |     def try_from4(cls, x: str) -> None: ...
-51 |     @classmethod
-52 |     def try_from4(cls, x: int | str) -> CheckClassMethod | None:
-53 |         if isinstance(x, int):
-54 |             return cls(x)
-55 |         return None
+42 |             # error: [call-non-callable] "Object of type `CheckClassMethod` is not callable"
+43 |             return cls(x)
+44 |         return None
+45 | 
+46 |     @overload
+47 |     @classmethod
+48 |     def try_from4(cls, x: int) -> CheckClassMethod: ...
+49 |     @overload
+50 |     @classmethod
+51 |     def try_from4(cls, x: str) -> None: ...
+52 |     @classmethod
+53 |     def try_from4(cls, x: int | str) -> CheckClassMethod | None:
+54 |         if isinstance(x, int):
+55 |             return cls(x)
+56 |         return None
 ```
 
 # Diagnostics
@@ -124,8 +125,22 @@ error[invalid-overload]: Overloaded function `try_from3` does not use the `@clas
    |         |
    |         Missing here
 41 |         if isinstance(x, int):
-42 |             return cls(x)
+42 |             # error: [call-non-callable] "Object of type `CheckClassMethod` is not callable"
    |
 info: rule `invalid-overload` is enabled by default
 
 ```
+
+```
+error[call-non-callable]: Object of type `CheckClassMethod` is not callable
+  --> src/mdtest_snippet.py:43:20
+   |
+41 |         if isinstance(x, int):
+42 |             # error: [call-non-callable] "Object of type `CheckClassMethod` is not callable"
+43 |             return cls(x)
+   |                    ^^^^^^
+44 |         return None
+   |
+info: rule `call-non-callable` is enabled by default
+
+```

crates/ty_python_semantic/src/types/function.rs

@@ -185,6 +185,16 @@ pub struct DataclassTransformerParams<'db> {
 
 impl get_size2::GetSize for DataclassTransformerParams<'_> {}
 
+/// Whether a function should implicitly be treated as a staticmethod based on its name.
+pub(crate) fn is_implicit_staticmethod(function_name: &str) -> bool {
+    matches!(function_name, "__new__")
+}
+
+/// Whether a function should implicitly be treated as a classmethod based on its name.
+pub(crate) fn is_implicit_classmethod(function_name: &str) -> bool {
+    matches!(function_name, "__init_subclass__" | "__class_getitem__")
+}
+
 /// Representation of a function definition in the AST: either a non-generic function, or a generic
 /// function that has not been specialized.
 ///
@@ -257,17 +267,15 @@ impl<'db> OverloadLiteral<'db> {
     /// Returns true if this overload is decorated with `@staticmethod`, or if it is implicitly a
     /// staticmethod.
     pub(crate) fn is_staticmethod(self, db: &dyn Db) -> bool {
-        self.has_known_decorator(db, FunctionDecorators::STATICMETHOD) || self.name(db) == "__new__"
+        self.has_known_decorator(db, FunctionDecorators::STATICMETHOD)
+            || is_implicit_staticmethod(self.name(db))
     }
 
     /// Returns true if this overload is decorated with `@classmethod`, or if it is implicitly a
     /// classmethod.
     pub(crate) fn is_classmethod(self, db: &dyn Db) -> bool {
         self.has_known_decorator(db, FunctionDecorators::CLASSMETHOD)
-            || matches!(
-                self.name(db).as_str(),
-                "__init_subclass__" | "__class_getitem__"
-            )
+            || is_implicit_classmethod(self.name(db))
     }
 
     fn node<'ast>(

crates/ty_python_semantic/src/types/infer/builder.rs

@@ -78,6 +78,7 @@ use crate::types::diagnostic::{
 };
 use crate::types::function::{
     FunctionDecorators, FunctionLiteral, FunctionType, KnownFunction, OverloadLiteral,
+    is_implicit_classmethod, is_implicit_staticmethod,
 };
 use crate::types::generics::{
     GenericContext, InferableTypeVars, LegacyGenericBase, SpecializationBuilder, bind_typevar,
@@ -2580,18 +2581,27 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
             return None;
         }
 
-        let method = infer_definition_types(db, method_definition)
-            .declaration_type(method_definition)
-            .inner_type()
-            .as_function_literal()?;
+        let function_node = function_definition.node(self.module());
+        let function_name = &function_node.name;
 
-        if method.is_classmethod(db) {
-            // TODO: set the type for `cls` argument
-            return None;
-        } else if method.is_staticmethod(db) {
+        // TODO: handle implicit type of `cls` for classmethods
+        if is_implicit_classmethod(function_name) || is_implicit_staticmethod(function_name) {
             return None;
         }
 
+        let inference = infer_definition_types(db, method_definition);
+        for decorator in &function_node.decorator_list {
+            let decorator_ty = inference.expression_type(&decorator.expression);
+            if decorator_ty.as_class_literal().is_some_and(|class| {
+                matches!(
+                    class.known(db),
+                    Some(KnownClass::Classmethod | KnownClass::Staticmethod)
+                )
+            }) {
+                return None;
+            }
+        }
+
         let class_definition = self.index.expect_single_definition(class);
         let class_literal = infer_definition_types(db, class_definition)
             .declaration_type(class_definition)