Refactor+Optimize color magic

colordiff.py

@@ -0,0 +1,47 @@
+import itertools
+import enum
+
+
+class Color(str, enum.Enum):
+    RESET = "\033[0m"
+
+    BLACK = "\033[30m"
+    RED = "\033[31m"
+    GREEN = "\033[32m"
+    YELLOW = "\033[33m"
+    BLUE = "\033[34m"
+    MAGENTA = "\033[35m"
+    CYAN = "\033[36m"
+    GRAY = "\033[37m"
+    WHITE = "\033[38m"
+
+
+class Diff:
+    def __init__(self, c1: str, c2: str):
+        self.prev = ""
+        self.c1 = c1
+        self.c2 = c2
+
+    @staticmethod
+    def get(s1: str, s2: str, color_same: str, color_diff: str, color_reset: str = Color.RESET):
+        out = ""
+        is_diff = None
+        for c1, c2 in itertools.zip_longest(s1, s2):
+            if c2 is None:
+                break
+            if c1 == c2:
+                if is_diff != False:
+                    out += color_same
+                    is_diff = False
+                out += c2
+            else:
+                if is_diff != True:
+                    out += color_diff
+                    is_diff = True
+                out += c2
+        out += color_reset
+        return out
+
+    def consume(self, s: str):
+        print(self.get(self.prev, s, self.c1, self.c2))
+        self.prev = s

handler.py

@@ -1,42 +1,19 @@
 import binascii
 import itertools
 
-RED = "\033[31m"
-GREEN = "\033[32m"
-YELLOW = "\033[33m"
-BLUE = "\033[34m"
-RESET = "\033[0m"
+from colordiff import Color, Diff
 
-
-class Diff:
-    def __init__(self, c1, c2):
-        self.state = ""
-        self.c1 = c1
-        self.c2 = c2
-
-    def consume(self, s: str):
-        for p, c in itertools.zip_longest(self.state, s):
-            if c is None:
-                break
-            if p == c:
-                print(f"{self.c1}{c}{RESET}", end="")
-            else:
-                print(f"{self.c2}{c}{RESET}", end="")
-        self.state = s
-        print()
-
-
-d_c = Diff(GREEN, RED)
-d_s = Diff(BLUE, YELLOW)
+d_c = Diff(Color.GREEN, Color.RED)
+d_s = Diff(Color.BLUE, Color.YELLOW)
 
 
 def handle(tag: str, data: bytes):
     msg = data.decode().strip()
     if msg.startswith("0c") or msg.startswith("0a"):
         if tag == "server":
-            print(f"{YELLOW}{msg}{RESET}")
+            print(f"{Color.YELLOW}{msg}{Color.RESET}")
         else:
-            print(f"{RED}{msg}{RESET}")
+            print(f"{Color.RED}{msg}{Color.RESET}")
     else:
         if tag == "server":
             d_s.consume(msg)