From a2025a074fa33d37206542379972a0cd96827e2a Mon Sep 17 00:00:00 2001
From: Jiahao Yuan <kaho0769@qq.com>
Date: Wed, 6 Nov 2024 21:03:03 +0800
Subject: [PATCH] feat: element label in plot (#221)

---
 examples/plot.py                | 13 +-----
 python/bosing/_bosing.pyi       | 36 ++++++++++++--
 python/bosing/_plot.py          | 34 +++++++++++---
 src/python.rs                   |  2 +-
 src/python/elements.rs          | 83 ++++++++++++++++++++++++---------
 src/python/elements/absolute.rs |  5 +-
 src/python/elements/grid.rs     |  5 +-
 src/python/elements/stack.rs    |  5 +-
 src/python/plot.rs              | 30 ++++++++++--
 src/quant.rs                    |  2 +
 src/schedule.rs                 | 13 +++++-
 stubtest-allowlist.txt          |  2 +
 12 files changed, 177 insertions(+), 53 deletions(-)

diff --git a/examples/plot.py b/examples/plot.py
index a56ec09..58bc59d 100644
--- a/examples/plot.py
+++ b/examples/plot.py
@@ -2,7 +2,7 @@
 
 from bosing import Barrier, Play, ShiftPhase, Stack
 
-xy = [Play(f"xy{i}", "hann", 1.0, 100e-9) for i in range(2)]
+xy = [Play(f"xy{i}", "hann", 1.0, 100e-9, label=f"xy{i}") for i in range(2)]
 z = [
     Stack(Play(f"z{i}", "hann", 1.0, 100e-9), ShiftPhase(f"xy{i}", 1.0))
     for i in range(2)
@@ -10,16 +10,7 @@
 m = Stack(*(Play(f"m{i}", "hann", 1.0, 100e-9, plateau=200e-9) for i in range(2)))
 b = Barrier()
 
-schedule = Stack(
-    xy[0],
-    xy[1],
-    b,
-    z[1],
-    b,
-    xy[1],
-    b,
-    m,
-)
+schedule = Stack(xy[0], xy[1], b, z[1], b, xy[1], b, m, label="root")
 
 schedule.plot()
 plt.show()
diff --git a/python/bosing/_bosing.pyi b/python/bosing/_bosing.pyi
index 5c6c5f6..18310a9 100644
--- a/python/bosing/_bosing.pyi
+++ b/python/bosing/_bosing.pyi
@@ -1,6 +1,6 @@
 # ruff: noqa: PLR0913
-from collections.abc import Iterable, Mapping, Sequence
-from typing import Any, ClassVar, Literal, final, type_check_only
+from collections.abc import Iterable, Iterator, Mapping, Sequence
+from typing import Any, ClassVar, Literal, final
 
 import numpy as np
 import numpy.typing as npt
@@ -34,6 +34,8 @@ __all__ = [
     "generate_waveforms",
     "generate_waveforms_with_states",
     "ItemKind",
+    "PlotArgs",
+    "PlotItem",
 ]
 
 _RichReprResult: TypeAlias = list[Any]
@@ -124,6 +126,8 @@ class Element:
     def max_duration(self) -> float: ...
     @property
     def min_duration(self) -> float: ...
+    @property
+    def label(self) -> str: ...
     def measure(self) -> float: ...
     def plot(
         self,
@@ -131,6 +135,7 @@ class Element:
         *,
         channels: Sequence[str] | None = ...,
         max_depth: int = ...,
+        show_label: bool = ...,
     ) -> Axes: ...
 
 @final
@@ -155,6 +160,7 @@ class Play(Element):
         duration: float | None = ...,
         max_duration: float = ...,
         min_duration: float = ...,
+        label: str | None = ...,
     ) -> Self: ...
     @property
     def channel_id(self) -> str: ...
@@ -191,6 +197,7 @@ class ShiftPhase(Element):
         duration: float | None = ...,
         max_duration: float = ...,
         min_duration: float = ...,
+        label: str | None = ...,
     ) -> Self: ...
     @property
     def channel_id(self) -> str: ...
@@ -213,6 +220,7 @@ class SetPhase(Element):
         duration: float | None = ...,
         max_duration: float = ...,
         min_duration: float = ...,
