Skip to content

Commit acd6bbc

Browse files
rcaneillrcaneill
committed
add check_func for attrs, add ndict for names
1 parent 1f41616 commit acd6bbc

File tree

5 files changed

+142
-42
lines changed

5 files changed

+142
-42
lines changed

‎gsw_xarray/_attributes.py

Lines changed: 0 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -1,192 +1,156 @@
11
_func_attrs = {
22
"CT_first_derivatives": (
33
{
4-
"name": "CT_SA",
54
"units": "K (g/kg)^-1",
65
},
76
{
8-
"name": "CT_pt",
97
"units": "1",
108
},
119
),
1210
"CT_from_enthalpy": {
13-
"name": "CT",
1411
"standard_name": "sea_water_conservative_temperature",
1512
"units": "degC",
1613
"reference_scale": "ITS-90",
1714
},
1815
"CT_from_enthalpy_exact": {
19-
"name": "CT",
2016
"standard_name": "sea_water_conservative_temperature",
2117
"units": "degC",
2218
"reference_scale": "ITS-90",
2319
},
2420
"CT_from_entropy": {
25-
"name": "CT",
2621
"standard_name": "sea_water_conservative_temperature",
2722
"units": "degC",
2823
"reference_scale": "ITS-90",
2924
},
3025
"CT_from_pt": {
31-
"name": "CT",
3226
"standard_name": "sea_water_conservative_temperature",
3327
"units": "degC",
3428
"reference_scale": "ITS-90",
3529
},
3630
"CT_from_rho": {
37-
"name": "CT",
3831
"standard_name": "sea_water_conservative_temperature",
3932
"units": "degC",
4033
"reference_scale": "ITS-90",
4134
},
4235
"CT_from_t": {
43-
"name": "CT",
4436
"standard_name": "sea_water_conservative_temperature",
4537
"units": "degC",
4638
"reference_scale": "ITS-90",
4739
},
4840
"C_from_SP": {
49-
"name": "C",
5041
"standard_name": "sea_water_electrical_conductivity",
5142
"units": "mS cm^-1",
5243
},
5344
"SA_from_SP": {
54-
"name": "SA",
5545
"standard_name": "sea_water_absolute_salinity",
5646
"units": "g/kg",
5747
},
5848
"SA_from_SP_Baltic": {
59-
"name": "SA",
6049
"standard_name": "sea_water_absolute_salinity",
6150
"units": "g/kg",
6251
},
6352
"SA_from_Sstar": {
64-
"name": "SA",
6553
"standard_name": "sea_water_absolute_salinity",
6654
"units": "g/kg",
6755
},
6856
"SA_from_rho": {
69-
"name": "SA",
7057
"standard_name": "sea_water_absolute_salinity",
7158
"units": "g/kg",
7259
},
7360
"SP_from_C": {
74-
"name": "SP",
7561
"standard_name": "sea_water_practical_salinity",
7662
"units": "1",
7763
},
7864
"SP_from_SA": {
79-
"name": "SP",
8065
"standard_name": "sea_water_practical_salinity",
8166
"units": "1",
8267
},
8368
"SP_from_SA_Baltic": {
84-
"name": "SP",
8569
"standard_name": "sea_water_practical_salinity",
8670
"units": "1",
8771
},
8872
"SP_from_SK": {
89-
"name": "SP",
9073
"standard_name": "sea_water_practical_salinity",
9174
"units": "1",
9275
},
9376
"SP_from_SR": {
94-
"name": "SP",
9577
"standard_name": "sea_water_practical_salinity",
9678
"units": "1",
9779
},
9880
"SP_from_Sstar": {
99-
"name": "SP",
10081
"standard_name": "sea_water_practical_salinity",
10182
"units": "1",
10283
},
10384
"SP_salinometer": {
104-
"name": "SP",
10585
"standard_name": "sea_water_practical_salinity",
10686
"units": "1",
10787
},
10888
"SR_from_SP": {
109-
"name": "SR",
11089
"standard_name": "sea_water_reference_salinity",
11190
"units": "g/kg",
11291
},
11392
"Sstar_from_SA": {
114-
"name": "Sstar",
11593
"standard_name": "sea_water_preformed_salinity",
11694
"units": "g/kg",
11795
},
11896
"Sstar_from_SP": {
119-
"name": "Sstar",
12097
"standard_name": "sea_water_preformed_salinity",
12198
"units": "g/kg",
12299
},
123100
"f": {
124-
"name": "f",
125101
"standard_name": "coriolis_parameter",
126102
"units": "s^-1",
127103
},
128104
"p_from_z": {
129-
"name": "p",
130105
"standard_name": "sea_water_pressure",
131106
"units": "dbar",
132107
},
133108
"pt0_from_t": {
134-
"name": "pt0",
135109
"standard_name": "sea_water_potential_temperature",
136110
"units": "degC",
137111
"reference_scale": "ITS-90",
138112
},
139113
"pt_from_CT": {
140-
"name": "pt",
141114
"standard_name": "sea_water_potential_temperature",
142115
"units": "degC",
143116
"reference_scale": "ITS-90",
144117
},
145118
"pt_from_entropy": {
146-
"name": "pt",
147119
"standard_name": "sea_water_potential_temperature",
148120
"units": "degC",
149121
"reference_scale": "ITS-90",
150122
},
151123
"pt_from_t": {
152-
"name": "pt",
153124
"standard_name": "sea_water_potential_temperature",
154125
"units": "degC",
155126
"reference_scale": "ITS-90",
156127
},
157128
"rho": {
158-
"name": "rho",
159129
"standard_name": "sea_water_density",
160130
"units": "kg m^-3",
161131
},
162132
"rho_t_exact": {
163-
"name": "rho",
164133
"standard_name": "sea_water_density",
165134
"units": "kg m^-3",
166135
},
167136
"sigma0": {
168-
"name": "sigma0",
169137
"standard_name": "sea_water_sigma_t",
170138
"units": "kg m^-3",
171139
},
172140
"sound_speed": {
173-
"name": "sound_speed",
174141
"standard_name": "speed_of_sound_in_sea_water",
175142
"units": "m s^-1",
176143
},
177144
"sound_speed_t_exact": {
178-
"name": "sound_speed",
179145
"standard_name": "speed_of_sound_in_sea_water",
180146
"units": "m s^-1",
181147
},
182148
"t_from_CT": {
183-
"name": "temperature",
184149
"standard_name": "sea_water_temperature",
185150
"units": "degC",
186151
"reference_scale": "ITS-90",
187152
},
188153
"z_from_p": {
189-
"name": "z",
190154
"standard_name": "height_above_mean_sea_level",
191155
"units": "m",
192156
},

‎gsw_xarray/_check_funcs.py

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,44 @@
1+
def check_z_from_p(attrs, args, kwargs):
2+
"""
3+
If the 2 optional arguments are not 0, removes the standard name.
4+
5+
Reminder of the function call:
6+
gsw.z_from_p(p, lat, geo_strf_dyn_height=0, sea_surface_geopotential=0)
7+
8+
Parameters
9+
----------
10+
args and kwargs : list and dict
11+
the arguments passed to gsw.z_from_p
12+
attrs : dict
13+
the attributes associated with the function z_from_p
14+
(e.g. standard_name, units, etc)
15+
"""
16+
# Getting the values of optional parameters
17+
# If the values are in kwargs, all good
18+
# Otherwise, we get them from args
19+
try:
20+
geo_strf_dyn_height = kwargs["geo_strf_dyn_height"]
21+
except KeyError:
22+
try:
23+
geo_strf_dyn_height = args[2]
24+
except IndexError:
25+
geo_strf_dyn_height = 0 # Could use inspect.signature to get default value
26+
27+
try:
28+
sea_surface_geopotential = kwargs["sea_surface_geopotential"]
29+
except KeyError:
30+
try:
31+
sea_surface_geopotential = args[3]
32+
except IndexError:
33+
sea_surface_geopotential = (
34+
0 # Could use inspect.signature to get default value
35+
)
36+
37+
if geo_strf_dyn_height != 0 or sea_surface_geopotential != 0:
38+
attrs = attrs.copy()
39+
# necessary to use a copy, or it will pop the item from the original dict
40+
attrs.pop("standard_name")
41+
return attrs
42+
43+
44+
_check_funcs = {"z_from_p": check_z_from_p}

‎gsw_xarray/_core.py

Lines changed: 12 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -4,24 +4,26 @@
44
import xarray as xr
55

66
from ._attributes import _func_attrs
7+
from ._names import _names
8+
from ._check_funcs import _check_funcs
79

810

9-
def add_attrs(rv, attrs):
10-
name = attrs.pop("name")
11+
def add_attrs(rv, attrs, name):
1112
rv.name = name
1213
rv.attrs = attrs
1314

1415

15-
def cf_attrs(attrs, extra=None):
16+
def cf_attrs(attrs, name, check_func):
1617
def cf_attrs_decorator(func):
1718
@wraps(func)
1819
def cf_attrs_wrapper(*args, **kwargs):
1920
rv = func(*args, **kwargs)
21+
attrs_checked = check_func(attrs, args, kwargs)
2022
if isinstance(rv, tuple):
2123
for (i, da) in enumerate(rv):
22-
add_attrs(da, attrs[i])
24+
add_attrs(da, attrs_checked[i], name[i])
2325
elif isinstance(rv, xr.DataArray):
24-
add_attrs(rv, attrs)
26+
add_attrs(rv, attrs_checked, name)
2527
return rv
2628

2729
return cf_attrs_wrapper
@@ -32,7 +34,11 @@ def cf_attrs_wrapper(*args, **kwargs):
3234
def _init_funcs():
3335
_wrapped_funcs = {}
3436
for func in _func_attrs.keys():
35-
_wrapped_funcs[func] = cf_attrs(_func_attrs[func])(getattr(gsw, func))
37+
_wrapped_funcs[func] = cf_attrs(
38+
_func_attrs[func],
39+
_names[func],
40+
_check_funcs.get(func, lambda attrs, *args, **kwargs: attrs),
41+
)(getattr(gsw, func))
3642
return _wrapped_funcs
3743

3844

‎gsw_xarray/_names.py

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,37 @@
1+
_names = {
2+
"CT_first_derivatives": ("CT_SA", "CT_pt"),
3+
"CT_from_enthalpy": "CT",
4+
"CT_from_enthalpy_exact": "CT",
5+
"CT_from_entropy": "CT",
6+
"CT_from_pt": "CT",
7+
"CT_from_rho": "CT",
8+
"CT_from_t": "CT",
9+
"C_from_SP": "C",
10+
"SA_from_SP": "SA",
11+
"SA_from_SP_Baltic": "SA",
12+
"SA_from_Sstar": "SA",
13+
"SA_from_rho": "SA",
14+
"SP_from_C": "SP",
15+
"SP_from_SA": "SP",
16+
"SP_from_SA_Baltic": "SP",
17+
"SP_from_SK": "SP",
18+
"SP_from_SR": "SP",
19+
"SP_from_Sstar": "SP",
20+
"SP_salinometer": "SP",
21+
"SR_from_SP": "SR",
22+
"Sstar_from_SA": "Sstar",
23+
"Sstar_from_SP": "Sstar",
24+
"f": "f",
25+
"p_from_z": "p",
26+
"pt0_from_t": "pt0",
27+
"pt_from_CT": "pt",
28+
"pt_from_entropy": "pt",
29+
"pt_from_t": "pt",
30+
"rho": "rho",
31+
"rho_t_exact": "rho",
32+
"sigma0": "sigma0",
33+
"sound_speed": "sound_speed",
34+
"sound_speed_t_exact": "sound_speed",
35+
"t_from_CT": "temperature",
36+
"z_from_p": "z",
37+
}

‎gsw_xarray/tests/test_gsw_xarray.py

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,7 @@
11
from gsw_xarray import __version__
22
import gsw_xarray as gsw
3+
import xarray as xr
4+
import pytest
35

46

57
def test_version():
@@ -18,3 +20,50 @@ def test_func_return_tuple(ds):
1820
(CT_SA, CT_pt) = gsw.CT_first_derivatives(ds.SA, 1)
1921
assert CT_SA.name == "CT_SA"
2022
assert CT_SA.attrs["units"] == "K (g/kg)^-1"
23+
24+
25+
@pytest.mark.parametrize("gsdh", [0, 1, None])
26+
@pytest.mark.parametrize("ssg", [0, 1, None])
27+
@pytest.mark.parametrize(
28+
"use_kw",
29+
[
30+
[
31+
False,
32+
]
33+
* (4 - i)
34+
+ [
35+
True,
36+
]
37+
* i
38+
for i in range(5)
39+
],
40+
)
41+
def test_check_func(gsdh, ssg, use_kw):
42+
ds = xr.Dataset()
43+
ds["p"] = xr.DataArray([0, 1, 2])
44+
ds["lat"] = xr.DataArray([40, 41, 21])
45+
args = []
46+
kwargs = {}
47+
if use_kw[0]:
48+
kwargs["p"] = ds["p"]
49+
else:
50+
args.append(ds["p"])
51+
if use_kw[1]:
52+
kwargs["lat"] = ds["lat"]
53+
else:
54+
args.append(ds["lat"])
55+
if gsdh is not None:
56+
if use_kw[2]:
57+
kwargs["geo_strf_dyn_height"] = gsdh
58+
else:
59+
args.append(gsdh)
60+
if ssg is not None:
61+
if use_kw[3]:
62+
kwargs["sea_surface_geopotential"] = ssg
63+
else:
64+
args.append(ssg)
65+
z = gsw.z_from_p(*args, **kwargs)
66+
if gsdh not in (0, None) or ssg not in (0, None):
67+
assert "standard_name" not in z.attrs.keys()
68+
else:
69+
assert "standard_name" in z.attrs.keys()

0 commit comments

Comments
 (0)