src/pyutils/dataclass_utils.py

   1 #!/usr/bin/env python3
   2
   3 """Utilities for dealing with Dataclasses.  A non-official type hint and some
   4 friendly wrappers around conversion to/from Dicts."""
   5
   6 import dataclasses
   7 from typing import Any, Dict, Protocol, Type
   8
   9
  10 class Dataclass(Protocol):
  11     """Dataclass isn't really a first class type and therefore there is no offical
  12     type hint for Dataclasses in Python (yet).  If you need one, here's a suitable
  13     stand in.  Example usage::
  14
  15         def f(d: Dataclass) -> Any:
  16             pass
  17
  18         def g(d: Dict[str, Any]) -> Dataclass:
  19             pass
  20     """
  21
  22     __dataclass_fields__: Dict
  23
  24
  25 def dataclass_from_dict(dataclass: Type[Dataclass], d: Dict[str, Any]) -> Dataclass:
  26     """Given a Dataclass type and a dict, return a populated instance.
  27
  28     Args:
  29         dataclass: the Class type to return an instance of
  30         d: the dict to be used to populate the new instance
  31
  32     Returns:
  33         A constructed and populated dataclass instance.
  34
  35     >>> from dataclasses import dataclass
  36     >>> from datetime import date
  37
  38     >>> @dataclass
  39     ... class Record:
  40     ...     name: str
  41     ...     phone: str
  42     ...     address: str
  43     ...     age: int
  44     ...     member_since: date
  45     ...
  46
  47     >>> d = {
  48     ...         'name': 'John Smith',
  49     ...         'phone': '555-1234',
  50     ...         'address': '994 Main St.',
  51     ...         'age': 26,
  52     ...         'member_since': date(2006, 5, 14),
  53     ...     }
  54
  55     >>> dataclass_from_dict(Record, d)
  56     Record(name='John Smith', phone='555-1234', address='994 Main St.', age=26, member_since=datetime.date(2006, 5, 14))
  57     """
  58     fields = {f.name for f in dataclasses.fields(dataclass) if f.init}
  59     filtered_args = {k: v for k, v in d.items() if k in fields}
  60     return dataclass(**filtered_args)
  61
  62
  63 def dataclass_to_dict(dataclass: Dataclass) -> Dict[str, Any]:
  64     """
  65     Returns:
  66         A dict-representation of a valid dataclass.
  67
  68     >>> from dataclasses import dataclass
  69     >>> from datetime import date
  70
  71     >>> @dataclass
  72     ... class Record:
  73     ...     name: str
  74     ...     phone: str
  75     ...     address: str
  76     ...     age: int
  77     ...     member_since: date
  78     ...
  79     >>> r = Record(name='Jane Doe', phone='555-1232', address='998 Main St.', age=23, member_since=date(2008, 3, 1))
  80     >>> dataclass_to_dict(r)
  81     {'name': 'Jane Doe', 'phone': '555-1232', 'address': '998 Main St.', 'age': 23, 'member_since': datetime.date(2008, 3, 1)}
  82     """
  83     assert dataclasses.is_dataclass(dataclass)
  84     return dataclasses.asdict(dataclass)
  85
  86
  87 if __name__ == '__main__':
  88     import doctest
  89
  90     doctest.testmod()