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)
