fixed typing

Author: sophia-chen-ttd

tests/test_identity_map_client.py

@@ -7,12 +7,6 @@
 from uid2_client import IdentityMapClient, IdentityMapInput, normalize_and_hash_email, normalize_and_hash_phone
 
 
-@unittest.skipIf(
-    os.getenv("UID2_BASE_URL") == None
-    or os.getenv("UID2_API_KEY") == None
-    or os.getenv("UID2_SECRET_KEY") == None,
-    reason="Environment variables UID2_BASE_URL, UID2_API_KEY, and UID2_SECRET_KEY must be set",
-)
 class IdentityMapIntegrationTests(unittest.TestCase):
     UID2_BASE_URL = None
     UID2_API_KEY = None

tests/test_identity_map_v3_client.py

@@ -8,12 +8,6 @@
 from uid2_client.unmapped_identity_reason import UnmappedIdentityReason
 
 
-@unittest.skipIf(
-    os.getenv("UID2_BASE_URL") == None
-    or os.getenv("UID2_API_KEY") == None
-    or os.getenv("UID2_SECRET_KEY") == None,
-    reason="Environment variables UID2_BASE_URL, UID2_API_KEY, and UID2_SECRET_KEY must be set",
-)
 class IdentityMapV3IntegrationTests(unittest.TestCase):
     UID2_BASE_URL = None
     UID2_API_KEY = None

tests/test_publisher_client.py

@@ -8,12 +8,6 @@
 from urllib.request import HTTPError
 
 
-@unittest.skipIf(
-    os.getenv("EUID_BASE_URL") == None
-    or os.getenv("EUID_API_KEY") == None
-    or os.getenv("EUID_SECRET_KEY") == None,
-    reason="Environment variables EUID_BASE_URL, EUID_API_KEY, and EUID_SECRET_KEY must be set",
-)
 class PublisherEuidIntegrationTests(unittest.TestCase):
 
     EUID_SECRET_KEY = None

uid2_client/identity_map_v3_input.py

@@ -1,5 +1,5 @@
 import json
-from typing import List, Dict
+from typing import List, Dict, Literal
 
 from uid2_client import normalize_and_hash_email, normalize_and_hash_phone
 
@@ -36,7 +36,7 @@ def with_emails(self, emails: List[str]) -> 'IdentityMapV3Input':
     def with_email(self, email: str) -> 'IdentityMapV3Input':
         hashed_email = normalize_and_hash_email(email)
         self._hashed_emails.append(hashed_email)
-        self._add_hashed_to_raw_dii_mapping(hashed_email, email)
+        self._add_to_dii_mappings(hashed_email, email)
         return self
 
     def with_phones(self, phones: List[str]) -> 'IdentityMapV3Input':
@@ -47,7 +47,7 @@ def with_phones(self, phones: List[str]) -> 'IdentityMapV3Input':
     def with_phone(self, phone: str) -> 'IdentityMapV3Input':
         hashed_phone = normalize_and_hash_phone(phone)
         self._hashed_phones.append(hashed_phone)
-        self._add_hashed_to_raw_dii_mapping(hashed_phone, phone)
+        self._add_to_dii_mappings(hashed_phone, phone)
         return self
 
     def with_hashed_emails(self, hashed_emails: List[str]) -> 'IdentityMapV3Input':
@@ -57,7 +57,7 @@ def with_hashed_emails(self, hashed_emails: List[str]) -> 'IdentityMapV3Input':
 
     def with_hashed_email(self, hashed_email: str) -> 'IdentityMapV3Input':
         self._hashed_emails.append(hashed_email)
-        self._add_hashed_to_raw_dii_mapping(hashed_email, hashed_email)
+        self._add_to_dii_mappings(hashed_email, hashed_email)
         return self
 
     def with_hashed_phones(self, hashed_phones: List[str]) -> 'IdentityMapV3Input':
@@ -67,19 +67,19 @@ def with_hashed_phones(self, hashed_phones: List[str]) -> 'IdentityMapV3Input':
 
     def with_hashed_phone(self, hashed_phone: str) -> 'IdentityMapV3Input':
         self._hashed_phones.append(hashed_phone)
-        self._add_hashed_to_raw_dii_mapping(hashed_phone, hashed_phone)
+        self._add_to_dii_mappings(hashed_phone, hashed_phone)
         return self
 
-    def get_input_diis(self, identity_type: str, index: int) -> List[str]:
+    def get_input_diis(self, identity_type: Literal['email_hash', 'phone_hash'], index: int) -> List[str]:
         hashed_dii = self._get_hashed_dii(identity_type, index)
         return self._hashed_dii_to_raw_diis.get(hashed_dii, [])
 
