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

zipline_polygon_bundle/config.py

@@ -1,8 +1,38 @@
-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
 
-import os
+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:
@@ -10,33 +40,37 @@ class PolygonConfig:
         self,
         environ: dict,
         calendar_name: str,
-        start_session: Date,
-        end_session: Date,
-        agg_time: str = "day",
+        start_date: Date,
+        end_date: Date,
+        agg_time: str = AGG_TIME_DAY,
     ):
-        if agg_time not in ["minute", "day"]:
-            raise ValueError(f"agg_time must be 'minute' or 'day', got '{agg_time}'")
         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_session, calendar=self.calendar)
-            if start_session
+            parse_date(start_date, calendar=self.calendar)
+            if start_date
             else self.calendar.first_session
         )
         self.end_timestamp = (
-            parse_date(end_session, calendar=self.calendar)
-            if end_session
+            parse_date(end_date, calendar=self.calendar)
+            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"))
         self.api_key = environ.get("POLYGON_API_KEY")
+        self.filesystem = LocalFileSystem()
         self.data_dir = environ.get("POLYGON_DATA_DIR", "data/files.polygon.io")
         self.cik_cusip_mapping_csv_path = environ.get(
             "CIK_CUSIP_MAPS_CSV", os.path.join(self.data_dir, "cik-cusip-maps.csv")
         )
-        self.asset_subdir = environ.get("POLYGON_ASSET_SUBDIR", "us_stocks_sip")
         self.market = environ.get("POLYGON_MARKET", "stocks")
+        self.asset_subdir = environ.get("POLYGON_ASSET_SUBDIR", "us_stocks_sip")
         self.tickers_dir = environ.get(
             "POLYGON_TICKERS_DIR",
             os.path.join(os.path.join(self.data_dir, "tickers"), self.asset_subdir),
@@ -51,34 +85,80 @@ class PolygonConfig:
         self.flat_files_dir = environ.get(
             "POLYGON_FLAT_FILES_DIR", os.path.join(self.data_dir, "flatfiles")
         )
-        self.csv_paths_pattern = environ.get("POLYGON_FLAT_FILES_CSV_PATTERN", "**/*.csv.gz")
-        self.agg_time = agg_time
+        # 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",
+        )
         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")
         self.day_aggs_dir = os.path.join(self.asset_files_dir, "day_aggs_v1")
-        self.aggs_dir = (
-            self.minute_aggs_dir if self.agg_time == "minute" else self.day_aggs_dir
-        )
+        self.trades_dir = os.path.join(self.asset_files_dir, "trades_v1")
+        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.asset_files_dir, "day_by_ticker_v1")
-        self.by_ticker_dir = (
-            self.minute_by_ticker_dir
-            if self.agg_time == "minute"
-            else self.day_by_ticker_dir
+        self.day_by_ticker_dir = os.path.join(
+            self.custom_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.csv_files_dir = self.trades_dir
+            self.custom_aggs_name_format = environ.get(
+                "CUSTOM_AGGS_NAME_FORMAT", "{config.agg_timedelta.seconds}sec_aggs"
+            )
+            self.aggs_dir = os.path.join(
+                self.custom_asset_files_dir,
+                self.custom_aggs_name_format.format(config=self),
+            )
+            self.by_ticker_dir = os.path.join(
+                self.custom_asset_files_dir,
+                (self.custom_aggs_name_format + "_by_ticker").format(config=self),
+            )
+        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 == 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.arrow_format = environ.get(
+            "POLYGON_ARROW_FORMAT", "parquet" if self.agg_time == AGG_TIME_DAY else "hive"
         )
-        self.arrow_format = environ.get("POLYGON_ARROW_FORMAT", "parquet" if self.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(
@@ -88,25 +168,76 @@ class PolygonConfig:
         return os.path.join(
             ticker_year_dir, f"tickers_{date.date().isoformat()}.parquet"
         )
-    
+
     def file_path_to_name(self, path: str):
+        # TODO: Use csv_paths_pattern to remove the suffixes
         return os.path.basename(path).removesuffix(".gz").removesuffix(".csv")
 
-    def by_ticker_aggs_arrow_dir(self, first_path: str, last_path: str):
-        return os.path.join(
-            self.by_ticker_dir,
-            f"{self.file_path_to_name(first_path)}_{self.file_path_to_name(last_path)}.arrow",
-        )
+    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):
+        # 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]:
+        for root, dirnames, filenames in os.walk(self.aggs_dir, topdown=True):
+            if dirnames:
+                dirnames[:] = sorted(dirnames)
+            # Filter out filenames that don't match the pattern.
+            filenames = fnmatch.filter(filenames, self.csv_paths_pattern)
+            if filenames:
+                for filename in sorted(filenames):
+                    yield os.path.join(root, filename)
+
+    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(aggs_dir, topdown=True):
+            if dirnames:
+                # We only want first and last in each directory.
+                sorted_dirs = sorted(dirnames)
+                dirnames[:] = (
+                    [sorted_dirs[0], sorted_dirs[-1]]
+                    if len(sorted_dirs) > 1
+                    else sorted_dirs
+                )
+            # Filter out filenames that don't match the 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)
+        return self.file_path_to_name(paths[0]), self.file_path_to_name(paths[-1])
+
 
 if __name__ == "__main__":
     config = PolygonConfig(os.environ, "XNYS", "2003-10-01", "2023-01-01")

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/polygon_file_reader.py

@@ -99,6 +99,6 @@ def process_all_csv_to_parquet(
 if __name__ == "__main__":
     # os.environ["POLYGON_DATA_DIR"] = "/Volumes/Oahu/Mirror/files.polygon.io"
     config = PolygonConfig(
-        environ=os.environ, calendar_name="XNYS", start_session=None, end_session=None
+        environ=os.environ, calendar_name="XNYS", start_date=None, end_date=None
     )
     process_all_csv_to_parquet(config.aggs_dir)

zipline_polygon_bundle/process_all_aggs.py

@@ -74,8 +74,8 @@ if __name__ == "__main__":
     config = PolygonConfig(
         environ=os.environ,
         calendar_name="XNYS",
-        start_session="2020-10-07",
-        end_session="2020-10-15",
+        start_date="2020-10-07",
+        end_date="2020-10-15",
     )
     print(f"{config.aggs_dir=}")
     max_ticker_lens = apply_to_all_aggs(

zipline_polygon_bundle/quotes.py

@@ -0,0 +1,101 @@
+from .config import PolygonConfig
+from .trades import cast_strings_to_list
+
+import os
+
+import pyarrow as pa
+from pyarrow import dataset as pa_ds
+from pyarrow import compute as pa_compute
+from pyarrow import fs as pa_fs
+from fsspec.implementations.arrow import ArrowFSWrapper
+from pyarrow import csv as pa_csv
+
+
+def quotes_schema(raw: bool = False) -> pa.Schema:
+    # There is some problem reading the timestamps as timestamps so we have to read as integer then change the schema.
+    # I make the timestamp timezone-aware because that's how Unix timestamps work and it may help avoid mistakes.
+    # timestamp_type = pa.timestamp("ns", tz="UTC")
+    timestamp_type = pa.int64() if raw else pa.timestamp("ns", tz="UTC")
+
+    # Polygon price scale is 4 decimal places (i.e. hundredths of a penny), but we'll use 10 because we have precision to spare.
+    # price_type = pa.decimal128(precision=38, scale=10)
+    # 64bit float a little overkill but avoids any plausible truncation error.
+    price_type = pa.float64()
+
+    # ticker: string
+    # ask_exchange: int64
+    # ask_price: double
+    # ask_size: int64
+    # bid_exchange: int64
+    # bid_price: double
+    # bid_size: int64
+    # conditions: string
+    # indicators: int64
+    # participant_timestamp: int64
+    # sequence_number: int64
+    # sip_timestamp: int64
+    # tape: int64
+    # trf_timestamp: int64
+
+    return pa.schema(
+            [
+                pa.field("ticker", pa.string(), nullable=False),
+                pa.field("ask_exchange", pa.int8(), nullable=False),
+                pa.field("ask_price", price_type, nullable=False),
+                pa.field("ask_size", pa.int64(), nullable=False),
+                pa.field("bid_exchange", pa.int8(), nullable=False),
+                pa.field("bid_price", price_type, nullable=False),
+                pa.field("bid_size", pa.int64(), nullable=False),
+                pa.field("conditions", pa.string(), nullable=False),
+                pa.field("indicators", pa.string(), nullable=False),
+                pa.field("participant_timestamp", timestamp_type, nullable=False),
+                pa.field("sequence_number", pa.int64(), nullable=False),
+                pa.field("sip_timestamp", timestamp_type, nullable=False),
+                pa.field("tape", pa.int8(), nullable=False),
+                pa.field("trf_timestamp", timestamp_type, nullable=False),
+            ]
+        )
+
+
+def quotes_dataset(config: PolygonConfig) -> pa_ds.Dataset:
+    """
+    Create a pyarrow dataset from the quotes files.
+    """
+
+    # https://arrow.apache.org/docs/python/filesystems.html#using-arrow-filesystems-with-fsspec
+    # https://filesystem-spec.readthedocs.io/en/latest/_modules/fsspec/spec.html#AbstractFileSystem.glob.
+    fsspec = ArrowFSWrapper(config.filesystem)
+
+    # We sort by path because they have the year and month in the dir names and the date in the filename.
+    paths = sorted(
+        fsspec.glob(os.path.join(config.quotes_dir, config.csv_paths_pattern))
+    )
+
+    return pa_ds.FileSystemDataset.from_paths(paths,
+                                              format=pa_ds.CsvFileFormat(),
+                                              schema=quotes_schema(raw=True),
+                                              filesystem=config.filesystem)
+
+
+def cast_strings_to_list(string_array, separator=",", default="0", value_type=pa.uint8()):
+    """Cast a PyArrow StringArray of comma-separated numbers to a ListArray of values."""
+
+    # Create a mask to identify empty strings
+    is_empty = pa_compute.equal(pa_compute.utf8_trim_whitespace(string_array), "")
+
+    # Use replace_with_mask to replace empty strings with the default ("0")
+    filled_column = pa_compute.replace_with_mask(string_array, is_empty, pa.scalar(default))
+
+    # Split the strings by comma
+    split_array = pa_compute.split_pattern(filled_column, pattern=separator)
+
+    # Cast each element in the resulting lists to integers
+    int_list_array = pa_compute.cast(split_array, pa.list_(value_type))
+
+    return int_list_array
+
+
+def cast_quotes(quotes):
+    quotes = quotes.cast(quotes_schema())
+    condition_values = cast_strings_to_list(quotes.column("conditions").combine_chunks())
+    return quotes.append_column('condition_values', condition_values)

zipline_polygon_bundle/tickers_and_names.py

@@ -3,6 +3,7 @@ from .config import PolygonConfig
 import datetime
 import os
 import pandas as pd
+import csv
 import polygon
 import logging
 from concurrent.futures import ProcessPoolExecutor
@@ -47,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.
@@ -383,40 +387,3 @@ def get_ticker_universe(config: PolygonConfig, fetch_missing: bool = False):
     merged_tickers = pd.read_parquet(parquet_path)
     merged_tickers.info()
     return merged_tickers
-
-
-# Initialize ticker files in __main__.  Use CLI args to specify start and end dates.
-if __name__ == "__main__":
-    import argparse
-
-    parser = argparse.ArgumentParser(description="Initialize ticker files.")
-    parser.add_argument(
-        "--start-date",
-        type=str,
-        help="Start date in ISO format (YYYY-MM-DD)",
-        default="2014-05-01",
-    )
-    parser.add_argument(
-        "--end-date",
-        type=str,
-        help="End date in ISO format (YYYY-MM-DD)",
-        default="2024-04-01",
-    )
-    args = parser.parse_args()
-
-    start_date = (
-        datetime.datetime.strptime(args.start_date, "%Y-%m-%d").date()
-        if args.start_date
-        else datetime.date.today()
-    )
-    end_date = (
-        datetime.datetime.strptime(args.end_date, "%Y-%m-%d").date()
-        if args.end_date
-        else datetime.date.today()
-    )
-
-    all_tickers = load_all_tickers(start_date, end_date, fetch_missing=True)
-    merged_tickers = merge_tickers(all_tickers)
-    merged_tickers.to_csv(f"data/tickers/us_tickers_{start_date}-{end_date}.csv")
-    ticker_names = ticker_names_from_merged_tickers(merged_tickers)
-    print(ticker_names)