Add process local objects ala threading.local

trytond/doc/ref/tools/index.rst

@@ -15,6 +15,7 @@ Tools API reference.
    immutabledict
    logging
    misc
+   multiprocessing
    qrcode
    singleton
    timezone

trytond/doc/ref/tools/multiprocessing.rst

@@ -0,0 +1,12 @@
+.. _ref-process_local:
+.. module:: trytond.tools.multiprocessing
+
+Process local data
+==================
+
+.. class:: local
+
+   Implements an object that stores process local informations.
+   It mimics `threading.local`_ from the standard library.
+
+.. _`threading.local`: https://docs.python.org/3/library/threading.html#thread-local-data

trytond/trytond/tests/test_tools.py

@@ -10,7 +10,10 @@ import unittest
 from copy import deepcopy
 from decimal import Decimal
 from io import BytesIO
+from multiprocessing import Process
+from multiprocessing.managers import SharedMemoryManager
 from unittest.mock import patch
+from uuid import uuid4
 
 import sql
 import sql.operators
@@ -35,6 +38,7 @@ from trytond.tools.domain_inversion import (
     prepare_reference_domain, simplify, sort, unique_value)
 from trytond.tools.immutabledict import ImmutableDict
 from trytond.tools.logging import format_args
+from trytond.tools.multiprocessing import local
 from trytond.tools.string_ import LazyString, StringPartitioned
 
 try:
@@ -1488,6 +1492,87 @@ class EmailTestCase(TestCase):
                 self.assertEqual(email_.format_address(address, name), result)
 
 
+class ProcessLocalTestCase(TestCase):
+
+    def test_without_fork(self):
+        mydata = local()
+        mydata.number = 42
+        self.assertEqual(mydata.number, 42)
+        self.assertEqual(mydata.__dict__, {'number': 42})
+
+    def test_fork(self):
+        mydata = local()
+        mydata.number = 42
+
+        def f(l, num):
+            mydata.number = num
+            l[0] = mydata.number
+
+        with SharedMemoryManager() as smm:
+            sl = smm.ShareableList([None])
+            p = Process(target=f, args=(sl, 1))
+            p.start()
+            p.join()
+            self.assertEqual(list(sl), [1])
+
+        self.assertEqual(mydata.number, 42)
+
+    def test_inheritance_setattr(self):
+        class MyLocal(local):
+            number = 2
+
+            def __init__(self):
+                self.id = uuid4()
+
+            def squared(self):
+                return self.number ** 2
+
+        mydata = MyLocal()
+
+        def f_setattr(l, num):
+            mydata.number = num
+            l[0] = mydata.squared()
+
+        with SharedMemoryManager() as smm:
+            sl = smm.ShareableList([None, 'x' * len(str(uuid4()))])
+            p = Process(target=f_setattr, args=(sl, 3))
+            p.start()
+            p.join()
+
+            sl = list(sl)
+
+        self.assertEqual(sl[0], 9)
+        self.assertEqual(mydata.number, 2)
+        self.assertEqual(mydata.squared(), 4)
+
+    def test_inheritance_getattr(self):
+        class MyLocal(local):
+            number = 2
+
+            def __init__(self):
+                self.id = uuid4()
+
+            def squared(self):
+                return self.number ** 2
+
+        mydata = MyLocal()
+
+        def f_getattr(l, num):
+            l[1] = str(mydata.id)
+
+        with SharedMemoryManager() as smm:
+            sl = smm.ShareableList([None, 'x' * len(str(uuid4()))])
+            p = Process(target=f_getattr, args=(sl, 3))
+            p.start()
+            p.join()
+
+            sl = list(sl)
+
+        self.assertNotEqual(sl[1], str(mydata.id))
+        self.assertEqual(mydata.number, 2)
+        self.assertEqual(mydata.squared(), 4)
+
+
 def load_tests(loader, tests, pattern):
     tests.addTest(doctest.DocTestSuite(decimal_))
     return tests

trytond/trytond/tools/multiprocessing.py

@@ -0,0 +1,56 @@
+# -*- coding: utf-8 -*-
+# This file is part of Tryton.  The COPYRIGHT file at the top level of
+# this repository contains the full copyright notices and license terms.
+
+import os
+from threading import Lock
+
+_lock = Lock()
+
+
+class _local_impl:
+    __slots__ = 'dict', 'pid', 'localargs'
+
+    def __init__(self, localargs):
+        self.localargs = localargs
+        self.reset()
+
+    def reset(self):
+        self.dict = {}
+        self.pid = os.getpid()
+
+
+class local:
+    __slots__ = '_local_impl', '__dict__'
+
+    def __new__(cls, /, *args, **kwargs):
+        self = super().__new__(cls)
+        # do not trigger __setattr__ below
+        object.__setattr__(self, '_local_impl', _local_impl((args, kwargs)))
+        return self
+
+    def __getattribute__(self, name):
+        impl = object.__getattribute__(self, '_local_impl')
+        if impl.pid == os.getpid():
+            d = impl.dict
+        else:
+            impl.reset()
+            d = impl.dict
+            args, kwargs = impl.localargs
+            with _lock:
+                self.__init__(*args, **kwargs)
+        object.__setattr__(self, '__dict__', d)
+        return object.__getattribute__(self, name)
+
+    def __setattr__(self, name, value):
+        impl = object.__getattribute__(self, '_local_impl')
+        if impl.pid == os.getpid():
+            d = impl.dict
+        else:
+            impl.reset()
+            d = impl.dict
+            args, kwargs = impl.localargs
+            with _lock:
+                self.__init__(*args, **kwargs)
+        object.__setattr__(self, '__dict__', d)
+        return object.__setattr__(self, name, value)