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

zipline_polygon_bundle/__init__.py

@@ -6,10 +6,11 @@ from .bundle import (
 )
 
 from .config import PolygonConfig
+from .nyse_all_hours_calendar import NYSE_ALL_HOURS, register_nyse_all_hours_calendar
 from .concat_all_aggs import concat_all_aggs_from_csv, generate_csv_agg_tables
 from .adjustments import load_splits, load_dividends, load_conditions
 from .trades import trades_schema, trades_dataset, cast_trades, date_to_path
-from .trades import custom_aggs_partitioning, custom_aggs_schema, trades_to_custom_aggs, convert_all_to_custom_aggs
+from .trades import custom_aggs_partitioning, custom_aggs_schema, trades_to_custom_aggs, convert_trades_to_custom_aggs
 from .trades import get_custom_aggs_dates, generate_csv_trades_tables, compute_signals_for_all_custom_aggs
 from .quotes import quotes_schema, quotes_dataset, cast_quotes
 # from .tickers_and_names import load_all_tickers, merge_tickers, ticker_names_from_merged_tickers, get_ticker_universe
@@ -18,6 +19,8 @@ from .tickers_and_names import PolygonAssets, get_ticker_universe
 
 __all__ = [
     "register_polygon_equities_bundle",
+    "register_nyse_all_hours_calendar",
+    "NYSE_ALL_HOURS",
     "symbol_to_upper",
     "polygon_equities_bundle_day",
     "polygon_equities_bundle_minute",
@@ -36,7 +39,7 @@ __all__ = [
     "custom_aggs_partitioning",
     "custom_aggs_schema",
     "trades_to_custom_aggs",
-    "convert_all_to_custom_aggs",
+    "convert_trades_to_custom_aggs",
     "compute_signals_for_all_custom_aggs",
     "quotes_schema",
     "quotes_dataset",

zipline_polygon_bundle/bundle.py

@@ -3,11 +3,13 @@ from zipline.data.bundles import register
 from zipline.data.resample import minute_frame_to_session_frame
 
 from exchange_calendars.calendar_helpers import parse_date
-from zipline.utils.calendar_utils import get_calendar
+from exchange_calendars.calendar_utils import get_calendar
 
-from .config import PolygonConfig
 from .concat_all_aggs import concat_all_aggs_from_csv, generate_csv_agg_tables
 from .adjustments import load_splits, load_dividends
+from .config import PolygonConfig
+from .nyse_all_hours_calendar import register_nyse_all_hours_calendar
+from .trades import convert_trades_to_custom_aggs, scatter_custom_aggs_to_by_ticker
 
 import pyarrow
 import pyarrow.compute
@@ -29,7 +31,7 @@ def symbol_to_upper(s: str) -> str:
 def generate_all_agg_tables_from_csv(
     config: PolygonConfig,
 ):
-    paths, schema, tables = generate_csv_agg_tables(config)
+    schema, tables = generate_csv_agg_tables(config)
     for table in tables:
         table = table.sort_by([("ticker", "ascending"), ("window_start", "ascending")])
         yield table
@@ -209,7 +211,19 @@ def polygon_equities_bundle_day(
         )
     )
 
-    table = aggregates.to_table()
+    # Only get the columns Zipline allows.
+    table = aggregates.to_table(
+        columns=[
+            "ticker",
+            "window_start",
+            "open",
+            "high",
+            "low",
+            "close",
+            "volume",
+            "transactions",
+        ]
+    )
     table = rename_polygon_to_zipline(table, "day")
     # Get all the symbols in the table by using value_counts to tabulate the unique values.
     # pyarrow.Table.column returns a pyarrow.ChunkedArray.
@@ -254,7 +268,19 @@ def process_minute_fragment(
     dates_with_data: set,
     agg_time: str,
 ):
-    table = fragment.to_table()
+    # Only get the columns Zipline allows.
+    table = fragment.to_table(
+        columns=[
+            "ticker",
+            "window_start",
+            "open",
+            "high",
+            "low",
+            "close",
+            "volume",
+            "transactions",
+        ]
+    )
     print(f" {table.num_rows=}")
     table = rename_polygon_to_zipline(table, "timestamp")
     table = table.sort_by([("symbol", "ascending"), ("timestamp", "ascending")])
@@ -486,6 +512,97 @@ def polygon_equities_bundle_minute(
     adjustment_writer.write(splits=splits, dividends=dividends)
 
 
+def polygon_equities_bundle_trades(
+    environ,
+    asset_db_writer,
+    minute_bar_writer,
+    daily_bar_writer,
+    adjustment_writer,
+    calendar,
+    start_date,
+    end_date,
+    cache,
+    show_progress,
+    output_dir,
+):
+    # TODO: Support agg durations other than `1min`.
+    config = PolygonConfig(
+        environ=environ,
+        calendar_name=calendar.name,
+        start_date=start_date,
+        end_date=end_date,
+        agg_time="1min",
+    )
+
+    convert_trades_to_custom_aggs(config, overwrite=False)
+    by_ticker_aggs_arrow_dir = scatter_custom_aggs_to_by_ticker(config)
+    aggregates = pyarrow.dataset.dataset(by_ticker_aggs_arrow_dir)
+    # 3.5 billion rows for 10 years of minute data.
+    # print(f"{aggregates.count_rows()=}")
+    # Can't sort the dataset because that reads it all into memory.
+    # aggregates = aggregates.sort_by([("ticker", "ascending"), ("window_start", "ascending")])
+    # print("Sorted")
+
+    # Zipline uses case-insensitive symbols, so we need to convert them to uppercase with a ^ prefix when lowercase.
+    # This is because the SQL schema zipline uses for symbols ignores case.
+    # We put the original symbol in the asset_name field.
+    metadata = pd.DataFrame(
+        columns=(
+            "start_date",
+            "end_date",
+            "auto_close_date",
+            "symbol",
+            "exchange",
+            "asset_name",
+        )
+    )
+
+    symbol_to_sid = {}
+    dates_with_data = set()
+
+    # Get data for all stocks and write to Zipline
+    daily_bar_writer.write(
+        process_minute_aggregates(
+            fragments=aggregates.get_fragments(),
+            sessions=calendar.sessions_in_range(start_date, end_date),
+            minutes=calendar.sessions_minutes(start_date, end_date),
+            metadata=metadata,
+            calendar=calendar,
+            symbol_to_sid=symbol_to_sid,
+            dates_with_data=dates_with_data,
+            agg_time="day",
+        ),
+        show_progress=show_progress,
+    )
+
+    # Get data for all stocks and write to Zipline
+    minute_bar_writer.write(
+        process_minute_aggregates(
+            fragments=aggregates.get_fragments(),
+            sessions=calendar.sessions_in_range(start_date, end_date),
+            minutes=calendar.sessions_minutes(start_date, end_date),
+            metadata=metadata,
+            calendar=calendar,
+            symbol_to_sid=symbol_to_sid,
+            dates_with_data=dates_with_data,
+            agg_time="minute",
+        ),
+        show_progress=show_progress,
+    )
+
+    # Write the metadata
+    asset_db_writer.write(equities=metadata)
+
+    # Load splits and dividends
+    first_start_end = min(dates_with_data)
+    last_end_date = max(dates_with_data)
+    splits = load_splits(config, first_start_end, last_end_date, symbol_to_sid)
+    dividends = load_dividends(config, first_start_end, last_end_date, symbol_to_sid)
+
+    # Write splits and dividends
+    adjustment_writer.write(splits=splits, dividends=dividends)
+
+
 def register_polygon_equities_bundle(
     bundlename,
     start_date=None,
@@ -496,10 +613,15 @@ def register_polygon_equities_bundle(
     # watchlists=None,
     # include_asset_types=None,
 ):
-    if agg_time not in ["day", "minute"]:
-        raise ValueError(f"agg_time must be 'day' or 'minute', not '{agg_time}'")
+    register_nyse_all_hours_calendar()
+
+    if agg_time not in ["day", "minute", "1min"]:
+        raise ValueError(
+            f"agg_time must be 'day', 'minute' (aggs), or '1min' (trades), not '{agg_time}'"
+        )
+
     # We need to know the start and end dates of the session before the bundle is
-    # registered because even though we only need it for ingest, the metadata in 
+    # registered because even though we only need it for ingest, the metadata in
     # the writer is initialized and written before our ingest function is called.
     if start_date is None or end_date is None:
         config = PolygonConfig(
@@ -509,23 +631,28 @@ def register_polygon_equities_bundle(
             end_date=end_date,
             agg_time=agg_time,
         )
-        first_aggs_date, last_aggs_date = config.find_first_and_last_aggs()
+        first_aggs_date, last_aggs_date = config.find_first_and_last_aggs(
+            config.aggs_dir if agg_time in ["day", "minute"] else config.trades_dir,
+            config.csv_paths_pattern,
+        )
         if start_date is None:
             start_date = first_aggs_date
         if end_date is None:
             end_date = last_aggs_date
 
-    calendar = get_calendar(calendar_name)
-
     register(
         bundlename,
         (
-            polygon_equities_bundle_minute
-            if agg_time == "minute"
-            else polygon_equities_bundle_day
+            polygon_equities_bundle_day
+            if agg_time == "day"
+            else (
+                polygon_equities_bundle_minute
+                if agg_time == "minute"
+                else polygon_equities_bundle_trades
+            )
         ),
-        start_session=parse_date(start_date, calendar=calendar),
-        end_session=parse_date(end_date, calendar=calendar),
+        start_session=parse_date(start_date, raise_oob=False) if start_date else None,
+        end_session=parse_date(end_date, raise_oob=False) if end_date else None,
         calendar_name=calendar_name,
         # minutes_per_day=390,
         # create_writers=True,

zipline_polygon_bundle/concat_all_aggs.py

@@ -1,10 +1,9 @@
-from .config import PolygonConfig
+from .config import PolygonConfig, PARTITION_COLUMN_NAME, to_partition_key
 
 import shutil
 from typing import Iterator, Tuple, List, Union
 
 import argparse
-import glob
 import os
 
 import pyarrow as pa
@@ -15,25 +14,6 @@ from pyarrow import compute as pa_compute
 import pandas as pd
 
 
-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
-
-
 def generate_tables_from_csv_files(
     paths: Iterator[Union[str, os.PathLike]],
     schema: pa.Schema,

zipline_polygon_bundle/config.py

@@ -1,7 +1,9 @@
-from exchange_calendars.calendar_helpers import Date, parse_date, parse_timestamp
-from zipline.utils.calendar_utils import get_calendar
+from exchange_calendars.calendar_helpers import Date, parse_date
+from exchange_calendars.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
@@ -10,6 +12,25 @@ import re
 import fnmatch
 
 
+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__(
         self,
@@ -61,7 +82,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")
@@ -114,7 +136,10 @@ class PolygonConfig:
 
     @property
     def calendar(self):
-        return get_calendar(self.calendar_name)
+        # If you don't give a start date you'll only get 20 years from today.
+        if self.calendar_name in [NYSE_ALL_HOURS, "us_futures", "CMES", "XNYS", "NYSE"]:
+            return get_calendar(self.calendar_name, side="right", start=pd.Timestamp("1990-01-01"))
+        return get_calendar(self.calendar_name, side="right")
 
     def ticker_file_path(self, date: pd.Timestamp):
         ticker_year_dir = os.path.join(
@@ -131,8 +156,16 @@ class PolygonConfig:
 
     @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",
+        )
 
     def api_cache_path(
         self, start_date: Date, end_date: Date, filename: str, extension=".parquet"
@@ -153,13 +186,13 @@ 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 +202,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.