Refactored TemporaryFile and GitRepository into context managers.

src/assemble_workflow/bundle.py

@@ -10,10 +10,10 @@
 import shutil
 import subprocess
 import tarfile
-import tempfile
 from abc import ABC, abstractmethod
 
 from paths.script_finder import ScriptFinder
+from system.temporary_directory import TemporaryDirectory
 
 """
 This class is responsible for executing the build of the full bundle and passing results to a bundle recorder.
@@ -34,7 +34,7 @@ def __init__(self, build_manifest, artifacts_dir, bundle_recorder):
         self.plugins = self.__get_plugins(build_manifest.components)
         self.artifacts_dir = artifacts_dir
         self.bundle_recorder = bundle_recorder
-        self.tmp_dir = tempfile.TemporaryDirectory()
+        self.tmp_dir = TemporaryDirectory()
         self.installed_plugins = []
         self.min_tarball_path = self._copy_component(self.min_tarball, "dist")
         self.__unpack_min_tarball(self.tmp_dir.name)

src/git/git_repository.py

@@ -7,9 +7,10 @@
 import logging
 import os
 import subprocess
-import tempfile
 from pathlib import Path
 
+from system.temporary_directory import TemporaryDirectory
+
 
 class GitRepository:
     """
@@ -22,7 +23,7 @@ def __init__(self, url, ref, directory=None, working_subdirectory=None):
         self.url = url
         self.ref = ref
         if directory is None:
-            self.temp_dir = tempfile.TemporaryDirectory()
+            self.temp_dir = TemporaryDirectory()
             self.dir = os.path.realpath(self.temp_dir.name)
         else:
             self.temp_dir = None
@@ -41,6 +42,13 @@ def __checkout__(self):
         self.sha = self.output("git rev-parse HEAD", self.dir)
         logging.info(f"Checked out {self.url}@{self.ref} into {self.dir} at {self.sha}")
 
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_value, exc_traceback):
+        if self.temp_dir:
+            self.temp_dir.__exit__(exc_type, exc_value, exc_traceback)
+
     @property
     def working_directory(self):
         if self.working_subdirectory:

src/manifests_workflow/component_opensearch.py

@@ -23,17 +23,18 @@ def checkout(
         snapshot=False,
         working_directory=None,
     ):
-        return ComponentOpenSearch(
-            name,
-            GitRepository(
-                f"https://github.com/opensearch-project/{name}.git",
-                branch,
-                path,
-                working_directory,
-            ),
-            opensearch_version,
-            snapshot,
-        )
+        with GitRepository(
+            f"https://github.com/opensearch-project/{name}.git",
+            branch,
+            path,
+            working_directory,
+        ) as repo:
+            return ComponentOpenSearch(
+                name,
+                repo,
+                opensearch_version,
+                snapshot,
+            )
 
     def __init__(self, name, repo, opensearch_version, snapshot=False):
         super().__init__(name, repo, snapshot)

src/manifests_workflow/component_opensearch_dashboards_min.py

@@ -21,14 +21,15 @@ def branches(self):
 
     @classmethod
     def checkout(self, path, branch="main", snapshot=False):
-        return ComponentOpenSearchDashboardsMin(
-            GitRepository(
-                "https://github.com/opensearch-project/OpenSearch-Dashboards.git",
-                branch,
-                path,
-            ),
-            snapshot,
-        )
+        with GitRepository(
+            "https://github.com/opensearch-project/OpenSearch-Dashboards.git",
+            branch,
+            path,
+        ) as repo:
+            return ComponentOpenSearchDashboardsMin(
+                repo,
+                snapshot,
+            )
 
     @property
     def properties(self):

src/manifests_workflow/input_manifests.py

@@ -41,15 +41,15 @@ def update(self, min_klass, component_klass, keep=False):
         logging.info(f"Known versions: {known_versions}")
         main_versions = {}
         with TemporaryDirectory(keep=keep) as work_dir:
-            logging.info(f"Checking out components into {work_dir}")
-            os.chdir(work_dir)
+            logging.info(f"Checking out components into {work_dir.name}")
+            os.chdir(work_dir.name)
 
             # check out and build #main, 1.x, etc.
             branches = min_klass.branches()
             logging.info(f"Checking {self.name} {branches} branches")
             for branch in branches:
                 c = min_klass.checkout(
-                    path=os.path.join(work_dir, f"{self.name.replace(' ', '')}/{branch}"),
+                    path=os.path.join(work_dir.name, f"{self.name.replace(' ', '')}/{branch}"),
                     branch=branch,
                 )
                 version = c.version
