From 71d8a5da2aef888e96e575f2a83c0b992e8ff36c Mon Sep 17 00:00:00 2001
From: Abhijeet Prasad Bodas
 <55339528+abhijeetbodas2001@users.noreply.github.com>
Date: Tue, 3 Jun 2025 22:20:29 +0530
Subject: [PATCH] [ty] dataclasses: Allow using dataclasses.dataclass as a
 function. (#18440)

## Summary

Part of https://github.com/astral-sh/ty/issues/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.
---
 .../resources/mdtest/dataclasses.md           | 15 ++++++++++++++-
 .../ty_python_semantic/src/types/call/bind.rs | 19 +++++++++++++++++--
 2 files changed, 31 insertions(+), 3 deletions(-)

diff --git a/crates/ty_python_semantic/resources/mdtest/dataclasses.md b/crates/ty_python_semantic/resources/mdtest/dataclasses.md
index a74c125b5d..4616e59e22 100644
--- a/crates/ty_python_semantic/resources/mdtest/dataclasses.md
+++ b/crates/ty_python_semantic/resources/mdtest/dataclasses.md
@@ -797,7 +797,20 @@ C(1) < C(2)  # ok
 
 ### Using `dataclass` as a function
 
-To do
+```py
+from dataclasses import dataclass
+
+class B:
+    x: int
+
+# error: [missing-argument]
+dataclass(B)()
+
+# error: [invalid-argument-type]
+dataclass(B)("a")
+
+reveal_type(dataclass(B)(3).x)  # revealed: int
+```
 
 ## Internals
 
diff --git a/crates/ty_python_semantic/src/types/call/bind.rs b/crates/ty_python_semantic/src/types/call/bind.rs
index b49343c12d..2eab02c9c7 100644
--- a/crates/ty_python_semantic/src/types/call/bind.rs
+++ b/crates/ty_python_semantic/src/types/call/bind.rs
@@ -22,8 +22,8 @@ use crate::types::function::{DataclassTransformerParams, FunctionDecorators, Kno
 use crate::types::generics::{Specialization, SpecializationBuilder, SpecializationError};
 use crate::types::signatures::{Parameter, ParameterForm};
 use crate::types::{
-    BoundMethodType, DataclassParams, KnownClass, KnownInstanceType, MethodWrapperKind,
-    PropertyInstanceType, SpecialFormType, TupleType, TypeMapping, UnionType,
+    BoundMethodType, ClassLiteral, DataclassParams, KnownClass, KnownInstanceType,
+    MethodWrapperKind, PropertyInstanceType, SpecialFormType, TupleType, TypeMapping, UnionType,
     WrapperDescriptorKind, ide_support, todo_type,
 };
 use ruff_db::diagnostic::{Annotation, Diagnostic, Severity, SubDiagnostic};
@@ -839,6 +839,21 @@ impl<'db> Bindings<'db> {
 
                                 overload.set_return_type(Type::DataclassDecorator(params));
                             }
+
+                            // `dataclass` being used as a non-decorator
+                            if let [Some(Type::ClassLiteral(class_literal))] =
+                                overload.parameter_types()
+                            {
+                                let params = DataclassParams::default();
+                                overload.set_return_type(Type::from(ClassLiteral::new(
+                                    db,
+                                    class_literal.name(db),
+                                    class_literal.body_scope(db),
+                                    class_literal.known(db),
+                                    Some(params),
+                                    class_literal.dataclass_transformer_params(db),
+                                )));
+                            }
                         }
 
                         Some(KnownFunction::DataclassTransform) => {