33"""
44
55import warnings
6+ from collections .abc import Sequence
67from typing import Literal
78
89from pygmt .exceptions import GMTInvalidInput
910
1011
1112def validate_output_table_type (
12- output_type : Literal ["pandas" , "numpy" , "file" ], outfile : str | None = None
13+ output_type : Literal ["pandas" , "numpy" , "file" ],
14+ valid_types : Sequence [str ] = ("pandas" , "numpy" , "file" ),
15+ outfile : str | None = None ,
1316) -> Literal ["pandas" , "numpy" , "file" ]:
1417 """
1518 Check if the ``output_type`` and ``outfile`` parameters are valid.
1619
1720 Parameters
1821 ----------
1922 output_type
20- Desired output type of tabular data. Valid values are ``"pandas"``,
21- ``"numpy"`` and ``"file"``.
23+ Desired output type of tabular data. Default valid values are ``"pandas"``,
24+ ``"numpy"`` and ``"file"``, but can be configured by parameter ``valid_types``.
25+ valid_types
26+ Tuple of valid desired output types.
2227 outfile
2328 File name for saving the result data. Required if ``output_type`` is ``"file"``.
2429 If specified, ``output_type`` will be forced to be ``"file"``.
@@ -36,23 +41,32 @@ def validate_output_table_type(
3641 'numpy'
3742 >>> validate_output_table_type(output_type="file", outfile="output-fname.txt")
3843 'file'
44+ >>> validate_output_table_type(output_type="pandas", valid_types=("pandas", "file"))
45+ 'pandas'
3946 >>> validate_output_table_type(output_type="invalid-type")
4047 Traceback (most recent call last):
4148 ...
42- pygmt.exceptions.GMTInvalidInput: Must specify 'output_type' either as 'file ', ...
49+ pygmt.exceptions.GMTInvalidInput: Must specify 'output_type' as 'pandas ', ...
4350 >>> validate_output_table_type("file", outfile=None)
4451 Traceback (most recent call last):
4552 ...
4653 pygmt.exceptions.GMTInvalidInput: Must specify 'outfile' for output_type='file'.
54+ >>> validate_output_table_type(output_type="numpy", valid_types=("pandas", "file"))
55+ Traceback (most recent call last):
56+ ...
57+ pygmt.exceptions.GMTInvalidInput: Must specify 'output_type' as 'pandas', or 'file'.
4758 >>> with warnings.catch_warnings(record=True) as w:
4859 ... validate_output_table_type("pandas", outfile="not-none.txt")
4960 ... assert len(w) == 1
5061 'file'
5162 """
52- if output_type not in ["file" , "numpy" , "pandas" ]:
53- raise GMTInvalidInput (
54- "Must specify 'output_type' either as 'file', 'numpy', or 'pandas'."
63+ if output_type not in valid_types :
64+ msg = (
65+ "Must specify 'output_type' as "
66+ + ", " .join (f"'{ v } '" for v in valid_types [:- 1 ])
67+ + f", or '{ valid_types [- 1 ]} '."
5568 )
69+ raise GMTInvalidInput (msg )
5670 if output_type == "file" and outfile is None :
5771 raise GMTInvalidInput ("Must specify 'outfile' for output_type='file'." )
5872 if output_type != "file" and outfile is not None :
0 commit comments