Advanced Typing

from collections.abc import Callable, Iterable, Sequence

def batched[T](items: Sequence[T], size: int) -> list[Sequence[T]]:
    return [items[i:i + size] for i in range(0, len(items), size)]

def map_concurrently[T, R](fn: Callable[[T], R], items: Iterable[T]) -> list[R]:
    return [fn(x) for x in items]

chunks = batched(["a", "b", "c"], 2)        # list[Sequence[str]]
lengths = map_concurrently(len, ["ab", "c"])  # list[int]

def mean(xs: list[float]) -> float: ...
scores: list[int] = [1, 2, 3]
mean(scores)          # error: list[int] is not list[float] (invariant)

def mean(xs: Sequence[float]) -> float: ...   # fix: read-only and covariant

from typing import Protocol, TypeVar

T_co = TypeVar("T_co", covariant=True)

class Source(Protocol[T_co]):
    def read(self) -> T_co: ...

# Functions
def first[T](xs: Sequence[T]) -> T: ...

# Classes (Generic base is implicit)
class Registry[K, V]:
    def __init__(self) -> None:
        self._items: dict[K, V] = {}
    def register(self, key: K, value: V) -> None:
        self._items[key] = value

# Aliases with parameters
type Scored[T] = list[tuple[T, float]]
type Handler[**P] = Callable[P, Awaitable[None]]
type Shape[*Ts] = tuple[*Ts]

# Defaults (3.13)
class Page[T = dict[str, object]]:
    items: list[T]

from collections.abc import Sequence
from typing import Protocol, runtime_checkable

class SupportsLen(Protocol):
    def __len__(self) -> int: ...

class Tokenizer(Protocol):
    @property
    def vocab_size(self) -> int: ...                       # read-only attribute
    def encode(self, text: str) -> list[int]: ...
    def decode(self, ids: Sequence[int]) -> str: ...

@runtime_checkable
class Closeable(Protocol):
    def close(self) -> None: ...

def count_tokens(tok: Tokenizer, texts: Sequence[str]) -> int:
    return sum(len(tok.encode(t)) for t in texts)

import functools
import time
from collections.abc import Callable

def timed[**P, R](fn: Callable[P, R]) -> Callable[P, R]:
    @functools.wraps(fn)
    def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
        start = time.perf_counter()
        try:
            return fn(*args, **kwargs)
        finally:
            print(f"{fn.__name__}: {time.perf_counter() - start:.3f}s")
    return wrapper

@timed
def embed(texts: list[str], model: str = "small") -> list[list[float]]: ...

embed(["a"], model="large")     # checked against the real signature
embed(["a"], modle="large")     # error: unexpected keyword argument

from typing import Concatenate

class Client:
    ...

def with_client[**P, R](fn: Callable[Concatenate[Client, P], R]) -> Callable[P, R]:
    """Inject a client as the first argument; callers no longer pass it."""
    def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
        return fn(Client(), *args, **kwargs)
    return wrapper

@with_client
def ask(client: Client, prompt: str, *, temperature: float = 0.0) -> str: ...

ask("hi", temperature=0.3)      # Client parameter removed from the signature

from typing import ParamSpec, TypeVar
P = ParamSpec("P")
R = TypeVar("R")
def timed(fn: Callable[P, R]) -> Callable[P, R]: ...

from typing import Literal, overload

class Stream: ...

@overload
def complete(prompt: str, *, stream: Literal[False] = False) -> str: ...
@overload
def complete(prompt: str, *, stream: Literal[True]) -> Stream: ...
@overload
def complete(prompt: str, *, stream: bool) -> str | Stream: ...
def complete(prompt: str, *, stream: bool = False) -> str | Stream:
    return Stream() if stream else "done"

text = complete("hi")                  # str
s = complete("hi", stream=True)        # Stream

from typing import Literal, NotRequired, ReadOnly, Required, TypedDict, Unpack

class Message(TypedDict):
    role: Literal["system", "user", "assistant", "tool"]
    content: str
    name: NotRequired[str]

class GenerationOptions(TypedDict, total=False):
    temperature: float
    max_tokens: int
    stop: list[str]
    model: Required[str]

def generate(prompt: str, **opts: Unpack[GenerationOptions]) -> str:   # PEP 692, 3.12
    return opts["model"]

generate("hi", model="small", temperature=0.2)     # keys and types checked
generate("hi", temprature=0.2)                     # error: unexpected keyword

