code review - cont.

Signed-off-by: imarom <imarom@cisco.com>

1 files changed, 6 insertions, 10 deletions

diff --git a/scripts/automation/trex_control_plane/stl/trex_stl_lib/utils/text_opts.py b/scripts/automation/trex_control_plane/stl/trex_stl_lib/utils/text_opts.py
index 3ffd07e2..477d81a6 100644
--- a/scripts/automation/trex_control_plane/stl/trex_stl_lib/utils/text_opts.py
+++ b/scripts/automation/trex_control_plane/stl/trex_stl_lib/utils/text_opts.py
@@ -133,16 +133,12 @@ def underline(text):
 
 # apply attribute on each non-empty line
 def text_attribute(text, attribute):
-    if isinstance(text, str):
-        return "{start}{txt}{stop}".format(start=TEXT_CODES[attribute]['start'],
-                                           txt=text,
-                                           stop=TEXT_CODES[attribute]['end'])
-    elif isinstance(text, unicode):
-        return u"{start}{txt}{stop}".format(start=TEXT_CODES[attribute]['start'],
-                                            txt=text,
-                                            stop=TEXT_CODES[attribute]['end'])
-    else:
-        raise Exception("not a string")
+    return '\n'.join(['{start}{txt}{end}'.format(
+        start = TEXT_CODES[attribute]['start'],
+        txt = line,
+        end = TEXT_CODES[attribute]['end'])
+                      if line else '' for line in ('%s' % text).split('\n')])
+
 
 FUNC_DICT = {'blue': blue,
              'bold': bold,