Sleds/scorewalker-utils/ConsoleUtils/ConsoleUtils.py

239 lines
8.2 KiB
Python

"""
Information
-----------
This is a small project to help with some console functionality. It allows for printing a progress bar to the console
as well as provides a utility for printing console output with the program name pre-pended to the output.
Python Module Usage
-------------------
"""
import os
import sys
import argparse
import linecache
from datetime import datetime as dt
_COLOR_START = '\e[38;5;%dm'
_COLOR_END = '\e[0m'
RED = 160
BLUE = 14
GREEN = 46
PINK = 165
YELLOW = 226
def make_color(text, color_num):
if color_num > 256 or color_num < 0:
return
colored_text = '%s%s%s' % (_COLOR_START % color_num, text, _COLOR_END)
return colored_text
class CustomPrintAction(argparse.Action):
def __init__(self, option_strings, print_fn, help=None, dest=None):
self.version_fn = print_fn
super(CustomPrintAction, self).__init__(option_strings, dest=argparse.SUPPRESS, default=argparse.SUPPRESS,
nargs=0, help=help)
def __call__(self, parser, args, values, option_string=None):
self.version_fn()
parser.exit()
class SLPrinter:
def __init__(self, prog_name, timestamps=True, out=sys.stdout, tabsize=3, maxlen=None):
self.old_stdout = out
self.prog_name = prog_name
self.log_time = timestamps
self.date_format = '%Y/%m/%d %H:%M:%S'
self.tab_str = ' ' * tabsize
self.max_len = maxlen
def write_tabbed(self, text, num_indents):
self._write_to_output(self._make_text(text, num_indents))
def write(self, text):
text = self._make_text(text)
if len(text) == 0:
return
self._write_to_output(text)
def write_no_prefix(self, text):
if not text.endswith('\n'):
text += '\n'
self._write_to_output(text)
def flush(self):
self.old_stdout.flush()
def write_line_break(self, break_char='-', line_size=80):
self._write_to_output('%s\n' % (break_char * line_size))
def _make_text(self, text, tab_factor=0):
text = text.replace('\n', '')
# If we just want a blank line...
raw_txt_len = len(text)
if raw_txt_len == 0:
return ''
if self.max_len is not None:
if raw_txt_len > self.max_len:
first_line = text[:self.max_len]
second_line = text[self.max_len:]
while second_line.startswith(' '):
second_line = second_line[1:]
text = '%s\n%s' % (first_line, self._make_text(second_line, tab_factor=tab_factor))
text = '[%s]%s %s\n' % (self.prog_name, (self.tab_str * tab_factor), text.rstrip())
if self.log_time:
text = '%s %s' % (dt.now().strftime(self.date_format), text)
return text
# This is the only part which needs to be overridden to add logging to files or other output piping/analysis
def _write_to_output(self, text):
self.old_stdout.write(text)
class SLLogger(SLPrinter):
def __init__(self, prog_name, log_path, timestamps=True, out=sys.stdout):
super().__init__(prog_name, timestamps, out)
self.log_file = log_path
# Overriding this here because we want to save the output to a file.
def __write_to_output__(self, text):
super().__write_to_output__(text)
with open(self.log_file, 'a+') as logger:
logger.write(text)
# Returns a green progress bar with a percentage on the left.
def get_prg_line(current, total, bar_size, line_size, percent_decimals=2):
num_to_repeat = int(bar_size*(current/total))
if num_to_repeat == 0:
num_to_repeat = 1
if current == total:
tmp_result = '{0:} '.format('Done!') + '\033[102m' + (' ' * num_to_repeat) + '\033[0m' + \
(' ' * (bar_size - num_to_repeat) + '\n')
else:
tmp_result = ('{1:.{0}%} '.format(percent_decimals, current/total) + '\033[102m' + (' ' * num_to_repeat) +
('\033[0m' + (' ' * (bar_size - num_to_repeat))))
final_result = ('{:^%d}' % line_size).format(tmp_result)
return final_result
# Writes a progress bar to the console with the given info and a % on the left.
# This will update if it is the only line printed!
def print_progress_bar(current, total, bar_size, line_size, precision=2):
sys.stdout.write('\r%s' % ' ' * line_size)
sys.stdout.write('\r%s' % get_prg_line(current, total, bar_size, line_size, precision))
def yes_or_no(message):
"""
Will prompt the user with the given message and accept either ``"y"``, ``"yes"``, ``"n"``, or ``"no"`` as inputs
ignoring case. The program will exit (with status 0) if the user enters ``"no"`` and will continue if the user
enters ``"yes"``.
Args:
``message`` -- ``str`` The message to display.
Returns:
``None``
*NOTE*: The user will be prompted to re-enter input if it is not valid with the following message (until valid
input is provided):
``"Invalid input, enter Y(es) or N(o):"``
"""
decision = input('%s yes/no: ' % message)
if decision.lower() == 'y' or decision.lower() == 'yes':
return
elif decision.lower() == 'n' or decision.lower() == 'no':
print('Exiting...')
exit(0)
else:
yes_or_no(' Invalid input, enter Y(es) or N(o):\n%s' % message)
def print_exception():
err_type, err_obj, tb = sys.exc_info()
if err_type is None or err_obj is None or tb is None:
return
print('%s: %s' % (err_type, str(err_obj)))
while tb.tb_next is not None:
tb = tb.tb_next
frame = tb.tb_frame
file_name = os.path.split(frame.f_code.co_filename)[1]
line_num = tb.tb_lineno
linecache.checkcache(file_name)
line = linecache.getline(file_name, line_num, frame.f_globals).strip()
line = make_color(line, PINK)
print('In "%s" on line %d: %s' % (file_name, line_num, line))
print('')
def get_header(program_name, program_version=None, build_date=None, program_author=None, line_width=80):
"""
Generates a header in the following format:
::
+----------------------------------------+
| program_name |
+----------------------------------------+
| Version: program_version |
| Build: build_date |
| Author: program_author |
| Copyright (c) 2017 Sequence Logic, LLC |
+----------------------------------------+
Args:
``program_name`` -- ``str`` The name of the program.
``program_version`` -- ``str`` The version of the program.
``build_date`` -- ``str`` The build date of the program.
``program_author`` -- ``str`` The author of the program.
``line_width`` -- ``int`` The width to make the header (default = 80).
Returns:
``str`` A formatted header for the program which can be printed.
If ``program_version``, ``build_date``, or ``program_author`` are not included, the line containing them will not be
in the result.
"""
# Do some set up
line_break = '-' * line_width
center_str = '|{:^%d}|\n' % (line_width - 2)
corner_str = '+%s+\n' % (line_break[:-2])
# Build the header string
result = corner_str
result += center_str.format(program_name)
result += corner_str
if program_version is not None:
result += center_str.format('Version: %s' % program_version)
if build_date is not None:
result += center_str.format('Build: %s' % build_date)
if program_author is not None:
result += center_str.format('Author: %s' % program_author)
result += center_str.format('Copyright (c) 2017 Sequence Logic, LLC')
result += corner_str + '\n'
# Give it back
return result
# if __name__ == '__main__':
# printer = SLPrinter('Testing', maxlen=50)
# sys.stdout = printer
# print('This is a test to see if this will work the way it should and cut the text off 50 characters in')
# print('111111111111111111112222222222222222222244444444444444444444')
# print('Hi')
#
# printer.write_tabbed('This is a test to make sure tabbing works the way it should while breaking lines!', 1)