feat: ship py.typed, with precise REST return types and corrected hints

AGENTS.md

@@ -51,6 +51,7 @@ The OpenAPI spec at https://api.massive.com/openapi is the source of truth for p
    - Single-object endpoints are named `get_<...>` and call `self._get(...)`.
    - Every query param is a keyword argument defaulting to `None`. Write range filters as `x_gt`, `x_gte`, `x_lt`, `x_lte`, `x_any_of`; `_get_params` turns them into `x.gt` and so on.
    - Always end the signature with `params: Optional[Dict[str, Any]] = None, raw: bool = False, options: Optional[RequestOptionBuilder] = None`.
+   - Precede the method with three `@overload`s that copy its signature, so callers get the precise type for each value of `raw`: `raw: Literal[False] = False` returns the model type (`Iterator[Model]`, `List[Model]` or `Model`); `*, raw: Literal[True]` (keyword-only) returns `HTTPResponse`; and `raw: bool = False` returns the full `Union`. Copy an existing method, such as `list_aggs` in `massive/rest/aggs.py`.
 4. **New domain only.** Create a `BaseClient` subclass and add it to `RESTClient`'s bases in `massive/rest/__init__.py`.
 5. **Tests, examples, docs.** For core endpoints, add a fixture under `test_rest/mocks/<url path>.json`, a test in `test_rest/test_<domain>.py`, an example in `examples/rest/`, and an `automethod` entry in `docs/source/*.rst`.
 6. Run `make lint` and `make test`.

Makefile

@@ -27,6 +27,9 @@ style:
 ## Check static types
 static:
 	poetry run mypy massive test_* examples
+	@# From test_rest/, massive resolves to the installed package, as it does
+	@# for users, so this also fails if the package stops shipping py.typed.
+	cd test_rest && poetry run mypy --strict typing_consumer.py
 
 ## Check code style and static types
 lint: style static

massive/__init__.py

@@ -1,6 +1,6 @@
-from .rest import RESTClient
+from .rest import RESTClient as RESTClient
 from .rest.base import version
-from .websocket import WebSocketClient
+from .websocket import WebSocketClient as WebSocketClient
 from .exceptions import *
 
 __version__ = version

massive/modelclass.py

@@ -1,10 +1,23 @@
 import inspect
+import sys
 import typing
 from dataclasses import dataclass
 
+if typing.TYPE_CHECKING:
+    if sys.version_info >= (3, 11):
+        from typing import dataclass_transform
+    else:
+        from typing_extensions import dataclass_transform
+else:
+
+    def dataclass_transform(**kwargs):
+        return lambda f: f
+
+
 _T = typing.TypeVar("_T")
 
 
+@dataclass_transform()
 def modelclass(cls: typing.Type[_T]) -> typing.Type[_T]:
     cls = dataclass(cls)
     type_hints = typing.get_type_hints(cls)

massive/rest/aggs.py

@@ -1,5 +1,5 @@
 from .base import BaseClient
-from typing import Optional, Any, Dict, List, Union, Iterator
+from typing import Optional, Any, Dict, List, Union, Iterator, Literal, overload
 from .models import Agg, GroupedDailyAgg, DailyOpenCloseAgg, PreviousCloseAgg, Sort
 from urllib3 import HTTPResponse
 from datetime import datetime, date
@@ -8,6 +8,52 @@
 
 
 class AggsClient(BaseClient):
