Skip to content

Commit 17f37b7

Browse files
authored
Don't apply mappers to DataArrays (#227)
1 parent 24bf4d3 commit 17f37b7

File tree

4 files changed

+45
-10
lines changed

4 files changed

+45
-10
lines changed

cf_xarray/accessor.py

Lines changed: 15 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,12 @@
2626

2727
from .criteria import coordinate_criteria, regex
2828
from .helpers import bounds_to_vertices
29-
from .utils import _is_datetime_like, invert_mappings, parse_cell_methods_attr
29+
from .utils import (
30+
_is_datetime_like,
31+
always_iterable,
32+
invert_mappings,
33+
parse_cell_methods_attr,
34+
)
3035

3136
#: Classes wrapped by cf_xarray.
3237
_WRAPPED_CLASSES = (
@@ -68,7 +73,7 @@
6873
def apply_mapper(
6974
mappers: Union[Mapper, Tuple[Mapper, ...]],
7075
obj: Union[DataArray, Dataset],
71-
key: str,
76+
key: Any,
7277
error: bool = True,
7378
default: Any = None,
7479
) -> List[Any]:
@@ -79,8 +84,13 @@ def apply_mapper(
7984
It should return a list in all other cases including when there are no
8085
results for a good key.
8186
"""
82-
if default is None:
83-
default = []
87+
88+
if not isinstance(key, str):
89+
if default is None:
90+
raise ValueError("`default` must be provided when `key` is not a string.")
91+
return list(always_iterable(default))
92+
93+
default = [] if default is None else list(always_iterable(default))
8494

8595
def _apply_single_mapper(mapper):
8696

@@ -917,8 +927,7 @@ def _rewrite_values(
917927
value = kwargs[key]
918928
mappers = all_mappers[key]
919929

920-
if isinstance(value, str):
921-
value = [value]
930+
value = always_iterable(value)
922931

923932
if isinstance(value, dict):
924933
# this for things like isel where **kwargs captures things like T=5

cf_xarray/tests/test_accessor.py

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -317,8 +317,13 @@ def test_weighted(obj):
317317
with raise_if_dask_computes(max_computes=2):
318318
# weights are checked for nans
319319
expected = obj.weighted(obj["cell_area"]).sum("lat")
320-
actual = obj.cf.weighted("area").sum("Y")
321-
assert_identical(expected, actual)
320+
actuals = [
321+
obj.cf.weighted("area").sum("Y"),
322+
obj.cf.weighted(obj["cell_area"]).sum("Y"),
323+
obj.cf.weighted(weights=obj["cell_area"]).sum("Y"),
324+
]
325+
for actual in actuals:
326+
assert_identical(expected, actual)
322327

323328

324329
@pytest.mark.parametrize("obj", objects)

cf_xarray/utils.py

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
from collections import defaultdict
2-
from typing import Dict
2+
from typing import Any, Dict, Iterable
33

44
from xarray import DataArray
55

@@ -53,3 +53,7 @@ def invert_mappings(*mappings):
5353
for name in v:
5454
merged[name] |= {k}
5555
return merged
56+
57+
58+
def always_iterable(obj: Any) -> Iterable:
59+
return [obj] if not isinstance(obj, (tuple, list, set, dict)) else obj

doc/examples/introduction.ipynb

Lines changed: 18 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,22 @@
2525
"import xarray as xr"
2626
]
2727
},
28+
{
29+
"cell_type": "markdown",
30+
"metadata": {},
31+
"source": [
32+
"`cf_xarray` works best when `xarray` keeps attributes by default.\n"
33+
]
34+
},
35+
{
36+
"cell_type": "code",
37+
"execution_count": null,
38+
"metadata": {},
39+
"outputs": [],
40+
"source": [
41+
"xr.set_options(keep_attrs=True)"
42+
]
43+
},
2844
{
2945
"cell_type": "markdown",
3046
"metadata": {},
@@ -914,6 +930,7 @@
914930
" * 110e3\n",
915931
")\n",
916932
"# and set proper attributes\n",
933+
"ds[\"cell_area\"].attrs = dict(standard_name=\"cell_area\", units=\"m2\")\n",
917934
"ds.air.attrs[\"cell_measures\"] = \"area: cell_area\""
918935
]
919936
},
@@ -1000,7 +1017,7 @@
10001017
"name": "python",
10011018
"nbconvert_exporter": "python",
10021019
"pygments_lexer": "ipython3",
1003-
"version": "3.9.1"
1020+
"version": "3.8.10"
10041021
},
10051022
"toc": {
10061023
"base_numbering": 1,

0 commit comments

Comments
 (0)