DAS-2067 - Migrate datatree io.py and common.py #9011

owenlittlejohns · 2024-05-06T23:37:41Z

This PR migrates what I believe is the last actual source code out of the xarray/datatree_/datatree directory into xarray/core. This is another piece of #8572.

I made a couple of decisions that I'd love some feedback/extra eyes on:

I moved the remaining functions from xarray/datatree_/datatree/io.py into their own module, rather than xarray/backends/api.py. It felt like these were a little separate, but I can definitely be convinced otherwise.
I created a subclass for TreeAttrAccessMixin. I really just wanted to swap over to using AttrAccessMixin, but I ran into issues with __slots__. Looking at it, I think I would need to work on the DataTree, TreeNode and NamedNode classes. I can go down that path, but thought I'd put up this PR for discussion before beginning that effort.

Continues Track merging datatree into xarray #8572 (Contributed to last subitem of "Merge in code for DataTree class.")
~~Tests added~~
User visible changes (including notable bug fixes) are documented in whats-new.rst
~~New functions/methods are listed in api.rst~~

xarray/core/types.py

xarray/core/datatree_io.py

owenlittlejohns · 2024-05-06T23:44:09Z

xarray/core/common.py

@@ -347,6 +347,34 @@ def _ipython_key_completions_(self) -> list[str]:
        return list(items)


+class TreeAttrAccessMixin(AttrAccessMixin):


So here's the biggest question I have... is this okay for now? (Maybe this is directed at @TomNicholas)

The use of __slots__ is a little new to me (but makes sense from reading up on it). I think to adopt __slots__ and remove __dict__ I would need to slightly rework DataTree, TreeNode and NamedNode. I wanted to check that the scope mainly seems to be self.children and self.parent attributes, and understand why those classes were implemented as they currently are before trying to alter things.

This is going to embarrass me, but why did you stop calling init_subclass on the Parent object? everything else with this merge seems reasonable to me.

Good catch. I accidentally commented out that last line. Will fix!

So, my last comment was true, but... uncommenting out that line means that the same method in the parent class (AttrAccessMixin) gets called by super in addition to this method, right? That would mean we end up calling the validation logic that this child method is trying to avoid for now (checking that __dict__ is not present on the class).

Now if only I'd said that the first time around 😉

flamingbear

just comments for now

xarray/core/types.py

flamingbear · 2024-05-07T14:59:36Z

xarray/core/common.py

@@ -347,6 +347,34 @@ def _ipython_key_completions_(self) -> list[str]:
        return list(items)


+class TreeAttrAccessMixin(AttrAccessMixin):


This is going to embarrass me, but why did you stop calling init_subclass on the Parent object? everything else with this merge seems reasonable to me.

xarray/core/datatree_io.py

Illviljan · 2024-05-13T20:21:24Z

xarray/core/datatree_io.py

    if kwargs.get("format", None) not in [None, "NETCDF4"]:
        raise ValueError("to_netcdf only supports the NETCDF4 format")

    engine = kwargs.get("engine", None)
-    if engine not in [None, "netcdf4", "h5netcdf"]:
+    if engine not in [None, *get_args(T_DataTreeNetcdfEngine)]:


No stopper but I'm not sure I like the kwargs.get-style here, engine wont be typed and might hide typing issues downstream. Is engine correctly narrowed from Any after this if-case?

If we want good typing coverage it's better to prefer keyword-only arguments (with default values) rather than using kwargs.get.

Interesting - yeah. I took a quick look at to_netcdf in xarray/backends/api.py, and it looks like making these explicitly named with defaults would be consistent with that, too.

