Add feedback.add_image to show plots with test results

quick_start.md

@@ -303,6 +303,7 @@ def test_with_partial_credit(sandbox: StudentFixture, feedback: FeedbackFixture)
 - `feedback.set_score(fraction)`: Set partial credit (0.0 to 1.0)
 - `feedback.set_score_final(fraction)`: Set final score (prevents further updates)
 - `feedback.add_message(msg)`: Add custom feedback message
+- `feedback.add_image(image, label=None)`: Show an image with the test's results. See [Showing Plots to Students](#showing-plots-to-students)
 
 ## Advanced Features
 
@@ -1108,6 +1109,28 @@ def test_student_plot(sandbox: StudentFixture) -> None:
     pt.assert_title_contains(["Data Visualization"])
 ```
 
+#### Showing Plots to Students
+
+To show the student's plot with a test's results, pass it to `feedback.add_image`:
+
+```python
+@pytest.mark.grading_data(name="Plot Test", points=5)
+def test_student_plot(sandbox: StudentFixture, feedback: FeedbackFixture) -> None:
+    plot = sandbox.query_function("create_plot", data)
+    feedback.add_image(plot, label="Your plot")
+    ...
+```
+
+- **Values:** `add_image` takes a matplotlib `Figure` or `Axes`, or the bytes of a PNG file, and
+  adds it to the test's `images` in `results.json`. `<pl-external-grader-results>` shows it under
+  the test, whether the test passes or fails.
+- **Call it first:** a failed assertion ends the test, so add the image before checking the plot.
+- **A copy of the plot:** the image is drawn from the `Figure` that the test received, which the
+  grader rebuilt from the student's plot. It shows the data, scales, labels and other parts the
+  grader copies, but other styling can look different from what the student saw.
+- **Size:** PrairieLearn limits the whole `results.json` to 1 MB. A typical plot is tens of
+  kilobytes, but don't add an image in every case of a test with many parameters.
+
 The grader automatically serializes/deserializes:
 
 - Matplotlib `Axes` and `Figure` values (received as a `Figure`)

src/pytest_prairielearn_grader/fixture.py

@@ -1,3 +1,4 @@
+import base64
 import json
 import logging
 import os
@@ -127,18 +128,40 @@ class FeedbackFixture:
 
     test_id: str
     messages: list[str]
+    images: list[dict[str, str]]
     score: float | None
     final_score_override: bool
 
     def __init__(self, test_id: str) -> None:
         self.test_id = test_id
         self.messages = []
+        self.images = []
         self.score = None
         self.final_score_override = False
 
     def add_message(self, message: str) -> None:
         self.messages.append(message)
 
+    def add_image(self, image: Any, *, label: str | None = None) -> None:
+        """
+        Shows an image with this test's results, such as a plot queried from student code.
+
+        `image` is a matplotlib Figure or Axes, or the bytes of a PNG file. PrairieLearn limits the
+        whole results file to 1 MB, so add only a few images.
+        """
+        if isinstance(image, bytes):
+            png = image
+        else:
+            # matplotlib is an optional dependency, so only import it when a plot is added.
+            from pytest_prairielearn_grader.matplotlib_json import plot_to_png  # noqa: PLC0415
+
+            png = plot_to_png(image)
+
+        entry = {"url": "data:image/png;base64," + base64.b64encode(png).decode("ascii")}
+        if label is not None:
+            entry["label"] = label
+        self.images.append(entry)
+
     def set_score(self, score: float) -> None:
         self.score = score
 
@@ -155,11 +178,14 @@ def set_score_final(self, score: float) -> None:
         self.final_score_override = True
 
     def to_dict(self) -> dict:
-        return {
+        result: dict[str, Any] = {
             "test_id": self.test_id,
             "message": os.linesep.join(self.messages),
             "points_frac": self.score,
         }
+        if self.images:
+            result["images"] = self.images
+        return result
 
 
 class _SandboxBase:

src/pytest_prairielearn_grader/matplotlib_json.py

@@ -8,6 +8,7 @@
 matplotlib is an optional dependency, so json_utils only imports this module when it's needed.
 """
 
+import io
 import math
 from typing import Any
 
@@ -40,6 +41,17 @@ def is_plot(value: Any) -> bool:
     return isinstance(value, Axes | Figure)
 
 
+def plot_to_png(value: Any) -> bytes:
+    """Render a Figure, or the figure that an Axes belongs to, as a PNG file."""
+    if isinstance(value, Axes):
+        value = value.figure
+    if not isinstance(value, Figure):
+        raise TypeError(f"add_image expects a matplotlib Figure or Axes, or PNG bytes, not {type(value).__name__}")
+    buffer = io.BytesIO()
+    value.savefig(buffer, format="png")
+    return buffer.getvalue()
+
+
 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."""
 

tests/scenario_root/test_plain_matplotlib/expected_outcome.json

@@ -1,6 +1,6 @@
 {
   "expected_data_object": {
-    "score": 0.8333333333333334,
+    "score": 0.8571428571428571,
     "tests": [
       {
         "max_points": 1,
@@ -50,6 +50,18 @@
         "outcome": "failed",
         "test_id": "test_plain_matplotlib.py::test_rejected_plot_shows_clear_message[student_code]",
         "message": "errorbar plots aren't supported"
+      },
+      {
+        "max_points": 1,
+        "name": "add_image",
+        "points": 1.0,
+        "points_frac": 1.0,
+        "outcome": "passed",
+        "test_id": "test_plain_matplotlib.py::test_add_image[student_code]",
+        "images": [
+          "Your plot",
+          null
+        ]
       }
     ]
   }

tests/scenario_root/test_plain_matplotlib/scenario.py

@@ -6,6 +6,7 @@
 from matplotlib.figure import Figure
 
 from pytest_prairielearn_grader.assertions import assert_true
+from pytest_prairielearn_grader.fixture import FeedbackFixture
 from pytest_prairielearn_grader.fixture import StudentFixture
 from pytest_prairielearn_grader.utils import GradingOutputLevel
 from pytest_prairielearn_grader.utils import QueryStatusCode
@@ -74,3 +75,9 @@ def test_rejected_plot_keeps_sandbox_usable(sandbox: StudentFixture) -> None:
 @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")
+
+
+@pytest.mark.grading_data(name="add_image", points=1)
+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"))

tests/test_autograder_scenarios.py

@@ -183,6 +183,12 @@ def test_autograder_scenario_with_pytester(pytester: pytest.Pytester, scenario_d
                 f"Message mismatch for test '{test_id}'.\nassert {normalized_expected!r} in {normalized_actual!r}"
             )
 
+        # The expected images are their labels, since the PNG bytes depend on the matplotlib version.
+        if "images" in expected_test:
+            actual_images = actual_test.get("images", [])
+            assert [image.get("label") for image in actual_images] == expected_test["images"], f"Images mismatch for test '{test_id}'."
+            assert all(image["url"].startswith("data:image/png;base64,") for image in actual_images)
+
         assert actual_test["max_points"] == expected_test["max_points"], f"Max points mismatch for test '{test_id}'."
         assert math.isclose(expected_test["points_frac"], actual_test["points_frac"]), f"Points fraction mismatch for test '{test_id}'."
         assert math.isclose(expected_test["points"], actual_test["points"]), f"Points mismatch for test '{test_id}'."

tests/test_feedback.py

@@ -0,0 +1,45 @@
+import base64
+
+import pytest
+from matplotlib.figure import Figure
+
+from pytest_prairielearn_grader.fixture import FeedbackFixture
+
+PNG_SIGNATURE = b"\x89PNG\r\n\x1a\n"
+
+
+def decode(url: str) -> bytes:
+    prefix = "data:image/png;base64,"
+    assert url.startswith(prefix)
+    return base64.b64decode(url.removeprefix(prefix))
+
+
+def test_add_image_from_figure_and_axes() -> None:
+    figure = Figure()
+    axes = figure.add_subplot()
+    axes.plot([1, 2, 3])
+    feedback = FeedbackFixture("test")
+
+    feedback.add_image(figure, label="Figure")
+    feedback.add_image(axes)
+
+    images = feedback.to_dict()["images"]
+    assert [image.get("label") for image in images] == ["Figure", None]
+    assert all(decode(image["url"]).startswith(PNG_SIGNATURE) for image in images)
+
+
+def test_add_image_from_bytes() -> None:
+    feedback = FeedbackFixture("test")
+    feedback.add_image(PNG_SIGNATURE + b"rest", label="Raw")
+    assert feedback.to_dict()["images"] == [
+        {"url": "data:image/png;base64," + base64.b64encode(PNG_SIGNATURE + b"rest").decode(), "label": "Raw"}
+    ]
+
+
+def test_no_images_key_without_images() -> None:
+    assert "images" not in FeedbackFixture("test").to_dict()
+
+
+def test_add_image_rejects_other_values() -> None:
+    with pytest.raises(TypeError, match="not str"):
+        FeedbackFixture("test").add_image("plot.png")