From 8b550a9fba4a169da463b2b034ef0960b46aa1bc Mon Sep 17 00:00:00 2001 From: Tristan Van Berkom Date: Sat, 8 Apr 2017 19:50:22 +0900 Subject: source.py: Added tempdir() context manager This is for Source implementations to use a temp directory, this should be used instead of python tempdir APIs because it will cleanup the tempdir automatically in the case of user provoked termination. --- buildstream/source.py | 28 +++++++++++++++++++++++++++- 1 file changed, 27 insertions(+), 1 deletion(-) (limited to 'buildstream/source.py') diff --git a/buildstream/source.py b/buildstream/source.py index fb4c9d95a..dcb5359bc 100644 --- a/buildstream/source.py +++ b/buildstream/source.py @@ -19,8 +19,11 @@ # Tristan Van Berkom import os +import tempfile +import shutil +from contextlib import contextmanager -from . import _yaml +from . import _yaml, _signals from . import ImplError from . import Plugin @@ -81,6 +84,29 @@ class Source(Plugin): os.makedirs(directory, exist_ok=True) return directory + @contextmanager + def tempdir(self): + """Context manager for working in a temporary directory + + Yields: + (str): A path to a temporary directory + + This should be used by source plugins directly instead of the + tempfile module, as it will take care of cleaning up the temporary + directory in the case of forced termination. + """ + mirrordir = self.get_mirror_directory() + tempdir = tempfile.mkdtemp(dir=mirrordir) + + def cleanup_tempdir(): + if os.path.isdir(tempdir): + shutil.rmtree(tempdir) + + with _signals.terminator(cleanup_tempdir): + yield tempdir + + cleanup_tempdir() + def get_consistency(self): """Report whether the source has a resolved reference -- cgit v1.2.1