Skip to content

Commit

Permalink
Let kwargs_to_strings work with default values and postional arguments
Browse files Browse the repository at this point in the history
  • Loading branch information
seisman committed Nov 22, 2023
1 parent 1ef9fb4 commit f53e5d7
Showing 1 changed file with 70 additions and 32 deletions.
102 changes: 70 additions & 32 deletions pygmt/helpers/decorators.py
Original file line number Diff line number Diff line change
Expand Up @@ -636,7 +636,6 @@ def kwargs_to_strings(**conversions):
Examples
--------
>>> @kwargs_to_strings(
... R="sequence", i="sequence_comma", files="sequence_space"
... )
Expand Down Expand Up @@ -686,58 +685,97 @@ 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": ",",
"sequence_plus": "+",
"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

Expand Down

0 comments on commit f53e5d7

Please sign in to comment.