+        label: str | None = ...,
     ) -> Self: ...
     @property
     def channel_id(self) -> str: ...
@@ -235,6 +243,7 @@ class ShiftFreq(Element):
         duration: float | None = ...,
         max_duration: float = ...,
         min_duration: float = ...,
+        label: str | None = ...,
     ) -> Self: ...
     @property
     def channel_id(self) -> str: ...
@@ -257,6 +266,7 @@ class SetFreq(Element):
         duration: float | None = ...,
         max_duration: float = ...,
         min_duration: float = ...,
+        label: str | None = ...,
     ) -> Self: ...
     @property
     def channel_id(self) -> str: ...
@@ -279,6 +289,7 @@ class SwapPhase(Element):
         duration: float | None = ...,
         max_duration: float = ...,
         min_duration: float = ...,
+        label: str | None = ...,
     ) -> Self: ...
     @property
     def channel_id1(self) -> str: ...
@@ -299,6 +310,7 @@ class Barrier(Element):
         duration: float | None = ...,
         max_duration: float = ...,
         min_duration: float = ...,
+        label: str | None = ...,
     ) -> Self: ...
     @property
     def channel_ids(self) -> Sequence[str]: ...
@@ -320,6 +332,7 @@ class Repeat(Element):
         duration: float | None = ...,
         max_duration: float = ...,
         min_duration: float = ...,
+        label: str | None = ...,
     ) -> Self: ...
     @property
     def child(self) -> Element: ...
@@ -350,6 +363,7 @@ class Stack(Element):
         duration: float | None = ...,
         max_duration: float = ...,
         min_duration: float = ...,
+        label: str | None = ...,
     ) -> Self: ...
     def with_children(self, *children: Element) -> Stack: ...
     @property
@@ -384,6 +398,7 @@ class Absolute(Element):
         duration: float | None = ...,
         max_duration: float = ...,
         min_duration: float = ...,
+        label: str | None = ...,
     ) -> Self: ...
     def with_children(self, *children: _AbsoluteEntryLike) -> Absolute: ...
     @property
@@ -444,6 +459,7 @@ class Grid(Element):
         duration: float | None = ...,
         max_duration: float = ...,
         min_duration: float = ...,