+    @overload
+    def list_aggs(
+        self,
+        ticker: str,
+        multiplier: int,
+        timespan: str,
+        from_: Union[str, int, datetime, date],
+        to: Union[str, int, datetime, date],
+        adjusted: Optional[bool] = None,
+        sort: Optional[Union[str, Sort]] = None,
+        limit: Optional[int] = None,
+        params: Optional[Dict[str, Any]] = None,
+        raw: Literal[False] = False,
+        options: Optional[RequestOptionBuilder] = None,
+    ) -> Iterator[Agg]: ...
+    @overload
+    def list_aggs(
+        self,
+        ticker: str,
+        multiplier: int,
+        timespan: str,
+        from_: Union[str, int, datetime, date],
+        to: Union[str, int, datetime, date],
+        adjusted: Optional[bool] = None,
+        sort: Optional[Union[str, Sort]] = None,
+        limit: Optional[int] = None,
+        params: Optional[Dict[str, Any]] = None,
+        *,
+        raw: Literal[True],
+        options: Optional[RequestOptionBuilder] = None,
+    ) -> HTTPResponse: ...
+    @overload
+    def list_aggs(
+        self,
+        ticker: str,
+        multiplier: int,
+        timespan: str,
+        from_: Union[str, int, datetime, date],
+        to: Union[str, int, datetime, date],
+        adjusted: Optional[bool] = None,
+        sort: Optional[Union[str, Sort]] = None,
+        limit: Optional[int] = None,
+        params: Optional[Dict[str, Any]] = None,
+        raw: bool = False,
+        options: Optional[RequestOptionBuilder] = None,
+    ) -> Union[Iterator[Agg], HTTPResponse]: ...
     def list_aggs(
         self,
         ticker: str,
@@ -53,6 +99,52 @@ def list_aggs(
             options=options,
         )
 
+    @overload
+    def get_aggs(
+        self,
+        ticker: str,
+        multiplier: int,
+        timespan: str,
+        from_: Union[str, int, datetime, date],
+        to: Union[str, int, datetime, date],
+        adjusted: Optional[bool] = None,
+        sort: Optional[Union[str, Sort]] = None,
+        limit: Optional[int] = None,
+        params: Optional[Dict[str, Any]] = None,
+        raw: Literal[False] = False,
+        options: Optional[RequestOptionBuilder] = None,
+    ) -> List[Agg]: ...
+    @overload
+    def get_aggs(
+        self,
+        ticker: str,
+        multiplier: int,
+        timespan: str,
+        from_: Union[str, int, datetime, date],
+        to: Union[str, int, datetime, date],
+        adjusted: Optional[bool] = None,
+        sort: Optional[Union[str, Sort]] = None,
+        limit: Optional[int] = None,
+        params: Optional[Dict[str, Any]] = None,
+        *,
+        raw: Literal[True],
+        options: Optional[RequestOptionBuilder] = None,
+    ) -> HTTPResponse: ...
+    @overload
+    def get_aggs(
+        self,
+        ticker: str,
+        multiplier: int,
+        timespan: str,
+        from_: Union[str, int, datetime, date],
+        to: Union[str, int, datetime, date],
+        adjusted: Optional[bool] = None,
+        sort: Optional[Union[str, Sort]] = None,
+        limit: Optional[int] = None,
+        params: Optional[Dict[str, Any]] = None,
+        raw: bool = False,
+        options: Optional[RequestOptionBuilder] = None,
+    ) -> Union[List[Agg], HTTPResponse]: ...
     def get_aggs(
         self,
         ticker: str,
@@ -99,6 +191,44 @@ def get_aggs(
             options=options,
         )
 
+    @overload
+    def get_grouped_daily_aggs(
+        self,
+        date: Union[str, date],
+        adjusted: Optional[bool] = None,
+        params: Optional[Dict[str, Any]] = None,
+        raw: Literal[False] = False,
+        locale: str = "us",
+        market_type: str = "stocks",
+        include_otc: bool = False,
+        options: Optional[RequestOptionBuilder] = None,
+    ) -> List[GroupedDailyAgg]: ...
+    @overload
+    def get_grouped_daily_aggs(
+        self,
+        date: Union[str, date],
+        adjusted: Optional[bool] = None,
+        params: Optional[Dict[str, Any]] = None,
+        *,
+        raw: Literal[True],
+        locale: str = "us",
+        market_type: str = "stocks",
+        include_otc: bool = False,
+        options: Optional[RequestOptionBuilder] = None,
+    ) -> HTTPResponse: ...
+    @overload
+    def get_grouped_daily_aggs(
+        self,
+        date: Union[str, date],
+        adjusted: Optional[bool] = None,
+        params: Optional[Dict[str, Any]] = None,
+        raw: bool = False,
+        locale: str = "us",
+        market_type: str = "stocks",
+        include_otc: bool = False,
+        options: Optional[RequestOptionBuilder] = None,
+    ) -> Union[List[GroupedDailyAgg], HTTPResponse]: ...
+
     # TODO: next breaking change release move "market_type" to be 2nd mandatory
     # param
     def get_grouped_daily_aggs(
@@ -132,6 +262,37 @@ def get_grouped_daily_aggs(
             options=options,
         )
 
+    @overload
+    def get_daily_open_close_agg(
+        self,
+        ticker: str,
+        date: Union[str, date],
+        adjusted: Optional[bool] = None,
+        params: Optional[Dict[str, Any]] = None,
+        raw: Literal[False] = False,
+        options: Optional[RequestOptionBuilder] = None,
+    ) -> DailyOpenCloseAgg: ...
+    @overload
+    def get_daily_open_close_agg(
+        self,
+        ticker: str,
+        date: Union[str, date],
+        adjusted: Optional[bool] = None,
+        params: Optional[Dict[str, Any]] = None,
+        *,
+        raw: Literal[True],
+        options: Optional[RequestOptionBuilder] = None,
+    ) -> HTTPResponse: ...
+    @overload
+    def get_daily_open_close_agg(
+        self,
+        ticker: str,
+        date: Union[str, date],
+        adjusted: Optional[bool] = None,
+        params: Optional[Dict[str, Any]] = None,
+        raw: bool = False,
+        options: Optional[RequestOptionBuilder] = None,
+    ) -> Union[DailyOpenCloseAgg, HTTPResponse]: ...
     def get_daily_open_close_agg(
         self,
         ticker: str,
@@ -161,22 +322,50 @@ def get_daily_open_close_agg(
             options=options,
         )
 
+    @overload
+    def get_previous_close_agg(
+        self,
+        ticker: str,
+        adjusted: Optional[bool] = None,
+        params: Optional[Dict[str, Any]] = None,
+        raw: Literal[False] = False,
+        options: Optional[RequestOptionBuilder] = None,
+    ) -> List[PreviousCloseAgg]: ...
+    @overload
+    def get_previous_close_agg(
+        self,
+        ticker: str,
+        adjusted: Optional[bool] = None,
+        params: Optional[Dict[str, Any]] = None,
+        *,
+        raw: Literal[True],
+        options: Optional[RequestOptionBuilder] = None,
+    ) -> HTTPResponse: ...
+    @overload
+    def get_previous_close_agg(
+        self,
+        ticker: str,
+        adjusted: Optional[bool] = None,
+        params: Optional[Dict[str, Any]] = None,
+        raw: bool = False,
+        options: Optional[RequestOptionBuilder] = None,
+    ) -> Union[List[PreviousCloseAgg], HTTPResponse]: ...
     def get_previous_close_agg(
         self,
         ticker: str,
         adjusted: Optional[bool] = None,
         params: Optional[Dict[str, Any]] = None,
         raw: bool = False,
         options: Optional[RequestOptionBuilder] = None,
-    ) -> Union[PreviousCloseAgg, HTTPResponse]:
+    ) -> Union[List[PreviousCloseAgg], HTTPResponse]:
         """
         Get the previous day's open, high, low, and close (OHLC) for the specified stock ticker.
 
         :param ticker: The ticker symbol of the stock/equity.
         :param adjusted: Whether or not the results are adjusted for splits. By default, results are adjusted. Set this to false to get results that are NOT adjusted for splits.
         :param params: Any additional query params
         :param raw: Return raw object instead of results object
-        :return: Previous close aggregate
+        :return: List of previous close aggregates (the API returns a one-element list)
         """
         url = f"/v2/aggs/ticker/{ticker}/prev"