Repository: pydata/xarray

Issue:
Consider supporting Dataset objects in DataTree.__setitem__
### What is your issue?

I think we could treat assigning a Dataset object as equivalent to assigning a DataTree object with the dataset's data without any children.

Currently this results in a rather bizarre error message:
```
In [35]: tree = DataTree()

In [36]: tree['foo'] = Dataset()
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-36-0aa35867bc31> in ?()
----> 1 tree['foo'] = Dataset()

~/dev/xarray/xarray/core/datatree.py in ?(self, key, value)
    923         elif isinstance(key, str):
    924             # TODO should possibly deal with hashables in general?
    925             # path-like: a name of a node/variable, or path to a node/variable
    926             path = NodePath(key)
--> 927             return self._set_item(path, value, new_nodes_along_path=True)
    928         else:
    929             raise ValueError("Invalid format for key")

~/dev/xarray/xarray/core/treenode.py in ?(self, path, item, new_nodes_along_path, allow_overwrite)
    556                 current_node._set(name, item)
    557             else:
    558                 raise KeyError(f"Already a node object at path {path}")
    559         else:
--> 560             current_node._set(name, item)

~/dev/xarray/xarray/core/datatree.py in ?(self, key, val)
    899             new_node._set_parent(new_parent=self, child_name=key)
    900         else:
    901             if not isinstance(val, DataArray | Variable):
    902                 # accommodate other types that can be coerced into Variables
--> 903                 val = DataArray(val)
    904
    905             self.update({key: val})

~/dev/xarray/xarray/core/dataarray.py in ?(self, data, coords, dims, name, attrs, indexes, fastpath)
    474             if attrs is None and not isinstance(data, PANDAS_TYPES):
    475                 attrs = getattr(data, "attrs", None)
    476
    477             data = _check_data_shape(data, coords, dims)
--> 478             data = as_compatible_data(data)
    479             coords, dims = _infer_coords_and_dims(data.shape, coords, dims)
    480             variable = Variable(dims, data, attrs, fastpath=True)
    481

~/dev/xarray/xarray/core/variable.py in ?(data, fastpath)
    316     ):
    317         return cast("T_DuckArray", data)
    318
    319     # validate whether the data is valid data types.
--> 320     data = np.asarray(data)
    321
    322     if isinstance(data, np.ndarray) and data.dtype.kind in "OMm":
    323         data = _possibly_convert_objects(data)

~/dev/xarray/xarray/core/dataset.py in ?(self, dtype, copy)
   1535         def __array__(self, dtype=None, copy=None):
-> 1536             raise TypeError(
   1537                 "cannot directly convert an xarray.Dataset into a "
   1538                 "numpy array. Instead, create an xarray.DataArray "
   1539                 "first, either with indexing on the Dataset or by "

TypeError: cannot directly convert an xarray.Dataset into a numpy array. Instead, create an xarray.DataArray first, either with indexing on the Dataset or by invoking the `to_dataarray()` method.
```

Discussion:

Patch:
diff --git a/doc/whats-new.rst b/doc/whats-new.rst
index 388d4d67..e4b2a06a 100644
--- a/doc/whats-new.rst
+++ b/doc/whats-new.rst
@@ -47,6 +47,11 @@ Bug fixes
 - Make illegal path-like variable names when constructing a DataTree from a Dataset
   (:issue:`9339`, :pull:`9378`)
   By `Etienne Schalk <https://github.com/etienneschalk>`_.
+- Work around `upstream pandas issue
+  <https://github.com/pandas-dev/pandas/issues/56996>`_ to ensure that we can
+  decode times encoded with small integer dtype values (e.g. ``np.int32``) in
+  environments with NumPy 2.0 or greater without needing to fall back to cftime
+  (:pull:`9518`). By `Spencer Clark <https://github.com/spencerkclark>`_.
 - Fix bug when encoding times with missing values as floats in the case when
   the non-missing times could in theory be encoded with integers
   (:issue:`9488`, :pull:`9497`). By `Spencer Clark
diff --git a/xarray/coding/times.py b/xarray/coding/times.py
index 5655bd20..2b9a5cf7 100644
--- a/xarray/coding/times.py
+++ b/xarray/coding/times.py
@@ -254,6 +254,15 @@ def _decode_datetime_with_pandas(
             "pandas."
         )
 
+    # Work around pandas.to_timedelta issue with dtypes smaller than int64 and
+    # NumPy 2.0 by casting all int and uint data to int64 and uint64,
+    # respectively. See https://github.com/pandas-dev/pandas/issues/56996 for
+    # more details.
+    if flat_num_dates.dtype.kind == "i":
+        flat_num_dates = flat_num_dates.astype(np.int64)
+    elif flat_num_dates.dtype.kind == "u":
+        flat_num_dates = flat_num_dates.astype(np.uint64)
+
     time_units, ref_date_str = _unpack_netcdf_time_units(units)
     time_units = _netcdf_to_numpy_timeunit(time_units)
     try:
diff --git a/xarray/core/datatree.py b/xarray/core/datatree.py
index 2b3179cb..bd583ac8 100644
--- a/xarray/core/datatree.py
+++ b/xarray/core/datatree.py
@@ -895,6 +895,8 @@ class DataTree(
             # TODO should possibly deal with hashables in general?
             # path-like: a name of a node/variable, or path to a node/variable
             path = NodePath(key)
+            if isinstance(value, Dataset):
+                value = DataTree(dataset=value)
             return self._set_item(path, value, new_nodes_along_path=True)
         else:
             raise ValueError("Invalid format for key")


Test patch:
diff --git a/xarray/tests/test_coding_times.py b/xarray/tests/test_coding_times.py
index 66caf25c..bb0dd1dd 100644
--- a/xarray/tests/test_coding_times.py
+++ b/xarray/tests/test_coding_times.py
@@ -8,7 +8,7 @@ from typing import Literal
 import numpy as np
 import pandas as pd
 import pytest
-from pandas.errors import OutOfBoundsDatetime
+from pandas.errors import OutOfBoundsDatetime, OutOfBoundsTimedelta
 
 from xarray import (
     DataArray,
@@ -1136,11 +1136,16 @@ def test_should_cftime_be_used_target_not_npable():
         _should_cftime_be_used(src, "noleap", False)
 
 
-@pytest.mark.parametrize("dtype", [np.uint8, np.uint16, np.uint32, np.uint64])
-def test_decode_cf_datetime_uint(dtype):
+@pytest.mark.parametrize(
+    "dtype",
+    [np.int8, np.int16, np.int32, np.int64, np.uint8, np.uint16, np.uint32, np.uint64],
+)
+def test_decode_cf_datetime_varied_integer_dtypes(dtype):
     units = "seconds since 2018-08-22T03:23:03Z"
     num_dates = dtype(50)
-    result = decode_cf_datetime(num_dates, units)
+    # Set use_cftime=False to ensure we cannot mask a failure by falling back
+    # to cftime.
+    result = decode_cf_datetime(num_dates, units, use_cftime=False)
     expected = np.asarray(np.datetime64("2018-08-22T03:23:53", "ns"))
     np.testing.assert_equal(result, expected)
 
@@ -1154,6 +1159,14 @@ def test_decode_cf_datetime_uint64_with_cftime():
     np.testing.assert_equal(result, expected)
 
 
+def test_decode_cf_datetime_uint64_with_pandas_overflow_error():
+    units = "nanoseconds since 1970-01-01"
+    calendar = "standard"
+    num_dates = np.uint64(1_000_000 * 86_400 * 360 * 500_000)
+    with pytest.raises(OutOfBoundsTimedelta):
+        decode_cf_datetime(num_dates, units, calendar, use_cftime=False)
+
+
 @requires_cftime
 def test_decode_cf_datetime_uint64_with_cftime_overflow_error():
     units = "microseconds since 1700-01-01"
@@ -1438,10 +1451,8 @@ _ENCODE_DATETIME64_VIA_DASK_TESTS = {
         "days since 1700-01-01",
         np.dtype("int32"),
     ),
-    "mixed-cftime-pandas-encoding-with-prescribed-units-and-dtype": (
-        "250YS",
-        "days since 1700-01-01",
-        np.dtype("int32"),
+    "mixed-cftime-pandas-encoding-with-prescribed-units-and-dtype": pytest.param(
+        "250YS", "days since 1700-01-01", np.dtype("int32"), marks=requires_cftime
     ),
     "pandas-encoding-with-default-units-and-dtype": ("250YS", None, None),
 }
diff --git a/xarray/tests/test_datatree.py b/xarray/tests/test_datatree.py
index 22c2cae7..4e22ba57 100644
--- a/xarray/tests/test_datatree.py
+++ b/xarray/tests/test_datatree.py
@@ -513,14 +513,12 @@ class TestSetItem:
         results["."] = data
         assert_identical(results.to_dataset(), data)
 
-    @pytest.mark.xfail(reason="assigning Datasets doesn't yet create new nodes")
     def test_setitem_dataset_as_new_node(self):
         data = xr.Dataset({"temp": [0, 50]})
         folder1 = DataTree(name="folder1")
         folder1["results"] = data
         assert_identical(folder1["results"].to_dataset(), data)
 
-    @pytest.mark.xfail(reason="assigning Datasets doesn't yet create new nodes")
     def test_setitem_dataset_as_new_node_requiring_intermediate_nodes(self):
         data = xr.Dataset({"temp": [0, 50]})
         folder1 = DataTree(name="folder1")