fix: pass before_send through module-level setup (#472)

Author: sipa-echo-zaoa

posthog/__init__.py

@@ -56,6 +56,7 @@
     SocketOptions as SocketOptions,
 )
 from posthog.types import (
+    BeforeSendCallback as BeforeSendCallback,
     FeatureFlag,
     FlagsAndPayloads,
 )
@@ -250,6 +251,7 @@ def get_tags() -> Dict[str, Any]:
 project_root = None  # type: Optional[str]
 # Used for our AI observability feature to not capture any prompt or output just usage + metadata
 privacy_mode = False  # type: bool
+before_send = None  # type: Optional[BeforeSendCallback]
 # Whether to enable feature flag polling for local evaluation by default. Defaults to True.
 # We recommend setting this to False if you are only using the personalApiKey for evaluating remote config payloads via `get_remote_config_payload` and not using local evaluation.
 enable_local_evaluation = True  # type: bool
@@ -875,6 +877,7 @@ def setup() -> Client:
             # or deprecate this proxy option fully (it's already in the process of deprecation, no new clients should be using this method since like 5-6 months)
             enable_exception_autocapture=enable_exception_autocapture,
             log_captured_exceptions=log_captured_exceptions,
+            before_send=before_send,
             enable_local_evaluation=enable_local_evaluation,
             flag_definition_cache_provider=flag_definition_cache_provider,
             capture_exception_code_variables=capture_exception_code_variables,
@@ -886,6 +889,7 @@ def setup() -> Client:
     # always set incase user changes it
     default_client.disabled = disabled
     default_client.debug = debug
+    default_client._set_before_send(before_send)
 
     return default_client
 

posthog/client.py

@@ -288,14 +288,7 @@ def __init__(
         else:
             self.log.setLevel(logging.WARNING)
 
-        if before_send is not None:
-            if callable(before_send):
-                self.before_send = before_send
-            else:
-                self.log.warning("before_send is not callable, it will be ignored")
-                self.before_send = None
-        else:
-            self.before_send = None
+        self._set_before_send(before_send)
 
         if self.enable_exception_autocapture:
             self.exception_capture = ExceptionCapture(self)
@@ -332,6 +325,16 @@ def __init__(
                 if send:
                     consumer.start()
 
+    def _set_before_send(self, before_send):
+        if before_send is not None:
+            if callable(before_send):
+                self.before_send = before_send
+            else:
+                self.log.warning("before_send is not callable, it will be ignored")
+                self.before_send = None
+        else:
+            self.before_send = None
+
     def new_context(self, fresh=False, capture_exceptions=True):
         """
         Create a new context for managing shared state. Learn more about [contexts](/docs/libraries/python#contexts).

posthog/test/test_before_send.py

@@ -1,6 +1,8 @@
+import importlib
 import unittest
 
 import mock
+import posthog
 
 from posthog.client import Client
 from posthog.test.test_utils import FAKE_TEST_API_KEY
@@ -216,3 +218,58 @@ def scrub_pii(event):
             self.assertEqual(enqueued_msg["properties"]["email"], "***@example.com")
             self.assertNotIn("credit_card", enqueued_msg["properties"])
             self.assertEqual(enqueued_msg["properties"]["form_name"], "contact")
+
+
+class TestModuleLevelBeforeSend(unittest.TestCase):
+    def setUp(self):
+        importlib.reload(posthog)
+
+    def tearDown(self):
+        if posthog.default_client:
+            posthog.shutdown()
+        importlib.reload(posthog)
+
+    def test_before_send_callback_used_during_module_level_setup(self):
+        def my_before_send(event):
+            event["properties"]["module_level_before_send"] = True
+            return event
+
+        with mock.patch("posthog.client.batch_post") as mock_post:
+            posthog.api_key = FAKE_TEST_API_KEY
+            posthog.before_send = my_before_send
+            posthog.sync_mode = True
+
+            msg_uuid = posthog.capture("test_event", distinct_id="user1")
+
+            self.assertIsNotNone(msg_uuid)
+            self.assertIs(posthog.default_client.before_send, my_before_send)
+
+            mock_post.assert_called_once()
+            batch_data = mock_post.call_args[1]["batch"]
+            enqueued_msg = batch_data[0]
+            self.assertTrue(enqueued_msg["properties"]["module_level_before_send"])
+
+    def test_before_send_callback_updates_after_client_initialization(self):
+        def my_before_send(event):
+            event["properties"]["updated_after_init"] = True
+            return event
+
+        with mock.patch("posthog.client.batch_post") as mock_post:
+            posthog.api_key = FAKE_TEST_API_KEY
+            posthog.sync_mode = True
+
+            first_msg_uuid = posthog.capture("first_event", distinct_id="user1")
+
+            posthog.before_send = my_before_send
+            second_msg_uuid = posthog.capture("second_event", distinct_id="user1")
+
+            self.assertIsNotNone(first_msg_uuid)
+            self.assertIsNotNone(second_msg_uuid)
+            self.assertIs(posthog.default_client.before_send, my_before_send)
+
+            self.assertEqual(mock_post.call_count, 2)
+            first_batch = mock_post.call_args_list[0][1]["batch"]
+            second_batch = mock_post.call_args_list[1][1]["batch"]
+
+            self.assertNotIn("updated_after_init", first_batch[0]["properties"])
+            self.assertTrue(second_batch[0]["properties"]["updated_after_init"])