Stronger types for @modelclass

[shooit]

Is your feature request related to a problem? Please describe.

The @modelclass decorator works at runtime as expected; however, creates issues at development time for users using editors backed by LSP mode with Pyright/Pylance such as VS Code. PyCharm has a bespoke way of determining types in the IDE and is not affected by this.

I have annotated the following code with the types/issues that appear in VSCode

import typing

from polygon import RESTClient
from polygon.rest.models.tickers import TickerDetails

client = RESTClient()

# type of ticker_details_response_uncast is `Unkown | HTTPResponse`
ticker_details_response_uncast = client.get_ticker_details(ticker="AAPL")

# error: Expected type expression but received "(Type[_T@dataclass]) -> Type[_T@dataclass]"
# this same error happens if you do ticker_details_response: TickerDetails = ...
ticker_details_response_cast = typing.cast(
    TickerDetails, client.get_ticker_details(ticker="AAPL")
)

# the type of TickerDetails.from_dict is `Any`; therefore its arguments are not type-checked and the type of ticker_details_from_dict is also `Any`
ticker_details_from_dict = TickerDetails.from_dict({})

# error: Argument of type "Literal[True]" cannot be assigned to parameter of type "Type[_T@dataclass]" \n Type "Literal[True]" cannot be assigned to type "Type[_T@dataclass]"
ticker_details_from_args = TickerDetails(True)

# error: Expected 1 more positional argument
ticker_details_from_kwargs = TickerDetails(active=True)

Describe the solution you'd like

With a minimum python version of 3.8 we are limited in what we can do to solve this problem but we can still make it better by adding type annotations to def modelclass(cls): in polygon/modelclass.py

import inspect
import typing
from dataclasses import dataclass


_T = typing.TypeVar("_T")


def modelclass(cls: typing.Type[_T]) -> typing.Type[_T]:
    cls = dataclass(cls)
    attributes = [
        a
        for a in cls.__dict__["__annotations__"].keys()
        if not a.startswith("__") and not inspect.isroutine(a)
    ]

    def init(self, *args, **kwargs):
        for (i, a) in enumerate(args):
            if i < len(attributes):
                self.__dict__[attributes[i]] = a
        for (k, v) in kwargs.items():
            if k in attributes:
                self.__dict__[k] = v

    cls.__init__ = init  # type: ignore[assignment]

    return cls

And when we go back to our example code we see much better type inference but still two errors when initializing classes in line. In my opinion, this would still be a huge win since consumers of this library will be most likely just consuming the responses from the client and not initializing data directly.

import typing

from polygon import RESTClient
from polygon.rest.models.tickers import TickerDetails

client = RESTClient()

# type of ticker_details_response_uncast is `TickerDetails | HTTPResponse`
ticker_details_response_uncast = client.get_ticker_details(ticker="AAPL")

# type of ticker_details_cast is `TickerDetails`
ticker_details_response_cast = typing.cast(
    TickerDetails, client.get_ticker_details(ticker="AAPL")
)

# type of TickerDetails.from_dict is `(d: Unknown) -> TickerDetails`
ticker_details_from_dict = TickerDetails.from_dict({})

# error: Expected no arguments to "TickerDetails" constructor
ticker_details_from_args = TickerDetails(True)

# Expected no arguments to "TickerDetails" constructor
ticker_details_from_kwargs = TickerDetails(active=True)

Python 3.11 adds support for the dataclass transform decorator which is directly aimed at this problem space. You can add it to modelclass like this:

@typing.dataclass_transform()
def modelclass(cls: typing.Type[_T]) -> typing.Type[_T]:
...

Now there are no issues in the example, even when initializing classes directly!

[justinpolygon]

This is amazing. Thanks for putting it all together and explaining in such detail.