UgridDataset.data_vars doesn't preserve type

[JoerivanEngelen]

Problem

I ran into an issue where I assumed a call to UgridDataset.data_vars would return a mapping of variable names to UgridDataArrays, but unfortunately for me it returned a mapping to regular xarray DataArrays.

import xarray as xr
import xugrid as xu
import numpy as np

grid = xu.data.xoxo()
ds = xr.DataArray(np.arange(283), dims=("mesh2d_nFaces")).to_dataset(name="var")
uds = xu.UgridDataset(ds, grids=[grid])
print(type(uds.data_vars["var"])) #prints xr.DataArray

Context

MyPy gave me some difficulties, I used data_vars to work around this. Python runs both examples without any issues (except for the bug described above).

def my_func(var: xu.UgridDataArray | xr.DataArray):
     return var

my_func(**dict(**uds.data_vars))  # MyPy is OK
my_func(**uds)  # MyPy fails

[Huite]

This is another one that would be fixed by #35

The wrapping code inspects the output any xarray method and only if the return type is a DataArray or Dataset will it re-attach the topology.
We could try fixing this consistently, we'd also wrap the Groupby objects, etc.

The explict indexes project isn't finished yet, but it appears the majority of work is done. It may be worthwhile to check the current status.