added universe saver and loader

alphamind/data/engines/universe.py

@@ -12,12 +12,13 @@ from sqlalchemy import or_
 from sqlalchemy import select
 from sqlalchemy import join
 from sqlalchemy import outerjoin
-from PyFin.api import pyFinAssert
 from alphamind.data.dbmodel.models import Universe as UniverseTable
 from alphamind.data.dbmodel.models import FullFactor
 from alphamind.data.engines.utilities import _map_factors
 from alphamind.data.engines.utilities import factor_tables
 from alphamind.data.transformer import Transformer
+from alphamind.utilities import encode
+from alphamind.utilities import decode
 
 
 class Universe(object):
@@ -25,8 +26,8 @@ class Universe(object):
     def __init__(self,
                  name: str,
                  base_universe: Iterable,
-                 exclude_universe: Iterable=None,
-                 special_codes: Iterable=None,
+                 exclude_universe: Iterable = None,
+                 special_codes: Iterable = None,
                  filter_cond=None):
         self.name = name
         self.base_universe = base_universe
@@ -59,7 +60,7 @@ class Universe(object):
             *and_conditions
         )
 
-    def query(self, engine, start_date: str=None, end_date: str=None, dates=None) -> pd.DataFrame:
+    def query(self, engine, start_date: str = None, end_date: str = None, dates=None) -> pd.DataFrame:
 
         universe_cond = self._query_statements(start_date, end_date, dates)
 
@@ -103,6 +104,29 @@ class Universe(object):
             df = df[df[filter_fields[0]] == 1].reset_index()[['trade_date', 'code']]
             return df
 
+    def save(self):
+        return dict(
+            name=self.name,
+            base_universe=self.base_universe,
+            exclude_universe=self.exclude_universe,
+            special_codes=self.special_codes,
+            filter_cond=encode(self.filter_cond)
+        )
+
+    @classmethod
+    def load(cls, universe_desc):
+        name = universe_desc['name']
+        base_universe = universe_desc['base_universe']
+        exclude_universe = universe_desc['exclude_universe']
+        special_codes = universe_desc['special_codes']
+        filter_cond = decode(universe_desc['filter_cond'])
+
+        return cls(name=name,
+                   base_universe=base_universe,
+                   exclude_universe=exclude_universe,
+                   special_codes=special_codes,
+                   filter_cond=filter_cond)
+
 
 if __name__ == '__main__':
     from PyFin.api import *