[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 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 ## Other special cases
### `dataclasses.dataclass` ### `dataclasses.dataclass`

View file

@ -2939,6 +2939,19 @@ impl<'db> Type<'db> {
)) ))
.into() .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 { Type::BoundMethod(bound_method) => match name_str {
"__self__" => Symbol::bound(bound_method.self_instance(db)).into(), "__self__" => Symbol::bound(bound_method.self_instance(db)).into(),
"__func__" => { "__func__" => {

View file

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