Annotate types

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
 __version__ = VERSION
 
 """Settings."""
@@ -403,7 +403,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:
     """
     Get feature flag variant for users. Used with experiments.
     Example:
@@ -446,7 +446,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:
     """
     Get all flags for a given user.
     Example:
@@ -477,7 +477,7 @@ def get_feature_flag_payload(
     only_evaluate_locally=False,
     send_feature_flag_events=True,
     disable_geoip=None,  # type: Optional[bool]
-):
+) -> str:
     return _proxy(
         "get_feature_flag_payload",
         key=key,
@@ -519,7 +519,7 @@ def get_all_flags_and_payloads(
     group_properties={},
     only_evaluate_locally=False,
     disable_geoip=None,  # type: Optional[bool]
-):
+) -> FlagsAndPayloads:
     return _proxy(
         "get_all_flags_and_payloads",
         distinct_id=distinct_id,

posthog/client.py

@@ -18,10 +18,10 @@
 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, normalize_decide_response, DecideResponse
+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.types import FlagsAndPayloads, FlagValue, to_values, to_payloads, to_flags_and_payloads, normalize_decide_response
 try:
     import queue
 except ImportError:
@@ -239,26 +239,23 @@ 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]:
         resp_data = self.get_decide(distinct_id, groups, person_properties, group_properties, disable_geoip)
-        return self.to_variants(resp_data)
+        return to_values(resp_data) or {}
 
     def get_feature_payloads(
         self, distinct_id, groups=None, person_properties=None, group_properties=None, disable_geoip=None
-    ):
+    ) -> dict[str, str]:
         resp_data = self.get_decide(distinct_id, groups, person_properties, group_properties, disable_geoip)
-        return self.to_payloads(resp_data)
+        return to_payloads(resp_data) or {}
 
     def get_feature_flags_and_payloads(
         self, distinct_id, groups=None, person_properties=None, group_properties=None, disable_geoip=None
-    ):
-        resp_data = self.get_decide(distinct_id, groups, person_properties, group_properties, disable_geoip)
-        return {
-            "featureFlags": self.to_variants(resp_data),
-            "featureFlagPayloads": self.to_payloads(resp_data),
-        }
+    ) -> FlagsAndPayloads:
+        resp = self.get_decide(distinct_id, groups, person_properties, group_properties, disable_geoip)
+        return to_flags_and_payloads(resp)
 
-    def get_decide(self, distinct_id, groups=None, person_properties=None, group_properties=None, disable_geoip=None):
+    def get_decide(self, distinct_id, groups=None, person_properties=None, group_properties=None, disable_geoip=None) -> DecideResponse:
         require("distinct_id", distinct_id, ID_TYPES)
 
         if disable_geoip is None:
@@ -317,8 +314,8 @@ def capture(
             require("groups", groups, dict)
             msg["properties"]["$groups"] = groups
 
-        extra_properties = {}
-        feature_variants = {}
+        extra_properties: dict[str, Any] = {}
+        feature_variants: dict[str, bool | str] | None = {}
         if send_feature_flags:
             try:
                 feature_variants = self.get_feature_variants(distinct_id, groups, disable_geoip=disable_geoip)
@@ -331,10 +328,10 @@ def capture(
                 distinct_id, groups=(groups or {}), disable_geoip=disable_geoip, only_evaluate_locally=True
             )
 
-        for feature, variant in feature_variants.items():
+        for feature, variant in (feature_variants or {}).items():
             extra_properties[f"$feature/{feature}"] = variant
 
-        active_feature_flags = [key for (key, value) in feature_variants.items() if value is not False]
+        active_feature_flags = [key for (key, value) in (feature_variants or {}).items() if value is not False]
         if active_feature_flags:
             extra_properties["$active_feature_flags"] = active_feature_flags
 
@@ -711,7 +708,7 @@ def _compute_flag_locally(
         person_properties={},
         group_properties={},
         warn_on_unknown_groups=True,
-    ):
+    ) -> FlagValue:
         if feature_flag.get("ensure_experience_continuity", False):
             raise InconclusiveMatchError("Flag has experience continuity enabled")
 
@@ -901,8 +898,13 @@ def get_feature_flag_payload(
                 responses_and_payloads = self.get_feature_flags_and_payloads(
                     distinct_id, groups, person_properties, group_properties, disable_geoip
                 )
-                response = responses_and_payloads["featureFlags"].get(key, None)
-                payload = responses_and_payloads["featureFlagPayloads"].get(str(key), None)
+                featureFlags = responses_and_payloads["featureFlags"]
+                if featureFlags is not None:
+                    response = featureFlags.get(key, None)
+                
+                featureFlagPayloads = responses_and_payloads["featureFlagPayloads"]
+                if featureFlagPayloads is not None:
+                    payload = featureFlagPayloads.get(str(key), None)
             except Exception as e:
                 self.log.exception(f"[FEATURE FLAGS] Unable to get feature flags and payloads: {e}")
 
@@ -949,7 +951,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, match_value):
+    def _compute_payload_locally(self, key: str, match_value: FlagValue) -> str | None:
         payload = None
 
         if self.feature_flags_by_key is None:
@@ -974,7 +976,7 @@ def get_all_flags(
         group_properties={},
         only_evaluate_locally=False,
         disable_geoip=None,
-    ) -> dict[str, bool | str]:
+    ) -> dict[str, bool | str] | None:
         response = self.get_all_flags_and_payloads(
             distinct_id,
             groups=groups,
@@ -984,7 +986,7 @@ def get_all_flags(
             disable_geoip=disable_geoip,
         )
 
-        return self.to_variants(response)
+        return response["featureFlags"]
 
     def get_all_flags_and_payloads(
         self,
@@ -995,46 +997,45 @@ def get_all_flags_and_payloads(
         group_properties={},
         only_evaluate_locally=False,
         disable_geoip=None,
-    ) -> DecideResponse:
+    ) -> FlagsAndPayloads:
         if self.disabled:
             return {"featureFlags": None, "featureFlagPayloads": None}
 
         person_properties, group_properties = self._add_local_person_and_group_properties(
             distinct_id, groups, person_properties, group_properties
         )
 
-        flags, payloads, fallback_to_decide = self._get_all_flags_and_payloads_locally(
+        response, fallback_to_decide = self._get_all_flags_and_payloads_locally(
             distinct_id, groups=groups, person_properties=person_properties, group_properties=group_properties
         )
-        
-        response = normalize_decide_response({"featureFlags": flags, "featureFlagPayloads": payloads})
 
         if fallback_to_decide and not only_evaluate_locally:
             try:
-                flags_and_payloads = self.get_decide(
+                decide_response = self.get_decide(
                     distinct_id,
                     groups=groups,
                     person_properties=person_properties,
                     group_properties=group_properties,
                     disable_geoip=disable_geoip,
                 )
-                response = flags_and_payloads
+                return to_flags_and_payloads(decide_response)
             except Exception as e:
                 self.log.exception(f"[FEATURE FLAGS] Unable to get feature flags and payloads: {e}")
 
         return response
 
+
     def _get_all_flags_and_payloads_locally(
         self, distinct_id, *, groups={}, person_properties={}, group_properties={}, warn_on_unknown_groups=False
-    ):
+    ) -> tuple[FlagsAndPayloads, bool]:
         require("distinct_id", distinct_id, ID_TYPES)
         require("groups", groups, dict)
 
         if self.feature_flags is None and self.personal_api_key:
             self.load_feature_flags()
 
-        flags = {}
-        payloads = {}
+        flags: dict[str, FlagValue] = {}
+        payloads: dict[str, str] = {}
         fallback_to_decide = False
         # If loading in previous line failed
         if self.feature_flags:
@@ -1060,7 +1061,7 @@ def _get_all_flags_and_payloads_locally(
         else:
             fallback_to_decide = True
 
-        return flags, payloads, fallback_to_decide
+        return {"featureFlags": flags, "featureFlagPayloads": payloads}, fallback_to_decide
 
     def feature_flag_definitions(self):
         return self.feature_flags
@@ -1078,18 +1079,6 @@ def _add_local_person_and_group_properties(self, distinct_id, groups, person_pro
 
         return all_person_properties, all_group_properties
 
-    def to_variants(self, response: any) -> dict[str, bool | str]:
-        if "flags" not in response:
-            return None
-
-        return {key: value.get_value() for key, value in response.get("flags", {}).items()}
-
-    def to_payloads(self, response: any) -> dict[str, str]:
-        if "flags" not in response:
-            return None
-
-        return {key: value.metadata.payload for key, value in response.get("flags", {}).items() if value.enabled}
-
 
 def require(name, field, data_type):
     """Require that the named `field` has the right `data_type`"""

posthog/feature_flags.py

@@ -8,6 +8,7 @@
 from dateutil.relativedelta import relativedelta
 
 from posthog import utils
+from posthog.types import FlagValue
 from posthog.utils import convert_to_datetime_aware, is_valid_regex
 
 __LONG_SCALE__ = float(0xFFFFFFFFFFFFFFF)
@@ -25,7 +26,7 @@ class InconclusiveMatchError(Exception):
 # Given the same distinct_id and key, it'll always return the same float. These floats are
 # uniformly distributed between 0 and 1, so if we want to show this feature to 20% of traffic
 # we can do _hash(key, distinct_id) < 0.2
-def _hash(key, distinct_id, salt=""):
+def _hash(key: str, distinct_id: str, salt: str = "") -> float:
     hash_key = f"{key}.{distinct_id}{salt}"
     hash_val = int(hashlib.sha1(hash_key.encode("utf-8")).hexdigest()[:15], 16)
     return hash_val / __LONG_SCALE__
@@ -50,7 +51,7 @@ def variant_lookup_table(feature_flag):
     return lookup_table
 
 
-def match_feature_flag_properties(flag, distinct_id, properties, cohort_properties=None):
+def match_feature_flag_properties(flag, distinct_id, properties, cohort_properties=None) -> FlagValue:
     flag_conditions = (flag.get("filters") or {}).get("groups") or []
     is_inconclusive = False
     cohort_properties = cohort_properties or {}
@@ -87,7 +88,7 @@ def match_feature_flag_properties(flag, distinct_id, properties, cohort_properti
     return False
 
 
-def is_condition_match(feature_flag, distinct_id, condition, properties, cohort_properties):
+def is_condition_match(feature_flag, distinct_id, condition, properties, cohort_properties) -> bool:
     rollout_percentage = condition.get("rollout_percentage")
     if len(condition.get("properties") or []) > 0:
         for prop in condition.get("properties"):

posthog/request.py

@@ -8,7 +8,7 @@
 import requests
 from dateutil.tz import tzutc
 
-from posthog.types import DecideResponse, FeatureFlag
+from posthog.types import DecideResponse, FeatureFlag, FlagValue
 from posthog.utils import remove_trailing_slash
 from posthog.version import VERSION
 
@@ -97,20 +97,6 @@ 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_flags.py

@@ -642,9 +642,8 @@ def test_get_all_flags_and_payloads_with_fallback(self, patch_decide, patch_capt
             },
         ]
         # beta-feature value overridden by /decide
-        flags_and_payloads = client.get_all_flags_and_payloads("distinct_id")
         self.assertEqual(
-            client.to_payloads(flags_and_payloads),
+            client.get_all_flags_and_payloads("distinct_id")["featureFlagPayloads"],
             {
                 "beta-feature": 100,
                 "beta-feature2": 300,
@@ -676,9 +675,8 @@ def test_get_all_flags_and_payloads_with_fallback_empty_local_flags(self, patch_
         client = self.client
         client.feature_flags = []
         # beta-feature value overridden by /decide
-        flags_and_payloads = client.get_all_flags_and_payloads("distinct_id")
         self.assertEqual(
-            client.to_payloads(flags_and_payloads),
+            client.get_all_flags_and_payloads("distinct_id")["featureFlagPayloads"],
             {"beta-feature": 100, "beta-feature2": 300},
         )
         self.assertEqual(patch_decide.call_count, 1)
@@ -768,9 +766,9 @@ def test_get_all_flags_and_payloads_with_no_fallback(self, patch_decide, patch_c
             basic_flag,
             disabled_flag,
         ]
-        all_flags_and_payloads = client.get_all_flags_and_payloads("distinct_id")
         self.assertEqual(
-            client.to_payloads(all_flags_and_payloads), {"beta-feature": "new"}
+            client.get_all_flags_and_payloads("distinct_id")["featureFlagPayloads"],
+            {"beta-feature": "new"}
         )
         # decide not called because this can be evaluated locally
         self.assertEqual(patch_decide.call_count, 0)
@@ -900,9 +898,8 @@ def test_get_all_flags_and_payloads_with_fallback_but_only_local_evaluation_set(
             flag_3,
         ]
         # beta-feature2 has no value
-        flags_and_payloads = client.get_all_flags_and_payloads("distinct_id", only_evaluate_locally=True)
         self.assertEqual(
-            client.to_payloads(flags_and_payloads),
+            client.get_all_flags_and_payloads("distinct_id", only_evaluate_locally=True)["featureFlagPayloads"],
             {"beta-feature": "some-payload"},
         )
         self.assertEqual(patch_decide.call_count, 0)