-    def _add_hashed_to_raw_dii_mapping(self, hashed_dii: str, raw_dii: str):
+    def _add_to_dii_mappings(self, hashed_dii: str, raw_dii: str) -> None:
         if hashed_dii not in self._hashed_dii_to_raw_diis:
             self._hashed_dii_to_raw_diis[hashed_dii] = []
         self._hashed_dii_to_raw_diis[hashed_dii].append(raw_dii)
 
-    def _get_hashed_dii(self, identity_type: str, index: int) -> str:
+    def _get_hashed_dii(self, identity_type: Literal['email_hash', 'phone_hash'], index: int) -> str:
         if identity_type == "email_hash":
             return self._hashed_emails[index]
         elif identity_type == "phone_hash":

uid2_client/identity_map_v3_response.py

@@ -1,30 +1,31 @@
 import json
 from datetime import datetime, timezone
-from typing import Dict, List, Optional, Any
+from typing import Dict, List, Optional, Any, Literal
 
+from .identity_map_v3_input import IdentityMapV3Input
 from .unmapped_identity_reason import UnmappedIdentityReason
 
 
 class IdentityMapV3Response:
-    def __init__(self, response: str, identity_map_input):
+    def __init__(self, response: str, identity_map_input: IdentityMapV3Input):
         self._mapped_identities: Dict[str, MappedIdentity] = {}
         self._unmapped_identities: Dict[str, UnmappedIdentity] = {}
+
         response_json = json.loads(response)
-        self._status = response_json["status"]
+        api_response = ApiResponse.from_json(response_json)
+        self._status = api_response.status
 
         if not self.is_success():
             raise ValueError("Got unexpected identity map status: " + self._status)
 
-        body = response_json["body"]
-        self._populate_identities(body, identity_map_input)
+        self._populate_identities(api_response.body, identity_map_input)
 
-    def _populate_identities(self, api_response: Dict[str, List[Dict]], identity_map_input):
+    def _populate_identities(self, api_response: Dict[Literal['email_hash', 'phone_hash'], List['ApiIdentity']], identity_map_input: IdentityMapV3Input) -> None:
         for identity_type, identities in api_response.items():
             self._populate_identities_for_type(identity_map_input, identity_type, identities)
 
-    def _populate_identities_for_type(self, identity_map_input, identity_type: str, identities: List[Dict]):
-        for i, api_identity_data in enumerate(identities):
-            api_identity = ApiIdentity.from_json(api_identity_data)
+    def _populate_identities_for_type(self, identity_map_input: IdentityMapV3Input, identity_type: Literal['email_hash', 'phone_hash'], identities: List['ApiIdentity']) -> None:
+        for i, api_identity in enumerate(identities):
             input_diis = identity_map_input.get_input_diis(identity_type, i)
 
             for input_dii in input_diis:
@@ -49,6 +50,22 @@ def status(self) -> str:
         return self._status
 
 
+class ApiResponse:
+    def __init__(self, status: str, body: Dict[Literal['email_hash', 'phone_hash'], List['ApiIdentity']]):
+        self.status = status
+        self.body = body
+
+    @classmethod
+    def from_json(cls, data) -> 'ApiResponse':
+        api_body: Dict[Literal['email_hash', 'phone_hash'], List['ApiIdentity']] = {
+            'email_hash': [ApiIdentity.from_json(item) for item in data.get('body').get('email_hash', [])] if data.get('body').get('email_hash') else [],
+            'phone_hash': [ApiIdentity.from_json(item) for item in data.get('body').get('phone_hash', [])] if data.get('body').get('phone_hash') else [],
+        }
+        return cls(
+            status=data.get("status"),
+            body=api_body
+        )
+
 class ApiIdentity:
     def __init__(self, current_uid: Optional[str], previous_uid: Optional[str], 
                  refresh_from_seconds: Optional[int], error: Optional[str]):
