Add color and label feature (#21)

* Add color and label feature Co-authored-by: Carinaqyy <[email protected]> Co-authored-by: Yuying Quan <[email protected]>
gaogaotiantian · Apr 24, 2021 · d00438b · d00438b
1 parent 55d81c4
commit d00438b
Show file tree

Hide file tree

Showing 8 changed files with 89 additions and 15 deletions.
diff --git a/src/objprint/color_util.py b/src/objprint/color_util.py
@@ -0,0 +1,18 @@
+# Licensed under the Apache License: http://www.apache.org/licenses/LICENSE-2.0
+# For details: https://github.com/gaogaotiantian/objprint/blob/master/NOTICE.txt
+
+
+class COLOR:
+    BLACK = '\033[30m'
+    RED = '\033[31m'
+    GREEN = '\033[32m'
+    YELLOW = '\033[33m'
+    BLUE = '\033[34m'
+    MAGENTA = '\033[35m'
+    CYAN = '\033[36m'
+    WHITE = '\033[37m'
+    DEFAULT = '\033[39m'
+
+
+def set_color(s, color):
+    return f"{color}{s}{COLOR.DEFAULT}"
diff --git a/src/objprint/objprint.py b/src/objprint/objprint.py
@@ -4,6 +4,7 @@
 
 import re
 from types import FunctionType
+from .color_util import COLOR, set_color
 
 
 class ObjPrint:
@@ -12,6 +13,8 @@ def __init__(self):
             "indent": 2,
             "depth": 3,
             "width": 80,
+            "color": True,
+            "label": [],
             "elements": None
         }
 
@@ -67,6 +70,15 @@ def objstr(self, obj, indent_level=0, include=[], exclude=[]):
         return self._get_pack_str(elems, type(obj), indent_level)
 
     def _get_custom_object_str(self, obj, indent_level, include=[], exclude=[]):
+
+        def _get_line(key):
+            if self.label and any(re.fullmatch(pattern, key) is not None for pattern in self.label):
+                return set_color(f".{key} = {self.objstr(obj.__dict__[key], indent_level + 1)}", COLOR.YELLOW)
+            elif self.color:
+                return f"{set_color('.'+key, COLOR.GREEN)} = {self.objstr(obj.__dict__[key], indent_level + 1)}"
+            else:
+                return f".{key} = {self.objstr(obj.__dict__[key], indent_level + 1)}"
+
         if hasattr(obj, "__dict__"):
             keys = []
             for key in obj.__dict__.keys():
@@ -77,7 +89,8 @@ def _get_custom_object_str(self, obj, indent_level, include=[], exclude=[]):
                     if any((re.fullmatch(pattern, key) is not None for pattern in exclude)):
                         continue
                 keys.append(key)
-            elems = (f".{key} = {self.objstr(obj.__dict__[key], indent_level + 1)}" for key in sorted(keys))
+
+            elems = (_get_line(key) for key in sorted(keys))
         else:
             return str(obj)
 
@@ -114,7 +127,10 @@ def _get_header_footer(self, obj_type):
             indicator = self.indicator_map[obj_type]
             return indicator[0], indicator[1]
         else:
-            return f"<{obj_type.__name__}", ">"
+            if self.color:
+                return set_color('<' + obj_type.__name__, COLOR.CYAN), set_color(">", COLOR.CYAN)
+            else:
+                return f"<{obj_type.__name__}", ">"
 
     def _get_ellipsis(self, obj):
         header, footer = self._get_header_footer(type(obj))

diff --git a/tests/objtest.py b/tests/objtest.py
@@ -2,7 +2,16 @@
 # For details: https://github.com/gaogaotiantian/objprint/blob/master/NOTICE.txt
 
 
+import unittest
+import objprint
+
+
 class ObjTest:
     def __init__(self, attrs):
         for key, value in attrs.items():
             setattr(self, key, value)
+
+
+class ObjprintTestCase(unittest.TestCase):
+    def setUp(self):
+        objprint.config(color=False)
diff --git a/tests/test_basic.py b/tests/test_basic.py
@@ -3,17 +3,17 @@
 
 
 import io
-import unittest
 import os
 from contextlib import redirect_stdout
 from objprint import objprint, objstr, config, install
+from .objtest import ObjprintTestCase
 
 
 class A:
     pass
 
 
