Using Data Class Transforms to advertise dataclasses compatibility

[ikonst]

Exposé

In PEP-681, Data Class Transforms were conceived as "a way for third-party libraries to indicate that certain decorator functions, classes, and metaclasses provide behaviors similar to dataclasses". Let's call this 📖 the intended use. The PEP was successful in that, e.g. in pydantic, but over time another use emerged — let's call it 🎁 the bonus use - where @dataclass_transform is used for decorators which don't merely return something with similar to, but actual dataclasses. The motivation for the latter is to make custom decorators seen as equivalent to @dataclasses.dataclass by type-checkers and IDEs:

@dataclass_transform()
def my_dataclass(kls):
  # ... do something with kls ...
  return dataclass(kls)

Real-world examples¹:

• pyserde
• strawberry
• ibis
• flax

So, @dataclass_transform is useful in more than one way, and that's great. Meanwhile, type checkers are also improving their understanding of dataclasses module's utility functions:

• dataclasses.is_dataclass
  • typeshed: use as type-guard
• dataclasses.fields
  • typeshed: check obj arg
  • Potential:
    • type retval (tuple[Field[int], Field[str], ...])
• dataclasses.as_dict
  • typeshed: check obj arg
  • Potential:
    • type retval (TypedDict?)
• dataclasses.as_tuple
  • typeshed: check obj arg
  • Potential:
    • type retval (tuple(42, 'spam'))
• dataclasses.replace
  • typeshed: check obj arg
  • mypy: check changes arg

The problem we're solving

For the purposes of the improved checks above, we should not treat any Data Class Transform'd class as a "true" dataclass, e.g. pydantic models do not function as dataclasses (i.e. is_dataclass(some_pydantic_model) is False). Straight reading of PEP-681 also does not suggest that. Meanwhile, some type checkers (mypy, pyright) are currently doing just that, and that's where we run into a conundrum:

1. If we keep the status quo, the 📖 intended use is broken.

   _ = replace(  # 💥 TypeError: replace() should be called on dataclass instances
     some_pydantic_model,
     my_str='bar',
   )

2. If we change type-checkers to stop treating Data Class Transforms as returning a "true" dataclass, then the 🎁 bonus use is broken, which would affect many projects.

   _ = replace(
     some_flax_model,  # E: Value of type variable "_DataclassT" of "replace" cannot be "SomeFlaxModel"
     foo='bar',
   )

3. If we revert Use a TypeGuard for dataclasses.is_dataclass(); refine asdict(), astuple(), fields(), replace() typeshed#9362, then we (a) lose type safety for obj, and (b) we keep inconsistent behavior in type-checkers whenever we have special treatment of dataclasses utility functions.

   _ = replace(  # 💥 TypeError: replace() should be called on dataclass instances
     some_pydantic_model,
     my_str=42,  # E: Argument 1 to "C" has incompatible type "int"; expected "str"
   )

4. If we remove all special treatment of dataclasses utility functions in type-checkers, we further degrade type-checking. As the saying goes, "no type checker, no problem".

Proposal

PEP-681 foresaw the need to add more parameters:

In the future, we may add additional parameters to dataclass_transform as needed to support common behaviors in user code. These additions will be made after reaching consensus on typing-sig rather than via additional PEPs.

I'm suggesting adding a returns_dataclass parameter to dataclass_transform:

 def dataclass_transform(
     *,
     eq_default: bool = True,
     order_default: bool = False,
     kw_only_default: bool = False,
+    returns_dataclass: bool = <...>,
     field_specifiers: tuple[type | Callable[..., Any], ...] = (),
     **kwargs: Any,
 ) -> Callable[[_T], _T]: ...

I'm not sure whether it should default to True or False. My unscientific exploration of uses in the wild suggests more 🎁 bonus uses than 📖 intended uses, which would tip the scale in favor of True. Meanwhile, adherence to PEP-681 contract would tip the scale in favor of False. But most importantly we'll gain a way to express this property of the returned class.

Footnotes

1. can probably find more if it makes the case stronger ↩

[JelleZijlstra]

A few thoughts:

• What does it mean to "return a dataclass"? I feel dataclasses are a structural concept, not a nominal one; it shouldn't matter whether the implementation calls @dataclass under the hood somewhere or does something else. Perhaps the right test is whether dataclasses.is_dataclass() will return true.
• The default for the new parameter should be false. PEP 681 as written did not envision that transformed classes would be actual dataclasses, and we should keep compatibility with the original spec by default. Users who want this new, additional behavior should opt in to it.

[ikonst]

Perhaps the right test is whether dataclasses.is_dataclass() will return true.

Yes, this one. Perhaps I should've used "Data Class" rather than "dataclass", or perhaps dataclasses should've picked a more humble name 😝

[erictraut]

Thanks for the great write-up Ilya.

I agree with Jelle's comments above.

Let's look at the motivating example, which I've expanded into a self-contained minimal sample:

from dataclasses import replace
from pydantic import BaseModel

class MyModel(BaseModel):
    my_str: str

some_pydantic_model = MyModel(my_str="foo")
_ = replace(some_pydantic_model, my_str="bar")

