weighted: small improvements

[mathause]

Small improvement to weighted - more accurate typing and using _to_temp_dataset & map. No user facing changes.

[max-sixty]

Nice!

[mathause]

The following now works well

import xarray as xr
da = xr.DataArray([1])
reveal_type(xr.core.weighted.DataArrayWeighted(da, da).mean())

However, what I cannot get to work is:

reveal_type(da.weighted(da))

(returns Any). What almost works is:

def weighted(self: T_DSorDA, weights: "DataArray") -> Weighted[T_DSorDA]:

    return self._weighted_cls(self, weights)

However, mypy does not like that I overwrite the type of self with T_DSorDA...

[mathause]

Finally figured it out - I had to use

T_DataWithCoords = TypeVar("T_DataWithCoords", bound="DataWithCoords")

DataWithCoords is a parent class of Dastaset and DataArray. Now the following works:

import xarray as xr

da = xr.DataArray([1])
ds = xr.Dataset(data_vars={"x": da})

reveal_type(da.weighted(da).mean()) # Revealed type is 'xarray.core.dataarray.DataArray*'
reveal_type(ds.weighted(da).mean()) # Revealed type is 'xarray.core.dataarray.Dataset*'