From 6f4428d20634d2aec94119e5c70c7c6e90d8efe7 Mon Sep 17 00:00:00 2001
From: Petr Heinz <petr@betterstack.com>
Date: Mon, 16 Sep 2024 14:57:30 +0200
Subject: [PATCH] Improve omitting circular references (#28)

---
 logtail/frame.py      | 10 ++++++++--
 tests/test_handler.py | 20 ++++++++++++++++++++
 2 files changed, 28 insertions(+), 2 deletions(-)

diff --git a/logtail/frame.py b/logtail/frame.py
index 90fe89d..8d730b5 100644
--- a/logtail/frame.py
+++ b/logtail/frame.py
@@ -60,6 +60,9 @@ def _parse_custom_events(record, include_extra_attributes):
 def _remove_circular_dependencies(obj, memo=None):
     if memo is None:
         memo = set()
+    # Skip immutable types, which can't contain circular dependencies
+    if isinstance(obj, (str, int, float, bool)):
+        return obj
     if id(obj) in memo:
         return "<omitted circular reference>"
     memo.add(id(obj))
@@ -69,8 +72,11 @@ def _remove_circular_dependencies(obj, memo=None):
             new_dict[key] = _remove_circular_dependencies(value, memo)
         return new_dict
     elif isinstance(obj, list):
-        new_list = [_remove_circular_dependencies(item, memo) for item in obj]
-        return new_list
+        return [_remove_circular_dependencies(item, memo) for item in obj]
+    elif isinstance(obj, tuple):
+        return tuple(_remove_circular_dependencies(item, memo) for item in obj)
+    elif isinstance(obj, set):
+        return {_remove_circular_dependencies(item, memo) for item in obj}
     else:
         return obj
 
diff --git a/tests/test_handler.py b/tests/test_handler.py
index ece326f..bd10255 100644
--- a/tests/test_handler.py
+++ b/tests/test_handler.py
@@ -203,6 +203,26 @@ def test_can_send_circular_dependency_in_extra_data(self, MockWorker):
         self.assertEqual(log_entry['data']['egg']['chicken'], "<omitted circular reference>")
         self.assertTrue(handler.pipe.empty())
 
+    @patch('logtail.handler.FlushWorker')
+    def test_can_have_multiple_instance_of_same_string_in_extra_data(self, MockWorker):
+        buffer_capacity = 1
+        handler = LogtailHandler(
+            source_token=self.source_token,
+            buffer_capacity=buffer_capacity
+        )
+
+        logger = logging.getLogger(__name__)
+        logger.handlers = []
+        logger.addHandler(handler)
+        test_string = 'this is a test string'
+        logger.info('hello', extra={'test1': test_string, 'test2': test_string})
+
+        log_entry = handler.pipe.get()
+
+        self.assertEqual(log_entry['message'], 'hello')
+        self.assertEqual(log_entry['test1'], 'this is a test string')
+        self.assertEqual(log_entry['test2'], 'this is a test string')
+        self.assertTrue(handler.pipe.empty())
 
     @patch('logtail.handler.FlushWorker')
     def test_can_send_circular_dependency_in_context(self, MockWorker):