Context manage connections (via BaseConn)

Signed-off-by: Johannes Berg <johannes@sipsolutions.net>
Reviewed-by: Rob Browning <rlb@defaultvalue.org>
[rlb@defaultvalue.org: split from larger commit; adjust commit message]
[rlb@defaultvalue.org: add pending_raise to exit handler]
Signed-off-by: Rob Browning <rlb@defaultvalue.org>
Tested-by: Rob Browning <rlb@defaultvalue.org>

diff --git a/lib/bup/cmd/on.py b/lib/bup/cmd/on.py
index bb5e3f3117a41eded776545426b66672bc63dd42..983283cebf45730406cb23034c0cb05a660bea26 100755 (executable)
--- 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 3d077dd65d90a7e0253f1c10fa2352883a4a670c..b4f047271858b054d936e2174b00446dec7fe3e5 100755 (executable)
--- 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 fdc683bd7c2ca739699e760a0b1eda5c3a8f156d..28d1c6d3595b612f17cfb14decadb0279ee4b8b9 100644 (file)
--- 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