UML class diagrams that interviewers can read

parking/draw.py

"""Draws the parking lot's class diagram from its code, at three levels of detail.

It reads model.py with inspect and typing, counts what each level would show, and
prints the D2 source of the level that suits an interview. A name that starts with
one underscore is private by convention: it is shown with "-" and without the
underscore. Every other name is shown with "+".
"""
import enum
import inspect
import types
import typing
from typing import Annotated, Protocol, get_args, get_origin

import model

CLASSES = [
    obj
    for obj in vars(model).values()
    if inspect.isclass(obj) and obj.__module__ == "model" and obj not in (model.Far, model.Owned)
]
LEVELS = ("names", "interview", "full")


def shown(name):
    """+name for public names; -name, without the underscore, for private ones."""
    private = name.startswith("_") and not name.startswith("__")
    return f"-{name[1:]}" if private else f"+{name}"


def unpack(hint):
    """(the type, the multiplicity at its end, the Annotated extras) of one hint."""
    extras = ()
    if get_origin(hint) is Annotated:
        hint, *extras = get_args(hint)
    origin, args = get_origin(hint), get_args(hint)
    if origin in (list, set, tuple):
        return args[0], "0..*", extras
    if origin is dict:
        return args[1], "0..*", extras
    if origin in (typing.Union, types.UnionType) and type(None) in args:
        return next(a for a in args if a is not type(None)), "0..1", extras
    return hint, "1", extras


def is_line(target):
    """Classes of the model are drawn as lines; enumerations stay attributes."""
    return target in CLASSES and not issubclass(target, enum.Enum)


def type_name(t):
    return getattr(t, "__name__", str(t))


def members(cls, level):
    if level == "names":
        return []
    if issubclass(cls, enum.Enum):
        return [(literal.name, "") for literal in cls]
    rows = []
    for name, hint in inspect.get_annotations(cls).items():
        target, many, _ = unpack(hint)
        if level == "interview" and (name.startswith("_") or is_line(target)):
            continue
        rows.append((shown(name), type_name(target) + ("" if many == "1" else f"[{many}]")))
    for name, fn in vars(cls).items():
        # Skip what Python adds by itself: annotation helpers and typing.Protocol's __init__.
        if not inspect.isfunction(fn) or name.startswith("__annotate") or fn.__module__ != model.__name__:
            continue
        if level == "interview" and (name.startswith("_") or name.startswith("get_")):
            continue
        params = ", ".join(p for p in inspect.signature(fn).parameters if p != "self")
        returns = inspect.get_annotations(fn).get("return")
        if returns is None:
            result = "None"
        else:
            target, many, _ = unpack(returns)
            result = type_name(target) + ("" if many == "1" else f"[{many}]")
        if getattr(fn, "__isabstractmethod__", False):
            result += " {abstract}"
        rows.append((f"{shown(name)}({params})", result))
    return rows


def lines():
    out = []
    for cls in CLASSES:
        for base in cls.__bases__:
            if base in CLASSES:
                kind = "realization" if Protocol in base.__bases__ else "generalization"
                out.append((kind, cls, base, "", "", ""))
        for name, hint in inspect.get_annotations(cls).items():
            target, many, extras = unpack(hint)
            if not is_line(target):
                continue
            owned = model.OWNED in extras
            far = next((e.multiplicity for e in extras if isinstance(e, model.Far)), "1" if owned else "")
            out.append(("composition" if owned else "association", cls, target, name, many, far))
    return out


def label(cls):
    if Protocol in cls.__bases__:
        return f"«interface»\\n{cls.__name__}"
    if issubclass(cls, enum.Enum):
        return f"«enumeration»\\n{cls.__name__}"
    if inspect.isabstract(cls):
        return f"{cls.__name__} {{abstract}}"
    return cls.__name__


