pythonpython-typingpython-dataclassespyright

Abbreviating dataclass decorator without losing IntelliSense

Scenario

Suppose I want to create an alias for a dataclasses.dataclass decorator with specific arguments. For example:

# Instead of repeating this decorator all the time:
@dataclasses.dataclass(frozen=True, kw_only=True)
class Entity:
    ...

# I just write something like this:
@struct
class Entity:
    ...

The static analyzer I am using is Pylance, in Visual Studio Code.

I am using Python 3.11.

Attempt 1: Direct Assignment (Runtime ✅, Static Analysis ❌)

My first instinct was to leverage the fact that functions are first-class citizens and simply assign the created decorator function to a custom name. This works at runtime, but Pylance no longer recognizes Entity as a dataclass, as evident from the static analysis error:

struct = dataclasses.dataclass(frozen=True, kw_only=True)

@struct
class Entity:
    name: str
    value: int

# STATIC ANALYZER:
# Expected no arguments to "Entity" constructor Pylance(reportCallIssue)
valid_entity = Entity(name="entity", value=42)

# RUNTIME:
# Entity(name='entity', value=42)
print(valid_entity)

Attempt 2: Wrapping (Runtime ❌, Static Analysis ❌)

I then thought that maybe some information was being lost somehow if I just assign to another name (though I don't see why that would be the case), so I looked to wrapping it with functools. However, this still has the same behavior in static analysis and even causes a runtime error, when I apply @struct:

import dataclasses
import functools

def struct(cls):
    decorator = dataclasses.dataclass(frozen=True, kw_only=True)
    decorated_cls = decorator(cls)
    functools.update_wrapper(decorated_cls, cls)
    return decorated_cls

# No error reported by static analyzer, but runtime error at `@struct`:
# AttributeError: 'mappingproxy' object has no attribute 'update'
@struct
class Entity:
    name: str
    value: int

# STATIC ANALYZER:
# Expected no arguments to "Entity" constructor Pylance(reportCallIssue)
# RUNTIME:
# (this line doesn't even get reached)
valid_entity = Entity(name="entity", value=42)

Full traceback:

Traceback (most recent call last):
  File "C:\Users\***\temp.py", line 12, in <module>
    @struct
     ^^^^^^
  File "C:\Users\***\temp.py", line 7, in struct
    functools.update_wrapper(decorated_cls, cls)
  File "C:\Users\***\AppData\Local\Programs\Python\Python311\Lib\functools.py", line 58, in update_wrapper
    getattr(wrapper, attr).update(getattr(wrapped, attr, {}))
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
AttributeError: 'mappingproxy' object has no attribute 'update'

Attempt 3: Wrapper Factory (Runtime ✅, Static Analysis ❌)

I then tried making struct a decorator factory instead and used functools.wraps() on a closure function that just forwards to the dataclass function. This now works at runtime, but Pylance still reports the same error as in Attempt 1:

def struct():
    decorator = dataclasses.dataclass(frozen=True, kw_only=True)

    @functools.wraps(decorator)
    def decorator_wrapper(*args, **kwargs):
        return decorator(*args, **kwargs)
    return decorator_wrapper

@struct()
class Entity:
    name: str
    value: int

# STATIC ANALYZER:
# Expected no arguments to "Entity" constructor Pylance(reportCallIssue)
valid_entity = Entity(name="entity", value=42)

# RUNTIME:
# Entity(name='entity', value=42)
print(valid_entity)

I also found that using the plain dataclasses.dataclass function itself (no ()) has the exact same problem across all 3 attempts.

Is there any way to get this to work without messing up IntelliSense?

Optional follow-up: why did Attempt 2 fail at runtime?

Solution

  • Decorate struct() with dataclass_transform(frozen_default = True, kw_only_default = True):

    (playgrounds: Mypy, Pyright)

    # 3.11+
from typing import dataclass_transform
# 3.10-
from typing_extensions import dataclass_transform

@dataclass_transform(frozen_default = True, kw_only_default = True)
def struct[T](cls: type[T]) -> type[T]:
    return dataclass(frozen = True, kw_only = True)(cls)

    # By the way, you can actually pass all of them
    # to dataclass() in just one call:
    #     dataclass(cls, frozen = True, kw_only = True)
    # It's just that this signature isn't defined statically.

    @struct
class Entity:
    name: str
    value: int

reveal_type(Entity.__init__)  # (self: Entity, *, name: str, value: int) -> None

valid_entity = Entity(name="entity", value=42)  # fine
valid_entity.name = ""                          # error: "Entity" is frozen

    dataclass_transform() is used to mark dataclass transformers (those that has similar behaviour to the built-in dataclasses.dataclass). It accepts a number of keyword arguments, in which:

    "By default" means that, unless otherwise specified via the frozen/kw_only arguments to the @struct decorator, the @struct-decorated class will behave as such. However, since struct itself takes no such arguments, "by default" here is the same as "always".