pygmt.set_display: Fix the bug that method=None doesn't reset to the default display method

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."
+            )

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