Spaces:
Building
Building
File size: 5,668 Bytes
f5f3483 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 |
# Copyright 2024 The etils Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Field utils."""
from __future__ import annotations
import dataclasses
import typing
from typing import Any, Callable, Generic, Optional, Type, TypeVar
from etils import epy
_Dataclass = Any
_In = Any
_Out = Any
_InT = TypeVar('_InT')
_OutT = TypeVar('_OutT')
def field(
*,
validate: Optional[Callable[[_In], _OutT]] = None,
**kwargs: Any,
) -> dataclasses.Field[_OutT]:
"""Like `dataclasses.field`, but allow `validator`.
Args:
validate: A callable `(x) -> x` called each time the variable is assigned.
**kwargs: Kwargs forwarded to `dataclasses.field`
Returns:
The field.
"""
if validate is None:
return dataclasses.field(**kwargs)
else:
field_ = _Field(validate=validate, field_kwargs=kwargs)
return typing.cast(dataclasses.Field, field_) # pylint: disable=g-bare-generic
class _Field(Generic[_InT, _OutT]):
"""Field descriptor."""
def __init__(
self,
validate: Callable[[_InT], _OutT],
field_kwargs: dict[str, Any],
) -> None:
"""Constructor.
Args:
validate: A callable called each time the variable is assigned.
field_kwargs: Kwargs forwarded to `dataclasses.field`
"""
# Attribute name and objtype refer to the object in which the descriptor
# is applied. E.g. if `A.x = edc.field()`:
# * _attribute_name = 'x'
# * _objtype = A
self._attribute_name: Optional[str] = None
self._objtype: Optional[Type[_Dataclass]] = None
self._validate_fn = validate
self._field_kwargs = field_kwargs
# Whether `__get__` has not been called yet. See `__get__` for details.
self._first_getattr_call: bool = True
def __set_name__(self, objtype: Type[_Dataclass], name: str) -> None:
"""Bind the descriptor to the class (PEP 487)."""
self._objtype = objtype
self._attribute_name = name
def __get__(
self,
obj: Optional[_Dataclass],
objtype: Optional[Type[_Dataclass]] = None,
) -> _OutT:
"""Called when `MyDataclass.x` or `my_dataclass.x`."""
# Called as `MyDataclass.my_attribute`
if obj is None:
if self._first_getattr_call:
# Count the number of times `dataclasses.dataclass(cls)` calls
# `getattr(cls, f.name)`.
# The first time, we return a `dataclasses.Field` to let dataclass
# do the magic.
# The second time, `dataclasses.dataclass` delete the descriptor if
# `isinstance(getattr(cls, f.name, None), Field)`. So it is very
# important to return anything except a `dataclasses.Field`.
# This rely on implementation detail, but seems to hold for python
# 3.6-3.10.
self._first_getattr_call = False
return dataclasses.field(**self._field_kwargs)
else:
# TODO(epot): Could better handle default value: Either by returning
# the default value, or raising an AttributeError. Currently, we just
# return the descriptor:
# assert isinstance(MyDataclass.my_attribute, _Field)
return self
else:
# Called as `my_dataclass.my_path`
return _getattr(obj, self._attribute_name)
def __set__(self, obj: _Dataclass, value: _InT) -> None:
"""Called as `my_dataclass.x = value`."""
# Validate the value during assignement
_setattr(obj, self._attribute_name, self._validate(value))
def _validate(self, value: _InT) -> _OutT:
try:
return self._validate_fn(value)
except Exception as e: # pylint: disable=broad-exception-caught
epy.reraise(e, prefix=f'Error assigning {self._attribute_name!r}: ')
# Because there is one instance of the `_Field` per class, shared across all
# class instances, we need to store the per-object state somewhere.
# The simplest is to attach the state in an extra `dict[str, value]`:
# `_dataclass_field_values`.
def _getattr(
obj: _Dataclass,
attribute_name: str,
) -> _Out:
"""Returns the `obj.attribute_name`."""
_init_dataclass_state(obj)
# Accessing the attribute before it was set (e.g. before super().__init__)
if attribute_name not in obj._dataclass_field_values: # pylint: disable=protected-access
raise AttributeError(
f"type object '{type(obj).__qualname__}' has no attribute "
f"'{attribute_name}'"
)
else:
return obj._dataclass_field_values[attribute_name] # pylint: disable=protected-access
def _setattr(
obj: _Dataclass,
attribute_name: str,
value: _In,
) -> None:
"""Set the `obj.attribute_name = value`."""
# Note: In `dataclasses.dataclass(frozen=True)`, obj.__setattr__ will
# correctly raise a `FrozenInstanceError` before `DataclassField.__set__` is
# called.
_init_dataclass_state(obj)
# fmt: off
obj._dataclass_field_values[attribute_name] = value # pylint: disable=protected-access
# fmt: on
def _init_dataclass_state(obj: _Dataclass) -> None:
"""Initialize the object state containing all DataclassField values."""
if not hasattr(obj, '_dataclass_field_values'):
# Use object.__setattr__ for frozen dataclasses
object.__setattr__(obj, '_dataclass_field_values', {})
|