# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved """ Helpers for converting python dataclasses to/from aws-friendly formats. There's likely already an existing library that exists somewhere that does this much better, but after spending 5 minutes looking, just try and write something on our own. How to use: import typing as t from hmalib.common.aws_dataclass import HasAWSSerialization from dataclasses import dataclass @dataclass class Item(HasAWSSerialization): a: int b: str c: t.List[str] item = Item(5, "5", ["five"]) more_portable_dict = item.to_aws() # Can write to dynamodb, or same_item = Item.from_aws(more_portable_dict) """ from decimal import Decimal from dataclasses import dataclass, field, fields, is_dataclass import json import typing as t T = t.TypeVar("T") class AWSSerializationFailure(ValueError): pass def py_to_aws(py_field: t.Any, in_type: t.Optional[t.Type[T]] = None) -> T: """ Convert a py item into its AWS equivalent. Should exactly inverse aws_to_py """ if in_type is None: in_type = type(py_field) origin = t.get_origin(in_type) args = t.get_args(in_type) check_type = origin or in_type if isinstance(check_type, t.ForwardRef): raise AWSSerializationFailure( "Serialization error: " f"Expected no forward refs, but detected {check_type}. " "Rework your dataclasses to avoid forward references." ) if not isinstance(py_field, check_type): raise AWSSerializationFailure( "Serialization error: " f"Expected {check_type} got {type(py_field)} ({py_field!r})" ) if in_type == int: # N # Technically, this also needs to be converted to decimal, # but the boto3 translater seems to handle it fine return py_field # type: ignore # mypy/issues/10003 if in_type == float: # N # WARNING WARNING # floating point is not truly supported in dynamodb # We can fake it for numbers without too much precision # but Decimal("3.4") != float(3.4) return Decimal(str(py_field)) # type: ignore # mypy/issues/10003 if in_type == Decimal: # N return py_field # type: ignore # mypy/issues/10003 if in_type == str: # S return py_field # type: ignore # mypy/issues/10003 if in_type == bool: # BOOL return py_field # type: ignore # mypy/issues/10003 if in_type == t.Set[str]: # SS return py_field # type: ignore # mypy/issues/10003 if in_type == t.Set[int]: # SN return {i for i in py_field} # type: ignore # mypy/issues/10003 if in_type == t.Set[float]: # SN # WARNING WARNING # floating point is not truly supported in dynamodb # See note above return {Decimal(str(s)) for s in py_field} # type: ignore # mypy/issues/10003 if origin is list: # L return [py_to_aws(v, args[0]) for v in py_field] # type: ignore # mypy/issues/10003 # various simple collections that don't fit into a # special cases above can likely be coerced into list. if origin is set: # L - Special case return [py_to_aws(v, args[0]) for v in py_field] # type: ignore # mypy/issues/10003 if origin is dict and args[0] is str: # M return {k: py_to_aws(v, args[1]) for k, v in py_field.items()} # type: ignore # mypy/issues/10003 if is_dataclass(in_type): return { field.name: py_to_aws(getattr(py_field, field.name), field.type) for field in fields(in_type) } # type: ignore # mypy/issues/10003 raise AWSSerializationFailure(f"Missing Serialization logic for {in_type!r}") def aws_to_py(in_type: t.Type[T], aws_field: t.Any) -> T: """ Convert an AWS item back into its py equivalent This might not even be strictly required, but we check that all the types are roughly what we expect, and convert Decimals back into ints/floats """ origin = t.get_origin(in_type) args = t.get_args(in_type) check_type = origin if in_type is float: check_type = Decimal elif in_type is int: check_type = (int, Decimal) elif is_dataclass(in_type): check_type = dict elif check_type is set and args: if args[0] not in (str, float, int, Decimal): check_type = list if not isinstance(aws_field, check_type or in_type): # If you are getting random deserialization errors in tests that you did # not touch, have a look at # https://github.com/facebook/ThreatExchange/issues/697 raise AWSSerializationFailure( "Deserialization error: " f"Expected {in_type} got {type(aws_field)} ({aws_field!r})" ) if in_type is int: # N return int(aws_field) # type: ignore # mypy/issues/10003 if in_type is float: # N return float(aws_field) # type: ignore # mypy/issues/10003 if in_type is Decimal: # N return aws_field # type: ignore # mypy/issues/10003 if in_type is str: # S return aws_field # type: ignore # mypy/issues/10003 if in_type is bool: # BOOL return aws_field # type: ignore # mypy/issues/10003 if in_type is t.Set[str]: # SS return aws_field # type: ignore # mypy/issues/10003 if in_type is t.Set[int]: # SN return {int(s) for s in aws_field} # type: ignore # mypy/issues/10003 if in_type is t.Set[float]: # SN return {float(s) for s in aws_field} # type: ignore # mypy/issues/10003 if origin is set: # L - special case return {aws_to_py(args[0], v) for v in aws_field} # type: ignore # mypy/issues/10003 if origin is list: # L return [aws_to_py(args[0], v) for v in aws_field] # type: ignore # mypy/issues/10003 # It would be possible to add support for nested dataclasses here, which # just become maps with the keys as their attributes # Another option would be adding a new class that adds methods to convert # to an AWS-friendly struct and back if origin is dict and args[0] is str: # M # check if value type of map origin is explicitly set return {k: aws_to_py(args[1], v) for k, v in aws_field.items()} # type: ignore # mypy/issues/10003 if is_dataclass(in_type): kwargs = {} for field in fields(in_type): if not field.init: continue val = aws_field.get(field.name) if val is None: continue # Hopefully missing b/c default or version difference kwargs[field.name] = aws_to_py(field.type, val) return in_type(**kwargs) # type: ignore # No idea how to correctly type this raise AWSSerializationFailure(f"Missing deserialization logic for {in_type!r}") class HasAWSSerialization: """Convenience mixin to add serialization to a class""" def to_aws(self): return py_to_aws(self) def to_aws_json(self): return json.dumps(self.to_aws()) @classmethod def from_aws(cls: t.Type[T], val: t.Dict[str, t.Any]) -> T: return aws_to_py(cls, val) @classmethod def from_aws_json(cls: t.Type[T], val: str) -> T: return aws_to_py(cls, json.loads(val))