From 64253a8290b073f5238ca4d49517de6cd7c500e5 Mon Sep 17 00:00:00 2001
From: Aanand Prasad <aanand.prasad@gmail.com>
Date: Wed, 18 Dec 2013 14:58:58 +0000
Subject: [PATCH] Basic log output

---
 plum/cli/colors.py      | 41 +++++++++++++++++++++++++++++++
 plum/cli/log_printer.py | 53 +++++++++++++++++++++++++++++++++++++++++
 plum/cli/main.py        | 13 ++++++++++
 plum/cli/multiplexer.py | 32 +++++++++++++++++++++++++
 4 files changed, 139 insertions(+)
 create mode 100644 plum/cli/colors.py
 create mode 100644 plum/cli/log_printer.py
 create mode 100644 plum/cli/multiplexer.py

diff --git a/plum/cli/colors.py b/plum/cli/colors.py
new file mode 100644
index 0000000000..09ec84bdb7
--- /dev/null
+++ b/plum/cli/colors.py
@@ -0,0 +1,41 @@
+NAMES = [
+    'grey',
+    'red',
+    'green',
+    'yellow',
+    'blue',
+    'magenta',
+    'cyan',
+    'white'
+]
+
+
+def get_pairs():
+    for i, name in enumerate(NAMES):
+        yield(name, str(30 + i))
+        yield('intense_' + name, str(30 + i) + ';1')
+
+
+def ansi(code):
+    return '\033[{0}m'.format(code)
+
+
+def ansi_color(code, s):
+    return '{0}{1}{2}'.format(ansi(code), s, ansi(0))
+
+
+def make_color_fn(code):
+    return lambda s: ansi_color(code, s)
+
+
+for (name, code) in get_pairs():
+    globals()[name] = make_color_fn(code)
+
+
+def rainbow():
+    cs = ['cyan', 'yellow', 'green', 'magenta', 'red', 'blue',
+          'intense_cyan', 'intense_yellow', 'intense_green',
+          'intense_magenta', 'intense_red', 'intense_blue']
+
+    for c in cs:
+        yield globals()[c]
diff --git a/plum/cli/log_printer.py b/plum/cli/log_printer.py
new file mode 100644
index 0000000000..9fced4f314
--- /dev/null
+++ b/plum/cli/log_printer.py
@@ -0,0 +1,53 @@
+import sys
+
+from itertools import cycle
+
+from ..service import get_container_name
+from .multiplexer import Multiplexer
+from . import colors
+
+
+class LogPrinter(object):
+    def __init__(self, client):
+        self.client = client
+
+    def attach(self, containers):
+        generators = self._make_log_generators(containers)
+        mux = Multiplexer(generators)
+        for line in mux.loop():
+            sys.stdout.write(line)
+
+    def _make_log_generators(self, containers):
+        color_fns = cycle(colors.rainbow())
+        generators = []
+
+        for container in containers:
+            color_fn = color_fns.next()
+            generators.append(self._make_log_generator(container, color_fn))
+
+        return generators
+
+    def _make_log_generator(self, container, color_fn):
+        container_name = get_container_name(container)
+        format = lambda line: color_fn(container_name + " | ") + line
+        return (format(line) for line in self._readlines(container))
+
+    def _readlines(self, container, logs=False, stream=True):
+        socket = self.client.attach_socket(
+            container['Id'],
+            params={
+                'stdin': 0,
+                'stdout': 1,
+                'stderr': 1,
+                'logs': 1 if logs else 0,
+                'stream': 1 if stream else 0
+            },
+        )
+
+        for line in iter(socket.makefile().readline, b''):
+            if not line.endswith('\n'):
+                line += '\n'
+
+            yield line
+
+        socket.close()
diff --git a/plum/cli/main.py b/plum/cli/main.py
index 93196edf41..9c14e8f59a 100644
--- a/plum/cli/main.py
+++ b/plum/cli/main.py
@@ -11,6 +11,7 @@ from .. import __version__
 from ..service import get_container_name
 from ..service_collection import ServiceCollection
 from .command import Command
+from .log_printer import LogPrinter
 
 from .errors import UserError
 from .docopt_command import NoSuchCommand
@@ -113,3 +114,15 @@ class TopLevelCommand(Command):
         """
         self.service_collection.stop()
 
+    def logs(self, options):
+        """
+        View containers' output
+
+        Usage: logs
+        """
+        containers = self._get_containers(all=False)
+        print "Attaching to", ", ".join(get_container_name(c) for c in containers)
+        LogPrinter(client=self.client).attach(containers)
+
+    def _get_containers(self, all):
+        return [c for s in self.service_collection for c in s.get_containers(all=all)]
diff --git a/plum/cli/multiplexer.py b/plum/cli/multiplexer.py
new file mode 100644
index 0000000000..cfcc3b6ce3
--- /dev/null
+++ b/plum/cli/multiplexer.py
@@ -0,0 +1,32 @@
+from threading import Thread
+
+try:
+    from Queue import Queue, Empty
+except ImportError:
+    from queue import Queue, Empty  # Python 3.x
+
+
+class Multiplexer(object):
+    def __init__(self, generators):
+        self.generators = generators
+        self.queue = Queue()
+
+    def loop(self):
+        self._init_readers()
+
+        while True:
+            try:
+                yield self.queue.get(timeout=0.1)
+            except Empty:
+                pass
+
+    def _init_readers(self):
+        for generator in self.generators:
+            t = Thread(target=_enqueue_output, args=(generator, self.queue))
+            t.daemon = True
+            t.start()
+
+
+def _enqueue_output(generator, queue):
+    for item in generator:
+        queue.put(item)