From baf295f1455702f6f82e879ed35a2b2cb33a54f3 Mon Sep 17 00:00:00 2001
From: jpic <jpic@yourlabs.org>
Date: Sat, 15 Feb 2020 22:38:07 +0100
Subject: [PATCH] Fix prefix and bugfix CLI args landing in Output constructor

---
 shlax/actions/base.py |  3 ++-
 shlax/output.py       | 61 ++++++++++++++++++++++++++++---------------
 shlax/proc.py         |  6 ++---
 3 files changed, 45 insertions(+), 25 deletions(-)

diff --git a/shlax/actions/base.py b/shlax/actions/base.py
index 6c1ecde..1b337aa 100644
--- a/shlax/actions/base.py
+++ b/shlax/actions/base.py
@@ -101,7 +101,7 @@ class Action:
 
     def output_factory(self, *args, **kwargs):
         kwargs.setdefault('regexps', self.regexps)
-        return Output(*args, **kwargs)
+        return Output(**kwargs)
 
     async def __call__(self, *args, **kwargs):
         self.call_args = args
@@ -126,6 +126,7 @@ class Action:
         finally:
             clean = getattr(self, 'clean', None)
             if clean:
+                self.output.clean(self)
                 await clean(*args, **kwargs)
         return result
 
diff --git a/shlax/output.py b/shlax/output.py
index 8ae6a4b..c9109ca 100644
--- a/shlax/output.py
+++ b/shlax/output.py
@@ -7,6 +7,14 @@ from .colors import colors
 class Output:
     prefixes = dict()
     colors = colors
+    prefix_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 color(self, code=None):
         if not code:
@@ -27,9 +35,7 @@ class Output:
 
     def __call__(self, line, highlight=True, flush=True):
         if self.prefix and self.prefix not in self.prefixes:
-            self.prefixes[self.prefix] = (
-                self.colors[len([*self.prefixes.keys()]) - 1]
-            )
+            self.prefixes[self.prefix] = self.prefix_colors[len(self.prefixes)]
             if len(self.prefix) > self.prefix_length:
                 self.prefix_length = len(self.prefix)
 
@@ -44,6 +50,7 @@ class Output:
                 + prefix_padding
                 + self.prefix
                 + ' '
+                + self.colors['reset']
                 + '| '
                 if self.prefix
                 else ''
@@ -88,26 +95,38 @@ class Output:
 
         return line
 
+    def clean(self, action):
+        if self.debug is True or 'visit' in str(self.debug):
+            self(''.join([
+                self.colors['bluebold'],
+                '+ CLEAN ',
+                self.colors['reset'],
+                action.colorized(),
+            ]))
+
     def start(self, action):
-        self(''.join([
-            self.colors['orangebold'],
-            '⚠ START ',
-            self.colors['reset'],
-            action.colorized(),
-        ]))
+        if self.debug is True or 'visit' in str(self.debug):
+            self(''.join([
+                self.colors['orangebold'],
+                '⚠ START ',
+                self.colors['reset'],
+                action.colorized(),
+            ]))
 
     def success(self, action):
-        self(''.join([
-            self.colors['greenbold'],
-            '✔ SUCCESS ',
-            self.colors['reset'],
-            action.colorized(),
-        ]))
+        if self.debug is True or 'visit' in str(self.debug):
+            self(''.join([
+                self.colors['greenbold'],
+                '✔ SUCCESS ',
+                self.colors['reset'],
+                action.colorized(),
+            ]))
 
     def fail(self, action, exception=None):
-        self(''.join([
-            self.colors['redbold'],
-            '✘ FAIL ',
-            self.colors['reset'],
-            action.colorized(),
-        ]))
+        if self.debug is True or 'visit' in str(self.debug):
+            self(''.join([
+                self.colors['redbold'],
+                '✘ FAIL ',
+                self.colors['reset'],
+                action.colorized(),
+            ]))
diff --git a/shlax/proc.py b/shlax/proc.py
index cacc54a..d50c0f8 100644
--- a/shlax/proc.py
+++ b/shlax/proc.py
@@ -22,9 +22,9 @@ class PrefixStreamProtocol(asyncio.subprocess.SubprocessStreamProtocol):
         super().__init__(*args, **kwargs)
 
     def pipe_data_received(self, fd, data):
-        if (self.output.debug is True or 'out' in str(self.output.debug)) and fd in (1, 2):
-            self.output(data, flush=False)
-            sys.stdout.flush()
+        if self.output.debug is True or 'out' in str(self.output.debug):
+            if fd in (1, 2):
+                self.output(data)
         super().pipe_data_received(fd, data)