From 7152f6d9156a19a9f1b31e313d2c19d840dc089d Mon Sep 17 00:00:00 2001
From: Jesse Noller <jnoller@gmail.com>
Date: Thu, 2 Apr 2009 05:17:26 +0000
Subject: Add custom initializer argument to multiprocess.Manager*, courtesy of
 lekma

---
 Doc/library/multiprocessing.rst  |  5 +++--
 Lib/multiprocessing/managers.py  | 13 ++++++++++---
 Lib/multiprocessing/pool.py      |  3 +++
 Lib/test/test_multiprocessing.py | 32 +++++++++++++++++++++++++++++++-
 Misc/NEWS                        |  3 +++
 5 files changed, 50 insertions(+), 6 deletions(-)

diff --git a/Doc/library/multiprocessing.rst b/Doc/library/multiprocessing.rst
index 18e81a7..a36836e 100644
--- a/Doc/library/multiprocessing.rst
+++ b/Doc/library/multiprocessing.rst
@@ -1130,9 +1130,10 @@ their parent process exits.  The manager classes are defined in the
    ``current_process().authkey``.  Otherwise *authkey* is used and it
    must be a string.
 
-   .. method:: start()
+   .. method:: start([initializer[, initargs]])
 
-      Start a subprocess to start the manager.
+      Start a subprocess to start the manager.  If *initializer* is not ``None``
+      then the subprocess will call ``initializer(*initargs)`` when it starts.
 
    .. method:: serve_forever()
 
diff --git a/Lib/multiprocessing/managers.py b/Lib/multiprocessing/managers.py
index 02e96b9..fde0e40 100644
--- a/Lib/multiprocessing/managers.py
+++ b/Lib/multiprocessing/managers.py
@@ -475,12 +475,15 @@ class BaseManager(object):
         dispatch(conn, None, 'dummy')
         self._state.value = State.STARTED
 
-    def start(self):
+    def start(self, initializer=None, initargs=()):
         '''
         Spawn a server process for this manager object
         '''
         assert self._state.value == State.INITIAL
 
+        if initializer is not None and not hasattr(initializer, '__call__'):
+            raise TypeError('initializer must be a callable')
+
         # pipe over which we will retrieve address of server
         reader, writer = connection.Pipe(duplex=False)
 
@@ -488,7 +491,7 @@ class BaseManager(object):
         self._process = Process(
             target=type(self)._run_server,
             args=(self._registry, self._address, self._authkey,
-                  self._serializer, writer),
+                  self._serializer, writer, initializer, initargs),
             )
         ident = ':'.join(str(i) for i in self._process._identity)
         self._process.name = type(self).__name__  + '-' + ident
@@ -509,10 +512,14 @@ class BaseManager(object):
             )
 
     @classmethod
-    def _run_server(cls, registry, address, authkey, serializer, writer):
+    def _run_server(cls, registry, address, authkey, serializer, writer,
+                    initializer=None, initargs=()):
         '''
         Create a server, report its address and run it
         '''
+        if initializer is not None:
+            initializer(*initargs)
+
         # create server
         server = cls._Server(registry, address, authkey, serializer)
 
diff --git a/Lib/multiprocessing/pool.py b/Lib/multiprocessing/pool.py
index 9da27d4..bc7e8f1 100644
--- a/Lib/multiprocessing/pool.py
+++ b/Lib/multiprocessing/pool.py
@@ -92,6 +92,9 @@ class Pool(object):
             except NotImplementedError:
                 processes = 1
 
+        if initializer is not None and not hasattr(initializer, '__call__'):
+            raise TypeError('initializer must be a callable')
+
         self._pool = []
         for i in range(processes):
             w = self.Process(
diff --git a/Lib/test/test_multiprocessing.py b/Lib/test/test_multiprocessing.py
index 8ef2e2f..cbcb630 100644
--- a/Lib/test/test_multiprocessing.py
+++ b/Lib/test/test_multiprocessing.py
@@ -1831,7 +1831,37 @@ class OtherTest(unittest.TestCase):
                           multiprocessing.connection.answer_challenge,
                           _FakeConnection(), b'abc')
 
-testcases_other = [OtherTest, TestInvalidHandle]
+#
+# Test Manager.start()/Pool.__init__() initializer feature - see issue 5585
+#
+
+def initializer(ns):
+    ns.test += 1
+
+class TestInitializers(unittest.TestCase):
+    def setUp(self):
+        self.mgr = multiprocessing.Manager()
+        self.ns = self.mgr.Namespace()
+        self.ns.test = 0
+
+    def tearDown(self):
+        self.mgr.shutdown()
+
+    def test_manager_initializer(self):
+        m = multiprocessing.managers.SyncManager()
+        self.assertRaises(TypeError, m.start, 1)
+        m.start(initializer, (self.ns,))
+        self.assertEqual(self.ns.test, 1)
+        m.shutdown()
+
+    def test_pool_initializer(self):
+        self.assertRaises(TypeError, multiprocessing.Pool, initializer=1)
+        p = multiprocessing.Pool(1, initializer, (self.ns,))
+        p.close()
+        p.join()
+        self.assertEqual(self.ns.test, 1)
+
+testcases_other = [OtherTest, TestInvalidHandle, TestInitializers]
 
 #
 #
diff --git a/Misc/NEWS b/Misc/NEWS
index 8a50059..7de7b8a 100644
--- a/Misc/NEWS
+++ b/Misc/NEWS
@@ -202,6 +202,9 @@ Core and Builtins
 Library
 -------
 
+- Issue 5585: Add the ability to call an initializer to mulitiprocessing.manager
+  so that users can install custonm handlers/etc.
+
 - Issue 3551: Patch multiprocessing to raise a proper exception if the size of the
   object when writefile is called causes a ERROR_NO_SYSTEM_RESOURCES. Added docs
   to note the limitation
-- 
cgit v0.12