3
3
"""
4
4
5
5
import warnings
6
+ from collections .abc import Sequence
6
7
from typing import Literal
7
8
8
9
from pygmt .exceptions import GMTInvalidInput
9
10
10
11
11
12
def 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 ,
13
16
) -> Literal ["pandas" , "numpy" , "file" ]:
14
17
"""
15
18
Check if the ``output_type`` and ``outfile`` parameters are valid.
16
19
17
20
Parameters
18
21
----------
19
22
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.
22
27
outfile
23
28
File name for saving the result data. Required if ``output_type`` is ``"file"``.
24
29
If specified, ``output_type`` will be forced to be ``"file"``.
@@ -36,23 +41,32 @@ def validate_output_table_type(
36
41
'numpy'
37
42
>>> validate_output_table_type(output_type="file", outfile="output-fname.txt")
38
43
'file'
44
+ >>> validate_output_table_type(output_type="pandas", valid_types=("pandas", "file"))
45
+ 'pandas'
39
46
>>> validate_output_table_type(output_type="invalid-type")
40
47
Traceback (most recent call last):
41
48
...
42
- pygmt.exceptions.GMTInvalidInput: Must specify 'output_type' either as 'file ', ...
49
+ pygmt.exceptions.GMTInvalidInput: Must specify 'output_type' as 'pandas ', ...
43
50
>>> validate_output_table_type("file", outfile=None)
44
51
Traceback (most recent call last):
45
52
...
46
53
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'.
47
58
>>> with warnings.catch_warnings(record=True) as w:
48
59
... validate_output_table_type("pandas", outfile="not-none.txt")
49
60
... assert len(w) == 1
50
61
'file'
51
62
"""
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 ]} '."
55
68
)
69
+ raise GMTInvalidInput (msg )
56
70
if output_type == "file" and outfile is None :
57
71
raise GMTInvalidInput ("Must specify 'outfile' for output_type='file'." )
58
72
if output_type != "file" and outfile is not None :
0 commit comments