From 4f152256652ee1c93e16e3fd2fc23e71fc7bcd3c Mon Sep 17 00:00:00 2001
From: William Ballenthin <william.ballenthin@fireeye.com>
Date: Fri, 27 Aug 2021 08:34:02 -0600
Subject: [PATCH] lint: handle calls to print within pbar

---
 scripts/lint.py | 49 ++++++++++++++++++++++++++++++++++++++++++++-----
 1 file changed, 44 insertions(+), 5 deletions(-)

diff --git a/scripts/lint.py b/scripts/lint.py
index c71c78c8..fcf3e627 100644
--- a/scripts/lint.py
+++ b/scripts/lint.py
@@ -19,12 +19,15 @@ import time
 import string
 import difflib
 import hashlib
+import inspect
 import logging
 import os.path
 import argparse
 import itertools
 import posixpath
+import contextlib
 
+import tqdm
 import termcolor
 import ruamel.yaml
 import tqdm.contrib.logging
@@ -686,6 +689,41 @@ def lint_rule(ctx, rule):
     return (lints_failed, lints_warned)
 
 
+def width(s, count):
+    if len(s) > count:
+        return s[: count - 3] + "..."
+    else:
+        return s.ljust(count)
+
+
+@contextlib.contextmanager
+def redirecting_print_to_tqdm():
+    """
+    tqdm (progress bar) expects to have fairly tight control over console output.
+    so calls to `print()` will break the progress bar and make things look bad.
+    so, this context manager temporarily replaces the `print` implementation
+    with one that is compatible with tqdm.
+
+    via: https://stackoverflow.com/a/42424890/87207
+    """
+    old_print = print
+
+    def new_print(*args, **kwargs):
+
+        # If tqdm.tqdm.write raises error, use builtin print
+        try:
+            tqdm.tqdm.write(*args, **kwargs)
+        except:
+            old_print(*args, **kwargs)
+
+    try:
+        # Globaly replace print with new_print
+        inspect.builtins.print = new_print
+        yield
+    finally:
+        inspect.builtins.print = old_print
+
+
 def lint(ctx, rules):
     """
     Args:
@@ -701,12 +739,13 @@ def lint(ctx, rules):
     ret = {}
 
     with tqdm.contrib.logging.tqdm_logging_redirect(rules.rules.items(), unit="rule") as pbar:
-        for name, rule in pbar:
-            if rule.meta.get("capa/subscope-rule", False):
-                continue
+        with redirecting_print_to_tqdm():
+            for name, rule in pbar:
+                if rule.meta.get("capa/subscope-rule", False):
+                    continue
 
-            pbar.set_description("linting rule: %s" % (name))
-            ret[name] = lint_rule(ctx, rule)
+                pbar.set_description(width("linting rule: %s" % (name), 48))
+                ret[name] = lint_rule(ctx, rule)
 
     return ret