[ty] Infer type of self for decorated methods and properties (#21123)
## Summary Infer a type of unannotated `self` parameters in decorated methods / properties. closes https://github.com/astral-sh/ty/issues/1448 ## Test Plan Existing tests, some new tests.
This commit is contained in:
@@ -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>(
|
||||
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user