Fix up type annotations, tests, and formatting

Author: haacked

posthog/__init__.py

@@ -4,8 +4,9 @@
 
 from posthog.client import Client
 from posthog.exception_capture import Integrations  # noqa: F401
+from posthog.types import FeatureFlag, FlagsAndPayloads
 from posthog.version import VERSION
-from posthog.types import FlagsAndPayloads
+
 __version__ = VERSION
 
 """Settings."""
@@ -403,7 +404,7 @@ def get_feature_flag(
     only_evaluate_locally=False,  # type: bool
     send_feature_flag_events=True,  # type: bool
     disable_geoip=None,  # type: Optional[bool]
-) -> str | bool | None:
+) -> Optional[FeatureFlag]:
     """
     Get feature flag variant for users. Used with experiments.
     Example:
@@ -446,7 +447,7 @@ def get_all_flags(
     group_properties={},  # type: dict
     only_evaluate_locally=False,  # type: bool
     disable_geoip=None,  # type: Optional[bool]
-) -> dict[str, str | bool] | None:
+) -> Optional[dict[str, FeatureFlag]]:
     """
     Get all flags for a given user.
     Example:
@@ -477,7 +478,7 @@ def get_feature_flag_payload(
     only_evaluate_locally=False,
     send_feature_flag_events=True,
     disable_geoip=None,  # type: Optional[bool]
-) -> str:
+) -> Optional[str]:
     return _proxy(
         "get_feature_flag_payload",
         key=key,

posthog/client.py

@@ -6,7 +6,7 @@
 import sys
 import warnings
 from datetime import datetime, timedelta
-from typing import Any
+from typing import Any, Optional, Union
 from uuid import UUID, uuid4
 
 import distro  # For Linux OS detection
@@ -18,28 +18,20 @@
 from posthog.exception_utils import exc_info_from_error, exceptions_from_error_tuple, handle_in_app
 from posthog.feature_flags import InconclusiveMatchError, match_feature_flag_properties
 from posthog.poller import Poller
-from posthog.request import (
-    DEFAULT_HOST,
-    APIError,
-    batch_post,
-    decide,
-    determine_server_host,
-    get,
-    remote_config,
-    DecideResponse,
-)
-from posthog.utils import SizeLimitedDict, clean, guess_timezone, remove_trailing_slash
-from posthog.version import VERSION
+from posthog.request import DEFAULT_HOST, APIError, batch_post, decide, determine_server_host, get, remote_config
 from posthog.types import (
-    FlagsAndPayloads,
+    DecideResponse,
     FeatureFlag,
-    FlagValue,
     FlagMetadata,
-    to_values,
-    to_payloads,
-    to_flags_and_payloads,
+    FlagsAndPayloads,
+    FlagValue,
     normalize_decide_response,
+    to_flags_and_payloads,
+    to_payloads,
+    to_values,
 )
+from posthog.utils import SizeLimitedDict, clean, guess_timezone, remove_trailing_slash
+from posthog.version import VERSION
 
 try:
     import queue
@@ -265,7 +257,7 @@ def identify(self, distinct_id=None, properties=None, context=None, timestamp=No
 
     def get_feature_variants(
         self, distinct_id, groups=None, person_properties=None, group_properties=None, disable_geoip=None
-    ) -> dict[str, str | bool]:
+    ) -> dict[str, Union[bool, str]]:
         """
         Get feature flag variants for a distinct_id by calling decide.
         """
@@ -352,7 +344,7 @@ def capture(
             msg["properties"]["$groups"] = groups
 
         extra_properties: dict[str, Any] = {}
-        feature_variants: dict[str, bool | str] | None = {}
+        feature_variants: Optional[dict[str, Union[bool, str]]] = {}
         if send_feature_flags:
             try:
                 feature_variants = self.get_feature_variants(distinct_id, groups, disable_geoip=disable_geoip)
@@ -820,7 +812,7 @@ def get_feature_flag(
         only_evaluate_locally=False,
         send_feature_flag_events=True,
         disable_geoip=None,
-    ) -> FlagValue | None:
+    ) -> Optional[FlagValue]:
         """
         Get a feature flag value for a key by evaluating locally or remotely
         depending on whether local evaluation is enabled and the flag can be
@@ -859,7 +851,7 @@ def get_feature_flag(
             self._capture_feature_flag_called(
                 distinct_id,
                 key,
-                response,
+                response or False,
                 None,
                 flag_was_locally_evaluated,
                 groups,
@@ -877,7 +869,7 @@ def _locally_evaluate_flag(
         groups: dict[str, str],
         person_properties: dict[str, str],
         group_properties: dict[str, str],
-    ) -> FlagValue | None:
+    ) -> Optional[FlagValue]:
         if self.feature_flags is None and self.personal_api_key:
             self.load_feature_flags()
         response = None
@@ -949,7 +941,7 @@ def get_feature_flag_payload(
             self._capture_feature_flag_called(
                 distinct_id,
                 key,
-                response,
+                response or False,
                 payload,
                 flag_was_locally_evaluated,
                 groups,
@@ -967,33 +959,33 @@ def _get_feature_flag_details_from_decide(
         groups: dict[str, str],
         person_properties: dict[str, str],
         group_properties: dict[str, str],
-        disable_geoip: bool | None,
-    ) -> tuple[FeatureFlag | None, str]:
+        disable_geoip: Optional[bool],
+    ) -> tuple[Optional[FeatureFlag], Optional[str]]:
         """
         Calls /decide and returns the flag details and request id
         """
         resp_data = self.get_decide(distinct_id, groups, person_properties, group_properties, disable_geoip)
         request_id = resp_data.get("requestId")
         flags = resp_data.get("flags")
-        flag_details = flags.get(key)
+        flag_details = flags.get(key) if flags else None
         return flag_details, request_id
 
     def _capture_feature_flag_called(
         self,
         distinct_id: str,
         key: str,
         response: FlagValue,
-        payload: str | None,
+        payload: Optional[str],
         flag_was_locally_evaluated: bool,
         groups: dict[str, str],
-        disable_geoip: bool | None,
-        request_id: str | None,
-        flag_details: FeatureFlag | None,
+        disable_geoip: Optional[bool],
+        request_id: Optional[str],
+        flag_details: Optional[FeatureFlag],
     ):
         feature_flag_reported_key = f"{key}_{str(response)}"
 
         if feature_flag_reported_key not in self.distinct_ids_feature_flags_reported[distinct_id]:
-            properties = {
+            properties: dict[str, Any] = {
                 "$feature_flag": key,
                 "$feature_flag_response": response,
                 "locally_evaluated": flag_was_locally_evaluated,
@@ -1043,7 +1035,7 @@ def get_remote_config_payload(self, key: str):
         except Exception as e:
             self.log.exception(f"[FEATURE FLAGS] Unable to get decrypted feature flag payload: {e}")
 
-    def _compute_payload_locally(self, key: str, match_value: FlagValue) -> str | None:
+    def _compute_payload_locally(self, key: str, match_value: FlagValue) -> Optional[str]:
         payload = None
 
         if self.feature_flags_by_key is None:
@@ -1068,7 +1060,7 @@ def get_all_flags(
         group_properties={},
         only_evaluate_locally=False,
         disable_geoip=None,
-    ) -> dict[str, bool | str] | None:
+    ) -> Optional[dict[str, Union[bool, str]]]:
         response = self.get_all_flags_and_payloads(
             distinct_id,
             groups=groups,

posthog/request.py

@@ -3,12 +3,11 @@
 from datetime import date, datetime
 from gzip import GzipFile
 from io import BytesIO
-from typing import Any, Optional, Union, cast
+from typing import Any, Optional, Union
 
 import requests
 from dateutil.tz import tzutc
 
-from posthog.types import DecideResponse, FeatureFlag, FlagValue
 from posthog.utils import remove_trailing_slash
 from posthog.version import VERSION
 

posthog/test/ai/openai/test_openai.py

@@ -83,6 +83,7 @@ def mock_openai_response_with_responses_api():
         usage=ResponseUsage(
             input_tokens=10,
             output_tokens=10,
+            input_tokens_details={"prompt_tokens": 10, "cached_tokens": 0},
             output_tokens_details={"reasoning_tokens": 15},
             total_tokens=20,
         ),

posthog/test/test_client.py

@@ -10,8 +10,9 @@
 from posthog.client import Client
 from posthog.request import APIError
 from posthog.test.test_utils import FAKE_TEST_API_KEY
-from posthog.version import VERSION
 from posthog.types import FeatureFlag, LegacyFlagMetadata
+from posthog.version import VERSION
+
 
 class TestClient(unittest.TestCase):
     @classmethod
@@ -1213,11 +1214,11 @@ def test_mock_system_context(
 
                 assert context == expected_context
 
-    @mock.patch("posthog.client.decide")    
+    @mock.patch("posthog.client.decide")
     def test_get_decide_returns_normalized_decide_response(self, patch_decide):
         patch_decide.return_value = {
             "featureFlags": {"beta-feature": "random-variant", "alpha-feature": True, "off-feature": False},
-            "featureFlagPayloads": {"beta-feature": "{\"some\": \"data\"}"},
+            "featureFlagPayloads": {"beta-feature": '{"some": "data"}'},
             "errorsWhileComputingFlags": False,
             "requestId": "test-id",
         }
@@ -1226,9 +1227,9 @@ def test_get_decide_returns_normalized_decide_response(self, patch_decide):
         distinct_id = "test_distinct_id"
         groups = {"test_group_type": "test_group_id"}
         person_properties = {"test_property": "test_value"}
-        
+
         response = client.get_decide(distinct_id, groups, person_properties)
-        
+
         assert response == {
             "flags": {
                 "beta-feature": FeatureFlag(
@@ -1237,8 +1238,8 @@ def test_get_decide_returns_normalized_decide_response(self, patch_decide):
                     variant="random-variant",
                     reason=None,
                     metadata=LegacyFlagMetadata(
-                        payload="{\"some\": \"data\"}",
-                    )
+                        payload='{"some": "data"}',
+                    ),
                 ),
                 "alpha-feature": FeatureFlag(
                     key="alpha-feature",
@@ -1247,7 +1248,7 @@ def test_get_decide_returns_normalized_decide_response(self, patch_decide):
                     reason=None,
                     metadata=LegacyFlagMetadata(
                         payload=None,
-                    )
+                    ),
                 ),
                 "off-feature": FeatureFlag(
                     key="off-feature",
@@ -1256,8 +1257,8 @@ def test_get_decide_returns_normalized_decide_response(self, patch_decide):
                     reason=None,
                     metadata=LegacyFlagMetadata(
                         payload=None,
-                    )
-                )
+                    ),
+                ),
             },
             "errorsWhileComputingFlags": False,
             "requestId": "test-id",