@@ -74,7 +91,9 @@ def __init__(self, current_uid: str, previous_uid: Optional[str], refresh_from_s
         self._refresh_from = refresh_from_seconds
 
     @classmethod
-    def from_api_identity(cls, api_identity: ApiIdentity):
+    def from_api_identity(cls, api_identity: ApiIdentity) -> 'MappedIdentity':
+        if api_identity.current_uid is None or api_identity.refresh_from_seconds is None:
+            raise ValueError("Mapped identity cannot be created from API identity with missing current_uid or refresh_from_seconds")
         return cls(api_identity.current_uid,
                               api_identity.previous_uid,
                               datetime.fromtimestamp(api_identity.refresh_from_seconds, tz=timezone.utc))

uid2_client/request_response_util.py

@@ -1,5 +1,7 @@
 import base64
 import os
+from datetime import datetime
+from typing import Dict, Optional
 from urllib import request
 
 import pkg_resources
@@ -10,10 +12,10 @@
 
 BINARY = 'application/octet-stream'
 
-def _make_url(base_url, path):
+def _make_url(base_url: str, path: str) -> str:
     return base_url + path
 
-def auth_headers(auth_key):
+def auth_headers(auth_key: str) -> Dict[str, str]:
     try:
         version = pkg_resources.get_distribution("uid2_client").version
     except Exception:
@@ -22,7 +24,7 @@ def auth_headers(auth_key):
     return {'Authorization': 'Bearer ' + auth_key,
             "X-UID2-Client-Version": "uid2-client-python-" + version}
 
-def create_envelope(secret_key, now, data=None) -> 'Envelope':
+def create_envelope(secret_key: bytes, now: datetime, data: Optional[bytes] = None) -> 'Envelope':
     payload = int.to_bytes(int(now.timestamp() * 1000), 8, 'big')
     nonce = os.urandom(8)
     payload += nonce
@@ -33,27 +35,33 @@ def create_envelope(secret_key, now, data=None) -> 'Envelope':
     envelope += _encrypt_gcm(payload, None, secret_key)
     return Envelope(envelope, nonce)
 
-def make_request(base_url, path, headers, envelope) -> 'Uid2Response':
+def make_request(base_url: str, path: str, headers: Dict[str, str], envelope: Envelope) -> 'Uid2Response':
     resp = post(base_url, path, headers, envelope.envelope)
     return Uid2Response.from_string(resp.read())
 
-def make_binary_request(base_url, path, headers, envelope) -> 'Uid2Response':
+def make_binary_request(base_url: str, path: str, headers: Dict[str, str], envelope: Envelope) -> 'Uid2Response':
     headers['Content-Type'] = BINARY
     resp = post(base_url, path, headers, envelope.binary_envelope)
     return Uid2Response.from_bytes(resp.read())
 
-def post(base_url, path, headers, data):
+def post(base_url: str, path: str, headers: Dict[str, str], data: bytes):
     req = request.Request(_make_url(base_url, path), headers=headers, method='POST', data=data)
     return request.urlopen(req)
 
-def parse_response(secret_key, uid2_response: Uid2Response, nonce):
+def parse_response(secret_key: bytes, uid2_response: Uid2Response, nonce: bytes) -> str:
     if uid2_response.is_binary():
-        return _decrypt_payload(secret_key, uid2_response.as_bytes, nonce)
+        as_bytes = uid2_response.as_bytes
+        if as_bytes is None:
+            raise ValueError("Binary response has no bytes data")
+        return _decrypt_payload(secret_key, as_bytes, nonce)
     else:
-        encrypted_string = base64.b64decode(uid2_response.as_string)
+        as_string = uid2_response.as_string
+        if as_string is None:
+            raise ValueError("String response has no string data")
+        encrypted_string = base64.b64decode(as_string)
         return _decrypt_payload(secret_key, encrypted_string, nonce)
 
-def _decrypt_payload(secret_key, encrypted, nonce):
+def _decrypt_payload(secret_key: bytes, encrypted: bytes, nonce: bytes) -> str:
     payload = _decrypt_gcm(encrypted, secret_key)
     if nonce != payload[8:16]:
         raise ValueError("nonce mismatch")

uid2_client/uid2_response.py

@@ -7,19 +7,19 @@ def __init__(self, as_string: Optional[str], as_bytes: Optional[bytes]):
         self._as_bytes = as_bytes
 
     @classmethod
-    def from_string(cls, as_string: str):
+    def from_string(cls, as_string: str) -> 'Uid2Response':
         return cls(as_string, None)
 
     @classmethod
-    def from_bytes(cls, as_bytes: bytes):
+    def from_bytes(cls, as_bytes: bytes) -> 'Uid2Response':
         return cls(None, as_bytes)
 
     @property
-    def as_string(self) -> str:
+    def as_string(self) -> Optional[str]:
         return self._as_string
 
     @property
-    def as_bytes(self) -> bytes:
+    def as_bytes(self) -> Optional[bytes]:
         return self._as_bytes
 
     def is_binary(self) -> bool:

uid2_client/unmapped_identity_reason.py

@@ -7,7 +7,7 @@ class UnmappedIdentityReason(Enum):
     UNKNOWN = "unknown"
 
     @classmethod
-    def from_string(cls, reason_str):
+    def from_string(cls, reason_str: str) -> 'UnmappedIdentityReason':
         try:
             return cls(reason_str)
         except ValueError: