Module functions to also return the same as its Client equivalent (#111)

Co-authored-by: Neil Kakkar <neilkakkar@gmail.com>

CHANGELOG.md

@@ -1,3 +1,16 @@
+## 3.3.2 - 2024-01-19
+
+1. Return success/failure with all capture calls from module functions
+
+
+## 3.3.1 - 2024-01-10
+
+1. Make sure we don't override any existing feature flag properties when adding locally evaluated feature flag properties.
+
+## 3.3.0 - 2024-01-09
+
+1. When local evaluation is enabled, we automatically add flag information to all events sent to PostHog, whenever possible. This makes it easier to use these events in experiments.
+
 ## 3.2.0 - 2024-01-09
 
 1. Numeric property handling for feature flags now does the expected: When passed in a number, we do a numeric comparison. When passed in a string, we do a string comparison. Previously, we always did a string comparison.

posthog/__init__.py

@@ -1,5 +1,5 @@
 import datetime  # noqa: F401
-from typing import Callable, Dict, Optional  # noqa: F401
+from typing import Callable, Dict, Optional, Tuple  # noqa: F401
 
 from posthog.client import Client
 from posthog.version import VERSION
@@ -33,7 +33,7 @@ def capture(
     send_feature_flags=False,
     disable_geoip=None,  # type: Optional[bool]
 ):
-    # type: (...) -> None
+    # type: (...) -> Tuple[bool, dict]
     """
     Capture allows you to capture anything a user does within your system, which you can later use in PostHog to find patterns in usage, work out which features to improve or where people are giving up.
 
@@ -54,7 +54,7 @@ def capture(
     posthog.capture('distinct id', 'purchase', groups={'company': 'id:5'})
     ```
     """
-    _proxy(
+    return _proxy(
         "capture",
         distinct_id=distinct_id,
         event=event,
@@ -76,7 +76,7 @@ def identify(
     uuid=None,  # type: Optional[str]
     disable_geoip=None,  # type: Optional[bool]
 ):
-    # type: (...) -> None
+    # type: (...) -> Tuple[bool, dict]
     """
     Identify lets you add metadata on your users so you can more easily identify who they are in PostHog, and even do things like segment users by these properties.
 
@@ -92,7 +92,7 @@ def identify(
     })
     ```
     """
-    _proxy(
+    return _proxy(
         "identify",
         distinct_id=distinct_id,
         properties=properties,
@@ -111,7 +111,7 @@ def set(
     uuid=None,  # type: Optional[str]
     disable_geoip=None,  # type: Optional[bool]
 ):
-    # type: (...) -> None
+    # type: (...) -> Tuple[bool, dict]
     """
     Set properties on a user record.
     This will overwrite previous people property values, just like `identify`.
@@ -127,7 +127,7 @@ def set(
      })
      ```
     """
-    _proxy(
+    return _proxy(
         "set",
         distinct_id=distinct_id,
         properties=properties,
@@ -146,7 +146,7 @@ def set_once(
     uuid=None,  # type: Optional[str]
     disable_geoip=None,  # type: Optional[bool]
 ):
-    # type: (...) -> None
+    # type: (...) -> Tuple[bool, dict]
     """
     Set properties on a user record, only if they do not yet exist.
     This will not overwrite previous people property values, unlike `identify`.
@@ -162,7 +162,7 @@ def set_once(
      })
      ```
     """
-    _proxy(
+    return _proxy(
         "set_once",
         distinct_id=distinct_id,
         properties=properties,
@@ -182,7 +182,7 @@ def group_identify(
     uuid=None,  # type: Optional[str]
     disable_geoip=None,  # type: Optional[bool]
 ):
-    # type: (...) -> None
+    # type: (...) -> Tuple[bool, dict]
     """
     Set properties on a group
 
@@ -198,7 +198,7 @@ def group_identify(
      })
      ```
     """
-    _proxy(
+    return _proxy(
         "group_identify",
         group_type=group_type,
         group_key=group_key,
@@ -218,7 +218,7 @@ def alias(
     uuid=None,  # type: Optional[str]
     disable_geoip=None,  # type: Optional[bool]
 ):
-    # type: (...) -> None
+    # type: (...) -> Tuple[bool, dict]
     """
     To marry up whatever a user does before they sign up or log in with what they do after you need to make an alias call. This will allow you to answer questions like "Which marketing channels leads to users churning after a month?" or "What do users do on our website before signing up?"
 
@@ -235,7 +235,7 @@ def alias(
     posthog.alias('anonymous session id', 'distinct id')
     ```
     """
-    _proxy(
+    return _proxy(
         "alias",
         previous_id=previous_id,
         distinct_id=distinct_id,

posthog/client.py

@@ -137,16 +137,6 @@ class Client(object):
         resp_data = self.get_decide(distinct_id, groups, person_properties, group_properties, disable_geoip)
         return resp_data["featureFlags"]
 
-    def _get_active_feature_variants(
-        self, distinct_id, groups=None, person_properties=None, group_properties=None, disable_geoip=None
-    ):
-        feature_variants = self.get_feature_variants(
-            distinct_id, groups, person_properties, group_properties, disable_geoip
-        )
-        return {
-            k: v for (k, v) in feature_variants.items() if v is not False
-        }  # explicitly test for false to account for values that may seem falsy (ex: 0)
-
     def get_feature_payloads(
         self, distinct_id, groups=None, person_properties=None, group_properties=None, disable_geoip=None
     ):
@@ -206,15 +196,29 @@ class Client(object):
             require("groups", groups, dict)
             msg["properties"]["$groups"] = groups
 
+        extra_properties = {}
+        feature_variants = {}
         if send_feature_flags:
             try:
-                feature_variants = self._get_active_feature_variants(distinct_id, groups, disable_geoip=disable_geoip)
+                feature_variants = self.get_feature_variants(distinct_id, groups, disable_geoip=disable_geoip)
             except Exception as e:
                 self.log.exception(f"[FEATURE FLAGS] Unable to get feature variants: {e}")
-            else:
-                for feature, variant in feature_variants.items():
-                    msg["properties"]["$feature/{}".format(feature)] = variant
-                msg["properties"]["$active_feature_flags"] = list(feature_variants.keys())
+
+        elif self.feature_flags:
+            # Local evaluation is enabled, flags are loaded, so try and get all flags we can without going to the server
+            feature_variants = self.get_all_flags(
+                distinct_id, groups=(groups or {}), disable_geoip=disable_geoip, only_evaluate_locally=True
+            )
+
+        for feature, variant in feature_variants.items():
+            extra_properties[f"$feature/{feature}"] = variant
+
+        active_feature_flags = [key for (key, value) in feature_variants.items() if value is not False]
+        if active_feature_flags:
+            extra_properties["$active_feature_flags"] = active_feature_flags
+
+        if extra_properties:
+            msg["properties"] = {**extra_properties, **msg["properties"]}
 
         return self._enqueue(msg, disable_geoip)
 
@@ -538,6 +542,10 @@ class Client(object):
         if self.disabled:
             return None
 
+        person_properties, group_properties = self._add_local_person_and_group_properties(
+            distinct_id, groups, person_properties, group_properties
+        )
+
         if self.feature_flags is None and self.personal_api_key:
             self.load_feature_flags()
         response = None
@@ -685,6 +693,10 @@ class Client(object):
         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(
             distinct_id, groups=groups, person_properties=person_properties, group_properties=group_properties
         )
@@ -743,6 +755,19 @@ class Client(object):
     def feature_flag_definitions(self):
         return self.feature_flags
 
+    def _add_local_person_and_group_properties(self, distinct_id, groups, person_properties, group_properties):
+        all_person_properties = {"$current_distinct_id": distinct_id, **(person_properties or {})}
+
+        all_group_properties = {}
+        if groups:
+            for group_name in groups:
+                all_group_properties[group_name] = {
+                    "$group_key": groups[group_name],
+                    **(group_properties.get(group_name) or {}),
+                }
+
+        return all_person_properties, all_group_properties
+
 
 def require(name, field, data_type):
     """Require that the named `field` has the right `data_type`"""

posthog/test/test_client.py

@@ -106,24 +106,187 @@ class TestClient(unittest.TestCase):
         self.assertEqual(patch_decide.call_count, 1)
 
     @mock.patch("posthog.client.decide")
-    def test_get_active_feature_flags(self, patch_decide):
-        patch_decide.return_value = {
-            "featureFlags": {"beta-feature": "random-variant", "alpha-feature": True, "off-feature": False}
-        }
-
+    def test_basic_capture_with_locally_evaluated_feature_flags(self, patch_decide):
+        patch_decide.return_value = {"featureFlags": {"beta-feature": "random-variant"}}
         client = Client(FAKE_TEST_API_KEY, on_error=self.set_fail, personal_api_key=FAKE_TEST_API_KEY)
-        variants = client._get_active_feature_variants("some_id", None, None, None, False)
-        self.assertEqual(variants, {"beta-feature": "random-variant", "alpha-feature": True})
-        patch_decide.assert_called_with(
-            "random_key",
-            None,
-            timeout=10,
-            distinct_id="some_id",
-            groups={},
-            person_properties=None,
-            group_properties=None,
-            disable_geoip=False,
+
+        multivariate_flag = {
+            "id": 1,
+            "name": "Beta Feature",
+            "key": "beta-feature-local",
+            "is_simple_flag": False,
+            "active": True,
+            "rollout_percentage": 100,
+            "filters": {
+                "groups": [
+                    {
+                        "properties": [
+                            {"key": "email", "type": "person", "value": "test@posthog.com", "operator": "exact"}
+                        ],
+                        "rollout_percentage": 100,
+                    },
+                    {
+                        "rollout_percentage": 50,
+                    },
+                ],
+                "multivariate": {
+                    "variants": [
+                        {"key": "first-variant", "name": "First Variant", "rollout_percentage": 50},
+                        {"key": "second-variant", "name": "Second Variant", "rollout_percentage": 25},
+                        {"key": "third-variant", "name": "Third Variant", "rollout_percentage": 25},
+                    ]
+                },
+                "payloads": {"first-variant": "some-payload", "third-variant": {"a": "json"}},
+            },
+        }
+        basic_flag = {
+            "id": 1,
+            "name": "Beta Feature",
+            "key": "person-flag",
+            "is_simple_flag": True,
+            "active": True,
+            "filters": {
+                "groups": [
+                    {
+                        "properties": [
+                            {
+                                "key": "region",
+                                "operator": "exact",
+                                "value": ["USA"],
+                                "type": "person",
+                            }
+                        ],
+                        "rollout_percentage": 100,
+                    }
+                ],
+                "payloads": {"true": 300},
+            },
+        }
+        false_flag = {
+            "id": 1,
+            "name": "Beta Feature",
+            "key": "false-flag",
+            "is_simple_flag": True,
+            "active": True,
+            "filters": {
+                "groups": [
+                    {
+                        "properties": [],
+                        "rollout_percentage": 0,
+                    }
+                ],
+                "payloads": {"true": 300},
+            },
+        }
+        client.feature_flags = [multivariate_flag, basic_flag, false_flag]
+
+        success, msg = client.capture("distinct_id", "python test event")
+        client.flush()
+        self.assertTrue(success)
+        self.assertFalse(self.failed)
+
+        self.assertEqual(msg["event"], "python test event")
+        self.assertTrue(isinstance(msg["timestamp"], str))
+        self.assertIsNone(msg.get("uuid"))
+        self.assertEqual(msg["distinct_id"], "distinct_id")
+        self.assertEqual(msg["properties"]["$lib"], "posthog-python")
+        self.assertEqual(msg["properties"]["$lib_version"], VERSION)
+        self.assertEqual(msg["properties"]["$feature/beta-feature-local"], "third-variant")
+        self.assertEqual(msg["properties"]["$feature/false-flag"], False)
+        self.assertEqual(msg["properties"]["$active_feature_flags"], ["beta-feature-local"])
+        assert "$feature/beta-feature" not in msg["properties"]
+
+        self.assertEqual(patch_decide.call_count, 0)
+
+        # test that flags are not evaluated without local evaluation
+        client.feature_flags = []
+        success, msg = client.capture("distinct_id", "python test event")
+        client.flush()
+        self.assertTrue(success)
+        self.assertFalse(self.failed)
+        assert "$feature/beta-feature" not in msg["properties"]
+        assert "$feature/beta-feature-local" not in msg["properties"]
+        assert "$feature/false-flag" not in msg["properties"]
+        assert "$active_feature_flags" not in msg["properties"]
+
+    @mock.patch("posthog.client.decide")
+    def test_dont_override_capture_with_local_flags(self, patch_decide):
+        patch_decide.return_value = {"featureFlags": {"beta-feature": "random-variant"}}
+        client = Client(FAKE_TEST_API_KEY, on_error=self.set_fail, personal_api_key=FAKE_TEST_API_KEY)
+
+        multivariate_flag = {
+            "id": 1,
+            "name": "Beta Feature",
+            "key": "beta-feature-local",
+            "is_simple_flag": False,
+            "active": True,
+            "rollout_percentage": 100,
+            "filters": {
+                "groups": [
+                    {
+                        "properties": [
+                            {"key": "email", "type": "person", "value": "test@posthog.com", "operator": "exact"}
+                        ],
+                        "rollout_percentage": 100,
+                    },
+                    {
+                        "rollout_percentage": 50,
+                    },
+                ],
+                "multivariate": {
+                    "variants": [
+                        {"key": "first-variant", "name": "First Variant", "rollout_percentage": 50},
+                        {"key": "second-variant", "name": "Second Variant", "rollout_percentage": 25},
+                        {"key": "third-variant", "name": "Third Variant", "rollout_percentage": 25},
+                    ]
+                },
+                "payloads": {"first-variant": "some-payload", "third-variant": {"a": "json"}},
+            },
+        }
+        basic_flag = {
+            "id": 1,
+            "name": "Beta Feature",
+            "key": "person-flag",
+            "is_simple_flag": True,
+            "active": True,
+            "filters": {
+                "groups": [
+                    {
+                        "properties": [
+                            {
+                                "key": "region",
+                                "operator": "exact",
+                                "value": ["USA"],
+                                "type": "person",
+                            }
+                        ],
+                        "rollout_percentage": 100,
+                    }
+                ],
+                "payloads": {"true": 300},
+            },
+        }
+        client.feature_flags = [multivariate_flag, basic_flag]
+
+        success, msg = client.capture(
+            "distinct_id", "python test event", {"$feature/beta-feature-local": "my-custom-variant"}
         )
+        client.flush()
+        self.assertTrue(success)
+        self.assertFalse(self.failed)
+
+        self.assertEqual(msg["event"], "python test event")
+        self.assertTrue(isinstance(msg["timestamp"], str))
+        self.assertIsNone(msg.get("uuid"))
+        self.assertEqual(msg["distinct_id"], "distinct_id")
+        self.assertEqual(msg["properties"]["$lib"], "posthog-python")
+        self.assertEqual(msg["properties"]["$lib_version"], VERSION)
+        self.assertEqual(msg["properties"]["$feature/beta-feature-local"], "my-custom-variant")
+        self.assertEqual(msg["properties"]["$active_feature_flags"], ["beta-feature-local"])
+        assert "$feature/beta-feature" not in msg["properties"]
+        assert "$feature/person-flag" not in msg["properties"]
+
+        self.assertEqual(patch_decide.call_count, 0)
 
     @mock.patch("posthog.client.decide")
     def test_basic_capture_with_feature_flags_returns_active_only(self, patch_decide):
@@ -620,7 +783,7 @@ class TestClient(unittest.TestCase):
             timeout=10,
             distinct_id="some_id",
             groups={},
-            person_properties={},
+            person_properties={"$current_distinct_id": "some_id"},
             group_properties={},
             disable_geoip=True,
         )
@@ -632,7 +795,7 @@ class TestClient(unittest.TestCase):
             timeout=10,
             distinct_id="feature_enabled_distinct_id",
             groups={},
-            person_properties={},
+            person_properties={"$current_distinct_id": "feature_enabled_distinct_id"},
             group_properties={},
             disable_geoip=True,
         )
@@ -644,7 +807,7 @@ class TestClient(unittest.TestCase):
             timeout=10,
             distinct_id="all_flags_payloads_id",
             groups={},
-            person_properties={},
+            person_properties={"$current_distinct_id": "all_flags_payloads_id"},
             group_properties={},
             disable_geoip=False,
         )
@@ -660,3 +823,70 @@ class TestClient(unittest.TestCase):
         client.feature_flags = [{"key": "example", "is_simple_flag": False}]
 
         self.assertFalse(client.feature_enabled("example", "distinct_id"))
+
+    @mock.patch("posthog.client.decide")
+    def test_default_properties_get_added_properly(self, patch_decide):
+        patch_decide.return_value = {
+            "featureFlags": {"beta-feature": "random-variant", "alpha-feature": True, "off-feature": False}
+        }
+        client = Client(FAKE_TEST_API_KEY, on_error=self.set_fail, disable_geoip=False)
+        client.get_feature_flag(
+            "random_key",
+            "some_id",
+            groups={"company": "id:5", "instance": "app.posthog.com"},
+            person_properties={"x1": "y1"},
+            group_properties={"company": {"x": "y"}},
+        )
+        patch_decide.assert_called_with(
+            "random_key",
+            None,
+            timeout=10,
+            distinct_id="some_id",
+            groups={"company": "id:5", "instance": "app.posthog.com"},
+            person_properties={"$current_distinct_id": "some_id", "x1": "y1"},
+            group_properties={
+                "company": {"$group_key": "id:5", "x": "y"},
+                "instance": {"$group_key": "app.posthog.com"},
+            },
+            disable_geoip=False,
+        )
+
+        patch_decide.reset_mock()
+        client.get_feature_flag(
+            "random_key",
+            "some_id",
+            groups={"company": "id:5", "instance": "app.posthog.com"},
+            person_properties={"$current_distinct_id": "override"},
+            group_properties={
+                "company": {
+                    "$group_key": "group_override",
+                }
+            },
+        )
+        patch_decide.assert_called_with(
+            "random_key",
+            None,
+            timeout=10,
+            distinct_id="some_id",
+            groups={"company": "id:5", "instance": "app.posthog.com"},
+            person_properties={"$current_distinct_id": "override"},
+            group_properties={
+                "company": {"$group_key": "group_override"},
+                "instance": {"$group_key": "app.posthog.com"},
+            },
+            disable_geoip=False,
+        )
+
+        patch_decide.reset_mock()
+        # test nones
+        client.get_all_flags_and_payloads("some_id", groups={}, person_properties=None, group_properties=None)
+        patch_decide.assert_called_with(
+            "random_key",
+            None,
+            timeout=10,
+            distinct_id="some_id",
+            groups={},
+            person_properties={"$current_distinct_id": "some_id"},
+            group_properties={},
+            disable_geoip=False,
+        )

posthog/test/test_module.py

@@ -6,6 +6,10 @@ from posthog import Posthog
 class TestModule(unittest.TestCase):
     posthog = None
 
+    def _assert_enqueue_result(self, result):
+        self.assertEqual(type(result[0]), bool)
+        self.assertEqual(type(result[1]), dict)
+
     def failed(self):
         self.failed = True
 
@@ -22,15 +26,18 @@ class TestModule(unittest.TestCase):
         self.assertRaises(Exception, self.posthog.capture)
 
     def test_track(self):
-        self.posthog.capture("distinct_id", "python module event")
+        res = self.posthog.capture("distinct_id", "python module event")
+        self._assert_enqueue_result(res)
         self.posthog.flush()
 
     def test_identify(self):
-        self.posthog.identify("distinct_id", {"email": "user@email.com"})
+        res = self.posthog.identify("distinct_id", {"email": "user@email.com"})
+        self._assert_enqueue_result(res)
         self.posthog.flush()
 
     def test_alias(self):
-        self.posthog.alias("previousId", "distinct_id")
+        res = self.posthog.alias("previousId", "distinct_id")
+        self._assert_enqueue_result(res)
         self.posthog.flush()
 
     def test_page(self):

posthog/version.py

@@ -1,4 +1,4 @@
-VERSION = "3.2.0"
+VERSION = "3.3.2"
 
 if __name__ == "__main__":
     print(VERSION, end="")  # noqa: T201