From b0a2c44e07e2515ebf4d5362369f2dea1e4eb1f2 Mon Sep 17 00:00:00 2001 From: Dongdong Tian Date: Sat, 16 Dec 2023 00:08:13 +0800 Subject: [PATCH] Let kwargs_to_strings work with default values and postional arguments (#2826) --- pygmt/helpers/decorators.py | 102 +++++++++++++++++++++++++----------- pygmt/src/subplot.py | 4 +- pygmt/src/timestamp.py | 12 ++--- pygmt/src/which.py | 6 +-- 4 files changed, 78 insertions(+), 46 deletions(-) diff --git a/pygmt/helpers/decorators.py b/pygmt/helpers/decorators.py index a6ab440f296..4e3d6926671 100644 --- a/pygmt/helpers/decorators.py +++ b/pygmt/helpers/decorators.py @@ -641,7 +641,6 @@ def kwargs_to_strings(**conversions): Examples -------- - >>> @kwargs_to_strings( ... R="sequence", i="sequence_comma", files="sequence_space" ... ) @@ -691,20 +690,35 @@ def kwargs_to_strings(**conversions): ... ] ... ) {'R': '2005-01-01T08:00:00.000000000/2015-01-01T12:00:00.123456'} + >>> # Here is a more realistic example + >>> # See https://github.com/GenericMappingTools/pygmt/issues/2361 + >>> @kwargs_to_strings( + ... files="sequence_space", + ... offset="sequence", + ... R="sequence", + ... i="sequence_comma", + ... ) + ... def module(files, offset=("-54p", "-54p"), **kwargs): + ... "A module that prints the arguments it received" + ... print(files, end=" ") + ... print(offset, end=" ") + ... print("{", end="") + ... print( + ... ", ".join(f"'{k}': {repr(kwargs[k])}" for k in sorted(kwargs)), + ... end="", + ... ) + ... print("}") + >>> module(files=["data1.txt", "data2.txt"]) + data1.txt data2.txt -54p/-54p {} + >>> module(["data1.txt", "data2.txt"]) + data1.txt data2.txt -54p/-54p {} + >>> module(files=["data1.txt", "data2.txt"], offset=("20p", "20p")) + data1.txt data2.txt 20p/20p {} + >>> module(["data1.txt", "data2.txt"], ("20p", "20p")) + data1.txt data2.txt 20p/20p {} + >>> module(["data1.txt", "data2.txt"], ("20p", "20p"), R=[1, 2, 3, 4]) + data1.txt data2.txt 20p/20p {'R': '1/2/3/4'} """ - valid_conversions = [ - "sequence", - "sequence_comma", - "sequence_plus", - "sequence_space", - ] - - for arg, fmt in conversions.items(): - if fmt not in valid_conversions: - raise GMTInvalidInput( - f"Invalid conversion type '{fmt}' for argument '{arg}'." - ) - separators = { "sequence": "/", "sequence_comma": ",", @@ -712,37 +726,61 @@ def kwargs_to_strings(**conversions): "sequence_space": " ", } + for arg, fmt in conversions.items(): + if fmt not in separators: + raise GMTInvalidInput( + f"Invalid conversion type '{fmt}' for argument '{arg}'." + ) + # Make the actual decorator function def converter(module_func): """ The decorator that creates our new function with the conversions. """ + sig = signature(module_func) @functools.wraps(module_func) def new_module(*args, **kwargs): """ New module instance that converts the arguments first. """ + # Inspired by https://stackoverflow.com/a/69170441 + bound = sig.bind(*args, **kwargs) + bound.apply_defaults() + for arg, fmt in conversions.items(): - if arg in kwargs: - value = kwargs[arg] - issequence = fmt in separators - if issequence and is_nonstr_iter(value): - for index, item in enumerate(value): - try: - # check if there is a space " " when converting - # a pandas.Timestamp/xr.DataArray to a string. - # If so, use np.datetime_as_string instead. - assert " " not in str(item) - except AssertionError: - # convert datetime-like item to ISO 8601 - # string format like YYYY-MM-DDThh:mm:ss.ffffff - value[index] = np.datetime_as_string( - np.asarray(item, dtype=np.datetime64) - ) - kwargs[arg] = separators[fmt].join(f"{item}" for item in value) + # The arg may be in args or kwargs + if arg in bound.arguments: + value = bound.arguments[arg] + elif arg in bound.arguments.get("kwargs"): + value = bound.arguments["kwargs"][arg] + else: + continue + + issequence = fmt in separators + if issequence and is_nonstr_iter(value): + for index, item in enumerate(value): + try: + # Check if there is a space " " when converting + # a pandas.Timestamp/xr.DataArray to a string. + # If so, use np.datetime_as_string instead. + assert " " not in str(item) + except AssertionError: + # Convert datetime-like item to ISO 8601 + # string format like YYYY-MM-DDThh:mm:ss.ffffff. + value[index] = np.datetime_as_string( + np.asarray(item, dtype=np.datetime64) + ) + newvalue = separators[fmt].join(f"{item}" for item in value) + # Changes in bound.arguments will reflect in bound.args + # and bound.kwargs. + if arg in bound.arguments: + bound.arguments[arg] = newvalue + elif arg in bound.arguments.get("kwargs"): + bound.arguments["kwargs"][arg] = newvalue + # Execute the original function and return its output - return module_func(*args, **kwargs) + return module_func(*bound.args, **bound.kwargs) return new_module diff --git a/pygmt/src/subplot.py b/pygmt/src/subplot.py index 6fcc9e4a00b..563bd655a96 100644 --- a/pygmt/src/subplot.py +++ b/pygmt/src/subplot.py @@ -8,7 +8,6 @@ from pygmt.helpers import ( build_arg_string, fmt_docstring, - is_nonstr_iter, kwargs_to_strings, use_alias, ) @@ -172,6 +171,7 @@ def subplot(self, nrows=1, ncols=1, **kwargs): @fmt_docstring @contextlib.contextmanager @use_alias(A="fixedlabel", C="clearance", V="verbose") +@kwargs_to_strings(panel="sequence_comma") def set_panel(self, panel=None, **kwargs): r""" Set the current subplot panel to plot on. @@ -221,8 +221,6 @@ def set_panel(self, panel=None, **kwargs): {verbose} """ kwargs = self._preprocess(**kwargs) - # convert tuple or list to comma-separated str - panel = ",".join(map(str, panel)) if is_nonstr_iter(panel) else panel with Session() as lib: arg_str = " ".join(["set", f"{panel}", build_arg_string(kwargs)]) diff --git a/pygmt/src/timestamp.py b/pygmt/src/timestamp.py index 1fe65c1bc4e..6ed3ee39fb4 100644 --- a/pygmt/src/timestamp.py +++ b/pygmt/src/timestamp.py @@ -5,11 +5,12 @@ from packaging.version import Version from pygmt.clib import Session, __gmt_version__ -from pygmt.helpers import build_arg_string, is_nonstr_iter +from pygmt.helpers import build_arg_string, kwargs_to_strings __doctest_skip__ = ["timestamp"] +@kwargs_to_strings(offset="sequence") def timestamp( self, text=None, @@ -83,14 +84,11 @@ def timestamp( kwdict["U"] += f"{label}" kwdict["U"] += f"+j{justification}" - if is_nonstr_iter(offset): # given a tuple - kwdict["U"] += "+o" + "/".join(f"{item}" for item in offset) - elif "/" not in offset and Version(__gmt_version__) <= Version("6.4.0"): + if Version(__gmt_version__) <= Version("6.4.0") and "/" not in str(offset): # Giving a single offset doesn't work in GMT <= 6.4.0. # See https://github.com/GenericMappingTools/gmt/issues/7107. - kwdict["U"] += f"+o{offset}/{offset}" - else: - kwdict["U"] += f"+o{offset}" + offset = f"{offset}/{offset}" + kwdict["U"] += f"+o{offset}" # The +t modifier was added in GMT 6.5.0. # See https://github.com/GenericMappingTools/gmt/pull/7127. diff --git a/pygmt/src/which.py b/pygmt/src/which.py index 3fdcafa91f5..d20059f0492 100644 --- a/pygmt/src/which.py +++ b/pygmt/src/which.py @@ -6,13 +6,14 @@ GMTTempFile, build_arg_string, fmt_docstring, - is_nonstr_iter, + kwargs_to_strings, use_alias, ) @fmt_docstring @use_alias(G="download", V="verbose") +@kwargs_to_strings(fname="sequence_space") def which(fname, **kwargs): r""" Find the full path to specified files. @@ -62,9 +63,6 @@ def which(fname, **kwargs): FileNotFoundError If the file is not found. """ - if is_nonstr_iter(fname): # Got a list of files - fname = " ".join(fname) - with GMTTempFile() as tmpfile: with Session() as lib: lib.call_module(