[chipkent]

Preview only. This is a draft in the fork showing the exact text of the PR planned for massive-com/client-python. It will be closed, not merged, once the upstream PR is opened.

Fixes #<issue: F4d>

Builds on #16 and #7. This branch is based on #7 and merges #16's branch (+56/-18 in 11 files from #14, #15 and #16), so that the overloads also cover #7's two new methods. The change this PR adds is the last two commits.

Every REST method returns Union[Model, HTTPResponse], whatever raw is, so once the package is typed, ordinary code needs isinstance checks. This PR gives each of the 94 methods that take raw three overloads:

@overload
def get_aggs(self, ticker: str, ..., raw: Literal[False] = False, ...) -> List[Agg]: ...
@overload
def get_aggs(self, ticker: str, ..., *, raw: Literal[True], ...) -> HTTPResponse: ...
@overload
def get_aggs(self, ticker: str, ..., raw: bool = False, ...) -> Union[List[Agg], HTTPResponse]: ...

• Default call: the model type. raw=True: HTTPResponse (keyword-only, so it does not overlap the first overload). Any other bool: the old Union, so existing code still checks.
• The overload commit is mechanical (+7,295/-452 in 14 files: reference.py +2,119/-278, benzinga.py +1,235/-73, financials.py +885/-61, futures.py +644/-17, etf_global.py +637/-14, economy.py +428/-1, and under +300 each in the other eight). It is 94 × 3 declarations plus two names on each typing import line, generated by the script below and then formatted with black massive. To review, regenerate rather than read: running it on the commit's parent reproduces the commit exactly, and running it again changes nothing. #<issue: F4d> asks whether a generated set like this is acceptable; if it is, I can commit the generator with a check that fails when running it would change anything.
• The last commit adds test_rest/typing_consumer.py, user code checked by mypy --strict in make static (never run), and documents the pattern in AGENTS.md. Without the overloads the consumer reports 20 errors; with them, none.
• typing_consumer.py types the WebSocket callback as List[WebSocketMessage], which is today's annotation and the one Fix WebSocketMessage type contract massive-com/client-python#1032 keeps, so the two do not conflict.
• No new dependency: Literal and overload come from typing; the consumer gets assert_type from typing_extensions only on 3.10 and only under mypy, which depends on it.
• Runtime behavior is unchanged: with the overloads removed, every module parses to the same AST as before, apart from the typing import line.

Generator: python gen_raw_overloads.py massive/rest/*.py && black massive

"""Add @overload declarations on `raw` to REST client methods.

For every method that takes `raw` and is annotated
`-> Union[X, HTTPResponse]`, insert three overloads before it:

    raw: Literal[False] = False  -> X
    *, raw: Literal[True]        -> HTTPResponse
    raw: bool = False            -> Union[X, HTTPResponse]

Every other parameter is copied from the method unchanged. Methods that
already have overloads, or any other decorator, are skipped, so running
it again changes nothing.

Usage: python gen_raw_overloads.py massive/rest/*.py && black massive
"""

import ast
import re
import sys


def param(src, arg, default, annotation=None, keep_default=True):
    text = arg.arg
    if annotation is None and arg.annotation is not None:
        annotation = ast.get_source_segment(src, arg.annotation)
    if annotation:
        text += f": {annotation}"
    if default is not None and keep_default:
        text += f" = {ast.get_source_segment(src, default)}"
    return text


def overloads(src, fn):
    args = fn.args.args
    defaults = [None] * (len(args) - len(fn.args.defaults)) + fn.args.defaults
    union = ast.get_source_segment(src, fn.returns)
    model = ast.get_source_segment(src, fn.returns.slice.elts[0])
    indent = " " * fn.col_offset
    variants = [
        ("Literal[False]", True, False, model),
        ("Literal[True]", False, True, "HTTPResponse"),
        (None, True, False, union),
    ]
    out = []
    for raw_type, keep_default, keyword_only, returns in variants:
        params = []
        for arg, default in zip(args, defaults):
            if arg.arg != "raw":
                params.append(param(src, arg, default))
                continue
            if keyword_only:
                params.append("*")
            params.append(param(src, arg, default, raw_type, keep_default))
        out.append(
            f"{indent}@overload\n"
            f"{indent}def {fn.name}({', '.join(params)}) -> {returns}: ...\n"
        )
    return "".join(out)


def is_target(fn):
    if not isinstance(fn, ast.FunctionDef):
        return False
    r = fn.returns
    return (
        not fn.decorator_list
        and "raw" in [a.arg for a in fn.args.args]
        and isinstance(r, ast.Subscript)
        and isinstance(r.value, ast.Name)
        and r.value.id == "Union"
        and isinstance(r.slice, ast.Tuple)
        and len(r.slice.elts) == 2
        and isinstance(r.slice.elts[1], ast.Name)
        and r.slice.elts[1].id == "HTTPResponse"
    )


def process(path):
    src = open(path).read()
    targets = []
    for cls in ast.parse(src).body:
        if isinstance(cls, ast.ClassDef):
            done = {f.name for f in cls.body if getattr(f, "decorator_list", None)}
            targets += [f for f in cls.body if is_target(f) and f.name not in done]
    if not targets:
        return 0
    lines = src.splitlines(keepends=True)
    for fn in sorted(targets, key=lambda f: f.lineno, reverse=True):
        i = fn.lineno - 1
        while i > 0 and lines[i - 1].strip().startswith("#"):
            i -= 1  # keep comments directly above the method attached to it
        lines.insert(i, overloads(src, fn))
    new = "".join(lines)
    m = re.search(r"^from typing import (.*)$", new, re.M)
    names = [n.strip() for n in m.group(1).split(",")]
    names += [n for n in ("Literal", "overload") if n not in names]
    new = new[: m.start(1)] + ", ".join(names) + new[m.end(1) :]
    open(path, "w").write(new)
    return len(targets)


total = 0
for path in sys.argv[1:]:
    n = process(path)
    print(path, n)
    total += n
print("total", total)

[coderabbitai]

Important

Draft PR not reviewed

Draft PRs are not automatically reviewed by default.

• Trigger a manual review

To automatically review draft PRs, update your CodeRabbit configuration:

reviews:
  auto_review:
    drafts: true

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

_{Comment @coderabbitai help to get the list of available commands.}