Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

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?

like image 601
vlin Avatar asked Sep 02 '25 03:09

vlin


1 Answers

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:

  • frozen_default = True means that the class decorated with @struct will be frozen "by default".
  • kw_only_default = True means that the constructor generated will only have keyword arguments (aside from self) "by default".

"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".

like image 154
InSync Avatar answered Sep 04 '25 18:09

InSync