From 1f5ef587dd1af9a86634101ef17efe2aec4f681e Mon Sep 17 00:00:00 2001 From: The MMGen Project Date: Thu, 6 Jan 2022 20:24:21 +0000 Subject: [PATCH] util.py: open_file_or_exit() -> open_or_die() --- mmgen/util.py | 39 +++++++++++++++++++++++---------------- 1 file changed, 23 insertions(+), 16 deletions(-) diff --git a/mmgen/util.py b/mmgen/util.py index 966f6838..660822f1 100755 --- a/mmgen/util.py +++ b/mmgen/util.py @@ -478,12 +478,15 @@ def shred_file(fn,verbose=False): + [fn], check=True ) -def open_file_or_exit(filename,mode,silent=False): +def open_or_die(filename,mode,silent=False): try: - return open(filename, mode) + return open(filename,mode) except: - op = ('writing','reading')['r' in mode] - die(2,'' if silent else f'Unable to open file {filename!r} for {op}') + die(2,'' if silent else + 'Unable to open file {!r} for {}'.format( + ({0:'STDIN',1:'STDOUT',2:'STDERR'}[filename] if type(filename) == int else filename), + ('reading' if 'r' in mode else 'writing') + )) def check_file_type_and_access(fname,ftype,blkdev_ok=False): @@ -646,13 +649,11 @@ def write_data_to_file( outfile,data,desc='data', # To maintain portability, always open files in binary mode # If 'binary' option not set, encode/decode data before writing and after reading - f = open_file_or_exit(outfile,'wb') - try: - f.write(data if binary else data.encode()) + with open_or_die(outfile,'wb') as fp: + fp.write(data if binary else data.encode()) except: die(2,f'Failed to write {desc} to file {outfile!r}') - f.close if not (hush or quiet): msg(f'{capfirst(desc)} written to file {outfile!r}') @@ -674,10 +675,15 @@ def get_words_from_user(prompt): def get_words_from_file(infile,desc,quiet=False): if not quiet: qmsg(f'Getting {desc} from file {infile!r}') - f = open_file_or_exit(infile, 'rb') - try: words = f.read().decode().split() - except: die(1,f'{capfirst(desc)} data must be UTF-8 encoded.') - f.close() + + with open_or_die(infile, 'rb') as fp: + data = fp.read() + + try: + words = data.decode().split() + except: + die(1,f'{capfirst(desc)} data must be UTF-8 encoded.') + dmsg('Sanitized input: [{}]'.format(' '.join(words))) return words @@ -715,10 +721,11 @@ def get_data_from_file(infile,desc='data',dash=False,silent=False,binary=False,q if not opt.quiet and not silent and not quiet and desc: qmsg(f'Getting {desc} from file {infile!r}') - if dash and infile == '-': - data = os.fdopen(0,'rb').read(g.max_input_size+1) - else: - data = open_file_or_exit(infile,'rb',silent=silent).read(g.max_input_size+1) + with open_or_die( + (0 if dash and infile == '-' else infile), + 'rb', + silent=silent) as fp: + data = fp.read(g.max_input_size+1) if not binary: data = data.decode()