+        label: str | None = ...,
     ) -> Self: ...
     def with_children(
         self,
@@ -471,7 +487,19 @@ class OscState:
     def with_time_shift(self, time: float) -> Self: ...
     def __rich_repr__(self) -> _RichReprResult: ...  # undocumented
 
-@type_check_only
+@final
+class PlotArgs:
+    @property
+    def ax(self) -> Axes | None: ...
+    @property
+    def blocks(self) -> Iterator[PlotItem]: ...
+    @property
+    def channels(self) -> list[str]: ...
+    @property
+    def max_depth(self) -> int: ...
+    @property
+    def show_label(self) -> bool: ...
+
 @final
 class PlotItem:
     @property
@@ -484,6 +512,8 @@ class PlotItem:
     def depth(self) -> int: ...
     @property
     def kind(self) -> ItemKind: ...
+    @property
+    def label(self) -> str | None: ...
 
 @final
 class ItemKind:
diff --git a/python/bosing/_plot.py b/python/bosing/_plot.py
index ef75f76..5105e01 100644
--- a/python/bosing/_plot.py
+++ b/python/bosing/_plot.py
@@ -5,6 +5,7 @@
 
 import matplotlib.pyplot as plt
 import numpy as np
+from matplotlib import patheffects
 from matplotlib.patches import PathPatch
 from matplotlib.path import Path
 from matplotlib.ticker import EngFormatter
@@ -17,10 +18,11 @@
     from matplotlib.axes import Axes
     from typing_extensions import TypeAlias
 
-    from bosing._bosing import PlotItem
+    from bosing._bosing import PlotArgs, PlotItem
 
     _RECTS: TypeAlias = defaultdict[ItemKind, list[tuple[float, float, float]]]
     _MARKERS: TypeAlias = defaultdict[ItemKind, tuple[list[float], list[float]]]
+    _TEXTS: TypeAlias = list[tuple[float, float, str]]
 
 
 COLORS = {
@@ -87,10 +89,11 @@ def process_blocks(
     channels: Sequence[str],
     max_depth: int,
     channels_ystart: dict[str, int],
-) -> tuple[_RECTS, _MARKERS]:
+) -> tuple[_RECTS, _MARKERS, _TEXTS]:
     ch_stack: list[list[str]] = []
     rects: _RECTS = defaultdict(list)
     markers: _MARKERS = defaultdict(lambda: ([], []))
+    texts: _TEXTS = []
 
     for x in blocks:
         manage_channel_stack(ch_stack, x)
@@ -105,17 +108,23 @@ def process_blocks(
                     my.append(y)
                 else:
                     rects[x.kind].append((x.start, y, x.span))
-    return rects, markers
+                if x.label is not None:
+                    texts.append((x.start, y, x.label))
+    return rects, markers, texts
 
 
-def plot(
-    ax: Axes | None, blocks: Iterator[PlotItem], channels: Sequence[str], max_depth: int
-) -> Axes:
+def plot(args: PlotArgs) -> Axes:
+    ax = args.ax
+    blocks = args.blocks
+    channels = args.channels
+    max_depth = args.max_depth
+    show_label = args.show_label
+
     if ax is None:
         ax = plt.gca()
 
     channels_ystart = {c: i * (max_depth + 1) for i, c in enumerate(channels)}
-    rects, markers = process_blocks(blocks, channels, max_depth, channels_ystart)
+    rects, markers, texts = process_blocks(blocks, channels, max_depth, channels_ystart)
 
     for k, r in rects.items():
         # numrects x [x, y, width]
@@ -143,6 +152,17 @@ def plot(
             markersize=12,
         )
 
+    if show_label:
+        for x, y, label in texts:
+            txt = ax.annotate(label, (x, y))  # pyright: ignore[reportUnknownMemberType]
+            # Add white outline to text
+            txt.set_path_effects(
+                [
+                    patheffects.Stroke(linewidth=2, foreground="white"),
+                    patheffects.Normal(),
+                ]
+            )
+
     _ = ax.set_yticks(list(channels_ystart.values()), channels_ystart.keys())  # pyright: ignore[reportUnknownMemberType]
     ax.xaxis.set_major_formatter(EngFormatter(places=3))
     _ = ax.set_xlabel("Time")  # pyright: ignore[reportUnknownMemberType]
diff --git a/src/python.rs b/src/python.rs
index d0faff3..ae5bd8d 100644
--- a/src/python.rs
+++ b/src/python.rs
@@ -17,7 +17,7 @@ mod export {
             GridLength, GridLengthUnit, Play, Repeat, SetFreq, SetPhase, ShiftFreq, ShiftPhase,
             Stack, SwapPhase,
         },
-        plot::ItemKind,
+        plot::{ItemKind, PlotArgs, PlotItem},
         shapes::{Hann, Interp, Shape},
         wavegen::{generate_waveforms, generate_waveforms_with_states, Channel, OscState},
     };
diff --git a/src/python/elements.rs b/src/python/elements.rs
index 51a31bd..981ed82 100644
--- a/src/python/elements.rs
+++ b/src/python/elements.rs
@@ -7,7 +7,7 @@ use std::{borrow::Borrow as _, fmt::Debug, sync::Arc};
 use pyo3::{exceptions::PyValueError, prelude::*, types::DerefToPyAny};
 
 use crate::{
-    quant::{Amplitude, ChannelId, Frequency, Phase, ShapeId, Time},
+    quant::{Amplitude, ChannelId, Frequency, Label, Phase, ShapeId, Time},
     schedule::{
         self, ElementCommon, ElementCommonBuilder, ElementRef, ElementVariant, Measure as _,
     },
@@ -87,24 +87,6 @@ impl Alignment {
     }
 }
 
-fn extract_alignment(obj: &Bound<PyAny>) -> PyResult<Alignment> {
-    Alignment::convert(obj).and_then(|x| x.extract(obj.py()))
-}
-
-fn extract_margin(obj: &Bound<PyAny>) -> PyResult<(Time, Time)> {
-    if let Ok(v) = obj.extract() {
-        let t = Time::new(v)?;
-        return Ok((t, t));
-    }
-    if let Ok((v1, v2)) = obj.extract() {
-        let t1 = Time::new(v1)?;
-        let t2 = Time::new(v2)?;
-        return Ok((t1, t2));
-    }
-    let msg = "Failed to convert the value to (float, float).";
-    Err(PyValueError::new_err(msg))
-}
-
 /// Base class for schedule elements.
 ///
 /// A schedule element is a node in the tree structure of a schedule similar to
@@ -179,6 +161,7 @@ fn extract_margin(obj: &Bound<PyAny>) -> PyResult<(Time, Time)> {
 ///     max_duration (float): Maximum duration of the element. Defaults to
 ///         ``inf``.
 ///     min_duration (float): Minimum duration of the element. Defaults to ``0``.
+///     label (str | None): Label of the element. Defaults to ``None``.
 #[pyclass(module = "bosing", subclass, frozen)]
 #[derive(Debug, Clone)]
 pub(crate) struct Element(pub(super) ElementRef);
@@ -215,6 +198,11 @@ impl Element {
         self.0.common.min_duration()
     }
 
+    #[getter]
+    fn label(&self) -> Option<&Label> {
+        self.0.common.label()
+    }
+
     /// Measure the minimum total duration required by the element.
     ///
     /// This value includes both inner `duration` and outer `margin` of the element.
@@ -234,18 +222,20 @@ impl Element {
     ///     channels (Sequence[str] | None): Channels to plot. If ``None``, all channels are
     ///         plotted.
     ///     max_depth (int): Maximum depth to plot. Defaults to ``5``.
+    ///     show_label (bool): Whether to show label of elements. Defaults to ``True``.
     ///
     /// Returns:
     ///     matplotlib.axes.Axes: Axes with the plot.
-    #[pyo3(signature = (ax=None, *, channels=None, max_depth=5))]
+    #[pyo3(signature = (ax=None, *, channels=None, max_depth=5, show_label=true))]
     fn plot(
         &self,
         py: Python,
         ax: Option<PyObject>,
         channels: Option<Vec<ChannelId>>,
         max_depth: usize,
+        show_label: bool,
     ) -> PyResult<PyObject> {
-        plot_element(py, self.0.clone(), ax, channels, max_depth)
+        plot_element(py, self.0.clone(), ax, channels, max_depth, show_label)
     }
 }
 
@@ -278,6 +268,7 @@ where
             .expect("Element should have a valid variant")
     }
 
+    #[allow(clippy::too_many_arguments)]
     fn build_element(
         variant: Self::Variant,
         margin: Option<&Bound<PyAny>>,
@@ -286,6 +277,7 @@ where
         duration: Option<Time>,
         max_duration: Time,
         min_duration: Time,
+        label: Option<Label>,
     ) -> PyResult<Element> {
         let mut builder = ElementCommonBuilder::new();
         if let Some(obj) = margin {
@@ -298,9 +290,28 @@ where
             .phantom(phantom)
             .duration(duration)
             .max_duration(max_duration)
-            .min_duration(min_duration);
+            .min_duration(min_duration)
+            .label(label);
         let common = builder.build()?;
-        Ok(Element(Arc::new(schedule::Element::new(common, variant))))
+        return Ok(Element(Arc::new(schedule::Element::new(common, variant))));
+
+        fn extract_alignment(obj: &Bound<PyAny>) -> PyResult<Alignment> {
+            Alignment::convert(obj).and_then(|x| x.extract(obj.py()))
+        }
+
+        fn extract_margin(obj: &Bound<PyAny>) -> PyResult<(Time, Time)> {
+            if let Ok(v) = obj.extract() {
+                let t = Time::new(v)?;
+                return Ok((t, t));
+            }
+            if let Ok((v1, v2)) = obj.extract() {
+                let t1 = Time::new(v1)?;
+                let t2 = Time::new(v2)?;
+                return Ok((t1, t2));
+            }
+            let msg = "Failed to convert the value to (float, float).";
+            Err(PyValueError::new_err(msg))
+        }
     }
 }
 
@@ -326,6 +337,7 @@ where
         push_repr!(res, py, "duration", slf.duration(), None);
         push_repr!(res, py, "max_duration", slf.max_duration(), Time::INFINITY);
         push_repr!(res, py, "min_duration", slf.min_duration(), Time::ZERO);
+        push_repr!(res, py, "label", slf.label(), None);
         res.into_iter()
     }
 }
