treewide: use ContextManagerMixin

__templates__/driver/pyproject.toml.tmpl

@@ -9,7 +9,7 @@ authors = [
 ]
 requires-python = ">=3.11"
 dependencies = [
-    "anyio>=4.6.2.post1",
+    "anyio>=4.10.0",
     "jumpstarter",
 ]
 

packages/jumpstarter-driver-energenie/pyproject.toml

@@ -9,7 +9,7 @@ authors = [
 ]
 requires-python = ">=3.11"
 dependencies = [
-    "anyio>=4.6.2.post1",
+    "anyio>=4.10.0",
     "jumpstarter",
     "jumpstarter-driver-power"
 ]

packages/jumpstarter-driver-flashers/pyproject.toml

@@ -11,7 +11,7 @@ authors = [
 requires-python = ">=3.11"
 dependencies = [
     "oras>=0.2.25",
-    "anyio>=4.6.2.post1",
+    "anyio>=4.10.0",
     "jumpstarter",
     "jumpstarter-driver-opendal",
     "jumpstarter-driver-pyserial",

packages/jumpstarter-driver-http-power/pyproject.toml

@@ -9,7 +9,7 @@ authors = [
 ]
 requires-python = ">=3.11"
 dependencies = [
-    "anyio>=4.6.2.post1",
+    "anyio>=4.10.0",
     "jumpstarter",
     "jumpstarter-driver-power",
 ]

packages/jumpstarter-driver-http/pyproject.toml

@@ -7,7 +7,7 @@ license = "Apache-2.0"
 authors = [{ name = "Benny Zlotnik", email = "bzlotnik@redhat.com" }]
 requires-python = ">=3.11"
 dependencies = [
-    "anyio>=4.6.2.post1",
+    "anyio>=4.10.0",
     "jumpstarter",
     "jumpstarter-driver-composite",
     "jumpstarter-driver-opendal",

packages/jumpstarter-driver-iscsi/pyproject.toml

@@ -6,7 +6,7 @@ readme = "README.md"
 authors = [{ name = "Benny Zlotnik", email = "bzlotnik@redhat.com" }]
 requires-python = ">=3.11"
 dependencies = [
-    "anyio>=4.6.2.post1",
+    "anyio>=4.10.0",
     "jumpstarter",
     "jumpstarter-driver-composite",
     "jumpstarter-driver-opendal",

packages/jumpstarter-driver-network/jumpstarter_driver_network/client.py

@@ -1,8 +1,11 @@
-from contextlib import AbstractContextManager
+from collections.abc import Generator
+from contextlib import contextmanager
 from ipaddress import IPv6Address, ip_address
 from threading import Event
+from typing import Any
 
 import click
+from anyio import ContextManagerMixin
 
 from .adapters import DbusAdapter, TcpPortforwardAdapter, UnixPortforwardAdapter
 from .driver import DbusNetwork
@@ -62,13 +65,11 @@ def forward_unix(path: str | None):
         return base
 
 
-class DbusNetworkClient(NetworkClient, AbstractContextManager):
-    def __enter__(self):
-        self.adapter = DbusAdapter(client=self)
-        self.adapter.__enter__()
-
-    def __exit__(self, exc_type, exc_value, traceback):
-        self.adapter.__exit__(exc_type, exc_value, traceback)
+class DbusNetworkClient(NetworkClient, ContextManagerMixin):
+    @contextmanager
+    def __contextmanager__(self) -> Generator[Any]:
+        with DbusAdapter(client=self) as value:
+            yield value
 
     @property
     def kind(self):

packages/jumpstarter-driver-probe-rs/pyproject.toml

@@ -7,7 +7,7 @@ license = "Apache-2.0"
 authors = [{ name = "Miguel Angel Ajo", email = "miguelangel@ajo.es" }]
 requires-python = ">=3.11"
 dependencies = [
-    "anyio>=4.6.2.post1",
+    "anyio>=4.10.0",
     "click>=8.1.7.2",
     "jumpstarter",
     "jumpstarter-driver-opendal",

packages/jumpstarter-driver-shell/pyproject.toml

@@ -6,7 +6,7 @@ readme = "README.md"
 authors = [{ name = "Miguel Angel Ajo", email = "miguelangel@ajo.es" }]
 requires-python = ">=3.11"
 license = "Apache-2.0"
-dependencies = ["anyio>=4.6.2.post1", "jumpstarter"]
+dependencies = ["anyio>=4.10.0", "jumpstarter"]
 
 [project.entry-points."jumpstarter.drivers"]
 Shell = "jumpstarter_driver_shell.driver:Shell"

packages/jumpstarter-driver-tasmota/pyproject.toml

@@ -7,7 +7,7 @@ license = "Apache-2.0"
 authors = [{ name = "Nick Cao", email = "nickcao@nichi.co" }]
 requires-python = ">=3.11"
 dependencies = [
-    "anyio>=4.6.2.post1",
+    "anyio>=4.10.0",
     "jumpstarter_driver_power",
     "jumpstarter",
     "paho-mqtt>=2.1.0",

packages/jumpstarter-driver-tftp/pyproject.toml

@@ -7,7 +7,7 @@ license = "Apache-2.0"
 authors = [{ name = "Benny Zlotnik", email = "bzlotnik@redhat.com" }]
 requires-python = ">=3.11"
 dependencies = [
-    "anyio>=4.6.2.post1",
+    "anyio>=4.10.0",
     "jumpstarter",
     "jumpstarter-driver-composite",
     "jumpstarter-driver-opendal",

packages/jumpstarter-driver-yepkit/pyproject.toml

@@ -7,7 +7,7 @@ license = "Apache-2.0"
 authors = [{ name = "Miguel Angel Ajo", email = "miguelangel@ajo.es" }]
 requires-python = ">=3.11"
 dependencies = [
-    "anyio>=4.6.2.post1",
+    "anyio>=4.10.0",
     "pyusb>=1.2.1",
     "jumpstarter_driver_power",
     "jumpstarter",

packages/jumpstarter/jumpstarter/client/lease.py

@@ -1,16 +1,15 @@
 import logging
+from collections.abc import AsyncGenerator, Generator
 from contextlib import (
-    AbstractAsyncContextManager,
-    AbstractContextManager,
     ExitStack,
     asynccontextmanager,
     contextmanager,
 )
 from dataclasses import dataclass, field
 from datetime import datetime, timedelta
-from typing import Any
+from typing import Any, Self
 
-from anyio import create_task_group, fail_after, sleep
+from anyio import AsyncContextManagerMixin, ContextManagerMixin, create_task_group, fail_after, sleep
 from anyio.from_thread import BlockingPortal
 from grpc.aio import Channel
 from jumpstarter_protocol import jumpstarter_pb2, jumpstarter_pb2_grpc
@@ -28,7 +27,7 @@
 
 
 @dataclass(kw_only=True)
-class Lease(AbstractContextManager, AbstractAsyncContextManager):
+class Lease(ContextManagerMixin, AsyncContextManagerMixin):
     channel: Channel
     duration: timedelta
     selector: str
@@ -48,7 +47,6 @@ def __post_init__(self):
 
         self.controller = jumpstarter_pb2_grpc.ControllerServiceStub(self.channel)
         self.svc = ClientService(channel=self.channel, namespace=self.namespace)
-        self.manager = self.portal.wrap_async_context_manager(self)
 
     async def _create(self):
         logger.debug("Creating lease request for selector %s for duration %s", self.selector, self.duration)
@@ -127,23 +125,22 @@ async def _acquire(self):
 
                 await sleep(1)
 
-    async def __aenter__(self):
-        return await self.request_async()
-
-    async def __aexit__(self, exc_type, exc_value, traceback):
-        if self.release:
-            logger.info("Releasing Lease %s", self.name)
-            await self.svc.DeleteLease(
-                name=self.name,
-            )
-
-    def __enter__(self):
-        # wraps the async context manager enter
-        return self.manager.__enter__()
+    @asynccontextmanager
+    async def __asynccontextmanager__(self) -> AsyncGenerator[Self]:
+        value = await self.request_async()
+        try:
+            yield value
+        finally:
+            if self.release:
+                logger.info("Releasing Lease %s", self.name)
+                await self.svc.DeleteLease(
+                    name=self.name,
+                )
 
-    def __exit__(self, exc_type, exc_value, traceback):
-        # wraps the async context manager exit
-        return self.manager.__exit__(exc_type, exc_value, traceback)
+    @contextmanager
+    def __contextmanager__(self) -> Generator[Self]:
+        with self.portal.wrap_async_context_manager(self) as value:
+            yield value
 
     async def handle_async(self, stream):
         logger.debug("Connecting to Lease with name %s", self.name)

packages/jumpstarter/jumpstarter/exporter/exporter.py

@@ -1,10 +1,19 @@
 import logging
-from collections.abc import Awaitable, Callable
-from contextlib import AbstractAsyncContextManager, asynccontextmanager
+from collections.abc import AsyncGenerator, Awaitable, Callable
+from contextlib import asynccontextmanager
 from dataclasses import dataclass, field
+from typing import Self
 
 import grpc
-from anyio import connect_unix, create_memory_object_stream, create_task_group, sleep
+from anyio import (
+    AsyncContextManagerMixin,
+    CancelScope,
+    connect_unix,
+    create_memory_object_stream,
+    create_task_group,
+    move_on_after,
+    sleep,
+)
 from anyio.abc import TaskGroup
 from google.protobuf import empty_pb2
 from jumpstarter_protocol import (
@@ -22,7 +31,7 @@
 
 
 @dataclass(kw_only=True)
-class Exporter(AbstractAsyncContextManager, Metadata):
+class Exporter(AsyncContextManagerMixin, Metadata):
     channel_factory: Callable[[], Awaitable[grpc.aio.Channel]]
     device_factory: Callable[[], Driver]
     lease_name: str = field(init=False, default="")
@@ -48,32 +57,34 @@ def stop(self, wait_for_lease_exit=False):
             self._stop_requested = True
             logger.info("Exporter marked for stop upon lease exit")
 
-    async def __aexit__(self, exc_type, exc_value, traceback):
-        import anyio
-
+    @asynccontextmanager
+    async def __asynccontextmanager__(self) -> AsyncGenerator[Self]:
         try:
-            if self.registered:
-                logger.info("Unregistering exporter with controller")
-                try:
-                    with anyio.move_on_after(10):  # 10 second timeout
-                        channel = await self.channel_factory()
-                        try:
-                            controller = jumpstarter_pb2_grpc.ControllerServiceStub(channel)
-                            await controller.Unregister(
-                                jumpstarter_pb2.UnregisterRequest(
-                                    reason="Exporter shutdown",
+            yield self
+        finally:
+            try:
+                if self.registered:
+                    logger.info("Unregistering exporter with controller")
+                    try:
+                        with move_on_after(10):  # 10 second timeout
+                            channel = await self.channel_factory()
+                            try:
+                                controller = jumpstarter_pb2_grpc.ControllerServiceStub(channel)
+                                await controller.Unregister(
+                                    jumpstarter_pb2.UnregisterRequest(
+                                        reason="Exporter shutdown",
+                                    )
                                 )
-                            )
-                            logger.info("Controller unregistration completed successfully")
-                        finally:
-                            with anyio.CancelScope(shield=True):
-                                await channel.close()
-                except Exception as e:
-                    logger.error("Error during controller unregistration: %s", e, exc_info=True)
-
-        except Exception as e:
-            logger.error("Error during exporter cleanup: %s", e, exc_info=True)
-            # Don't re-raise to avoid masking the original exception
+                                logger.info("Controller unregistration completed successfully")
+                            finally:
+                                with CancelScope(shield=True):
+                                    await channel.close()
+                    except Exception as e:
+                        logger.error("Error during controller unregistration: %s", e, exc_info=True)
+
+            except Exception as e:
+                logger.error("Error during exporter cleanup: %s", e, exc_info=True)
+                # Don't re-raise to avoid masking the original exception
 
     async def __handle(self, path, endpoint, token, tls_config, grpc_options):
         try:

packages/jumpstarter/jumpstarter/exporter/session.py

@@ -1,12 +1,14 @@
 import logging
 from collections import deque
-from contextlib import AbstractContextManager, asynccontextmanager, contextmanager, suppress
+from collections.abc import Generator
+from contextlib import asynccontextmanager, contextmanager, suppress
 from dataclasses import dataclass, field
 from logging.handlers import QueueHandler
+from typing import Self
 from uuid import UUID
 
 import grpc
-from anyio import Event, TypedAttributeLookupError, sleep
+from anyio import ContextManagerMixin, Event, TypedAttributeLookupError, sleep
 from anyio.from_thread import start_blocking_portal
 from jumpstarter_protocol import (
     jumpstarter_pb2,
@@ -30,32 +32,33 @@ class Session(
     jumpstarter_pb2_grpc.ExporterServiceServicer,
     router_pb2_grpc.RouterServiceServicer,
     Metadata,
-    AbstractContextManager,
+    ContextManagerMixin,
 ):
     root_device: Driver
     mapping: dict[UUID, Driver]
 
     _logging_queue: deque = field(init=False)
     _logging_handler: QueueHandler = field(init=False)
 
-    def __enter__(self):
+    @contextmanager
+    def __contextmanager__(self) -> Generator[Self]:
         logging.getLogger().addHandler(self._logging_handler)
         self.root_device.reset()
-        return self
-
-    def __exit__(self, exc_type, exc_value, traceback):
         try:
-            self.root_device.close()
-        except Exception as e:
-            # Get driver name from report for more descriptive logging
-            try:
-                report = self.root_device.report()
-                driver_name = report.labels.get('jumpstarter.dev/name', self.root_device.__class__.__name__)
-            except Exception:
-                driver_name = self.root_device.__class__.__name__
-            logger.error("Error closing driver %s: %s", driver_name, e, exc_info=True)
+            yield self
         finally:
-            logging.getLogger().removeHandler(self._logging_handler)
+            try:
+                self.root_device.close()
+            except Exception as e:
+                # Get driver name from report for more descriptive logging
+                try:
+                    report = self.root_device.report()
+                    driver_name = report.labels.get("jumpstarter.dev/name", self.root_device.__class__.__name__)
+                except Exception:
+                    driver_name = self.root_device.__class__.__name__
+                logger.error("Error closing driver %s: %s", driver_name, e, exc_info=True)
+            finally:
+                logging.getLogger().removeHandler(self._logging_handler)
 
     def __init__(self, *args, root_device, **kwargs):
         super().__init__(*args, **kwargs)