fixes #4

phelps-sg · Sep 22, 2023 · 5533742 · 5533742
1 parent e59fb77
commit 5533742
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 1 deletion.
diff --git a/tests/test_tardis_bundle.py b/tests/test_tardis_bundle.py
@@ -34,6 +34,7 @@
 from ray.util.client import ray
 from zipline import Blotter, TradingAlgorithm, get_calendar
 from zipline.assets import AssetDBWriter
+from zipline.country import CountryCode
 from zipline.data import bundles
 from zipline.data.adjustments import SQLiteAdjustmentWriter
 from zipline.data.bcolz_daily_bars import BcolzDailyBarReader, BcolzDailyBarWriter
@@ -284,13 +285,15 @@ def test_generate_empty_metadata():
         "symbol",
         "calendar_name",
         "exchange",
+        "country_code",
     }
     assert result.start_date.dtype == dtype("<M8[ns]")
     assert result.end_date.dtype == dtype("<M8[ns]")
     assert result.auto_close_date.dtype == dtype("<M8[ns]")
     assert result.symbol.dtype == dtype("O")
     assert result.calendar_name.dtype == dtype("O")
     assert result.exchange.dtype == dtype("O")
+    assert result.country_code.dtype == dtype("O")
 
 
 def test_generate_metadata():
@@ -307,6 +310,7 @@ def test_generate_metadata():
             pair.symbol,
             CALENDAR_24_7,
             CALENDAR_24_7,
+            CountryCode.UNITED_STATES,
         )
 
 

diff --git a/zipline_tardis_bundle/bundle.py b/zipline_tardis_bundle/bundle.py
@@ -39,6 +39,7 @@
 from ray.util.client import RayAPIStub, ray
 from tardis_dev import datasets
 from zipline.assets import AssetDBWriter
+from zipline.country import CountryCode
 from zipline.data.adjustments import SQLiteAdjustmentWriter
 from zipline.data.bcolz_daily_bars import BcolzDailyBarWriter
 from zipline.data.bcolz_minute_bars import BcolzMinuteBarWriter
@@ -60,7 +61,7 @@
 
 MINUTES_PER_DAY = 60 * 24
 
-_Metadata = Tuple[pd.Timestamp, pd.Timestamp, pd.Timestamp, str, str, str]
+_Metadata = Tuple[pd.Timestamp, pd.Timestamp, pd.Timestamp, str, str, str, str]
 _IngestPipeline = Iterator[Tuple[int, pd.DataFrame, _Metadata]]
 
 logger = logging.getLogger(__name__)
@@ -342,6 +343,7 @@ def _generate_empty_metadata(pairs: Sized) -> pd.DataFrame:
         ("symbol", "object"),
         ("calendar_name", "object"),
         ("exchange", "object"),
+        ("country_code", "object"),
     ]
     return pd.DataFrame(np.empty(len(pairs), dtype=data_types))
 
@@ -363,6 +365,7 @@ def _generate_metadata(pricing_data: pd.DataFrame, asset: Asset) -> _Metadata:
         asset.symbol,
         CALENDAR_24_7,
         CALENDAR_24_7,
+        CountryCode.UNITED_STATES,
     )