Skip to content

Instantly share code, notes, and snippets.

@rnag

rnag/dec.py Secret

Last active August 4, 2021 17:50
Show Gist options
  • Save rnag/e8c641d7fed3e2967dd138d1ae249915 to your computer and use it in GitHub Desktop.
Save rnag/e8c641d7fed3e2967dd138d1ae249915 to your computer and use it in GitHub Desktop.
Dataclass Support for Default Properties (modified)
from dataclasses import MISSING, Field
from functools import wraps
from typing import Dict, Any, get_type_hints
def dataclass_property_support(*args, **kwargs):
"""Adds support for using properties with default values in dataclasses."""
cls = type(*args, **kwargs)
cls_dict: Dict[str, Any] = args[2]
annotations = get_type_hints(cls)
def get_default_from_annotation(field_: str):
"""Get the default value for the type annotated on a field"""
default_type = annotations.get(field_)
try:
return default_type()
except TypeError:
return None
# For each property, we want to replace the annotation for the underscore-
# leading field associated with that property with the 'public' field
# name, and this mapping helps us keep a track of that.
annotation_repls = {}
for f, val in cls_dict.items():
if isinstance(val, property):
if val.fset is None:
# The property is read-only, not settable
continue
if f.startswith('_'):
# This metaclass should only work for 'public' properties
continue
# The field with a leading underscore
under_f = '_' + f
try:
# Get the value of the underscored field
default = getattr(cls, under_f)
except AttributeError:
# The public field is probably type-annotated but not defined
# i.e. my_var: str
default = get_default_from_annotation(under_f)
else:
# Check if the value of underscored field is a dataclass
# Field. If so, we can use the `default` if one is set.
if isinstance(default, Field):
if default.default is not MISSING:
default = default.default
else:
default = get_default_from_annotation(under_f)
def wrapper(fset, initial_val):
"""
Wraps the property `setter` method to check if we are passed
in a property object itself, which will be true when no
initial value is specified (thanks to @Martin CR).
"""
@wraps(fset)
def new_fset(self, value):
if isinstance(value, property):
value = initial_val
fset(self, value)
return new_fset
# Wraps the `setter` for the property
val = val.setter(wrapper(val.fset, default))
# Set the field that does not start with an underscore
setattr(cls, f, val)
# Also add it to the list of class annotations to replace later
# (this is what `dataclasses` uses to add the field to the
# constructor)
annotation_repls[under_f] = f
# Delete the field name that starts with an underscore
try:
delattr(cls, under_f)
except AttributeError:
pass
if annotation_repls and getattr(cls, '__annotations__', None):
# Use a comprehension approach because we want to replace a
# key while preserving the insertion order, because the order
# of fields does matter when the constructor is called.
cls.__annotations__ = {annotation_repls.get(f, f): type_
for f, type_ in cls.__annotations__.items()}
return cls
from dataclasses import dataclass
from typing import Union
from dec import dataclass_property_support
@dataclass
class Vehicle(metaclass=dataclass_property_support):
_wheels: Union[int, str] = 4
@property
def wheels(self) -> int:
return self._wheels
@wheels.setter
def wheels(self, wheels: Union[int, str]):
self._wheels = int(wheels)
def main():
v = Vehicle()
print(v)
# My IDE does complain on this part - it also suggests `_wheels` as a
# keyword argument to the constructor, and that's expected, but it will
# error if you try it that way.
v = Vehicle(wheels=3)
print(v)
# Passing positional arguments seems to be preferable as the IDE does not
# complain here.
v = Vehicle('6')
print(v)
assert v.wheels == 6, 'The constructor should use our setter method'
# Confirm that we go through our setter method
v.wheels = '123'
assert v.wheels == 123
if __name__ == '__main__':
main()
@rnag
Copy link
Author

rnag commented Aug 4, 2021

Edit: I've managed to encapsulate the above logic, including support for additional edge cases, to a helper package in case it is useful to others.
https://pypi.org/project/dataclass-wizard/

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment