From d75a9b7fafc1ff95b91b9ee5f38153e767144d1a Mon Sep 17 00:00:00 2001 From: Johannes Berg Date: Sun, 5 Dec 2021 21:19:52 +0100 Subject: [PATCH] Context manage connections (via BaseConn) Signed-off-by: Johannes Berg Reviewed-by: Rob Browning [rlb@defaultvalue.org: split from larger commit; adjust commit message] [rlb@defaultvalue.org: add pending_raise to exit handler] Signed-off-by: Rob Browning Tested-by: Rob Browning --- lib/bup/cmd/on.py | 6 +++--- lib/bup/cmd/server.py | 6 +++--- lib/bup/helpers.py | 8 +++++++- 3 files changed, 13 insertions(+), 7 deletions(-) diff --git a/lib/bup/cmd/on.py b/lib/bup/cmd/on.py index bb5e3f3..983283c 100755 --- a/lib/bup/cmd/on.py +++ b/lib/bup/cmd/on.py @@ -56,9 +56,9 @@ def main(argv): p.stdin.close() p.stdout.close() # Demultiplex remote client's stderr (back to stdout/stderr). - dmc = DemuxConn(p.stderr.fileno(), open(os.devnull, "wb")) - for line in iter(dmc.readline, b''): - out.write(line) + with DemuxConn(p.stderr.fileno(), open(os.devnull, "wb")) as dmc: + for line in iter(dmc.readline, b''): + out.write(line) finally: while 1: # if we get a signal while waiting, we have to keep waiting, just diff --git a/lib/bup/cmd/server.py b/lib/bup/cmd/server.py index 3d077dd..b4f0472 100755 --- a/lib/bup/cmd/server.py +++ b/lib/bup/cmd/server.py @@ -301,10 +301,10 @@ def main(argv): # FIXME: this protocol is totally lame and not at all future-proof. # (Especially since we abort completely as soon as *anything* bad happens) sys.stdout.flush() - conn = Conn(byte_stream(sys.stdin), byte_stream(sys.stdout)) - lr = linereader(conn) - with finalized(None, lambda _: repo and repo.close()), \ + with Conn(byte_stream(sys.stdin), byte_stream(sys.stdout)) as conn, \ + finalized(None, lambda _: repo and repo.close()), \ finalized(None, lambda _: suspended_w and suspended_w.close()): + lr = linereader(conn) for _line in lr: line = _line.strip() if not line: diff --git a/lib/bup/helpers.py b/lib/bup/helpers.py index fdc683b..28d1c6d 100644 --- a/lib/bup/helpers.py +++ b/lib/bup/helpers.py @@ -459,7 +459,13 @@ class BaseConn: def close(self): self._base_closed = True - while self._read(65536): pass + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_value, tb): + with pending_raise(exc_value, rethrow=False): + self.close() def __del__(self): assert self._base_closed -- 2.39.2