enhance the performance

alphamind/data/engines/sqlengine.py

@@ -106,44 +106,9 @@ class SqlEngine(object):
         return strategy_names
 
     def fetch_codes(self, ref_date: str, univ: Universe) -> List[int]:
-
-        def get_universe(univ, ref_date):
-            query = select([UniverseTable.Code]).distinct().where(
-                and_(
-                    UniverseTable.Date == ref_date,
-                    UniverseTable.universe.in_(univ)
-                )
-            )
-
-            cursor = self.engine.execute(query)
-            codes_set = {c[0] for c in cursor.fetchall()}
-            return codes_set
-
-        codes_set = None
-
-        if univ.include_universe:
-            include_codes_set = get_universe(univ.include_universe, ref_date)
-            codes_set = include_codes_set
-
-        if univ.exclude_universe:
-            exclude_codes_set = get_universe(univ.exclude_universe, ref_date)
-            codes_set -= exclude_codes_set
-
-        if univ.include_codes:
-            codes_set = codes_set.union(univ.include_codes)
-
-        if univ.exclude_codes:
-            codes_set -= set(univ.exclude_codes)
-
-        if univ.filter_cond is not None:
-            query = select([UniverseTable.Code]).distinct().where(
-                and_(UniverseTable.Date == ref_date,
-                     univ.filter_cond)
-            )
-
-            cursor = self.engine.execute(query)
-            filtered_codes = {c[0] for c in cursor.fetchall()}
-            codes_set = codes_set.intersection(filtered_codes)
+        query = univ.query(ref_date)
+        cursor = self.engine.execute(query)
+        codes_set = {c[0] for c in cursor.fetchall()}
 
         return sorted(codes_set)
 

alphamind/data/engines/universe.py

@@ -7,6 +7,7 @@ Created on 2017-7-7
 
 from typing import Iterable
 from sqlalchemy import and_
+from sqlalchemy import select
 from alphamind.data.dbmodel.models import Uqer
 from alphamind.data.dbmodel.models import Universe as UniverseTable
 
@@ -26,16 +27,37 @@ class Universe(object):
         self.exclude_universe = exclude_universe
         self.include_codes = include_codes
         self.exclude_codes = exclude_codes
+        self.filter_cond = filter_cond
 
-        if filter_cond is not None:
-            self.filter_cond = self.format_cond(filter_cond)
-        else:
-            self.filter_cond = None
+    def query(self, ref_date):
+        query = select([UniverseTable.Code]).distinct()
 
-    def format_cond(self, filter_cond):
-        filter_cond = and_(
-            filter_cond,
-            Uqer.Code == UniverseTable.Code,
-            Uqer.Date == UniverseTable.Date
+        all_conditions = [UniverseTable.Date == ref_date]
+        if self.include_universe:
+            univ_in = UniverseTable.universe.in_(self.include_universe)
+            all_conditions.append(univ_in)
+
+        if self.exclude_universe:
+            univ_out = UniverseTable.universe.notin_(self.exclude_universe)
+            all_conditions.append(univ_out)
+
+        if self.include_codes:
+            codes_in = UniverseTable.Code.in_(self.include_codes)
+            all_conditions.append(codes_in)
+
+        if self.exclude_codes:
+            codes_out = UniverseTable.Code.notin_(self.exclude_codes)
+            all_conditions.append(codes_out)
+
+        if self.filter_cond is not None:
+            all_conditions.extend([self.filter_cond,
+                                   UniverseTable.Code == Uqer.Code,
+                                   UniverseTable.Date == Uqer.Date])
+
+        query = query.where(
+            and_(
+                *all_conditions
+            )
         )
-        return filter_cond
+
+        return query