Repository: astral-sh/ruff
File: crates/ty_python_semantic/src/types/call/bind.rs (Rust)

Code before:
/// Infer the return type for a call to `asynccontextmanager`.
///
/// The `@asynccontextmanager` decorator transforms a function that returns (a subtype of) `AsyncIterator[T]`
/// into a function that returns `_AsyncGeneratorContextManager[T]`.
///
/// TODO: This function only handles the most basic case. It should be removed once we have
/// full support for generic protocols in the solver.
fn asynccontextmanager_return_type<'db>(db: &'db dyn Db, func_ty: Type<'db>) -> Option<Type<'db>> {
    let bindings = func_ty.bindings(db);
    let binding = bindings
        .single_element()?
        .overloads
        .iter()
        .exactly_one()
        .ok()?;
    let signature = &binding.signature;
    let return_ty = signature.return_ty?;

    let yield_ty = return_ty
        .try_iterate_with_mode(db, EvaluationMode::Async)
        .ok()?
        .homogeneous_element_type(db);

    if yield_ty.is_divergent() {
        return Some(yield_ty);
    }

    let context_manager =
        known_module_symbol(db, KnownModule::Contextlib, "_AsyncGeneratorContextManager")
            .place
            .ignore_possibly_undefined()?
            .as_class_literal()?;

    let context_manager = context_manager.apply_specialization(db, |generic_context| {
        generic_context.specialize_partial(db, [Some(yield_ty), None])
    });

    let new_return_ty = Type::from(context_manager).to_instance(db)?;
    let new_signature = Signature::new(signature.parameters().clone(), Some(new_return_ty));

    Some(Type::Callable(CallableType::new(
        db,
        CallableSignature::single(new_signature),
        CallableTypeKind::FunctionLike,
    )))
}


Review comment (suggestion):
This took me a while to figure out. Looks related to https://github.com/astral-sh/ty/issues/1821 / https://github.com/astral-sh/ty/issues/1821. If we don't do this, we end up with unions of the actual type that we want to have and `Divergent`

Code after:
/// Infer the return type for a call to `asynccontextmanager`.
///
/// The `@asynccontextmanager` decorator transforms a function that returns (a subtype of) `AsyncIterator[T]`
/// into a function that returns `_AsyncGeneratorContextManager[T]`.
///
/// TODO: This function only handles the most basic case. It should be removed once we have
/// full support for generic protocols in the solver.
fn asynccontextmanager_return_type<'db>(db: &'db dyn Db, func_ty: Type<'db>) -> Option<Type<'db>> {
    let bindings = func_ty.bindings(db);
    let binding = bindings
        .single_element()?
        .overloads
        .iter()
        .exactly_one()
        .ok()?;
    let signature = &binding.signature;
    let return_ty = signature.return_ty?;

    let yield_ty = return_ty
        .try_iterate_with_mode(db, EvaluationMode::Async)
        .ok()?
        .homogeneous_element_type(db);

    if yield_ty.is_divergent()
        || signature
            .parameters()
            .iter()
            .any(|param| param.annotated_type().is_some_and(|ty| ty.is_divergent()))
    {
        return Some(yield_ty);
    }

    let context_manager =
        known_module_symbol(db, KnownModule::Contextlib, "_AsyncGeneratorContextManager")
            .place
            .ignore_possibly_undefined()?
            .as_class_literal()?;

    let context_manager = context_manager.apply_specialization(db, |generic_context| {
        generic_context.specialize_partial(db, [Some(yield_ty), None])
    });

    let new_return_ty = Type::from(context_manager).to_instance(db)?;
    let new_signature = Signature::new(signature.parameters().clone(), Some(new_return_ty));

    Some(Type::Callable(CallableType::new(
        db,
        CallableSignature::single(new_signature),
        CallableTypeKind::FunctionLike,
    )))
}