def __init__(self, key, mode='r', connection=None, encrypt=True, version_id=None): from baiji.connection import S3Connection self.encrypt = encrypt self.key = key if path.islocal(key): self.should_upload_on_close = False self.mode = FileMode(mode, allowed_modes='arwxb+t') from six.moves import builtins local_path = path.parse(key).path if self.mode.is_output and not os.path.exists( os.path.dirname(local_path)): from baiji.util.shutillib import mkdir_p mkdir_p(os.path.dirname(local_path)) try: # Use os.open to catch exclusive access to the file, but use open to get a nice, useful file object self.fd = os.open(local_path, self.mode.flags) self.f = builtins.open(local_path, self.mode.mode.replace('x', 'w')) os.close(self.fd) except OSError as e: import errno if e.errno is errno.EEXIST: raise KeyExists("Local file exists: %s" % local_path) elif e.errno is errno.ENOENT: raise KeyNotFound("Local file does not exist: %s" % local_path) else: raise IOError(e.errno, "%s: %s" % (e.strerror, e.filename)) else: if connection is None: connection = S3Connection() self.connection = connection self.mode = FileMode(mode, allowed_modes='rwxbt') self.should_upload_on_close = self.mode.is_output if self.mode.creating_exclusively: if self.connection.exists(self.key): raise KeyExists("Key exists in bucket: %s" % self.key) else: self.connection.touch(self.key, encrypt=self.encrypt) # Use w+ so we can read back the contents in upload() new_mode = ('w+' + (self.mode.binary and 'b' or '') + (self.mode.text and 't' or '')) from baiji.util import tempfile self.f = tempfile.NamedTemporaryFile( mode=new_mode, suffix=os.path.splitext(path.parse(self.key).path)[1]) self.name = self.f.name self.remotename = key # Used by some serialization code to find files which sit along side the file in question, like textures which sit next to a mesh file if self.mode.reading: self.connection.cp(self.key, self.name, force=True, version_id=version_id)
def test_that_NamedTemporaryFile_honors_TMP_env_var(self): env = EnvironmentVarGuard() test_temp_dir = 'test_tempfile_' + str(uuid.uuid4()) env.set('BAIJI_TMP', test_temp_dir) with env: with tempfile.NamedTemporaryFile('w') as tf: self.assertEquals(os.path.dirname(tf.name), os.path.join(os.getcwd(), test_temp_dir)) shutil.rmtree(test_temp_dir)
def create_random_temporary_file(sz=None): ''' Create a temporary file with random contents, and return its path. The caller is responsible for removing the file when done. sz is optional and approximate ''' from baiji.util import tempfile with tempfile.NamedTemporaryFile('w', delete=False) as f: f.write(random_data(sz)) return f.name
def test_md5(self): import os import struct from baiji.util import tempfile from baiji.util.md5 import md5_for_file f = tempfile.NamedTemporaryFile('wb', delete=False) for _ in range(2**16): f.write(struct.pack("<I", 42)) f.close() self.assertEqual(md5_for_file(f.name), "9b0cbf6ba5c0ad5606df9b9630a44db6") os.remove(f.name)
def upload(self): if self.gzip: import gzip from baiji.util import tempfile with tempfile.NamedTemporaryFile() as compressed_src: with open(self.src.path, 'rb') as f: with gzip.open(compressed_src.name, 'wb') as tf: tf.writelines(f) self.src.path = compressed_src.name self.encoding = 'gzip' self._upload() else: self._upload()
def download(self): ''' Download to local file If `validate` is set, check etags and retry once if not match. Raise TransientError when download is corrupted again after retry ''' import shutil from baiji.util import tempfile from baiji.util.with_progressbar import FileTransferProgressbar # We create, close, and delete explicitly rather than using # a `with` block, since on windows we can't have a file open # twice by the same process. tf = tempfile.NamedTemporaryFile(delete=False) try: key = self.src.lookup(version_id=self.version_id) with FileTransferProgressbar(supress=(not self.progress)) as cb: key.get_contents_to_file(tf, cb=cb) tf.close() if self.validate: self.ensure_integrity(tf.name) # We only actually write to dst.path if the download succeeds and # if necessary is validated. This avoids leaving partially # downloaded files if something goes wrong. shutil.copy(tf.name, self.dst.path) except (get_transient_error_class(), KeyNotFound) as retryable_error: # Printed here so that papertrail can alert us when this occurs print(retryable_error) # retry once or raise if self.retries_made < self.retries_allowed: self._retries += 1 self.download() else: raise finally: self.connection.rm(tf.name)