Skip to content

Commit fcdffe4

Browse files
authored
[ty] Pass down specialization to generic dataclass bases (#19472)
## Summary closes astral-sh/ty#853 ## Test Plan Regression test
1 parent 88de572 commit fcdffe4

File tree

2 files changed

+47
-6
lines changed

2 files changed

+47
-6
lines changed

crates/ty_python_semantic/resources/mdtest/dataclasses/dataclasses.md

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -640,6 +640,8 @@ reveal_type(C.__init__) # revealed: (self: C, normal: int, conditionally_presen
640640
python-version = "3.12"
641641
```
642642

643+
### Basic
644+
643645
```py
644646
from dataclasses import dataclass
645647

@@ -658,6 +660,34 @@ reveal_type(d_int.description) # revealed: str
658660
DataWithDescription[int](None, "description")
659661
```
660662

663+
### Deriving from generic dataclasses
664+
665+
This is a regression test for <https://github.com/astral-sh/ty/issues/853>.
666+
667+
```py
668+
from dataclasses import dataclass
669+
670+
@dataclass
671+
class Wrap[T]:
672+
data: T
673+
674+
reveal_type(Wrap[int].__init__) # revealed: (self: Wrap[int], data: int) -> None
675+
676+
@dataclass
677+
class WrappedInt(Wrap[int]):
678+
other_field: str
679+
680+
reveal_type(WrappedInt.__init__) # revealed: (self: WrappedInt, data: int, other_field: str) -> None
681+
682+
# Make sure that another generic type parameter does not affect the `data` field
683+
@dataclass
684+
class WrappedIntAndExtraData[T](Wrap[int]):
685+
extra_data: T
686+
687+
# revealed: (self: WrappedIntAndExtraData[bytes], data: int, extra_data: bytes) -> None
688+
reveal_type(WrappedIntAndExtraData[bytes].__init__)
689+
```
690+
661691
## Descriptor-typed fields
662692

663693
### Same type in `__get__` and `__set__`

crates/ty_python_semantic/src/types/class.rs

Lines changed: 17 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1668,16 +1668,16 @@ impl<'db> ClassLiteral<'db> {
16681668
if field_policy == CodeGeneratorKind::NamedTuple {
16691669
// NamedTuples do not allow multiple inheritance, so it is sufficient to enumerate the
16701670
// fields of this class only.
1671-
return self.own_fields(db);
1671+
return self.own_fields(db, specialization);
16721672
}
16731673

16741674
let matching_classes_in_mro: Vec<_> = self
16751675
.iter_mro(db, specialization)
16761676
.filter_map(|superclass| {
16771677
if let Some(class) = superclass.into_class() {
1678-
let class_literal = class.class_literal(db).0;
1678+
let (class_literal, specialization) = class.class_literal(db);
16791679
if field_policy.matches(db, class_literal) {
1680-
Some(class_literal)
1680+
Some((class_literal, specialization))
16811681
} else {
16821682
None
16831683
}
@@ -1691,7 +1691,7 @@ impl<'db> ClassLiteral<'db> {
16911691
matching_classes_in_mro
16921692
.into_iter()
16931693
.rev()
1694-
.flat_map(|class| class.own_fields(db))
1694+
.flat_map(|(class, specialization)| class.own_fields(db, specialization))
16951695
// We collect into a FxOrderMap here to deduplicate attributes
16961696
.collect()
16971697
}
@@ -1707,7 +1707,11 @@ impl<'db> ClassLiteral<'db> {
17071707
/// y: str = "a"
17081708
/// ```
17091709
/// we return a map `{"x": (int, None), "y": (str, Some(Literal["a"]))}`.
1710-
fn own_fields(self, db: &'db dyn Db) -> FxOrderMap<Name, (Type<'db>, Option<Type<'db>>)> {
1710+
fn own_fields(
1711+
self,
1712+
db: &'db dyn Db,
1713+
specialization: Option<Specialization<'db>>,
1714+
) -> FxOrderMap<Name, (Type<'db>, Option<Type<'db>>)> {
17111715
let mut attributes = FxOrderMap::default();
17121716

17131717
let class_body_scope = self.body_scope(db);
@@ -1747,7 +1751,14 @@ impl<'db> ClassLiteral<'db> {
17471751
let bindings = use_def.end_of_scope_bindings(place_id);
17481752
let default_ty = place_from_bindings(db, bindings).ignore_possibly_unbound();
17491753

1750-
attributes.insert(place_expr.expect_name().clone(), (attr_ty, default_ty));
1754+
attributes.insert(
1755+
place_expr.expect_name().clone(),
1756+
(
1757+
attr_ty.apply_optional_specialization(db, specialization),
1758+
default_ty
1759+
.map(|ty| ty.apply_optional_specialization(db, specialization)),
1760+
),
1761+
);
17511762
}
17521763
}
17531764
}

0 commit comments

Comments
 (0)