# coding: utf-8
"""IO related utilities."""
import os
from contextlib import ExitStack
from subprocess import call
from monty.termcolor import cprint
from monty.string import list_strings
[docs]class ExitStackWithFiles(ExitStack):
"""
Context manager for dynamic management of a stack of file-like objects.
Mainly used in a callee that needs to return files to the caller
Usage example:
.. code-block:: python
exit_stack = ExitStackWithFiles()
exit_stack.enter_context(phbst_file)
return exit_stack
"""
def __init__(self):
self.files = []
super().__init__()
[docs] def enter_context(self, myfile):
# If my file is None, we add it to files but without registering the callback.
self.files.append(myfile)
if myfile is not None:
return super().enter_context(myfile)
def __iter__(self):
return self.files.__iter__()
def __next__(self):
return self.files.__next__()
def __getitem__(self, slice):
return self.files.__getitem__(slice)
[docs]def ask_yes_no(prompt, default=None): # pragma: no cover
"""
Ask a question and return a boolean (y/n) answer.
If default is given (one of 'y','n'), it is used if the user input is
empty. Otherwise the question is repeated until an answer is given.
An EOF is treated as the default answer. If there is no default, an
exception is raised to prevent infinite loops.
Valid answers are: y/yes/n/no (match is not case sensitive).
"""
# Fixes py2.x
try:
my_input = raw_input
except NameError:
my_input = input
answers = {'y': True, 'n': False, 'yes': True, 'no': False}
ans = None
while ans not in answers.keys():
try:
ans = my_input(prompt + ' ').lower()
if not ans:
# response was an empty string
ans = default
except KeyboardInterrupt:
pass
except EOFError:
if default in answers.keys():
ans = default
print("")
else:
raise
return answers[ans]
def _user_wants_to_exit(): # pragma: no cover
# Fixes py2.x
try:
my_input = raw_input
except NameError:
my_input = input
try:
answer = my_input("Do you want to continue [Y/n]")
except EOFError:
return True
if answer.lower().strip() in ["n", "no"]: return True
return False
[docs]class EditorError(Exception):
"""Base class for exceptions raised by `Editor`"""
[docs]class Editor(object): # pragma: no cover
DEFAULT_EDITOR = "vi"
Error = EditorError
def __init__(self, editor=None):
if editor is None:
self.editor = os.getenv("EDITOR", self.DEFAULT_EDITOR)
else:
self.editor = str(editor)
[docs] def edit_file(self, fname):
retcode = call([self.editor, fname])
if retcode != 0:
cprint("Retcode %s while editing file: %s" % (retcode, fname), "red")
return retcode
[docs] def edit_files(self, fnames, ask_for_exit=True):
for idx, fname in enumerate(list_strings(fnames)):
exit_status = self.edit_file(fname)
if exit_status != 0:
return exit_status
if ask_for_exit and idx != len(fnames) - 1 and _user_wants_to_exit():
break
return 0