Mypy does not recognize aliases of dataclasses.dataclass

[wkschwartz]

In Python 3.7.0 and mypy 0.630+dev-3fb16a2b8c5439074e39b75feebc109a439eede3, if I save the following code to mypy-test.py

from dataclasses import dataclass

dataclass_alias = dataclass
dataclass_alias2 = dataclass(order=False)

def dataclass_wrapper(cls):
	return dataclass(cls)

@dataclass
class A: arg1: int

@dataclass_alias
class B: arg1: int

@dataclass_alias2
class C: arg1: int

@dataclass_wrapper
class D: arg1: int

a = A(arg1=1)
b = B(arg1=1)
c = C(arg1=1)
d = D(arg1=1)

I get the following mypy output

$ mypy mypy-test.py 
mypy-test.py:22: error: Unexpected keyword argument "arg1" for "B"
mypy-test.py:23: error: Unexpected keyword argument "arg1" for "C"
mypy-test.py:24: error: Unexpected keyword argument "arg1" for "D"

Notice that mypy has no trouble recognizing the signature of A's constructor. I can understand how mypy would have a hard time with dataclass_wrapper, and dataclass_alias is probably not critical in terms of use case. However, dataclass_alias2 is important if you want (as I do) to define many data classes with the same options, especially if you want to make these options available as part of your module's public interface.

[JukkaL]

It's unlikely that mypy will support dataclass_wrapper, but dataclass_alias and dataclass_alias2 are something that mypy could support. This is low priority since adding support for these use cases would require some pretty non-trivial new machinery. Also, dataclasses is still new and we don't know how common this use case is. We can increase the priority if many users encounter this limitation.

[wkschwartz]

In the mean time, perhaps a single sentence added to the documentation might be helpful.

[gvanrossum]

@wkschwartz

In the mean time, perhaps a single sentence added to the documentation might be helpful.

That makes sense. Since you know where you (would have) looked in the docs, can you help us by sending a Pull Request?

[edaqa-uncountable]

Support for the dataclass_wrapper case would be helpful. I was trying to create a wrapped class that adds API information to my dataclasses.

The workaround I've resorted too is double-wrapping:

def my_class(desc: Optional[str] = None):
    def decorate(orig_class):
        orig_class.__unc_data = {
            'desc': desc,
        }
        return orig_class
    return decorate

@my_class(desc='Standard Layout')
@dataclass
class LayoutSimple:
    num_cols: int = my_field(desc='Number of columns')
    num_rows: int = my_field(desc='Number of rows')
    type: Literal['simple'] = 'simple'

I'd also be wrapping the fields in my_field to provide a similarly simplified interface.

[bukzor]

Does there not exist any annotation or cast that will tell mypy to treat some variable as if it were dataclasses.dataclass ? I had expected typing.cast to do the trick, but apparently there's some magic happening here that doesn't involve the type system -- the type of dataclass is simply Callable.

Here's some of my attempts, with corresponding errors:

demo.py:

import dataclasses
immutable = dataclasses.dataclass(frozen=True, order=True)

@immutable
class MyClass:
  myint: int
  
myobject = MyClass(3)
print(myobject.myint)

$ mypy demo.py 
demo.py:8: error: Too many arguments for "MyClass"

Next is my attempt to use annotations to request that mypy treat immutable as dataclass. This won't be perfectly correct, but still much preferable to the above behavior. To my eye, it actually seems to almost work. The necessary changes would be to allow dataclasses.dataclass? to be a valid type, and also to have the the corresponding mypy dataclass plugin act on it.

immutable: dataclasses.dataclass = dataclasses.dataclass(frozen=True, order=True)

demo.py:3: error: Function "dataclasses.dataclass" is not valid as a type
demo.py:3: note: Perhaps you need "Callable[...]" or a callback protocol?
demo.py:6: error: dataclasses.dataclass? not callable
demo.py:9: error: Too many arguments for "MyClass"

Using cast rather than an annotation actually produces the exact same errors:

immutable = typing.cast(dataclasses.dataclass, dataclasses.dataclass(frozen=True, order=True))

demo.py:4: error: Function "dataclasses.dataclass" is not valid as a type
demo.py:4: note: Perhaps you need "Callable[...]" or a callback protocol?
demo.py:7: error: dataclasses.dataclass? not callable
demo.py:10: error: Too many arguments for "MyClass"

From my perspective, the problem here is that the special-case handling for dataclasses.dataclass isn't attached to its type. What is it in fact attached to? Its name? I don't know, but I expect it's the name. Adding one level of abstraction would fix the problem, I believe: define a type for dataclass (I'd name it DataclassMeta) and have the mypy plugin act on anything with that type. Then other objects (not just the one stdlib object) can have the behavior of dataclass.

---

P.S. I hope my use case isn't relevant, but:

My intent is to ease the use of immutable-by-default design in my team by making a clean alias of dataclasses.dataclass(frozen=True, order=True) -- otherwise, the "good" code looks worse than the "worse".

[gvanrossum]

If you need help figuring this out today, I recommend asking in the typing discussions group: https://github.com/python/typing/discussions

[parched]

If this proposal was implemented I think you'd be able to do something like this

@typing.dataclass_transform()
def dataclass_wrapper(cls):
	return dataclass(cls)

which would be neat.

[bukzor]

@typing.dataclass_transform()

To me, this improvement seems much too specific to dataclass. I'd rather see a general solution that works for enum.unique, functools.total_ordering and the rest.

I think the Right Way to fix this is dataclass() should return something more specific than Callable, and mypy plugins should be prevented from seeing the name of variables, seeing only the types. Then some (most?) of my attempts above would Just Work.

[NeilGirdhar]

I would like this too, but it might have been easier if there were a dataclass base class (say, DataClass) in the standard library that applies dataclass in its __init_subclass__. Then mypy could just check for subclasses of DataClass) and custom subclasses would work fine.

[NeilGirdhar]

I think I just got this working by adding a custom mypy plugin like this:

from typing import Type

import mypy.plugins.dataclasses
from mypy.plugin import Plugin


class CustomPlugin(Plugin):
    pass

def plugin(version: str) -> Type[CustomPlugin]:
    mypy.plugins.dataclasses.dataclass_makers.add('tjax._src.dataclasses.dataclass.dataclass')
    # mypy.plugins.dataclasses.field_makers.add('tjax._src.dataclasses.helpers.field')
    return CustomPlugin

[NeilGirdhar]

The dataclass transforms PEP (681) would be a great addition to MyPy 😄 It would solve this problem, and the problem of base classes that define dataclasses.

[NeilGirdhar]

Perhaps this issue should be renamed "Implement PEP 681" to give it a bit more visibility? It's a huge deal for projects like Flax that extensively use dataclasses (both using a modified decorator flax.struct.dataclass and a special base class flax.linen.Module).

[hauntsaninja]

Use PEP 681, which is implemented.

If you need something that is truly, exactly a dataclass alias, you can do something like from dataclasses import dataclass as dataclass_alias (instead of assignment, possibly also mix in an if TYPE_CHECKING) and mypy will do the right thing