Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Support __getitem__ type inference for subscripts #13579

Merged
merged 2 commits into from
Oct 1, 2024
Merged

Support __getitem__ type inference for subscripts #13579

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
0ef8def
Support __getitem__ type inference for subscripts
charliermarsh Oct 1, 2024
11acc07
Add tests; support unions
charliermarsh Oct 1, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 10 additions & 0 deletions crates/red_knot_python_semantic/src/types.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -401,6 +401,16 @@ impl<'db> Type<'db> {
}
}

/// Return true if the type is a class or a union of classes.
pub fn is_class(&self, db: &'db dyn Db) -> bool {
match self {
Type::Union(union) => union.elements(db).iter().all(|ty| ty.is_class(db)),
Type::Class(_) => true,
// / TODO include type[X], once we add that type
_ => false,
charliermarsh marked this conversation as resolved.
Show resolved Hide resolved
}
}
Comment on lines +404 to +412
Copy link
Member

@AlexWaygood AlexWaygood Oct 1, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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?

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Member

@AlexWaygood AlexWaygood Oct 1, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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


/// Return true if this type is a [subtype of] type `target`.
///
/// [subtype of]: https://typing.readthedocs.io/en/latest/spec/concepts.html#subtype-supertype-and-type-equivalence
Expand Down
301 changes: 300 additions & 1 deletion crates/red_knot_python_semantic/src/types/infer.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1322,6 +1322,22 @@ impl<'db> TypeInferenceBuilder<'db> {
);
}

/// Emit a diagnostic declaring that a type does not support subscripting.
pub(super) fn non_subscriptable_diagnostic(
&mut self,
node: AnyNodeRef,
non_subscriptable_ty: Type<'db>,
) {
self.add_diagnostic(
node,
"non-subscriptable",
format_args!(
"Cannot subscript object of type '{}' with no `__getitem__` method.",
non_subscriptable_ty.display(self.db)
),
);
}

fn infer_for_statement_definition(
&mut self,
target: &ast::ExprName,
Expand Down Expand Up @@ -2588,7 +2604,35 @@ impl<'db> TypeInferenceBuilder<'db> {
Type::Unknown
})
}
_ => Type::Todo,
(value_ty, slice_ty) => {
// Resolve the value to its class.
let value_meta_ty = value_ty.to_meta_type(self.db);

// If the class defines `__getitem__`, return its return type.
//
// See: https://docs.python.org/3/reference/datamodel.html#class-getitem-versus-getitem
charliermarsh marked this conversation as resolved.
Show resolved Hide resolved
let dunder_getitem_method = value_meta_ty.member(self.db, "__getitem__");
if !dunder_getitem_method.is_unbound() {
return dunder_getitem_method
.call(self.db, &[slice_ty])
.unwrap_with_diagnostic(self.db, value.as_ref().into(), self);
}

// Otherwise, if the value is itself a class and defines `__class_getitem__`,
// return its return type.
if value_ty.is_class(self.db) {
let dunder_class_getitem_method = value_ty.member(self.db, "__class_getitem__");
if !dunder_class_getitem_method.is_unbound() {
return dunder_class_getitem_method
.call(self.db, &[slice_ty])
.unwrap_with_diagnostic(self.db, value.as_ref().into(), self);
}
}

// Otherwise, emit a diagnostic.
self.non_subscriptable_diagnostic((&**value).into(), value_ty);
Type::Unknown
}
}
}

Expand Down Expand Up @@ -6723,6 +6767,261 @@ mod tests {
Ok(())
}

#[test]
fn subscript_getitem_unbound() -> anyhow::Result<()> {
let mut db = setup_db();

db.write_dedented(
"/src/a.py",
"
class NotSubscriptable:
pass

a = NotSubscriptable()[0]
",
)?;

assert_public_ty(&db, "/src/a.py", "a", "Unknown");
assert_file_diagnostics(
&db,
"/src/a.py",
&["Cannot subscript object of type 'NotSubscriptable' with no `__getitem__` method."],
);

Ok(())
}

