Fix up type annotation

Author: haacked

posthog/__init__.py

@@ -5,7 +5,7 @@
 from posthog.client import Client
 from posthog.exception_capture import Integrations  # noqa: F401
 from posthog.version import VERSION
-from posthog.types import FlagsAndPayloads
+from posthog.types import FlagsAndPayloads, FeatureFlag
 
 __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:
+) -> FeatureFlag | None:
     """
     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:
+) -> dict[str, FeatureFlag] | None:
     """
     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:
+) -> str | None:
     return _proxy(
         "get_feature_flag_payload",
         key=key,

posthog/test/test_feature_flag.py

@@ -65,20 +65,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,22 +83,15 @@ 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)
@@ -123,14 +110,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 +126,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 +137,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/types.py

@@ -20,7 +20,7 @@ class FlagReason:
     description: str
 
     @classmethod
-    def from_json(cls, resp: Any) -> "FlagReason":
+    def from_json(cls, resp: Any) -> "FlagReason" | None:
         if not resp:
             return None
         return cls(
@@ -33,12 +33,12 @@ def from_json(cls, resp: Any) -> "FlagReason":
 @dataclass(frozen=True)
 class FlagMetadata:
     id: int
-    payload: Any
+    payload: str | None
     version: int
     description: str
 
     @classmethod
-    def from_json(cls, resp: Any) -> "FlagMetadata":
+    def from_json(cls, resp: Any) -> "FlagMetadata" | None:
         if not resp:
             return None
         return cls(