Improve typehints of xr.Dataset.__getitem__

Resolves pydata#4125
nbren12 · Jun 10, 2020 · 34e3ce0 · 34e3ce0
1 parent 2a288f6
commit 34e3ce0
Showing 1 changed file with 9 additions and 1 deletion.
diff --git a/xarray/core/dataset.py b/xarray/core/dataset.py
@@ -27,6 +27,7 @@
     TypeVar,
     Union,
     cast,
+    overload,
 )
 
 import numpy as np
@@ -1240,8 +1241,15 @@ def loc(self) -> _LocIndexer:
         and only when the key is a dict of the form {dim: labels}.
         """
         return _LocIndexer(self)
+
 
-    def __getitem__(self, key: Any) -> "Union[DataArray, Dataset]":
+    @overload
+    def __getitem__(self, key: Hashable) -> DataArray: ...
+
+    @overload
+    def __getitem__(self, key: Any) -> Dataset: ...
+
+    def __getitem__(self, key):
         """Access variables or coordinates this dataset as a
         :py:class:`~xarray.DataArray`.