@@ -414,6 +426,7 @@ impl Play {
         duration=None,
         max_duration=Time::INFINITY,
         min_duration=Time::ZERO,
+        label=None,
     ))]
     #[allow(clippy::too_many_arguments)]
     fn new(
@@ -432,6 +445,7 @@ impl Play {
         duration: Option<Time>,
         max_duration: Time,
         min_duration: Time,
+        label: Option<Label>,
     ) -> PyResult<(Self, Element)> {
         let variant = schedule::Play::new(channel_id, shape_id, amplitude, width)?
             .with_plateau(plateau)?
@@ -449,6 +463,7 @@ impl Play {
                 duration,
                 max_duration,
                 min_duration,
+                label,
             )?,
         ))
     }
@@ -549,6 +564,7 @@ impl ShiftPhase {
         duration=None,
         max_duration=Time::INFINITY,
         min_duration=Time::ZERO,
+        label=None,
     ))]
     #[allow(clippy::too_many_arguments)]
     fn new(
@@ -560,6 +576,7 @@ impl ShiftPhase {
         duration: Option<Time>,
         max_duration: Time,
         min_duration: Time,
+        label: Option<Label>,
     ) -> PyResult<(Self, Element)> {
         let variant = schedule::ShiftPhase::new(channel_id, phase)?;
         Ok((
@@ -572,6 +589,7 @@ impl ShiftPhase {
                 duration,
                 max_duration,
                 min_duration,
+                label,
             )?,
         ))
     }
@@ -645,6 +663,7 @@ impl SetPhase {
         duration=None,
         max_duration=Time::INFINITY,
         min_duration=Time::ZERO,
+        label=None,
     ))]
     #[allow(clippy::too_many_arguments)]
     fn new(
@@ -656,6 +675,7 @@ impl SetPhase {
         duration: Option<Time>,
         max_duration: Time,
         min_duration: Time,
+        label: Option<Label>,
     ) -> PyResult<(Self, Element)> {
         let variant = schedule::SetPhase::new(channel_id, phase)?;
         Ok((
@@ -668,6 +688,7 @@ impl SetPhase {
                 duration,
                 max_duration,
                 min_duration,
+                label,
             )?,
         ))
     }
@@ -729,6 +750,7 @@ impl ShiftFreq {
         duration=None,
         max_duration=Time::INFINITY,
         min_duration=Time::ZERO,
+        label=None,
     ))]
     #[allow(clippy::too_many_arguments)]
     fn new(
@@ -740,6 +762,7 @@ impl ShiftFreq {
         duration: Option<Time>,
         max_duration: Time,
         min_duration: Time,
+        label: Option<Label>,
     ) -> PyResult<(Self, Element)> {
         let variant = schedule::ShiftFreq::new(channel_id, frequency)?;
         Ok((
@@ -752,6 +775,7 @@ impl ShiftFreq {
                 duration,
                 max_duration,
                 min_duration,
+                label,
             )?,
         ))
     }
@@ -814,6 +838,7 @@ impl SetFreq {
         duration=None,
         max_duration=Time::INFINITY,
         min_duration=Time::ZERO,
+        label=None,
     ))]
     #[allow(clippy::too_many_arguments)]
     fn new(
@@ -825,6 +850,7 @@ impl SetFreq {
         duration: Option<Time>,
         max_duration: Time,
         min_duration: Time,
+        label: Option<Label>,
     ) -> PyResult<(Self, Element)> {
         let variant = schedule::SetFreq::new(channel_id, frequency)?;
         Ok((
@@ -837,6 +863,7 @@ impl SetFreq {
                 duration,
                 max_duration,
                 min_duration,
+                label,
             )?,
         ))
     }
@@ -901,6 +928,7 @@ impl SwapPhase {
         duration=None,
         max_duration=Time::INFINITY,
         min_duration=Time::ZERO,
+        label=None,
     ))]
     #[allow(clippy::too_many_arguments)]
     fn new(
@@ -912,6 +940,7 @@ impl SwapPhase {
         duration: Option<Time>,
         max_duration: Time,
         min_duration: Time,
+        label: Option<Label>,
     ) -> PyResult<(Self, Element)> {
         let variant = schedule::SwapPhase::new(channel_id1, channel_id2);
         Ok((
@@ -924,6 +953,7 @@ impl SwapPhase {
                 duration,
                 max_duration,
                 min_duration,
+                label,
             )?,
         ))
     }
@@ -986,7 +1016,9 @@ impl Barrier {
         duration=None,
         max_duration=Time::INFINITY,
         min_duration=Time::ZERO,
+        label=None,
     ))]
