Support __getitem__ type inference for subscripts

[charliermarsh]

Summary

Follow-up to #13562, to add support for "arbitrary" subscript operations.

[github-actions]

ruff-ecosystem results

Linter (stable)

✅ ecosystem check detected no linter changes.

Linter (preview)

✅ ecosystem check detected no linter changes.

[charliermarsh]

We're getting new errors in the benchmark due to expressions like frozenset[str] or dict[str]:

Dunder method `__class_getitem__` is not callable on type 'Literal[frozenset]'

[charliermarsh]

Those seem to have a type that's a union of Unknown and the __class_getitem__ function definition.

[charliermarsh]

Okay, I think (?) that's because in typeshed they're defined as:

if sys.version_info >= (3, 9):
    def __class_getitem__(cls, item: Any, /) -> GenericAlias: ...

[AlexWaygood]

Those seem to have a type that's a union of Unknown and the __class_getitem__ function definition.

Hmm, that's kinda odd that it's creating issues, though. Both Unknown and the __class_getitem__ function should be treated as callable, because Unknown should have all possible attributes

[AlexWaygood]

Those seem to have a type that's a union of Unknown and the __class_getitem__ function definition.

Hmm, that's kinda odd that it's creating issues, though. Both Unknown and the __class_getitem__ function should be treated as callable, because Unknown should have all possible attributes

Okay @charliermarsh I think this is because you're doing this:

                    let CallOutcome::Callable { return_ty } =
                        dunder_getitem_method.call(self.db, &[slice_ty])
                    else

But the result of Union[Unknown, <some function>].call() is not CallOutcome::Callable (even though all items in the union are actually callable), it's CallOutcome::Union

[AlexWaygood]

I think you can fix this by making use of one of these two helper methods we have on CallOutcome, both of which abstract away the difference between CallOutcome::Callable and CallOutcome::Union:

ruff/crates/red_knot_python_semantic/src/types.rs

Lines 851 to 971 in cfd5d63

	/// Get the return type of the call, or `None` if not callable.
	fn return_ty(&self, db: &'db dyn Db) -> Option<Type<'db>> {
	match self {
	Self::Callable { return_ty } => Some(*return_ty),
	Self::RevealType {
	return_ty,
	revealed_ty: _,
	} => Some(*return_ty),
	Self::NotCallable { not_callable_ty: _ } => None,
	Self::Union {
	outcomes,
	called_ty: _,
	} => outcomes
	.iter()
	// If all outcomes are NotCallable, we return None; if some outcomes are callable
	// and some are not, we return a union including Unknown.
	.fold(None, |acc, outcome| {
	let ty = outcome.return_ty(db);
	match (acc, ty) {
	(None, None) => None,
	(None, Some(ty)) => Some(UnionBuilder::new(db).add(ty)),
	(Some(builder), ty) => Some(builder.add(ty.unwrap_or(Type::Unknown))),
	}
	})
	.map(UnionBuilder::build),
	}
	}
	/// Get the return type of the call, emitting diagnostics if needed.
	fn unwrap_with_diagnostic<'a>(
	&self,
	db: &'db dyn Db,
	node: ast::AnyNodeRef,
	builder: &'a mut TypeInferenceBuilder<'db>,
	) -> Type<'db> {
	match self {
	Self::Callable { return_ty } => *return_ty,
	Self::RevealType {
	return_ty,
	revealed_ty,
	} => {
	builder.add_diagnostic(
	node,
	"revealed-type",
	format_args!("Revealed type is '{}'.", revealed_ty.display(db)),
	);
	*return_ty
	}
	Self::NotCallable { not_callable_ty } => {
	builder.add_diagnostic(
	node,
	"call-non-callable",
	format_args!(
	"Object of type '{}' is not callable.",
	not_callable_ty.display(db)
	),
	);
	Type::Unknown
	}
	Self::Union {
	outcomes,
	called_ty,
	} => {
	let mut not_callable = vec![];
	let mut union_builder = UnionBuilder::new(db);
	let mut revealed = false;
	for outcome in &**outcomes {
	let return_ty = match outcome {
	Self::NotCallable { not_callable_ty } => {
	not_callable.push(*not_callable_ty);
	Type::Unknown
	}
	Self::RevealType {
	return_ty,
	revealed_ty: _,
	} => {
	if revealed {
	*return_ty
	} else {
	revealed = true;
	outcome.unwrap_with_diagnostic(db, node, builder)
	}
	}
	_ => outcome.unwrap_with_diagnostic(db, node, builder),
	};
	union_builder = union_builder.add(return_ty);
	}
	match not_callable[..] {
	[] => {}
	[elem] => builder.add_diagnostic(
	node,
	"call-non-callable",
	format_args!(
	"Object of type '{}' is not callable (due to union element '{}').",
	called_ty.display(db),
	elem.display(db),
	),
	),
	_ if not_callable.len() == outcomes.len() => builder.add_diagnostic(
	node,
	"call-non-callable",
	format_args!(
	"Object of type '{}' is not callable.",
	called_ty.display(db)
	),
	),
	_ => builder.add_diagnostic(
	node,
	"call-non-callable",
	format_args!(
	"Object of type '{}' is not callable (due to union elements {}).",
	called_ty.display(db),
	not_callable.display(db),
	),
	),
	}
	union_builder.build()
	}
	}
	}
	}

[charliermarsh]

@carljm -- I believe all the feedback is resolved with the exception that the dunder diagnostic regressed due to the use of unwrap_with_diagnostic (as you mentioned). I may look at that in a separate PR since it doesn't necessarily need to be coupled to these changes.

[carljm]

Looks great!

[AlexWaygood]

I'm not sure we'll do the right thing with this for something like

flag = True

if flag:
    class Foo:
        def __class_getitem__(self, x: int) -> str:
            pass
else:
    class Foo:
        pass

Foo[42]

Here we want to:

1. Infer a union type for the Foo variable (it could be the first class definition, or the second)
2. Emit a diagnostic because not all items in the union support being subscripted
3. Infer str | Unknown for the result of the subscription

[charliermarsh]

I think that is working as described? I added a test.

[carljm]

Yeah, I was debating whether to bring this up earlier, or whether it's adequate to just say "not subscriptable" in this case. I agree the behavior you're describing is probably more ideal. It might require explicitly handling unions as a special case here, and calling the class-getitem resolution method per element of the union? I'd also be fine with a TODO for this so this PR can land.

[charliermarsh]

Are you instead thinking of something like this:

flag = True

if flag:
    class Foo:
        def __class_getitem__(self, x: int) -> str:
            pass
else:
    Foo = 1

Foo[42]

Where not all members of the union are classes?

[AlexWaygood]

Huh, your tests are pretty convincing. Thanks for adding them!

[charliermarsh]

I added a test for the case I described above, which does resolve to Unknown (with a TODO).

[AlexWaygood]

Ahh, I think I did spot a bug but didn't accurately identify where the bug would manifest? Anyway, thanks again, a TODO for now is fine by me

[AlexWaygood]

Thanks!

[carljm]

This case also highlights that perhaps our "not subscriptable" diagnostic should mention __class_getitem__ instead of __getitem__ if Type::is_class is true? But I don't think that's critical to fix now.

[charliermarsh]

#13596