Allow validate_output_table_type to specify the supported output types (

#3191)
GenericMappingTools · May 7, 2024 · a4d2b8e · a4d2b8e
1 parent c783a79
commit a4d2b8e
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 8 deletions.
diff --git a/pygmt/helpers/validators.py b/pygmt/helpers/validators.py
@@ -3,22 +3,27 @@
 """
 
 import warnings
+from collections.abc import Sequence
 from typing import Literal
 
 from pygmt.exceptions import GMTInvalidInput
 
 
 def validate_output_table_type(
-    output_type: Literal["pandas", "numpy", "file"], outfile: str | None = None
+    output_type: Literal["pandas", "numpy", "file"],
+    valid_types: Sequence[str] = ("pandas", "numpy", "file"),
+    outfile: str | None = None,
 ) -> Literal["pandas", "numpy", "file"]:
     """
     Check if the ``output_type`` and ``outfile`` parameters are valid.
 
     Parameters
     ----------
     output_type
-        Desired output type of tabular data. Valid values are ``"pandas"``,
-        ``"numpy"`` and ``"file"``.
+        Desired output type of tabular data. Default valid values are ``"pandas"``,
+        ``"numpy"`` and ``"file"``, but can be configured by parameter ``valid_types``.
+    valid_types
+        Tuple of valid desired output types.
     outfile
         File name for saving the result data. Required if ``output_type`` is ``"file"``.
         If specified, ``output_type`` will be forced to be ``"file"``.
@@ -36,23 +41,32 @@ def validate_output_table_type(
     'numpy'
     >>> validate_output_table_type(output_type="file", outfile="output-fname.txt")
     'file'
+    >>> validate_output_table_type(output_type="pandas", valid_types=("pandas", "file"))
+    'pandas'
     >>> validate_output_table_type(output_type="invalid-type")
     Traceback (most recent call last):
         ...
-    pygmt.exceptions.GMTInvalidInput: Must specify 'output_type' either as 'file', ...
+    pygmt.exceptions.GMTInvalidInput: Must specify 'output_type' as 'pandas', ...
     >>> validate_output_table_type("file", outfile=None)
     Traceback (most recent call last):
         ...
     pygmt.exceptions.GMTInvalidInput: Must specify 'outfile' for output_type='file'.
+    >>> validate_output_table_type(output_type="numpy", valid_types=("pandas", "file"))
+    Traceback (most recent call last):
+        ...
+    pygmt.exceptions.GMTInvalidInput: Must specify 'output_type' as 'pandas', or 'file'.
     >>> with warnings.catch_warnings(record=True) as w:
     ...     validate_output_table_type("pandas", outfile="not-none.txt")
     ...     assert len(w) == 1
     'file'
     """
-    if output_type not in ["file", "numpy", "pandas"]:
-        raise GMTInvalidInput(
-            "Must specify 'output_type' either as 'file', 'numpy', or 'pandas'."
+    if output_type not in valid_types:
+        msg = (
+            "Must specify 'output_type' as "
+            + ", ".join(f"'{v}'" for v in valid_types[:-1])
+            + f", or '{valid_types[-1]}'."
         )
+        raise GMTInvalidInput(msg)
     if output_type == "file" and outfile is None:
         raise GMTInvalidInput("Must specify 'outfile' for output_type='file'.")
     if output_type != "file" and outfile is not None:

diff --git a/pygmt/src/triangulate.py b/pygmt/src/triangulate.py
@@ -233,7 +233,7 @@ def delaunay_triples(
         ``triangulate`` is a Cartesian or small-geographic area operator and is
         unaware of periodic or polar boundary conditions.
         """
-        output_type = validate_output_table_type(output_type, outfile)
+        output_type = validate_output_table_type(output_type, outfile=outfile)
 
         with Session() as lib:
             with (