zipline_polygon_bundle 0.1.8__py3-none-any.whl → 0.2.0__py3-none-any.whl

zipline_polygon_bundle/config.py

@@ -1,14 +1,39 @@
-from exchange_calendars.calendar_helpers import Date, parse_date, parse_timestamp
+from exchange_calendars.calendar_helpers import Date, parse_date
 from zipline.utils.calendar_utils import get_calendar
 
-from typing import Iterator, Tuple, List
+from .nyse_all_hours_calendar import NYSE_ALL_HOURS
+
+from typing import Iterator, Tuple
 
 import pandas as pd
 from pyarrow.fs import LocalFileSystem
 import os
+import datetime
 import re
 import fnmatch
 
+AGG_TIME_DAY = "day"
+AGG_TIME_MINUTE = "minute"
+AGG_TIME_TRADES = "1minute"
+
+PARTITION_COLUMN_NAME = "part"
+PARTITION_KEY_LENGTH = 2
+
+
+def to_partition_key(s: str) -> str:
+    """
+    Partition key is low cardinality and must be filesystem-safe.
+    The reason for partitioning is to keep the parquet files from getting too big.
+    10 years of minute aggs for US stocks is 83GB gzipped.  A single parquet would be 62GB on disk.
+    Currently the first two characters so files stay under 1GB.  Weird characters are replaced with "A".
+    """
+    k = (s + "A")[0:PARTITION_KEY_LENGTH].upper()
+    if k.isalpha():
+        return k
+    # Replace non-alpha characters with "A".
+    k = "".join([c if c.isalpha() else "A" for c in k])
+    return k
+
 
 class PolygonConfig:
     def __init__(
@@ -17,12 +42,12 @@ class PolygonConfig:
         calendar_name: str,
         start_date: Date,
         end_date: Date,
-        agg_time: str = "day",
-        custom_aggs_format: str = "{config.agg_timedelta.seconds}sec_aggs",
+        agg_time: str = AGG_TIME_DAY,
     ):
         self.calendar_name = calendar_name
         self.start_date = start_date
         self.end_date = end_date
+        # TODO: We can't use PolygonConfig.calendar because it gets these properties for start/end session.
         self.start_timestamp = (
             parse_date(start_date, calendar=self.calendar)
             if start_date
@@ -33,6 +58,8 @@ class PolygonConfig:
             if end_date
             else self.calendar.last_session
         )
+        self.agg_time = agg_time
+
         self.max_workers = None
         if environ.get("POLYGON_MAX_WORKERS", "").strip() != "":
             self.max_workers = int(environ.get("POLYGON_MAX_WORKERS"))
@@ -61,7 +88,8 @@ class PolygonConfig:
         # TODO: Restore non-recusive option.  Always recursive for now.
         self.csv_paths_pattern = environ.get(
             # "POLYGON_FLAT_FILES_CSV_PATTERN", "**/*.csv.gz"
-            "POLYGON_FLAT_FILES_CSV_PATTERN", "*.csv.gz"
+            "POLYGON_FLAT_FILES_CSV_PATTERN",
+            "*.csv.gz",
         )
         self.asset_files_dir = os.path.join(self.flat_files_dir, self.asset_subdir)
         self.minute_aggs_dir = os.path.join(self.asset_files_dir, "minute_aggs_v1")
@@ -70,51 +98,67 @@ class PolygonConfig:
         self.quotes_dir = os.path.join(self.asset_files_dir, "quotes_v1")
 
         # TODO: The "by ticker" files are temporary/intermediate and should/could be in the zipline data dir.
+        self.custom_asset_files_dir = environ.get(
+            "CUSTOM_ASSET_FILES_DIR", self.asset_files_dir
+        )
+
+        self.cache_dir = os.path.join(self.custom_asset_files_dir, "api_cache")
+
+        self.lock_file_path = os.path.join(self.custom_asset_files_dir, "ingest.lock")
+        self.custom_aggs_dates_path = os.path.join(self.custom_asset_files_dir, "aggs_dates.json")
+        self.by_ticker_dates_path = os.path.join(self.custom_asset_files_dir, "by_ticker_dates.json")
+
         self.minute_by_ticker_dir = os.path.join(
-            self.asset_files_dir, "minute_by_ticker_v1"
+            self.custom_asset_files_dir, "minute_by_ticker_v1"
+        )
+        self.day_by_ticker_dir = os.path.join(
+            self.custom_asset_files_dir, "day_by_ticker_v1"
         )
