Skip to content

Commit e887777

Browse files
committed
also allow overriding the dtype by variable
1 parent ba8f77c commit e887777

File tree

2 files changed

+22
-3
lines changed

2 files changed

+22
-3
lines changed

xarray/core/common.py

Lines changed: 15 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1327,8 +1327,9 @@ def full_like(other, fill_value, dtype: DTypeLike = None):
13271327
Value to fill the new object with before returning it. If
13281328
other is a Dataset, may also be a dict-like mapping data
13291329
variables to fill values.
1330-
dtype : dtype, optional
1331-
dtype of the new array. If omitted, it defaults to other.dtype.
1330+
dtype : dtype or dict-like of dtype, optional
1331+
dtype of the new array. If a dict-like, maps dtypes to
1332+
variables. If omitted, it defaults to other.dtype.
13321333
13331334
Returns
13341335
-------
@@ -1407,6 +1408,14 @@ def full_like(other, fill_value, dtype: DTypeLike = None):
14071408
Data variables:
14081409
a (x) int64 1 1 1
14091410
b (x) int64 2 2 2
1411+
>>> xr.full_like(ds, fill_value={"a": 1, "b": 2}, dtype={"a": bool, "b": float})
1412+
<xarray.Dataset>
1413+
Dimensions: (x: 3)
1414+
Coordinates:
1415+
* x (x) int64 2 4 6
1416+
Data variables:
1417+
a (x) bool True True True
1418+
b (x) float64 2.0 2.0 2.0
14101419
14111420
See also
14121421
--------
@@ -1430,8 +1439,11 @@ def full_like(other, fill_value, dtype: DTypeLike = None):
14301439
if not isinstance(fill_value, dict):
14311440
fill_value = {k: fill_value for k in other.data_vars.keys()}
14321441

1442+
if not isinstance(dtype, dict):
1443+
dtype = {k: dtype for k in other.data_vars.keys()}
1444+
14331445
data_vars = {
1434-
k: _full_like_variable(v, fill_value.get(k, dtypes.NA), dtype)
1446+
k: _full_like_variable(v, fill_value.get(k, dtypes.NA), dtype.get(k, None))
14351447
for k, v in other.data_vars.items()
14361448
}
14371449
return Dataset(data_vars, coords=other.coords, attrs=other.attrs)

xarray/tests/test_dataset.py

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5461,6 +5461,13 @@ def test_full_like(self):
54615461
assert expected["d2"].dtype == float
54625462
assert_identical(expected, actual)
54635463

5464+
# override multiple dtypes
5465+
actual = full_like(ds, fill_value={"d1": 1, "d2": 2.3}, dtype={"d1": bool})
5466+
expected = ds.assign(d1=("x", [True, True, True]), d2=("y", [2.3, 2.3, 2.3]))
5467+
assert expected["d1"].dtype == bool
5468+
assert expected["d2"].dtype == float
5469+
assert_identical(expected, actual)
5470+
54645471
def test_combine_first(self):
54655472
dsx0 = DataArray([0, 0], [("x", ["a", "b"])]).to_dataset(name="dsx0")
54665473
dsx1 = DataArray([1, 1], [("x", ["b", "c"])]).to_dataset(name="dsx1")

0 commit comments

Comments
 (0)