Looks like the potential list of arguments is: format, engine, group and compute (with group and compute for _datatree_to_zarr). These changes would also presumably require similar updates to DataTree.to_netcdf and DataTree.to_zarr(https://github.com/pydata/xarray/blob/main/xarray/core/datatree.py#L1516).

Just to check @TomNicholas - am I missing a reason for keeping the kwargs.get implementation?

I can't see a valid reason not to replace them, but I'm not 100% familiar.

I like the idea of making them consistent with the other to_netcdf functions in xarray/backends/api.py. I'll make this change and push it up as I get chance.

Just to check @TomNicholas - am I missing a reason for keeping the kwargs.get implementation?

Not that I can think of.

I've made these changes in this commit

TomNicholas · 2024-05-28T14:36:04Z

xarray/core/common.py

@@ -347,6 +347,34 @@ def _ipython_key_completions_(self) -> list[str]:
        return list(items)


+class TreeAttrAccessMixin(AttrAccessMixin):


Why is this class necessary? Can't DataTree inherit directly from AttrAccessMixin? And override _attr_sources and so on?

I think the reason this class existed separately was just to do with hacking around limitations of datatree being in a separate repository...

I initially tried to use AttrAccessMixin. Just dropping it in as a replacement for TreeAttrAccessMixin doesn't work because of the checking done inside the AttrAccessMixin.__init_subclass__ method. The check for if __dict__ exists fails. DataTree, TreeNode and NamedNode do all define __slots__, but I think there are attributes being defined in places like DataTree.__init__, for example DataTree.parent and DataTree.children (versus the DataTree._parent and DataTree._children attributes defined in DataTree.__slots__).

I think we talked about this a little bit in one of our calls - things like DataTree.parent and DataTree.children seemed like tricky things to change over.

TomNicholas · 2024-05-28T14:40:04Z

xarray/core/datatree_io.py

    if kwargs.get("format", None) not in [None, "NETCDF4"]:
        raise ValueError("to_netcdf only supports the NETCDF4 format")

    engine = kwargs.get("engine", None)
-    if engine not in [None, "netcdf4", "h5netcdf"]:
+    if engine not in [None, *get_args(T_DataTreeNetcdfEngine)]:


Just to check @TomNicholas - am I missing a reason for keeping the kwargs.get implementation?

Not that I can think of.

owenlittlejohns added 3 commits May 6, 2024 16:58

DAS-2067 - Migrate datatree_io.py.

9d554da

DAS-2067 - Migrate TreeAttrAccessMixin to xarray/core/common.py

b3a4662

DAS-2067 - Update whats-new.rst.

fecb5d7

owenlittlejohns commented May 6, 2024

View reviewed changes

TomNicholas added the topic-DataTree Related to the implementation of a DataTree class label May 7, 2024

TomNicholas added this to In progress in DataTree integration via automation May 7, 2024

flamingbear reviewed May 7, 2024

View reviewed changes

Illviljan reviewed May 7, 2024

View reviewed changes

xarray/core/datatree_io.py Outdated Show resolved Hide resolved

DAS-2067 - Implement T_DataTreeNetcdfEngine Literal type.

6ed1491

Illviljan reviewed May 13, 2024

View reviewed changes

owenlittlejohns added 2 commits May 20, 2024 20:47

Merge branch 'main' into DAS-2067-io-and-common

9565b41

DAS-2067 - Move PR content to 2024.05.1.

f3261e7

TomNicholas reviewed May 28, 2024

View reviewed changes

owenlittlejohns added 2 commits May 28, 2024 11:16

DAS-2067 - Additional typing on DataTree.to_netcdf and DataTree.to_zarr

e968631

DAS-2067 - Update DataTree.to_zarr documentation string.

3d4ad40

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

DAS-2067 - Migrate datatree io.py and common.py #9011

DAS-2067 - Migrate datatree io.py and common.py #9011

owenlittlejohns commented May 6, 2024

owenlittlejohns May 6, 2024

flamingbear May 7, 2024

owenlittlejohns May 7, 2024

owenlittlejohns May 8, 2024

flamingbear left a comment

flamingbear May 7, 2024

Illviljan May 13, 2024 •

edited

owenlittlejohns May 14, 2024

flamingbear May 21, 2024

owenlittlejohns May 22, 2024 •

edited

TomNicholas May 28, 2024

owenlittlejohns May 28, 2024

TomNicholas May 28, 2024

owenlittlejohns May 28, 2024

TomNicholas May 28, 2024

		@@ -347,6 +347,34 @@ def _ipython_key_completions_(self) -> list[str]:
		return list(items)


		class TreeAttrAccessMixin(AttrAccessMixin):

DAS-2067 - Migrate datatree io.py and common.py #9011

Are you sure you want to change the base?

DAS-2067 - Migrate datatree io.py and common.py #9011

Conversation

owenlittlejohns commented May 6, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

flamingbear left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Illviljan May 13, 2024 • edited

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

owenlittlejohns May 22, 2024 • edited

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Illviljan May 13, 2024 •

edited

owenlittlejohns May 22, 2024 •

edited