@@ -67,7 +67,7 @@ def update(self, min_klass, component_klass, keep=False):
                     logging.info(f"Checking out {component.name}#main")
                     component = component_klass.checkout(
                         name=component.name,
-                        path=os.path.join(work_dir, component.name),
+                        path=os.path.join(work_dir.name, component.name),
                         version=manifest.build.version,
                         branch="main",
                     )

src/run_assemble.py

@@ -10,12 +10,12 @@
 import logging
 import os
 import sys
-import tempfile
 
 from assemble_workflow.bundle_recorder import BundleRecorder
 from assemble_workflow.bundles import Bundles
 from manifests.build_manifest import BuildManifest
 from system import console
+from system.temporary_directory import TemporaryDirectory
 
 
 def main():
@@ -40,10 +40,10 @@ def main():
     output_dir = os.path.join(os.getcwd(), "bundle")
     os.makedirs(output_dir, exist_ok=True)
 
-    with tempfile.TemporaryDirectory() as work_dir:
+    with TemporaryDirectory() as work_dir:
         logging.info(f"Bundling {build.name} ({build.architecture}) on {build.platform} into {output_dir} ...")
 
-        os.chdir(work_dir)
+        os.chdir(work_dir.name)
 
         bundle_recorder = BundleRecorder(build, output_dir, artifacts_dir)
 

src/run_build.py

@@ -28,9 +28,9 @@ def main():
     with TemporaryDirectory(keep=args.keep) as work_dir:
         output_dir = os.path.join(os.getcwd(), "artifacts")
 
-        logging.info(f"Building in {work_dir}")
+        logging.info(f"Building in {work_dir.name}")
 
-        os.chdir(work_dir)
+        os.chdir(work_dir.name)
 
         target = BuildTarget(
             name=manifest.build.name,
@@ -52,20 +52,19 @@ def main():
                 continue
 
             logging.info(f"Building {component.name}")
-            repo = GitRepository(
+            with GitRepository(
                 component.repository,
                 component.ref,
-                os.path.join(work_dir, component.name),
+                os.path.join(work_dir.name, component.name),
                 component.working_directory,
-            )
-
-            try:
-                builder = Builder(component.name, repo, build_recorder)
-                builder.build(target)
-                builder.export_artifacts()
-            except:
-                logging.error(f"Error building {component.name}, retry with: {args.component_command(component.name)}")
-                raise
+            ) as repo:
+                try:
+                    builder = Builder(component.name, repo, build_recorder)
+                    builder.build(target)
+                    builder.export_artifacts()
+                except:
+                    logging.error(f"Error building {component.name}, retry with: {args.component_command(component.name)}")
+                    raise
 
         build_recorder.write_manifest()
 

src/run_bwc_test.py

@@ -19,15 +19,15 @@ def main():
     args = TestArgs()
     console.configure(level=args.logging_level)
     with TemporaryDirectory(keep=args.keep) as work_dir:
-        logging.info("Switching to temporary work_dir: " + work_dir)
+        logging.info(f"Switching to temporary work_dir: {work_dir.name}")
         bundle_manifest = BundleManifest.from_s3(
             args.s3_bucket,
             args.build_id,
             args.opensearch_version,
             args.architecture,
-            work_dir,
+            work_dir.name,
         )
-        BwcTestSuite(bundle_manifest, work_dir, args.component, args.keep).execute()
+        BwcTestSuite(bundle_manifest, work_dir.name, args.component, args.keep).execute()
 
 
 if __name__ == "__main__":

src/run_checkout.py

@@ -23,21 +23,21 @@ def main():
     manifest = InputManifest.from_file(args.manifest)
 
     with TemporaryDirectory(keep=True) as work_dir:
-        logging.info(f"Checking out into {work_dir}")
+        logging.info(f"Checking out into {work_dir.name}")
 
-        os.chdir(work_dir)
+        os.chdir(work_dir.name)
 
         for component in manifest.components:
 
             logging.info(f"Checking out {component.name}")
             GitRepository(
                 component.repository,
                 component.ref,
-                os.path.join(work_dir, component.name),
+                os.path.join(work_dir.name, component.name),
                 component.working_directory,
             )
 
-    logging.info(f"Done, checked out into {work_dir}.")
+    logging.info(f"Done, checked out into {work_dir.name}.")
 
 
 if __name__ == "__main__":

src/run_ci.py

@@ -27,9 +27,9 @@ def main():
     target = CiTarget(version=manifest.build.version, snapshot=args.snapshot)
 
     with TemporaryDirectory(keep=args.keep) as work_dir:
-        logging.info(f"Sanity-testing in {work_dir}")
+        logging.info(f"Sanity-testing in {work_dir.name}")
 
-        os.chdir(work_dir)
+        os.chdir(work_dir.name)
 
         logging.info(f"Sanity testing {manifest.build.name}")
 
@@ -40,19 +40,18 @@ def main():
                 continue
 
             logging.info(f"Sanity checking {component.name}")
-            repo = GitRepository(
+            with GitRepository(
                 component.repository,
                 component.ref,
-                os.path.join(work_dir, component.name),
+                os.path.join(work_dir.name, component.name),
                 component.working_directory,
-            )
-
-            try:
-                ci = Ci(component, repo, target)
-                ci.check()
-            except:
-                logging.error(f"Error checking {component.name}, retry with: {args.component_command(component.name)}")
-                raise
+            ) as repo:
+                try:
+                    ci = Ci(component, repo, target)
+                    ci.check()
+                except:
+                    logging.error(f"Error checking {component.name}, retry with: {args.component_command(component.name)}")
+                    raise
 
     logging.info("Done.")
 

src/run_integ_test.py

@@ -41,26 +41,26 @@ def main():
         if component.integ_test is not None:
             integ_test_config[component.name] = component
     with TemporaryDirectory(keep=args.keep) as work_dir:
-        logging.info("Switching to temporary work_dir: " + work_dir)
-        test_recorder = TestRecorder(args.test_run_id, "integ-test", work_dir)
-        os.chdir(work_dir)
+        logging.info(f"Switching to temporary work_dir: {work_dir.name}")
+        test_recorder = TestRecorder(args.test_run_id, "integ-test", work_dir.name)
+        os.chdir(work_dir.name)
         bundle_manifest = BundleManifest.from_s3(
             args.s3_bucket,
             args.build_id,
             args.opensearch_version,
             args.platform,
             args.architecture,
-            work_dir,
+            work_dir.name,
         )
         build_manifest = BuildManifest.from_s3(
             args.s3_bucket,
             args.build_id,
             args.opensearch_version,
             args.platform,
             args.architecture,
-            work_dir,
+            work_dir.name,
         )
-        pull_build_repo(work_dir)
+        pull_build_repo(work_dir.name)
         DependencyInstaller(build_manifest.build).install_all_maven_dependencies()
         all_results = TestSuiteResults()
         for component in bundle_manifest.components:
@@ -70,7 +70,7 @@ def main():
                     integ_test_config[component.name],
                     bundle_manifest,
                     build_manifest,
-                    work_dir,
+                    work_dir.name,
                     args.s3_bucket,
                     test_recorder,
                 )

