diff options
-rw-r--r-- | morphlib/ostreeartifactcache.py | 48 |
1 files changed, 25 insertions, 23 deletions
diff --git a/morphlib/ostreeartifactcache.py b/morphlib/ostreeartifactcache.py index 64c6e8f8..0c00c659 100644 --- a/morphlib/ostreeartifactcache.py +++ b/morphlib/ostreeartifactcache.py @@ -15,6 +15,7 @@ import collections +import contextlib import logging import os import shutil @@ -40,6 +41,7 @@ class OSTreeArtifactCache(object): if self.status_cb is not None: self.status_cb(*args, **kwargs) + @contextlib.contextmanager def _get_file_from_remote(self, artifact, remote, metadata_name=None): if metadata_name: handle = remote.get_artifact_metadata(artifact, metadata_name) @@ -52,10 +54,12 @@ class OSTreeArtifactCache(object): msg='Downloading %(name)s as a tarball.', chatty=True, name=artifact.name) - fd, path = tempfile.mkstemp() - with open(path, 'w+') as temp: - shutil.copyfileobj(handle, temp) - return path + try: + temporary_download = tempfile.NamedTemporaryFile(dir=self.cachedir) + shutil.copyfileobj(handle, temporary_download) + yield temporary_download.name + finally: + temporary_download.close() def _get_artifact_cache_name(self, artifact): logging.debug('LAC: %s' % artifact.basename()) @@ -90,25 +94,23 @@ class OSTreeArtifactCache(object): else: filename = self.artifact_filename(artifact) shutil.copy(location, filename) - os.remove(location) def copy_from_remote(self, artifact, remote): """Get 'artifact' from remote artifact cache and store it locally.""" if remote.method == 'tarball': - location = self._get_file_from_remote(artifact, remote) - try: - tempdir = tempfile.mkdtemp() - with tarfile.open(name=location) as tf: - tf.extractall(path=tempdir) + with self._get_file_from_remote(artifact, remote) as location: try: - self.put(tempdir, artifact) - finally: - os.remove(location) - shutil.rmtree(tempdir) - except tarfile.ReadError: - # Reading the artifact as a tarball failed, so it must be a - # single file (for example a stratum artifact). - self.put_non_ostree_artifact(artifact, location) + tempdir = tempfile.mkdtemp(dir=self.cachedir) + try: + with tarfile.open(name=location) as tf: + tf.extractall(path=tempdir) + self.put(tempdir, artifact) + finally: + shutil.rmtree(tempdir) + except tarfile.ReadError: + # Reading the artifact as a tarball failed, so it must be a + # single file (for example a stratum artifact). + self.put_non_ostree_artifact(artifact, location) elif remote.method == 'ostree': self.status(msg='Pulling artifact for %(name)s from remote.', @@ -118,14 +120,14 @@ class OSTreeArtifactCache(object): except Exception: # if we can't split the name properly, we must want metadata a, name = artifact.basename().split('.', 1) - location = self._get_file_from_remote( - ArtifactCacheReference(a), remote, name) - self.put_non_ostree_artifact(artifact, location, name) + with self._get_file_from_remote(ArtifactCacheReference(a), + remote, name) as location: + self.put_non_ostree_artifact(artifact, location, name) return if artifact.basename().split('.', 2)[1] == 'stratum': - location = self._get_file_from_remote(artifact, remote) - self.put_non_ostree_artifact(artifact, location) + with self._get_file_from_remote(artifact, remote) as location: + self.put_non_ostree_artifact(artifact, location) return try: |