Skip to content

Commit

Permalink
Add argument name check in preprocess_and_wrap, and fix dewpoint_from…
Browse files Browse the repository at this point in the history
…_specific_humidity
  • Loading branch information
jthielen committed Aug 31, 2021
1 parent 5ab290b commit 472407d
Show file tree
Hide file tree
Showing 3 changed files with 19 additions and 2 deletions.
2 changes: 1 addition & 1 deletion src/metpy/calc/thermo.py
Original file line number Diff line number Diff line change
Expand Up @@ -3053,7 +3053,7 @@ def static_stability(pressure, temperature, vertical_dim=0):
@exporter.export
@preprocess_and_wrap(
wrap_like='temperature',
broadcast=('pressure', 'temperature', 'specific_humdiity')
broadcast=('pressure', 'temperature', 'specific_humidity')
)
@check_units('[pressure]', '[temperature]', '[dimensionless]')
def dewpoint_from_specific_humidity(pressure, temperature, specific_humidity):
Expand Down
11 changes: 10 additions & 1 deletion src/metpy/xarray.py
Original file line number Diff line number Diff line change
Expand Up @@ -1170,9 +1170,18 @@ def preprocess_and_wrap(broadcast=None, wrap_like=None, match_unit=False, to_mag
xarray arguments to Quantity, and do not change other array-like arguments.
"""
def decorator(func):
sig = signature(func)
if broadcast is not None:
for arg_name in broadcast:
if arg_name not in sig.parameters:
raise ValueError(
f'Cannot broadcast argument {arg_name} as it is not in function '
'signature'
)

@functools.wraps(func)
def wrapper(*args, **kwargs):
bound_args = signature(func).bind(*args, **kwargs)
bound_args = sig.bind(*args, **kwargs)

# Auto-broadcast select xarray arguments, and update bound_args
if broadcast is not None:
Expand Down
8 changes: 8 additions & 0 deletions tests/test_xarray.py
Original file line number Diff line number Diff line change
Expand Up @@ -1298,6 +1298,14 @@ def func(a, b):
assert_array_equal(func(data, data2), [[0, 1, 2], [0, 0, 0], [0, 0, 0]] * units('N m'))


def test_preprocess_and_wrap_broadcasting_error():
"""Test that decorator with bad arguments specified to broadcast errors out."""
with pytest.raises(ValueError):
@preprocess_and_wrap(broadcast=('a', 'c'))
def func(a, b):
return a + b


def test_preprocess_and_wrap_with_to_magnitude():
"""Test preprocessing and wrapping with casting to magnitude."""
data = xr.DataArray([1, 0, 1] * units.m)
Expand Down

0 comments on commit 472407d

Please sign in to comment.