Flesh out Decide response types

Author: haacked

posthog/request.py

@@ -3,11 +3,12 @@
 from datetime import date, datetime
 from gzip import GzipFile
 from io import BytesIO
-from typing import Any, Optional, Union
+from typing import Any, Optional, Union, cast
 
 import requests
 from dateutil.tz import tzutc
 
+from posthog.types import DecideResponse, FeatureFlag
 from posthog.utils import remove_trailing_slash
 from posthog.version import VERSION
 
@@ -96,6 +97,20 @@ def decide(api_key: str, host: Optional[str] = None, gzip: bool = False, timeout
     res = post(api_key, host, "/decide/?v=3", gzip, timeout, **kwargs)
     return _process_response(res, success_message="Feature flags decided successfully")
 
+def normalize_decide_response(resp: any) -> DecideResponse:
+    if "requestId" not in resp:
+        resp["requestId"] = None
+    if "flags" not in resp:
+        featureFlags = resp.get("featureFlags", {})
+        featureFlagPayloads = resp.get("featureFlagPayloads", {})
+        resp.pop("featureFlags", None)
+        resp.pop("featureFlagPayloads", None)
+        # look at each key in featureFlags and create a FeatureFlag object
+        flags = {}
+        for key, value in featureFlags.items():
+            flags[key] = FeatureFlag.from_value_and_payload(key, value, featureFlagPayloads.get(key, None))
+        resp["flags"] = flags
+    return cast(DecideResponse, resp)
 
 def remote_config(personal_api_key: str, host: Optional[str] = None, key: str = "", timeout: int = 15) -> Any:
     """Get remote config flag value from remote_config API endpoint"""

posthog/test/test_feature_flag.py

@@ -0,0 +1,33 @@
+import pytest
+from posthog.types import FeatureFlag
+
+
+def test_from_value_and_payload_boolean_true_value():
+    # Test with boolean value
+    flag = FeatureFlag.from_value_and_payload(key="my-flag", value=True, payload='{"some": "data"}')
+
+    assert flag.key == "my-flag"
+    assert flag.enabled is True
+    assert flag.variant is None
+    assert flag.metadata.payload == '{"some": "data"}'
+
+
+def test_from_value_and_payload_boolean_false_value():
+    # Test with False value
+    flag = FeatureFlag.from_value_and_payload(key="my-flag", value=False, payload=None)
+
+    assert flag.key == "my-flag"
+    assert flag.enabled is False
+    assert flag.variant is None
+    assert flag.metadata.payload is None
+
+
+def test_from_value_and_payload_string_variant():
+    flag = FeatureFlag.from_value_and_payload(
+        key="my-flag", value="test-variant", payload='{"variant": "test-variant"}'
+    )
+
+    assert flag.key == "my-flag"
+    assert flag.enabled is True  # String values should make enabled True
+    assert flag.variant == "test-variant"
+    assert flag.metadata.payload == '{"variant": "test-variant"}'

posthog/test/test_request.py

@@ -1,13 +1,23 @@
 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,
+    normalize_decide_response,
+)
 from posthog.test.test_utils import TEST_API_KEY
+from posthog.types import FeatureFlag, FlagMetadata, FlagReason, LegacyFlagMetadata
 
 
 class TestRequests(unittest.TestCase):
@@ -75,6 +85,90 @@ def test_normal_decide_response(self):
             response = decide("fake_key", "fake_host")
             self.assertEqual(response["featureFlags"], {"flag1": True})
 
