From 9c149f2a893f94ed1469f57e2a256f197ca0f58d Mon Sep 17 00:00:00 2001 From: Illviljan <14371165+Illviljan@users.noreply.github.com> Date: Sun, 13 Nov 2022 20:22:19 +0100 Subject: [PATCH 01/14] Switch to T_DataArray in .coords --- xarray/core/coordinates.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/xarray/core/coordinates.py b/xarray/core/coordinates.py index 47350b9403f..78f8515f690 100644 --- a/xarray/core/coordinates.py +++ b/xarray/core/coordinates.py @@ -16,6 +16,7 @@ if TYPE_CHECKING: from .dataarray import DataArray from .dataset import Dataset + from .types import T_DataArray # Used as the key corresponding to a DataArray's variable when converting # arbitrary DataArray objects to datasets @@ -343,7 +344,7 @@ class DataArrayCoordinates(Coordinates): __slots__ = ("_data",) - def __init__(self, dataarray: DataArray): + def __init__(self, dataarray: T_DataArray): self._data = dataarray @property @@ -366,7 +367,7 @@ def dtypes(self) -> Frozen[Hashable, np.dtype]: def _names(self) -> set[Hashable]: return set(self._data._coords) - def __getitem__(self, key: Hashable) -> DataArray: + def __getitem__(self, key: Hashable) -> T_DataArray: return self._data._getitem_coord(key) def _update_coords( From 816dbb5ce56ada21011843f318980b408538d918 Mon Sep 17 00:00:00 2001 From: Illviljan <14371165+Illviljan@users.noreply.github.com> Date: Sun, 13 Nov 2022 20:37:49 +0100 Subject: [PATCH 02/14] Update coordinates.py --- xarray/core/coordinates.py | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/xarray/core/coordinates.py b/xarray/core/coordinates.py index 78f8515f690..097e60c7b96 100644 --- a/xarray/core/coordinates.py +++ b/xarray/core/coordinates.py @@ -14,7 +14,6 @@ from .variable import Variable, calculate_dimensions if TYPE_CHECKING: - from .dataarray import DataArray from .dataset import Dataset from .types import T_DataArray @@ -23,10 +22,10 @@ _THIS_ARRAY = ReprObject("") -class Coordinates(Mapping[Hashable, "DataArray"]): +class Coordinates(Mapping[Hashable, "T_DataArray"]): __slots__ = () - def __getitem__(self, key: Hashable) -> DataArray: + def __getitem__(self, key: Hashable) -> T_DataArray: raise NotImplementedError() def __setitem__(self, key: Hashable, value: Any) -> None: @@ -276,10 +275,10 @@ def variables(self) -> Mapping[Hashable, Variable]: {k: v for k, v in self._data.variables.items() if k in self._names} ) - def __getitem__(self, key: Hashable) -> DataArray: + def __getitem__(self, key: Hashable) -> T_DataArray: if key in self._data.data_vars: raise KeyError(key) - return cast("DataArray", self._data[key]) + return cast("T_DataArray", self._data[key]) def to_dataset(self) -> Dataset: """Convert these coordinates into a new Dataset""" @@ -453,7 +452,7 @@ def drop_coords( def assert_coordinate_consistent( - obj: DataArray | Dataset, coords: Mapping[Any, Variable] + obj: T_DataArray | Dataset, coords: Mapping[Any, Variable] ) -> None: """Make sure the dimension coordinate of obj is consistent with coords. From f85901d8b289cc1991eb66c29e053ad1f7b3e0ab Mon Sep 17 00:00:00 2001 From: Illviljan <14371165+Illviljan@users.noreply.github.com> Date: Mon, 14 Nov 2022 00:12:35 +0100 Subject: [PATCH 03/14] Update coordinates.py --- xarray/core/coordinates.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/xarray/core/coordinates.py b/xarray/core/coordinates.py index 097e60c7b96..a57d1208d0c 100644 --- a/xarray/core/coordinates.py +++ b/xarray/core/coordinates.py @@ -238,6 +238,8 @@ class DatasetCoordinates(Coordinates): objects. """ + _data: Dataset + __slots__ = ("_data",) def __init__(self, dataset: Dataset): @@ -341,6 +343,8 @@ class DataArrayCoordinates(Coordinates): dimensions and the values given by corresponding DataArray objects. """ + _data: T_DataArray + __slots__ = ("_data",) def __init__(self, dataarray: T_DataArray): From c6a571302fce90d06ab86271a71c5c144391818a Mon Sep 17 00:00:00 2001 From: Illviljan <14371165+Illviljan@users.noreply.github.com> Date: Wed, 16 Nov 2022 23:46:57 +0100 Subject: [PATCH 04/14] mypy understands the type from items better apparanetly --- xarray/core/dataarray.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/xarray/core/dataarray.py b/xarray/core/dataarray.py index 15d1777b270..4db2d55c5e5 100644 --- a/xarray/core/dataarray.py +++ b/xarray/core/dataarray.py @@ -3993,8 +3993,8 @@ def to_dict(self, data: bool = True, encoding: bool = False) -> dict[str, Any]: """ d = self.variable.to_dict(data=data) d.update({"coords": {}, "name": self.name}) - for k in self.coords: - d["coords"][k] = self.coords[k].variable.to_dict(data=data) + for k, coord in self.coords.items(): + d["coords"][k] = coord.variable.to_dict(data=data) if encoding: d["encoding"] = dict(self.encoding) return d From 047ef1b08d1a0c2297d4c682288953ca4801104d Mon Sep 17 00:00:00 2001 From: Illviljan <14371165+Illviljan@users.noreply.github.com> Date: Wed, 16 Nov 2022 23:47:05 +0100 Subject: [PATCH 05/14] Update coordinates.py --- xarray/core/coordinates.py | 21 +++++++++++++++------ 1 file changed, 15 insertions(+), 6 deletions(-) diff --git a/xarray/core/coordinates.py b/xarray/core/coordinates.py index a57d1208d0c..10676712e0f 100644 --- a/xarray/core/coordinates.py +++ b/xarray/core/coordinates.py @@ -2,7 +2,17 @@ import warnings from contextlib import contextmanager -from typing import TYPE_CHECKING, Any, Hashable, Iterator, Mapping, Sequence, cast +from typing import ( + TYPE_CHECKING, + Any, + Generic, + Hashable, + Iterator, + Mapping, + Protocol, + Sequence, + cast, +) import numpy as np import pandas as pd @@ -14,6 +24,7 @@ from .variable import Variable, calculate_dimensions if TYPE_CHECKING: + from .dataarray import DataArray from .dataset import Dataset from .types import T_DataArray @@ -277,10 +288,10 @@ def variables(self) -> Mapping[Hashable, Variable]: {k: v for k, v in self._data.variables.items() if k in self._names} ) - def __getitem__(self, key: Hashable) -> T_DataArray: + def __getitem__(self, key: Hashable) -> DataArray: if key in self._data.data_vars: raise KeyError(key) - return cast("T_DataArray", self._data[key]) + return self._data[key] def to_dataset(self) -> Dataset: """Convert these coordinates into a new Dataset""" @@ -343,11 +354,9 @@ class DataArrayCoordinates(Coordinates): dimensions and the values given by corresponding DataArray objects. """ - _data: T_DataArray - __slots__ = ("_data",) - def __init__(self, dataarray: T_DataArray): + def __init__(self, dataarray: T_DataArray) -> None: self._data = dataarray @property From 130754d7c3ed69625252c4115004d88e5e2de36a Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 16 Nov 2022 22:48:21 +0000 Subject: [PATCH 06/14] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- xarray/core/coordinates.py | 12 +----------- 1 file changed, 1 insertion(+), 11 deletions(-) diff --git a/xarray/core/coordinates.py b/xarray/core/coordinates.py index 10676712e0f..c9ab8711a6a 100644 --- a/xarray/core/coordinates.py +++ b/xarray/core/coordinates.py @@ -2,17 +2,7 @@ import warnings from contextlib import contextmanager -from typing import ( - TYPE_CHECKING, - Any, - Generic, - Hashable, - Iterator, - Mapping, - Protocol, - Sequence, - cast, -) +from typing import TYPE_CHECKING, Any, Hashable, Iterator, Mapping, Sequence import numpy as np import pandas as pd From 5d3319fcad0cb8a6a03f576f9c638d116b65172f Mon Sep 17 00:00:00 2001 From: Michael Niklas Date: Fri, 18 Nov 2022 16:49:39 +0100 Subject: [PATCH 07/14] resolve DataArrayCoords generic type --- xarray/core/coordinates.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/xarray/core/coordinates.py b/xarray/core/coordinates.py index ba6e7db001b..e566f64738e 100644 --- a/xarray/core/coordinates.py +++ b/xarray/core/coordinates.py @@ -14,6 +14,7 @@ from .variable import Variable, calculate_dimensions if TYPE_CHECKING: + from .common import DataWithCoords from .dataarray import DataArray from .dataset import Dataset from .types import T_DataArray @@ -24,6 +25,7 @@ class Coordinates(Mapping[Hashable, "T_DataArray"]): + _data: DataWithCoords __slots__ = () def __getitem__(self, key: Hashable) -> T_DataArray: @@ -337,13 +339,15 @@ def _ipython_key_completions_(self): ] -class DataArrayCoordinates(Coordinates): +class DataArrayCoordinates(Coordinates[T_DataArray]): """Dictionary like container for DataArray coordinates. Essentially a dict with keys given by the array's dimensions and the values given by corresponding DataArray objects. """ + _data: T_DataArray + __slots__ = ("_data",) def __init__(self, dataarray: T_DataArray) -> None: From d7ae38d519a4e424125d80a4b1d0f4c438d91bad Mon Sep 17 00:00:00 2001 From: Michael Niklas Date: Fri, 18 Nov 2022 17:12:11 +0100 Subject: [PATCH 08/14] fix import --- xarray/core/coordinates.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xarray/core/coordinates.py b/xarray/core/coordinates.py index e566f64738e..4226176247b 100644 --- a/xarray/core/coordinates.py +++ b/xarray/core/coordinates.py @@ -339,7 +339,7 @@ def _ipython_key_completions_(self): ] -class DataArrayCoordinates(Coordinates[T_DataArray]): +class DataArrayCoordinates(Coordinates["T_DataArray"]): """Dictionary like container for DataArray coordinates. Essentially a dict with keys given by the array's From d8bef27e54aa9e81873d5d64fca6a1d4d324ca62 Mon Sep 17 00:00:00 2001 From: Illviljan <14371165+Illviljan@users.noreply.github.com> Date: Sun, 20 Nov 2022 09:07:45 +0100 Subject: [PATCH 09/14] Test adding __class_getitem__ --- xarray/core/coordinates.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/xarray/core/coordinates.py b/xarray/core/coordinates.py index 4226176247b..0fd660bef8f 100644 --- a/xarray/core/coordinates.py +++ b/xarray/core/coordinates.py @@ -23,11 +23,17 @@ # arbitrary DataArray objects to datasets _THIS_ARRAY = ReprObject("") +# TODO: Remove when min python version >= 3.9: +GenericAlias = type(list[int]) + class Coordinates(Mapping[Hashable, "T_DataArray"]): _data: DataWithCoords __slots__ = () + # TODO: Remove when min python version >= 3.9: + __class_getitem__ = classmethod(GenericAlias) + def __getitem__(self, key: Hashable) -> T_DataArray: raise NotImplementedError() From db6270d9825984f89dbd752f9b0d0aeae9c3f5da Mon Sep 17 00:00:00 2001 From: Illviljan <14371165+Illviljan@users.noreply.github.com> Date: Sun, 20 Nov 2022 09:12:54 +0100 Subject: [PATCH 10/14] Update coordinates.py --- xarray/core/coordinates.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/xarray/core/coordinates.py b/xarray/core/coordinates.py index 0fd660bef8f..8ee8842098f 100644 --- a/xarray/core/coordinates.py +++ b/xarray/core/coordinates.py @@ -2,7 +2,7 @@ import warnings from contextlib import contextmanager -from typing import TYPE_CHECKING, Any, Hashable, Iterator, Mapping, Sequence +from typing import TYPE_CHECKING, Any, Hashable, Iterator, List, Mapping, Sequence import numpy as np import pandas as pd @@ -24,7 +24,7 @@ _THIS_ARRAY = ReprObject("") # TODO: Remove when min python version >= 3.9: -GenericAlias = type(list[int]) +GenericAlias = type(List[int]) class Coordinates(Mapping[Hashable, "T_DataArray"]): From bbb2ca3e3ca9af6b8a2a38917e4569677f441cf6 Mon Sep 17 00:00:00 2001 From: Illviljan <14371165+Illviljan@users.noreply.github.com> Date: Sun, 20 Nov 2022 09:42:00 +0100 Subject: [PATCH 11/14] Test adding a _data slot. --- xarray/core/coordinates.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/xarray/core/coordinates.py b/xarray/core/coordinates.py index 8ee8842098f..a489dc8ecdd 100644 --- a/xarray/core/coordinates.py +++ b/xarray/core/coordinates.py @@ -29,10 +29,10 @@ class Coordinates(Mapping[Hashable, "T_DataArray"]): _data: DataWithCoords - __slots__ = () + __slots__ = ("_data",) - # TODO: Remove when min python version >= 3.9: - __class_getitem__ = classmethod(GenericAlias) + # # TODO: Remove when min python version >= 3.9: + # __class_getitem__ = classmethod(GenericAlias) def __getitem__(self, key: Hashable) -> T_DataArray: raise NotImplementedError() From af0afa468475e03fcdf591c341728b9201826c37 Mon Sep 17 00:00:00 2001 From: Illviljan <14371165+Illviljan@users.noreply.github.com> Date: Sun, 20 Nov 2022 09:46:05 +0100 Subject: [PATCH 12/14] Adding class_getitem seems to work. --- xarray/core/coordinates.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/xarray/core/coordinates.py b/xarray/core/coordinates.py index a489dc8ecdd..52110d5bfee 100644 --- a/xarray/core/coordinates.py +++ b/xarray/core/coordinates.py @@ -31,8 +31,8 @@ class Coordinates(Mapping[Hashable, "T_DataArray"]): _data: DataWithCoords __slots__ = ("_data",) - # # TODO: Remove when min python version >= 3.9: - # __class_getitem__ = classmethod(GenericAlias) + # TODO: Remove when min python version >= 3.9: + __class_getitem__ = classmethod(GenericAlias) def __getitem__(self, key: Hashable) -> T_DataArray: raise NotImplementedError() From 12fbd72221e16efa8854034273e93884ea459128 Mon Sep 17 00:00:00 2001 From: Illviljan <14371165+Illviljan@users.noreply.github.com> Date: Sun, 20 Nov 2022 18:31:11 +0100 Subject: [PATCH 13/14] test mypy on 3.8 --- .github/workflows/ci-additional.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/ci-additional.yaml b/.github/workflows/ci-additional.yaml index dd18aa365d5..0eb1ce2f60d 100644 --- a/.github/workflows/ci-additional.yaml +++ b/.github/workflows/ci-additional.yaml @@ -89,7 +89,7 @@ jobs: shell: bash -l {0} env: CONDA_ENV_FILE: ci/requirements/environment.yml - PYTHON_VERSION: "3.10" + PYTHON_VERSION: "3.8" steps: - uses: actions/checkout@v3 From eceeacfbcb746966c7423b01d27150fdfdf20dbb Mon Sep 17 00:00:00 2001 From: Illviljan <14371165+Illviljan@users.noreply.github.com> Date: Mon, 21 Nov 2022 22:36:38 +0100 Subject: [PATCH 14/14] Update ci-additional.yaml --- .github/workflows/ci-additional.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/ci-additional.yaml b/.github/workflows/ci-additional.yaml index 46e051cc4da..bdae56ae6db 100644 --- a/.github/workflows/ci-additional.yaml +++ b/.github/workflows/ci-additional.yaml @@ -89,7 +89,7 @@ jobs: shell: bash -l {0} env: CONDA_ENV_FILE: ci/requirements/environment.yml - PYTHON_VERSION: "3.8" + PYTHON_VERSION: "3.10" steps: - uses: actions/checkout@v3