Got crazy on output coloration

podctl/proc.py

@@ -3,10 +3,87 @@ Asynchronous process execution wrapper.
 """
 
 import asyncio
+from colorama import Fore, Back, Style
+import os
 import shlex
 import sys
 
 from .exceptions import WrongResult
+import pygments
+from pygments import lexers
+from pygments.formatters import TerminalFormatter
+from pygments.formatters import Terminal256Formatter
+
+
+class Output:
+    colors = (
+        '\x1b[1;36;45m',
+        '\x1b[1;36;41m',
+        '\x1b[1;36;40m',
+        '\x1b[1;37;45m',
+        '\x1b[1;32m',
+        '\x1b[1;37;44m',
+    )
+    def __init__(self):
+        self.prefixes = dict()
+        self.prefix_length = 0
+
+    def __call__(self, line, prefix, highlight=True):
+        if prefix not in self.prefixes:
+            self.prefixes[prefix] = (
+                self.colors[len([*self.prefixes.keys()]) - 1]
+            )
+            if len(prefix) > self.prefix_length:
+                self.prefix_length = len(prefix)
+
+        prefix_color = self.prefixes[prefix]
+        prefix_padding = '.' * (self.prefix_length - len(prefix) - 2)
+        if prefix_padding:
+            prefix_padding = ' ' + prefix_padding + ' '
+
+        if not prefix:
+            breakpoint()
+        #breakpoint()
+        sys.stdout.buffer.write((
+            (
+                prefix_color
+                + prefix_padding
+                + prefix
+                + ' '
+                + Back.RESET
+                + Style.RESET_ALL
+                + Fore.LIGHTBLACK_EX
+                + '| '
+                + Style.RESET_ALL
+                if prefix
+                else ''
+            )
+            + self.highlight(line, highlight)
+        ).encode('utf8'))
+
+    def cmd(self, line, prefix):
+        self(
+            Fore.LIGHTBLACK_EX
+            + '+ '
+            + Style.RESET_ALL
+            + self.highlight(line, 'bash'),
+            prefix,
+            highlight=False
+        )
+
+    def highlight(self, line, highlight=True):
+        if not highlight:
+            return line
+        elif isinstance(highlight, str):
+            lexer = lexers.get_lexer_by_name(highlight)
+        else:
+            lexer = lexers.get_lexer_by_name('python')
+        formatter = Terminal256Formatter(
+            style=os.getenv('PODCTL_STYLE', 'fruity'))
+        return pygments.highlight(line, lexer, formatter)
+
+
+output = Output()
 
 
 class PrefixStreamProtocol(asyncio.subprocess.SubprocessStreamProtocol):
@@ -27,10 +104,7 @@ class PrefixStreamProtocol(asyncio.subprocess.SubprocessStreamProtocol):
             for line in data.split(b'\n'):
                 if not line:
                     continue
-                sys.stdout.buffer.write(
+                output(line, self.prefix)
-                    self.prefix.encode('utf8') + b' | ' + line + b'\n'
-                    if self.prefix else line + b'\n'
-                )
             sys.stdout.flush()
         super().pipe_data_received(fd, data)
 
@@ -61,13 +135,14 @@ class Proc:
     """
 
     def __init__(self, *args, prefix=None, raises=True):
+        args = [str(a) for a in args]
+        self.cmd = ' '.join(args)
         if len(args) == 1:
             if isinstance(args[0], (list, tuple)):
-                args = args[0]
+                args = self.cmd = args[0]
             else:
                 args = ['sh', '-euc', ' '.join(args)]
-        self.args = [str(a) for a in args]
+        self.args = args
-        self.cmd = shlex.join(self.args)
         self.prefix = prefix
         self.raises = raises
         self.called = False
@@ -80,10 +155,7 @@ class Proc:
         from .console_script import console_script
         debug = console_script.options.get('debug', False)
         if debug is True or 'cmd' in str(debug):
-            if self.prefix:
+            output.cmd(self.cmd, self.prefix)
-                print(f'{self.prefix} | + {self.cmd}')
-            else:
-                print(f'+ {self.cmd}')
 
         loop = asyncio.events.get_event_loop()
         transport, protocol = await loop.subprocess_exec(

podctl/script.py

@@ -3,8 +3,9 @@ import copy
 import cli2
 import os
 import textwrap
+import sys
 
-from .proc import Proc
+from .proc import output, Proc
 
 
 class Script:
@@ -52,13 +53,18 @@ class Script:
                     method = 'clean_' + self.name
                     result = getattr(visitor, method)(self)
                     if debug is True or 'visit' in str(debug):
-                        print(
+                        output(
-                            getattr(visitable, 'name', '') + ' | ',
+                            ''.join([
-                            '.'.join([type(visitor).__name__, method]),
+                                '.'.join([type(visitor).__name__, method]),
-                            ' '.join(f'{k}={v}' for k, v in visitor.__dict__.items())
+                                '(',
+                                ', '.join(f'{k}={v}' for k, v in visitor.__dict__.items()),
+                                ')'
+                            ]),
+                            getattr(visitable, 'name', None)
                         )
                     if result:
                         await result
+            sys.stdout.flush()
 
         for prefix in ('init_', 'pre_', '', 'post_', 'clean_'):
             method = prefix + self.name
@@ -67,11 +73,16 @@ class Script:
                     continue
 
                 if debug is True or 'visit' in str(debug):
-                    print(
+                    output(
-                        getattr(visitable, 'name', '') + ' | ',
+                        ''.join([
-                        '.'.join([type(visitor).__name__, method]),
+                            '.'.join([type(visitor).__name__, method]),
-                        ' '.join(f'{k}={v}' for k, v in visitor.__dict__.items())
+                            '(',
+                            ', '.join(f'{k}={v}' for k, v in visitor.__dict__.items()),
+                            ')'
+                        ]),
+                        getattr(visitable, 'name', None)
                     )
+
                 result = getattr(visitor, method)(self)
                 if result:
                     try:
@@ -79,6 +90,7 @@ class Script:
                     except Exception as e:
                         await clean()
                         raise
+        sys.stdout.flush()
 
     async def run(self, *args, **kwargs):
         if self.unshare and os.getuid() != 0: