Add type hints across the Pretzel package

Annotate all function signatures, mark module-level constants as Final,
and reuse the geometry aliases from pretzel.engine. The make_assets
script gets its own local aliases. No behavior change.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Author: realpython-bot

python315-sampling-profiler/examples/fetch_textures.py

@@ -9,19 +9,19 @@
 }
 
 
-async def download_texture(name, latency):
+async def download_texture(name: str, latency: float) -> bytes:
     await asyncio.sleep(latency)
     return name.encode() * 100_000
 
 
-async def decode_texture(payload):
+async def decode_texture(payload: bytes) -> int:
     checksum = 0
     for byte in payload:
         checksum = (checksum * 31 + byte) % 1_000_003
     return checksum
 
 
-async def main():
+async def main() -> None:
     async with asyncio.TaskGroup() as group:
         downloads = {
             name: group.create_task(

python315-sampling-profiler/make_assets.py

@@ -8,51 +8,56 @@
 
 import math
 import pathlib
+from typing import Final
+
+type Vector = tuple[float, float, float]
+type Face = tuple[int, int, int]
+type Color = tuple[int, int, int]
 
 ASSETS_DIR = pathlib.Path(__file__).parent / "src" / "pretzel" / "assets"
 
-SEGMENTS = 240
-RING_POINTS = 16
-TUBE_RADIUS = 0.62
+SEGMENTS: Final = 240
+RING_POINTS: Final = 16
+TUBE_RADIUS: Final = 0.62
 
-WALLPAPER_SIZE = 512
-WALLPAPERS = {
+WALLPAPER_SIZE: Final = 512
+WALLPAPERS: Final = {
     "bakery_dawn.ppm": ((252, 210, 153), (146, 90, 118), (255, 236, 179)),
     "bakery_noon.ppm": ((214, 235, 251), (245, 232, 201), (255, 255, 224)),
     "bakery_dusk.ppm": ((94, 63, 107), (233, 156, 90), (255, 214, 138)),
 }
 
 
-def trace_trefoil(t):
+def trace_trefoil(t: float) -> Vector:
     return (
         math.sin(t) + 2.0 * math.sin(2.0 * t),
         math.cos(t) - 2.0 * math.cos(2.0 * t),
         -math.sin(3.0 * t) * 1.2,
     )
 
 
-def subtract(a, b):
+def subtract(a: Vector, b: Vector) -> Vector:
     return (a[0] - b[0], a[1] - b[1], a[2] - b[2])
 
 
-def cross(a, b):
+def cross(a: Vector, b: Vector) -> Vector:
     return (
         a[1] * b[2] - a[2] * b[1],
         a[2] * b[0] - a[0] * b[2],
         a[0] * b[1] - a[1] * b[0],
     )
 
 
-def dot(a, b):
+def dot(a: Vector, b: Vector) -> float:
     return a[0] * b[0] + a[1] * b[1] + a[2] * b[2]
 
 
-def normalize(vector):
+def normalize(vector: Vector) -> Vector:
     length = math.sqrt(dot(vector, vector)) or 1.0
     return (vector[0] / length, vector[1] / length, vector[2] / length)
 
 
-def sweep_tube():
+def sweep_tube() -> tuple[list[Vector], list[Face]]:
     step = 2.0 * math.pi / SEGMENTS
     centers = [trace_trefoil(index * step) for index in range(SEGMENTS)]
     tangents = [
@@ -96,7 +101,7 @@ def sweep_tube():
     return vertices, faces
 
 
-def write_mesh(path):
+def write_mesh(path: pathlib.Path) -> None:
     vertices, faces = sweep_tube()
     with path.open("w", encoding="utf-8") as file:
         file.write(f"# Trefoil-knot pretzel: {len(vertices)} vertices,")
@@ -108,14 +113,16 @@ def write_mesh(path):
     print(f"Wrote {path} ({len(vertices)} vertices, {len(faces)} faces)")
 
 
-def blend(low, high, amount):
+def blend(low: Color, high: Color, amount: float) -> Color:
     return tuple(
         round(low[channel] + (high[channel] - low[channel]) * amount)
         for channel in range(3)
     )
 
 
-def write_wallpaper(path, top, bottom, glow):
+def write_wallpaper(
+    path: pathlib.Path, top: Color, bottom: Color, glow: Color
+) -> None:
     size = WALLPAPER_SIZE
     lights = [
         (size * 0.25, size * 0.3, size * 0.22),
@@ -139,7 +146,7 @@ def write_wallpaper(path, top, bottom, glow):
     print(f"Wrote {path} ({size}x{size})")
 
 
-def main():
+def main() -> None:
     ASSETS_DIR.mkdir(parents=True, exist_ok=True)
     write_mesh(ASSETS_DIR / "pretzel.mdl")
     for name, (top, bottom, glow) in WALLPAPERS.items():

python315-sampling-profiler/src/pretzel/assets.py

@@ -2,11 +2,12 @@
 
 from dataclasses import dataclass
 from importlib.resources import files
+from typing import BinaryIO
 
 import numpy as np
 import numpy.typing as npt
 
-from pretzel.engine import Mesh
+from pretzel.engine import Face, Mesh, Vertex
 
 ASSETS_DIR = files("pretzel") / "assets"
 
@@ -18,12 +19,13 @@ class Background:
     pixels: npt.NDArray[np.float32]
 
 
-def asset_path(name):
+def asset_path(name: str) -> str:
     return str(ASSETS_DIR / name)
 
 
-def load_mesh(name):
-    vertices, faces = [], []
+def load_mesh(name: str) -> Mesh:
+    vertices: list[Vertex] = []
+    faces: list[Face] = []
     with open(asset_path(name), encoding="utf-8") as file:
         for line in file:
             kind, _, rest = line.partition(" ")
@@ -36,13 +38,13 @@ def load_mesh(name):
     return Mesh(vertices, faces)
 
 
-def load_background(name, fast=False):
+def load_background(name: str, fast: bool = False) -> Background:
     path = asset_path(name)
     decode = decode_ppm if fast else decode_ppm_pixel_by_pixel
     return Background(name=name, path=path, pixels=decode(path))
 
 
-def decode_ppm_pixel_by_pixel(path):
+def decode_ppm_pixel_by_pixel(path: str) -> npt.NDArray[np.float32]:
     with open(path, "rb") as file:
         width, height = read_ppm_header(file)
         rows = []
@@ -58,11 +60,11 @@ def decode_ppm_pixel_by_pixel(path):
     return np.array(rows, dtype=np.float32)
 
 
-def to_linear(value):
+def to_linear(value: int) -> float:
     return 255.0 * (value / 255.0) ** 2.2
 
 
-def decode_ppm(path):
+def decode_ppm(path: str) -> npt.NDArray[np.float32]:
     with open(path, "rb") as file:
         width, height = read_ppm_header(file)
         data = file.read(width * height * 3)
@@ -71,7 +73,7 @@ def decode_ppm(path):
     return 255.0 * (pixels / 255.0) ** 2.2
 
 
-def read_ppm_header(file):
+def read_ppm_header(file: BinaryIO) -> tuple[int, int]:
     magic = file.readline().strip()
     if magic != b"P6":
         raise ValueError(f"{file.name} is not a binary PPM file")

python315-sampling-profiler/src/pretzel/cli.py

@@ -2,16 +2,18 @@
 
 import argparse
 import time
+from typing import Final
 
 from pretzel import assets
+from pretzel.engine import Mesh
 from pretzel.render import compute_frame
 from pretzel.streaming import BackgroundStreamer
 from pretzel.telemetry import FrameLog
 
-MODEL = "pretzel.mdl"
+MODEL: Final = "pretzel.mdl"
 
 
-def main():
+def main() -> None:
     parser = argparse.ArgumentParser(prog="pretzel", description=__doc__)
     parser.add_argument(
         "--frames",
@@ -47,15 +49,22 @@ def main():
         frame_log.close()
 
 
-def show_window(mesh, streamer, frame_log):
+def show_window(
+    mesh: Mesh, streamer: BackgroundStreamer, frame_log: FrameLog
+) -> None:
     from pretzel.viewer import Viewer
 
     Viewer(mesh, streamer, frame_log).run()
 
 
 def render_headlessly(
-    mesh, streamer, frame_log, frames, cache_colors=False, size=512
-):
+    mesh: Mesh,
+    streamer: BackgroundStreamer,
+    frame_log: FrameLog,
+    frames: int,
+    cache_colors: bool = False,
+    size: int = 512,
+) -> None:
     started = time.perf_counter()
     for frame in range(frames):
         frame_started = time.perf_counter()

python315-sampling-profiler/src/pretzel/engine.py

@@ -19,7 +19,7 @@ class Mesh:
     faces: list[Face]
 
 
-def rotation_matrix(pitch, yaw, roll):
+def rotation_matrix(pitch: float, yaw: float, roll: float) -> Matrix:
     sin_p, cos_p = math.sin(pitch), math.cos(pitch)
     sin_y, cos_y = math.sin(yaw), math.cos(yaw)
     sin_r, cos_r = math.sin(roll), math.cos(roll)
@@ -38,7 +38,9 @@ def rotation_matrix(pitch, yaw, roll):
     )
 
 
-def transform_vertices(vertices, matrix, distance):
+def transform_vertices(
+    vertices: list[Vertex], matrix: Matrix, distance: float
+) -> list[Vertex]:
     (xx, xy, xz), (yx, yy, yz), (zx, zy, zz) = matrix
     transformed = []
     for x, y, z in vertices:
@@ -49,7 +51,9 @@ def transform_vertices(vertices, matrix, distance):
     return transformed
 
 
-def project_vertices(transformed, width, height, focal_length):
+def project_vertices(
+    transformed: list[Vertex], width: int, height: int, focal_length: float
+) -> list[Point]:
     half_width, half_height = width / 2, height / 2
     projected = []
     for x, y, z in transformed:
@@ -58,7 +62,7 @@ def project_vertices(transformed, width, height, focal_length):
     return projected
 
 
-def cull_backfaces(faces, projected):
+def cull_backfaces(faces: list[Face], projected: list[Point]) -> list[Face]:
     visible = []
     for face in faces:
         a, b, c = face
@@ -70,8 +74,10 @@ def cull_backfaces(faces, projected):
     return visible
 
 
-def sort_back_to_front(faces, transformed):
-    def face_depth(face):
+def sort_back_to_front(
+    faces: list[Face], transformed: list[Vertex]
+) -> list[Face]:
+    def face_depth(face: Face) -> float:
         a, b, c = face
         return transformed[a][2] + transformed[b][2] + transformed[c][2]
 

python315-sampling-profiler/src/pretzel/render.py

@@ -1,13 +1,25 @@
 """Turn a mesh, a background, and a clock tick into shaded polygons."""
 
 import math
+from typing import Final
 
 from pretzel import engine, shading
+from pretzel.assets import Background
+from pretzel.engine import Mesh
 
-CAMERA_DISTANCE = 7.0
+CAMERA_DISTANCE: Final = 7.0
 
+type Polygon = tuple[list[float], str]
 
-def compute_frame(mesh, background, tick, width, height, cache_colors=False):
+
+def compute_frame(
+    mesh: Mesh,
+    background: Background,
+    tick: float,
+    width: int,
+    height: int,
+    cache_colors: bool = False,
+) -> list[Polygon]:
     matrix = engine.rotation_matrix(
         pitch=0.45 * math.sin(tick * 0.6),
         yaw=tick * 0.9,
@@ -23,9 +35,9 @@ def compute_frame(mesh, background, tick, width, height, cache_colors=False):
     ordered = engine.sort_back_to_front(visible, transformed)
     ambient = shading.sample_ambient_light(background.pixels, tick)
     channels = shading.shade_faces(ordered, transformed, ambient)
-    polygons = []
+    polygons: list[Polygon] = []
     for face, (red, green, blue) in zip(ordered, channels):
-        coordinates = []
+        coordinates: list[float] = []
         for index in face:
             coordinates.extend(projected[index])
         polygons.append(
@@ -39,10 +51,10 @@ def compute_frame(mesh, background, tick, width, height, cache_colors=False):
     return polygons
 
 
-_HEX_CACHE = {}
+_HEX_CACHE: dict[tuple[int, int, int], str] = {}
 
 
-def hex_color(red, green, blue):
+def hex_color(red: int, green: int, blue: int) -> str:
     key = (red, green, blue)
     color = _HEX_CACHE.get(key)
     if color is None:

python315-sampling-profiler/src/pretzel/shading.py

@@ -1,14 +1,19 @@
 """NumPy-powered lighting. Most of this work runs in native code."""
 
 import numpy as np
+import numpy.typing as npt
+
+from pretzel.engine import Face, Vertex
 
 LIGHT_DIRECTION = np.array([0.35, 0.55, -0.75])
 LIGHT_DIRECTION /= np.linalg.norm(LIGHT_DIRECTION)
 
 CRUST_COLOR = np.array([224.0, 147.0, 65.0])
 
 
-def sample_ambient_light(pixels, tick):
+def sample_ambient_light(
+    pixels: npt.NDArray[np.float32], tick: float
+) -> float:
     samples = pixels[::2, ::2]
     height, width = samples.shape[:2]
     rows = np.arange(height, dtype=np.float32)[:, np.newaxis]
@@ -18,7 +23,9 @@ def sample_ambient_light(pixels, tick):
     return 0.15 + 0.3 * float(lit.mean()) / 255.0
 
 
-def shade_faces(faces, transformed, ambient):
+def shade_faces(
+    faces: list[Face], transformed: list[Vertex], ambient: float
+) -> list[list[int]]:
     vertices = np.asarray(transformed)
     corners = vertices[np.asarray(faces)]
     edges_ab = corners[:, 1] - corners[:, 0]