From 2bf98417a624f7b2037e0f0a11e0d549e91706ed Mon Sep 17 00:00:00 2001 From: desaleo <43293180+desaleo@users.noreply.github.com> Date: Wed, 22 Nov 2023 15:35:21 +0100 Subject: [PATCH 1/3] add context manager add a context manager for lighter syntax and guarantee connection closure --- factorio_rcon/_factorio_rcon.py | 27 ++++++++++++++++++++++++++- 1 file changed, 26 insertions(+), 1 deletion(-) diff --git a/factorio_rcon/_factorio_rcon.py b/factorio_rcon/_factorio_rcon.py index 5820b5b..35beb07 100644 --- a/factorio_rcon/_factorio_rcon.py +++ b/factorio_rcon/_factorio_rcon.py @@ -3,7 +3,8 @@ import functools import socket import struct -from typing import Any, Callable, Dict, NamedTuple, Optional, TypeVar, cast +from types import TracebackType +from typing import Any, Callable, Dict, NamedTuple, Optional, TypeVar, cast, Self try: import anyio @@ -402,6 +403,19 @@ def send_commands(self, commands: Dict[T, str]) -> Dict[T, Optional[str]]: else: results[id_map[response.id]] = response.body.rstrip() return results + + def __enter__(self) -> Self: + if self.rcon_socket is None: + raise RCONNotConnected(NOT_CONNECTED) + return self + + def __exit__( + self, + exc_type: Optional[type[BaseException]], + exc_value: Optional[BaseException], + traceback: Optional[TracebackType] + ) -> None: + self.close() class AsyncRCONClient(RCONSharedBase): @@ -616,6 +630,17 @@ async def send_commands(self, commands: Dict[T, str]) -> Dict[T, Optional[str]]: else: results[id_map[response.id]] = response.body.rstrip() return results + + async def __aenter__(self) -> Self: + return self + + async def __aexit__( + self, + exc_type: Optional[type[BaseException]], + exc_value: Optional[BaseException], + traceback: Optional[TracebackType] + ) -> None: + await self.close() INVALID_PASS = "The RCON password is incorrect" From f4644a0aa6a263156ad5ba040744bcae8bfa5792 Mon Sep 17 00:00:00 2001 From: desaleo Date: Sat, 25 Nov 2023 21:34:21 +0100 Subject: [PATCH 2/3] added requested changes - __enter__ and __ aenter__ now rely on handle_socket_errors to make sure socket is alive - __enter__ and __aenter__ are now typed with the class name instead of Self to maintain compatibility with python 3.8+ - formatted the whole file with black --- factorio_rcon/_factorio_rcon.py | 39 +++++++++++++++++---------------- 1 file changed, 20 insertions(+), 19 deletions(-) diff --git a/factorio_rcon/_factorio_rcon.py b/factorio_rcon/_factorio_rcon.py index 35beb07..89fb9ff 100644 --- a/factorio_rcon/_factorio_rcon.py +++ b/factorio_rcon/_factorio_rcon.py @@ -1,10 +1,11 @@ """RCON client for Factorio servers""" +from __future__ import annotations import enum import functools import socket import struct from types import TracebackType -from typing import Any, Callable, Dict, NamedTuple, Optional, TypeVar, cast, Self +from typing import Any, Callable, Dict, NamedTuple, Optional, TypeVar, cast try: import anyio @@ -403,18 +404,17 @@ def send_commands(self, commands: Dict[T, str]) -> Dict[T, Optional[str]]: else: results[id_map[response.id]] = response.body.rstrip() return results - - def __enter__(self) -> Self: - if self.rcon_socket is None: - raise RCONNotConnected(NOT_CONNECTED) + + @handle_socket_errors(alive_socket_required=True) + def __enter__(self) -> RCONClient: return self - + def __exit__( - self, - exc_type: Optional[type[BaseException]], - exc_value: Optional[BaseException], - traceback: Optional[TracebackType] - ) -> None: + self, + exc_type: Optional[type[BaseException]], + exc_value: Optional[BaseException], + traceback: Optional[TracebackType], + ) -> None: self.close() @@ -630,16 +630,17 @@ async def send_commands(self, commands: Dict[T, str]) -> Dict[T, Optional[str]]: else: results[id_map[response.id]] = response.body.rstrip() return results - - async def __aenter__(self) -> Self: + + @handle_socket_errors(alive_socket_required=True) + async def __aenter__(self) -> AsyncRCONClient: return self - + async def __aexit__( - self, - exc_type: Optional[type[BaseException]], - exc_value: Optional[BaseException], - traceback: Optional[TracebackType] - ) -> None: + self, + exc_type: Optional[type[BaseException]], + exc_value: Optional[BaseException], + traceback: Optional[TracebackType], + ) -> None: await self.close() From 7610c6d49bd9410c29cfbc6e1521458e0c85c40d Mon Sep 17 00:00:00 2001 From: desaleo Date: Tue, 23 Jan 2024 23:33:13 +0100 Subject: [PATCH 3/3] added requested changes - typing now use string instead of __future__ annotations - synchronous context manager now ignore the connect_on_init parameter and connect the client by default - context managers will attempt to connect the client if not already connected - context managers will attempt to reconnect the client if in an error state --- factorio_rcon/_factorio_rcon.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/factorio_rcon/_factorio_rcon.py b/factorio_rcon/_factorio_rcon.py index 89fb9ff..9f6545b 100644 --- a/factorio_rcon/_factorio_rcon.py +++ b/factorio_rcon/_factorio_rcon.py @@ -1,5 +1,4 @@ """RCON client for Factorio servers""" -from __future__ import annotations import enum import functools import socket @@ -405,8 +404,10 @@ def send_commands(self, commands: Dict[T, str]) -> Dict[T, Optional[str]]: results[id_map[response.id]] = response.body.rstrip() return results - @handle_socket_errors(alive_socket_required=True) - def __enter__(self) -> RCONClient: + @handle_socket_errors(alive_socket_required=False) + def __enter__(self) -> "RCONClient": + if self.rcon_socket is None or self.rcon_failure: + self.connect() return self def __exit__( @@ -631,8 +632,10 @@ async def send_commands(self, commands: Dict[T, str]) -> Dict[T, Optional[str]]: results[id_map[response.id]] = response.body.rstrip() return results - @handle_socket_errors(alive_socket_required=True) - async def __aenter__(self) -> AsyncRCONClient: + @async_handle_socket_errors(alive_socket_required=False) + async def __aenter__(self) -> "AsyncRCONClient": + if self.rcon_socket is None or self.rcon_failure: + await self.connect() return self async def __aexit__(