+    #[allow(clippy::too_many_arguments)]
     fn new(
         channel_ids: Vec<ChannelId>,
         margin: Option<&Bound<PyAny>>,
@@ -995,6 +1027,7 @@ impl Barrier {
         duration: Option<Time>,
         max_duration: Time,
         min_duration: Time,
+        label: Option<Label>,
     ) -> PyResult<(Self, Element)> {
         let variant = schedule::Barrier::new(channel_ids);
         Ok((
@@ -1007,6 +1040,7 @@ impl Barrier {
                 duration,
                 max_duration,
                 min_duration,
+                label,
             )?,
         ))
     }
@@ -1066,6 +1100,7 @@ impl Repeat {
         duration=None,
         max_duration=Time::INFINITY,
         min_duration=Time::ZERO,
+        label=None,
     ))]
     #[allow(clippy::too_many_arguments)]
     fn new(
@@ -1078,6 +1113,7 @@ impl Repeat {
         duration: Option<Time>,
         max_duration: Time,
         min_duration: Time,
+        label: Option<Label>,
     ) -> PyResult<(Self, Element)> {
         let rust_child = child.get().0.clone();
         let variant = schedule::Repeat::new(rust_child, count).with_spacing(spacing)?;
@@ -1091,6 +1127,7 @@ impl Repeat {
                 duration,
                 max_duration,
                 min_duration,
+                label,
             )?,
         ))
     }
