reduce for _thread._local

Author: pwnslinger

dill/_dill.py

@@ -47,6 +47,11 @@ def _trace(boolean):
         from _thread import RLock as RLockType
     else:
         from threading import _RLock as RLockType
+    # monkey patch for _thread._local
+    try:
+        from _thread import _local as LocalType
+    except ImportError:
+        from _threading_local import local as LocalType
    #from io import IOBase
     from types import CodeType, FunctionType, MethodType, GeneratorType, \
         TracebackType, FrameType, ModuleType, BuiltinMethodType
@@ -85,7 +90,7 @@ def _trace(boolean):
     FileNotFoundError = IOError
 if PY3 and sys.hexversion < 0x03040000:
     GENERATOR_FAIL = True
-else: GENERATOR_FAIL = False    
+else: GENERATOR_FAIL = False
 try:
     import ctypes
     HAS_CTYPES = True
@@ -596,6 +601,13 @@ def _create_ftype(ftypeobj, func, args, kwds):
         args = ()
     return ftypeobj(func, *args, **kwds)
 
+def _create_local(impl, *args, **kwargs): #XXX: ignores 'blocking'
+    local = LocalType()
+    if local:
+        if impl:
+            local.__setattr__('_local__impl', impl)
+    return local
+
 def _create_lock(locked, *args): #XXX: ignores 'blocking'
     from threading import Lock
     lock = Lock()
@@ -927,6 +939,23 @@ def save_classobj(pickler, obj): #FIXME: enable pickler._byref
         log.info("# C2")
     return
 
+@register(LocalType)
+def save_local(pickler, obj):
+    log.info("Loc: %s" % obj)
+    impl = obj.__getattribute__('_local__impl')
+    if impl:
+        try:
+            dct = impl.get_dict()
+        except KeyError:
+            dct = impl.create_dict()
+            args, kwargs = impl.localargs
+            obj.__init__(*args, **kwargs)
+        with impl.locallock:
+            obj.__setattr__('__dict__', dct)
+    pickler.save_reduce(_create_local, (impl,), obj=obj)
+    log.info("# Loc")
+    return
+
 @register(LockType)
 def save_lock(pickler, obj):
     log.info("Lo: %s" % obj)