Use an ExitStack in the main block.

Simplifies cleanup and will be more useful in future commits.

1 files changed, 31 insertions, 30 deletions

diff --git a/fshare.py b/fshare.py
index 72a8d27..6bf4f0c 100755
--- a/fshare.py
+++ b/fshare.py
@@ -17,6 +17,7 @@
 # fshare.  If not, see <http://www.gnu.org/licenses/>.
 
 import argparse
+import contextlib
 import fcntl
 import hmac
 import json
@@ -375,33 +376,33 @@ def excepthook(t, v, tb, logger = logger):
     logger.error('Uncaught top-level exception', exc_info = (t, v, tb))
 sys.excepthook = excepthook
 
-# open the state dir
-try:
-    state_dir_fd = os.open(args.state_dir, os.O_RDONLY | os.O_DIRECTORY)
-except (FileNotFoundError, NotADirectoryError) as e:
-    logger.error('The state directory "%s" is not an existing directory: %s',
-                 args.state_dir, e)
-    sys.exit(1)
-
-# lock the state dir
-try:
-    fcntl.flock(state_dir_fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
-except BlockingIOError:
-    logger.error('The state directory is already locked by another process')
-    os.close(state_dir_fd)
-    sys.exit(1)
-
-try:
-    # read the state file
-    state = PersistentState(state_dir_fd, logger)
-
-    # launch the server
-    server = FShareServer((args.address, args.port), args.ipv4, args.ipv6,
-                          state, args.data_dir, logger)
-    server.serve_forever()
-except StateCorruptError:
-    logger.error('Corrupted state file')
-    sys.exit(1)
-finally:
-    fcntl.flock(state_dir_fd, fcntl.LOCK_UN)
-    os.close(state_dir_fd)
+with contextlib.ExitStack() as stack:
+    # open the state dir
+    try:
+        state_dir_fd = os.open(args.state_dir, os.O_RDONLY | os.O_DIRECTORY)
+    except (FileNotFoundError, NotADirectoryError) as e:
+        logger.error('The state directory "%s" is not an existing directory: %s',
+                     args.state_dir, e)
+        sys.exit(1)
+    stack.callback(os.close, state_dir_fd)
+
+    # lock the state dir
+    try:
+        fcntl.flock(state_dir_fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
+    except BlockingIOError:
+        logger.error('The state directory is already locked by another process')
+        sys.exit(1)
+
+    try:
+        # read the state file
+        state = PersistentState(state_dir_fd, logger)
+
+        # launch the server
+        server = FShareServer((args.address, args.port), args.ipv4, args.ipv6,
+                              state, args.data_dir, logger)
+        server.serve_forever()
+    except StateCorruptError:
+        logger.error('Corrupted state file')
+        sys.exit(1)
+    finally:
+        fcntl.flock(state_dir_fd, fcntl.LOCK_UN)