diff --git a/spalloc_client/protocol_client.py b/spalloc_client/protocol_client.py index 5706ecb3d..2da314e42 100644 --- a/spalloc_client/protocol_client.py +++ b/spalloc_client/protocol_client.py @@ -19,8 +19,6 @@ import json import socket from threading import current_thread, RLock, local -from spinn_utilities.abstract_context_manager import AbstractContextManager -from spinn_utilities.overrides import overrides from spalloc_client._utils import time_left, timed_out, make_timeout @@ -51,7 +49,7 @@ def __init__(self): self.sock = None -class ProtocolClient(AbstractContextManager): +class ProtocolClient(object): """ A simple (blocking) client implementation of the `spalloc-server `_ protocol. @@ -105,9 +103,13 @@ def __init__(self, hostname, port=22244, timeout=None): self._notifications_lock = RLock() self._default_timeout = timeout - @overrides(AbstractContextManager._context_entered) - def _context_entered(self): # pragma: no cover + def __enter__(self): self.connect(self._default_timeout) + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + self.close() + return False def _get_connection(self, timeout): if self._dead: