Verified Commit acb9b6ac authored by Dom Sekotill's avatar Dom Sekotill
Browse files

Refactor BaseClient.attach()

parent 02a4afb6
Loading
Loading
Loading
Loading
+22 −23
Original line number Diff line number Diff line
@@ -21,9 +21,11 @@ from __future__ import annotations
import enum
import logging
import os
import sys
from collections.abc import AsyncIterator
from contextlib import asynccontextmanager
from re import compile as regex
from types import TracebackType as Traceback
from typing import AsyncContextManager
from typing import Callable
from typing import Optional
from typing import Tuple
@@ -211,11 +213,28 @@ class BaseClient:
			)
		return None

	def attach(self) -> AsyncContextManager[None]:
	@asynccontextmanager
	async def attach(self) -> AsyncIterator[None]:
		"""
		Return a context manager that handles attaching to the daemon's message queue
		"""
		return self._AttachContext(self)
		assert self._eventcount >= 0
		self._eventcount += 1
		if self._eventcount == 1:
			await self.send_command(consts.COMMAND_ATTACH)
		try:
			yield
		except:
			if __debug__:
				exc_type, *_ = sys.exc_info()
				assert exc_type is not None
				self.logger.debug("Detaching due to %s", exc_type.__name__)
			raise
		finally:
			assert self._eventcount > 0
			self._eventcount -= 1
			if self._eventcount == 0:
				await self.send_command(consts.COMMAND_DETACH)

	async def event(self, *events: str) -> EventInfo:
		"""
@@ -279,23 +298,3 @@ class BaseClient:
			return

		self._event = (prio, name, msg or None)

	class _AttachContext:
		def __init__(self, client: BaseClient) -> None:
			self.client = client

		async def __aenter__(self) -> None:
			client = self.client
			assert client._eventcount >= 0
			if client._eventcount == 0:
				await client.send_command(consts.COMMAND_ATTACH)
			client._eventcount += 1

		async def __aexit__(self, exc_type: type[BaseException]|None, *exc_info: object) -> None:
			client = self.client
			assert client._eventcount > 0
			client._eventcount -= 1
			if client._eventcount == 0:
				if __debug__ and exc_type:
					client.logger.debug(f"Detaching due to {exc_type.__name__}")
				await client.send_command(consts.COMMAND_DETACH)