Improve color.py

color.py

@@ -20,26 +20,34 @@ class Color:
     BRIGHT_CYAN = 14
     BRIGHT_WHITE = 15
 
-    def __init__(self, foreground: Optional[int],
+    def __init__(
-                 background: Optional[int] = None) -> None:
+        self, foreground: Optional[int], background: Optional[int] = None
+    ) -> None:
         self.foreground = foreground
         self.background = background
 
-    def print(self, *args: Any, end: str = '\n') -> None:
+    def print(self, *args: Any, end: str = "\n") -> None:
         print(self.format(*args), end=end)
 
     def format(self, *args: Any) -> str:
-        colors = ['0']
+        # Create the color string.
-
+        colors = []
         if self.foreground is not None:
             colors.append(f"38;5;{self.foreground}")
-
         if self.background is not None:
             colors.append(f"48;5;{self.background}")
+        color_str = ";".join(colors)
 
-        color_str = ';'.join(colors)
+        text = " ".join(str(a) for a in args)
-        text = ' '.join(str(a) for a in args)
         if not text:
-            return ''
+            return ""
 
-        return f'\033[{color_str}m{text}\033[0m'
+        # Reset the color as necessary.
+        reset = []
+        if self.foreground is not None:
+            reset.append("39")
+        if self.background is not None:
+            reset.append("49")
+        reset_str = ";".join(reset)
+
+        return f"\033[{color_str}m{text}\033[{reset_str}m"