Improve typehints of xr.Dataset.__getitem__

Resolves #4125

Author: nbren12

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`.