PB-1920: Fix AttributeError: '_thread._local' object has no attribute 'data'

Author: msom

logging_utilities/context/thread_context.py

@@ -14,30 +14,42 @@ class ThreadMappingContext(BaseContext):
 
     def __init__(self):
         self.__local = threading.local()
-        self.__local.data = {}
+        self.ensure_data()
+
+    def ensure_data(self):
+        if not hasattr(self.__local, 'data'):
+            self.__local.data = {}
 
     def __str__(self):
+        self.ensure_data()
         return str(self.__local.data)
 
     def __getitem__(self, __key):
+        self.ensure_data()
         return self.__local.data[__key]
 
     def __setitem__(self, __key, __value):
+        self.ensure_data()
         self.__local.data[__key] = __value
 
     def __delitem__(self, __key):
+        self.ensure_data()
         del self.__local.data[__key]
 
     def __len__(self):
+        self.ensure_data()
         return len(self.__local.data)
 
     def __iter__(self):
+        self.ensure_data()
         return self.__local.data.__iter__()
 
     def __contains__(self, __o):
+        self.ensure_data()
         return self.__local.data.__contains__(__o)
 
     def init(self, data=None):
+        self.ensure_data()
         if data is None:
             self.__local.data = {}
         else:
@@ -46,18 +58,23 @@ def init(self, data=None):
             self.__local.data = data
 
     def get(self, key, default=None):
+        self.ensure_data()
         return self.__local.data.get(key, default)
 
     def pop(self, key, default=__marker):
+        self.ensure_data()
         if default == self.__marker:
             return self.__local.data.pop(key)
         return self.__local.data.pop(key, default)
 
     def set(self, key, value):
+        self.ensure_data()
         self.__local.data[key] = value
 
     def delete(self, key):
+        self.ensure_data()
         del self.__local.data[key]
 
     def clear(self):
+        self.ensure_data()
         self.__local.data = {}

tests/test_logging_context.py

@@ -2,6 +2,7 @@
 import unittest
 from concurrent.futures import ThreadPoolExecutor
 from concurrent.futures import as_completed
+from threading import Thread
 
 from logging_utilities.context import get_logging_context
 from logging_utilities.context import remove_logging_context
@@ -125,6 +126,21 @@ def test_thread_context_str(self):
         ctx.init({'a': 1, 'b': 2, 'c': 'my string'})
         self.assertEqual(str(ctx), "{'a': 1, 'b': 2, 'c': 'my string'}")
 
+    def test_thread_context_local_data(self):
+        ctx = ThreadMappingContext()
+        ctx['thread'] = 'main'
+        results = {}
+
+        def worker():
+            assert 'thread' not in ctx
+            ctx['thread'] = 'worker'
+
+        t = Thread(target=worker)
+        t.start()
+        t.join()
+
+        assert ctx['thread'] == 'main'
+
 
 class LoggingContextTest(unittest.TestCase):