update: refine traffic allocation type hints and key naming in bucketer and decision service

Author: FarhanAnjum-opti

optimizely/bucketer.py

@@ -169,7 +169,7 @@ def bucket_to_entity_id(
         self,
         bucketing_id: str,
         experiment: Experiment,
-        traffic_allocations: list,
+        traffic_allocations: list[TrafficAllocation],
         group: Optional[Group] = None
     ) -> tuple[Optional[str], list[str]]:
         """
@@ -193,8 +193,8 @@ def bucket_to_entity_id(
 
             matched = False
             for allocation in group.trafficAllocation:
-                end_of_range = cast(int, allocation.get("end_of_range", 0))
-                entity_id = cast(Optional[str], allocation.get("entity_id"))
+                end_of_range = allocation['endOfRange']
+                entity_id = allocation['entityId']
                 if bucket_val < end_of_range:
                     matched = True
                     if entity_id != experiment.id:
@@ -218,8 +218,8 @@ def bucket_to_entity_id(
         decide_reasons.append(f'Generated experiment bucket value {bucket_val} for key "{bucket_key}".')
 
         for allocation in traffic_allocations:
-            end_of_range = allocation.get("end_of_range", 0)
-            entity_id = allocation.get("entity_id")
+            end_of_range = allocation['endOfRange']
+            entity_id = allocation['entityId']
             if bucket_val < end_of_range:
                 decide_reasons.append(
                     f'User bucketed into entity id "{entity_id}".'

optimizely/decision_service.py

@@ -30,6 +30,7 @@
     # prevent circular dependenacy by skipping import at runtime
     from .project_config import ProjectConfig
     from .logger import Logger
+    from .helpers.types import TrafficAllocation
 
 
 class CmabDecisionResult(TypedDict):
@@ -391,9 +392,9 @@ def get_variation(
         if experiment.cmab:
             CMAB_DUMMY_ENTITY_ID = "$"
             # Build the CMAB-specific traffic allocation
-            cmab_traffic_allocation = [{
-                "entity_id": CMAB_DUMMY_ENTITY_ID,
-                "end_of_range": experiment.cmab['trafficAllocation']
+            cmab_traffic_allocation: list[TrafficAllocation] = [{
+                "entityId": CMAB_DUMMY_ENTITY_ID,
+                "endOfRange": experiment.cmab['trafficAllocation']
             }]
 
             # Check if user is in CMAB traffic allocation