From bd666c85c8f5b4688b7c64ab56983809ed32543f Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Wed, 11 Feb 2026 11:21:29 -0500 Subject: [PATCH 1/6] [ty] Third try at detecting Self --- .../resources/mdtest/annotations/self.md | 239 +++++++++++++++++- crates/ty_python_semantic/src/types.rs | 10 + .../src/types/special_form.rs | 59 ++++- 3 files changed, 297 insertions(+), 11 deletions(-) diff --git a/crates/ty_python_semantic/resources/mdtest/annotations/self.md b/crates/ty_python_semantic/resources/mdtest/annotations/self.md index 2f60d65e973b04..2b1e92d78e6868 100644 --- a/crates/ty_python_semantic/resources/mdtest/annotations/self.md +++ b/crates/ty_python_semantic/resources/mdtest/annotations/self.md @@ -796,20 +796,253 @@ class Foo: return Foo() @staticmethod - # TODO: The usage of `Self` here should be rejected because this is a static method + # error: [invalid-type-form] "`Self` cannot be used in a static method" def make() -> Self: - # error: [invalid-return-type] return Foo() class Bar(Generic[T]): ... # error: [invalid-type-form] class Baz(Bar[Self]): ... +``` + +## Self usage in static methods + +`Self` cannot be used anywhere in a static method, including parameters, return types, nested +functions, and default argument values. + +```py +from typing import Self + +class StaticMethodTests: + @staticmethod + # error: [invalid-type-form] "`Self` cannot be used in a static method" + def with_self_return() -> Self: + pass + + @staticmethod + # error: [invalid-type-form] "`Self` cannot be used in a static method" + def with_self_param(x: Self) -> None: + pass + + @staticmethod + def with_nested_function() -> None: + # `Self` in nested function inside static method is also invalid + # because `Self` binds to the outermost method (the static method). + # error: [invalid-type-form] "`Self` cannot be used in a static method" + def inner() -> Self: + pass + + @staticmethod + # error: [invalid-type-form] "`Self` cannot be used in a static method" + def with_self_default(x: int = 0, y: "Self | None" = None) -> None: + pass +``` + +## Aliased staticmethod decorator + +Using an aliased `staticmethod` decorator should still be detected: + +```py +from typing import Self + +sm = staticmethod + +class AliasedStaticMethod: + @sm + # error: [invalid-type-form] "`Self` cannot be used in a static method" + def aliased_static() -> Self: + pass +``` + +## `__new__` allows `Self` + +`__new__` is a static method even without an explicit `@staticmethod` decorator, but at runtime it +is heavily special-cased by the interpreter to behave more like a classmethod. It always receives a +`cls` parameter with type `type[Self]` and typically returns an object of type `Self`, so `Self` is +permitted in `__new__`: + +```py +from typing import Self + +class WithNew: + def __new__(cls) -> Self: + instance = object.__new__(cls) + return instance + +reveal_type(WithNew()) # revealed: WithNew + +class SubclassWithNew(WithNew): + def __new__(cls) -> Self: + return super().__new__(cls) + +reveal_type(SubclassWithNew()) # revealed: SubclassWithNew +``` + +## Stacked decorators with staticmethod + +When `@staticmethod` is stacked with other decorators, `Self` should still be invalid: + +```py +from typing import Self, Callable, TypeVar + +T = TypeVar("T") + +def identity(f: T) -> T: + return f + +class StackedDecorators: + @staticmethod + @identity + # error: [invalid-type-form] "`Self` cannot be used in a static method" + def static_then_identity() -> Self: + pass + # TODO: On Python <3.10, this should ideally be rejected, because `staticmethod` objects were not callable. + @identity + @staticmethod + # error: [invalid-type-form] "`Self` cannot be used in a static method" + def identity_then_static() -> Self: + pass +``` + +## Self usage in metaclasses + +`Self` cannot be used in a metaclass because the semantics are confusing: in a metaclass, `self` +refers to a class (the metaclass instance), not a regular object instance. + +```py +from typing import Self class MyMetaclass(type): - # TODO: reject the Self usage. because self cannot be used within a metaclass. + # error: [invalid-type-form] "`Self` cannot be used in a metaclass" + registry: list[Self] + + # error: [invalid-type-form] "`Self` cannot be used in a metaclass" def __new__(cls, name, bases, dct) -> Self: return cls(name, bases, dct) + # error: [invalid-type-form] "`Self` cannot be used in a metaclass" + def instance_method(self) -> Self: + return self + + @classmethod + # error: [invalid-type-form] "`Self` cannot be used in a metaclass" + def metaclass_classmethod(cls) -> Self: + return cls("", (), {}) + # Note: static methods in metaclasses get the static method error, not metaclass error + @staticmethod + # error: [invalid-type-form] "`Self` cannot be used in a static method" + def metaclass_staticmethod() -> Self: + pass +``` + +## Runtime use of `self` parameter in metaclass + +Using the `self` parameter as a runtime value (e.g. in `Union[self, other]`) should not be flagged, +even in a metaclass. Only the literal `Self` type form should be disallowed. + +```py +from typing import Union + +class AnnotableMeta(type): + def __or__(self, other): + return Union[self, other] # No error: runtime use of `self`, not the `Self` type form +``` + +## Indirect metaclass inheritance + +Classes that inherit from `type` indirectly (through another metaclass) are also metaclasses: + +```py +from typing import Self +from abc import ABCMeta + +class IndirectMetaclass(ABCMeta): + # error: [invalid-type-form] "`Self` cannot be used in a metaclass" + def method(self) -> Self: + return self + +class MultiLevelMeta(IndirectMetaclass): + # error: [invalid-type-form] "`Self` cannot be used in a metaclass" + def another_method(self) -> Self: + return self +``` + +## Classes using a metaclass are not metaclasses + +A class that uses a metaclass (via `metaclass=...`) is _not_ itself a metaclass. `Self` should be +valid in such classes: + +```py +from typing import Self + +class SomeMeta(type): + pass + +class UsesMetaclass(metaclass=SomeMeta): + def method(self) -> Self: + reveal_type(self) # revealed: Self@method + return self + +reveal_type(UsesMetaclass().method()) # revealed: UsesMetaclass + +class SubclassOfMetaclassUser(UsesMetaclass): + def another(self) -> Self: + return self + +reveal_type(SubclassOfMetaclassUser().another()) # revealed: SubclassOfMetaclassUser +``` + +## Nested class inside a metaclass + +A nested class inside a metaclass is _not_ a metaclass (unless it also inherits from `type`): + +```py +from typing import Self + +class OuterMeta(type): + # error: [invalid-type-form] "`Self` cannot be used in a metaclass" + def meta_method(self) -> Self: + return self + + class NestedRegularClass: + # This is fine - NestedRegularClass is not a metaclass + def method(self) -> Self: + reveal_type(self) # revealed: Self@method + return self + + class NestedMetaclass(type): + # error: [invalid-type-form] "`Self` cannot be used in a metaclass" + def nested_meta_method(self) -> Self: + return self +``` + +## `builtins.staticmethod` + +Using the fully qualified `builtins.staticmethod` should also be detected: + +```py +from typing import Self +import builtins + +class BuiltinsStaticMethod: + @builtins.staticmethod + # error: [invalid-type-form] "`Self` cannot be used in a static method" + def method() -> Self: + pass +``` + +## EnumMeta is a metaclass + +`enum.EnumMeta` (or `enum.EnumType` in Python 3.11+) is a metaclass, so `Self` should be invalid: + +```py +from typing import Self +from enum import EnumMeta + +class CustomEnumMeta(EnumMeta): + # error: [invalid-type-form] "`Self` cannot be used in a metaclass" + def custom_method(self) -> Self: + return self ``` ## Explicit annotations override implicit `Self` diff --git a/crates/ty_python_semantic/src/types.rs b/crates/ty_python_semantic/src/types.rs index f7f1dd4c17e61f..b32442e4a9b135 100644 --- a/crates/ty_python_semantic/src/types.rs +++ b/crates/ty_python_semantic/src/types.rs @@ -6725,6 +6725,10 @@ enum InvalidTypeExpression<'db> { /// Type qualifiers that are invalid in type expressions, /// and which would require exactly one argument even if they appeared in an annotation expression TypeQualifierRequiresOneArgument(TypeQualifier), + /// `typing.Self` cannot be used in `@staticmethod` definitions. + TypingSelfInStaticMethod, + /// `typing.Self` cannot be used in metaclass definitions. + TypingSelfInMetaclass, /// Some types are always invalid in type expressions InvalidType(Type<'db>, ScopeId<'db>), InvalidBareParamSpec(TypeVarInstance<'db>), @@ -6794,6 +6798,12 @@ impl<'db> InvalidTypeExpression<'db> { "Type qualifier `{qualifier}` is not allowed in type expressions \ (only in annotation expressions, and only with exactly one argument)", ), + InvalidTypeExpression::TypingSelfInStaticMethod => { + f.write_str("`Self` cannot be used in a static method") + } + InvalidTypeExpression::TypingSelfInMetaclass => { + f.write_str("`Self` cannot be used in a metaclass") + } InvalidTypeExpression::InvalidType(Type::FunctionLiteral(function), _) => { write!( f, diff --git a/crates/ty_python_semantic/src/types/special_form.rs b/crates/ty_python_semantic/src/types/special_form.rs index 7a70efb211ad2b..d05e1335800879 100644 --- a/crates/ty_python_semantic/src/types/special_form.rs +++ b/crates/ty_python_semantic/src/types/special_form.rs @@ -5,12 +5,14 @@ use super::{ClassType, Type, class::KnownClass}; use crate::db::Db; use crate::semantic_index::place::ScopedPlaceId; use crate::semantic_index::{ - FileScopeId, definition::Definition, place_table, scope::ScopeId, semantic_index, use_def_map, + FileScopeId, definition::{Definition, DefinitionKind}, place_table, scope::ScopeId, + semantic_index, use_def_map, }; use crate::types::IntersectionType; use crate::types::{ - CallableType, InvalidTypeExpression, InvalidTypeExpressionError, TypeDefinition, - TypeQualifiers, generics::typing_self, infer::nearest_enclosing_class, + CallableType, FunctionDecorators, InvalidTypeExpression, InvalidTypeExpressionError, + TypeDefinition, TypeQualifiers, generics::typing_self, + infer::{infer_definition_types, nearest_enclosing_class}, }; use ruff_db::files::File; use strum_macros::EnumString; @@ -678,11 +680,52 @@ impl SpecialFormType { }); }; - Ok( - typing_self(db, scope_id, typevar_binding_context, class.into()) - .map(Type::TypeVar) - .unwrap_or(Type::SpecialForm(self)), - ) + let typing_self = + typing_self(db, scope_id, typevar_binding_context, class.into()); + + let in_staticmethod = typing_self.is_some_and(|typing_self| { + let Some(binding_definition) = typing_self.binding_context(db).definition() + else { + return false; + }; + + let DefinitionKind::Function(_) = binding_definition.kind(db) else { + return false; + }; + + infer_definition_types(db, binding_definition) + .declaration_type(binding_definition) + .inner_type() + .as_function_literal() + .is_some_and(|function| { + function.name(db).as_str() != "__new__" + && function + .has_known_decorator(db, FunctionDecorators::STATICMETHOD) + }) + }); + if in_staticmethod { + return Err(InvalidTypeExpressionError { + fallback_type: Type::unknown(), + invalid_expressions: smallvec::smallvec_inline![ + InvalidTypeExpression::TypingSelfInStaticMethod + ], + }); + } + + let is_in_metaclass = KnownClass::Type + .to_class_literal(db) + .to_class_type(db) + .is_some_and(|type_class| class.is_subclass_of(db, None, type_class)); + if is_in_metaclass { + return Err(InvalidTypeExpressionError { + fallback_type: Type::unknown(), + invalid_expressions: smallvec::smallvec_inline![ + InvalidTypeExpression::TypingSelfInMetaclass + ], + }); + } + + Ok(typing_self.map(Type::TypeVar).unwrap_or(Type::SpecialForm(self))) } // We ensure that `typing.TypeAlias` used in the expected position (annotating an // annotated assignment statement) doesn't reach here. Using it in any other type From abab660689580b0805b73218bfba1d27762f5981 Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Sun, 15 Feb 2026 15:49:29 -0500 Subject: [PATCH 2/6] Use a dedicated Salsa query --- .../resources/mdtest/annotations/self.md | 19 +++++------- crates/ty_python_semantic/src/types/infer.rs | 27 ++++++++++++++++- .../src/types/infer/builder.rs | 18 +++++++++++ .../src/types/special_form.rs | 30 +++++++++---------- 4 files changed, 66 insertions(+), 28 deletions(-) diff --git a/crates/ty_python_semantic/resources/mdtest/annotations/self.md b/crates/ty_python_semantic/resources/mdtest/annotations/self.md index 2b1e92d78e6868..0ea49cdaeb409b 100644 --- a/crates/ty_python_semantic/resources/mdtest/annotations/self.md +++ b/crates/ty_python_semantic/resources/mdtest/annotations/self.md @@ -884,11 +884,9 @@ reveal_type(SubclassWithNew()) # revealed: SubclassWithNew When `@staticmethod` is stacked with other decorators, `Self` should still be invalid: ```py -from typing import Self, Callable, TypeVar +from typing import Self, Callable -T = TypeVar("T") - -def identity(f: T) -> T: +def identity[**P, R](f: Callable[P, R]) -> Callable[P, R]: return f class StackedDecorators: @@ -907,8 +905,7 @@ class StackedDecorators: ## Self usage in metaclasses -`Self` cannot be used in a metaclass because the semantics are confusing: in a metaclass, `self` -refers to a class (the metaclass instance), not a regular object instance. +The spec [prohibits the use of `Self` in metaclasses][spec], so we emit a diagnostic for this. ```py from typing import Self @@ -937,15 +934,13 @@ class MyMetaclass(type): ## Runtime use of `self` parameter in metaclass -Using the `self` parameter as a runtime value (e.g. in `Union[self, other]`) should not be flagged, -even in a metaclass. Only the literal `Self` type form should be disallowed. +Using the `self` parameter as a runtime value should not be flagged, even in a metaclass. Only the +literal `Self` type form should be disallowed. ```py -from typing import Union - class AnnotableMeta(type): def __or__(self, other): - return Union[self, other] # No error: runtime use of `self`, not the `Self` type form + return self # No error: runtime use of `self`, not the `Self` type form ``` ## Indirect metaclass inheritance @@ -1239,3 +1234,5 @@ class Child(Base): reveal_type(Child.attr) # revealed: Child reveal_type(Child().attr) # revealed: Child ``` + +[spec]: https://typing.python.org/en/latest/spec/generics.html#valid-locations-for-self diff --git a/crates/ty_python_semantic/src/types/infer.rs b/crates/ty_python_semantic/src/types/infer.rs index b51edc09a6af85..2065dcc8d9b6d1 100644 --- a/crates/ty_python_semantic/src/types/infer.rs +++ b/crates/ty_python_semantic/src/types/infer.rs @@ -49,7 +49,7 @@ use crate::semantic_index::expression::Expression; use crate::semantic_index::scope::ScopeId; use crate::semantic_index::{SemanticIndex, semantic_index}; use crate::types::diagnostic::TypeCheckDiagnostics; -use crate::types::function::FunctionType; +use crate::types::function::{FunctionDecorators, FunctionType}; use crate::types::generics::Specialization; use crate::types::unpacker::{UnpackResult, Unpacker}; use crate::types::{ @@ -102,6 +102,31 @@ fn definition_cycle_initial<'db>( DefinitionInference::cycle_initial(definition.scope(db), Type::divergent(id)) } +/// Infer just the known-decorator flags for a function definition. +/// +/// This is a lightweight query that avoids the cycle risk of calling +/// `infer_definition_types` when we need to check decorators while +/// already inside definition inference (e.g. checking `Self` in a +/// `@staticmethod`). +/// +/// TODO: This results in double inference of decorator expressions, since +/// `infer_function_definition` also infers them independently. Ideally we'd +/// reuse this query there, but that would require returning a full +/// `TypeInference` (with expression types) rather than just `FunctionDecorators`, +/// increasing memory usage for every function. +#[salsa::tracked(heap_size=ruff_memory_usage::heap_size)] +pub(crate) fn function_known_decorators<'db>( + db: &'db dyn Db, + definition: Definition<'db>, +) -> FunctionDecorators { + let file = definition.file(db); + let module = parsed_module(db, file).load(db); + let index = semantic_index(db, file); + + TypeInferenceBuilder::new(db, InferenceRegion::Definition(definition), index, &module) + .finish_function_decorators() +} + /// Infer types for all deferred type expressions in a [`Definition`]. /// /// Deferred expressions are type expressions (annotations, base classes, aliases...) in a stub diff --git a/crates/ty_python_semantic/src/types/infer/builder.rs b/crates/ty_python_semantic/src/types/infer/builder.rs index 72da0bb9154b99..4afd5084962302 100644 --- a/crates/ty_python_semantic/src/types/infer/builder.rs +++ b/crates/ty_python_semantic/src/types/infer/builder.rs @@ -9026,6 +9026,24 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { } } + pub(super) fn finish_function_decorators(mut self) -> FunctionDecorators { + let InferenceRegion::Definition(definition) = self.region else { + panic!("Expected Definition region"); + }; + let DefinitionKind::Function(func_ref) = definition.kind(self.db()) else { + let _ = self.context.finish(); + return FunctionDecorators::empty(); + }; + let func = func_ref.node(self.module()); + let mut decorators = FunctionDecorators::empty(); + for decorator in &func.decorator_list { + let decorator_type = self.infer_decorator(decorator); + decorators |= FunctionDecorators::from_decorator_type(self.db(), decorator_type); + } + let _ = self.context.finish(); + decorators + } + pub(super) fn finish_definition(mut self) -> DefinitionInference<'db> { self.infer_region(); diff --git a/crates/ty_python_semantic/src/types/special_form.rs b/crates/ty_python_semantic/src/types/special_form.rs index d05e1335800879..4381170ffedfbc 100644 --- a/crates/ty_python_semantic/src/types/special_form.rs +++ b/crates/ty_python_semantic/src/types/special_form.rs @@ -5,14 +5,17 @@ use super::{ClassType, Type, class::KnownClass}; use crate::db::Db; use crate::semantic_index::place::ScopedPlaceId; use crate::semantic_index::{ - FileScopeId, definition::{Definition, DefinitionKind}, place_table, scope::ScopeId, + FileScopeId, + definition::{Definition, DefinitionKind}, + place_table, + scope::ScopeId, semantic_index, use_def_map, }; use crate::types::IntersectionType; use crate::types::{ CallableType, FunctionDecorators, InvalidTypeExpression, InvalidTypeExpressionError, TypeDefinition, TypeQualifiers, generics::typing_self, - infer::{infer_definition_types, nearest_enclosing_class}, + infer::{function_known_decorators, nearest_enclosing_class}, }; use ruff_db::files::File; use strum_macros::EnumString; @@ -680,8 +683,7 @@ impl SpecialFormType { }); }; - let typing_self = - typing_self(db, scope_id, typevar_binding_context, class.into()); + let typing_self = typing_self(db, scope_id, typevar_binding_context, class.into()); let in_staticmethod = typing_self.is_some_and(|typing_self| { let Some(binding_definition) = typing_self.binding_context(db).definition() @@ -689,19 +691,13 @@ impl SpecialFormType { return false; }; - let DefinitionKind::Function(_) = binding_definition.kind(db) else { + if !matches!(binding_definition.kind(db), DefinitionKind::Function(_)) { return false; - }; + } - infer_definition_types(db, binding_definition) - .declaration_type(binding_definition) - .inner_type() - .as_function_literal() - .is_some_and(|function| { - function.name(db).as_str() != "__new__" - && function - .has_known_decorator(db, FunctionDecorators::STATICMETHOD) - }) + binding_definition.name(db).as_deref() != Some("__new__") + && function_known_decorators(db, binding_definition) + .contains(FunctionDecorators::STATICMETHOD) }); if in_staticmethod { return Err(InvalidTypeExpressionError { @@ -725,7 +721,9 @@ impl SpecialFormType { }); } - Ok(typing_self.map(Type::TypeVar).unwrap_or(Type::SpecialForm(self))) + Ok(typing_self + .map(Type::TypeVar) + .unwrap_or(Type::SpecialForm(self))) } // We ensure that `typing.TypeAlias` used in the expected position (annotating an // annotated assignment statement) doesn't reach here. Using it in any other type From 02e1c2bbc74f3324d968477194805747fdd612b7 Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Wed, 25 Feb 2026 12:29:22 -0500 Subject: [PATCH 3/6] Return a TypeInference --- crates/ty_python_semantic/src/types/infer.rs | 40 ++++++--- .../src/types/infer/builder.rs | 85 ++++++++++++++++--- .../src/types/infer/builder/function.rs | 7 +- .../src/types/special_form.rs | 13 ++- 4 files changed, 116 insertions(+), 29 deletions(-) diff --git a/crates/ty_python_semantic/src/types/infer.rs b/crates/ty_python_semantic/src/types/infer.rs index 2065dcc8d9b6d1..61d4793ec2123a 100644 --- a/crates/ty_python_semantic/src/types/infer.rs +++ b/crates/ty_python_semantic/src/types/infer.rs @@ -44,7 +44,7 @@ use salsa::plumbing::AsId; use crate::Db; use crate::semantic_index::ast_ids::node_key::ExpressionNodeKey; -use crate::semantic_index::definition::Definition; +use crate::semantic_index::definition::{Definition, DefinitionKind}; use crate::semantic_index::expression::Expression; use crate::semantic_index::scope::ScopeId; use crate::semantic_index::{SemanticIndex, semantic_index}; @@ -102,29 +102,47 @@ fn definition_cycle_initial<'db>( DefinitionInference::cycle_initial(definition.scope(db), Type::divergent(id)) } -/// Infer just the known-decorator flags for a function definition. +/// Infer decorator expression types for a function definition. /// /// This is a lightweight query that avoids the cycle risk of calling /// `infer_definition_types` when we need to check decorators while /// already inside definition inference (e.g. checking `Self` in a /// `@staticmethod`). -/// -/// TODO: This results in double inference of decorator expressions, since -/// `infer_function_definition` also infers them independently. Ideally we'd -/// reuse this query there, but that would require returning a full -/// `TypeInference` (with expression types) rather than just `FunctionDecorators`, -/// increasing memory usage for every function. -#[salsa::tracked(heap_size=ruff_memory_usage::heap_size)] +#[salsa::tracked(returns(ref), heap_size=ruff_memory_usage::heap_size)] pub(crate) fn function_known_decorators<'db>( db: &'db dyn Db, definition: Definition<'db>, -) -> FunctionDecorators { +) -> DefinitionInference<'db> { let file = definition.file(db); let module = parsed_module(db, file).load(db); let index = semantic_index(db, file); TypeInferenceBuilder::new(db, InferenceRegion::Definition(definition), index, &module) - .finish_function_decorators() + .finish_function_decorator_types() +} + +pub(crate) fn function_known_decorator_flags<'db>( + db: &'db dyn Db, + definition: Definition<'db>, +) -> FunctionDecorators { + let DefinitionKind::Function(function) = definition.kind(db) else { + return FunctionDecorators::empty(); + }; + + let file = definition.file(db); + let module = parsed_module(db, file).load(db); + let decorator_inference = function_known_decorators(db, definition); + + function.node(&module).decorator_list.iter().fold( + FunctionDecorators::empty(), + |decorators, decorator| { + decorators + | FunctionDecorators::from_decorator_type( + db, + decorator_inference.expression_type(&decorator.expression), + ) + }, + ) } /// Infer types for all deferred type expressions in a [`Definition`]. diff --git a/crates/ty_python_semantic/src/types/infer/builder.rs b/crates/ty_python_semantic/src/types/infer/builder.rs index 4afd5084962302..b937a0b0a3b62b 100644 --- a/crates/ty_python_semantic/src/types/infer/builder.rs +++ b/crates/ty_python_semantic/src/types/infer/builder.rs @@ -85,6 +85,7 @@ use crate::types::diagnostic::{ report_unsupported_comparison, }; use crate::types::enums::{enum_ignored_names, is_enum_class_by_inheritance}; +use crate::types::context::InNoTypeCheck; use crate::types::function::{FunctionType, KnownFunction, report_revealed_type}; use crate::types::generics::{InferableTypeVars, SpecializationBuilder, bind_typevar}; use crate::types::infer::builder::named_tuple::NamedTupleKind; @@ -9026,22 +9027,82 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { } } - pub(super) fn finish_function_decorators(mut self) -> FunctionDecorators { + pub(super) fn finish_function_decorator_types(mut self) -> DefinitionInference<'db> { let InferenceRegion::Definition(definition) = self.region else { panic!("Expected Definition region"); }; - let DefinitionKind::Function(func_ref) = definition.kind(self.db()) else { - let _ = self.context.finish(); - return FunctionDecorators::empty(); - }; - let func = func_ref.node(self.module()); - let mut decorators = FunctionDecorators::empty(); - for decorator in &func.decorator_list { - let decorator_type = self.infer_decorator(decorator); - decorators |= FunctionDecorators::from_decorator_type(self.db(), decorator_type); + if let DefinitionKind::Function(func_ref) = definition.kind(self.db()) { + let func = func_ref.node(self.module()); + for decorator in &func.decorator_list { + let decorator_type = self.infer_decorator(decorator); + if let Type::FunctionLiteral(function) = decorator_type + && let Some(KnownFunction::NoTypeCheck) = function.known(self.db()) + { + // Match `infer_function_definition`: suppress diagnostics that follow + // `@no_type_check`, including later decorators. + self.context.set_in_no_type_check(InNoTypeCheck::Yes); + } + } + } + + let Self { + context, + mut expressions, + string_annotations, + scope, + bindings, + declarations, + deferred, + cycle_recovery, + undecorated_type, + called_functions, + // builder only state + dataclass_field_specifiers: _, + all_definitely_bound: _, + typevar_binding_context: _, + inference_flags: _, + deferred_state: _, + inferring_vararg_annotation: _, + multi_inference_state: _, + inner_expression_inference_state: _, + index: _, + region: _, + return_types_and_ranges: _, + } = self; + + let _ = scope; + let diagnostics = context.finish(); + + let extra = (!diagnostics.is_empty() + || !string_annotations.is_empty() + || cycle_recovery.is_some() + || undecorated_type.is_some() + || !deferred.is_empty() + || !called_functions.is_empty()) + .then(|| { + Box::new(DefinitionInferenceExtra { + string_annotations, + called_functions: called_functions + .into_iter() + .collect::>() + .into_boxed_slice(), + cycle_recovery, + deferred: deferred.into_boxed_slice(), + diagnostics, + undecorated_type, + }) + }); + + expressions.shrink_to_fit(); + + DefinitionInference { + expressions, + #[cfg(debug_assertions)] + scope, + bindings: bindings.into_boxed_slice(), + declarations: declarations.into_boxed_slice(), + extra, } - let _ = self.context.finish(); - decorators } pub(super) fn finish_definition(mut self) -> DefinitionInference<'db> { diff --git a/crates/ty_python_semantic/src/types/infer/builder/function.rs b/crates/ty_python_semantic/src/types/infer/builder/function.rs index 7ab3e601058940..fe8d391ac41b33 100644 --- a/crates/ty_python_semantic/src/types/infer/builder/function.rs +++ b/crates/ty_python_semantic/src/types/infer/builder/function.rs @@ -21,7 +21,7 @@ use crate::{ }, generics::{enclosing_generic_contexts, typing_self}, infer::{ - TypeInferenceBuilder, + TypeInferenceBuilder, function_known_decorators, builder::{ DeclaredAndInferredType, DeferredExpressionState, TypeAndRange, validate_paramspec_components, @@ -220,6 +220,9 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { let db = self.db(); + let decorator_inference = function_known_decorators(db, definition); + self.extend_definition(decorator_inference); + let mut decorator_types_and_nodes = Vec::with_capacity(decorator_list.len()); let mut function_decorators = FunctionDecorators::empty(); let mut deprecated = None; @@ -227,7 +230,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { let mut final_decorator = None; for decorator in decorator_list { - let decorator_type = self.infer_decorator(decorator); + let decorator_type = decorator_inference.expression_type(&decorator.expression); let decorator_function_decorator = FunctionDecorators::from_decorator_type(db, decorator_type); function_decorators |= decorator_function_decorator; diff --git a/crates/ty_python_semantic/src/types/special_form.rs b/crates/ty_python_semantic/src/types/special_form.rs index 4381170ffedfbc..2e7a6e38bf15a5 100644 --- a/crates/ty_python_semantic/src/types/special_form.rs +++ b/crates/ty_python_semantic/src/types/special_form.rs @@ -14,8 +14,9 @@ use crate::semantic_index::{ use crate::types::IntersectionType; use crate::types::{ CallableType, FunctionDecorators, InvalidTypeExpression, InvalidTypeExpressionError, - TypeDefinition, TypeQualifiers, generics::typing_self, - infer::{function_known_decorators, nearest_enclosing_class}, + TypeDefinition, TypeQualifiers, + generics::typing_self, + infer::{function_known_decorator_flags, nearest_enclosing_class}, }; use ruff_db::files::File; use strum_macros::EnumString; @@ -696,7 +697,7 @@ impl SpecialFormType { } binding_definition.name(db).as_deref() != Some("__new__") - && function_known_decorators(db, binding_definition) + && function_known_decorator_flags(db, binding_definition) .contains(FunctionDecorators::STATICMETHOD) }); if in_staticmethod { @@ -711,7 +712,11 @@ impl SpecialFormType { let is_in_metaclass = KnownClass::Type .to_class_literal(db) .to_class_type(db) - .is_some_and(|type_class| class.is_subclass_of(db, None, type_class)); + .is_some_and(|type_class| { + class + .default_specialization(db) + .is_subclass_of(db, type_class) + }); if is_in_metaclass { return Err(InvalidTypeExpressionError { fallback_type: Type::unknown(), From c3e81cecc7abb795ca20af2513dfdd30cf15b5ec Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Wed, 25 Feb 2026 16:41:32 -0500 Subject: [PATCH 4/6] Use a smaller struct --- crates/ty_python_semantic/src/types/infer.rs | 55 +++++++++----- .../src/types/infer/builder.rs | 76 +++++-------------- .../src/types/infer/builder/function.rs | 16 +++- 3 files changed, 65 insertions(+), 82 deletions(-) diff --git a/crates/ty_python_semantic/src/types/infer.rs b/crates/ty_python_semantic/src/types/infer.rs index 61d4793ec2123a..95267804ffd790 100644 --- a/crates/ty_python_semantic/src/types/infer.rs +++ b/crates/ty_python_semantic/src/types/infer.rs @@ -44,7 +44,7 @@ use salsa::plumbing::AsId; use crate::Db; use crate::semantic_index::ast_ids::node_key::ExpressionNodeKey; -use crate::semantic_index::definition::{Definition, DefinitionKind}; +use crate::semantic_index::definition::Definition; use crate::semantic_index::expression::Expression; use crate::semantic_index::scope::ScopeId; use crate::semantic_index::{SemanticIndex, semantic_index}; @@ -112,37 +112,52 @@ fn definition_cycle_initial<'db>( pub(crate) fn function_known_decorators<'db>( db: &'db dyn Db, definition: Definition<'db>, -) -> DefinitionInference<'db> { +) -> FunctionDecoratorInference<'db> { let file = definition.file(db); let module = parsed_module(db, file).load(db); let index = semantic_index(db, file); TypeInferenceBuilder::new(db, InferenceRegion::Definition(definition), index, &module) - .finish_function_decorator_types() + .finish_function_decorator_inference() } pub(crate) fn function_known_decorator_flags<'db>( db: &'db dyn Db, definition: Definition<'db>, ) -> FunctionDecorators { - let DefinitionKind::Function(function) = definition.kind(db) else { - return FunctionDecorators::empty(); - }; + function_known_decorators(db, definition).known_decorators() +} - let file = definition.file(db); - let module = parsed_module(db, file).load(db); - let decorator_inference = function_known_decorators(db, definition); - - function.node(&module).decorator_list.iter().fold( - FunctionDecorators::empty(), - |decorators, decorator| { - decorators - | FunctionDecorators::from_decorator_type( - db, - decorator_inference.expression_type(&decorator.expression), - ) - }, - ) +/// A compact inference result for function decorators. +/// +/// Unlike [`DefinitionInference`], this stores only decorator expression types and +/// diagnostics, plus precomputed known-decorator flags. +#[derive(Debug, Eq, PartialEq, salsa::Update, get_size2::GetSize)] +pub(crate) struct FunctionDecoratorInference<'db> { + expression_types: FxHashMap>, + known_decorators: FunctionDecorators, + diagnostics: TypeCheckDiagnostics, +} + +impl<'db> FunctionDecoratorInference<'db> { + pub(crate) fn expression_type( + &self, + expression: impl Into, + ) -> Option> { + self.expression_types.get(&expression.into()).copied() + } + + pub(crate) fn expression_types(&self) -> &FxHashMap> { + &self.expression_types + } + + pub(crate) fn known_decorators(&self) -> FunctionDecorators { + self.known_decorators + } + + pub(crate) fn diagnostics(&self) -> &TypeCheckDiagnostics { + &self.diagnostics + } } /// Infer types for all deferred type expressions in a [`Definition`]. diff --git a/crates/ty_python_semantic/src/types/infer/builder.rs b/crates/ty_python_semantic/src/types/infer/builder.rs index b937a0b0a3b62b..191189e0180621 100644 --- a/crates/ty_python_semantic/src/types/infer/builder.rs +++ b/crates/ty_python_semantic/src/types/infer/builder.rs @@ -20,9 +20,9 @@ use ty_module_resolver::{KnownModule, ModuleName, resolve_module}; use super::deferred; use super::{ DefinitionInference, DefinitionInferenceExtra, ExpressionInference, ExpressionInferenceExtra, - InferenceRegion, ScopeInference, ScopeInferenceExtra, infer_deferred_types, - infer_definition_types, infer_expression_types, infer_same_file_expression_type, - infer_unpack_types, + FunctionDecoratorInference, InferenceRegion, ScopeInference, ScopeInferenceExtra, + infer_deferred_types, infer_definition_types, infer_expression_types, + infer_same_file_expression_type, infer_unpack_types, }; use crate::diagnostic::format_enumeration; use crate::node_key::NodeKey; @@ -58,6 +58,7 @@ use crate::types::class::{ DynamicMetaclassConflict, MethodDecorator, }; use crate::types::constraints::{ConstraintSetBuilder, PathBounds, Solutions}; +use crate::types::context::InNoTypeCheck; use crate::types::context::InferContext; use crate::types::diagnostic::{ self, CALL_NON_CALLABLE, CONFLICTING_DECLARATIONS, CYCLIC_CLASS_DEFINITION, @@ -85,8 +86,9 @@ use crate::types::diagnostic::{ report_unsupported_comparison, }; use crate::types::enums::{enum_ignored_names, is_enum_class_by_inheritance}; -use crate::types::context::InNoTypeCheck; -use crate::types::function::{FunctionType, KnownFunction, report_revealed_type}; +use crate::types::function::{ + FunctionDecorators, FunctionType, KnownFunction, report_revealed_type, +}; use crate::types::generics::{InferableTypeVars, SpecializationBuilder, bind_typevar}; use crate::types::infer::builder::named_tuple::NamedTupleKind; use crate::types::infer::builder::paramspec_validation::validate_paramspec_components; @@ -9027,14 +9029,18 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { } } - pub(super) fn finish_function_decorator_types(mut self) -> DefinitionInference<'db> { + pub(super) fn finish_function_decorator_inference(mut self) -> FunctionDecoratorInference<'db> { let InferenceRegion::Definition(definition) = self.region else { panic!("Expected Definition region"); }; + + let mut known_decorators = FunctionDecorators::empty(); if let DefinitionKind::Function(func_ref) = definition.kind(self.db()) { let func = func_ref.node(self.module()); for decorator in &func.decorator_list { let decorator_type = self.infer_decorator(decorator); + known_decorators |= + FunctionDecorators::from_decorator_type(self.db(), decorator_type); if let Type::FunctionLiteral(function) = decorator_type && let Some(KnownFunction::NoTypeCheck) = function.known(self.db()) { @@ -9047,61 +9053,15 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { let Self { context, - mut expressions, - string_annotations, - scope, - bindings, - declarations, - deferred, - cycle_recovery, - undecorated_type, - called_functions, - // builder only state - dataclass_field_specifiers: _, - all_definitely_bound: _, - typevar_binding_context: _, - inference_flags: _, - deferred_state: _, - inferring_vararg_annotation: _, - multi_inference_state: _, - inner_expression_inference_state: _, - index: _, - region: _, - return_types_and_ranges: _, + expressions, + .. } = self; - - let _ = scope; let diagnostics = context.finish(); - let extra = (!diagnostics.is_empty() - || !string_annotations.is_empty() - || cycle_recovery.is_some() - || undecorated_type.is_some() - || !deferred.is_empty() - || !called_functions.is_empty()) - .then(|| { - Box::new(DefinitionInferenceExtra { - string_annotations, - called_functions: called_functions - .into_iter() - .collect::>() - .into_boxed_slice(), - cycle_recovery, - deferred: deferred.into_boxed_slice(), - diagnostics, - undecorated_type, - }) - }); - - expressions.shrink_to_fit(); - - DefinitionInference { - expressions, - #[cfg(debug_assertions)] - scope, - bindings: bindings.into_boxed_slice(), - declarations: declarations.into_boxed_slice(), - extra, + FunctionDecoratorInference { + expression_types: expressions, + known_decorators, + diagnostics, } } diff --git a/crates/ty_python_semantic/src/types/infer/builder/function.rs b/crates/ty_python_semantic/src/types/infer/builder/function.rs index fe8d391ac41b33..47b8e9b1d901ef 100644 --- a/crates/ty_python_semantic/src/types/infer/builder/function.rs +++ b/crates/ty_python_semantic/src/types/infer/builder/function.rs @@ -21,12 +21,12 @@ use crate::{ }, generics::{enclosing_generic_contexts, typing_self}, infer::{ - TypeInferenceBuilder, function_known_decorators, + TypeInferenceBuilder, builder::{ DeclaredAndInferredType, DeferredExpressionState, TypeAndRange, validate_paramspec_components, }, - nearest_enclosing_function, + function_known_decorators, nearest_enclosing_function, }, infer_definition_types, infer_scope_types, todo_type, }, @@ -221,7 +221,13 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { let db = self.db(); let decorator_inference = function_known_decorators(db, definition); - self.extend_definition(decorator_inference); + self.context.extend(decorator_inference.diagnostics()); + self.expressions.extend( + decorator_inference + .expression_types() + .iter() + .map(|(expression, ty)| (*expression, *ty)), + ); let mut decorator_types_and_nodes = Vec::with_capacity(decorator_list.len()); let mut function_decorators = FunctionDecorators::empty(); @@ -230,7 +236,9 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { let mut final_decorator = None; for decorator in decorator_list { - let decorator_type = decorator_inference.expression_type(&decorator.expression); + let decorator_type = decorator_inference + .expression_type(&decorator.expression) + .unwrap_or_else(Type::unknown); let decorator_function_decorator = FunctionDecorators::from_decorator_type(db, decorator_type); function_decorators |= decorator_function_decorator; From eb848d137a1521a650e45b038ec19b50d5b3b657 Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Wed, 25 Mar 2026 10:33:57 -0400 Subject: [PATCH 5/6] Use InferenceRegion::FunctionDecorators --- crates/ty_python_semantic/src/types/infer.rs | 20 +++++-- .../src/types/infer/builder.rs | 60 +++++++++++++------ 2 files changed, 55 insertions(+), 25 deletions(-) diff --git a/crates/ty_python_semantic/src/types/infer.rs b/crates/ty_python_semantic/src/types/infer.rs index 95267804ffd790..29f7b643dec3c5 100644 --- a/crates/ty_python_semantic/src/types/infer.rs +++ b/crates/ty_python_semantic/src/types/infer.rs @@ -1,5 +1,6 @@ //! We have Salsa queries for inferring types at three different granularities: scope-level, -//! definition-level, and expression-level. +//! definition-level, and expression-level, plus lightweight queries for focused subregions like +//! function decorators. //! //! Scope-level inference is for when we are actually checking a file, and need to check types for //! everything in that file's scopes, or give a linter access to types of arbitrary expressions @@ -117,8 +118,13 @@ pub(crate) fn function_known_decorators<'db>( let module = parsed_module(db, file).load(db); let index = semantic_index(db, file); - TypeInferenceBuilder::new(db, InferenceRegion::Definition(definition), index, &module) - .finish_function_decorator_inference() + TypeInferenceBuilder::new( + db, + InferenceRegion::FunctionDecorators(definition), + index, + &module, + ) + .finish_function_decorator_inference() } pub(crate) fn function_known_decorator_flags<'db>( @@ -571,6 +577,8 @@ pub(crate) enum InferenceRegion<'db> { Expression(Expression<'db>, TypeContext<'db>), /// infer types for a [`Definition`] Definition(Definition<'db>), + /// infer types for the decorators on a function [`Definition`] + FunctionDecorators(Definition<'db>), /// infer deferred types for a [`Definition`] Deferred(Definition<'db>), /// infer types for an entire [`ScopeId`] @@ -581,9 +589,9 @@ impl<'db> InferenceRegion<'db> { fn scope(self, db: &'db dyn Db) -> ScopeId<'db> { match self { InferenceRegion::Expression(expression, _) => expression.scope(db), - InferenceRegion::Definition(definition) | InferenceRegion::Deferred(definition) => { - definition.scope(db) - } + InferenceRegion::Definition(definition) + | InferenceRegion::FunctionDecorators(definition) + | InferenceRegion::Deferred(definition) => definition.scope(db), InferenceRegion::Scope(scope, _) => scope, } } diff --git a/crates/ty_python_semantic/src/types/infer/builder.rs b/crates/ty_python_semantic/src/types/infer/builder.rs index 191189e0180621..92be14402917a4 100644 --- a/crates/ty_python_semantic/src/types/infer/builder.rs +++ b/crates/ty_python_semantic/src/types/infer/builder.rs @@ -174,7 +174,9 @@ const NUM_FIELD_SPECIFIERS_INLINE: usize = 1; /// /// The `finish` methods call [`infer_region`](TypeInferenceBuilder::infer_region), which delegates /// to one of [`infer_region_scope`](TypeInferenceBuilder::infer_region_scope), -/// [`infer_region_definition`](TypeInferenceBuilder::infer_region_definition), or +/// [`infer_region_definition`](TypeInferenceBuilder::infer_region_definition), +/// [`infer_region_function_decorators`](TypeInferenceBuilder::infer_region_function_decorators), +/// [`infer_region_deferred`](TypeInferenceBuilder::infer_region_deferred), or /// [`infer_region_expression`](TypeInferenceBuilder::infer_region_expression), depending which /// kind of [`InferenceRegion`] we are inferring types for. /// @@ -590,6 +592,9 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { match self.region { InferenceRegion::Scope(scope, tcx) => self.infer_region_scope(scope, tcx), InferenceRegion::Definition(definition) => self.infer_region_definition(definition), + InferenceRegion::FunctionDecorators(definition) => { + self.infer_region_function_decorators(definition); + } InferenceRegion::Deferred(definition) => self.infer_region_deferred(definition), InferenceRegion::Expression(expression, tcx) => { self.infer_region_expression(expression, tcx); @@ -832,6 +837,23 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { } } + fn infer_region_function_decorators(&mut self, definition: Definition<'db>) { + let DefinitionKind::Function(function) = definition.kind(self.db()) else { + return; + }; + + for decorator in &function.node(self.module()).decorator_list { + let decorator_type = self.infer_decorator(decorator); + if let Type::FunctionLiteral(function) = decorator_type + && let Some(KnownFunction::NoTypeCheck) = function.known(self.db()) + { + // Match `infer_function_definition`: suppress diagnostics that follow + // `@no_type_check`, including later decorators. + self.context.set_in_no_type_check(InNoTypeCheck::Yes); + } + } + } + fn infer_region_deferred(&mut self, definition: Definition<'db>) { // N.B. We don't defer the types for an annotated assignment here because it is done in // the same definition query. It utilizes the deferred expression state instead. @@ -9030,26 +9052,26 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { } pub(super) fn finish_function_decorator_inference(mut self) -> FunctionDecoratorInference<'db> { - let InferenceRegion::Definition(definition) = self.region else { - panic!("Expected Definition region"); - }; + self.infer_region(); - let mut known_decorators = FunctionDecorators::empty(); - if let DefinitionKind::Function(func_ref) = definition.kind(self.db()) { - let func = func_ref.node(self.module()); - for decorator in &func.decorator_list { - let decorator_type = self.infer_decorator(decorator); - known_decorators |= - FunctionDecorators::from_decorator_type(self.db(), decorator_type); - if let Type::FunctionLiteral(function) = decorator_type - && let Some(KnownFunction::NoTypeCheck) = function.known(self.db()) - { - // Match `infer_function_definition`: suppress diagnostics that follow - // `@no_type_check`, including later decorators. - self.context.set_in_no_type_check(InNoTypeCheck::Yes); + let known_decorators = match self.region { + InferenceRegion::FunctionDecorators(definition) => match definition.kind(self.db()) { + DefinitionKind::Function(function) => { + function.node(self.module()).decorator_list.iter().fold( + FunctionDecorators::empty(), + |known_decorators, decorator| { + known_decorators + | FunctionDecorators::from_decorator_type( + self.db(), + self.expression_type(&decorator.expression), + ) + }, + ) } - } - } + _ => FunctionDecorators::empty(), + }, + _ => FunctionDecorators::empty(), + }; let Self { context, From 034f9e6230aef0f0e045f79e035fa0633416c353 Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Wed, 25 Mar 2026 10:45:06 -0400 Subject: [PATCH 6/6] Track bindings --- .../resources/mdtest/decorators.md | 16 ++++++++++++ crates/ty_python_semantic/src/types/infer.rs | 15 ++++++++++- .../src/types/infer/builder.rs | 25 ++++++++++++++++++- .../src/types/infer/builder/function.rs | 4 +++ 4 files changed, 58 insertions(+), 2 deletions(-) diff --git a/crates/ty_python_semantic/resources/mdtest/decorators.md b/crates/ty_python_semantic/resources/mdtest/decorators.md index 941e5d7a74c687..c4b6e5fc63ed46 100644 --- a/crates/ty_python_semantic/resources/mdtest/decorators.md +++ b/crates/ty_python_semantic/resources/mdtest/decorators.md @@ -58,6 +58,22 @@ reveal_type(even) # revealed: (int, /) -> bool reveal_type(even(14)) # revealed: bool ``` +Decorator expressions can also introduce bindings that remain visible after the decorated +definition: + +```py +def decorator_factory(flag: bool): + def decorator(func): + return func + return decorator + +@decorator_factory(seen := True) +def f(): + pass + +reveal_type(seen) # revealed: Literal[True] +``` + ## Multiple decorators Multiple decorators can be applied to a single function. They are applied in "bottom-up" order, diff --git a/crates/ty_python_semantic/src/types/infer.rs b/crates/ty_python_semantic/src/types/infer.rs index 29f7b643dec3c5..5a2ee4441422a2 100644 --- a/crates/ty_python_semantic/src/types/infer.rs +++ b/crates/ty_python_semantic/src/types/infer.rs @@ -137,10 +137,13 @@ pub(crate) fn function_known_decorator_flags<'db>( /// A compact inference result for function decorators. /// /// Unlike [`DefinitionInference`], this stores only decorator expression types and -/// diagnostics, plus precomputed known-decorator flags. +/// diagnostics, plus the expression-side state that needs to be merged back into +/// function-definition inference. #[derive(Debug, Eq, PartialEq, salsa::Update, get_size2::GetSize)] pub(crate) struct FunctionDecoratorInference<'db> { expression_types: FxHashMap>, + bindings: Box<[(Definition<'db>, Type<'db>)]>, + called_functions: Box<[FunctionType<'db>]>, known_decorators: FunctionDecorators, diagnostics: TypeCheckDiagnostics, } @@ -157,6 +160,16 @@ impl<'db> FunctionDecoratorInference<'db> { &self.expression_types } + pub(crate) fn bindings( + &self, + ) -> impl ExactSizeIterator, Type<'db>)> + '_ { + self.bindings.iter().copied() + } + + pub(crate) fn called_functions(&self) -> &[FunctionType<'db>] { + &self.called_functions + } + pub(crate) fn known_decorators(&self) -> FunctionDecorators { self.known_decorators } diff --git a/crates/ty_python_semantic/src/types/infer/builder.rs b/crates/ty_python_semantic/src/types/infer/builder.rs index 92be14402917a4..d7940755c9258c 100644 --- a/crates/ty_python_semantic/src/types/infer/builder.rs +++ b/crates/ty_python_semantic/src/types/infer/builder.rs @@ -9076,12 +9076,35 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { let Self { context, expressions, - .. + bindings, + called_functions, + declarations: _, + deferred: _, + scope: _, + string_annotations: _, + return_types_and_ranges: _, + dataclass_field_specifiers: _, + undecorated_type: _, + typevar_binding_context: _, + inference_flags: _, + deferred_state: _, + multi_inference_state: _, + inner_expression_inference_state: _, + inferring_vararg_annotation: _, + index: _, + region: _, + cycle_recovery: _, + all_definitely_bound: _, } = self; let diagnostics = context.finish(); FunctionDecoratorInference { expression_types: expressions, + bindings: bindings.into_boxed_slice(), + called_functions: called_functions + .into_iter() + .collect::>() + .into_boxed_slice(), known_decorators, diagnostics, } diff --git a/crates/ty_python_semantic/src/types/infer/builder/function.rs b/crates/ty_python_semantic/src/types/infer/builder/function.rs index 47b8e9b1d901ef..841a4dc23dca25 100644 --- a/crates/ty_python_semantic/src/types/infer/builder/function.rs +++ b/crates/ty_python_semantic/src/types/infer/builder/function.rs @@ -228,6 +228,10 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { .iter() .map(|(expression, ty)| (*expression, *ty)), ); + self.bindings + .extend(decorator_inference.bindings(), self.multi_inference_state); + self.called_functions + .extend(decorator_inference.called_functions().iter().copied()); let mut decorator_types_and_nodes = Vec::with_capacity(decorator_list.len()); let mut function_decorators = FunctionDecorators::empty();