make sure that self._start and self.stop calls are contained within TaskGroup acmgr (#25)

* make sure that self._start and self.stop calls are contained within TaskGroup acmgr

* ruff format

* restore stack to None after exiting it

* declare __stack

* fix mypy

Author: graingert

src/zmq_anyio/_socket.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from contextlib import AsyncExitStack
 import pickle
 import selectors
 from collections import deque
@@ -156,7 +157,7 @@ class Socket(zmq.Socket):
     _fd = None
     _exit_stack = None
     _task_group = None
-    __task_group = None
+    __stack: AsyncExitStack | None = None
     _thread = None
     started = None
     stopped = None
@@ -191,6 +192,7 @@ def __init__(
         self._exited = Event()
         self.stopped = Event()
         self._task_group = task_group
+        self.__stack = None
 
     def get(self, key):
         result = super().get(key)
@@ -831,17 +833,26 @@ async def __aenter__(self) -> Socket:
             return
 
         self._starting = True
-        if self._task_group is None:
-            self.__task_group = create_task_group()
-            self._task_group = await self.__task_group.__aenter__()
-            await self._task_group.start(self._start)
+        if self._task_group is not None:
+            return self
+
+        async with AsyncExitStack() as stack:
+            self._task_group = task_group = await stack.enter_async_context(
+                create_task_group()
+            )
+            await task_group.start(self._start)
+            stack.push_async_callback(self.stop)
+            self.__stack = stack.pop_all()
 
         return self
 
     async def __aexit__(self, exc_type, exc_value, exc_tb):
+        if self.__stack is not None:
+            try:
+                return await self.__stack.__aexit__(exc_type, exc_value, exc_tb)
+            finally:
+                self.__stack = None
         await self.stop()
-        if self.__task_group is not None:
-            return await self.__task_group.__aexit__(exc_type, exc_value, exc_tb)
 
     async def start(
         self,