From a8782730543a56d8a7a595fe5169f44024ecd465 Mon Sep 17 00:00:00 2001
From: Elliott Sales de Andrade <quantum.analyst@gmail.com>
Date: Mon, 12 Feb 2024 19:36:32 -0500
Subject: [PATCH] TST: Ensure Matplotlib is always cleaned up

The seaborn test also uses Matplotlib but was not wrapped in the cleanup
fixture, As there are now 3 files that need this fixture, refactor to
reduce code duplication.
---
 pandas/conftest.py                            | 34 +++++++++++++++++++
 .../tests/io/formats/style/test_matplotlib.py | 22 +-----------
 pandas/tests/plotting/conftest.py             | 21 ++----------
 pandas/tests/test_downstream.py               |  2 +-
 4 files changed, 38 insertions(+), 41 deletions(-)

diff --git a/pandas/conftest.py b/pandas/conftest.py
index 868bd365fa0ba..254d605e13460 100644
--- a/pandas/conftest.py
+++ b/pandas/conftest.py
@@ -28,6 +28,7 @@
     timezone,
 )
 from decimal import Decimal
+import gc
 import operator
 import os
 from typing import (
@@ -1863,6 +1864,39 @@ def ip():
     return InteractiveShell(config=c)
 
 
+@pytest.fixture
+def mpl_cleanup():
+    """
+    Ensure Matplotlib is cleaned up around a test.
+
+    Before a test is run:
+
+    1) Set the backend to "template" to avoid requiring a GUI.
+
+    After a test is run:
+
+    1) Reset units registry
+    2) Reset rc_context
+    3) Close all figures
+
+    See matplotlib/testing/decorators.py#L24.
+    """
+    mpl = pytest.importorskip("matplotlib")
+    mpl_units = pytest.importorskip("matplotlib.units")
+    plt = pytest.importorskip("matplotlib.pyplot")
+    orig_units_registry = mpl_units.registry.copy()
+    try:
+        with mpl.rc_context():
+            mpl.use("template")
+            yield
+    finally:
+        mpl_units.registry.clear()
+        mpl_units.registry.update(orig_units_registry)
+        plt.close("all")
+        # https://matplotlib.org/stable/users/prev_whats_new/whats_new_3.6.0.html#garbage-collection-is-no-longer-run-on-figure-close  # noqa: E501
+        gc.collect(1)
+
+
 @pytest.fixture(
     params=[
         getattr(pd.offsets, o)
diff --git a/pandas/tests/io/formats/style/test_matplotlib.py b/pandas/tests/io/formats/style/test_matplotlib.py
index ef7bfb11d81d8..70ddd65c02d14 100644
--- a/pandas/tests/io/formats/style/test_matplotlib.py
+++ b/pandas/tests/io/formats/style/test_matplotlib.py
@@ -1,5 +1,3 @@
-import gc
-
 import numpy as np
 import pytest
 
@@ -16,25 +14,7 @@
 
 from pandas.io.formats.style import Styler
 
-
-@pytest.fixture(autouse=True)
-def mpl_cleanup():
-    # matplotlib/testing/decorators.py#L24
-    # 1) Resets units registry
-    # 2) Resets rc_context
-    # 3) Closes all figures
-    mpl = pytest.importorskip("matplotlib")
-    mpl_units = pytest.importorskip("matplotlib.units")
-    plt = pytest.importorskip("matplotlib.pyplot")
-    orig_units_registry = mpl_units.registry.copy()
-    with mpl.rc_context():
-        mpl.use("template")
-        yield
-    mpl_units.registry.clear()
-    mpl_units.registry.update(orig_units_registry)
-    plt.close("all")
-    # https://matplotlib.org/stable/users/prev_whats_new/whats_new_3.6.0.html#garbage-collection-is-no-longer-run-on-figure-close  # noqa: E501
-    gc.collect(1)
+pytestmark = pytest.mark.usefixtures("mpl_cleanup")
 
 
 @pytest.fixture
diff --git a/pandas/tests/plotting/conftest.py b/pandas/tests/plotting/conftest.py
index d688bbd47595c..eb5a1f1f6382e 100644
--- a/pandas/tests/plotting/conftest.py
+++ b/pandas/tests/plotting/conftest.py
@@ -1,5 +1,3 @@
-import gc
-
 import numpy as np
 import pytest
 
@@ -10,23 +8,8 @@
 
 
 @pytest.fixture(autouse=True)
-def mpl_cleanup():
-    # matplotlib/testing/decorators.py#L24
-    # 1) Resets units registry
-    # 2) Resets rc_context
-    # 3) Closes all figures
-    mpl = pytest.importorskip("matplotlib")
-    mpl_units = pytest.importorskip("matplotlib.units")
-    plt = pytest.importorskip("matplotlib.pyplot")
-    orig_units_registry = mpl_units.registry.copy()
-    with mpl.rc_context():
-        mpl.use("template")
-        yield
-    mpl_units.registry.clear()
-    mpl_units.registry.update(orig_units_registry)
-    plt.close("all")
-    # https://matplotlib.org/stable/users/prev_whats_new/whats_new_3.6.0.html#garbage-collection-is-no-longer-run-on-figure-close  # noqa: E501
-    gc.collect(1)
+def autouse_mpl_cleanup(mpl_cleanup):
+    pass
 
 
 @pytest.fixture
diff --git a/pandas/tests/test_downstream.py b/pandas/tests/test_downstream.py
index ead36ee08b407..7186056f90299 100644
--- a/pandas/tests/test_downstream.py
+++ b/pandas/tests/test_downstream.py
@@ -155,7 +155,7 @@ def test_scikit_learn():
     clf.predict(digits.data[-1:])
 
 
-def test_seaborn():
+def test_seaborn(mpl_cleanup):
     seaborn = pytest.importorskip("seaborn")
     tips = DataFrame(
         {"day": pd.date_range("2023", freq="D", periods=5), "total_bill": range(5)}