Fix Kubernetes backend utils.py typing (#3889)

See: lexdene/kubernetes-stubs#11

Author: un-def

pyproject.toml

@@ -152,7 +152,7 @@ dev = [
     "testcontainers>=4.9.2",
     "pytest-xdist>=3.6.1",
     "pyinstrument>=5.0.0",
-    "kubernetes-stubs-elephant-fork",
+    "kubernetes-stubs-elephant-fork>=35.0.0.post3",
     # docs
     "dstack[server]; python_version >= '3.11'",
     "dstack-plugin-server; python_version >= '3.11'",

src/dstack/_internal/core/backends/kubernetes/utils.py

@@ -1,6 +1,17 @@
 from collections.abc import Generator
 from contextlib import contextmanager
-from typing import Annotated, Any, Callable, Generic, Literal, Optional, Protocol, TypeVar, Union
+from typing import (
+    Annotated,
+    Any,
+    Callable,
+    Generic,
+    Literal,
+    Optional,
+    Protocol,
+    TypeVar,
+    Union,
+    cast,
+)
 
 import yaml
 from kubernetes.client import CoreV1Api, V1Status
@@ -9,9 +20,7 @@
     # XXX: This function is missing in the stubs package
     new_client_from_config_dict,  # pyright: ignore[reportAttributeAccessIssue]
 )
-
-# XXX: The watch module is missing in the stubs package
-from kubernetes.watch import Watch  # pyright: ignore[reportMissingImports]
+from kubernetes.watch import Watch
 from pydantic import Field
 from typing_extensions import ParamSpec, TypedDict
 from urllib3.exceptions import HTTPError
@@ -157,7 +166,8 @@ def watch_events(
     method: Callable[P, ObjectList[T]], *args: P.args, **kwargs: P.kwargs
 ) -> Generator[Generator[tuple[str, T], None, None], None, None]:
     watch = Watch()
-    gen = _watch_events_gen(watch.stream(method, *args, **kwargs))
+    inner_gen = cast(Generator[_EventDict[T], None, None], watch.stream(method, *args, **kwargs))
+    gen = _watch_events_gen(inner_gen)
     try:
         yield gen
     finally:
@@ -182,8 +192,11 @@ class _ErrorEventDict(TypedDict):
     object: V1Status
 
 
+_EventDict = Union[_StateEventDict[T], _BookmarkEventDict[T], _ErrorEventDict]
+
+
 def _watch_events_gen(
-    gen: Generator[Union[_StateEventDict[T], _BookmarkEventDict[T], _ErrorEventDict], None, None],
+    gen: Generator[_EventDict[T], None, None],
 ) -> Generator[tuple[str, T], None, None]:
     try:
         for event in gen: