From 6bd742e6aa2c0bb48e2b7b9adaa056cff06d96d2 Mon Sep 17 00:00:00 2001 From: Dongdong Tian Date: Thu, 8 Aug 2024 13:25:36 +0800 Subject: [PATCH] pygmt.set_display: Fix a bug when setting `method=None`, add type hints and refactor --- pygmt/figure.py | 27 ++++++++++++++------------- 1 file changed, 14 insertions(+), 13 deletions(-) diff --git a/pygmt/figure.py b/pygmt/figure.py index 7e503667a17..f2135d0aa61 100644 --- a/pygmt/figure.py +++ b/pygmt/figure.py @@ -568,22 +568,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 -------- @@ -606,10 +604,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 + case None: + SHOW_CONFIG["method"] = _get_default_display_method() + case _: + raise GMTInvalidInput( + f"Invalid display method '{method}'. Valid values are 'external'," + "'notebook', 'none' or None." + )