-class TestBasic(unittest.TestCase):
+class TestBasic(ObjprintTestCase):
     def test_print(self):
         with io.StringIO() as buf, redirect_stdout(buf):
             objprint(A())

diff --git a/tests/test_config.py b/tests/test_config.py
@@ -2,15 +2,14 @@
 # For details: https://github.com/gaogaotiantian/objprint/blob/master/NOTICE.txt
 
 
-import unittest
 import io
 from contextlib import redirect_stdout
 
 from objprint import config, objprint
-from .objtest import ObjTest
+from .objtest import ObjTest, ObjprintTestCase
 
 
-class TestConfig(unittest.TestCase):
+class TestConfig(ObjprintTestCase):
     def test_config_none_exist(self):
         self.assertRaises(TypeError, lambda: config(height=50))
 

diff --git a/tests/test_decorator.py b/tests/test_decorator.py
@@ -2,12 +2,12 @@
 # For details: https://github.com/gaogaotiantian/objprint/blob/master/NOTICE.txt
 
 
-import unittest
 import io
 from contextlib import redirect_stdout
 
 from objprint import add_objprint
 from objprint import objprint
+from .objtest import ObjprintTestCase
 
 
 @add_objprint
@@ -31,7 +31,7 @@ def __init__(self):
         self.color3 = "yellow"
 
 
-class TestDecorator(unittest.TestCase):
+class TestDecorator(ObjprintTestCase):
     def test_two_class(self):
         actual = []
         with io.StringIO() as buf, redirect_stdout(buf):

diff --git a/tests/test_objprint.py b/tests/test_objprint.py
@@ -2,15 +2,15 @@
 # For details: https://github.com/gaogaotiantian/objprint/blob/master/NOTICE.txt
 
 
-import unittest
 import io
 from contextlib import redirect_stdout
 
 from objprint import objprint
-from .objtest import ObjTest
+from objprint.color_util import COLOR
+from .objtest import ObjTest, ObjprintTestCase
 
 
-class TestObjprint(unittest.TestCase):
+class TestObjprint(ObjprintTestCase):
     def test_objprint(self):
         with io.StringIO() as buf, redirect_stdout(buf):
             b = ObjTest({})
@@ -81,3 +81,36 @@ def test_config_exclude(self):
             output = buf.getvalue()
         expected = "<ObjTest\n .elem1 = 1,\n .elem2 = 2\n>\n"
         self.assertEqual(expected, output)
+
+    def test_color_without_label(self):
+        with io.StringIO() as buf, redirect_stdout(buf):
+            obj = ObjTest({"name": "Lisa"})
+            objprint(obj, color=True)
+            output = buf.getvalue()
+        self.assertIn(COLOR.CYAN, output)
+        self.assertIn(COLOR.GREEN, output)
+        self.assertIn(COLOR.DEFAULT, output)
+
+    def test_color_with_label(self):
+        with io.StringIO() as buf, redirect_stdout(buf):
+            obj = ObjTest({"name": "Lisa", "age": 19})
+            objprint(obj, color=True, label=["age"])
+            output = buf.getvalue()
+        self.assertIn(COLOR.YELLOW, output)
+        self.assertIn(COLOR.DEFAULT, output)
+
+    def test_no_color(self):
+        with io.StringIO() as buf, redirect_stdout(buf):
+            obj = ObjTest({"Age": 10, "grade": 5})
+            objprint(obj, color=False)
+            output = buf.getvalue()
+        self.assertNotIn(COLOR.CYAN, output)
+        self.assertNotIn(COLOR.DEFAULT, output)
+
+    def test_label_only(self):
+        with io.StringIO() as buf, redirect_stdout(buf):
+            obj = ObjTest({"Age": 10, "grade": 5})
+            objprint(obj, color=False, label=['grade'])
+            output = buf.getvalue()
+        self.assertNotIn(COLOR.CYAN, output)
+        self.assertIn(COLOR.YELLOW, output)
diff --git a/tests/test_objstr.py b/tests/test_objstr.py
@@ -3,13 +3,12 @@
 
 
 import io
-import unittest
 import random
 from objprint import objstr
-from .objtest import ObjTest
+from .objtest import ObjTest, ObjprintTestCase
 
 
-class TestObjStr(unittest.TestCase):
+class TestObjStr(ObjprintTestCase):
     def test_list(self):
         lsts = (
             ([], "[]"),