diff --git a/src/python/elements/absolute.rs b/src/python/elements/absolute.rs
index 0f881e1..dd3ce2d 100644
--- a/src/python/elements/absolute.rs
+++ b/src/python/elements/absolute.rs
@@ -4,7 +4,7 @@ use pyo3::{exceptions::PyValueError, prelude::*};
 
 use crate::{quant::Time, schedule};
 
-use super::{Arg, Element, ElementSubclass, RichRepr};
+use super::{Arg, Element, ElementSubclass, Label, RichRepr};
 
 /// An absolute layout element.
 ///
@@ -59,6 +59,7 @@ impl Absolute {
         duration=None,
         max_duration=Time::INFINITY,
         min_duration=Time::ZERO,
+        label=None,
     ))]
     #[allow(clippy::too_many_arguments)]
     fn new(
@@ -70,6 +71,7 @@ impl Absolute {
         duration: Option<Time>,
         max_duration: Time,
         min_duration: Time,
+        label: Option<Label>,
     ) -> PyResult<(Self, Element)> {
         let children: Vec<AbsoluteEntry> = children
             .into_iter()
@@ -93,6 +95,7 @@ impl Absolute {
                 duration,
                 max_duration,
                 min_duration,
+                label,
             )?,
         ))
     }
diff --git a/src/python/elements/grid.rs b/src/python/elements/grid.rs
index 5f646d5..3e05779 100644
--- a/src/python/elements/grid.rs
+++ b/src/python/elements/grid.rs
@@ -4,7 +4,7 @@ use pyo3::{exceptions::PyValueError, prelude::*};
 
 use crate::{quant::Time, schedule};
 
-use super::{Arg, Element, ElementSubclass, RichRepr};
+use super::{Arg, Element, ElementSubclass, Label, RichRepr};
 
 /// A grid layout element.
 ///
