def download_if_needed(uri, download_dir, fs=None):
"""Download a file into a directory if it's remote.
If uri is local, there is no need to download the file.
Args:
uri: (string) URI of file
download_dir: (string) local directory to download file into
fs: Optional FileSystem to use.
Returns:
(string) path to local file
Raises:
NotReadableError if URI cannot be read from
"""
if uri is None:
return None
if not fs:
fs = FileSystem.get_file_system(uri, 'r')
path = get_local_path(uri, download_dir, fs=fs)
make_dir(path, use_dirname=True)
if path != uri:
log.debug('Downloading {} to {}'.format(uri, path))
fs.copy_from(uri, path)
return path
def set_tmp_dir(tmp_dir=None):
"""Set RVConfig.tmp_dir to well-known value.
This static method sets the value of RVConfig.tmp_dir to some
well-known value. The value is chosen from one of the
following (in order of preference): an explicit value
(presumably from the command line) is considered first, then
values from the environment are considered, then the current
value of RVConfig.tmp_dir is considered, then a directory from
tempfile.TemporaryDirectory() is considered.
Args:
tmp_dir: Either a string or None.
"""
DEFAULT_DIR = '/opt/data/tmp/'
# Check the various possibilities in order of priority.
tmp_dir_array = [tmp_dir]
env_array = [
os.environ.get(k) for k in ['TMPDIR', 'TEMP', 'TMP']
if k in os.environ
]
current_array = [RVConfig.tmp_dir]
it = tmp_dir_array + env_array + current_array
it = list(filter(lambda p: p is not None, it))
if it:
explicit_tmp_dir = it[0]
else:
explicit_tmp_dir = tempfile.TemporaryDirectory().name
try:
# Try to create directory
if not os.path.exists(explicit_tmp_dir):
os.makedirs(explicit_tmp_dir, exist_ok=True)
# Check that it is actually a directory
if not os.path.isdir(explicit_tmp_dir):
raise Exception(
'{} is not a directory.'.format(explicit_tmp_dir))
# Can we interact with directory?
Path.touch(Path(os.path.join(explicit_tmp_dir, '.can_touch')))
# All checks have passed by this point
RVConfig.tmp_dir = explicit_tmp_dir
# If directory cannot be made and/or cannot be interacted
# with, fall back to default.
except Exception as e:
log.warning(
'Root temporary directory cannot be used: {}. Using root: {}'.
format(explicit_tmp_dir, DEFAULT_DIR))
RVConfig.tmp_dir = DEFAULT_DIR
finally:
make_dir(RVConfig.tmp_dir)
log.debug('Temporary directory is: {}'.format(RVConfig.tmp_dir))
def zipdir(dir, zip_path):
"""Save a zip file with contents of directory.
Contents of directory will be at root of zip file.
Args:
dir: (str) directory to zip
zip_path: (str) path to zip file to create
"""
make_dir(zip_path, use_dirname=True)
with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as ziph:
for dirpath, dirnames, filenames in os.walk(dir):
for fn in filenames:
ziph.write(join(dirpath, fn), join(dirpath[len(dir):], fn))
def unzip(zip_path, target_dir):
"""Unzip contents of zip file at zip_path into target_dir"""
make_dir(target_dir)
zip_ref = zipfile.ZipFile(zip_path, 'r')
zip_ref.extractall(target_dir)
zip_ref.close()
