Better handling of printing results board

python/toolbox.py

@@ -2,12 +2,35 @@ from typing import Any, Callable, List, Union, Final
 from time import perf_counter_ns
 from numba import njit
 import numpy as np
+from sys import stderr
 import pickle
 import os
 from config import MODEL_DIR, OUT_DIR
 from decorators import njit
 
 time_formats: Final = ["ns", "µs", "ms", "s", "m", "h", "j", "w", "M", "y", "c"]
+def formatted_row(gaps: list[int], titles: list[str], separator: str = '│') -> None:
+	for gap, title in zip(gaps, titles):
+		print(f"{separator} {title:{'>' if gap < 0 else '<'}{abs(gap)}} ", end = '')
+	print(separator)
+
+def formatted_line(gaps: list[int], right: str, middle: str, separator: str, left: str) -> None:
+	print(right, end = '')
+	last_gap = len(gaps) - 1
+	for i, gap in enumerate(gaps):
+		print(f'{separator * (abs(gap) + 2)}', end = '')
+		if i != last_gap:
+			print(middle, end = '')
+	print(left)
+
+def header(titles: list[str], gaps: list[int]) -> None:
+	formatted_line(gaps, '┌', '┬', '─', '┐')
+	formatted_row(gaps, titles)
+	formatted_line(gaps, '├', '┼', '─', '┤')
+
+def footer(gaps: list[int]) -> None:
+	formatted_line(gaps, '└', '┴', '─', '┘')
+
 time_numbers: Final = np.array([1, 1e3, 1e6, 1e9, 6e10, 36e11, 864e11, 6048e11, 26784e11, 31536e12, 31536e14], dtype = np.uint64)
 @njit('str(uint64)')
 def format_time_ns(time: int) -> str:
@@ -53,7 +76,7 @@ def picke_multi_loader(filenames: List[str], save_dir: str = MODEL_DIR) -> List[
 			b.append(None)
 	return b
 
-def benchmark_function(step_name: str, fnc: Callable) -> Any:
+def benchmark_function(step_name: str, column_width: int, fnc: Callable) -> Any:
 	"""Benchmark a function and display the result of stdout.
 
 	Args:
@@ -63,14 +86,15 @@ def benchmark_function(step_name: str, fnc: Callable) -> Any:
 	Returns:
 		Any: Result of the function.
 	"""
-	print(f"{step_name}...", end = "\r")
+	print(f'{step_name}...', file = stderr, end = '\r')
 	s = perf_counter_ns()
 	b = fnc()
 	e = perf_counter_ns() - s
-	print(f"| {step_name:<49} | {e:>18,} | {format_time_ns(e):<29} |")
+	print(f'│ {step_name:<{column_width}} │ {e:>18,} │ {format_time_ns(e):<29} │')
 	return b
 
-def state_saver(step_name: str, filename: Union[str, List[str]], fnc, force_redo: bool = False, save_state: bool = True, save_dir: str = OUT_DIR) -> Any:
+def state_saver(step_name: str, column_width: int, filename: Union[str, List[str]], fnc, force_redo: bool = False,
+				save_state: bool = True, save_dir: str = OUT_DIR) -> Any:
 	"""Either execute a function then saves the result or load the already existing result.
 
 	Args:
@@ -85,18 +109,18 @@ def state_saver(step_name: str, filename: Union[str, List[str]], fnc, force_redo
 	"""
 	if isinstance(filename, str):
 		if not os.path.exists(f"{save_dir}/{filename}.pkl") or force_redo:
-			b = benchmark_function(step_name, fnc)
+			b = benchmark_function(step_name, column_width, fnc)
 			if save_state:
-				print(f"Saving results of {step_name}", end = '\r')
 				with open(f"{save_dir}/{filename}.pkl", 'wb') as f:
+				print(f'Saving results of {step_name}', file = stderr, end = '\r')
 					pickle.dump(b, f)
-				print(' ' * 100, end = '\r')
+				print(' ' * 100, file = stderr, end = '\r')
 			return b
 		else:
-			print(f"Loading results of {step_name}", end = '\r')
 			with open(f"{save_dir}/{filename}.pkl", "rb") as f:
+			print(f'Loading results of {step_name}', file = stderr, end = '\r')
 				res = pickle.load(f)
-			print(f"| {step_name:<49} | {'None':>18} | {'loaded saved state':<29} |")
+			print(f"│ {step_name:<{column_width}} │ {'None':>18} │ {'loaded saved state':<29} │")
 			return res
 	elif isinstance(filename, list):
 		abs = False
@@ -105,22 +129,22 @@ def state_saver(step_name: str, filename: Union[str, List[str]], fnc, force_redo
 				abs = True
 				break
 		if abs or force_redo:
-			b = benchmark_function(step_name, fnc)
+			b = benchmark_function(step_name, column_width, fnc)
 			if save_state:
-				print(f"Saving results of {step_name}", end = '\r')
+				print(f'Saving results of {step_name}', file = stderr, end = '\r')
 				for bi, fnI in zip(b, filename):
 					with open(f"{save_dir}/{fnI}.pkl", 'wb') as f:
 						pickle.dump(bi, f)
-				print(' ' * 100, end = '\r')
+				print(' ' * 100, file = stderr, end = '\r')
 			return b
 
-		print(f"| {step_name:<49} | {'None':>18} | {'loaded saved state':<29} |")
+		print(f"│ {step_name:<{column_width}} │ {'None':>18} │ {'loaded saved state':<29} │")
 		b = []
-		print(f"Loading results of {step_name}", end = '\r')
+		print(f'Loading results of {step_name}', file = stderr, end = '\r')
 		for fn in filename:
 			with open(f"{save_dir}/{fn}.pkl", "rb") as f:
 				b.append(pickle.load(f))
-		print(' ' * 100, end = '\r')
+		print(' ' * 100, file = stderr, end = '\r')
 		return b
 	else:
 		assert False, f"Incompatible filename type = {type(filename)}"