This produces no errors in pyright but produces an error in mypy 1.5.1 (on the last line). I'm guessing this is because mypy has special-case logic for the dataclasses.replace function, and it's overriding the typeshed definition of this function. Without special-case logic, the typeshed definition would indicate this is type safe, so no errors should be produced.

A potential solution here is to remove that error condition from mypy's special-case handling. This would eliminate the false positive error in cases where dataclass_transform is used with an actual dataclass or an implementation that is compatible with dataclasses.replace. This would admittedly leave a potential false negative in the case where dataclass_transform is used with a dataclass-like class that is not compatible with replace, but I think that's reasonable given that we have no evidence (to my knowledge) that developers are stumbling on this and unknowingly passing such objects to replace or the other four functions listed above.

With this straightforward fix, no changes would be required to typeshed, type checkers other than mypy, any third-party libraries, or any specs/documentation. If we later see evidence that developers are tripping over the false negative with some regularity, that could motivate a more complex and invasive change.

Another option that wasn't listed above that is worth considering...

Type checkers could synthesize a __dataclass_fields__ (and therefore claim compatibility with the five functions listed above) only in the case where the field_specifiers for the dataclass_transform are limited to dataclasses.Field and/or dataclasses.field. This is admittedly a heuristic, but I think it's defensible. Libraries like pydantic and attrs have their own field classes, whereas a dataclass_transform that uses stdlib dataclass under the hood will naturally use the stdlib dataclasses.Field for its field specifiers.

This solution would require a small change in mypy, pyright and pyre (and any other type checkers that have added support for PEP 681), but it wouldn't require any change to typeshed, any third-party libraries, or any specs/documentation.

[ikonst]

A potential solution here is to remove that error condition from mypy's special-case handling.

This is effectively scenario (3) after python/mypy#15962 is merged (on a side note, that error handling took a bunch of time to implement and proved gratuitous, so glad we're reverting it), ... assuming we don't revert the indiscriminate __dataclass_fields__ synthesizing from the type-checkers.

This is admittedly a heuristic, but I think it's defensible.

Very interesting... 🤔

[NeilGirdhar]

Type checkers could synthesize a __dataclass_fields__ (and therefore claim compatibility with the five functions listed above) only in the case where the field_specifiers for the dataclass_transform are limited to dataclasses.Field

I really like this idea! It would definitely solve my problem. A couple notes: I think it should support field specifiers that return any subclasses of dataclasses.Field (to stay flexible; the Jax family of custom dataclasses like Flax's use a custom field specifier to specify "staticness" in the Jax sense). And currently type checkers struggle with field specifiers that return Field[T] and so the field specifier functions in the typeshed and user code tend to return T to avoid type errors. So, I think you might need to special case Field[T] as an acceptable static type for RHS values of dataclass attributes where the LHS has static type T?

[ikonst]

Project	fields=	compatible?
flax	(flax.struct.field1,)	flax.struct.field would have to be made a derived class rather than a factory
strawberry-graphql	(dataclass.field, StrawberryField2)	assuming derived StrawberryField is ok, this would work
serde	(,)	would need trivial update
ibis	(,)	would need trivial update

... so this might work, although not quite "no changes to third-party libraries". Also, if we go down this route, I think it'll be prudent to document it in somewhere, either in cpython's typing.rst or in python/typing's docs, or in each type checker's docs, but certainly not solely in this discussion 🙏 . IMO the first is preferable, since it's the main and only documentation on the fields parameter of @dataclass_transform. I can see the concern about enshrining this behavior into 'spec', but...

Also, given that libs would have to make some updates anyway, this would be clearly the more complex syntax (vs. bool param), and I wonder to what advantage (fewer params? some future flexibility? some merit in avoiding special-casing dataclass more than it already is?).

Footnotes

1. a function returning dataclasses.field(metadata={'pytree_node': pytree_node}, **kwargs) ↩

2. class StrawberryField(dataclasses.Field): ... ↩

[erictraut]

Thanks for doing that investigation. I had assumed that anyone using a real stdlib dataclass would also use dataclasses.Field or dataclasses.field. That appears to be a bad assumption.

Then I'd be in favor of simply fixing the issue in mypy and calling it good enough. That wouldn't require any further changes, and I think it largely addresses the problem as stated. Unless/until we see evidence that developers are mistakenly passing incompatible dataclass-like objects to the five functions listed above, I don't see a strong need to take further action.

[NeilGirdhar]

Then I'd be in favor of simply fixing the issue in mypy and calling it good enough.

Sorry, but what does that mean exactly?

[ikonst]

I think Eric means merging python/mypy#15962.

Something notable about mypy: this was a problem because of non-uniform treatment of Data Class Transforms (DCT for short). We would synthesize __dataclass_fields__ for them, so typeshed signature of replace would accept both "true" DCs and DCTs, but won't generate a __mypy-replace symbol for them, so DCT instances would get the custom "Argument 1 to "replace" has incompatible type "..."; expected a dataclass" error.

By removing the custom error code, now both "true" DCs and DCTs get accepted by replace.

To summarize, that'll take us to a point where pydantic models (a) won't be rejected by DC utils, (b) would have the same special handling of DC utils (e.g. replace changes arg checking, maybe in the future fields retval etc.) and the user would be unlikely to stumble into it as one would be unlikely to try and use a wine opener as a shoe horn.