@@ -85,6 +85,7 @@ impl Grid {
         duration=None,
         max_duration=Time::INFINITY,
         min_duration=Time::ZERO,
+        label=None,
     ))]
     #[allow(clippy::too_many_arguments)]
     fn new(
@@ -97,6 +98,7 @@ impl Grid {
         duration: Option<Time>,
         max_duration: Time,
         min_duration: Time,
+        label: Option<Label>,
     ) -> PyResult<(Self, Element)> {
         let children: Vec<_> = children
             .into_iter()
@@ -129,6 +131,7 @@ impl Grid {
                 duration,
                 max_duration,
                 min_duration,
+                label,
             )?,
         ))
     }
diff --git a/src/python/elements/stack.rs b/src/python/elements/stack.rs
index cf90f3f..941ea6d 100644
--- a/src/python/elements/stack.rs
+++ b/src/python/elements/stack.rs
@@ -4,7 +4,7 @@ use pyo3::{exceptions::PyValueError, prelude::*};
 
 use crate::{quant::Time, schedule};
 
-use super::{Arg, Element, ElementSubclass, RichRepr as _};
+use super::{Arg, Element, ElementSubclass, Label, RichRepr as _};
 
 /// A stack layout element.
 ///
@@ -57,6 +57,7 @@ impl Stack {
         duration=None,
         max_duration=Time::INFINITY,
         min_duration=Time::ZERO,
+        label=None,
     ))]
     #[allow(clippy::too_many_arguments)]
     fn new(
@@ -68,6 +69,7 @@ impl Stack {
         duration: Option<Time>,
         max_duration: Time,
         min_duration: Time,
+        label: Option<Label>,
     ) -> PyResult<(Self, Element)> {
         let rust_children = children.iter().map(|x| x.get().0.clone()).collect();
         let variant = schedule::Stack::new().with_children(rust_children);
@@ -86,6 +88,7 @@ impl Stack {
                 duration,
                 max_duration,
                 min_duration,
+                label,
             )?,
         ))
     }
diff --git a/src/python/plot.rs b/src/python/plot.rs
index c8b39ee..cf52d21 100644
--- a/src/python/plot.rs
+++ b/src/python/plot.rs
@@ -2,7 +2,7 @@ use itertools::Itertools as _;
 use pyo3::{prelude::*, sync::GILOnceCell, types::PyList};
 
 use crate::{
-    quant::{ChannelId, Time},
+    quant::{ChannelId, Label, Time},
     schedule::{Arrange as _, Arranged, ElementRef, ElementVariant, Measure, TimeRange},
     util::{pre_order_iter, IterVariant},
 };
