From 9cf58b361d17377e5d2a8c3711fea9b185e4712c Mon Sep 17 00:00:00 2001 From: Eliot Robson Date: Thu, 5 Nov 2026 15:27:33 -0500 Subject: [PATCH 1/2] Accept plain matplotlib Axes and Figure values as JSON Student code had to build plots with MatplotlibSerializer. It can now provide a plain Axes or Figure, which the grader receives as a Figure, sent as validated JSON data with no pickling. An Axes brings its whole figure, so subplots aren't lost. Lines, scatter points, bars, histograms, scales, labels, titles, limits and legend labels are carried. Anything else, such as errorbar plots, images or text, is rejected with a clear error rather than losing data. The sandbox now defaults matplotlib to the Agg backend, because student code runs on a worker thread with no display. Importing matplotlib in the new unit tests means the sandbox has to build its own font cache in pytester's temporary home directory, so test_matplotlib_plots needs a longer timeout. The same change is in #27. Co-Authored-By: Claude Opus 5.5 --- .github/copilot-instructions.md | 3 +- quick_start.md | 42 +- .../_student_code_runner.py | 2 + src/pytest_prairielearn_grader/json_utils.py | 13 + .../matplotlib_json.py | 375 ++++++++++++++ .../test_matplotlib_plots/scenario.py | 3 + .../test_plain_matplotlib/__init__.py | 0 .../expected_outcome.json | 56 ++ .../test_plain_matplotlib/scenario.py | 73 +++ .../test_plain_matplotlib/student_code.py | 38 ++ tests/test_matplotlib_json.py | 480 ++++++++++++++++++ 11 files changed, 1070 insertions(+), 15 deletions(-) create mode 100644 src/pytest_prairielearn_grader/matplotlib_json.py create mode 100644 tests/scenario_root/test_plain_matplotlib/__init__.py create mode 100644 tests/scenario_root/test_plain_matplotlib/expected_outcome.json create mode 100644 tests/scenario_root/test_plain_matplotlib/scenario.py create mode 100644 tests/scenario_root/test_plain_matplotlib/student_code.py create mode 100644 tests/test_matplotlib_json.py diff --git a/.github/copilot-instructions.md b/.github/copilot-instructions.md index 9e88357..5ae9fef 100644 --- a/.github/copilot-instructions.md +++ b/.github/copilot-instructions.md @@ -143,7 +143,8 @@ result = sandbox.query_function("student_func", arg1, arg2, timeout=5) ### 5. Serialization -- Uses `dill` for complex Python objects (numpy arrays, pandas DataFrames, matplotlib plots) +- Student values reach the grader as JSON through `to_json` and `from_json`, including numpy arrays, pandas DataFrames and matplotlib plots (see [src/pytest_prairielearn_grader/matplotlib_json.py](src/pytest_prairielearn_grader/matplotlib_json.py)) +- Uses `dill` for grader-to-student function arguments and `query_setup` values - Base64-encoded in JSON for transport - See [src/pytest_prairielearn_grader/json_utils.py](src/pytest_prairielearn_grader/json_utils.py) for custom serialization diff --git a/quick_start.md b/quick_start.md index 594cbee..70be60b 100644 --- a/quick_start.md +++ b/quick_start.md @@ -264,7 +264,7 @@ result = sandbox.query_function("slow_computation", data, query_timeout=2.0) **Note**: Student code must define the queried symbols, and return values must be JSON-serializable. Supported types include: `int`, `float`, `str`, `list`, `dict`, `bool`, `None`, numpy arrays, pandas DataFrames, -and plots built with `MatplotlibSerializer` (see [Testing Matplotlib Plots](#testing-matplotlib-plots)). +and matplotlib plots (see [Testing Matplotlib Plots](#testing-matplotlib-plots)). ## Partial Credit and Feedback @@ -952,25 +952,39 @@ since students can learn from reading full tracebacks. ### Testing Matplotlib Plots -To grade a plot, student code builds it with `MatplotlibSerializer` from the `plot-serializer` package and returns the serializer. The test receives a matplotlib `Figure`. - -`plot-serializer` is included in the `:latest` image but not in `:lite`. Returning a plain matplotlib `Figure` from student code isn't supported. +Student code can use matplotlib as usual and provide the plot's `Axes` or `Figure`. The test receives a matplotlib `Figure`. ```python # Student code -from plot_serializer.matplotlib.serializer import MatplotlibSerializer +import matplotlib.pyplot as plt def create_plot(data): - serializer = MatplotlibSerializer() - fig, ax = serializer.subplots() - ax.plot(data) - ax.set_xlabel("Time") - ax.set_ylabel("Value") - ax.set_title("Data Visualization") - return serializer + plt.plot(data) + plt.xlabel("Time") + plt.ylabel("Value") + plt.title("Data Visualization") + return plt.gca() ``` +An `Axes` value brings its whole figure, so subplots are never lost: use `figure.axes` in your test to get each one, in the order the student created them. + +What's sent to the test: + +- **Lines**, with their data, label, color, width, style, marker and draw style. This includes `plot`, `semilogy`, `loglog` and `step`. +- **Scatter points**, with their positions, sizes and colors. +- **Bars and histograms**, with their positions, sizes, colors and orientation. +- **Axes details:** scales, labels, titles, limits, position and legend labels. Twin axes come back as separate axes in `figure.axes`. +- **Figure details:** size and overall title. + +Other plot content makes the query fail with a message saying what isn't supported, rather than losing data. This includes `errorbar`, `stem`, `fill_between`, `pie`, `imshow`, `contour`, `vlines` and `hlines`, `axhline` and `axvline`, `axvspan`, text and annotations, colormapped scatter plots, polar axes, inset axes and subfigures. + +Some details aren't sent, so don't grade them: grid lines, tick positions and labels, line and marker transparency set separately from color, scatter marker shapes, the base of a log scale, and `symlog` settings. Values are sent as decimal numbers, and dates come back as matplotlib date numbers (use `matplotlib.dates.num2date` to convert them). Categories on a line come back as strings. + +Importing matplotlib can take most of the default 1-second sandbox timeout, so set a longer one in plot questions, such as `sandbox_timeout = 15` at the top of the test file. + +Student code can instead build the plot with `MatplotlibSerializer` from the `plot-serializer` package and return the serializer, which also arrives as a `Figure`. `matplotlib` and `plot-serializer` are included in the `:latest` image but not in `:lite`. + ```python # Test code import matplotlib @@ -985,7 +999,7 @@ from pytest_prairielearn_grader.fixture import StudentFixture @pytest.mark.grading_data(name="Plot Test", points=5) def test_student_plot(sandbox: StudentFixture) -> None: """Test that student creates a correct plot.""" - # Student function returns a MatplotlibSerializer, which arrives here as a Figure + # Student function returns an Axes or Figure, which arrives here as a Figure plot = sandbox.query_function("create_plot", data) assert isinstance(plot, Figure) @@ -1004,7 +1018,7 @@ def test_student_plot(sandbox: StudentFixture) -> None: The grader automatically serializes/deserializes: -- Plots built with `MatplotlibSerializer` (received as a `Figure`) +- Matplotlib `Axes` and `Figure` values, and plots built with `MatplotlibSerializer` (received as a `Figure`) - NumPy arrays - Pandas DataFrames - Standard Python types (int, float, str, list, dict, bool, None) diff --git a/src/pytest_prairielearn_grader/_student_code_runner.py b/src/pytest_prairielearn_grader/_student_code_runner.py index 77b34ea..14e75a8 100644 --- a/src/pytest_prairielearn_grader/_student_code_runner.py +++ b/src/pytest_prairielearn_grader/_student_code_runner.py @@ -622,4 +622,6 @@ async def main(): if __name__ == "__main__": + # Student code runs on a worker thread with no display, so matplotlib must not pick a GUI backend. + os.environ.setdefault("MPLBACKEND", "Agg") asyncio.run(main()) diff --git a/src/pytest_prairielearn_grader/json_utils.py b/src/pytest_prairielearn_grader/json_utils.py index 8f84265..f469169 100644 --- a/src/pytest_prairielearn_grader/json_utils.py +++ b/src/pytest_prairielearn_grader/json_utils.py @@ -50,6 +50,12 @@ def to_json(v: Any) -> Any: pure_json_df = json.loads(encoded_json_str_df) return {"_type": "dataframe_v2", "_value": pure_json_df} + elif module == "matplotlib": + from pytest_prairielearn_grader.matplotlib_json import is_plot + from pytest_prairielearn_grader.matplotlib_json import plot_to_json + + if is_plot(v): + return plot_to_json(v) elif module == "plot_serializer": from plot_serializer.matplotlib.serializer import MatplotlibSerializer @@ -114,6 +120,13 @@ def from_json(v_json) -> Any: # pandas read_json() can process it. value_str = StringIO(json.dumps(v_json["_value"])) return pd.read_json(value_str, orient="table") + elif v_json["_type"] == "matplotlib_figure": + from pytest_prairielearn_grader.matplotlib_json import plot_from_json + + if "_value" in v_json: + return plot_from_json(v_json["_value"]) + else: + raise ValueError("variable of type matplotlib_figure should have value") elif v_json["_type"] == "matplotlib_serializer": from plot_serializer.matplotlib.deserializer import deserialize_from_json diff --git a/src/pytest_prairielearn_grader/matplotlib_json.py b/src/pytest_prairielearn_grader/matplotlib_json.py new file mode 100644 index 0000000..238ffdb --- /dev/null +++ b/src/pytest_prairielearn_grader/matplotlib_json.py @@ -0,0 +1,375 @@ +""" +Convert plain matplotlib figures to JSON data and back. + +Student code runs in the sandbox, so everything in the JSON data is untrusted. Conversion to JSON +uses an allow-list: a plot containing anything not handled here raises UnsupportedPlotError rather +than losing data. Rebuilding only passes validated strings and numbers to matplotlib. + +matplotlib is an optional dependency, so json_utils only imports this module when it's needed. +""" + +import math +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.container import ErrorbarContainer +from matplotlib.container import StemContainer +from matplotlib.figure import Figure +from matplotlib.lines import Line2D +from matplotlib.patches import Rectangle + +MATPLOTLIB_FIGURE_TYPE = "matplotlib_figure" + +_SCALES = {"linear", "log", "symlog", "logit"} +_LINESTYLES = {"-", "--", "-.", ":", "None", " ", ""} +_DRAWSTYLES = {"default", "steps", "steps-pre", "steps-mid", "steps-post"} +_ORIENTATIONS = {"vertical", "horizontal"} +_AXIS_UNITS = {"number", "date", "category"} +_MAX_FIGURE_SIZE_INCHES = 1000.0 + + +class UnsupportedPlotError(ValueError): + """The plot contains something that can't be sent to the grader without losing data.""" + + +def is_plot(value: Any) -> bool: + return isinstance(value, Axes | Figure) + + +def plot_to_json(value: Any) -> dict[str, Any]: + """Return JSON data for a Figure, or for the whole figure that an Axes belongs to.""" + + figure: Any = value.figure if isinstance(value, Axes) else value + if not isinstance(figure, Figure): + raise UnsupportedPlotError("these axes aren't part of a top-level figure") + _check_figure(figure) + width, height = figure.get_size_inches() + return { + "_type": MATPLOTLIB_FIGURE_TYPE, + "_value": { + "size": [float(width), float(height)], + "suptitle": figure.get_suptitle(), # type: ignore[attr-defined] + "axes": [_axes_to_json(axes) for axes in figure.axes], + }, + } + + +def _check_figure(figure: Any) -> None: + suptitle = figure.get_suptitle() + extra_texts = [text for text in figure.texts if text.get_text() != suptitle] + if figure.subfigs: + raise UnsupportedPlotError("subfigures aren't supported") + if extra_texts or figure.legends or figure.images or figure.lines or figure.patches or figure.artists: + raise UnsupportedPlotError("text, legends and other items added directly to a figure aren't supported") + + +def _check_axes(axes: Any) -> None: + if axes.name != "rectilinear": + raise UnsupportedPlotError(f"{axes.name} axes aren't supported") + if axes.child_axes: + raise UnsupportedPlotError("inset axes aren't supported") + if axes.images: + raise UnsupportedPlotError("images aren't supported") + if axes.texts: + raise UnsupportedPlotError("text and annotations aren't supported") + if axes.tables: + raise UnsupportedPlotError("tables aren't supported") + if axes.artists: + raise UnsupportedPlotError("custom artists aren't supported") + + bar_patches: set[int] = set() + for container in axes.containers: + if isinstance(container, ErrorbarContainer): + raise UnsupportedPlotError("errorbar plots aren't supported") + if isinstance(container, StemContainer): + raise UnsupportedPlotError("stem plots aren't supported") + if not isinstance(container, BarContainer): + raise UnsupportedPlotError(f"{type(container).__name__} isn't supported") + if container.errorbar is not None: + raise UnsupportedPlotError("bars with error bars aren't supported") + bar_patches.update(id(patch) for patch in container.patches) + + for patch in axes.patches: + if type(patch) is not Rectangle or id(patch) not in bar_patches: + raise UnsupportedPlotError(f"{type(patch).__name__} shapes other than bars and histograms aren't supported") + if patch.get_data_transform() is not axes.transData: + raise UnsupportedPlotError("bars that don't use data coordinates aren't supported") + + for collection in axes.collections: + if type(collection) is not PathCollection: + raise UnsupportedPlotError(f"{type(collection).__name__} plots aren't supported") + if collection.get_array() is not None: + raise UnsupportedPlotError("scatter plots colored by a colormap aren't supported") + # Compare by identity: matplotlib's == treats some other transforms as equal to transData. + if collection.get_offset_transform() is not axes.transData: + raise UnsupportedPlotError("scatter plots that don't use data coordinates aren't supported") + + for line in axes.lines: + if type(line) is not Line2D: + raise UnsupportedPlotError(f"{type(line).__name__} lines aren't supported") + if line.get_transform() is not axes.transData: + raise UnsupportedPlotError("lines that don't use data coordinates, such as axhline and axvline, aren't supported") + if not isinstance(line.get_marker(), str | None): + raise UnsupportedPlotError("custom line markers aren't supported") + + +def _axis_to_json(axis: Any) -> dict[str, Any]: + scale = axis.get_scale() + if scale not in _SCALES: + raise UnsupportedPlotError(f"the {scale} axis scale isn't supported") + + converter = axis.get_converter() + converter_module = "" if converter is None else type(converter).__module__ + units = "number" + categories: list[str] = [] + if converter_module == "matplotlib.dates": + units = "date" + elif converter_module == "matplotlib.category": + units = "category" + mapping = axis.units._mapping + categories = [str(category) for category in sorted(mapping, key=mapping.__getitem__)] + elif converter is not None: + raise UnsupportedPlotError(f"axis units from {type(converter).__name__} aren't supported") + + return {"scale": scale, "label": axis.get_label_text(), "units": units, "categories": categories} + + +def _floats(values: Any) -> list[float]: + return np.asarray(values, dtype=float).ravel().tolist() + + +def _line_to_json(line: Any) -> dict[str, Any]: + return { + "x": _floats(line.get_xdata(orig=False)), + "y": _floats(line.get_ydata(orig=False)), + "label": str(line.get_label()), + "color": to_hex(line.get_color(), keep_alpha=True), + "linewidth": float(line.get_linewidth()), + "linestyle": line.get_linestyle(), + "marker": line.get_marker(), + "drawstyle": line.get_drawstyle(), + } + + +def _scatter_to_json(collection: Any) -> dict[str, Any]: + offsets = collection.get_offsets() + return { + "x": _floats(offsets[:, 0]), + "y": _floats(offsets[:, 1]), + "sizes": _floats(collection.get_sizes()), + "colors": [to_hex(color, keep_alpha=True) for color in collection.get_facecolor()], + "label": str(collection.get_label()), + } + + +def _bars_to_json(container: Any) -> dict[str, Any]: + return { + "x": _floats([patch.get_x() for patch in container.patches]), + "y": _floats([patch.get_y() for patch in container.patches]), + "width": _floats([patch.get_width() for patch in container.patches]), + "height": _floats([patch.get_height() for patch in container.patches]), + "colors": [to_hex(patch.get_facecolor(), keep_alpha=True) for patch in container.patches], + "orientation": container.orientation, + "label": str(container.get_label()), + } + + +def _axes_to_json(axes: Any) -> dict[str, Any]: + _check_axes(axes) + legend = axes.get_legend() + return { + "position": [float(value) for value in axes.get_position().bounds], + "x_axis": _axis_to_json(axes.xaxis), + "y_axis": _axis_to_json(axes.yaxis), + "title": {location: axes.get_title(loc=location) for location in ("left", "center", "right")}, + "x_limits": [float(value) for value in axes.get_xlim()], + "y_limits": [float(value) for value in axes.get_ylim()], + "lines": [_line_to_json(line) for line in axes.lines], + "scatters": [_scatter_to_json(collection) for collection in axes.collections], + "bars": [_bars_to_json(container) for container in axes.containers], + "legend": None + if legend is None + else {"title": legend.get_title().get_text(), "labels": [text.get_text() for text in legend.get_texts()]}, + } + + +def _as_dict(value: Any, name: str) -> dict[str, Any]: + if type(value) is not dict: + raise ValueError(f"{name} should be a dictionary") + return value + + +def _as_list(value: Any, name: str) -> list[Any]: + if type(value) is not list: + raise ValueError(f"{name} should be a list") + return value + + +def _as_str(value: Any, name: str) -> str: + if type(value) is not str: + raise ValueError(f"{name} should be a string") + return value + + +def _as_number(value: Any, name: str) -> float: + if type(value) not in (int, float): + raise ValueError(f"{name} should be a number") + return float(value) + + +def _as_finite_number(value: Any, name: str) -> float: + number = _as_number(value, name) + if not math.isfinite(number): + raise ValueError(f"{name} should be finite") + return number + + +def _as_choice(value: Any, choices: set[str], name: str) -> str: + if type(value) is not str or value not in choices: + raise ValueError(f"{name} has an unsupported value") + return value + + +def _as_numbers(value: Any, name: str) -> list[float]: + return [_as_number(item, name) for item in _as_list(value, name)] + + +def _as_strs(value: Any, name: str) -> list[str]: + return [_as_str(item, name) for item in _as_list(value, name)] + + +def _as_finite_numbers(value: Any, length: int, name: str) -> list[float]: + numbers = [_as_finite_number(item, name) for item in _as_list(value, name)] + if len(numbers) != length: + raise ValueError(f"{name} should have {length} values") + return numbers + + +def _same_length(name: str, *columns: list[Any]) -> None: + if len({len(column) for column in columns}) > 1: + raise ValueError(f"the values in {name} should have the same length") + + +def _category_values(values: list[float], categories: list[str], name: str) -> list[str]: + result = [] + for value in values: + if not value.is_integer() or not 0 <= value < len(categories): + raise ValueError(f"{name} has a value that isn't one of its categories") + result.append(categories[int(value)]) + return result + + +def _apply_axis_units(axis: Any, axis_data: dict[str, Any], categories: list[str]) -> None: + if axis_data["units"] == "date": + axis.axis_date() + elif axis_data["units"] == "category" and categories: + axis.update_units(categories) + + +def _add_bars(axes: Any, data: Any) -> None: + bars = _as_dict(data, "bars") + x = _as_numbers(bars["x"], "bar x") + y = _as_numbers(bars["y"], "bar y") + width = _as_numbers(bars["width"], "bar width") + height = _as_numbers(bars["height"], "bar height") + colors = _as_strs(bars["colors"], "bar colors") + _same_length("bars", x, y, width, height, colors) + label = _as_str(bars["label"], "bar label") + if _as_choice(bars["orientation"], _ORIENTATIONS, "bar orientation") == "vertical": + axes.bar(x, height, width, bottom=y, align="edge", color=colors, label=label) + else: + axes.barh(y, width, height, left=x, align="edge", color=colors, label=label) + + +def _add_scatter(axes: Any, data: Any) -> None: + scatter = _as_dict(data, "scatter") + x = _as_numbers(scatter["x"], "scatter x") + y = _as_numbers(scatter["y"], "scatter y") + _same_length("scatter", x, y) + sizes = _as_numbers(scatter["sizes"], "scatter sizes") + colors = _as_strs(scatter["colors"], "scatter colors") + axes.scatter(x, y, s=sizes, c=colors if colors else "none", label=_as_str(scatter["label"], "scatter label")) + + +def _add_line(axes: Any, data: Any, x_categories: list[str], y_categories: list[str]) -> None: + line = _as_dict(data, "line") + x: list[Any] = _as_numbers(line["x"], "line x") + y: list[Any] = _as_numbers(line["y"], "line y") + _same_length("line", x, y) + if x_categories: + x = _category_values(x, x_categories, "line x") + if y_categories: + y = _category_values(y, y_categories, "line y") + marker = line["marker"] + if marker is not None: + marker = _as_str(marker, "line marker") + axes.plot( + x, + y, + label=_as_str(line["label"], "line label"), + color=_as_str(line["color"], "line color"), + linewidth=_as_finite_number(line["linewidth"], "line width"), + linestyle=_as_choice(line["linestyle"], _LINESTYLES, "line style"), + marker=marker, + drawstyle=_as_choice(line["drawstyle"], _DRAWSTYLES, "line drawstyle"), + ) + + +def _add_axes(figure: Any, data: Any) -> None: + axes_data = _as_dict(data, "axes") + axes = figure.add_axes(_as_finite_numbers(axes_data["position"], 4, "axes position")) + x_axis = _as_dict(axes_data["x_axis"], "x axis") + y_axis = _as_dict(axes_data["y_axis"], "y axis") + for axis_data in (x_axis, y_axis): + _as_choice(axis_data["units"], _AXIS_UNITS, "axis units") + _as_str(axis_data["label"], "axis label") + x_categories = _as_strs(x_axis["categories"], "x categories") + y_categories = _as_strs(y_axis["categories"], "y categories") + + axes.set_xscale(_as_choice(x_axis["scale"], _SCALES, "x scale")) + axes.set_yscale(_as_choice(y_axis["scale"], _SCALES, "y scale")) + + # Bars and scatter offsets are numbers even on a category axis, so add them before setting units. + # Lines keep their category names, as they do in the student's plot, so add them after. + for bars in _as_list(axes_data["bars"], "bars"): + _add_bars(axes, bars) + for scatter in _as_list(axes_data["scatters"], "scatters"): + _add_scatter(axes, scatter) + _apply_axis_units(axes.xaxis, x_axis, x_categories) + _apply_axis_units(axes.yaxis, y_axis, y_categories) + for line in _as_list(axes_data["lines"], "lines"): + _add_line(axes, line, x_categories, y_categories) + + axes.set_xlabel(x_axis["label"]) + axes.set_ylabel(y_axis["label"]) + titles = _as_dict(axes_data["title"], "title") + for location in ("left", "center", "right"): + axes.set_title(_as_str(titles[location], "title"), loc=location) + axes.set_xlim(_as_finite_numbers(axes_data["x_limits"], 2, "x limits")) + axes.set_ylim(_as_finite_numbers(axes_data["y_limits"], 2, "y limits")) + + legend = axes_data["legend"] + if legend is not None: + legend = _as_dict(legend, "legend") + axes.legend(_as_strs(legend["labels"], "legend labels"), title=_as_str(legend["title"], "legend title")) + + +def plot_from_json(data: Any) -> Any: + """Rebuild a Figure from JSON data, raising ValueError if the data is malformed.""" + + try: + figure_data = _as_dict(data, "figure") + size = _as_finite_numbers(figure_data["size"], 2, "figure size") + if not all(0 < value <= _MAX_FIGURE_SIZE_INCHES for value in size): + raise ValueError("figure size is out of range") + figure: Any = Figure(figsize=(size[0], size[1])) + figure.suptitle(_as_str(figure_data["suptitle"], "figure title")) + for axes_data in _as_list(figure_data["axes"], "figure axes"): + _add_axes(figure, axes_data) + except (KeyError, TypeError, ValueError, OverflowError) as e: + raise ValueError(f"variable of type {MATPLOTLIB_FIGURE_TYPE} has invalid data: {e!r}") from e + return figure diff --git a/tests/scenario_root/test_matplotlib_plots/scenario.py b/tests/scenario_root/test_matplotlib_plots/scenario.py index dcaa940..b66658c 100644 --- a/tests/scenario_root/test_matplotlib_plots/scenario.py +++ b/tests/scenario_root/test_matplotlib_plots/scenario.py @@ -8,6 +8,9 @@ from pytest_prairielearn_grader.fixture import StudentFixture +# In a fresh environment, the first sandbox's imports can take longer than the default timeout. +sandbox_timeout = 30.0 + @pytest.mark.grading_data(name="line_plot_basic", points=2) def test_line_plot_basic(sandbox: StudentFixture) -> None: diff --git a/tests/scenario_root/test_plain_matplotlib/__init__.py b/tests/scenario_root/test_plain_matplotlib/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/scenario_root/test_plain_matplotlib/expected_outcome.json b/tests/scenario_root/test_plain_matplotlib/expected_outcome.json new file mode 100644 index 0000000..b60f77b --- /dev/null +++ b/tests/scenario_root/test_plain_matplotlib/expected_outcome.json @@ -0,0 +1,56 @@ +{ + "expected_data_object": { + "score": 0.8333333333333334, + "tests": [ + { + "max_points": 1, + "name": "query_axes", + "points": 1.0, + "points_frac": 1.0, + "outcome": "passed", + "test_id": "test_plain_matplotlib.py::test_query_axes[student_code]" + }, + { + "max_points": 1, + "name": "query_figure", + "points": 1.0, + "points_frac": 1.0, + "outcome": "passed", + "test_id": "test_plain_matplotlib.py::test_query_figure[student_code]" + }, + { + "max_points": 1, + "name": "query_function", + "points": 1.0, + "points_frac": 1.0, + "outcome": "passed", + "test_id": "test_plain_matplotlib.py::test_function_returning_axes_works_with_plot_tester[student_code]" + }, + { + "max_points": 1, + "name": "subplots", + "points": 1.0, + "points_frac": 1.0, + "outcome": "passed", + "test_id": "test_plain_matplotlib.py::test_function_returning_subplots[student_code]" + }, + { + "max_points": 1, + "name": "rejected_raw", + "points": 1.0, + "points_frac": 1.0, + "outcome": "passed", + "test_id": "test_plain_matplotlib.py::test_rejected_plot_keeps_sandbox_usable[student_code]" + }, + { + "max_points": 1, + "name": "rejected_message", + "points": 0.0, + "points_frac": 0.0, + "outcome": "failed", + "test_id": "test_plain_matplotlib.py::test_rejected_plot_shows_clear_message[student_code]", + "message": "errorbar plots aren't supported" + } + ] + } +} diff --git a/tests/scenario_root/test_plain_matplotlib/scenario.py b/tests/scenario_root/test_plain_matplotlib/scenario.py new file mode 100644 index 0000000..0c2e9b5 --- /dev/null +++ b/tests/scenario_root/test_plain_matplotlib/scenario.py @@ -0,0 +1,73 @@ +import numpy as np +import pytest +from matplotcheck.base import PlotTester +from matplotlib.figure import Figure + +from pytest_prairielearn_grader.assertions import assert_true +from pytest_prairielearn_grader.fixture import StudentFixture +from pytest_prairielearn_grader.utils import GradingOutputLevel +from pytest_prairielearn_grader.utils import QueryStatusCode + +# Importing matplotlib in student code can take longer than the default timeout. +sandbox_timeout = 30.0 + +EXPECTED_ERROR = ( + "Variable 'errorbar_plot' has a value of type Axes, which can't be sent back to the grader: errorbar plots aren't supported" +) + + +def check_semilogy(figure: Figure) -> None: + assert isinstance(figure, Figure) + assert len(figure.axes) == 1 + axes = figure.axes[0] + assert axes.get_xscale() == "linear" + assert axes.get_yscale() == "log" + x, y = axes.get_lines()[0].get_data() + np.testing.assert_allclose(np.asarray(x, dtype=float), np.arange(0, 11)) + np.testing.assert_allclose(np.asarray(y, dtype=float), np.exp(np.arange(0, 11))) + + +@pytest.mark.grading_data(name="query_axes", points=1) +def test_query_axes(sandbox: StudentFixture) -> None: + check_semilogy(sandbox.query("plot")) + + +@pytest.mark.grading_data(name="query_figure", points=1) +def test_query_figure(sandbox: StudentFixture) -> None: + check_semilogy(sandbox.query("figure")) + + +@pytest.mark.grading_data(name="query_function", points=1) +def test_function_returning_axes_works_with_plot_tester(sandbox: StudentFixture) -> None: + tester = PlotTester(sandbox.query_function("make_axes").axes[0]) + + tester.assert_plot_type("line") + tester.assert_axis_label_contains(axis="x", strings_expected=["Input"]) + tester.assert_axis_label_contains(axis="y", strings_expected=["Output"]) + tester.assert_title_contains(["Doubling"]) + tester.assert_legend_labels(["doubled"]) + + +@pytest.mark.grading_data(name="subplots", points=1) +def test_function_returning_subplots(sandbox: StudentFixture) -> None: + figure = sandbox.query_function("make_subplots") + + assert len(figure.axes) == 2 + np.testing.assert_allclose(figure.axes[0].collections[0].get_offsets(), [[1, 3], [2, 4]]) + assert [patch.get_height() for patch in figure.axes[1].patches] == [5, 6] + + +@pytest.mark.grading_data(name="rejected_raw", points=1) +def test_rejected_plot_keeps_sandbox_usable(sandbox: StudentFixture) -> None: + response = sandbox.query_raw("errorbar_plot") + + assert response["status"] == QueryStatusCode.NOT_SERIALIZABLE + assert response["error"] is not None + assert EXPECTED_ERROR in response["error"] + check_semilogy(sandbox.query("plot")) + + +@pytest.mark.output(level=GradingOutputLevel.Friendly) +@pytest.mark.grading_data(name="rejected_message", points=1) +def test_rejected_plot_shows_clear_message(sandbox: StudentFixture) -> None: + assert_true(isinstance(sandbox.query("errorbar_plot"), Figure), msg="plot is not a Figure") diff --git a/tests/scenario_root/test_plain_matplotlib/student_code.py b/tests/scenario_root/test_plain_matplotlib/student_code.py new file mode 100644 index 0000000..7948ab1 --- /dev/null +++ b/tests/scenario_root/test_plain_matplotlib/student_code.py @@ -0,0 +1,38 @@ +import matplotlib.pyplot as plt +import numpy as np + +x = np.arange(0, 11) +fx = np.exp(x) +plt.semilogy(x, fx, label="exp") +plt.xlabel("x") +plt.ylabel("f(x)") +plt.title("Growth") +plt.legend() +plot = plt.gca() +figure = plt.gcf() + + +def make_axes(): + _, axes = plt.subplots() + axes.plot([1, 2, 3], [2, 4, 6], label="doubled") + axes.set_xlabel("Input") + axes.set_ylabel("Output") + axes.set_title("Doubling") + axes.legend() + return axes + + +def make_subplots(): + figure, (left, right) = plt.subplots(1, 2) + left.scatter([1, 2], [3, 4]) + right.bar(["a", "b"], np.array([5, 6])) + return figure + + +def make_errorbar(): + _, axes = plt.subplots() + axes.errorbar([1, 2], [3, 4], yerr=[0.1, 0.2]) + return axes + + +errorbar_plot = make_errorbar() diff --git a/tests/test_matplotlib_json.py b/tests/test_matplotlib_json.py new file mode 100644 index 0000000..3b507d0 --- /dev/null +++ b/tests/test_matplotlib_json.py @@ -0,0 +1,480 @@ +import json +from collections.abc import Callable +from datetime import UTC +from datetime import datetime +from typing import Any + +import matplotlib + +matplotlib.use("Agg") +import matplotlib.pyplot as plt +import numpy as np +import pytest +from matplotcheck.base import PlotTester +from matplotlib.colors import to_hex +from matplotlib.figure import Figure + +from pytest_prairielearn_grader.json_utils import from_json +from pytest_prairielearn_grader.json_utils import to_json +from pytest_prairielearn_grader.matplotlib_json import UnsupportedPlotError +from pytest_prairielearn_grader.matplotlib_json import plot_from_json + + +@pytest.fixture(autouse=True) +def close_figures() -> Any: + # Other tests in this process can leave pyplot figures open. + plt.close("all") + yield + plt.close("all") + + +def round_trip(value: Any) -> Any: + """Send a value through real JSON text, as the sandbox does.""" + return from_json(json.loads(json.dumps(to_json(value)))) + + +def test_semilogy_axes() -> None: + x = np.arange(0, 11) + plt.semilogy(x, np.exp(x), label="f", color="red", linewidth=3, linestyle="--", marker="o") + plt.xlabel("x") + plt.ylabel("f(x)") + plt.title("Growth") + + figure = round_trip(plt.gca()) + + assert isinstance(figure, Figure) + assert len(figure.axes) == 1 + axes = figure.axes[0] + assert axes.get_xscale() == "linear" + assert axes.get_yscale() == "log" + assert (axes.get_xlabel(), axes.get_ylabel(), axes.get_title()) == ("x", "f(x)", "Growth") + assert axes.get_xlim() == plt.gca().get_xlim() + assert axes.get_ylim() == plt.gca().get_ylim() + line = axes.get_lines()[0] + np.testing.assert_allclose(line.get_xdata(), x) + np.testing.assert_allclose(line.get_ydata(), np.exp(x)) + assert line.get_label() == "f" + assert to_hex(line.get_color()) == to_hex("red") + assert line.get_linewidth() == 3 + assert line.get_linestyle() == "--" + assert line.get_marker() == "o" + + +def test_figure_value_and_all_subplots_kept() -> None: + figure, (left, right) = plt.subplots(1, 2) + left.plot([1, 2], [3, 4]) + right.plot([5, 6], [7, 8]) + figure.suptitle("Both") + + from_figure = round_trip(figure) + from_axes = round_trip(right) + + for rebuilt in (from_figure, from_axes): + assert isinstance(rebuilt, Figure) + assert rebuilt.get_suptitle() == "Both" + assert len(rebuilt.axes) == 2 + np.testing.assert_allclose(rebuilt.axes[0].get_lines()[0].get_ydata(), [3, 4]) + np.testing.assert_allclose(rebuilt.axes[1].get_lines()[0].get_ydata(), [7, 8]) + assert rebuilt.axes[1].get_position().bounds == right.get_position().bounds + + +def test_twin_axes_kept_as_separate_axes() -> None: + _, axes = plt.subplots() + axes.plot([1, 2], [3, 4]) + twin = axes.twinx() + twin.plot([1, 2], [50, 60]) + + figure = round_trip(axes) + + assert len(figure.axes) == 2 + np.testing.assert_allclose(figure.axes[1].get_lines()[0].get_ydata(), [50, 60]) + + +def test_scatter_values_sizes_and_colors() -> None: + _, axes = plt.subplots() + axes.scatter([1, 2, 3], [4, 5, 6], s=[10, 20, 30], c=["red", "green", "blue"], label="points") + + collection = round_trip(axes).axes[0].collections[0] + + np.testing.assert_allclose(collection.get_offsets(), [[1, 4], [2, 5], [3, 6]]) + np.testing.assert_allclose(collection.get_sizes(), [10, 20, 30]) + assert [to_hex(color) for color in collection.get_facecolor()] == [to_hex(color) for color in ("red", "green", "blue")] + assert collection.get_label() == "points" + + +def test_bar_with_numpy_values() -> None: + _, axes = plt.subplots() + axes.bar(np.array([0, 1, 2]), np.array([5, 6, 7], dtype=np.int64), bottom=np.array([1, 1, 1]), color="green", label="bars") + + rebuilt = round_trip(axes).axes[0] + + assert [patch.get_height() for patch in rebuilt.patches] == [5, 6, 7] + assert [patch.get_y() for patch in rebuilt.patches] == [1, 1, 1] + assert [patch.get_x() for patch in rebuilt.patches] == pytest.approx([-0.4, 0.6, 1.6]) + assert to_hex(rebuilt.patches[0].get_facecolor()) == to_hex("green") + assert rebuilt.containers[0].get_label() == "bars" + + +def test_horizontal_bars() -> None: + _, axes = plt.subplots() + axes.barh([0, 1], [5, 6]) + + rebuilt = round_trip(axes).axes[0] + + assert rebuilt.containers[0].orientation == "horizontal" + assert [patch.get_width() for patch in rebuilt.patches] == [5, 6] + + +def test_histogram() -> None: + _, axes = plt.subplots() + counts, edges, _ = axes.hist([1, 1, 2, 3, 3, 3], bins=3) + + rebuilt = round_trip(axes).axes[0] + + np.testing.assert_allclose([patch.get_height() for patch in rebuilt.patches], counts) + np.testing.assert_allclose([patch.get_x() for patch in rebuilt.patches], edges[:-1]) + + +def test_categorical_line_and_bar() -> None: + _, axes = plt.subplots() + axes.bar(["a", "b"], [1, 2]) + axes.plot(["a", "b"], [3, 4]) + + rebuilt = round_trip(axes).axes[0] + + assert list(rebuilt.get_lines()[0].get_xdata()) == ["a", "b"] + assert [patch.get_height() for patch in rebuilt.patches] == [1, 2] + assert [patch.get_x() for patch in rebuilt.patches] == pytest.approx([-0.4, 0.6]) + rebuilt.figure.canvas.draw() + assert [label.get_text() for label in rebuilt.get_xticklabels()] == ["a", "b"] + + +def test_dates_come_back_as_date_numbers() -> None: + _, axes = plt.subplots() + axes.plot([datetime(2024, 1, 1, tzinfo=UTC), datetime(2024, 1, 2, tzinfo=UTC)], [1, 2]) + expected = axes.get_lines()[0].get_xdata(orig=False) + + rebuilt = round_trip(axes).axes[0] + + np.testing.assert_allclose(rebuilt.get_lines()[0].get_xdata(), expected) + assert type(rebuilt.xaxis.get_converter()).__module__ == "matplotlib.dates" + + +def test_nan_infinity_and_masked_values() -> None: + _, axes = plt.subplots() + axes.plot([1, 2, 3], [1.0, float("nan"), float("inf")]) + axes.plot([1, 2], np.ma.masked_array([1.0, 2.0], mask=[False, True])) + axes.set_ylim(0, 5) + + lines = round_trip(axes).axes[0].get_lines() + + np.testing.assert_array_equal(lines[0].get_ydata(), [1.0, float("nan"), float("inf")]) + np.testing.assert_array_equal(lines[1].get_ydata(), [1.0, float("nan")]) + + +def test_empty_axes() -> None: + _, axes = plt.subplots() + + rebuilt = round_trip(axes) + + assert len(rebuilt.axes) == 1 + assert not rebuilt.axes[0].get_lines() + + +def test_legend_labels_kept() -> None: + _, axes = plt.subplots() + axes.plot([1, 2], [3, 4], label="data") + axes.legend(title="Key") + + legend = round_trip(axes).axes[0].get_legend() + + assert [text.get_text() for text in legend.get_texts()] == ["data"] + assert legend.get_title().get_text() == "Key" + + +def test_rebuilt_figure_can_be_drawn() -> None: + _, axes = plt.subplots() + axes.plot([1, 2], [3, 4], label="data") + axes.scatter([1], [2]) + axes.bar([1], [2]) + axes.legend() + + round_trip(axes).canvas.draw() + + +def test_rebuild_does_not_touch_pyplot_figures() -> None: + plt.plot([1, 2], [3, 4]) + before = plt.get_fignums() + + round_trip(plt.gca()) + + assert plt.get_fignums() == before + + +def test_plot_tester_works_on_rebuilt_figure() -> None: + plt.plot([1, 2, 3], [2, 4, 6], label="doubled") + plt.xlabel("Input") + plt.ylabel("Output") + plt.title("Doubling") + plt.legend() + + tester = PlotTester(round_trip(plt.gca()).axes[0]) + + tester.assert_plot_type("line") + tester.assert_axis_label_contains(axis="x", strings_expected=["Input"]) + tester.assert_axis_label_contains(axis="y", strings_expected=["Output"]) + tester.assert_title_contains(["Doubling"]) + tester.assert_legend_labels(["doubled"]) + + +def _errorbar(axes: Any) -> None: + axes.errorbar([1, 2], [3, 4], yerr=[0.1, 0.2]) + + +def _bar_with_errors(axes: Any) -> None: + axes.bar([1, 2], [3, 4], yerr=[0.1, 0.2]) + + +def _stem(axes: Any) -> None: + axes.stem([1, 2], [3, 4]) + + +def _image(axes: Any) -> None: + axes.imshow([[1, 2], [3, 4]]) + + +def _text(axes: Any) -> None: + axes.plot([1, 2], [3, 4]) + axes.text(1, 3, "note") + + +def _annotation(axes: Any) -> None: + axes.annotate("note", (1, 3)) + + +def _axhline(axes: Any) -> None: + axes.plot([1, 2], [3, 4]) + axes.axhline(3) + + +def _axline(axes: Any) -> None: + axes.axline((0, 0), slope=1) + + +def _axes_coordinates(axes: Any) -> None: + axes.plot([0, 1], [0, 1], transform=axes.transAxes) + + +def _blended_coordinates(axes: Any) -> None: + axes.plot([0, 1], [0, 1], transform=axes.get_xaxis_transform()) + + +def _bar_axes_coordinates(axes: Any) -> None: + axes.bar([0, 1], [0, 1], transform=axes.transAxes) + + +def _scatter_axes_coordinates(axes: Any) -> None: + axes.scatter([0, 1], [0, 1], transform=axes.transAxes) + + +def _fill_between(axes: Any) -> None: + axes.fill_between([1, 2], [3, 4]) + + +def _pie(axes: Any) -> None: + axes.pie([1, 2]) + + +def _colormapped_scatter(axes: Any) -> None: + axes.scatter([1, 2], [3, 4], c=[1.0, 2.0]) + + +def _vlines(axes: Any) -> None: + axes.vlines([1, 2], 0, 1) + + +def _contour(axes: Any) -> None: + axes.contourf([[1, 2], [3, 4]]) + + +def _step_histogram(axes: Any) -> None: + axes.hist([1, 2, 3], histtype="step") + + +def _rectangle_patch(axes: Any) -> None: + axes.axvspan(1, 2) + + +def _inset(axes: Any) -> None: + axes.inset_axes([0.5, 0.5, 0.2, 0.2]) + + +def _function_scale(axes: Any) -> None: + axes.set_yscale("function", functions=(lambda v: v, lambda v: v)) + + +@pytest.mark.parametrize( + ("draw", "message"), + [ + (_errorbar, "errorbar"), + (_bar_with_errors, "errorbar"), + (_stem, "stem"), + (_image, "images"), + (_text, "text"), + (_annotation, "text"), + (_axhline, "axhline"), + (_axline, "AxLine"), + (_axes_coordinates, "data coordinates"), + (_blended_coordinates, "data coordinates"), + (_bar_axes_coordinates, "data coordinates"), + (_scatter_axes_coordinates, "data coordinates"), + (_fill_between, "plots"), + (_pie, "aren't supported"), + (_colormapped_scatter, "colormap"), + (_vlines, "LineCollection"), + (_contour, "aren't supported"), + (_step_histogram, "Polygon"), + (_rectangle_patch, "Rectangle"), + (_inset, "inset"), + (_function_scale, "function"), + ], +) +def test_unsupported_plots_are_rejected(draw: Callable[[Any], None], message: str) -> None: + _, axes = plt.subplots() + draw(axes) + + with pytest.raises(UnsupportedPlotError, match=message): + to_json(axes) + with pytest.raises(UnsupportedPlotError, match=message): + to_json(axes.figure) + + +def test_polar_axes_rejected() -> None: + _, axes = plt.subplots(subplot_kw={"projection": "polar"}) + + with pytest.raises(UnsupportedPlotError, match="polar"): + to_json(axes) + + +def test_rejection_in_any_subplot_rejects_whole_figure() -> None: + _, (left, right) = plt.subplots(1, 2) + left.plot([1, 2], [3, 4]) + _errorbar(right) + + with pytest.raises(UnsupportedPlotError, match="errorbar"): + to_json(left) + + +def test_figure_level_items_rejected() -> None: + figure = plt.figure() + figure.text(0.5, 0.5, "note") + + with pytest.raises(UnsupportedPlotError, match="directly to a figure"): + to_json(figure) + + +def test_subfigures_rejected() -> None: + figure = plt.figure() + sub = figure.subfigures(1, 2) + axes = sub[0].subplots() + + with pytest.raises(UnsupportedPlotError, match="subfigures"): + to_json(figure) + with pytest.raises(UnsupportedPlotError, match="top-level figure"): + to_json(axes) + + +def valid_data() -> dict[str, Any]: + _, axes = plt.subplots() + axes.plot([1, 2], [3, 4], label="line") + axes.scatter([1], [2]) + axes.bar([1], [2]) + axes.legend() + return json.loads(json.dumps(to_json(axes)))["_value"] + + +def _set(path: list[Any], value: Any) -> Callable[[dict[str, Any]], None]: + def change(data: dict[str, Any]) -> None: + target: Any = data + for key in path[:-1]: + target = target[key] + target[path[-1]] = value + + return change + + +def _delete(path: list[Any]) -> Callable[[dict[str, Any]], None]: + def change(data: dict[str, Any]) -> None: + target: Any = data + for key in path[:-1]: + target = target[key] + del target[path[-1]] + + return change + + +LINE = ["axes", 0, "lines", 0] + + +@pytest.mark.parametrize( + "change", + [ + _set(["axes"], "not a list"), + _set(["size"], [1]), + _set(["size"], [0, 1]), + _set(["size"], [1e9, 1]), + _set(["size"], [float("nan"), 1]), + _set(["suptitle"], 5), + _delete(["axes", 0, "lines"]), + _set(["axes", 0], []), + _set(["axes", 0, "position"], [0, 0, 1]), + _set(["axes", 0, "x_axis", "scale"], "__import__('os')"), + _set(["axes", 0, "x_axis", "units"], "other"), + _set(["axes", 0, "x_limits"], [float("nan"), 1]), + _set(["axes", 0, "x_limits"], ["0", "1"]), + _set([*LINE, "x"], "12"), + _set([*LINE, "x"], ["1", "2"]), + _set([*LINE, "x"], [True, False]), + _set([*LINE, "x"], [[1], [2]]), + _set([*LINE, "x"], [1]), + _set([*LINE, "color"], "not a color"), + _set([*LINE, "color"], {"a": 1}), + _set([*LINE, "linestyle"], "wavy"), + _set([*LINE, "drawstyle"], "wavy"), + _set([*LINE, "marker"], "not a marker"), + _set([*LINE, "marker"], 5), + _set([*LINE, "linewidth"], "wide"), + _set(["axes", 0, "bars", 0, "orientation"], "diagonal"), + _set(["axes", 0, "bars", 0, "height"], []), + _set(["axes", 0, "scatters", 0, "x"], []), + _set(["axes", 0, "legend", "labels"], "line"), + _set(["axes", 0, "title"], "title"), + ], +) +def test_malformed_data_raises_value_error(change: Callable[[dict[str, Any]], None]) -> None: + data = valid_data() + change(data) + + with pytest.raises(ValueError, match="matplotlib_figure has invalid data"): + plot_from_json(data) + + +def test_category_value_out_of_range_rejected() -> None: + _, axes = plt.subplots() + axes.plot(["a", "b"], [1, 2]) + data = json.loads(json.dumps(to_json(axes)))["_value"] + data["axes"][0]["lines"][0]["x"] = [0, 5] + + with pytest.raises(ValueError, match="matplotlib_figure has invalid data"): + plot_from_json(data) + + +def test_missing_value_rejected() -> None: + with pytest.raises(ValueError, match="should have value"): + from_json({"_type": "matplotlib_figure"}) + + +def test_other_matplotlib_values_are_not_converted() -> None: + line = plt.plot([1, 2], [3, 4])[0] + + assert to_json(line) is line From ad970de1670df433d20e8837271f42e007a0facc Mon Sep 17 00:00:00 2001 From: Eliot Robson Date: Thu, 5 Nov 2026 15:35:50 -0500 Subject: [PATCH 2/2] Fix mypy errors in the plain matplotlib tests matplotlib-stubs types get_yscale as an ellipsis and leaves out Figure.get_suptitle, so annotate those values as Any. Co-Authored-By: Claude Opus 5.5 --- tests/scenario_root/test_plain_matplotlib/scenario.py | 5 ++++- tests/test_matplotlib_json.py | 10 ++++++++-- 2 files changed, 12 insertions(+), 3 deletions(-) diff --git a/tests/scenario_root/test_plain_matplotlib/scenario.py b/tests/scenario_root/test_plain_matplotlib/scenario.py index 0c2e9b5..7f37cc8 100644 --- a/tests/scenario_root/test_plain_matplotlib/scenario.py +++ b/tests/scenario_root/test_plain_matplotlib/scenario.py @@ -1,3 +1,5 @@ +from typing import Any + import numpy as np import pytest from matplotcheck.base import PlotTester @@ -19,7 +21,8 @@ def check_semilogy(figure: Figure) -> None: assert isinstance(figure, Figure) assert len(figure.axes) == 1 - axes = figure.axes[0] + # The matplotlib stubs don't type get_yscale correctly. + axes: Any = figure.axes[0] assert axes.get_xscale() == "linear" assert axes.get_yscale() == "log" x, y = axes.get_lines()[0].get_data() diff --git a/tests/test_matplotlib_json.py b/tests/test_matplotlib_json.py index 3b507d0..48f2979 100644 --- a/tests/test_matplotlib_json.py +++ b/tests/test_matplotlib_json.py @@ -28,6 +28,11 @@ def close_figures() -> Any: plt.close("all") +def get_suptitle(figure: Any) -> str: + # The matplotlib stubs don't include Figure.get_suptitle. + return figure.get_suptitle() + + def round_trip(value: Any) -> Any: """Send a value through real JSON text, as the sandbox does.""" return from_json(json.loads(json.dumps(to_json(value)))) @@ -44,7 +49,8 @@ def test_semilogy_axes() -> None: assert isinstance(figure, Figure) assert len(figure.axes) == 1 - axes = figure.axes[0] + # The matplotlib stubs don't type get_yscale correctly. + axes: Any = figure.axes[0] assert axes.get_xscale() == "linear" assert axes.get_yscale() == "log" assert (axes.get_xlabel(), axes.get_ylabel(), axes.get_title()) == ("x", "f(x)", "Growth") @@ -71,7 +77,7 @@ def test_figure_value_and_all_subplots_kept() -> None: for rebuilt in (from_figure, from_axes): assert isinstance(rebuilt, Figure) - assert rebuilt.get_suptitle() == "Both" + assert get_suptitle(rebuilt) == "Both" assert len(rebuilt.axes) == 2 np.testing.assert_allclose(rebuilt.axes[0].get_lines()[0].get_ydata(), [3, 4]) np.testing.assert_allclose(rebuilt.axes[1].get_lines()[0].get_ydata(), [7, 8])