#[test]
fn subscript_not_callable_getitem() -> anyhow::Result<()> {
let mut db = setup_db();

db.write_dedented(
"/src/a.py",
"
class NotSubscriptable:
__getitem__ = None

a = NotSubscriptable()[0]
",
)?;

assert_public_ty(&db, "/src/a.py", "a", "Unknown");
assert_file_diagnostics(
&db,
"/src/a.py",
&["Object of type 'None' is not callable."],
);

Ok(())
}

#[test]
fn subscript_str_literal() -> anyhow::Result<()> {
let mut db = setup_db();

db.write_dedented(
"/src/a.py",
"
def add(x: int, y: int) -> int:
return x + y

a = 'abcde'[add(0, 1)]
",
)?;

assert_public_ty(&db, "/src/a.py", "a", "str");

Ok(())
}

#[test]
fn subscript_getitem() -> anyhow::Result<()> {
let mut db = setup_db();

db.write_dedented(
"/src/a.py",
"
class Identity:
def __getitem__(self, index: int) -> int:
return index

a = Identity()[0]
",
)?;

assert_public_ty(&db, "/src/a.py", "a", "int");

Ok(())
}

#[test]
fn subscript_class_getitem() -> anyhow::Result<()> {
let mut db = setup_db();

db.write_dedented(
"/src/a.py",
"
class Identity:
def __class_getitem__(cls, item: int) -> str:
return item

a = Identity[0]
",
)?;

assert_public_ty(&db, "/src/a.py", "a", "str");

Ok(())
}

#[test]
fn subscript_getitem_union() -> anyhow::Result<()> {
let mut db = setup_db();

db.write_dedented(
"/src/a.py",
"
flag = True

class Identity:
if flag:
def __getitem__(self, index: int) -> int:
return index
else:
def __getitem__(self, index: int) -> str:
return str(index)

a = Identity()[0]
",
)?;

assert_public_ty(&db, "/src/a.py", "a", "int | str");

Ok(())
}

#[test]
fn subscript_class_getitem_union() -> anyhow::Result<()> {
let mut db = setup_db();

db.write_dedented(
"/src/a.py",
"
flag = True

class Identity:
if flag:
def __class_getitem__(cls, item: int) -> str:
return item
else:
def __class_getitem__(cls, item: int) -> int:
return item

a = Identity[0]
",
)?;

assert_public_ty(&db, "/src/a.py", "a", "str | int");

Ok(())
}

#[test]
fn subscript_class_getitem_class_union() -> anyhow::Result<()> {
let mut db = setup_db();

db.write_dedented(
"/src/a.py",
"
flag = True

class Identity1:
def __class_getitem__(cls, item: int) -> str:
return item

class Identity2:
def __class_getitem__(cls, item: int) -> int:
return item

if flag:
a = Identity1
else:
a = Identity2

b = a[0]
",
)?;

assert_public_ty(&db, "/src/a.py", "a", "Literal[Identity1, Identity2]");
assert_public_ty(&db, "/src/a.py", "b", "str | int");

Ok(())
}

#[test]
fn subscript_class_getitem_unbound_method_union() -> anyhow::Result<()> {
let mut db = setup_db();

db.write_dedented(
"/src/a.py",
"
flag = True

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

a = Identity[42]
",
)?;

assert_public_ty(&db, "/src/a.py", "a", "str | Unknown");

assert_file_diagnostics(
&db,
"/src/a.py",
&["Object of type 'Literal[__class_getitem__] | Unbound' is not callable (due to union element 'Unbound')."],
);

Ok(())
}

#[test]
fn subscript_class_getitem_non_class_union() -> anyhow::Result<()> {
let mut db = setup_db();

db.write_dedented(
"/src/a.py",
"
flag = True

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

a = Identity[42]
",
)?;

// TODO this should _probably_ emit `str | Unknown` instead of `Unknown`.
assert_public_ty(&db, "/src/a.py", "a", "Unknown");

assert_file_diagnostics(
&db,
"/src/a.py",
&["Cannot subscript object of type 'Literal[Identity] | Literal[1]' with no `__getitem__` method."],
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

#13596

);

Ok(())
}

#[test]
fn dunder_call() -> anyhow::Result<()> {
let mut db = setup_db();
Expand Down
Loading