Accept plain matplotlib Axes and Figure values as JSON

.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
 

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)

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())

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