File size: 4,272 Bytes
35cdf53 | 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 |
"""Config for the protein folding model and experiment."""
from collections.abc import Mapping
import copy
import dataclasses
import types
import typing
from typing import Any, ClassVar, TypeVar
_T = TypeVar('_T')
_ConfigT = TypeVar('_ConfigT', bound='BaseConfig')
def _strip_optional(t: type[Any]) -> type[Any]:
"""Transforms type annotations of the form `T | None` to `T`."""
if typing.get_origin(t) in (typing.Union, types.UnionType):
args = set(typing.get_args(t)) - {types.NoneType}
if len(args) == 1:
return args.pop()
return t
_NO_UPDATE = object()
class _Autocreate:
def __init__(self, **defaults: Any):
self.defaults = defaults
def autocreate(**defaults: Any) -> Any:
"""Marks a field as having a default factory derived from its type."""
return _Autocreate(**defaults)
def _clone_field(
field: dataclasses.Field[_T], new_default: _T
) -> dataclasses.Field[_T]:
if new_default is _NO_UPDATE:
return copy.copy(field)
return dataclasses.field(
default=new_default,
init=True,
kw_only=True,
repr=field.repr,
hash=field.hash,
compare=field.compare,
metadata=field.metadata,
)
@typing.dataclass_transform()
class ConfigMeta(type):
"""Metaclass that synthesizes a __post_init__ that coerces dicts to Config subclass instances."""
def __new__(mcs, name, bases, classdict):
cls = super().__new__(mcs, name, bases, classdict)
def _coercable_fields(self) -> Mapping[str, tuple[ConfigMeta, Any]]:
type_hints = typing.get_type_hints(self.__class__)
fields = dataclasses.fields(self.__class__)
field_to_type_and_default = {
field.name: (_strip_optional(type_hints[field.name]), field.default)
for field in fields
}
coercable_fields = {
f: t
for f, t in field_to_type_and_default.items()
if issubclass(type(t[0]), ConfigMeta)
}
return coercable_fields
cls._coercable_fields = property(_coercable_fields)
old_post_init = getattr(cls, '__post_init__', None)
def _post_init(self) -> None:
# Use get_type_hints instead of Field.type to ensure that forward
# references are resolved.
for field_name, (
field_type,
field_default,
) in self._coercable_fields.items(): # pylint: disable=protected-access
field_value = getattr(self, field_name)
if field_value is None:
continue
try:
match field_value:
case _Autocreate():
# Construct from field defaults.
setattr(self, field_name, field_type(**field_value.defaults))
case Mapping():
# Field value is not yet a `Config` instance; Assume we can create
# one by splatting keys and values.
args = {}
# Apply default args first, if present.
if isinstance(field_default, _Autocreate):
args.update(field_default.defaults)
args.update(field_value)
setattr(self, field_name, field_type(**args))
case _:
pass
except TypeError as e:
raise TypeError(
f'Failure while coercing field {field_name!r} of'
f' {self.__class__.__qualname__}'
) from e
if old_post_init:
old_post_init(self)
cls.__post_init__ = _post_init
return dataclasses.dataclass(kw_only=True)(cls)
class BaseConfig(metaclass=ConfigMeta):
"""Config base class.
Subclassing Config automatically makes the subclass a kw_only dataclass with
a `__post_init__` that coerces Config-subclass field values from mappings to
instances of the right type.
"""
# Provided by dataclasses.make_dataclass
__dataclass_fields__: ClassVar[dict[str, dataclasses.Field[Any]]]
# Overridden by metaclass
@property
def _coercable_fields(self) -> Mapping[str, tuple[type['BaseConfig'], Any]]:
return {}
def as_dict(self) -> Mapping[str, Any]:
result = dataclasses.asdict(self)
for field_name in self._coercable_fields:
field_value = getattr(self, field_name, None)
if isinstance(field_value, BaseConfig):
result[field_name] = field_value.as_dict()
return result
|