[ty] Implement DataClassInstance protocol for dataclasses. (#18018)

Fixes: https://github.com/astral-sh/ty/issues/92

## Summary

We currently get a `invalid-argument-type` error when using
`dataclass.fields` on a dataclass, because we do not synthesize the
`__dataclass_fields__` member.

This PR fixes this diagnostic.

Note that we do not yet model the `Field` type correctly. After that is
done, we can assign a more precise `tuple[Field, ...]` type to this new
member.

## Test Plan
New mdtest.

---------

Co-authored-by: David Peter <mail@david-peter.de>
This commit is contained in:
Abhijeet Prasad Bodas 2025-05-13 14:01:26 +05:30 committed by GitHub
parent 0ae07cdd1f
commit 68b0386007
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 47 additions and 5 deletions

View file

@ -616,6 +616,25 @@ reveal_type(C.__init__) # revealed: (field: str | int = int) -> None
To do
## `dataclass.fields`
Dataclasses have `__dataclass_fields__` in them, which makes them a subtype of the
`DataclassInstance` protocol.
Here, we verify that dataclasses can be passed to `dataclasses.fields` without any errors, and that
the return type of `dataclasses.fields` is correct.
```py
from dataclasses import dataclass, fields
@dataclass
class Foo:
x: int
reveal_type(Foo.__dataclass_fields__) # revealed: dict[str, Field[Any]]
reveal_type(fields(Foo)) # revealed: tuple[Field[Any], ...]
```
## Other special cases
### `dataclasses.dataclass`

View file

@ -2939,6 +2939,19 @@ impl<'db> Type<'db> {
))
.into()
}
Type::ClassLiteral(class)
if name == "__dataclass_fields__" && class.dataclass_params(db).is_some() =>
{
// Make this class look like a subclass of the `DataClassInstance` protocol
Symbol::bound(KnownClass::Dict.to_specialized_instance(
db,
[
KnownClass::Str.to_instance(db),
KnownClass::Field.to_specialized_instance(db, [Type::any()]),
],
))
.with_qualifiers(TypeQualifiers::CLASS_VAR)
}
Type::BoundMethod(bound_method) => match name_str {
"__self__" => Symbol::bound(bound_method.self_instance(db)).into(),
"__func__" => {

View file

@ -1958,6 +1958,8 @@ pub enum KnownClass {
// backported as `builtins.ellipsis` by typeshed on Python <=3.9
EllipsisType,
NotImplementedType,
// dataclasses
Field,
}
impl<'db> KnownClass {
@ -2037,7 +2039,8 @@ impl<'db> KnownClass {
// and raises a `TypeError` in Python >=3.14
// (see https://docs.python.org/3/library/constants.html#NotImplemented)
| Self::NotImplementedType
| Self::Classmethod => Truthiness::Ambiguous,
| Self::Classmethod
| Self::Field => Truthiness::Ambiguous,
}
}
@ -2108,7 +2111,8 @@ impl<'db> KnownClass {
| Self::VersionInfo
| Self::EllipsisType
| Self::NotImplementedType
| Self::UnionType => false,
| Self::UnionType
| Self::Field => false,
}
}
@ -2181,6 +2185,7 @@ impl<'db> KnownClass {
}
}
Self::NotImplementedType => "_NotImplementedType",
Self::Field => "Field",
}
}
@ -2405,6 +2410,7 @@ impl<'db> KnownClass {
| Self::DefaultDict
| Self::Deque
| Self::OrderedDict => KnownModule::Collections,
Self::Field => KnownModule::Dataclasses,
}
}
@ -2464,7 +2470,8 @@ impl<'db> KnownClass {
| Self::ABCMeta
| Self::Super
| Self::NamedTuple
| Self::NewType => false,
| Self::NewType
| Self::Field => false,
}
}
@ -2526,7 +2533,8 @@ impl<'db> KnownClass {
| Self::Super
| Self::UnionType
| Self::NamedTuple
| Self::NewType => false,
| Self::NewType
| Self::Field => false,
}
}
@ -2596,6 +2604,7 @@ impl<'db> KnownClass {
Self::EllipsisType
}
"_NotImplementedType" => Self::NotImplementedType,
"Field" => Self::Field,
_ => return None,
};
@ -2647,7 +2656,8 @@ impl<'db> KnownClass {
| Self::UnionType
| Self::GeneratorType
| Self::AsyncGeneratorType
| Self::WrapperDescriptorType => module == self.canonical_module(db),
| Self::WrapperDescriptorType
| Self::Field => module == self.canonical_module(db),
Self::NoneType => matches!(module, KnownModule::Typeshed | KnownModule::Types),
Self::SpecialForm
| Self::TypeVar