changeset 459:2ef04e16f0b9

internal: Add support for fake pickling of date/time structures.
author Ludovic Chabant <ludovic@chabant.com>
date Sat, 11 Jul 2015 00:44:58 -0700
parents 612ed0526afd
children 55fc8918cb75
files piecrust/fastpickle.py
diffstat 1 files changed, 72 insertions(+), 16 deletions(-) [+]
line wrap: on
line diff
--- a/piecrust/fastpickle.py	Thu Jul 09 22:30:45 2015 -0700
+++ b/piecrust/fastpickle.py	Sat Jul 11 00:44:58 2015 -0700
@@ -11,34 +11,77 @@
     return _unpickle_object(data)
 
 
-def _tuple_dispatch(obj, func):
+_PICKLING = 0
+_UNPICKLING = 1
+
+
+def _tuple_dispatch(obj, func, op):
     res = [None] * len(obj)
     for i, c in enumerate(obj):
         res[i] = func(c)
     return tuple(res)
 
 
-def _list_dispatch(obj, func):
+def _list_dispatch(obj, func, op):
     res = [None] * len(obj)
     for i, c in enumerate(obj):
         res[i] = func(c)
     return res
 
 
-def _dict_dispatch(obj, func):
+def _dict_dispatch(obj, func, op):
     res = {}
     for k, v in obj.items():
         res[k] = func(v)
     return res
 
 
-def _set_dispatch(obj, func):
+def _set_dispatch(obj, func, op):
     res = set()
     for v in obj:
         res.add(func(v))
     return res
 
 
+def _date_convert(obj, op):
+    if op == _PICKLING:
+        return {'__class__': 'date',
+                'year': obj.year,
+                'month': obj.month,
+                'day': obj.day}
+    elif op == _UNPICKLING:
+        return datetime.date(
+                obj['year'], obj['month'], obj['day'])
+
+
+def _datetime_convert(obj, op):
+    if op == _PICKLING:
+        return {'__class__': 'datetime',
+                'year': obj.year,
+                'month': obj.month,
+                'day': obj.day,
+                'hour': obj.hour,
+                'minute': obj.minute,
+                'second': obj.second,
+                'microsecond': obj.microsecond}
+    elif op == _UNPICKLING:
+        return datetime.datetime(
+                obj['year'], obj['month'], obj['day'],
+                obj['hour'], obj['minute'], obj['second'], obj['microsecond'])
+
+
+def _time_convert(obj, op):
+    if op == _PICKLING:
+        return {'__class__': 'time',
+                'hour': obj.hour,
+                'minute': obj.minute,
+                'second': obj.second,
+                'microsecond': obj.microsecond}
+    elif op == _UNPICKLING:
+        return datetime.time(
+                obj['hour'], obj['minute'], obj['second'], obj['microsecond'])
+
+
 _identity_dispatch = object()
 
 _type_dispatch = {
@@ -47,9 +90,6 @@
         int: _identity_dispatch,
         float: _identity_dispatch,
         str: _identity_dispatch,
-        datetime.date: _identity_dispatch,
-        datetime.datetime: _identity_dispatch,
-        datetime.time: _identity_dispatch,
         tuple: _tuple_dispatch,
         list: _list_dispatch,
         dict: _dict_dispatch,
@@ -58,6 +98,20 @@
         }
 
 
+_type_convert = {
+        datetime.date: _date_convert,
+        datetime.datetime: _datetime_convert,
+        datetime.time: _time_convert
+        }
+
+
+_type_unconvert = {
+        'date': _date_convert,
+        'datetime': _datetime_convert,
+        'time': _time_convert
+        }
+
+
 def _pickle_object(obj):
     t = type(obj)
     disp = _type_dispatch.get(t)
@@ -65,10 +119,11 @@
         return obj
 
     if disp is not None:
-        return disp(obj, _pickle_object)
+        return disp(obj, _pickle_object, _PICKLING)
 
-    if isinstance(obj, Exception):
-        return obj
+    conv = _type_convert.get(t)
+    if conv is not None:
+        return conv(obj, _PICKLING)
 
     getter = getattr(obj, '__getstate__', None)
     if getter is not None:
@@ -76,7 +131,7 @@
     else:
         state = obj.__dict__
 
-    state = _dict_dispatch(state, _pickle_object)
+    state = _dict_dispatch(state, _pickle_object, _PICKLING)
     state['__class__'] = obj.__class__.__name__
     state['__module__'] = obj.__class__.__module__
 
@@ -90,15 +145,16 @@
         return state
 
     if (disp is not None and
-            (t != dict or '__module__' not in state)):
-        return disp(state, _unpickle_object)
+            (t != dict or '__class__' not in state)):
+        return disp(state, _unpickle_object, _UNPICKLING)
 
-    if isinstance(state, Exception):
-        return state
+    class_name = state['__class__']
+    conv = _type_unconvert.get(class_name)
+    if conv is not None:
+        return conv(state, _UNPICKLING)
 
     mod_name = state['__module__']
     mod = sys.modules[mod_name]
-    class_name = state['__class__']
     class_def = getattr(mod, class_name)
     obj = class_def.__new__(class_def)