refactor serialization

alphamind/formula/utilities.py

@@ -5,13 +5,12 @@ Created on 2017-11-27
 @author: cheng.li
 """
 
-import pickle
-import base64
+from alphamind.utilities import encode
+from alphamind.utilities import decode
 
 
 def encode_formula(formula):
-    encoded = base64.encodebytes(pickle.dumps(formula))
-    str_repr = encoded.decode('ascii')
+    str_repr = encode(formula)
     return {'desc': str_repr,
             'formula_type': formula.__class__.__module__ + "." + formula.__class__.__name__,
             'dependency': formula.fields,
@@ -20,7 +19,7 @@ def encode_formula(formula):
 
 def decode_formula(str_repr):
     encoded = str_repr.encode('ascii')
-    formula = pickle.loads(base64.decodebytes(encoded))
+    formula = decode(encoded)
     return formula
 
 

alphamind/model/linearmodel.py

@@ -5,8 +5,6 @@ Created on 2017-5-10
 @author: cheng.li
 """
 
-import base64
-import pickle
 import numpy as np
 import arrow
 from distutils.version import LooseVersion
@@ -15,6 +13,8 @@ from sklearn.linear_model import LinearRegression as LinearRegressionImpl
 from PyFin.api import pyFinAssert
 from alphamind.model.modelbase import ModelBase
 from alphamind.utilities import alpha_logger
+from alphamind.utilities import encode
+from alphamind.utilities import decode
 
 
 class ConstLinearModel(ModelBase):
@@ -65,8 +65,7 @@ class LinearRegression(ModelBase):
     def save(self) -> dict:
         model_desc = super().save()
         model_desc['internal_model'] = self.impl.__class__.__module__ + "." + self.impl.__class__.__name__
-        encoded = base64.encodebytes(pickle.dumps(self.impl))
-        model_desc['desc'] = encoded.decode('ascii')
+        model_desc['desc'] = encode(self.impl)
         model_desc['sklearn_version'] = sklearn_version
         model_desc['trained_time'] = self.trained_time
         model_desc['weight'] = self.impl.coef_.tolist()
@@ -86,8 +85,7 @@ class LinearRegression(ModelBase):
                                  'Loaded model may work incorrectly.'.format(
                 sklearn_version, model_desc['sklearn_version']))
 
-        encoded = model_desc['desc'].encode('ascii')
-        obj_layout.impl = pickle.loads(base64.decodebytes(encoded))
+        obj_layout.impl = decode(model_desc['desc'])
         return obj_layout
 
     @property

alphamind/utilities.py

@@ -5,6 +5,8 @@ Created on 2017-4-25
 @author: cheng.li
 """
 
+import base64
+import pickle
 import math
 from simpleutils import CustomLogger
 import numpy as np
@@ -316,3 +318,13 @@ def aggregate(groups, x, func, ddof=1):
         raise ValueError('({0}) is not recognized as valid functor'.format(func))
 
     return value_data
+
+
+def encode(obj: object) -> str:
+    encoded = base64.encodebytes(pickle.dumps(obj))
+    return encoded.decode('ascii')
+
+
+def decode(str_repr: str):
+    encoded = str_repr.encode('ascii')
+    return pickle.loads(base64.decodebytes(encoded))
\ No newline at end of file