diff --git a/pygmt/figure.py b/pygmt/figure.py index 05cba50ed02..9ad6f737553 100644 --- a/pygmt/figure.py +++ b/pygmt/figure.py @@ -562,22 +562,20 @@ def _repr_html_(self): ) -def set_display(method=None): +def set_display(method: Literal["external", "notebook", "none", None] = None): """ Set the display method when calling :meth:`pygmt.Figure.show`. Parameters ---------- - method : str or None + method The method to display an image preview. Choose from: - ``"external"``: External PDF preview using the default PDF viewer - ``"notebook"``: Inline PNG preview in the current notebook - ``"none"``: Disable image preview - - ``None``: Reset to the default display method - - The default display method is ``"external"`` in Python consoles or - ``"notebook"`` in Jupyter notebooks. + - ``None``: Reset to the default display method, which is either ``"external"`` + in Python consoles or ``"notebook"`` in Jupyter notebooks. Examples -------- @@ -600,10 +598,13 @@ def set_display(method=None): >>> pygmt.set_display(method=None) >>> fig.show() # again, will show a PNG image in the current notebook """ - if method in {"notebook", "external", "none"}: - SHOW_CONFIG["method"] = method - elif method is not None: - raise GMTInvalidInput( - f"Invalid display mode '{method}', " - "should be either 'notebook', 'external', 'none' or None." - ) + match method: + case "external" | "notebook" | "none": + SHOW_CONFIG["method"] = method # type: ignore[assignment] + case None: + SHOW_CONFIG["method"] = _get_default_display_method() # type: ignore[assignment] + case _: + raise GMTInvalidInput( + f"Invalid display method '{method}'. Valid values are 'external'," + "'notebook', 'none' or None." + ) diff --git a/pygmt/tests/test_figure.py b/pygmt/tests/test_figure.py index aadaaacfad3..4c8e184d118 100644 --- a/pygmt/tests/test_figure.py +++ b/pygmt/tests/test_figure.py @@ -11,7 +11,7 @@ import pytest from pygmt import Figure, set_display from pygmt.exceptions import GMTError, GMTInvalidInput -from pygmt.figure import _get_default_display_method +from pygmt.figure import SHOW_CONFIG, _get_default_display_method from pygmt.helpers import GMTTempFile try: @@ -373,12 +373,31 @@ def test_figure_display_external(): fig.show(method="external") -def test_figure_set_display_invalid(): +class TestSetDisplay: """ - Test to check if an error is raised when an invalid method is passed to set_display. + Test the pygmt.set_display method. """ - with pytest.raises(GMTInvalidInput): - set_display(method="invalid") + + def test_set_display(self): + """ + Test if pygmt.set_display updates the SHOW_CONFIG variable correctly. + """ + default_method = SHOW_CONFIG["method"] # Current default method + + for method in ("notebook", "external", "none"): + set_display(method=method) + assert SHOW_CONFIG["method"] == method + + # Setting method to None should revert it to the default method. + set_display(method=None) + assert SHOW_CONFIG["method"] == default_method + + def test_invalid_method(self): + """ + Test if an error is raised when an invalid method is passed. + """ + with pytest.raises(GMTInvalidInput): + set_display(method="invalid") def test_figure_unsupported_xshift_yshift():