A fork-friendly threading.local().

process_local.py

import os
import threading


class ObjectProxy(object):

    """Wrap a target object and provide an interface to its underlying object.

    This is especially useful when you're defining an interface which overrides
    the attribute access magic methods, but you still need to access named
    slots on the underlying object.

    Example
    -------

    Define an object which overrides attribute access:

        >>> class X(object):
        ...     def __getattr__(self, attr):
        ...         return attr.upper()
        ...     def __setattr__(self, attr, value):
        ...         print "Tried to set %s to %r" % (attr, value)
        >>> x = X()
        >>> x.someattr = 123
        Tried to set someattr to 123
        >>> x.someattr
        'SOMEATTR'

    Create an `ObjectProxy` over `x` -- this will give you access to the
    underlying object, not the overridden attribute interface:

        >>> proxy = ObjectProxy(x)
        >>> hasattr(proxy, 'someattr')
        False
        >>> proxy.someattr = 123
        >>> proxy.someattr
        123

    The attribute you set will appear on the original object, since you defined
    `__getattr__` (which is only ever used as a fallback):

        >>> x.someattr
        123
    """

    def __init__(self, target):
        object.__setattr__(self, 'target', target)

    def __getattribute__(self, attr):
        return object.__getattribute__(object.__getattribute__(self, 'target'),
                                       attr)

    def __setattr__(self, attr, value):
        return object.__setattr__(object.__getattribute__(self, 'target'),
                                  attr, value)

    def __delattr__(self, attr):
        return object.__delattr__(object.__getattribute__(self, 'target'),
                                  attr)


class ProcessLocal(object):

    """A fork-friendly `threading.local()`.

    This object behaves much like `threading.local()`, only it maintains
    separate objects for different processes. This enables you to call
    `os.fork()` and not have to worry about visibility of shared state, like
    sockets or database connections.

    Example
    -------

    Create a `ProcessLocal` and check that you can set and retrieve attributes
    on it:

        >>> import multiprocessing, sys, time
        >>> loc = ProcessLocal()
        >>> loc.A = 123
        >>> print loc.A
        123

    Spawn a process which attempts to access `loc.A`, and puts the value it
    sees on a queue back to the parent process:

        >>> q = multiprocessing.Queue()
        >>> def show_A(queue):
        ...     queue.put(getattr(loc, 'A', 'NOT SET'))
        >>> proc = multiprocessing.Process(target=show_A, args=(q,))

    As you'll see, the child process doesn't see the `A` local:

        >>> proc.start()
        >>> print q.get()
        NOT SET
        >>> proc.join()

    But `A` remains defined in the parent process:

        >>> print loc.A
        123

    The same works with threading:

        >>> import Queue, threading
        >>> q2 = Queue.Queue()
        >>> t = threading.Thread(target=show_A, args=(q2,))
        >>> t.start()
        >>> print q2.get()
        NOT SET
        >>> t.join()
        >>> print loc.A
        123
    """

    def __init__(self):
        self = ObjectProxy(self)
        self.state = {}
        self.lock = threading.Lock()

    def _get_locals(self):
        self = ObjectProxy(self)
        pid = os.getpid()
        # The following check-lock-check is for performance reasons: we don't
        # want to acquire the lock every time we check, but if we just check
        # and acquire the lock and don't check again, the key may already have
        # been set in the dictionary.
        if pid not in self.state:
            with self.lock:
                if pid not in self.state:
                    self.state[pid] = threading.local()
        return self.state[pid]

    def __getattribute__(self, attr):
        return getattr(ObjectProxy(self)._get_locals(), attr)

    def __setattr__(self, attr, value):
        setattr(ObjectProxy(self)._get_locals(), attr, value)

    def __delattr__(self, attr):
        delattr(ObjectProxy(self)._get_locals(), attr)