+    @parameterized.expand([(True,), (False,)])
+    def test_normalize_decide_response_v4(self, has_errors: bool):
+        resp = {
+            "flags": {
+                "my-flag": FeatureFlag(
+                    key="my-flag",
+                    enabled=True,
+                    variant="test-variant",
+                    reason=FlagReason(
+                        code="matched_condition", condition_index=0, description="Matched condition set 1"
+                    ),
+                    metadata=FlagMetadata(id=1, payload='{"some": "json"}', version=2, description="test-description"),
+                )
+            },
+            "errorsWhileComputingFlags": has_errors,
+            "requestId": "test-id",
+        }
+
+        result = normalize_decide_response(resp)
+        
+        flag = result["flags"]["my-flag"]
+        self.assertEqual(flag.key, "my-flag")
+        self.assertTrue(flag.enabled)
+        self.assertEqual(flag.variant, "test-variant")
+        self.assertEqual(flag.get_value(), "test-variant")
+        self.assertEqual(
+            flag.reason, FlagReason(code="matched_condition", condition_index=0, description="Matched condition set 1")
+        )
+        self.assertEqual(
+            flag.metadata, FlagMetadata(id=1, payload='{"some": "json"}', version=2, description="test-description")
+        )
+        self.assertEqual(result["errorsWhileComputingFlags"], has_errors)
+        self.assertEqual(result["requestId"], "test-id")
+
+    def test_normalize_decide_response_legacy(self):
+        # Test legacy response format with "featureFlags" and "featureFlagPayloads"
+        resp = {
+            "featureFlags": {"my-flag": "test-variant"},
+            "featureFlagPayloads": {"my-flag": "{\"some\": \"json-payload\"}"},
+            "errorsWhileComputingFlags": False,
+            "requestId": "test-id",
+        }
+
+        result = normalize_decide_response(resp)
+
+        flag = result["flags"]["my-flag"]
+        self.assertEqual(flag.key, "my-flag")
+        self.assertTrue(flag.enabled)
+        self.assertEqual(flag.variant, "test-variant")
+        self.assertEqual(flag.get_value(), "test-variant")
+        self.assertIsNone(flag.reason)
+        self.assertEqual(
+            flag.metadata, LegacyFlagMetadata(payload='{"some": "json-payload"}')
+        )
+        self.assertFalse(result["errorsWhileComputingFlags"])
+        self.assertEqual(result["requestId"], "test-id")
+        # Verify legacy fields are removed
+        self.assertNotIn("featureFlags", result)
+        self.assertNotIn("featureFlagPayloads", result)
+
+    def test_normalize_decide_response_boolean_flag(self):
+        # Test legacy response with boolean flag
+        resp = {
+            "featureFlags": {"my-flag": True},
+            "errorsWhileComputingFlags": False
+        }
+
+        result = normalize_decide_response(resp)
+
+        self.assertIn("requestId", result)
+        self.assertIsNone(result["requestId"])
+
+        flag = result["flags"]["my-flag"]
+        self.assertEqual(flag.key, "my-flag")
+        self.assertTrue(flag.enabled)
+        self.assertIsNone(flag.variant)
+        self.assertIsNone(flag.reason)
+        self.assertEqual(
+            flag.metadata, LegacyFlagMetadata(payload=None)
+        )
+        self.assertFalse(result["errorsWhileComputingFlags"])
+        self.assertNotIn("featureFlags", result)
+        self.assertNotIn("featureFlagPayloads", result)
+
 
 @pytest.mark.parametrize(
     "host, expected",

posthog/types.py

@@ -0,0 +1,54 @@
+from dataclasses import dataclass
+from typing import Dict, List, TypedDict, Any
+
+
+@dataclass(frozen=True)
+class FlagReason(TypedDict):
+    code: str
+    condition_index: int
+    description: str
+
+
+@dataclass(frozen=True)
+class FlagMetadata:
+    id: int
+    payload: Any
+    version: int
+    description: str
+
+
+@dataclass(frozen=True)
+class LegacyFlagMetadata:
+    payload: Any
+
+
+@dataclass(frozen=True)
+class FeatureFlag:
+    key: str
+    enabled: bool
+    variant: str | None
+    reason: FlagReason | None
+    metadata: FlagMetadata | LegacyFlagMetadata
+
+    def get_value(self):
+        assert self.variant is None or self.enabled
+        return self.variant or self.enabled
+
+    @classmethod
+    def from_value_and_payload(cls, key: str, value: any, payload: any) -> "FeatureFlag":
+        return cls(
+            key=key,
+            enabled=True if isinstance(value, str) else value,
+            variant=value if isinstance(value, str) else None,
+            reason=None,
+            metadata=LegacyFlagMetadata(
+                payload=payload,
+            ),
+        )
+
+
+class DecideResponse(TypedDict, total=False):
+    flags: Dict[str, FeatureFlag]
+    errorsWhileComputingFlags: bool
+    requestId: str
+    quotaLimit: List[str] | None