google / mesop

Rapidly build AI apps in Python
https://google.github.io/mesop/
Apache License 2.0
5.65k stars 275 forks source link

Decorate `stateclass` with `@dataclass_transform()` #1113

Open kkom opened 1 week ago

kkom commented 1 week ago

Is your feature request related to a problem? Please describe.

I have pretty strict type checking enabled in my project. This results in hitting Pyright's reportUninitializedInstanceVariable rule when defining a Mesop stateclass:

@me.stateclass
class State:
    first_name: str

That's the error:

Instance variable "first_name" is not initialized in the class body or __init__ method reportUninitializedInstanceVariable

Describe the solution you'd like

I wonder if it's possible to annotate stateclass with dataclass_transform(), to make the type checker okay with things that normally look sketchy to it?

I.e. go from this:

https://github.com/google/mesop/blob/85fac407747073d9a3ae71a2e755304070fce275/mesop/__init__.py#L258-L262

To this:

if sys.version_info >= (3, 10):
    from typing import dataclass_transform
else:
    from typing_extensions import dataclass_transform

...

@dataclass_transform()
def stateclass(cls: type[_T] | None, **kw_args: Any) -> type[_T]:
  """
  Similar as dataclass, but it also registers with Mesop runtime().
  """

Describe alternatives you've considered

Can't think of any.

Additional context

Note – I'm vaguely aware of dataclass transforms. I am not 100% sure what other things need to be considered when adding this annotation, i.e. if the runtime code must then implement specific parts of the dataclass spec.

But I can say that making that change fixed the static errors for me.

wwwillchen commented 5 days ago

Looks like a reasonable change - feel free to open a PR