add functionality+tests for set and set_once

Author: neilkakkar

README.md

@@ -14,5 +14,3 @@ Specifically, the [Python integration](https://posthog.com/docs/integrations/pyt
 2. Run `source env/bin/activate` (activates the virtual environment)
 3. Run `python3 -m pip install -e ".[test]"` (installs the package in develop mode, along with test dependencies)
 4. Run `make test`
-
-To run tests specific to a file, use 

example.py

@@ -6,7 +6,7 @@
 import posthog
 
 # You can find this key on the /setup page in PostHog
-posthog.api_key = ""
+posthog.api_key = "LXP6nQXvo-2TCqGVrWvPah8uJIyVykoMmhnEkEBi5PA"
 posthog.personal_api_key = ""
 
 # Where you host PostHog, with no trailing /.
@@ -19,19 +19,29 @@
 print(posthog.feature_enabled("beta-feature", "distinct_id"))
 
 print("sleeping")
-time.sleep(45)
+time.sleep(5)
 
 print(posthog.feature_enabled("beta-feature", "distinct_id"))
 
 # # Alias a previous distinct id with a new one
+
 posthog.alias("distinct_id", "new_distinct_id")
 
+posthog.capture("new_distinct_id", "event2", {"property1": "value", "property2": "value"})
+
 # # Add properties to the person
-posthog.identify("distinct_id", {"email": "[email protected]"})
+posthog.identify("new_distinct_id", {"email": "[email protected]"})
+
+# properties set only once to the person
+posthog.set_once("new_distinct_id", {"self_serve_signup": True})
 
+time.sleep(3)
 
-# properties only once to the person
-posthog.set_once("distinct_id", {"self_serve_signup": True})
 posthog.set_once(
-    "distinct_id", {"self_serve_signup": False}
+    "new_distinct_id", {"self_serve_signup": False}
 )  # this will not change the property (because it was already set)
+
+posthog.set("new_distinct_id", {"current_browser": "Chrome"})
+posthog.set("new_distinct_id", {"current_browser": "Firefox"})
+
+posthog.shutdown()

posthog/__init__.py

@@ -87,6 +87,38 @@ def identify(
     )
 
 
+def set(
+    distinct_id,  # type: str,
+    properties=None,  # type: Optional[Dict]
+    context=None,  # type: Optional[Dict]
+    timestamp=None,  # type: Optional[datetime.datetime]
+    message_id=None,  # type: Optional[str]
+):
+    # type: (...) -> None
+    """
+    Set properties on a user record.
+    This will overwrite previous people property values, just like `identify`.
+
+     A `set` call requires
+     - `distinct id` which uniquely identifies your user
+     - `properties` with a dict with any key: value pairs
+
+     For example:
+     ```python
+     posthog.set('distinct id', {
+         'current_browser': 'Chrome',
+     })
+     ```
+    """
+    _proxy(
+        "set",
+        distinct_id=distinct_id,
+        properties=properties,
+        context=context,
+        timestamp=timestamp,
+        message_id=message_id,
+    )
+
 def set_once(
     distinct_id,  # type: str,
     properties=None,  # type: Optional[Dict]

posthog/client.py

@@ -138,7 +138,7 @@ def capture(self, distinct_id=None, event=None, properties=None, context=None, t
 
         return self._enqueue(msg)
 
-    def set_once(self, distinct_id=None, properties=None, context=None, timestamp=None, message_id=None):
+    def set(self, distinct_id=None, properties=None, context=None, timestamp=None, message_id=None):
         properties = properties or {}
         context = context or {}
         require("distinct_id", distinct_id, ID_TYPES)
@@ -149,6 +149,23 @@ def set_once(self, distinct_id=None, properties=None, context=None, timestamp=No
             "context": context,
             "distinct_id": distinct_id,
             "$set": properties,
+            "event": "$set",
+            "messageId": message_id,
+        }
+
+        return self._enqueue(msg)
+
+    def set_once(self, distinct_id=None, properties=None, context=None, timestamp=None, message_id=None):
+        properties = properties or {}
+        context = context or {}
+        require("distinct_id", distinct_id, ID_TYPES)
+        require("properties", properties, dict)
+
+        msg = {
+            "timestamp": timestamp,
+            "context": context,
+            "distinct_id": distinct_id,
+            "$set_once": properties,
             "event": "$set_once",
             "messageId": message_id,
         }

posthog/test/test_client.py

@@ -105,6 +105,66 @@ def test_advanced_identify(self):
         self.assertEqual(msg["messageId"], "messageId")
         self.assertEqual(msg["distinct_id"], "distinct_id")
 
+
+    def test_basic_set(self):
+        client = self.client
+        success, msg = client.set("distinct_id", {"trait": "value"})
+        client.flush()
+        self.assertTrue(success)
+        self.assertFalse(self.failed)
+
+        self.assertEqual(msg["$set"]["trait"], "value")
+        self.assertTrue(isinstance(msg["timestamp"], str))
+        self.assertTrue(isinstance(msg["messageId"], str))
+        self.assertEqual(msg["distinct_id"], "distinct_id")
+
+    def test_advanced_set(self):
+        client = self.client
+        success, msg = client.set(
+            "distinct_id", {"trait": "value"}, {"ip": "192.168.0.1"}, datetime(2014, 9, 3), "messageId"
+        )
+
+        self.assertTrue(success)
+
+        self.assertEqual(msg["timestamp"], "2014-09-03T00:00:00+00:00")
+        self.assertEqual(msg["context"]["ip"], "192.168.0.1")
+        self.assertEqual(msg["$set"]["trait"], "value")
+        self.assertEqual(msg["properties"]["$lib"], "posthog-python")
+        self.assertEqual(msg["properties"]["$lib_version"], VERSION)
+        self.assertTrue(isinstance(msg["timestamp"], str))
+        self.assertEqual(msg["messageId"], "messageId")
+        self.assertEqual(msg["distinct_id"], "distinct_id")
+
+
+    def test_basic_set_once(self):
+        client = self.client
+        success, msg = client.set_once("distinct_id", {"trait": "value"})
+        client.flush()
+        self.assertTrue(success)
+        self.assertFalse(self.failed)
+
+        self.assertEqual(msg["$set_once"]["trait"], "value")
+        self.assertTrue(isinstance(msg["timestamp"], str))
+        self.assertTrue(isinstance(msg["messageId"], str))
+        self.assertEqual(msg["distinct_id"], "distinct_id")
+
+    def test_advanced_set_once(self):
+        client = self.client
+        success, msg = client.set_once(
+            "distinct_id", {"trait": "value"}, {"ip": "192.168.0.1"}, datetime(2014, 9, 3), "messageId"
+        )
+
+        self.assertTrue(success)
+
+        self.assertEqual(msg["timestamp"], "2014-09-03T00:00:00+00:00")
+        self.assertEqual(msg["context"]["ip"], "192.168.0.1")
+        self.assertEqual(msg["$set_once"]["trait"], "value")
+        self.assertEqual(msg["properties"]["$lib"], "posthog-python")
+        self.assertEqual(msg["properties"]["$lib_version"], VERSION)
+        self.assertTrue(isinstance(msg["timestamp"], str))
+        self.assertEqual(msg["messageId"], "messageId")
+        self.assertEqual(msg["distinct_id"], "distinct_id")
+
     def test_basic_alias(self):
         client = self.client
         success, msg = client.alias("previousId", "distinct_id")
@@ -125,7 +185,7 @@ def test_basic_page(self):
 
     def test_basic_page_distinct_uuid(self):
         client = self.client
-        distinct_id = uuid4()
+        distinct_id = str(uuid4())
         success, msg = client.page(distinct_id, url="https://posthog.com/contact")
         self.assertFalse(self.failed)
         client.flush()