def d2(level):
    out = ["direction: down"]
    for cls in CLASSES:
        out.append(f"{cls.__name__}: {{")
        out.append("  shape: class")
        out.append(f'  label: "{label(cls)}"')
        out.extend(f'  "{key}": "{value}"' for key, value in members(cls, level))
        out.append("}")
    for kind, a, b, role, many, far in lines():
        if kind in ("generalization", "realization"):
            dash = " style.stroke-dash: 4;" if kind == "realization" else ""
            out.append(f"{a.__name__} -> {b.__name__}: {{{dash} target-arrowhead: {{shape: triangle; style.filled: false}} }}")
        elif kind == "composition":
            out.append(
                f'{a.__name__} <-> {b.__name__}: "{role}" {{ source-arrowhead: {{shape: diamond; style.filled: true; label: "{far}"}};'
                f' target-arrowhead: {{shape: arrow; label: "{many}"}} }}'
            )
        else:
            out.append(f'{a.__name__} -> {b.__name__}: "{role}" {{ source-arrowhead.label: "{far}"; target-arrowhead: {{shape: arrow; label: "{many}"}} }}')
    return "\n".join(out)


print(f"{'level':<10} {'boxes':>5} {'members':>8} {'lines':>6}")
for level in LEVELS:
    count = sum(len(members(cls, level)) for cls in CLASSES)
    print(f"{level:<10} {len(CLASSES):>5} {count:>8} {len(lines()):>6}")
print()
print("D2 source of the interview level:")
print(d2("interview"))

level      boxes  members  lines
names          9        0      8
interview      9       21      8
full           9       38      8

D2 source of the interview level:
direction: down
SpotSize: {
  shape: class
  label: "«enumeration»\nSpotSize"
  "SMALL": ""
  "MEDIUM": ""
  "LARGE": ""
}
Vehicle: {
  shape: class
  label: "Vehicle {abstract}"
  "+plate": "str"
  "+size()": "SpotSize {abstract}"
}
Car: {
  shape: class
  label: "Car"
  "+size()": "SpotSize"
}
Motorcycle: {
  shape: class
  label: "Motorcycle"
  "+size()": "SpotSize"
}
Spot: {
  shape: class
  label: "Spot"
  "+number": "int"
  "+size": "SpotSize"
  "+fits(vehicle)": "bool"
  "+park(vehicle)": "None"
  "+leave()": "Vehicle[0..1]"
}
Ticket: {
  shape: class
  label: "Ticket"
  "+number": "int"
  "+entered_at": "int"
  "+minutes_until(now)": "int"
}
FeePolicy: {
  shape: class
  label: "«interface»\nFeePolicy"
  "+fee(minutes)": "int"
}
HourlyFee: {
  shape: class
  label: "HourlyFee"
  "+rate": "int"
  "+fee(minutes)": "int"
}
ParkingLot: {
  shape: class
  label: "ParkingLot"
  "+enter(vehicle, now)": "Ticket"
  "+exit(ticket_number, now)": "int"
  "+free_spots(size)": "int"
}
Car -> Vehicle: { target-arrowhead: {shape: triangle; style.filled: false} }
Motorcycle -> Vehicle: { target-arrowhead: {shape: triangle; style.filled: false} }
Spot -> Vehicle: "vehicle" { source-arrowhead.label: "0..1"; target-arrowhead: {shape: arrow; label: "0..1"} }
Ticket -> Spot: "spot" { source-arrowhead.label: "0..1"; target-arrowhead: {shape: arrow; label: "1"} }
HourlyFee -> FeePolicy: { style.stroke-dash: 4; target-arrowhead: {shape: triangle; style.filled: false} }
ParkingLot <-> Spot: "spots" { source-arrowhead: {shape: diamond; style.filled: true; label: "1"}; target-arrowhead: {shape: arrow; label: "0..*"} }
ParkingLot <-> Ticket: "tickets" { source-arrowhead: {shape: diamond; style.filled: true; label: "1"}; target-arrowhead: {shape: arrow; label: "0..*"} }
ParkingLot -> FeePolicy: "fee_policy" { source-arrowhead.label: "0..*"; target-arrowhead: {shape: arrow; label: "1"} }

parking/model.py

"""The parking lot of the lesson, written so that a program can read its structure.

Type hints give the type of every attribute. Annotated adds the two things a hint
cannot say: OWNED marks parts that the class creates and keeps (composition), and
Far gives the multiplicity at the other end of an association.
"""
from abc import ABC, abstractmethod
from dataclasses import dataclass
from enum import Enum
from typing import Annotated, Protocol


@dataclass(frozen=True)
class Far:
    multiplicity: str


class Owned:
    pass


OWNED = Owned()


class SpotSize(Enum):
    SMALL = 1
    MEDIUM = 2
    LARGE = 3


class Vehicle(ABC):
    plate: str

    def __init__(self, plate: str) -> None:
        self.plate = plate

    @abstractmethod
    def size(self) -> SpotSize:
        """The smallest spot the vehicle fits in."""

    def get_plate(self) -> str:
        return self.plate


class Car(Vehicle):
    def size(self) -> SpotSize:
        return SpotSize.MEDIUM


class Motorcycle(Vehicle):
    def size(self) -> SpotSize:
        return SpotSize.SMALL


class Spot:
    number: int
    size: SpotSize
    vehicle: Annotated[Vehicle | None, Far("0..1")]

    def __init__(self, number: int, size: SpotSize) -> None:
        self.number = number
        self.size = size
        self.vehicle = None

    def fits(self, vehicle: Vehicle) -> bool:
        return self.vehicle is None and vehicle.size().value <= self.size.value

    def park(self, vehicle: Vehicle) -> None:
        self.vehicle = vehicle

    def leave(self) -> Vehicle | None:
        vehicle, self.vehicle = self.vehicle, None
        return vehicle

    def get_number(self) -> int:
        return self.number

    def __repr__(self) -> str:
        return f"Spot({self.number}, {self.size.name})"


class Ticket:
    number: int
    entered_at: int
    spot: Annotated[Spot, Far("0..1")]

    def __init__(self, number: int, spot: Spot, entered_at: int) -> None:
        self.number = number
        self.spot = spot
        self.entered_at = entered_at

    def minutes_until(self, now: int) -> int:
        return now - self.entered_at


class FeePolicy(Protocol):
    def fee(self, minutes: int) -> int: ...


class HourlyFee(FeePolicy):
    rate: int
    _free_minutes: int

    def __init__(self, rate: int, free_minutes: int = 15) -> None:
        self.rate = rate
        self._free_minutes = free_minutes

    def fee(self, minutes: int) -> int:
        if minutes <= self._free_minutes:
            return 0
        return -(-minutes // 60) * self.rate


class ParkingLot:
    spots: Annotated[list[Spot], OWNED]
    tickets: Annotated[dict[int, Ticket], OWNED]
    fee_policy: Annotated[FeePolicy, Far("0..*")]
    _next_number: int

    def __init__(self, sizes: list[SpotSize], fee_policy: FeePolicy) -> None:
        self.spots = [Spot(n, size) for n, size in enumerate(sizes, start=1)]
        self.tickets = {}
        self.fee_policy = fee_policy
        self._next_number = 1

    def enter(self, vehicle: Vehicle, now: int) -> Ticket:
        spot = self._find_spot(vehicle)
        if spot is None:
            raise LookupError(f"no free spot for {vehicle.plate}")
        spot.park(vehicle)
        ticket = Ticket(self._next_number, spot, now)
        self._next_number += 1
        self.tickets[ticket.number] = ticket
        return ticket

    def exit(self, ticket_number: int, now: int) -> int:
        ticket = self.tickets.pop(ticket_number)
        ticket.spot.leave()
        return self.fee_policy.fee(ticket.minutes_until(now))

    def free_spots(self, size: SpotSize) -> int:
        return sum(1 for spot in self.spots if spot.vehicle is None and spot.size is size)

    def _find_spot(self, vehicle: Vehicle) -> Spot | None:
        return next((spot for spot in self.spots if spot.fits(vehicle)), None)

    def _log(self, message: str) -> None:
        print(message)