Skip to content

Commit 71d8a5d

Browse files
[ty] dataclasses: Allow using dataclasses.dataclass as a function. (#18440)
## Summary Part of astral-sh/ty#111 Using `dataclass` as a function, instead of as a decorator did not work as expected prior to this. Fix that by modifying the dataclass overload's return type. ## Test Plan New mdtests, fixing the existing TODO.
1 parent 2c3b3d3 commit 71d8a5d

File tree

2 files changed

+31
-3
lines changed

2 files changed

+31
-3
lines changed

crates/ty_python_semantic/resources/mdtest/dataclasses.md

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -797,7 +797,20 @@ C(1) < C(2) # ok
797797

798798
### Using `dataclass` as a function
799799

800-
To do
800+
```py
801+
from dataclasses import dataclass
802+
803+
class B:
804+
x: int
805+
806+
# error: [missing-argument]
807+
dataclass(B)()
808+
809+
# error: [invalid-argument-type]
810+
dataclass(B)("a")
811+
812+
reveal_type(dataclass(B)(3).x) # revealed: int
813+
```
801814

802815
## Internals
803816

crates/ty_python_semantic/src/types/call/bind.rs

Lines changed: 17 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -22,8 +22,8 @@ use crate::types::function::{DataclassTransformerParams, FunctionDecorators, Kno
2222
use crate::types::generics::{Specialization, SpecializationBuilder, SpecializationError};
2323
use crate::types::signatures::{Parameter, ParameterForm};
2424
use crate::types::{
25-
BoundMethodType, DataclassParams, KnownClass, KnownInstanceType, MethodWrapperKind,
26-
PropertyInstanceType, SpecialFormType, TupleType, TypeMapping, UnionType,
25+
BoundMethodType, ClassLiteral, DataclassParams, KnownClass, KnownInstanceType,
26+
MethodWrapperKind, PropertyInstanceType, SpecialFormType, TupleType, TypeMapping, UnionType,
2727
WrapperDescriptorKind, ide_support, todo_type,
2828
};
2929
use ruff_db::diagnostic::{Annotation, Diagnostic, Severity, SubDiagnostic};
@@ -839,6 +839,21 @@ impl<'db> Bindings<'db> {
839839

840840
overload.set_return_type(Type::DataclassDecorator(params));
841841
}
842+
843+
// `dataclass` being used as a non-decorator
844+
if let [Some(Type::ClassLiteral(class_literal))] =
845+
overload.parameter_types()
846+
{
847+
let params = DataclassParams::default();
848+
overload.set_return_type(Type::from(ClassLiteral::new(
849+
db,
850+
class_literal.name(db),
851+
class_literal.body_scope(db),
852+
class_literal.known(db),
853+
Some(params),
854+
class_literal.dataclass_transformer_params(db),
855+
)));
856+
}
842857
}
843858

844859
Some(KnownFunction::DataclassTransform) => {

0 commit comments

Comments
 (0)