From: Daan De Meyer <daan.j.demeyer@gmail.com>
Date: Fri, 21 Jul 2023 14:59:00 +0000 (+0200)
Subject: Make MkosiState a context manager
X-Git-Tag: v15~66^2~3
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=d69a0bc61551eb9c19511725c5128fd0ba755d14;p=thirdparty%2Fmkosi.git

Make MkosiState a context manager

We have to make sure we properly cleanup the workspace object we
allocate so let's make MkosiState a context manager so we do just
that.
---

diff --git a/mkosi/__init__.py b/mkosi/__init__.py
index 98113ade9..16fc2d8fc 100644
--- a/mkosi/__init__.py
+++ b/mkosi/__init__.py
@@ -47,7 +47,6 @@ from mkosi.types import PathString
 from mkosi.util import (
     InvokingUser,
     flatten,
-    flock,
     format_bytes,
     format_rlimit,
     scopedenv,
@@ -1749,12 +1748,11 @@ def finalize_staging(state: MkosiState) -> None:
 
 
 def build_image(args: MkosiArgs, config: MkosiConfig) -> None:
-    state = MkosiState(args, config)
     manifest = Manifest(config)
 
     # Make sure tmpfiles' aging doesn't interfere with our workspace
     # while we are working on it.
-    with flock(state.workspace), scopedenv({"TMPDIR" : str(state.workspace)}):
+    with MkosiState(args, config) as state, scopedenv({"TMPDIR" : str(state.workspace)}):
         install_package_manager_trees(state)
 
         with mount_image(state):
diff --git a/mkosi/state.py b/mkosi/state.py
index f83b74ba2..4beb19351 100644
--- a/mkosi/state.py
+++ b/mkosi/state.py
@@ -3,6 +3,8 @@
 import importlib
 import tempfile
 from pathlib import Path
+from types import TracebackType
+from typing import Optional, Type
 
 from mkosi.btrfs import btrfs_maybe_make_subvolume
 from mkosi.config import MkosiArgs, MkosiConfig
@@ -17,8 +19,6 @@ class MkosiState:
         self.args = args
         self.config = config
 
-        self._workspace = tempfile.TemporaryDirectory(dir=config.workspace_dir or Path.cwd(), prefix=".mkosi.tmp")
-
         try:
             distro = str(self.config.distribution)
             mod = importlib.import_module(f"mkosi.distributions.{distro}")
@@ -30,12 +30,22 @@ class MkosiState:
             die("No installer for this distribution.")
         self.installer = instance
 
+    def __enter__(self) -> "MkosiState":
+        self._workspace = tempfile.TemporaryDirectory(dir=self.config.workspace_dir or Path.cwd(), prefix=".mkosi.tmp")
         btrfs_maybe_make_subvolume(self.config, self.root, mode=0o755)
         self.staging.mkdir()
         self.pkgmngr.mkdir()
         self.install_dir.mkdir(exist_ok=True)
-
         self.cache_dir.mkdir(parents=True, exist_ok=True)
+        return self
+
+    def __exit__(
+        self,
+        exc_type: Optional[Type[BaseException]],
+        exc: Optional[BaseException],
+        traceback: Optional[TracebackType]
+    ) -> None:
+        self._workspace.cleanup()
 
     @property
     def workspace(self) -> Path: