Proposal: Inferring the fields of a dataclass type from a statically-typed __init__ signature.

[rsokl]

I would love for type checkers to have better support dynamically-generated dataclass types. Specifically, I would like them to be able to infer the fields on a dynamically-generated dataclass type based on a statically-typed __init__ signature. E.g.

from typing import Any, ClassVar, ParamSpec, Protocol, Callable

P = ParamSpec("P")

# This doesn't work, but this roughly what I want
class DataClass(Protocol[P]):
    __dataclass_fields__: ClassVar[dict[str, Any]]
    def __init__(self, *a: P.args, **k:P.kwargs) -> None:
        ...

# some function that dynamically generates a dataclass type
def make_dataclass_from_func_sig(fn: Callable[P, Any]) -> type[DataClass[P]]:  ...

def foo_fn(x: int, y: str = "a"):  ...

Foo = make_dataclass_from_func_sig(foo_fn)  # type[DataClass[(x: int, y: str = "a")]] 

foo = Foo(x=1)  # Foo's attributes inferred from ParamSpec
foo.x  # type checker should see `foo.x` as int
foo.y  # type checker should see `foo.y` as str
foo.z  # type checker should report Member "z" is unknown

PEP 681 enables you to statically declare fields on a class to inform the __init__ of that class in the eyes of type checkers. What I am proposing is enabling this to go in the opposite direction, so to speak.

Proposal

A PEP that specifies, for type checkers, a dataclass-type protocol decorator

from typing import dataclass_protocol

@dataclass_protocol  # this is just a marker for type checkers
class DataClass(Protocol[P]):
    ...

In the context of a @dataclass_protocol-decorated protocol:

• The type's __init__ and attributes are inferred from ParamSpec
• Additional fields can be declared statically here to indicate fields that are known to "live" on this type.
  These would represent no-init fields on the returned type and would have to be non-overlapping with ParamSpec.

@dataclass_protocol would not have any runtime details, it is meant to serve as a marker for type checkers to infer attributes from ParamSpec. Otherwise the behavior of ParamSpec remains unchanged for non-decorated protocols.

Limitations

In this proposal, P is only valid in this context if there are no *args, **kw, or positional-only entries in the function's signature - these are not supported by dataclass types. Furthermore, declaring InitVar (i.e. init-only) entries in the signature would not be supported. (There may be other dataclass features that I am overlooking here.)

Motivation

Many experiment / configuration framework libraries (including hydra-zen) have gravitated towards using dataclass types to represent "configs" for functions. The following pattern is quite common:

from dataclasses import dataclass

@dataclass
class Config:
    x: int
    y: str


def func(config: Config):
    # this might be a very expensive function
    ...

Config is acting like a floating, mutable signature for func. Its defaults can be modified/overwritten via inheritance, and you can create and store many instances of Config as a way to call func in various ways later. These are the upsides to this approach. The downsides, however, are numerous - such as decoupling Config from the implementation func that you care about, and no longer being able to call func in an ergonomic way. Soon, your code base is littered with this pattern instead of having simple functions.

This pattern can be seen throughout nerfstudio and many many ML code bases.

Far preferable would be to leave func as a normal function and then generate the Config type dynamically, based off of its signature.

def func(x: int, y: str):
    ...

Config = make_dataclass_from_func_sig(func)

then it is simple enough to write a wrapper that lets you do call_via_config(foo)(Config(x=1, y="b")) to call the function using your config instance. Now, the simple function foo is the star player in your code base once again and your desire to create a type that configure it can be done automatically and as an afterthought.

Indeed this is the pattern that hydra-zen enables. But a major downside to this dynamic approach is the lack of type-checking/autocomplete support for the attributes on the generated type, which is where the present proposal comes in.

I have to imagine that libraries like pydantic and other heavy users of PEP 681 would have use for this feature as well for purposes well beyond that of code configuration.

Closing Remarks

The introduction of @typing.dataclass_protocol is meant to enable basic type check support for the attributes of dynamically generated dataclass types. It is not meant to provide the same level of type information / specification as a statically-defined dataclass.

I am keen to get feedback on this. Thank you for your attention! 😄

[erictraut]

@rsokl, thanks for posting here.

I've already commented on an earlier iteration of this proposal in a pyright discussion thread, so I won't repeat my full response here. In summary, the specifics of this proposal strike me as very problematic. The solution doesn't compose well with other features in the type system.

My recommendation is to try to reframe the problem or look for an alternative solution that doesn't involve the use of a ParamSpec.