-        self.day_by_ticker_dir = os.path.join(self.asset_files_dir, "day_by_ticker_v1")
 
+        # If agg_time begins with a digit, it is a timedelta string and we're using custom aggs from trades.
         if bool(re.match(r"^\d", agg_time)):
             self.agg_timedelta = pd.to_timedelta(agg_time)
-            self.custom_asset_files_dir = environ.get(
-                "CUSTOM_ASSET_FILES_DIR", self.asset_files_dir
+            self.csv_files_dir = self.trades_dir
+            self.custom_aggs_name_format = environ.get(
+                "CUSTOM_AGGS_NAME_FORMAT", "{config.agg_timedelta.seconds}sec_aggs"
             )
-            self.custom_aggs_dir = os.path.join(
-                self.custom_asset_files_dir, custom_aggs_format.format(config=self)
+            self.aggs_dir = os.path.join(
+                self.custom_asset_files_dir,
+                self.custom_aggs_name_format.format(config=self),
             )
-            self.custom_aggs_by_ticker_dir = os.path.join(
+            self.by_ticker_dir = os.path.join(
                 self.custom_asset_files_dir,
-                (custom_aggs_format + "_by_ticker").format(config=self),
+                (self.custom_aggs_name_format + "_by_ticker").format(config=self),
             )
-            self.aggs_dir = self.custom_aggs_dir
-            self.by_ticker_dir = self.custom_aggs_by_ticker_dir
-        elif agg_time == "minute":
+        elif agg_time == AGG_TIME_MINUTE:
             self.agg_timedelta = pd.to_timedelta("1minute")
             self.aggs_dir = self.minute_aggs_dir
+            self.csv_files_dir = self.aggs_dir
             self.by_ticker_dir = self.minute_by_ticker_dir
-        elif agg_time == "day":
+        elif agg_time == AGG_TIME_DAY:
             self.agg_timedelta = pd.to_timedelta("1day")
             self.aggs_dir = self.day_aggs_dir
+            self.csv_files_dir = self.aggs_dir
             self.by_ticker_dir = self.day_by_ticker_dir
         else:
             raise ValueError(
                 f"agg_time must be 'minute', 'day', or a timedelta string; got '{agg_time=}'"
             )
-        self.agg_time = agg_time
 
         self.arrow_format = environ.get(
-            "POLYGON_ARROW_FORMAT", "parquet" if self.agg_time == "day" else "hive"
+            "POLYGON_ARROW_FORMAT", "parquet" if self.agg_time == AGG_TIME_DAY else "hive"
         )
         # self.by_ticker_hive_dir = os.path.join(
         #     self.by_ticker_dir,
         #     f"{self.agg_time}_{self.start_timestamp.date().isoformat()}_{self.end_timestamp.date().isoformat()}.hive",
         # )
-        self.cache_dir = os.path.join(self.asset_files_dir, "api_cache")
 
     @property
     def calendar(self):
-        return get_calendar(self.calendar_name)
+        # print call stack
+        # import traceback
+        # traceback.print_stack()
+        return get_calendar(self.calendar_name, start_session=self.start_date, end_session=self.end_date)
 
     def ticker_file_path(self, date: pd.Timestamp):
         ticker_year_dir = os.path.join(
@@ -129,18 +173,30 @@ class PolygonConfig:
         # TODO: Use csv_paths_pattern to remove the suffixes
         return os.path.basename(path).removesuffix(".gz").removesuffix(".csv")
 
+    def date_to_csv_file_path(self, date: datetime.date, ext=".csv.gz"):
+        return f"{self.csv_files_dir}/{date.strftime('%Y/%m/%Y-%m-%d') + ext}"
+
     @property
     def by_ticker_aggs_arrow_dir(self):
-        return os.path.join(self.by_ticker_dir,
-                            f"{self.start_timestamp.date().isoformat()}_{self.end_timestamp.date().isoformat()}.arrow")
+        # TODO: Don't split these up by ingestion range.  They're already time indexed.
+        # Only reason to separate them is if we're worried about (or want) data being different across ingestions.
+        # This scattering is really slow and is usually gonna be redundant.
+        # This wasn't a problem when start/end dates were the calendar bounds when omitted.
+        # Can't just drop this because concat_all_aggs_from_csv will skip if it exists.
+        # return os.path.join(
+        #     self.by_ticker_dir,
+        #     f"{self.start_timestamp.date().isoformat()}_{self.end_timestamp.date().isoformat()}.arrow",
+        #     # "aggs.arrow",
+        # )
+        return self.by_ticker_dir
 
     def api_cache_path(
-        self, start_date: Date, end_date: Date, filename: str, extension=".parquet"
+        self, first_day: pd.Timestamp, last_day: pd.Timestamp, filename: str, extension=".parquet"
     ):
-        start_str = parse_date(start_date, calendar=self.calendar).date().isoformat()
-        end_str = parse_date(end_date, calendar=self.calendar).date().isoformat()
+        first_day_str = first_day.date().isoformat()
+        last_day_str = last_day.date().isoformat()
         return os.path.join(
-            self.cache_dir, f"{start_str}_{end_str}/{filename}{extension}"
+            self.cache_dir, f"{first_day_str}_{last_day_str}/{filename}{extension}"
         )
 
     def csv_paths(self) -> Iterator[str]:
@@ -153,13 +209,15 @@ class PolygonConfig:
                 for filename in sorted(filenames):
                     yield os.path.join(root, filename)
 
-    def find_first_and_last_aggs(self) -> Tuple[str, str]:
+    def find_first_and_last_aggs(
+        self, aggs_dir, file_pattern
+    ) -> Tuple[str | None, str | None]:
         # Find the path to the lexically first and last paths in aggs_dir that matches csv_paths_pattern.
         # Would like to use Path.walk(top_down=True) but it is only availble in Python 3.12+.
         # This needs to be efficient because it is called on every init, even though we only need it for ingest.
         # But we can't call it in ingest because the writer initializes and writes the metadata before it is called.
         paths = []
-        for root, dirnames, filenames in os.walk(self.aggs_dir, topdown=True):
+        for root, dirnames, filenames in os.walk(aggs_dir, topdown=True):
             if dirnames:
                 # We only want first and last in each directory.
                 sorted_dirs = sorted(dirnames)
@@ -169,15 +227,15 @@ class PolygonConfig:
                     else sorted_dirs
                 )
             # Filter out filenames that don't match the pattern.
-            filenames = fnmatch.filter(filenames, self.csv_paths_pattern)
+            filenames = fnmatch.filter(filenames, file_pattern)
             if filenames:
                 filenames = sorted(filenames)
                 paths.append(os.path.join(root, filenames[0]))
                 if len(filenames) > 1:
                     paths.append(os.path.join(root, filenames[-1]))
+        if not paths:
+            return None, None
         paths = sorted(paths)
-        if len(paths) < 2:
-            raise ValueError(f"Need more than one aggs file but found {len(paths)} paths in {self.aggs_dir}")
         return self.file_path_to_name(paths[0]), self.file_path_to_name(paths[-1])
 
 

zipline_polygon_bundle/nyse_all_hours_calendar.py

@@ -0,0 +1,25 @@
+import datetime
+from exchange_calendars.calendar_utils import get_calendar_names, register_calendar_type
+from exchange_calendars.exchange_calendar_xnys import XNYSExchangeCalendar
+
+
+NYSE_ALL_HOURS = "NYSE_ALL_HOURS"
+
+
+class USExtendedHoursExchangeCalendar(XNYSExchangeCalendar):
+    """
+    A calendar for extended hours which runs from 4 AM to 8 PM.
+    """
+
+    name = NYSE_ALL_HOURS
+
+    open_times = ((None, datetime.time(4)),)
+
+    close_times = ((None, datetime.time(20)),)
+
+    regular_early_close = datetime.time(13)
+
+
+def register_nyse_all_hours_calendar():
+    if NYSE_ALL_HOURS not in get_calendar_names():
+        register_calendar_type(NYSE_ALL_HOURS, USExtendedHoursExchangeCalendar)

zipline_polygon_bundle/tickers_and_names.py

@@ -48,7 +48,10 @@ class PolygonAssets:
         active: bool = True,
     ):
         response = self.polygon_client.list_tickers(
-            market=self.config.market, active=active, date=date.date(), limit=500
+            market=self.config.market,
+            active=active,
+            date=date.date().isoformat(),
+            limit=500,
         )
         tickers_df = pd.DataFrame(list(response))
         # The currency info is for crypto.  The source_feed is always NA.