diff --git a/opstool/utils/__init__.py b/opstool/utils/__init__.py index 5e9a317..4021729 100644 --- a/opstool/utils/__init__.py +++ b/opstool/utils/__init__.py @@ -1,7 +1,7 @@ from ._load_ops_examples import load_ops_examples, run_model from .ops_ele_class_tags import OPS_ELE_TAGS, OPS_ELE_CLASSTAG2TYPE, OPS_ELE_TYPES from .consts import CONSOLE, PKG_NAME, PKG_PREFIX, SHAPE_MAP, RESULTS_DIR -from ._util_funcs import add_ops_hints_file, print_version, check_file_type +from ._util_funcs import add_ops_hints_file, print_version, check_file_type, suppress_ops_print from ._util_funcs import get_color_rich, get_cycle_color_rich, get_random_color from ._util_funcs import get_cycle_color, get_random_color_rich, gram_schmidt @@ -10,6 +10,7 @@ "run_model", "print_version", "add_ops_hints_file", + "suppress_ops_print", # -------------------- "OPS_ELE_TAGS", "OPS_ELE_CLASSTAG2TYPE", diff --git a/opstool/utils/_util_funcs.py b/opstool/utils/_util_funcs.py index ebda08f..cc3bb22 100644 --- a/opstool/utils/_util_funcs.py +++ b/opstool/utils/_util_funcs.py @@ -1,8 +1,10 @@ +import os import sys import numpy as np from pathlib import Path from typing import Union from itertools import cycle +from contextlib import contextmanager from .consts import CONSOLE, PKG_PREFIX @@ -146,3 +148,21 @@ def gram_schmidt(v1, v2): y = y / np.linalg.norm(y) z = z / np.linalg.norm(z) return x, y, z + + +# Context manager to temporarily suppress stdout and stderr +@contextmanager +def suppress_ops_print(): + # Save the original stdout and stderr + stdout = sys.stdout + stderr = sys.stderr + try: + # Redirect stdout and stderr to null (discard output) + with open(os.devnull, 'w') as fnull: + sys.stdout = fnull + sys.stderr = fnull + yield + finally: + # Restore the original stdout and stderr + sys.stdout = stdout + sys.stderr = stderr \ No newline at end of file