add fieldType for context fields

spotify · May 28, 2024 · ffff91b · ffff91b
1 parent 057d637
commit ffff91b
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 12 deletions.
diff --git a/confidence/confidence.py b/confidence/confidence.py
@@ -65,14 +65,12 @@ class ResolveResult(object):
 
 
 class Confidence:
-    context: Dict[str, Union[str, int, float, bool]] = {}
+    context: Dict[str, FieldType] = {}
 
-    def put_context(self, key: str, value: Union[str, int, float, bool]) -> None:
+    def put_context(self, key: str, value: FieldType) -> None:
         self.context[key] = value
 
-    def with_context(
-        self, context: Dict[str, Union[str, int, float, bool]]
-    ) -> "Confidence":
+    def with_context(self, context: Dict[str, FieldType]) -> "Confidence":
         new_confidence = Confidence(
             self._client_secret, self._region, self._apply_on_resolve
         )
@@ -124,7 +122,7 @@ def _evaluate(
         flag_key: str,
         value_type: Type[FieldType],
         default_value: FieldType,
-        context: Dict[str, Union[str, int, float, bool]],
+        context: Dict[str, FieldType],
     ) -> FlagResolutionDetails[Any]:
         if "." in flag_key:
             flag_id, value_path = flag_key.split(".", 1)
@@ -152,18 +150,22 @@ def _evaluate(
             flag_metadata={"flag_key": flag_key},
         )
 
-    def track(self, event_name: str, data: Dict[str, str]):
+    # type-arg: ignore
+    def track(self, event_name: str, data: Dict[str, FieldType]):
         return asyncio.create_task(self._send_event(event_name, data))
 
-    async def _send_event(self, event_name: str, data: Dict[str, str]) -> None:
+    async def track_async(self, event_name: str, data: Dict[str, FieldType]):
+        return await self._send_event(event_name, data)
+
+    async def _send_event(self, event_name: str, data: Dict[str, FieldType]) -> None:
         current_time = datetime.utcnow().isoformat() + "Z"
         request_body = {
             "clientSecret": self._client_secret,
             "sendTime": current_time,
             "events": [
                 {
                     "eventDefinition": f"eventDefinitions/{event_name}",
-                    "payload": {**self.context, **data},
+                    "payload": {"context": {**self.context}, **data},
                     "eventTime": current_time,
                 }
             ],
@@ -176,7 +178,7 @@ async def _send_event(self, event_name: str, data: Dict[str, str]) -> None:
         response.raise_for_status()
 
     def _resolve(
-        self, flag_name: FlagName, context: Dict[str, Union[str, int, float, bool]]
+        self, flag_name: FlagName, context: Dict[str, FieldType]
     ) -> ResolveResult:
         request_body = {
             "clientSecret": self._client_secret,

diff --git a/confidence/provider/provider.py b/confidence/provider/provider.py
@@ -64,7 +64,7 @@ class ResolveResult(object):
 
 
 def _to_openfeature_error_code(
-        error_code: Optional[ErrorCode]
+    error_code: Optional[ErrorCode],
 ) -> Optional[open_feature_exception.ErrorCode]:
     """
     Convert a confidence error code to an openfeature error code
@@ -196,7 +196,7 @@ def resolve_object_details(
     def _confidence_with_context(
         self, evaluation_context: Optional[EvaluationContext]
     ) -> confidence.confidence.Confidence:
-        eval_context: Dict[str, Union[str, int, float, bool]] = {}
+        eval_context: Dict[str, FieldType] = {}
         if evaluation_context:
             if evaluation_context.targeting_key:
                 eval_context["targeting_key"] = evaluation_context.targeting_key