Coding style.

sandbox_fs.py

@@ -3,7 +3,6 @@
 # Copyright (C) 2016 Andrew Hamilton. All rights reserved.
 # Licensed under the Artistic License 2.0.
 
-import contextlib
 import os
 import subprocess
 import sys
@@ -15,6 +14,12 @@ class OverlayfsMount():
     def __init__(self, lower_dir, mount_point):
         self.lower_dir = lower_dir
         self.mount_point = mount_point
+
+    def __repr__(self):
+        return "<OverlayfsMount:%r over %r>" % (self.mount_point,
+                                                self.lower_dir)
+
+    def __enter__(self):
         self.upper_dir = tempfile.TemporaryDirectory()
         self.work_dir = tempfile.TemporaryDirectory()
         option_string = ("lowerdir=%s,upperdir=%s,workdir=%s" %
@@ -23,12 +28,9 @@ class OverlayfsMount():
         subprocess.check_call(["mount", "-t", "overlay", "-o",
                                option_string, "overlay", self.mount_point],
                               stderr=subprocess.PIPE)
+        return self
 
-    def __repr__(self):
-        return "<OverlayfsMount:%r over %r>" % (self.mount_point,
-                                                self.lower_dir)
-
-    def umount(self):
+    def __exit__(self, exc_type, exc_value, traceback):
         subprocess.check_call(["umount", "--lazy", self.mount_point])
 
 
@@ -55,31 +57,34 @@ def _find_mounts():
 
 class SandboxFs:
 
-    def __init__(self, mount_point, holes=None):
-        self.mount_point = mount_point
+    def __init__(self, holes=None):
         self.holes = [] if holes is None else holes
         for hole in self.holes:
             if not hole.startswith("/"):
                 raise ValueError("Holes must be absolute paths: %r" % hole)
+        self.temp_dir = tempfile.TemporaryDirectory()
+        self.mount_point = self.temp_dir.name
         self.overlay_mounts = []
 
     def __repr__(self):
         return ("<SandboxFs:%r mounts:%r>" %
-                (self.mount_point, len(self.overlay_mounts)))
+                (self.temp_dir.name, len(self.overlay_mounts)))
 
-    def mount(self):
-        self.overlay_mounts = [OverlayfsMount(mount_point,
-                                              self.mount_point + mount_point)
-                               for mount_point in sorted(_find_mounts())]
+    def __enter__(self):
+        self.overlay_mounts = [
+            OverlayfsMount(mount_point,
+                           self.mount_point + mount_point).__enter__()
+            for mount_point in sorted(_find_mounts())]
         for hole in self.holes:
             subprocess.check_call(["mount", "--bind", hole,
                                    self.mount_point + hole])
+        return self
 
-    def umount(self):
+    def __exit__(self, exc_type, exc_value, traceback):
         for hole in reversed(self.holes):
             subprocess.check_call(["umount", self.mount_point + hole])
         for mount in reversed(self.overlay_mounts):
-            mount.umount()
+            mount.__exit__(None, None, None)
         self.overlay_mounts = []
 
     def command(self, command, env=None):
@@ -87,22 +92,11 @@ class SandboxFs:
                 _in_directory(os.getcwd(), command))
 
 
-@contextlib.contextmanager
-def sandbox_(holes=None):
-    temp_dir = tempfile.TemporaryDirectory()
-    sandbox = SandboxFs(temp_dir.name, holes)
-    sandbox.mount()
-    try:
-        yield sandbox
-    finally:
-        sandbox.umount()
-
-
 if __name__ == "__main__":
     try:
         divider_index = sys.argv.index("--")
         holes, command = sys.argv[1:divider_index], sys.argv[divider_index+1:]
     except ValueError:
         holes, command = None, sys.argv[1:]
-    with sandbox_(holes) as sandbox:
+    with SandboxFs(holes) as sandbox:
         subprocess.check_call(sandbox.command(command))