T8154: make catch_broken_pipe decorator to common vyos.utils.io function

This adds the capability to make it re-usable in other code paths.

1 files changed, 11 insertions, 0 deletions

diff --git a/python/vyos/utils/io.py b/python/vyos/utils/io.py
index 0883376d1..3efb55ddc 100644
--- a/python/vyos/utils/io.py
+++ b/python/vyos/utils/io.py
@@ -111,3 +111,14 @@ def select_entry(l: list, list_msg: str = '', prompt_msg: str = '',
     select = ask_input(prompt_msg, default=default_entry, numeric_only=True,
                        valid_responses=valid_entry)
     return next(filter(lambda x: x[0] == select, en))[1]
+
+def catch_broken_pipe(func):
+    import os
+    import sys
+    def wrapped(*args, **kwargs):
+        try:
+            func(*args, **kwargs)
+        except (BrokenPipeError, KeyboardInterrupt):
+            # Flush output to /dev/null and bail out.
+            os.dup2(os.open(os.devnull, os.O_WRONLY), sys.stdout.fileno()) # pylint: disable = no-member
+    return wrapped