HTTP constant literals for typeshed

make-http-constants.py

import ast
import sys
import textwrap
from pathlib import Path
from typing import NamedTuple

from typeshed_client import get_search_context, get_stub_ast

try:
    typeshed_root = Path(sys.argv[1])
except IndexError:
    typeshed_root = Path(__file__).parent.parent

http_init_module: ast.Module = get_stub_ast(
    "http",
    search_context=get_search_context(typeshed=typeshed_root / "stdlib"),
)  # pyright: ignore[reportAssignmentType]


class Member(NamedTuple):
    name: str
    value: int
    since: tuple[int, int, int] | None = None

    @property
    def annotation_ast(self) -> ast.Subscript:
        return ast.Subscript(value=ast.Name(id="Literal", ctx=ast.Load()), slice=ast.Constant(self.value))

    @property
    def version_check_ast(self) -> ast.Compare | None:
        if not self.since:
            return None
        return ast.Compare(
            left=ast.Attribute(ast.Name("sys"), attr="version_info"),
            ops=[ast.GtE()],
            comparators=[ast.Tuple(elts=[*map(ast.Constant, self.since)])],
        )

    @property
    def definition_ast(self) -> ast.AnnAssign:
        assignment = ast.AnnAssign(
            target=ast.Name(id=self.name),
            annotation=self.annotation_ast,
            value=None,
            simple=1,
        )
        return assignment


def collect_members() -> list[Member]:
    members = []
    for node in http_init_module.body:
        match node:
            case ast.ClassDef(body=body):
                for stmt in body:
                    match stmt:
                        case ast.If(
                            test=ast.Compare(ops=[ast.GtE()], comparators=[version_expr]),
                            body=attrs,
                        ):
                            members.extend(
                                Member(
                                    name=target.id,
                                    value=ast.literal_eval(attr.value),
                                    since=ast.literal_eval(version_expr),
                                )
                                for attr in attrs
                                if isinstance(attr, ast.Assign) and isinstance((target := attr.targets[0]), ast.Name)
                            )
                        case ast.Assign(targets=[ast.Name(id=name)], value=value):
                            members.append(Member(name=name, value=ast.literal_eval(value)))
                        case _:
                            continue
            case _:
                continue
    return members


def build_annotated_constants() -> list[str]:
    constants = []
    prev_member = None
    prev_since = None
    prev_version_check = None
    version_checked_members: list[Member] = []

    for member in *collect_members(), None:
        if member is None or prev_since and member.since != prev_since:
            assert version_checked_members
            assert prev_version_check
            if_stmt = (
                ast.unparse(ast.If(test=prev_version_check, body=[]))
                + "\n"
                + "\n".join(
                    textwrap.indent(ast.unparse(version_checked_member.definition_ast), " " * 4)
                    for version_checked_member in version_checked_members
                )
            )
            constants.extend(("", if_stmt, *(() if member is None else ("",))))
            version_checked_members.clear()
        if member is None:
            continue
        prev_since = member.since
        if prev_version_check := member.version_check_ast:
            version_checked_members.append(member)
        else:
            if prev_member and prev_member.value // 100 != member.value // 100:
                constants.append("")
            constants.append(ast.unparse(member.definition_ast))
        prev_member = member

    return constants


if __name__ == "__main__":
    print("\n".join(build_annotated_constants()))