diff --git a/astropy/tests/figures/py311-test-image-mpl360-cov.json b/astropy/tests/figures/py311-test-image-mpl360-cov.json index 38fbc9f68f66..4c0f42f67984 100644 --- a/astropy/tests/figures/py311-test-image-mpl360-cov.json +++ b/astropy/tests/figures/py311-test-image-mpl360-cov.json @@ -55,5 +55,6 @@ "astropy.visualization.wcsaxes.tests.test_transform_coord_meta.TestTransformCoordMeta.test_coords_overlay_auto_coord_meta": "2f737bb70fb1a5452cb0efa79010376614dc559e9aff607f638f044ac6b04448", "astropy.visualization.wcsaxes.tests.test_transform_coord_meta.TestTransformCoordMeta.test_direct_init": "1f24c5243bfdf0f30e88afc4f2f5d66db85954cea50a2910af94975ff8765b45", "astropy.visualization.wcsaxes.tests.test_wcsapi.test_wcsapi_5d_with_names": "c90ae6f3b0f9f407ca7a866fa648dc3ef4bea78369315292bc82f16104ed5736", - "astropy.visualization.wcsaxes.tests.test_wcsapi.test_wcsapi_2d_celestial_arcsec": "4b743d645a85d7516decbcf4a831c127af5a1800072597c2a1299d17fb186adb" + "astropy.visualization.wcsaxes.tests.test_wcsapi.test_wcsapi_2d_celestial_arcsec": "4b743d645a85d7516decbcf4a831c127af5a1800072597c2a1299d17fb186adb", + "astropy.visualization.wcsaxes.tests.test_images.test_astropy_support": "aee3dd8cfc0998ec2d904549c6527fa55f3b16cbe9a30c222a2c0ac6a6db155e" } diff --git a/astropy/tests/figures/py311-test-image-mpldev-cov.json b/astropy/tests/figures/py311-test-image-mpldev-cov.json index b83783c6b23c..15b96718b64e 100644 --- a/astropy/tests/figures/py311-test-image-mpldev-cov.json +++ b/astropy/tests/figures/py311-test-image-mpldev-cov.json @@ -55,5 +55,6 @@ "astropy.visualization.wcsaxes.tests.test_transform_coord_meta.TestTransformCoordMeta.test_coords_overlay_auto_coord_meta": "2f737bb70fb1a5452cb0efa79010376614dc559e9aff607f638f044ac6b04448", "astropy.visualization.wcsaxes.tests.test_transform_coord_meta.TestTransformCoordMeta.test_direct_init": "1f24c5243bfdf0f30e88afc4f2f5d66db85954cea50a2910af94975ff8765b45", "astropy.visualization.wcsaxes.tests.test_wcsapi.test_wcsapi_5d_with_names": "c90ae6f3b0f9f407ca7a866fa648dc3ef4bea78369315292bc82f16104ed5736", - "astropy.visualization.wcsaxes.tests.test_wcsapi.test_wcsapi_2d_celestial_arcsec": "4b743d645a85d7516decbcf4a831c127af5a1800072597c2a1299d17fb186adb" + "astropy.visualization.wcsaxes.tests.test_wcsapi.test_wcsapi_2d_celestial_arcsec": "4b743d645a85d7516decbcf4a831c127af5a1800072597c2a1299d17fb186adb", + "astropy.visualization.wcsaxes.tests.test_images.test_astropy_support": "aee3dd8cfc0998ec2d904549c6527fa55f3b16cbe9a30c222a2c0ac6a6db155e" } diff --git a/astropy/visualization/astropy_support.py b/astropy/visualization/astropy_support.py new file mode 100644 index 000000000000..dc880d6999a7 --- /dev/null +++ b/astropy/visualization/astropy_support.py @@ -0,0 +1,44 @@ +from contextlib import ExitStack, contextmanager + +from astropy.visualization.time import time_support +from astropy.visualization.units import quantity_support + + +@contextmanager +def astropy_support(*, quantity_support_kwargs=None, time_support_kwargs=None): + """ + Enable support for plotting `astropy.units.Quantity` and `astropy.time.Time` instances in + matplotlib. + + It can be used as a decorator or with a ``with`` statement. + + Examples + -------- + + .. plot:: + :include-source: + + import matplotlib.pyplot as plt + import astropy.units as u + from astropy.time import Time + from astropy.visualization.astropy_support import astropy_support + + @astropy_support() + def plot_example(): + plt.figure() + plt.plot([1, 2, 3] * u.m) + plt.plot(Time(['2000-01-01', '2000-01-02', '2000-01-03']).plot_date) + plt.draw() + plt.show() + + with astropy_support(): # doctest: +IGNORE_OUTPUT + plt.figure() + plt.plot([1, 2, 3] * u.m) + plt.plot(Time(['2000-01-01', '2000-01-02', '2000-01-03']).plot_date) + plt.draw() + + """ + with ExitStack() as stack: + stack.enter_context(quantity_support(**(quantity_support_kwargs or {}))) + stack.enter_context(time_support(**(time_support_kwargs or {}))) + yield diff --git a/astropy/visualization/wcsaxes/tests/test_images.py b/astropy/visualization/wcsaxes/tests/test_images.py index 26a109559569..97c68ecbd146 100644 --- a/astropy/visualization/wcsaxes/tests/test_images.py +++ b/astropy/visualization/wcsaxes/tests/test_images.py @@ -19,9 +19,11 @@ ) from astropy.io import fits from astropy.tests.figures import figure_test +from astropy.time import Time, TimeDelta from astropy.utils import isiterable from astropy.utils.data import get_pkg_data_filename from astropy.utils.exceptions import AstropyUserWarning +from astropy.visualization.astropy_support import astropy_support from astropy.visualization.wcsaxes import WCSAxes, add_beam, add_scalebar from astropy.visualization.wcsaxes.frame import EllipticalFrame from astropy.visualization.wcsaxes.patches import Quadrangle, SphericalCircle @@ -1233,6 +1235,33 @@ def test_allsky_labels_wrap(): return fig +@figure_test(tolerance=1) +def test_astropy_support(): + x_meters = np.linspace(1, 10, 50) * u.m # Distances in meters + x_kilometers = np.linspace(0.002, 0.011, 50) * u.km # Distances in kilometers + + y_seconds = Time("2000-01-01T00:10:00", scale="utc") + TimeDelta( + np.linspace(0, 3600, 50), format="sec" + ) + + start_time = Time("2000-01-01T00:00:00", scale="utc") + end_time = start_time + TimeDelta(2 / 24, format="jd") + + with astropy_support(): + fig, ax = plt.subplots() + + ax.plot(x_meters, y_seconds, label="Meters vs. Seconds") + ax.plot(x_kilometers, y_seconds, label="Kilometers vs. Seconds") + + ax.set_ylim(start_time, Time((end_time), format="jd")) + + ax.set_xlabel("Distance (m)") + ax.set_ylabel("Time") + ax.legend() + + return fig + + @figure_test def test_tickable_gridlines(): wcs = WCS( diff --git a/docs/changes/visualization/17575.feature.rst b/docs/changes/visualization/17575.feature.rst new file mode 100644 index 000000000000..b0b30fc7732b --- /dev/null +++ b/docs/changes/visualization/17575.feature.rst @@ -0,0 +1,2 @@ +Added a new ``astropy_support`` context manager in ``astropy.visualization``. +This enables ``time_support`` and ``quantity_support`` for plotting ``astropy.units.Quantity`` and ``astropy.time.Time`` instances in matplotlib.