Repository: astral-sh/ruff
File: crates/ty_python_semantic/src/types/infer/builder/type_expression.rs (Rust)

Code before:
};
use crate::types::signatures::Signature;
use crate::types::string_annotation::parse_string_annotation;
use crate::types::tuple::{TupleSpecBuilder, TupleType};
use crate::types::visitor::any_over_type;
use crate::types::{
    CallableType, DynamicType, IntersectionBuilder, KnownClass, KnownInstanceType,
    LintDiagnosticGuard, Parameter, Parameters, SpecialFormType, SubclassOfType, Type,
    TypeAliasType, TypeContext, TypeIsType, UnionBuilder, UnionType, todo_type,
};

/// Type expressions
impl<'db> TypeInferenceBuilder<'db, '_> {
    /// Infer the type of a type expression.
    pub(super) fn infer_type_expression(&mut self, expression: &ast::Expr) -> Type<'db> {
        let mut ty = self.infer_type_expression_no_store(expression);
        let divergent = Type::divergent(Some(self.scope()));
        if ty.has_divergent_type(self.db(), divergent) {
            ty = divergent;
        }
        self.store_expression_type(expression, ty);
        ty
    }

    /// Similar to [`infer_type_expression`], but accepts a [`DeferredExpressionState`].
    ///
    /// [`infer_type_expression`]: TypeInferenceBuilder::infer_type_expression
    fn infer_type_expression_with_state(
        &mut self,
        expression: &ast::Expr,
        deferred_state: DeferredExpressionState,
    ) -> Type<'db> {
        let previous_deferred_state = std::mem::replace(&mut self.deferred_state, deferred_state);
        let annotation_ty = self.infer_type_expression(expression);
        self.deferred_state = previous_deferred_state;
        annotation_ty
    }

    /// Similar to [`infer_type_expression`], but accepts an optional expression.
    ///
    /// [`infer_type_expression`]: TypeInferenceBuilder::infer_type_expression_with_state
    pub(super) fn infer_optional_type_expression(
        &mut self,
        expression: Option<&ast::Expr>,
    ) -> Option<Type<'db>> {
        expression.map(|expr| self.infer_type_expression(expr))
    }

    fn report_invalid_type_expression(
        &self,
        expression: &ast::Expr,

Review comment (refactor):
In general, setting the deferred state without a paired revert back to the previous deferred state is not right; it will cause the new deferred state to leak back up the stack, outside of this type expression inference.

I feel like there might be some refactor called for here between this method and `infer_type_expression_with_state`, but I haven't quite sorted yet what it should be...

Code after:
};
use crate::types::signatures::Signature;
use crate::types::string_annotation::parse_string_annotation;
use crate::types::tuple::{TupleSpecBuilder, TupleType};
use crate::types::visitor::any_over_type;
use crate::types::{
    CallableType, DynamicType, IntersectionBuilder, KnownClass, KnownInstanceType,
    LintDiagnosticGuard, Parameter, Parameters, SpecialFormType, SubclassOfType, Type,
    TypeAliasType, TypeContext, TypeIsType, UnionBuilder, UnionType, todo_type,
};

/// Type expressions
impl<'db> TypeInferenceBuilder<'db, '_> {
    /// Infer the type of a type expression.
    pub(super) fn infer_type_expression(&mut self, expression: &ast::Expr) -> Type<'db> {
        // `DeferredExpressionState::InStringAnnotation` takes precedence over other states.
        // However, if it's not a stringified annotation, we must still ensure that annotation expressions
        // are always deferred in stub files.
        match self.deferred_state {
            DeferredExpressionState::None => {
                if self.in_stub() {
                    self.deferred_state = DeferredExpressionState::Deferred;
                }
            }
            DeferredExpressionState::InStringAnnotation(_) | DeferredExpressionState::Deferred => {}
        }

        let mut ty = self.infer_type_expression_no_store(expression);
        let divergent = Type::divergent(Some(self.scope()));
        if ty.has_divergent_type(self.db(), divergent) {
            ty = divergent;
        }
        self.store_expression_type(expression, ty);
        ty
    }

    /// Similar to [`infer_type_expression`], but accepts a [`DeferredExpressionState`].
    ///
    /// [`infer_type_expression`]: TypeInferenceBuilder::infer_type_expression
    fn infer_type_expression_with_state(
        &mut self,
        expression: &ast::Expr,
        deferred_state: DeferredExpressionState,
    ) -> Type<'db> {
        let previous_deferred_state = std::mem::replace(&mut self.deferred_state, deferred_state);
        let annotation_ty = self.infer_type_expression(expression);
        self.deferred_state = previous_deferred_state;
        annotation_ty
    }

    /// Similar to [`infer_type_expression`], but accepts an optional expression.