From 7a5030d05426e9ed81559eef9d22c70e0184c248 Mon Sep 17 00:00:00 2001 From: Eliot Robson Date: Wed, 11 Nov 2026 14:00:22 -0500 Subject: [PATCH 1/2] Add assert_plot_equal to compare a plot with a reference plot It replaces the legacy grader's Feedback.check_plot. By default it compares the number of subplots, axis scales, and line, scatter and bar data, with options to check labels, titles, legends, limits and styles, require order, or allow extra series. matplotlib is still imported only when a plot is checked. Co-Authored-By: Claude Opus 5.5 --- README.md | 1 + quick_start.md | 64 ++++ src/pytest_prairielearn_grader/assertions.py | 97 ++++++ src/pytest_prairielearn_grader/plot_checks.py | 294 ++++++++++++++++++ .../expected_outcome.json | 19 +- .../test_plain_matplotlib/scenario.py | 30 ++ tests/test_plot_assertions.py | 263 ++++++++++++++++ 7 files changed, 767 insertions(+), 1 deletion(-) create mode 100644 src/pytest_prairielearn_grader/plot_checks.py create mode 100644 tests/test_plot_assertions.py diff --git a/README.md b/README.md index faf1b15..156797c 100644 --- a/README.md +++ b/README.md @@ -78,6 +78,7 @@ Available helpers: - `assert_fn_equal(sandbox, func_name, args=..., expected=...)` — call + compare - `assert_fn_approx_equal(sandbox, func_name, args=..., expected=...)` — call + approx compare - `assert_true(condition)` / `assert_false(condition)` — boolean checks +- `assert_plot_equal(plot, expected)` — compare a matplotlib plot with a reference plot ### Ungradable Submission Detection diff --git a/quick_start.md b/quick_start.md index 4f19665..fb30c0f 100644 --- a/quick_start.md +++ b/quick_start.md @@ -1017,6 +1017,10 @@ Same as `assert_fn_equal` but uses approximate comparison: assert_fn_approx_equal(sandbox, "compute_area", args=(3.0,), expected=28.274, rtol=1e-2) ``` +**`assert_plot_equal(plot, expected, *, name="your plot", rtol=1e-5, atol=1e-8, ..., msg=None)`** + +Compares a student's matplotlib plot with a reference plot. See [Comparing with a Reference Plot](#comparing-with-a-reference-plot). + #### Combining with Partial Credit Friendly mode works seamlessly with `FeedbackFixture`: @@ -1109,6 +1113,66 @@ def test_student_plot(sandbox: StudentFixture) -> None: pt.assert_title_contains(["Data Visualization"]) ``` +#### Comparing with a Reference Plot + +`assert_plot_equal` checks a student's plot against a reference plot that the test draws itself: + +```python +import numpy as np +import pytest +from matplotlib.figure import Figure +from pytest_prairielearn_grader.assertions import assert_plot_equal +from pytest_prairielearn_grader.fixture import StudentFixture +from pytest_prairielearn_grader.utils import GradingOutputLevel + + +@pytest.mark.output(level=GradingOutputLevel.Friendly) +@pytest.mark.grading_data(name="plot", points=1) +def test_plot(sandbox: StudentFixture) -> None: + reference = Figure() + axes = reference.subplots() + x = np.arange(0, 11) + axes.semilogy(x, np.exp(x)) + + assert_plot_equal(sandbox.query("plot"), reference) +``` + +Draw the reference on a `Figure` made directly, as here, rather than with `plt`, so that it doesn't +share pyplot's current figure with anything else in the test. Both plots can be a `Figure` or an +`Axes`, and subplots are compared in order. + +By default, it checks: + +- **Subplots:** the number of axes. +- **Scales:** the x and y axis scales, such as linear or log. +- **Data:** the lines, scatter plots and bar plots, which must have the same number of each kind. Data is compared with `numpy.allclose`, using `rtol` and `atol`. + +The series can be in any order, and a line also matches when its points are in reverse order, since +it looks the same. The points in a scatter plot and the bars in a bar plot can be in any order. + +Keyword arguments change what's checked: + +| Argument | Default | Checks | +| -------------- | ------- | ---------------------------------------------------------------------------------------------------------- | +| `check_scales` | `True` | The x and y axis scales. | +| `check_data` | `True` | The data in lines, scatter plots and bar plots. | +| `ordered` | `False` | Requires lines, scatter plots, bar plots and legend labels to be in the reference's order. | +| `allow_extra` | `False` | Allows more lines, scatter plots or bar plots than the reference, such as a student's extra reference line. | +| `check_labels` | `False` | The x and y axis labels, exactly. | +| `check_titles` | `False` | The title of each subplot and the figure, exactly. | +| `check_legend` | `False` | The legend's labels and title, or that neither plot has a legend. | +| `check_limits` | `False` | The x and y axis limits, within `rtol` and `atol`. | +| `check_styles` | `False` | Line colors, styles and markers, scatter point colors, and bar colors. | + +A failure describes the first difference found, such as `Expected a log scale on the y-axis of your +plot, but found a linear scale.` or `No line in your plot matches the line labeled 'exp'.` Lines +are named by their labels when the reference has them. Use `name` to call the plot something else +in messages, such as `name="plot"` for a variable, or `msg` to show your own message instead. + +Anything `assert_plot_equal` doesn't cover, such as checking that a label contains a word, can be +checked directly on the `Figure`, or with `matplotcheck` as in the earlier example. Both can be +used in the same test. + #### Showing Plots to Students To show the student's plot with a test's results, pass it to `feedback.add_image`: diff --git a/src/pytest_prairielearn_grader/assertions.py b/src/pytest_prairielearn_grader/assertions.py index 5ced73d..6b98aae 100644 --- a/src/pytest_prairielearn_grader/assertions.py +++ b/src/pytest_prairielearn_grader/assertions.py @@ -188,6 +188,103 @@ def assert_false( raise AssertionError(msg or "Expected the condition to be false, but it was true.") +def assert_plot_equal( + plot: Any, + expected: Any, + *, + name: str = "your plot", + rtol: float = 1e-5, + atol: float = 1e-8, + check_scales: bool = True, + check_data: bool = True, + ordered: bool = False, + allow_extra: bool = False, + check_labels: bool = False, + check_titles: bool = False, + check_legend: bool = False, + check_limits: bool = False, + check_styles: bool = False, + msg: str | None = None, +) -> None: + """Assert that a student's matplotlib plot matches a reference plot. + + Both plots can be a ``Figure`` or an ``Axes``. ``sandbox.query`` returns + a student's plot as a ``Figure``; the reference is usually drawn by the + test itself. Subplots are compared in order. + + By default this checks the number of subplots, the axis scales, and the + data in each line, scatter plot and bar plot, which must have the same + number of each kind. Series can be in any order, and a line also matches + when its points are in reverse order. Scatter points and bars can be in + any order. Numbers are compared with ``numpy.allclose``. + + Parameters + ---------- + plot: + The student's plot. Anything other than a ``Figure`` or ``Axes`` + fails the assertion. + expected: + The reference plot. + name: + What to call the student's plot in failure messages. + rtol, atol: + Tolerances for comparing data and axis limits. + check_scales: + Compare the x and y axis scales, such as ``"linear"`` and ``"log"``. + check_data: + Compare the data in lines, scatter plots and bar plots. + ordered: + Require lines, scatter plots, bar plots and legend labels to be in + the same order as in the reference. + allow_extra: + Allow the student's plot to have more lines, scatter plots or bar + plots than the reference. + check_labels: + Compare the x and y axis labels exactly. + check_titles: + Compare the titles of each subplot and the figure exactly. + check_legend: + Compare the legend's labels and title, or that neither plot has a + legend. + check_limits: + Compare the x and y axis limits. + check_styles: + Compare line colors, line styles and markers, scatter point colors, + and bar colors and orientation. + msg: + Optional message to show instead of the description of the + difference. + + Raises + ------ + AssertionError + Describing the first difference found. + TypeError + When *expected* is not a ``Figure`` or ``Axes``. + """ + # matplotlib is an optional dependency, so only import it when a plot is checked. + from .plot_checks import plot_mismatch # noqa: PLC0415 + + mismatch = plot_mismatch( + plot, + expected, + name=name, + rtol=rtol, + atol=atol, + check_scales=check_scales, + check_data=check_data, + ordered=ordered, + allow_extra=allow_extra, + check_labels=check_labels, + check_titles=check_titles, + check_legend=check_legend, + check_limits=check_limits, + check_styles=check_styles, + ) + if mismatch is not None: + raise AssertionError(msg or mismatch) + + # --------------------------------------------------------------------------- # Sandbox-integrated assertions (call student function + compare result) # --------------------------------------------------------------------------- diff --git a/src/pytest_prairielearn_grader/plot_checks.py b/src/pytest_prairielearn_grader/plot_checks.py new file mode 100644 index 0000000..f650466 --- /dev/null +++ b/src/pytest_prairielearn_grader/plot_checks.py @@ -0,0 +1,294 @@ +""" +Compare a student's matplotlib plot with a reference plot, for assertions.assert_plot_equal. + +matplotlib is an optional dependency, so assertions only imports this module when a plot is checked. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +import numpy as np +from matplotlib.axes import Axes +from matplotlib.collections import PathCollection +from matplotlib.colors import to_hex +from matplotlib.container import BarContainer +from matplotlib.figure import Figure +from matplotlib.lines import Line2D + + +@dataclass +class _Series: + kind: str + label: str + data: list[np.ndarray] + # Every series can be compared as given; lines can also be drawn in reverse. + reversible: bool + styles: dict[str, Any] + + @property + def has_label(self) -> bool: + # matplotlib gives unlabeled artists names such as "_child0". + return bool(self.label) and not self.label.startswith("_") + + +@dataclass +class _Options: + name: str + rtol: float + atol: float + check_scales: bool + check_data: bool + ordered: bool + allow_extra: bool + check_labels: bool + check_titles: bool + check_legend: bool + check_limits: bool + check_styles: bool + + +class _Mismatch(Exception): + pass + + +def plot_mismatch(plot: Any, expected: Any, **options: Any) -> str | None: + """Return a message describing how plot differs from expected, or None if they match.""" + opts = _Options(**options) + if not isinstance(expected, Axes | Figure): + raise TypeError(f"assert_plot_equal expects a matplotlib Figure or Axes as the reference plot, not {type(expected).__name__}") + try: + _compare(plot, expected, opts) + except _Mismatch as e: + return str(e) + return None + + +def _axes_of(value: Any) -> list[Any]: + return [value] if isinstance(value, Axes) else list(value.axes) + + +def _plural(count: int, word: str) -> str: + return f"{count} {word}" if count == 1 else f"{count} {word}s" + + +def _compare(plot: Any, expected: Any, opts: _Options) -> None: + if not isinstance(plot, Axes | Figure): + raise _Mismatch(f"Expected {opts.name} to be a Matplotlib plot, such as plt.gca(), but found {type(plot).__name__}.") + + plot_axes = _axes_of(plot) + expected_axes = _axes_of(expected) + if len(plot_axes) != len(expected_axes): + raise _Mismatch(f"Expected {_plural(len(expected_axes), 'set')} of axes in {opts.name}, but found {len(plot_axes)}.") + + if opts.check_titles: + actual_title = _figure_of(plot).get_suptitle() + expected_title = _figure_of(expected).get_suptitle() + if actual_title != expected_title: + raise _Mismatch(f"Expected the figure title {expected_title!r} in {opts.name}, but found {actual_title!r}.") + + for index, (actual, reference) in enumerate(zip(plot_axes, expected_axes, strict=True)): + where = opts.name if len(expected_axes) == 1 else f"subplot {index + 1} of {opts.name}" + _compare_axes(actual, reference, where, opts) + + +def _figure_of(value: Any) -> Any: + return value.figure if isinstance(value, Axes) else value + + +def _compare_axes(actual: Any, reference: Any, where: str, opts: _Options) -> None: + if opts.check_scales: + for axis in ("x", "y"): + actual_scale = getattr(actual, f"get_{axis}scale")() + expected_scale = getattr(reference, f"get_{axis}scale")() + if actual_scale != expected_scale: + raise _Mismatch(f"Expected a {expected_scale} scale on the {axis}-axis of {where}, but found a {actual_scale} scale.") + + if opts.check_data or opts.check_styles: + for kind, extract in (("line", _lines), ("scatter plot", _scatters), ("bar plot", _bars)): + _compare_series(extract(actual), extract(reference), kind, where, opts) + + if opts.check_limits: + for axis in ("x", "y"): + actual_limits = getattr(actual, f"get_{axis}lim")() + expected_limits = getattr(reference, f"get_{axis}lim")() + if not np.allclose(actual_limits, expected_limits, rtol=opts.rtol, atol=opts.atol): + raise _Mismatch( + f"Expected {axis}-axis limits {_format_limits(expected_limits)} in {where}, but found {_format_limits(actual_limits)}." + ) + + if opts.check_labels: + for axis in ("x", "y"): + actual_label = getattr(actual, f"get_{axis}label")() + expected_label = getattr(reference, f"get_{axis}label")() + if actual_label != expected_label: + raise _Mismatch(f"Expected the {axis}-axis label {expected_label!r} in {where}, but found {actual_label!r}.") + + if opts.check_titles: + for location in ("center", "left", "right"): + actual_title = actual.get_title(loc=location) + expected_title = reference.get_title(loc=location) + if actual_title != expected_title: + title = "title" if location == "center" else f"{location} title" + raise _Mismatch(f"Expected the {title} {expected_title!r} in {where}, but found {actual_title!r}.") + + if opts.check_legend: + actual_legend = _legend(actual) + expected_legend = _legend(reference) + if not opts.ordered and actual_legend is not None and expected_legend is not None: + actual_legend = (actual_legend[0], sorted(actual_legend[1])) + expected_legend = (expected_legend[0], sorted(expected_legend[1])) + if actual_legend != expected_legend: + raise _Mismatch(f"Expected {_describe_legend(expected_legend)} in {where}, but found {_describe_legend(actual_legend)}.") + + +def _format_limits(limits: Any) -> str: + low, high = limits + return f"({low:g}, {high:g})" + + +def _legend(axes: Any) -> tuple[str, list[str]] | None: + legend = axes.get_legend() + if legend is None: + return None + return legend.get_title().get_text(), [text.get_text() for text in legend.get_texts()] + + +def _describe_legend(legend: tuple[str, list[str]] | None) -> str: + if legend is None: + return "no legend" + title, labels = legend + description = f"a legend with labels {labels}" + return f"{description} and title {title!r}" if title else description + + +def _floats(values: Any) -> np.ndarray: + return np.asarray(values, dtype=float).ravel() + + +def _hex(color: Any) -> str: + return to_hex(color, keep_alpha=True) + + +def _lines(axes: Any) -> list[_Series]: + return [ + _Series( + kind="line", + label=str(line.get_label()), + data=[_floats(line.get_xdata(orig=False)), _floats(line.get_ydata(orig=False))], + reversible=True, + styles={ + "color": _hex(line.get_color()), + "line style": line.get_linestyle(), + "marker": line.get_marker(), + }, + ) + for line in axes.lines + if type(line) is Line2D + ] + + +def _sorted_points(columns: list[np.ndarray], colors: list[str]) -> tuple[list[np.ndarray], list[str]]: + """Sort points so that the order they were plotted in doesn't matter.""" + order = np.lexsort(columns[::-1]) if len(columns[0]) else np.arange(0) + return [column[order] for column in columns], [colors[i] for i in order] + + +def _scatters(axes: Any) -> list[_Series]: + series = [] + for collection in axes.collections: + if type(collection) is not PathCollection: + continue + offsets = np.asarray(collection.get_offsets(), dtype=float).reshape(-1, 2) + count = len(offsets) + facecolors = collection.get_facecolor() + colors = [_hex(facecolors[i % len(facecolors)]) for i in range(count)] if len(facecolors) else ["none"] * count + data, colors = _sorted_points([offsets[:, 0], offsets[:, 1]], colors) + series.append( + _Series(kind="scatter plot", label=str(collection.get_label()), data=data, reversible=False, styles={"colors": colors}) + ) + return series + + +def _bars(axes: Any) -> list[_Series]: + series = [] + # Typed as Any because the matplotlib stubs don't let get_label be called on a container. + containers: list[Any] = [container for container in axes.containers if isinstance(container, BarContainer)] + for container in containers: + patches = container.patches + columns = [ + _floats([patch.get_x() for patch in patches]), + _floats([patch.get_y() for patch in patches]), + _floats([patch.get_width() for patch in patches]), + _floats([patch.get_height() for patch in patches]), + ] + data, colors = _sorted_points(columns, [_hex(patch.get_facecolor()) for patch in patches]) + series.append( + _Series( + kind="bar plot", + label=str(container.get_label()), + data=data, + reversible=False, + styles={"orientation": container.orientation, "colors": colors}, + ) + ) + return series + + +def _data_matches(actual: _Series, expected: _Series, opts: _Options) -> bool: + def close(columns: list[np.ndarray]) -> bool: + return all( + column.shape == reference.shape and np.allclose(column, reference, rtol=opts.rtol, atol=opts.atol, equal_nan=True) + for column, reference in zip(columns, expected.data, strict=True) + ) + + if actual.styles.get("orientation") != expected.styles.get("orientation"): + return False + if close(actual.data): + return True + return expected.reversible and close([column[::-1] for column in actual.data]) + + +def _compare_series(actual: list[_Series], expected: list[_Series], kind: str, where: str, opts: _Options) -> None: + if opts.allow_extra: + if len(actual) < len(expected): + raise _Mismatch(f"Expected at least {_plural(len(expected), kind)} in {where}, but found {len(actual)}.") + elif len(actual) != len(expected): + raise _Mismatch(f"Expected {_plural(len(expected), kind)} in {where}, but found {len(actual)}.") + + def describe(index: int) -> str: + series = expected[index] + if series.has_label: + return f"the {kind} labeled {series.label!r}" + return f"the expected {kind}" if len(expected) == 1 else f"expected {kind} {index + 1}" + + pairs: list[tuple[_Series, int]] = [] + if opts.ordered: + for index, series in enumerate(expected): + if opts.check_data and not _data_matches(actual[index], series, opts): + raise _Mismatch(f"{kind.capitalize()} {index + 1} in {where} doesn't match {describe(index)}.") + pairs.append((actual[index], index)) + else: + remaining = list(actual) + for index, series in enumerate(expected): + match = next((candidate for candidate in remaining if not opts.check_data or _data_matches(candidate, series, opts)), None) + if match is None: + raise _Mismatch(f"No {kind} in {where} matches {describe(index)}.") + remaining.remove(match) + pairs.append((match, index)) + + if opts.check_styles: + for match, index in pairs: + for style, expected_value in expected[index].styles.items(): + if match.styles[style] != expected_value: + raise _Mismatch( + f"Expected {describe(index)} in {where} to have {style} {_format_style(expected_value)}, " + f"but found {_format_style(match.styles[style])}." + ) + + +def _format_style(value: Any) -> str: + if isinstance(value, list) and value and all(item == value[0] for item in value): + return str(value[0]) + return str(value) diff --git a/tests/scenario_root/test_plain_matplotlib/expected_outcome.json b/tests/scenario_root/test_plain_matplotlib/expected_outcome.json index 5eb5187..30f8980 100644 --- a/tests/scenario_root/test_plain_matplotlib/expected_outcome.json +++ b/tests/scenario_root/test_plain_matplotlib/expected_outcome.json @@ -1,6 +1,6 @@ { "expected_data_object": { - "score": 0.8571428571428571, + "score": 0.7777777777777778, "tests": [ { "max_points": 1, @@ -62,6 +62,23 @@ "Your plot", null ] + }, + { + "max_points": 1, + "name": "assert_plot_equal", + "points": 1.0, + "points_frac": 1.0, + "outcome": "passed", + "test_id": "test_plain_matplotlib.py::test_assert_plot_equal[student_code]" + }, + { + "max_points": 1, + "name": "assert_plot_equal_message", + "points": 0.0, + "points_frac": 0.0, + "outcome": "failed", + "test_id": "test_plain_matplotlib.py::test_assert_plot_equal_message[student_code]", + "message": "Expected a log scale on the y-axis of your plot, but found a linear scale." } ] } diff --git a/tests/scenario_root/test_plain_matplotlib/scenario.py b/tests/scenario_root/test_plain_matplotlib/scenario.py index 275c80b..d42fedd 100644 --- a/tests/scenario_root/test_plain_matplotlib/scenario.py +++ b/tests/scenario_root/test_plain_matplotlib/scenario.py @@ -5,6 +5,7 @@ from matplotcheck.base import PlotTester from matplotlib.figure import Figure +from pytest_prairielearn_grader.assertions import assert_plot_equal from pytest_prairielearn_grader.assertions import assert_true from pytest_prairielearn_grader.fixture import FeedbackFixture from pytest_prairielearn_grader.fixture import StudentFixture @@ -81,3 +82,32 @@ def test_rejected_plot_shows_clear_message(sandbox: StudentFixture) -> None: def test_add_image(sandbox: StudentFixture, feedback: FeedbackFixture) -> None: feedback.add_image(sandbox.query("plot"), label="Your plot") feedback.add_image(sandbox.query_function("make_axes")) + + +def reference_plot() -> Figure: + figure = Figure() + axes = figure.subplots() + axes.semilogy(np.arange(0, 11), np.exp(np.arange(0, 11)), label="exp") + axes.set_xlabel("x") + axes.set_ylabel("f(x)") + axes.set_title("Growth") + axes.legend() + return figure + + +@pytest.mark.grading_data(name="assert_plot_equal", points=1) +def test_assert_plot_equal(sandbox: StudentFixture) -> None: + assert_plot_equal( + sandbox.query("plot"), + reference_plot(), + check_labels=True, + check_titles=True, + check_legend=True, + check_styles=True, + ) + + +@pytest.mark.output(level=GradingOutputLevel.Friendly) +@pytest.mark.grading_data(name="assert_plot_equal_message", points=1) +def test_assert_plot_equal_message(sandbox: StudentFixture) -> None: + assert_plot_equal(sandbox.query_function("make_axes"), reference_plot()) diff --git a/tests/test_plot_assertions.py b/tests/test_plot_assertions.py new file mode 100644 index 0000000..9134d26 --- /dev/null +++ b/tests/test_plot_assertions.py @@ -0,0 +1,263 @@ +import subprocess +import sys +from collections.abc import Callable +from typing import Any + +import numpy as np +import pytest +from matplotlib.figure import Figure + +from pytest_prairielearn_grader.assertions import assert_plot_equal +from pytest_prairielearn_grader.matplotlib_json import plot_from_json +from pytest_prairielearn_grader.matplotlib_json import plot_to_json + +X = np.arange(0, 11) + +Draw = Callable[[Any], object] + + +def make_plot(draw: Draw, subplots: int = 1) -> Figure: + figure = Figure() + axes = figure.subplots(1, subplots, squeeze=False)[0] + for subplot in axes: + draw(subplot) + return figure + + +def semilogy(axes: Any) -> Any: + return axes.semilogy(X, np.exp(X), label="exp") + + +def two_lines(axes: Any) -> None: + axes.plot(X, X, label="linear") + axes.plot(X, X**2, label="square") + + +def full_plot(axes: Any) -> None: + axes.plot(X, X**2, "r--o", label="square") + axes.scatter([3, 1, 2], [6, 2, 4], color="green", label="points") + axes.bar([4, 6, 8], [1, 2, 3], color="purple", label="bars") + axes.set_xlabel("Input") + axes.set_ylabel("Output") + axes.set_title("Results") + axes.set_title("left", loc="left") + axes.set_xlim(-1, 12) + axes.legend(title="Key") + + +@pytest.mark.parametrize( + ("student", "reference", "options"), + [ + pytest.param(semilogy, semilogy, {}, id="same"), + pytest.param(lambda a: a.semilogy(X, np.exp(X) * (1 + 1e-7)), semilogy, {}, id="within-tolerance"), + pytest.param(lambda a: a.semilogy(X, np.exp(X) * 1.01), semilogy, {"rtol": 0.02}, id="custom-tolerance"), + pytest.param(lambda a: a.semilogy(X[::-1], np.exp(X[::-1])), semilogy, {}, id="reversed-line"), + pytest.param(lambda a: (a.plot(X, X**2), a.plot(X, X)), two_lines, {}, id="lines-any-order"), + pytest.param(lambda a: a.scatter([2, 1, 3], [4, 2, 6]), lambda a: a.scatter([1, 2, 3], [2, 4, 6]), {}, id="scatter-any-order"), + pytest.param(lambda a: a.bar([2, 1], [5, 4]), lambda a: a.bar([1, 2], [4, 5]), {}, id="bars-any-order"), + pytest.param(lambda a: (semilogy(a), a.plot(X, X)), semilogy, {"allow_extra": True}, id="allow-extra"), + pytest.param(lambda a: (a.semilogy(X, np.exp(X), "r:"), a.set_xlabel("other")), semilogy, {}, id="labels-and-styles-unchecked"), + pytest.param(lambda a: a.plot(X, X), semilogy, {"check_scales": False, "check_data": False}, id="nothing-checked"), + pytest.param( + full_plot, + full_plot, + {"check_labels": True, "check_titles": True, "check_legend": True, "check_limits": True, "check_styles": True}, + id="everything", + ), + ], +) +def test_matching_plots(student: Draw, reference: Draw, options: dict[str, Any]) -> None: + assert_plot_equal(make_plot(student), make_plot(reference), **options) + + +def test_axes_and_figures_can_be_mixed() -> None: + student = make_plot(semilogy) + reference = make_plot(semilogy) + assert_plot_equal(student.axes[0], reference) + assert_plot_equal(student, reference.axes[0]) + + +def test_plot_sent_from_the_sandbox_matches() -> None: + reference = make_plot(full_plot, subplots=2) + student = plot_from_json(plot_to_json(make_plot(full_plot, subplots=2))["_value"]) + assert_plot_equal( + student, + reference, + ordered=True, + check_labels=True, + check_titles=True, + check_legend=True, + check_limits=True, + check_styles=True, + ) + + +@pytest.mark.parametrize( + ("student", "reference", "options", "message"), + [ + pytest.param( + lambda a: a.plot(X, np.exp(X)), + semilogy, + {}, + "Expected a log scale on the y-axis of your plot, but found a linear scale.", + id="scale", + ), + pytest.param( + lambda a: (semilogy(a), a.semilogy(X, X)), semilogy, {}, "Expected 1 line in your plot, but found 2.", id="line-count" + ), + pytest.param( + lambda a: a.semilogy(X, np.exp(X) * 1.01), semilogy, {}, "No line in your plot matches the line labeled 'exp'.", id="data" + ), + pytest.param( + lambda a: a.semilogy(X, np.exp(X) * 2), + lambda a: a.semilogy(X, np.exp(X)), + {}, + "No line in your plot matches the expected line.", + id="unlabeled", + ), + pytest.param( + lambda a: a.semilogy(X[[1, 0, *range(2, 11)]], np.exp(X[[1, 0, *range(2, 11)]])), + semilogy, + {}, + "No line in your plot matches", + id="scrambled-line", + ), + pytest.param( + lambda a: (a.plot(X, X**2), a.plot(X, X)), + two_lines, + {"ordered": True}, + "Line 1 in your plot doesn't match the line labeled 'linear'.", + id="ordered", + ), + pytest.param( + lambda a: a.set_yscale("log"), + semilogy, + {"allow_extra": True}, + "Expected at least 1 line in your plot, but found 0.", + id="allow-extra-still-needs-reference", + ), + pytest.param( + lambda a: a.scatter([1, 2], [2, 5]), + lambda a: a.scatter([1, 2], [2, 4]), + {}, + "No scatter plot in your plot matches the expected scatter plot.", + id="scatter", + ), + pytest.param( + lambda a: a.barh([1, 2], [4, 5]), lambda a: a.bar([1, 2], [4, 5]), {}, "No bar plot in your plot matches", id="bar-orientation" + ), + pytest.param( + lambda a: (semilogy(a), a.set_xlabel("X")), + lambda a: (semilogy(a), a.set_xlabel("x")), + {"check_labels": True}, + "Expected the x-axis label 'x' in your plot, but found 'X'.", + id="label", + ), + pytest.param( + semilogy, + lambda a: (semilogy(a), a.set_title("Growth")), + {"check_titles": True}, + "Expected the title 'Growth' in your plot, but found ''.", + id="title", + ), + pytest.param( + semilogy, + lambda a: (semilogy(a), a.figure.suptitle("All")), + {"check_titles": True}, + "Expected the figure title 'All' in your plot, but found ''.", + id="suptitle", + ), + pytest.param( + semilogy, + lambda a: (semilogy(a), a.legend()), + {"check_legend": True}, + "Expected a legend with labels ['exp'] in your plot, but found no legend.", + id="legend", + ), + pytest.param( + lambda a: (semilogy(a), a.set_xlim(0, 12)), + lambda a: (semilogy(a), a.set_xlim(0, 10)), + {"check_limits": True}, + "Expected x-axis limits (0, 10) in your plot, but found (0, 12).", + id="limits", + ), + pytest.param( + lambda a: a.semilogy(X, np.exp(X), color="red", label="exp"), + lambda a: a.semilogy(X, np.exp(X), color="blue", label="exp"), + {"check_styles": True}, + "Expected the line labeled 'exp' in your plot to have color #0000ffff, but found #ff0000ff.", + id="line-color", + ), + pytest.param( + lambda a: a.semilogy(X, np.exp(X), "o"), + lambda a: a.semilogy(X, np.exp(X), "x"), + {"check_styles": True, "check_data": False}, + "to have marker x, but found o.", + id="marker", + ), + pytest.param( + lambda a: a.bar([1, 2], [4, 5], color="red"), + lambda a: a.bar([1, 2], [4, 5]), + {"check_styles": True}, + "Expected the expected bar plot in your plot to have colors #1f77b4ff, but found #ff0000ff.", + id="bar-color", + ), + pytest.param( + lambda a: a.plot(X, X), + semilogy, + {"name": "plot"}, + "Expected a log scale on the y-axis of plot, but found a linear scale.", + id="name", + ), + pytest.param(lambda a: a.plot(X, X), semilogy, {"msg": "Check your plot."}, "Check your plot.", id="msg"), + ], +) +def test_mismatched_plots(student: Draw, reference: Draw, options: dict[str, Any], message: str) -> None: + with pytest.raises(AssertionError) as e: + assert_plot_equal(make_plot(student), make_plot(reference), **options) + assert message in str(e.value) + + +def test_subplots_are_named_in_messages() -> None: + student = make_plot(semilogy, subplots=2) + student.axes[1].set_yscale("linear") + with pytest.raises(AssertionError, match="y-axis of subplot 2 of your plot"): + assert_plot_equal(student, make_plot(semilogy, subplots=2)) + + +def test_subplot_count() -> None: + with pytest.raises(AssertionError, match=r"^Expected 1 set of axes in your plot, but found 2\.$"): + assert_plot_equal(make_plot(semilogy, subplots=2), make_plot(semilogy)) + + +@pytest.mark.parametrize("value", [5, None, "plot"]) +def test_student_value_that_is_not_a_plot(value: Any) -> None: + with pytest.raises( + AssertionError, match=f"Expected your plot to be a Matplotlib plot, such as plt.gca\\(\\), but found {type(value).__name__}" + ): + assert_plot_equal(value, make_plot(semilogy)) + + +def test_reference_that_is_not_a_plot() -> None: + with pytest.raises(TypeError, match="reference plot"): + assert_plot_equal(make_plot(semilogy), [1, 2, 3]) + + +def test_grader_works_without_matplotlib_or_numpy() -> None: + # The lite image has neither, so importing the grader mustn't import them. + code = """ +import sys +sys.modules["matplotlib"] = None +sys.modules["numpy"] = None +import pytest_prairielearn_grader.plugin +from pytest_prairielearn_grader.assertions import assert_equal, assert_plot_equal +from pytest_prairielearn_grader.fixture import FeedbackFixture +assert_equal(1, 1) +FeedbackFixture("test").add_image(b"png") +try: + assert_plot_equal(None, None) +except ImportError: + print("ok") +""" + result = subprocess.run([sys.executable, "-c", code], capture_output=True, text=True, check=True) + assert result.stdout == "ok\n" From 5a84e8d75725b1ac467953c7b16ccbb5e7e4e794 Mon Sep 17 00:00:00 2001 From: Eliot Robson Date: Wed, 11 Nov 2026 14:14:37 -0500 Subject: [PATCH 2/2] Make plot_checks classes public, frozen and slotted Also remove a matched series by position instead of with list.remove, which compared series with == and raised for two lines with the same label. Co-Authored-By: Claude Opus 5.5 --- src/pytest_prairielearn_grader/plot_checks.py | 72 ++++++++++--------- tests/test_plot_assertions.py | 6 ++ 2 files changed, 43 insertions(+), 35 deletions(-) diff --git a/src/pytest_prairielearn_grader/plot_checks.py b/src/pytest_prairielearn_grader/plot_checks.py index f650466..17b4962 100644 --- a/src/pytest_prairielearn_grader/plot_checks.py +++ b/src/pytest_prairielearn_grader/plot_checks.py @@ -18,8 +18,8 @@ from matplotlib.lines import Line2D -@dataclass -class _Series: +@dataclass(frozen=True, slots=True) +class Series: kind: str label: str data: list[np.ndarray] @@ -33,8 +33,8 @@ def has_label(self) -> bool: return bool(self.label) and not self.label.startswith("_") -@dataclass -class _Options: +@dataclass(frozen=True, slots=True) +class Options: name: str rtol: float atol: float @@ -49,18 +49,18 @@ class _Options: check_styles: bool -class _Mismatch(Exception): +class Mismatch(Exception): pass def plot_mismatch(plot: Any, expected: Any, **options: Any) -> str | None: """Return a message describing how plot differs from expected, or None if they match.""" - opts = _Options(**options) + opts = Options(**options) if not isinstance(expected, Axes | Figure): raise TypeError(f"assert_plot_equal expects a matplotlib Figure or Axes as the reference plot, not {type(expected).__name__}") try: _compare(plot, expected, opts) - except _Mismatch as e: + except Mismatch as e: return str(e) return None @@ -73,20 +73,20 @@ def _plural(count: int, word: str) -> str: return f"{count} {word}" if count == 1 else f"{count} {word}s" -def _compare(plot: Any, expected: Any, opts: _Options) -> None: +def _compare(plot: Any, expected: Any, opts: Options) -> None: if not isinstance(plot, Axes | Figure): - raise _Mismatch(f"Expected {opts.name} to be a Matplotlib plot, such as plt.gca(), but found {type(plot).__name__}.") + raise Mismatch(f"Expected {opts.name} to be a Matplotlib plot, such as plt.gca(), but found {type(plot).__name__}.") plot_axes = _axes_of(plot) expected_axes = _axes_of(expected) if len(plot_axes) != len(expected_axes): - raise _Mismatch(f"Expected {_plural(len(expected_axes), 'set')} of axes in {opts.name}, but found {len(plot_axes)}.") + raise Mismatch(f"Expected {_plural(len(expected_axes), 'set')} of axes in {opts.name}, but found {len(plot_axes)}.") if opts.check_titles: actual_title = _figure_of(plot).get_suptitle() expected_title = _figure_of(expected).get_suptitle() if actual_title != expected_title: - raise _Mismatch(f"Expected the figure title {expected_title!r} in {opts.name}, but found {actual_title!r}.") + raise Mismatch(f"Expected the figure title {expected_title!r} in {opts.name}, but found {actual_title!r}.") for index, (actual, reference) in enumerate(zip(plot_axes, expected_axes, strict=True)): where = opts.name if len(expected_axes) == 1 else f"subplot {index + 1} of {opts.name}" @@ -97,13 +97,13 @@ def _figure_of(value: Any) -> Any: return value.figure if isinstance(value, Axes) else value -def _compare_axes(actual: Any, reference: Any, where: str, opts: _Options) -> None: +def _compare_axes(actual: Any, reference: Any, where: str, opts: Options) -> None: if opts.check_scales: for axis in ("x", "y"): actual_scale = getattr(actual, f"get_{axis}scale")() expected_scale = getattr(reference, f"get_{axis}scale")() if actual_scale != expected_scale: - raise _Mismatch(f"Expected a {expected_scale} scale on the {axis}-axis of {where}, but found a {actual_scale} scale.") + raise Mismatch(f"Expected a {expected_scale} scale on the {axis}-axis of {where}, but found a {actual_scale} scale.") if opts.check_data or opts.check_styles: for kind, extract in (("line", _lines), ("scatter plot", _scatters), ("bar plot", _bars)): @@ -114,7 +114,7 @@ def _compare_axes(actual: Any, reference: Any, where: str, opts: _Options) -> No actual_limits = getattr(actual, f"get_{axis}lim")() expected_limits = getattr(reference, f"get_{axis}lim")() if not np.allclose(actual_limits, expected_limits, rtol=opts.rtol, atol=opts.atol): - raise _Mismatch( + raise Mismatch( f"Expected {axis}-axis limits {_format_limits(expected_limits)} in {where}, but found {_format_limits(actual_limits)}." ) @@ -123,7 +123,7 @@ def _compare_axes(actual: Any, reference: Any, where: str, opts: _Options) -> No actual_label = getattr(actual, f"get_{axis}label")() expected_label = getattr(reference, f"get_{axis}label")() if actual_label != expected_label: - raise _Mismatch(f"Expected the {axis}-axis label {expected_label!r} in {where}, but found {actual_label!r}.") + raise Mismatch(f"Expected the {axis}-axis label {expected_label!r} in {where}, but found {actual_label!r}.") if opts.check_titles: for location in ("center", "left", "right"): @@ -131,7 +131,7 @@ def _compare_axes(actual: Any, reference: Any, where: str, opts: _Options) -> No expected_title = reference.get_title(loc=location) if actual_title != expected_title: title = "title" if location == "center" else f"{location} title" - raise _Mismatch(f"Expected the {title} {expected_title!r} in {where}, but found {actual_title!r}.") + raise Mismatch(f"Expected the {title} {expected_title!r} in {where}, but found {actual_title!r}.") if opts.check_legend: actual_legend = _legend(actual) @@ -140,7 +140,7 @@ def _compare_axes(actual: Any, reference: Any, where: str, opts: _Options) -> No actual_legend = (actual_legend[0], sorted(actual_legend[1])) expected_legend = (expected_legend[0], sorted(expected_legend[1])) if actual_legend != expected_legend: - raise _Mismatch(f"Expected {_describe_legend(expected_legend)} in {where}, but found {_describe_legend(actual_legend)}.") + raise Mismatch(f"Expected {_describe_legend(expected_legend)} in {where}, but found {_describe_legend(actual_legend)}.") def _format_limits(limits: Any) -> str: @@ -171,9 +171,9 @@ def _hex(color: Any) -> str: return to_hex(color, keep_alpha=True) -def _lines(axes: Any) -> list[_Series]: +def _lines(axes: Any) -> list[Series]: return [ - _Series( + Series( kind="line", label=str(line.get_label()), data=[_floats(line.get_xdata(orig=False)), _floats(line.get_ydata(orig=False))], @@ -195,7 +195,7 @@ def _sorted_points(columns: list[np.ndarray], colors: list[str]) -> tuple[list[n return [column[order] for column in columns], [colors[i] for i in order] -def _scatters(axes: Any) -> list[_Series]: +def _scatters(axes: Any) -> list[Series]: series = [] for collection in axes.collections: if type(collection) is not PathCollection: @@ -206,12 +206,12 @@ def _scatters(axes: Any) -> list[_Series]: colors = [_hex(facecolors[i % len(facecolors)]) for i in range(count)] if len(facecolors) else ["none"] * count data, colors = _sorted_points([offsets[:, 0], offsets[:, 1]], colors) series.append( - _Series(kind="scatter plot", label=str(collection.get_label()), data=data, reversible=False, styles={"colors": colors}) + Series(kind="scatter plot", label=str(collection.get_label()), data=data, reversible=False, styles={"colors": colors}) ) return series -def _bars(axes: Any) -> list[_Series]: +def _bars(axes: Any) -> list[Series]: series = [] # Typed as Any because the matplotlib stubs don't let get_label be called on a container. containers: list[Any] = [container for container in axes.containers if isinstance(container, BarContainer)] @@ -225,7 +225,7 @@ def _bars(axes: Any) -> list[_Series]: ] data, colors = _sorted_points(columns, [_hex(patch.get_facecolor()) for patch in patches]) series.append( - _Series( + Series( kind="bar plot", label=str(container.get_label()), data=data, @@ -236,7 +236,7 @@ def _bars(axes: Any) -> list[_Series]: return series -def _data_matches(actual: _Series, expected: _Series, opts: _Options) -> bool: +def _data_matches(actual: Series, expected: Series, opts: Options) -> bool: def close(columns: list[np.ndarray]) -> bool: return all( column.shape == reference.shape and np.allclose(column, reference, rtol=opts.rtol, atol=opts.atol, equal_nan=True) @@ -250,12 +250,12 @@ def close(columns: list[np.ndarray]) -> bool: return expected.reversible and close([column[::-1] for column in actual.data]) -def _compare_series(actual: list[_Series], expected: list[_Series], kind: str, where: str, opts: _Options) -> None: +def _compare_series(actual: list[Series], expected: list[Series], kind: str, where: str, opts: Options) -> None: if opts.allow_extra: if len(actual) < len(expected): - raise _Mismatch(f"Expected at least {_plural(len(expected), kind)} in {where}, but found {len(actual)}.") + raise Mismatch(f"Expected at least {_plural(len(expected), kind)} in {where}, but found {len(actual)}.") elif len(actual) != len(expected): - raise _Mismatch(f"Expected {_plural(len(expected), kind)} in {where}, but found {len(actual)}.") + raise Mismatch(f"Expected {_plural(len(expected), kind)} in {where}, but found {len(actual)}.") def describe(index: int) -> str: series = expected[index] @@ -263,26 +263,28 @@ def describe(index: int) -> str: return f"the {kind} labeled {series.label!r}" return f"the expected {kind}" if len(expected) == 1 else f"expected {kind} {index + 1}" - pairs: list[tuple[_Series, int]] = [] + pairs: list[tuple[Series, int]] = [] if opts.ordered: for index, series in enumerate(expected): if opts.check_data and not _data_matches(actual[index], series, opts): - raise _Mismatch(f"{kind.capitalize()} {index + 1} in {where} doesn't match {describe(index)}.") + raise Mismatch(f"{kind.capitalize()} {index + 1} in {where} doesn't match {describe(index)}.") pairs.append((actual[index], index)) else: remaining = list(actual) for index, series in enumerate(expected): - match = next((candidate for candidate in remaining if not opts.check_data or _data_matches(candidate, series, opts)), None) - if match is None: - raise _Mismatch(f"No {kind} in {where} matches {describe(index)}.") - remaining.remove(match) - pairs.append((match, index)) + position = next( + (i for i, candidate in enumerate(remaining) if not opts.check_data or _data_matches(candidate, series, opts)), + None, + ) + if position is None: + raise Mismatch(f"No {kind} in {where} matches {describe(index)}.") + pairs.append((remaining.pop(position), index)) if opts.check_styles: for match, index in pairs: for style, expected_value in expected[index].styles.items(): if match.styles[style] != expected_value: - raise _Mismatch( + raise Mismatch( f"Expected {describe(index)} in {where} to have {style} {_format_style(expected_value)}, " f"but found {_format_style(match.styles[style])}." ) diff --git a/tests/test_plot_assertions.py b/tests/test_plot_assertions.py index 9134d26..146a034 100644 --- a/tests/test_plot_assertions.py +++ b/tests/test_plot_assertions.py @@ -53,6 +53,12 @@ def full_plot(axes: Any) -> None: pytest.param(lambda a: a.semilogy(X, np.exp(X) * 1.01), semilogy, {"rtol": 0.02}, id="custom-tolerance"), pytest.param(lambda a: a.semilogy(X[::-1], np.exp(X[::-1])), semilogy, {}, id="reversed-line"), pytest.param(lambda a: (a.plot(X, X**2), a.plot(X, X)), two_lines, {}, id="lines-any-order"), + pytest.param( + lambda a: (a.plot(X, X**2, label="same"), a.plot(X, X, label="same")), + lambda a: (a.plot(X, X, label="same"), a.plot(X, X**2, label="same")), + {}, + id="same-labels-any-order", + ), pytest.param(lambda a: a.scatter([2, 1, 3], [4, 2, 6]), lambda a: a.scatter([1, 2, 3], [2, 4, 6]), {}, id="scatter-any-order"), pytest.param(lambda a: a.bar([2, 1], [5, 4]), lambda a: a.bar([1, 2], [4, 5]), {}, id="bars-any-order"), pytest.param(lambda a: (semilogy(a), a.plot(X, X)), semilogy, {"allow_extra": True}, id="allow-extra"),