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

zipline_polygon_bundle/__init__.py

@@ -6,11 +6,21 @@ 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
+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_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
+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",
@@ -19,4 +29,24 @@ __all__ = [
     "generate_csv_agg_tables",
     "load_splits",
     "load_dividends",
+    "load_conditions",
+    "trades_schema",
+    "trades_dataset",
+    "cast_trades",
+    "date_to_path",
+    "get_custom_aggs_dates",
+    "generate_csv_trades_tables",
+    "custom_aggs_partitioning",
+    "custom_aggs_schema",
+    "trades_to_custom_aggs",
+    "convert_trades_to_custom_aggs",
+    "compute_signals_for_all_custom_aggs",
+    "quotes_schema",
+    "quotes_dataset",
+    "cast_quotes",
+    # "load_all_tickers", 
+    # "merge_tickers",
+    # "ticker_names_from_merged_tickers",
+    "PolygonAssets",
+    "get_ticker_universe",
 ]

zipline_polygon_bundle/adjustments.py

@@ -153,3 +153,37 @@ def load_dividends(
     return dividends[
         ["sid", "ex_date", "declared_date", "record_date", "pay_date", "amount"]
     ]
+
+
+def load_conditions(config: PolygonConfig) -> pd.DataFrame:
+    # The API doesn't use dates for the condition codes but this is a way to provide control over caching.
+    # Main thing is to get the current conditions list but we don't want to call more than once a day.
+    conditions_path = config.api_cache_path(
+        start_date=config.start_timestamp.date(), end_date=config.end_timestamp.date(), filename="conditions"
+    )
+    expected_conditions_count = 100
+    if not os.path.exists(conditions_path):
+        client = polygon.RESTClient(api_key=config.api_key)
+        conditions_response = client.list_conditions(
+            limit=1000,
+        )
+        if conditions_response is HTTPResponse:
+            raise ValueError(f"Polygon.list_splits bad HTTPResponse: {conditions_response}")
+        conditions = pd.DataFrame(conditions_response)
+        print(f"Got {len(conditions)=} from Polygon list_conditions.")
+        os.makedirs(os.path.dirname(conditions_path), exist_ok=True)
+        conditions.to_parquet(conditions_path)
+        if len(conditions) < expected_conditions_count:
+            logging.warning(
+                f"Only got {len(conditions)=} from Polygon list_splits (expected {expected_conditions_count=}).  "
+            )
+        # We will always load from the file to avoid any chance of weird errors.
+    if os.path.exists(conditions_path):
+        conditions = pd.read_parquet(conditions_path)
+        print(f"Loaded {len(conditions)=} from {conditions_path}")
+        if len(conditions) < expected_conditions_count:
+            logging.warning(
+                f"Only got {len(conditions)=} from cached conditions (expected {expected_conditions_count=}).  "
+            )
+        return conditions
+    raise ValueError(f"Failed to load splits from {conditions_path}")

zipline_polygon_bundle/bundle.py

@@ -1,18 +1,23 @@
+import os
 from zipline.data.bundles import register
 from zipline.data.resample import minute_frame_to_session_frame
 
-from .config import PolygonConfig
+from exchange_calendars.calendar_helpers import parse_date
+from exchange_calendars.calendar_utils import get_calendar
+
 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
+import pyarrow.dataset
 
 import pandas as pd
 import logging
 
-import concurrent.futures
-
 
 # TODO: Change warnings to be relative to number of days in the range.
 
@@ -26,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
@@ -175,8 +180,8 @@ def polygon_equities_bundle_day(
     daily_bar_writer,
     adjustment_writer,
     calendar,
-    start_session,
-    end_session,
+    start_date,
+    end_date,
     cache,
     show_progress,
     output_dir,
@@ -184,8 +189,8 @@ def polygon_equities_bundle_day(
     config = PolygonConfig(
         environ=environ,
         calendar_name=calendar.name,
-        start_session=start_session,
-        end_session=end_session,
+        start_date=start_date,
+        end_date=end_date,
         agg_time="day",
     )
 
@@ -206,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.
@@ -219,7 +236,7 @@ def polygon_equities_bundle_day(
     daily_bar_writer.write(
         process_day_aggregates(
             table=table,
-            sessions=calendar.sessions_in_range(start_session, end_session),
+            sessions=calendar.sessions_in_range(start_date, end_date),
             metadata=metadata,
             calendar=calendar,
             symbol_to_sid=symbol_to_sid,
@@ -251,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")])
@@ -400,8 +429,8 @@ def polygon_equities_bundle_minute(
     daily_bar_writer,
     adjustment_writer,
     calendar,
-    start_session,
-    end_session,
+    start_date,
+    end_date,
     cache,
     show_progress,
     output_dir,
@@ -409,8 +438,8 @@ def polygon_equities_bundle_minute(
     config = PolygonConfig(
         environ=environ,
         calendar_name=calendar.name,
-        start_session=start_session,
-        end_session=end_session,
+        start_date=start_date,
+        end_date=end_date,
         agg_time="minute",
     )
 
@@ -444,8 +473,99 @@ def polygon_equities_bundle_minute(
     daily_bar_writer.write(
         process_minute_aggregates(
             fragments=aggregates.get_fragments(),
-            sessions=calendar.sessions_in_range(start_session, end_session),
-            minutes=calendar.sessions_minutes(start_session, end_session),
+            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 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,
@@ -459,8 +579,8 @@ def polygon_equities_bundle_minute(
     minute_bar_writer.write(
         process_minute_aggregates(
             fragments=aggregates.get_fragments(),
-            sessions=calendar.sessions_in_range(start_session, end_session),
-            minutes=calendar.sessions_minutes(start_session, end_session),
+            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,
@@ -485,25 +605,54 @@ def polygon_equities_bundle_minute(
 
 def register_polygon_equities_bundle(
     bundlename,
-    start_session=None,
-    end_session=None,
+    start_date=None,
+    end_date=None,
     calendar_name="XNYS",
     agg_time="day",
     # ticker_list=None,
     # 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
+    # the writer is initialized and written before our ingest function is called.
+    if start_date is None or end_date is None:
+        config = PolygonConfig(
+            environ=os.environ,
+            calendar_name=calendar_name,
+            start_date=start_date,
+            end_date=end_date,
+            agg_time=agg_time,
+        )
+        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
+
     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=start_session,
-        end_session=end_session,
+        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,
@@ -517,12 +666,12 @@ def register_polygon_equities_bundle(
 #     config = PolygonConfig(
 #         environ=os.environ,
 #         calendar_name="XNYS",
-#         # start_session="2003-10-01",
-#         # start_session="2018-01-01",
-#         start_session="2023-01-01",
-#         # end_session="2023-01-12",
-#         end_session="2023-12-31",
-#         # end_session="2024-06-30",
+#         # start_date="2003-10-01",
+#         # start_date="2018-01-01",
+#         start_date="2023-01-01",
+#         # end_date="2023-01-12",
+#         end_date="2023-12-31",
+#         # end_date="2024-06-30",
 #     )
 #     splits = load_polygon_splits(config)
 #     splits.info()

zipline_polygon_bundle/concat_all_aggs.py

@@ -1,40 +1,21 @@
-from .config import PolygonConfig
+from .config import PolygonConfig, PARTITION_COLUMN_NAME, to_partition_key
 
 import shutil
-from typing import Iterator, Tuple
+from typing import Iterator, Tuple, List, Union
 
 import argparse
-import glob
 import os
 
 import pyarrow as pa
 from pyarrow import dataset as pa_ds
 from pyarrow import csv as pa_csv
+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: list,
+    paths: Iterator[Union[str, os.PathLike]],
     schema: pa.Schema,
     start_timestamp: pd.Timestamp,
     limit_timestamp: pd.Timestamp,
@@ -57,7 +38,7 @@ def generate_tables_from_csv_files(
             quoted_strings_can_be_null=False,
         )
 
-        table = pa.csv.read_csv(path, convert_options=convert_options)
+        table = pa_csv.read_csv(path, convert_options=convert_options)
         tables_read_count += 1
         table = table.set_column(
             table.column_names.index("window_start"),
@@ -75,10 +56,10 @@ def generate_tables_from_csv_files(
                 ),
             )
         expr = (
-            pa.compute.field("window_start")
+            pa_compute.field("window_start")
             >= pa.scalar(start_timestamp, type=schema.field("window_start").type)
         ) & (
-            pa.compute.field("window_start")
+            pa_compute.field("window_start")
             < pa.scalar(
                 limit_timestamp,
                 type=schema.field("window_start").type,
@@ -101,22 +82,8 @@ def generate_tables_from_csv_files(
 
 def generate_csv_agg_tables(
     config: PolygonConfig,
-) -> Tuple[list[str], pa.Schema, Iterator[pa.Table]]:
+) -> Tuple[pa.Schema, Iterator[pa.Table]]:
     """zipline does bundle ingestion one ticker at a time."""
-    # We sort by path because they have the year and month in the dir names and the date in the filename.
-    paths = sorted(
-        list(
-            glob.glob(
-                os.path.join(config.aggs_dir, config.csv_paths_pattern),
-                recursive="**" in config.csv_paths_pattern,
-            )
-        )
-    )
-
-    print(f"{len(paths)=}")
-    if len(paths) > 0:
-        print(f"{paths[0]=}")
-        print(f"{paths[-1]=}")
 
     # Polygon Aggregate flatfile timestamps are in nanoseconds (like trades), not milliseconds as the docs say.
     # I make the timestamp timezone-aware because that's how Unix timestamps work and it may help avoid mistakes.
@@ -154,11 +121,11 @@ def generate_csv_agg_tables(
             pa.field(PARTITION_COLUMN_NAME, pa.string(), nullable=False)
         )
 
+    # TODO: Use generator like os.walk for paths.
     return (
-        paths,
         polygon_aggs_schema,
         generate_tables_from_csv_files(
-            paths=paths,
+            paths=config.csv_paths(),
             schema=polygon_aggs_schema,
             start_timestamp=config.start_timestamp,
             limit_timestamp=config.end_timestamp + pd.to_timedelta(1, unit="day"),
@@ -176,11 +143,9 @@ def concat_all_aggs_from_csv(
     config: PolygonConfig,
     overwrite: bool = False,
 ) -> str:
-    paths, schema, tables = generate_csv_agg_tables(config)
+    schema, tables = generate_csv_agg_tables(config)
 
-    if len(paths) < 1:
-        raise ValueError(f"No Polygon CSV flat files found in {config.aggs_dir=}")
-    by_ticker_aggs_arrow_dir = config.by_ticker_aggs_arrow_dir(paths[0], paths[-1])
+    by_ticker_aggs_arrow_dir = config.by_ticker_aggs_arrow_dir
     if os.path.exists(by_ticker_aggs_arrow_dir):
         if overwrite:
             print(f"Removing {by_ticker_aggs_arrow_dir=}")
@@ -212,10 +177,10 @@ if __name__ == "__main__":
     parser = argparse.ArgumentParser()
     parser.add_argument("--calendar_name", default="XNYS")
 
-    parser.add_argument("--start_session", default="2014-06-16")
-    parser.add_argument("--end_session", default="2024-09-06")
-    # parser.add_argument("--start_session", default="2020-01-01")
-    # parser.add_argument("--end_session", default="2020-12-31")
+    parser.add_argument("--start_date", default="2014-06-16")
+    parser.add_argument("--end_date", default="2024-09-06")
+    # parser.add_argument("--start_date", default="2020-01-01")
+    # parser.add_argument("--end_date", default="2020-12-31")
 
     parser.add_argument("--agg_time", default="day")
 
@@ -235,8 +200,8 @@ if __name__ == "__main__":
     config = PolygonConfig(
         environ=os.environ,
         calendar_name=args.calendar_name,
-        start_session=args.start_session,
-        end_session=args.end_session,
+        start_date=args.start_date,
+        end_date=args.end_date,
         agg_time=args.agg_time,
     )
 

zipline_polygon_bundle/concat_all_aggs_partitioned.py

@@ -138,10 +138,10 @@ if __name__ == "__main__":
     parser = argparse.ArgumentParser()
     parser.add_argument("--calendar_name", default="XNYS")
 
-    parser.add_argument("--start_session", default="2014-06-16")
-    parser.add_argument("--end_session", default="2024-09-06")
-    # parser.add_argument("--start_session", default="2020-10-07")
-    # parser.add_argument("--end_session", default="2020-10-15")
+    parser.add_argument("--start_date", default="2014-06-16")
+    parser.add_argument("--end_date", default="2024-09-06")
+    # parser.add_argument("--start_date", default="2020-10-07")
+    # parser.add_argument("--end_date", default="2020-10-15")
     # parser.add_argument("--aggs_pattern", default="2020/10/**/*.csv.gz")
     parser.add_argument("--aggs_pattern", default="**/*.csv.gz")
 
@@ -163,8 +163,8 @@ if __name__ == "__main__":
     config = PolygonConfig(
         environ=os.environ,
         calendar_name=args.calendar_name,
-        start_session=args.start_session,
-        end_session=args.end_session,
+        start_date=args.start_date,
+        end_date=args.end_date,
     )
 
     concat_all_aggs_from_csv(