Fix up type annotation

Author: haacked

posthog/__init__.py

@@ -4,8 +4,8 @@
 
 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
 
@@ -404,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:
@@ -447,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:
@@ -478,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/test_client.py

@@ -10,8 +10,8 @@
 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):

posthog/test/test_feature_flag.py

@@ -1,12 +1,6 @@
 import unittest
-from posthog.types import (
-    FeatureFlag,
-    FlagMetadata,
-    FlagReason,
-    LegacyFlagMetadata,
-    normalize_decide_response,
-    to_flags_and_payloads,
-)
+
+from posthog.types import FeatureFlag, FlagMetadata, FlagReason, LegacyFlagMetadata
 
 
 class TestFeatureFlag(unittest.TestCase):
@@ -65,20 +59,14 @@ def test_feature_flag_from_json_without_metadata(self):
 
     def test_flag_reason_from_json(self):
         # Test with complete data
-        resp = {
-            "code": "user_in_segment",
-            "condition_index": 1,
-            "description": "User is in segment 'beta_users'"
-        }
+        resp = {"code": "user_in_segment", "condition_index": 1, "description": "User is in segment 'beta_users'"}
         reason = FlagReason.from_json(resp)
         self.assertEqual(reason.code, "user_in_segment")
         self.assertEqual(reason.condition_index, 1)
         self.assertEqual(reason.description, "User is in segment 'beta_users'")
 
         # Test with partial data
-        resp = {
-            "code": "user_in_segment"
-        }
+        resp = {"code": "user_in_segment"}
         reason = FlagReason.from_json(resp)
         self.assertEqual(reason.code, "user_in_segment")
         self.assertEqual(reason.condition_index, 0)  # default value
@@ -89,30 +77,23 @@ def test_flag_reason_from_json(self):
 
     def test_flag_metadata_from_json(self):
         # Test with complete data
-        resp = {
-            "id": 123,
-            "payload": {"key": "value"},
-            "version": 1,
-            "description": "Test flag"
-        }
+        resp = {"id": 123, "payload": {"key": "value"}, "version": 1, "description": "Test flag"}
         metadata = FlagMetadata.from_json(resp)
         self.assertEqual(metadata.id, 123)
         self.assertEqual(metadata.payload, {"key": "value"})
         self.assertEqual(metadata.version, 1)
         self.assertEqual(metadata.description, "Test flag")
 
         # Test with partial data
-        resp = {
-            "id": 123
-        }
+        resp = {"id": 123}
         metadata = FlagMetadata.from_json(resp)
         self.assertEqual(metadata.id, 123)
         self.assertIsNone(metadata.payload)
         self.assertEqual(metadata.version, 0)  # default value
         self.assertEqual(metadata.description, "")  # default value
 
         # Test with None
-        self.assertIsNone(FlagMetadata.from_json(None))
+        self.assertIsInstance(FlagMetadata.from_json(None), LegacyFlagMetadata)
 
     def test_feature_flag_from_json_complete(self):
         # Test with complete data
@@ -123,14 +104,9 @@ def test_feature_flag_from_json_complete(self):
             "reason": {
                 "code": "user_in_segment",
                 "condition_index": 1,
-                "description": "User is in segment 'beta_users'"
+                "description": "User is in segment 'beta_users'",
             },
-            "metadata": {
-                "id": 123,
-                "payload": {"key": "value"},
-                "version": 1,
-                "description": "Test flag"
-            }
+            "metadata": {"id": 123, "payload": {"key": "value"}, "version": 1, "description": "Test flag"},
         }
         flag = FeatureFlag.from_json(resp)
         self.assertEqual(flag.key, "test-flag")
@@ -144,10 +120,7 @@ def test_feature_flag_from_json_complete(self):
 
     def test_feature_flag_from_json_minimal_data(self):
         # Test with minimal data
-        resp = {
-            "key": "test-flag",
-            "enabled": False
-        }
+        resp = {"key": "test-flag", "enabled": False}
         flag = FeatureFlag.from_json(resp)
         self.assertEqual(flag.key, "test-flag")
         self.assertFalse(flag.enabled)
@@ -158,13 +131,7 @@ def test_feature_flag_from_json_minimal_data(self):
 
     def test_feature_flag_from_json_with_reason(self):
         # Test with reason but no metadata
-        resp = {
-            "key": "test-flag",
-            "enabled": True,
-            "reason": {
-                "code": "user_in_segment"
-            }
-        }
+        resp = {"key": "test-flag", "enabled": True, "reason": {"code": "user_in_segment"}}
         flag = FeatureFlag.from_json(resp)
         self.assertEqual(flag.key, "test-flag")
         self.assertTrue(flag.enabled)

posthog/test/test_request.py

@@ -1,22 +1,13 @@
 import json
 import unittest
 from datetime import date, datetime
-from unittest.mock import patch, MagicMock
-from parameterized import parameterized
 
 import mock
 import pytest
 import requests
 
-from posthog.request import (
-    DatetimeSerializer,
-    QuotaLimitError,
-    batch_post,
-    decide,
-    determine_server_host,
-)
+from posthog.request import DatetimeSerializer, QuotaLimitError, batch_post, decide, determine_server_host
 from posthog.test.test_utils import TEST_API_KEY
-from posthog.types import FeatureFlag, FlagMetadata, FlagReason, LegacyFlagMetadata
 
 
 class TestRequests(unittest.TestCase):

posthog/test/test_types.py

@@ -1,5 +1,7 @@
 import unittest
+
 from parameterized import parameterized
+
 from posthog.types import (
     FeatureFlag,
     FlagMetadata,