Skip to content

Commit cb3663d

Browse files
owenlittlejohnsowenlittlejohns
authored
Migrate datatree io.py and common.py into xarray/core (#9011)
1 parent d9e4de6 commit cb3663d

File tree

9 files changed

+150
-152
lines changed

9 files changed

+150
-152
lines changed

‎doc/whats-new.rst

Lines changed: 8 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@ What's New
1717
1818
.. _whats-new.2024.05.1:
1919

20-
v2024.05.1 (unreleased)
20+
v2024.06 (unreleased)
2121
-----------------------
2222

2323
New Features
@@ -59,6 +59,10 @@ Documentation
5959

6060
Internal Changes
6161
~~~~~~~~~~~~~~~~
62+
- Migrates remainder of ``io.py`` to ``xarray/core/datatree_io.py`` and
63+
``TreeAttrAccessMixin`` into ``xarray/core/common.py`` (:pull: `9011`)
64+
By `Owen Littlejohns <https://github.com/owenlittlejohns>`_ and
65+
`Tom Nicholas <https://github.com/TomNicholas>`_.
6266

6367

6468
.. _whats-new.2024.05.0:
@@ -141,10 +145,9 @@ Internal Changes
141145
By `Owen Littlejohns <https://github.com/owenlittlejohns>`_, `Matt Savoie
142146
<https://github.com/flamingbear>`_ and `Tom Nicholas <https://github.com/TomNicholas>`_.
143147
- ``transpose``, ``set_dims``, ``stack`` & ``unstack`` now use a ``dim`` kwarg
144-
rather than ``dims`` or ``dimensions``. This is the final change to unify
145-
xarray functions to use ``dim``. Using the existing kwarg will raise a
146-
warning.
147-
By `Maximilian Roos <https://github.com/max-sixty>`_
148+
rather than ``dims`` or ``dimensions``. This is the final change to make xarray methods
149+
consistent with their use of ``dim``. Using the existing kwarg will raise a
150+
warning. By `Maximilian Roos <https://github.com/max-sixty>`_
148151

149152
.. _whats-new.2024.03.0:
150153

‎xarray/backends/api.py

Lines changed: 9 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,7 @@
3636
from xarray.core.dataarray import DataArray
3737
from xarray.core.dataset import Dataset, _get_chunk, _maybe_chunk
3838
from xarray.core.indexes import Index
39-
from xarray.core.types import ZarrWriteModes
39+
from xarray.core.types import NetcdfWriteModes, ZarrWriteModes
4040
from xarray.core.utils import is_remote_uri
4141
from xarray.namedarray.daskmanager import DaskManager
4242
from xarray.namedarray.parallelcompat import guess_chunkmanager
@@ -1120,7 +1120,7 @@ def open_mfdataset(
11201120
def to_netcdf(
11211121
dataset: Dataset,
11221122
path_or_file: str | os.PathLike | None = None,
1123-
mode: Literal["w", "a"] = "w",
1123+
mode: NetcdfWriteModes = "w",
11241124
format: T_NetcdfTypes | None = None,
11251125
group: str | None = None,
11261126
engine: T_NetcdfEngine | None = None,
@@ -1138,7 +1138,7 @@ def to_netcdf(
11381138
def to_netcdf(
11391139
dataset: Dataset,
11401140
path_or_file: None = None,
1141-
mode: Literal["w", "a"] = "w",
1141+
mode: NetcdfWriteModes = "w",
11421142
format: T_NetcdfTypes | None = None,
11431143
group: str | None = None,
11441144
engine: T_NetcdfEngine | None = None,
@@ -1155,7 +1155,7 @@ def to_netcdf(
11551155
def to_netcdf(
11561156
dataset: Dataset,
11571157
path_or_file: str | os.PathLike,
1158-
mode: Literal["w", "a"] = "w",
1158+
mode: NetcdfWriteModes = "w",
11591159
format: T_NetcdfTypes | None = None,
11601160
group: str | None = None,
11611161
engine: T_NetcdfEngine | None = None,
@@ -1173,7 +1173,7 @@ def to_netcdf(
11731173
def to_netcdf(
11741174
dataset: Dataset,
11751175
path_or_file: str | os.PathLike,
1176-
mode: Literal["w", "a"] = "w",
1176+
mode: NetcdfWriteModes = "w",
11771177
format: T_NetcdfTypes | None = None,
11781178
group: str | None = None,
11791179
engine: T_NetcdfEngine | None = None,
@@ -1191,7 +1191,7 @@ def to_netcdf(
11911191
def to_netcdf(
11921192
dataset: Dataset,
11931193
path_or_file: str | os.PathLike,
1194-
mode: Literal["w", "a"] = "w",
1194+
mode: NetcdfWriteModes = "w",
11951195
format: T_NetcdfTypes | None = None,
11961196
group: str | None = None,
11971197
engine: T_NetcdfEngine | None = None,
@@ -1209,7 +1209,7 @@ def to_netcdf(
12091209
def to_netcdf(
12101210
dataset: Dataset,
12111211
path_or_file: str | os.PathLike,
1212-
mode: Literal["w", "a"] = "w",
1212+
mode: NetcdfWriteModes = "w",
12131213
format: T_NetcdfTypes | None = None,
12141214
group: str | None = None,
12151215
engine: T_NetcdfEngine | None = None,
@@ -1226,7 +1226,7 @@ def to_netcdf(
12261226
def to_netcdf(
12271227
dataset: Dataset,
12281228
path_or_file: str | os.PathLike | None,
1229-
mode: Literal["w", "a"] = "w",
1229+
mode: NetcdfWriteModes = "w",
12301230
format: T_NetcdfTypes | None = None,
12311231
group: str | None = None,
12321232
engine: T_NetcdfEngine | None = None,
@@ -1241,7 +1241,7 @@ def to_netcdf(
12411241
def to_netcdf(
12421242
dataset: Dataset,
12431243
path_or_file: str | os.PathLike | None = None,
1244-
mode: Literal["w", "a"] = "w",
1244+
mode: NetcdfWriteModes = "w",
12451245
format: T_NetcdfTypes | None = None,
12461246
group: str | None = None,
12471247
engine: T_NetcdfEngine | None = None,

‎xarray/core/common.py

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -347,6 +347,24 @@ def _ipython_key_completions_(self) -> list[str]:
347347
return list(items)
348348

349349

350+
class TreeAttrAccessMixin(AttrAccessMixin):
351+
"""Mixin class that allows getting keys with attribute access"""
352+
353+
# TODO: Ensure ipython tab completion can include both child datatrees and
354+
# variables from Dataset objects on relevant nodes.
355+
356+
__slots__ = ()
357+
358+
def __init_subclass__(cls, **kwargs):
359+
"""This method overrides the check from ``AttrAccessMixin`` that ensures
360+
``__dict__`` is absent in a class, with ``__slots__`` used instead.
361+
``DataTree`` has some dynamically defined attributes in addition to those
362+
defined in ``__slots__``. (GH9068)
363+
"""
364+
if not hasattr(object.__new__(cls), "__dict__"):
365+
pass
366+
367+
350368
def get_squeeze_dims(
351369
xarray_obj,
352370
dim: Hashable | Iterable[Hashable] | None = None,

‎xarray/core/dataarray.py

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,7 @@
5454
from xarray.core.options import OPTIONS, _get_keep_attrs
5555
from xarray.core.types import (
5656
DaCompatible,
57+
NetcdfWriteModes,
5758
T_DataArray,
5859
T_DataArrayOrSet,
5960
ZarrWriteModes,
@@ -3945,7 +3946,7 @@ def to_masked_array(self, copy: bool = True) -> np.ma.MaskedArray:
39453946
def to_netcdf(
39463947
self,
39473948
path: None = None,
3948-
mode: Literal["w", "a"] = "w",
3949+
mode: NetcdfWriteModes = "w",
39493950
format: T_NetcdfTypes | None = None,
39503951
group: str | None = None,
39513952
engine: T_NetcdfEngine | None = None,
@@ -3960,7 +3961,7 @@ def to_netcdf(
39603961
def to_netcdf(
39613962
self,
39623963
path: str | PathLike,
3963-
mode: Literal["w", "a"] = "w",
3964+
mode: NetcdfWriteModes = "w",
39643965
format: T_NetcdfTypes | None = None,
39653966
group: str | None = None,
39663967
engine: T_NetcdfEngine | None = None,
@@ -3976,7 +3977,7 @@ def to_netcdf(
39763977
def to_netcdf(
39773978
self,
39783979
path: str | PathLike,
3979-
mode: Literal["w", "a"] = "w",
3980+
mode: NetcdfWriteModes = "w",
39803981
format: T_NetcdfTypes | None = None,
39813982
group: str | None = None,
39823983
engine: T_NetcdfEngine | None = None,
@@ -3992,7 +3993,7 @@ def to_netcdf(
39923993
def to_netcdf(
39933994
self,
39943995
path: str | PathLike,
3995-
mode: Literal["w", "a"] = "w",
3996+
mode: NetcdfWriteModes = "w",
39963997
format: T_NetcdfTypes | None = None,
39973998
group: str | None = None,
39983999
engine: T_NetcdfEngine | None = None,
@@ -4005,7 +4006,7 @@ def to_netcdf(
40054006
def to_netcdf(
40064007
self,
40074008
path: str | PathLike | None = None,
4008-
mode: Literal["w", "a"] = "w",
4009+
mode: NetcdfWriteModes = "w",
40094010
format: T_NetcdfTypes | None = None,
40104011
group: str | None = None,
40114012
engine: T_NetcdfEngine | None = None,

‎xarray/core/dataset.py

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -88,6 +88,7 @@
8888
from xarray.core.missing import get_clean_interp_index
8989
from xarray.core.options import OPTIONS, _get_keep_attrs
9090
from xarray.core.types import (
91+
NetcdfWriteModes,
9192
QuantileMethods,
9293
Self,
9394
T_ChunkDim,
@@ -2171,7 +2172,7 @@ def dump_to_store(self, store: AbstractDataStore, **kwargs) -> None:
21712172
def to_netcdf(
21722173
self,
21732174
path: None = None,
2174-
mode: Literal["w", "a"] = "w",
2175+
mode: NetcdfWriteModes = "w",
21752176
format: T_NetcdfTypes | None = None,
21762177
group: str | None = None,
21772178
engine: T_NetcdfEngine | None = None,
@@ -2186,7 +2187,7 @@ def to_netcdf(
21862187
def to_netcdf(
21872188
self,
21882189
path: str | PathLike,
2189-
mode: Literal["w", "a"] = "w",
2190+
mode: NetcdfWriteModes = "w",
21902191
format: T_NetcdfTypes | None = None,
21912192
group: str | None = None,
21922193
engine: T_NetcdfEngine | None = None,
@@ -2202,7 +2203,7 @@ def to_netcdf(
22022203
def to_netcdf(
22032204
self,
22042205
path: str | PathLike,
2205-
mode: Literal["w", "a"] = "w",
2206+
mode: NetcdfWriteModes = "w",
22062207
format: T_NetcdfTypes | None = None,
22072208
group: str | None = None,
22082209
engine: T_NetcdfEngine | None = None,
@@ -2218,7 +2219,7 @@ def to_netcdf(
22182219
def to_netcdf(
22192220
self,
22202221
path: str | PathLike,
2221-
mode: Literal["w", "a"] = "w",
2222+
mode: NetcdfWriteModes = "w",
22222223
format: T_NetcdfTypes | None = None,
22232224
group: str | None = None,
22242225
engine: T_NetcdfEngine | None = None,
@@ -2231,7 +2232,7 @@ def to_netcdf(
22312232
def to_netcdf(
22322233
self,
22332234
path: str | PathLike | None = None,
2234-
mode: Literal["w", "a"] = "w",
2235+
mode: NetcdfWriteModes = "w",
22352236
format: T_NetcdfTypes | None = None,
22362237
group: str | None = None,
22372238
engine: T_NetcdfEngine | None = None,

‎xarray/core/datatree.py

Lines changed: 47 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -9,12 +9,14 @@
99
Any,
1010
Callable,
1111
Generic,
12+
Literal,
1213
NoReturn,
1314
Union,
1415
overload,
1516
)
1617

1718
from xarray.core import utils
19+
from xarray.core.common import TreeAttrAccessMixin
1820
from xarray.core.coordinates import DatasetCoordinates
1921
from xarray.core.dataarray import DataArray
2022
from xarray.core.dataset import Dataset, DataVariables
@@ -46,7 +48,6 @@
4648
maybe_wrap_array,
4749
)
4850
from xarray.core.variable import Variable
49-
from xarray.datatree_.datatree.common import TreeAttrAccessMixin
5051

5152
try:
5253
from xarray.core.variable import calculate_dimensions
@@ -57,8 +58,9 @@
5758
if TYPE_CHECKING:
5859
import pandas as pd
5960

61+
from xarray.core.datatree_io import T_DataTreeNetcdfEngine, T_DataTreeNetcdfTypes
6062
from xarray.core.merge import CoercibleValue
61-
from xarray.core.types import ErrorOptions
63+
from xarray.core.types import ErrorOptions, NetcdfWriteModes, ZarrWriteModes
6264

6365
# """
6466
# DEVELOPERS' NOTE
@@ -1475,7 +1477,16 @@ def groups(self):
14751477
return tuple(node.path for node in self.subtree)
14761478

14771479
def to_netcdf(
1478-
self, filepath, mode: str = "w", encoding=None, unlimited_dims=None, **kwargs
1480+
self,
1481+
filepath,
1482+
mode: NetcdfWriteModes = "w",
1483+
encoding=None,
1484+
unlimited_dims=None,
1485+
format: T_DataTreeNetcdfTypes | None = None,
1486+
engine: T_DataTreeNetcdfEngine | None = None,
1487+
group: str | None = None,
1488+
compute: bool = True,
1489+
**kwargs,
14791490
):
14801491
"""
14811492
Write datatree contents to a netCDF file.
@@ -1499,26 +1510,47 @@ def to_netcdf(
14991510
By default, no dimensions are treated as unlimited dimensions.
15001511
Note that unlimited_dims may also be set via
15011512
``dataset.encoding["unlimited_dims"]``.
1513+
format : {"NETCDF4", }, optional
1514+
File format for the resulting netCDF file:
1515+
1516+
* NETCDF4: Data is stored in an HDF5 file, using netCDF4 API features.
1517+
engine : {"netcdf4", "h5netcdf"}, optional
1518+
Engine to use when writing netCDF files. If not provided, the
1519+
default engine is chosen based on available dependencies, with a
1520+
preference for "netcdf4" if writing to a file on disk.
1521+
group : str, optional
1522+
Path to the netCDF4 group in the given file to open as the root group
1523+
of the ``DataTree``. Currently, specifying a group is not supported.
1524+
compute : bool, default: True
1525+
If true compute immediately, otherwise return a
1526+
``dask.delayed.Delayed`` object that can be computed later.
1527+
Currently, ``compute=False`` is not supported.
15021528
kwargs :
15031529
Addional keyword arguments to be passed to ``xarray.Dataset.to_netcdf``
15041530
"""
1505-
from xarray.datatree_.datatree.io import _datatree_to_netcdf
1531+
from xarray.core.datatree_io import _datatree_to_netcdf
15061532

15071533
_datatree_to_netcdf(
15081534
self,
15091535
filepath,
15101536
mode=mode,
15111537
encoding=encoding,
15121538
unlimited_dims=unlimited_dims,
1539+
format=format,
1540+
engine=engine,
1541+
group=group,
1542+
compute=compute,
15131543
**kwargs,
15141544
)
15151545

15161546
def to_zarr(
15171547
self,
15181548
store,
1519-
mode: str = "w-",
1549+
mode: ZarrWriteModes = "w-",
15201550
encoding=None,
15211551
consolidated: bool = True,
1552+
group: str | None = None,
1553+
compute: Literal[True] = True,
15221554
**kwargs,
15231555
):
15241556
"""
@@ -1541,17 +1573,26 @@ def to_zarr(
15411573
consolidated : bool
15421574
If True, apply zarr's `consolidate_metadata` function to the store
15431575
after writing metadata for all groups.
1576+
group : str, optional
1577+
Group path. (a.k.a. `path` in zarr terminology.)
1578+
compute : bool, default: True
1579+
If true compute immediately, otherwise return a
1580+
``dask.delayed.Delayed`` object that can be computed later. Metadata
1581+
is always updated eagerly. Currently, ``compute=False`` is not
1582+
supported.
15441583
kwargs :
15451584
Additional keyword arguments to be passed to ``xarray.Dataset.to_zarr``
15461585
"""
1547-
from xarray.datatree_.datatree.io import _datatree_to_zarr
1586+
from xarray.core.datatree_io import _datatree_to_zarr
15481587

15491588
_datatree_to_zarr(
15501589
self,
15511590
store,
15521591
mode=mode,
15531592
encoding=encoding,
15541593
consolidated=consolidated,
1594+
group=group,
1595+
compute=compute,
15551596
**kwargs,
15561597
)
15571598

0 commit comments

Comments
 (0)