from typing import Literal, TypedDict

class TextBlock(TypedDict):
    type: Literal["text"]
    text: str

class ToolUseBlock(TypedDict):
    type: Literal["tool_use"]
    name: str
    input: dict[str, object]

type Block = TextBlock | ToolUseBlock

def render(block: Block) -> str:
    if block["type"] == "text":
        return block["text"]            # narrowed to TextBlock
    return f"<tool {block['name']}>"    # narrowed to ToolUseBlock

from typing import TypeGuard, TypeIs

def is_str_list(xs: list[object]) -> TypeGuard[list[str]]:   # 3.10: narrows the True branch only
    return all(isinstance(x, str) for x in xs)

def is_text_block(b: Block) -> TypeIs[TextBlock]:            # 3.13: narrows both branches
    return b["type"] == "text"

def handle(b: Block) -> None:
    if is_text_block(b):
        print(b["text"])          # TextBlock
    else:
        print(b["name"])          # ToolUseBlock (TypeGuard would leave Block here)

from typing import Self

class Pipeline:
    def __init__(self) -> None:
        self.steps: list[str] = []

    def add(self, step: str) -> Self:          # subclasses keep their own type
        self.steps.append(step)
        return self

    @classmethod
    def from_config(cls, steps: list[str]) -> Self:
        p = cls()
        for s in steps:
            p.add(s)
        return p

class RagPipeline(Pipeline):
    def rerank(self) -> Self:
        return self.add("rerank")

RagPipeline().add("retrieve").rerank()          # ok: add returns RagPipeline

from enum import Enum
from typing import NoReturn, assert_never

def fail(msg: str) -> NoReturn:
    raise RuntimeError(msg)

class Finish(Enum):
    STOP = "stop"
    LENGTH = "length"
    TOOL = "tool_calls"

def describe(reason: Finish) -> str:
    if reason is Finish.STOP:
        return "complete"
    elif reason is Finish.LENGTH:
        return "truncated"
    elif reason is Finish.TOOL:
        return "needs tool"
    else:
        assert_never(reason)      # error here if a member is added and not handled

from collections.abc import Callable

def call_later[*Ts](fn: Callable[[*Ts], None], *args: *Ts) -> None:
    fn(*args)

def log(level: int, msg: str) -> None: ...

call_later(log, 10, "hi")       # ok
call_later(log, "hi", 10)       # error: argument types do not match

class Array[*Shape]:
    def __init__(self, *shape: *Shape) -> None:
        self.shape = shape

from typing import Annotated

from fastapi import Query
from pydantic import BaseModel, Field

class SearchRequest(BaseModel):
    query: Annotated[str, Field(min_length=1, max_length=2000)]
    k: Annotated[int, Field(ge=1, le=100)] = 10

def search(q: Annotated[str, Query(max_length=200)], k: Annotated[int, Query(ge=1)] = 5) -> list[str]: ...

type Temperature = Annotated[float, Field(ge=0.0, le=2.0)]   # reusable constrained type

from typing import final, override          # override: 3.12
from warnings import deprecated             # 3.13 (typing_extensions before)

class BaseRetriever:
    def retrieve(self, query: str) -> list[str]:
        return []

    @final
    def close(self) -> None: ...              # subclasses may not override

class BM25Retriever(BaseRetriever):
    @override
    def retrieve(self, query: str) -> list[str]:   # error if the base method is renamed
        return [query]

@deprecated("use embed_batch()")
def embed_one(text: str) -> list[float]: ...      # callers get a checker warning

type JSON = dict[str, JSON] | list[JSON] | str | int | float | bool | None   # 3.12

import json

def load_json(text: str) -> JSON:
    return json.loads(text)        # json.loads returns Any; the annotation documents the boundary

def get_path(doc: JSON, *keys: str) -> JSON:
    for k in keys:
        if not isinstance(doc, dict):
            raise KeyError(k)
        doc = doc[k]
    return doc

# stubs/legacy_vectordb/__init__.pyi
from collections.abc import Sequence

class Collection:
    def add(self, ids: Sequence[str], embeddings: Sequence[Sequence[float]]) -> None: ...
    def query(self, embedding: Sequence[float], n_results: int = ...) -> dict[str, list[list[str]]]: ...

def connect(url: str, *, timeout: float = ...) -> Collection: ...