src/run_perf_test.py

@@ -42,8 +42,8 @@ def get_infra_repo_url():
 
 def main():
     with TemporaryDirectory(keep=args.keep) as work_dir:
-        os.chdir(work_dir)
-        current_workspace = os.path.join(work_dir, "infra")
+        os.chdir(work_dir.name)
+        current_workspace = os.path.join(work_dir.name, "infra")
         GitRepository(get_infra_repo_url(), "main", current_workspace)
         security = False
         for component in manifest.components:

src/system/temporary_directory.py

@@ -7,16 +7,19 @@
 import logging
 import shutil
 import tempfile
-from contextlib import contextmanager
 
 
-@contextmanager
-def TemporaryDirectory(keep=False):
-    name = tempfile.mkdtemp()
-    try:
-        yield name
-    finally:
-        if keep:
-            logging.info(f"Keeping {name}")
+class TemporaryDirectory:
+    def __init__(self, keep=False):
+        self.keep = keep
+        self.name = tempfile.mkdtemp()
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_value, exc_traceback):
+        if self.keep:
+            logging.info(f"Keeping {self.name}")
         else:
-            shutil.rmtree(name)
+            logging.debug(f"Removing {self.name}")
+            shutil.rmtree(self.name)

tests/test_run_assemble.py

@@ -35,10 +35,10 @@ def test_usage(self, *mocks):
     @patch("argparse._sys.argv", ["run_assemble.py", BUILD_MANIFEST])
     @patch("run_assemble.Bundles", return_value=MagicMock())
     @patch("run_assemble.BundleRecorder", return_value=MagicMock())
-    @patch("tempfile.TemporaryDirectory")
+    @patch("run_assemble.TemporaryDirectory")
     @patch("shutil.copy2")
     def test_main(self, mock_copy, mock_temp, mock_recorder, mock_bundles, *mocks):
-        mock_temp.return_value.__enter__.return_value = tempfile.gettempdir()
+        mock_temp.return_value.__enter__.return_value.name = tempfile.gettempdir()
         mock_bundle = MagicMock(archive_path="path")
         mock_bundles.create.return_value = mock_bundle