How to annotate a function that returns a modified dataclass type

[jrheard-seesaw]

Hi,

I've encountered some code that looks like this:

def dataclass_with_additional_method(cls):
    result_cls = dataclass(cls)  # apply the dataclass decorator

    def from_dict(cls, some_dict):
        return cls(**some_dict)

    result_cls.from_dict = classmethod(from_dict)

    return result_cls

I'm glossing over some details here, but that's the overall shape of this function. It's used as a decorator for a number of classes, like this:

@dataclass_with_additional_method
class MyClass:
    foo: int

I'm having a lot of trouble figuring out the right way to annotate this dataclass_with_additional_method function. I've been playing around with TypeVars and _typeshed.DataclassInstance, but haven't quite found the right thing to do.

I think that the desired behavior would be for this function to take type[T] and return type[T & DataclassInstance], but that doesn't seem possible in Python.

Just about every type signature I've tried causes one or more of these problems:

• I can't do MyClass(foo=5) because Pyright can't see that MyClass takes a foo field as input
• I can't pass an instance of MyClass to dataclasses.asdict()
• Fields of MyClass instances are inferred to be of type Any

Is there anything I can do to the signature of dataclass_with_additional_method that would solve all of these problems? Thanks!

[jrheard-seesaw]

I'm thinking about nuking this decorator, updating all the callsites to use @dataclass instead, and writing a function that looks like this, which would be used like from_dict(MyClass, {"foo": 5}) rather than MyClass.from_dict({"foo": 5}).

if TYPE_CHECKING:
    from _typeshed import DataclassInstance

T = TypeVar("T", bound="DataclassInstance")

def from_dict(cls: type[T], data: dict[str, Any]) -> T:
    return cls(**dict)

(again, I'm handwaving over some details, mainly focusing on the overall types/shapes.)

Does that seem like a better way forward?

[erictraut]

How about a mix-in class?

from dataclasses import dataclass
from typing import Any, Self

class MyMethodMixin:
    @classmethod
    def my_method(cls, some_dict: dict[str, Any]) -> Self:
        return cls(**some_dict)

@dataclass
class MyClass(MyMethodMixin):
    foo: int

Or you could use __init_subclass__ together with dataclass_transform:

from dataclasses import dataclass, field
from typing import Any, Self, dataclass_transform

@dataclass_transform(field_specifiers=(field,))
class MyMethodBase:
    @classmethod
    def my_method(cls, some_dict: dict[str, Any]) -> Self:
        return cls(**some_dict)

    def __init_subclass__(cls) -> None:
        dataclass(cls)

class MyClass(MyMethodBase):
    foo: int

[jrheard-seesaw]

Those are great suggestions, I'll poke around and choose one to go with. Thank you!!