@@ -32,6 +32,7 @@ pub(super) fn plot_element(
     ax: Option<PyObject>,
     channels: Option<Vec<ChannelId>>,
     max_depth: usize,
+    show_label: bool,
 ) -> PyResult<PyObject> {
     let channels = match channels {
         Some(channels) => PyList::new_bound(py, channels),
@@ -39,7 +40,7 @@ pub(super) fn plot_element(
     };
     let plot_items = Box::new(arrange_to_plot(root));
     let blocks = PlotIter { inner: plot_items };
-    call_plot(py, ax, blocks, channels, max_depth)
+    call_plot(py, ax, blocks, channels, max_depth, show_label)
 }
 
 impl ItemKind {
@@ -60,6 +61,16 @@ impl ItemKind {
     }
 }
 
+#[pyclass(module = "bosing._bosing", frozen, get_all)]
+#[derive(Debug)]
+pub(super) struct PlotArgs {
+    ax: Option<PyObject>,
+    blocks: Py<PlotIter>,
+    channels: Py<PyList>,
+    max_depth: usize,
+    show_label: bool,
+}
+
 #[pyclass(module = "bosing._bosing")]
 struct PlotIter {
     inner: Box<dyn Iterator<Item = PlotItem> + Send>,
@@ -78,12 +89,13 @@ impl PlotIter {
 
 #[pyclass(module = "bosing._bosing", frozen, get_all)]
 #[derive(Debug)]
-struct PlotItem {
+pub(super) struct PlotItem {
     channels: Vec<ChannelId>,
     start: Time,
     span: Time,
     depth: usize,
     kind: ItemKind,
+    label: Option<Label>,
 }
 
 fn call_plot(
@@ -92,6 +104,7 @@ fn call_plot(
     blocks: PlotIter,
     channels: Bound<'_, PyList>,
     max_depth: usize,
+    show_label: bool,
 ) -> PyResult<PyObject> {
     static PLOT: GILOnceCell<PyObject> = GILOnceCell::new();
     let plot = PLOT.get_or_try_init(py, || {
@@ -99,7 +112,14 @@ fn call_plot(
             .getattr(BOSING_PLOT_PLOT)
             .map(Into::into)
     })?;
-    plot.call1(py, (ax, blocks, channels, max_depth))
+    let args = PlotArgs {
+        ax,
+        blocks: Py::new(py, blocks)?,
+        channels: channels.unbind(),
+        max_depth,
+        show_label,
+    };
+    plot.call1(py, (args,))
 }
 
 fn arrange_to_plot(root: ElementRef) -> impl Iterator<Item = PlotItem> {
@@ -115,12 +135,14 @@ fn arrange_to_plot(root: ElementRef) -> impl Iterator<Item = PlotItem> {
          }| {
             let kind = ItemKind::from_variant(&item.variant);
             let channels = item.channels().to_vec();
+            let label = item.common.label().cloned();
             PlotItem {
                 channels,
                 start,
                 span,
                 depth,
                 kind,
+                label,
             }
         },
     )
diff --git a/src/quant.rs b/src/quant.rs
index f33213e..a0ead14 100644
--- a/src/quant.rs
+++ b/src/quant.rs
@@ -43,6 +43,7 @@ macro_rules! def_id {
 
 def_id!(ChannelId);
 def_id!(ShapeId);
+def_id!(Label);
 
 type Result<T> = std::result::Result<T, Error>;
 
@@ -326,3 +327,4 @@ macro_rules! impl_id {
 
 impl_id!(ChannelId);
 impl_id!(ShapeId);
+impl_id!(Label);
diff --git a/src/schedule.rs b/src/schedule.rs
index 0193f47..9335ca5 100644
--- a/src/schedule.rs
+++ b/src/schedule.rs
@@ -14,7 +14,7 @@ use mockall::automock;
 
 use crate::{
     python::Alignment,
-    quant::{ChannelId, Time},
+    quant::{ChannelId, Label, Time},
 };
 
 pub(crate) use absolute::{Absolute, AbsoluteEntry};
@@ -40,6 +40,7 @@ pub(crate) struct ElementCommon {
     duration: Option<Time>,
     max_duration: Time,
     min_duration: Time,
+    label: Option<Label>,
 }
 
 #[derive(Debug, Clone)]
@@ -175,6 +176,10 @@ impl ElementCommon {
         self.min_duration
     }
 
+    pub(crate) fn label(&self) -> Option<&Label> {
+        self.label.as_ref()
+    }
+
     fn min_max_duration(&self) -> MinMax {
         let min_max = MinMax::new(self.min_duration, self.max_duration);
         let max = min_max.clamp(self.duration.unwrap_or(Time::INFINITY));
@@ -222,6 +227,11 @@ impl ElementCommonBuilder {
         self
     }
 
+    pub(crate) fn label(&mut self, label: Option<Label>) -> &mut Self {
+        self.0.label = label;
+        self
+    }
+
     pub(crate) fn validate(&self) -> Result<()> {
         let v = &self.0;
         if !(v.margin.0.value().is_finite() && v.margin.1.value().is_finite()) {
@@ -266,6 +276,7 @@ impl Default for ElementCommonBuilder {
             duration: None,
             max_duration: Time::INFINITY,
             min_duration: Default::default(),
+            label: None,
         })
     }
 }
diff --git a/stubtest-allowlist.txt b/stubtest-allowlist.txt
index d1043ef..c73b786 100644
--- a/stubtest-allowlist.txt
+++ b/stubtest-allowlist.txt
@@ -1,5 +1,7 @@
 bosing._plot.Axes
 bosing._plot.Iterator
 bosing._plot.PathPatch.set
+bosing._plot.PlotArgs
+bosing._plot.PlotItem
 bosing._plot.Sequence
 bosing._plot.TypeAlias