diff --git a/arachnid_shield_sdk/models/scanned_media.py b/arachnid_shield_sdk/models/scanned_media.py index 7298d18..d73d476 100644 --- a/arachnid_shield_sdk/models/scanned_media.py +++ b/arachnid_shield_sdk/models/scanned_media.py @@ -21,6 +21,13 @@ class MatchType(str, Enum): Exact = "exact" Near = "near" + @classmethod + def from_value(cls, value: typing.Optional[str]) -> typing.Optional["MatchType"]: + if value is None: + return None + return cls(value) + + @dataclasses.dataclass class NearMatchDetail: """A record of a near match (based on perceptual hashing) to a known image in our database. @@ -109,7 +116,7 @@ def from_dict(cls, src_dict: typing.Dict[str, typing.Any]) -> "ScannedMedia": sha1_base32=src_dict["sha1_base32"], size_bytes=src_dict["size_bytes"], classification=classification, - match_type=src_dict["match_type"], + match_type=MatchType.from_value(src_dict["match_type"]), near_match_details=near_match_details, ) diff --git a/arachnid_shield_sdk/models/scanned_pdq.py b/arachnid_shield_sdk/models/scanned_pdq.py index 9c0558e..f57b239 100644 --- a/arachnid_shield_sdk/models/scanned_pdq.py +++ b/arachnid_shield_sdk/models/scanned_pdq.py @@ -32,7 +32,7 @@ def from_dict(cls, src_dict: typing.Dict[str, typing.Any]) -> "PDQMatch": classification = MediaClassification(classification) return cls( classification=classification, - match_type=MatchType(src_dict["match_type"]), + match_type=MatchType.from_value(src_dict["match_type"]), near_match_details=near_match_details, ) @@ -62,7 +62,7 @@ def from_dict(cls, src_dict: typing.Dict[str, typing.Any]) -> "ScannedPDQHashes" near_match_details = NearMatchDetail.from_dict(near_match_details) scanned_hashes[key] = PDQMatch( classification=value['classification'], - match_type=MatchType(value['match_type']), + match_type=MatchType.from_value(value['match_type']), near_match_details=near_match_details ) return cls(