update universe

alphamind/data/engines/sqlengine.py

@@ -143,20 +143,15 @@ class SqlEngine(object):
         return strategy_names
 
     def fetch_codes(self, ref_date: str, universe: Universe) -> List[int]:
-        cond = universe.query(ref_date)
-        query = select([UniverseTable.trade_date, UniverseTable.code]).distinct().where(cond)
-        cursor = self.engine.execute(query)
-        codes_set = {c[1] for c in cursor.fetchall()}
-        return sorted(codes_set)
+        df = universe.query(self, ref_date, ref_date)
+        return sorted(df.code.tolist())
 
     def fetch_codes_range(self,
                           universe: Universe,
                           start_date: str = None,
                           end_date: str = None,
                           dates: Iterable[str] = None) -> pd.DataFrame:
-        cond = universe.query_range(start_date, end_date, dates)
-        query = select([UniverseTable.trade_date, UniverseTable.code]).distinct().where(cond)
-        return pd.read_sql(query, self.engine)
+        return universe.query(self, start_date, end_date, dates)
 
     def fetch_dx_return(self,
                         ref_date: str,

alphamind/data/engines/universe.py

@@ -51,7 +51,7 @@ class Universe(object):
 
         and_conditions = []
         if self.exclude_universe:
-            and_conditions.append(~UniverseTable.universe.in_(self.exclude_universe))
+            and_conditions.append(UniverseTable.universe.notin_(self.exclude_universe))
 
         return and_(
             query,
@@ -67,7 +67,7 @@ class Universe(object):
             # simple case
             query = select([UniverseTable.trade_date, UniverseTable.code]).where(
                 universe_cond
-            )
+            ).distinct()
             return pd.read_sql(query, engine.engine)
         else:
             if isinstance(self.filter_cond, Transformer):
@@ -109,7 +109,7 @@ if __name__ == '__main__':
     from alphamind.data.engines.sqlengine import SqlEngine
 
     engine = SqlEngine()
-    universe = Universe('ss', ['hs300'], special_codes=[603138])
+    universe = Universe('ss', ['ashare_ex'], exclude_universe=['hs300', 'zz500'], special_codes=[603138])
     print(universe.query(engine,
                          start_date='2017-12-21',
                          end_date='2017-12-25'))