MOD: Update corporate actions API

Author: cjdsellers

databento/reference/api/corporate.py

@@ -13,11 +13,9 @@
 from databento.common.http import BentoHttpAPI
 from databento.common.parsing import convert_date_columns
 from databento.common.parsing import convert_datetime_columns
-from databento.common.parsing import datetime_to_date_string
-from databento.common.parsing import optional_date_to_string
+from databento.common.parsing import datetime_to_string
+from databento.common.parsing import optional_datetime_to_string
 from databento.common.parsing import optional_symbols_list_to_list
-from databento.common.publishers import Dataset
-from databento.common.validation import validate_semantic_string
 
 
 class CorporateActionsHttpAPI(BentoHttpAPI):
@@ -31,12 +29,12 @@ def __init__(self, key: str, gateway: str) -> None:
 
     def get_range(
         self,
-        start_date: date | str,
-        end_date: date | str | None = None,
-        dataset: Dataset | str | None = None,
+        start: pd.Timestamp | date | str | int,
+        end: pd.Timestamp | date | str | int | None = None,
         symbols: Iterable[str] | str | None = None,
         stype_in: SType | str = "raw_symbol",
         events: Iterable[str] | str | None = None,
+        us_only: bool = False,
     ) -> pd.DataFrame:
         """
         Request a new corporate actions time series from Databento.
@@ -45,12 +43,16 @@ def get_range(
 
         Parameters
         ----------
-        start_date : date or str
-            The start date (UTC) of the request time range (inclusive).
-        end_date : date or str, optional
-            The end date (UTC) of the request time range (exclusive).
-        dataset : Dataset or str, optional
-            The dataset code (string identifier) for the request.
+        start : pd.Timestamp or date or str or int
+            The start datetime of the request time range (inclusive).
+            Assumes UTC as timezone unless passed a tz-aware object.
+            If an integer is passed, then this represents nanoseconds since the UNIX epoch.
+        end : pd.Timestamp or date or str or int, optional
+            The end datetime of the request time range (exclusive).
+            Assumes UTC as timezone unless passed a tz-aware object.
+            If an integer is passed, then this represents nanoseconds since the UNIX epoch.
+            Values are forward filled based on the resolution provided.
+            Defaults to the same value as `start`.
         symbols : Iterable[str] or str, optional
             The symbols to filter for. Takes up to 2,000 symbols per request.
             If more than 1 symbol is specified, the data is merged and sorted by time.
@@ -64,26 +66,27 @@ def get_range(
             Takes any number of event types per request.
             If not specified then will be for **all** event types.
             See [EVENT](https://databento.com/docs/standards-and-conventions/reference-data-enums#event) enum.
+        us_only : bool, default False
+            If filtering for US markets only.
 
         Returns
         -------
         pandas.DataFrame
             The data converted into a data frame.
 
         """
-        dataset = validate_semantic_string(dataset, "dataset") if dataset is not None else None
         symbols_list = optional_symbols_list_to_list(symbols, SType.RAW_SYMBOL)
 
         if isinstance(events, str):
             events = events.strip().strip(",").split(",")
 
         data: dict[str, object | None] = {
-            "start_date": datetime_to_date_string(start_date),
-            "end_date": optional_date_to_string(end_date),
-            "dataset": dataset,
+            "start": datetime_to_string(start),
+            "end": optional_datetime_to_string(end),
             "symbols": ",".join(symbols_list),
             "stype_in": stype_in,
             "events": ",".join(events) if events else None,
+            "us_only": us_only,
         }
 
         response = self._post(

examples/reference_corporate_actions.py

@@ -9,12 +9,12 @@
     client = Reference(key=key)
 
     response: pd.DataFrame = client.corporate_actions.get_range(
-        dataset="XNAS.ITCH",
         symbols="AAPL,MSFT,TSLA",
         stype_in="raw_symbol",
-        start_date="2023",
-        end_date="2024-04",
+        start="2023",
+        end="2024-04",
         events="DIV,LIQ",
+        us_only=True,
     )
 
     pprint(response.head())

tests/test_reference_corporate.py

@@ -52,11 +52,10 @@ def test_corporate_actions_get_range_sends_expected_request(
 
     # Act
     reference_client.corporate_actions.get_range(
-        dataset=None,
         symbols="AAPL",
         stype_in="raw_symbol",
-        start_date="2024-01",
-        end_date="2024-04",
+        start="2024-01",
+        end="2024-04",
         events=events,
     )
 
@@ -69,12 +68,12 @@ def test_corporate_actions_get_range_sends_expected_request(
     assert call["headers"]["accept"] == "application/json"
     assert all(v in call["headers"]["user-agent"] for v in ("Databento/", "Python/"))
     assert call["data"] == {
-        "dataset": None,
-        "start_date": "2024-01",
-        "end_date": "2024-04",
+        "start": "2024-01",
+        "end": "2024-04",
         "symbols": "AAPL",
         "stype_in": "raw_symbol",
         "events": data_events,
+        "us_only": False,
     }
     assert call["timeout"] == (100, 100)
     assert isinstance(call["auth"], requests.auth.HTTPBasicAuth)
@@ -94,11 +93,10 @@ def test_corporate_actions_get_range_response_parsing(
 
     # Act
     df_raw = reference_client.corporate_actions.get_range(
-        dataset=None,
         symbols="AAPL",
         stype_in="raw_symbol",
-        start_date="2024-01",
-        end_date="2024-04",
+        start="2024-01",
+        end="2024-04",
     )
 
     # Assert