""" This submodule includes some useful utility functions
for exporting data.
"""
from traceback import print_exc
import os
import random
[docs]def file_writer_function(function):
""" Wrapper for file writers to safely overwrite/hash duplicate files.
Keyword arguments:
overwrite (bool): whether or not to overwrite colliding files.
hash_dupe (bool): whether or not to create a unique filename for
any colliding files, or just skip writing them.
"""
from functools import wraps
@wraps(function)
def wrapped_writer(doc, path, overwrite=False, hash_dupe=False, **kwargs):
""" Wrap and return the writer function. """
try:
overwrite = kwargs.pop('overwrite', overwrite)
hash_dupe = kwargs.pop('hash_dupe', hash_dupe)
flines, ext = function(doc, path, overwrite=overwrite, hash_dupe=hash_dupe, **kwargs)
if ext is not None and not path.endswith('.' + ext):
path += '.{}'.format(ext)
if os.path.isfile(path):
if overwrite:
os.remove(path)
elif hash_dupe:
print('File name already exists, generating hash...')
path = '{}-{}.{}'.format(path.replace(ext, ''), generate_hash(), ext)
else:
print('File name {} already exists! Skipping!'.format(path))
return False
with open(path, 'w') as f:
for line in flines:
f.write(line + '\n')
except Exception as exc:
print_exc()
raise type(exc)('Failed to write {}: {}'.format(path, exc))
return wrapped_writer
[docs]def generate_hash(hash_len=6):
""" Quick hash generator, based on implementation in PyAIRSS by J. Wynn.
Keyword arguments:
hash_len (int): desired length of hash.
"""
hash_chars = [str(x) for x in range(10)] + [chr(97+i) for i in range(25)]
_hash = ''
for _ in range(hash_len):
_hash += random.choice(hash_chars)
return _hash
[docs]def generate_relevant_path(**args):
""" Generates a suitable path name based on query. """
dirname = ''
if args.get('subcmd') is not None:
dirname += args.get('subcmd') + '-'
else:
dirname = 'query'
if args.get('composition') is not None:
for comp in args['composition']:
dirname += comp
elif args.get('formula') is not None:
dirname += args.get('formula')[0]
if args.get('db') is not None:
dirname += '-' + args.get('db')[0]
if args.get('swap') is not None:
for swap in args['swap']:
dirname += '-' + swap
if args.get('hull_cutoff') is not None:
dirname += '-hull-' + str(args.get('hull_cutoff')) + 'eV'
if args.get('id') is not None:
dirname += '-' + args.get('id')[0] + '_' + args.get('id')[1]
dirname = dirname.replace('--', '-')
return dirname