Source code for pytyche.viz._cells

"""Cell-allocation bar chart."""

from __future__ import annotations

from collections.abc import Sequence
from typing import Any

import matplotlib.pyplot as plt
from matplotlib.axes import Axes




def plot_cells(cells: Sequence[Any], ax: Axes | None = None) -> Axes:
    """Horizontal bar chart of per-cell traffic weights.

    Args:
        cells: Objects carrying ``id`` (str) and ``weight`` (float)
            attributes — the shipped L1 ``Cell`` shape, consumed
            duck-typed.
        ax: Existing axes to draw into; ``None`` creates a new figure.

    Returns:
        The axes containing one bar per cell, top-to-bottom in list order.
    """
    if ax is None:
        _, ax = plt.subplots()
    positions = range(len(cells))
    ax.barh(positions, [float(c.weight) for c in cells])
    ax.set_yticks(positions, [c.id for c in cells])
    ax.invert_yaxis()
    ax.set_xlim(0.0, 1.0)
    ax.set_xlabel("traffic share")
    ax.set_title("cell allocation")
    return ax