zipline_polygon_bundle 0.1.6__tar.gz → 0.1.8__tar.gz

{zipline_polygon_bundle-0.1.6 → zipline_polygon_bundle-0.1.8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: zipline_polygon_bundle
-Version: 0.1.6
+Version: 0.1.8
 Summary: A zipline-reloaded data provider bundle for Polygon.io
 License: GNU AFFERO GENERAL PUBLIC LICENSE
                                 Version 3, 19 November 2007
@@ -671,10 +671,13 @@ Classifier: Programming Language :: Python :: 3
 Classifier: License :: OSI Approved :: GNU Affero General Public License v3
 Classifier: Operating System :: OS Independent
 Requires-Dist: bcolz-zipline (>=1.2.11)
+Requires-Dist: fsspec (>=2024.10)
 Requires-Dist: numpy (<2)
 Requires-Dist: pandas (>=2.2,<3)
-Requires-Dist: polygon-api-client
-Requires-Dist: pyarrow
+Requires-Dist: pandas-market-calendars (>=4.4.2)
+Requires-Dist: pandas_ta (>=0.3)
+Requires-Dist: polygon-api-client (>=1.14.2)
+Requires-Dist: pyarrow (>=18.1.0,<19)
 Requires-Dist: pytz (>=2018.5)
 Requires-Dist: requests (>=2.9.1)
 Requires-Dist: toolz (>=0.8.2)

{zipline_polygon_bundle-0.1.6 → zipline_polygon_bundle-0.1.8}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = 'zipline_polygon_bundle'
-version = '0.1.6'
+version = '0.1.8'
 description = 'A zipline-reloaded data provider bundle for Polygon.io'
 authors = [
     { name = 'Jim White', email = 'jim@fovi.com' },
@@ -19,7 +19,7 @@ Repository = 'https://github.com/fovi-llc/zipline-polygon-bundle'
 
 [tool.poetry]
 name = 'zipline-polygon-bundle'
-version = '0.1.6'
+version = '0.1.8'
 description = 'A zipline-reloaded data provider bundle for Polygon.io'
 authors = ['Jim White <jim@fovi.com>']
 license = 'AGPL-3.0'
@@ -32,13 +32,16 @@ classifiers = [
 ]
 
 [tool.poetry.dependencies]
+fsspec = ">=2024.10"
 python = ">=3.9,<4.0"
-polygon-api-client = "*"
+polygon-api-client = ">=1.14.2"
 pandas = ">=2.2,<3"
+pandas-market-calendars = ">=4.4.2"
+pandas_ta = ">=0.3"
 pytz = ">=2018.5"
 requests = ">=2.9.1"
 bcolz-zipline = ">=1.2.11"
-pyarrow = "*"
+pyarrow = ">=18.1.0,<19"
 numpy = "<2"
 toolz = ">=0.8.2"
 zipline-reloaded = ">=3.1"

zipline_polygon_bundle-0.1.8/zipline_polygon_bundle/__init__.py

@@ -0,0 +1,49 @@
+from .bundle import (
+    register_polygon_equities_bundle,
+    symbol_to_upper,
+    polygon_equities_bundle_day,
+    polygon_equities_bundle_minute,
+)
+
+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, 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 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",
+    "symbol_to_upper",
+    "polygon_equities_bundle_day",
+    "polygon_equities_bundle_minute",
+    "PolygonConfig",
+    "concat_all_aggs_from_csv",
+    "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_all_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-0.1.6 → zipline_polygon_bundle-0.1.8}/zipline_polygon_bundle/adjustments.py

@@ -64,8 +64,9 @@ def load_splits(
     splits["split_from"] = splits["split_from"].astype(float)
     splits["split_to"] = splits["split_to"].astype(float)
     splits["ratio"] = splits["split_from"] / splits["split_to"]
-    splits.drop(columns=["ticker", "split_from", "split_to"], inplace=True)
-    return splits
+    # Only return columns Zipline wants.
+    # Polygon may add more columns in the future (as they did with `id`).
+    return splits[["sid", "effective_date", "ratio"]]
 
 
 def load_polygon_dividends(
@@ -87,7 +88,9 @@ def load_polygon_dividends(
         dividends = pd.DataFrame(dividends)
         os.makedirs(os.path.dirname(dividends_path), exist_ok=True)
         dividends.to_parquet(dividends_path)
-        print(f"Wrote {len(dividends)=} from Polygon list_dividends to {dividends_path=}")
+        print(
+            f"Wrote {len(dividends)=} from Polygon list_dividends to {dividends_path=}"
+        )
         # if len(dividends) < 10000:
         #     logging.error(f"Only got {len(dividends)=} from Polygon list_dividends.")
     # We will always load from the file to avoid any chance of weird errors.
@@ -116,9 +119,9 @@ def load_chunked_polygon_dividends(
         next_end_date = first_of_next_month - datetime.timedelta(days=1)
         if next_end_date > last_end_date:
             next_end_date = last_end_date
-        dividends_list.append(load_polygon_dividends(
-            config, next_start_end, next_end_date
-        ))
+        dividends_list.append(
+            load_polygon_dividends(config, next_start_end, next_end_date)
+        )
         next_start_end = next_end_date + datetime.timedelta(days=1)
     return pd.concat(dividends_list)
 
@@ -145,7 +148,42 @@ def load_dividends(
         },
         inplace=True,
     )
-    dividends.drop(
-        columns=["ticker", "frequency", "currency", "dividend_type"], inplace=True
+    # Only return columns Zipline wants.
+    # Polygon may add more columns in the future (as they did with `id`).
+    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"
     )
-    return dividends
+    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-0.1.6 → zipline_polygon_bundle-0.1.8}/zipline_polygon_bundle/bundle.py

@@ -1,18 +1,21 @@
+import os
 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 .config import PolygonConfig
 from .concat_all_aggs import concat_all_aggs_from_csv, generate_csv_agg_tables
 from .adjustments import load_splits, load_dividends
 
 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.
 
@@ -175,8 +178,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 +187,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",
     )
 
@@ -219,7 +222,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,
@@ -400,8 +403,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 +412,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 +447,8 @@ 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,
@@ -459,8 +462,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,8 +488,8 @@ 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,
@@ -495,6 +498,25 @@ def register_polygon_equities_bundle(
 ):
     if agg_time not in ["day", "minute"]:
         raise ValueError(f"agg_time must be 'day' or 'minute', 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()
+        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,
         (
@@ -502,8 +524,8 @@ def register_polygon_equities_bundle(
             if agg_time == "minute"
             else polygon_equities_bundle_day
         ),
-        start_session=start_session,
-        end_session=end_session,
+        start_session=parse_date(start_date, calendar=calendar),
+        end_session=parse_date(end_date, calendar=calendar),
         calendar_name=calendar_name,
         # minutes_per_day=390,
         # create_writers=True,
@@ -517,12 +539,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-0.1.6 → zipline_polygon_bundle-0.1.8}/zipline_polygon_bundle/concat_all_aggs.py

@@ -1,7 +1,7 @@
 from .config import PolygonConfig
 
 import shutil
-from typing import Iterator, Tuple
+from typing import Iterator, Tuple, List, Union
 
 import argparse
 import glob
@@ -10,6 +10,7 @@ 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
 
@@ -34,7 +35,7 @@ def to_partition_key(s: str) -> str:
 
 
 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 +58,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 +76,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 +102,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 +141,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 +163,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 +197,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 +220,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-0.1.6 → zipline_polygon_bundle-0.1.8}/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(