From e8472be6d5b21246303f2a78eb1cdfeb62ea202a Mon Sep 17 00:00:00 2001 From: Aanand Prasad <aanand.prasad@gmail.com> Date: Wed, 22 Jan 2014 17:44:04 +0000 Subject: [PATCH] Fig bug in split_buffer where input was being discarded Also, write some tests for it. --- fig/cli/log_printer.py | 21 ++------------------- fig/cli/utils.py | 25 +++++++++++++++++++++++++ tests/split_buffer_test.py | 37 +++++++++++++++++++++++++++++++++++++ 3 files changed, 64 insertions(+), 19 deletions(-) create mode 100644 tests/split_buffer_test.py diff --git a/fig/cli/log_printer.py b/fig/cli/log_printer.py index a3ad15eb3..0fe3215e6 100644 --- a/fig/cli/log_printer.py +++ b/fig/cli/log_printer.py @@ -6,6 +6,7 @@ from itertools import cycle from .multiplexer import Multiplexer from . import colors +from .utils import split_buffer class LogPrinter(object): @@ -33,7 +34,7 @@ class LogPrinter(object): prefix = color_fn(container.name + " | ") # Attach to container before log printer starts running line_generator = split_buffer(self._attach(container), '\n') - return (prefix + line for line in line_generator) + return (prefix + line.decode('utf-8') for line in line_generator) def _attach(self, container): params = { @@ -44,21 +45,3 @@ class LogPrinter(object): params.update(self.attach_params) params = dict((name, 1 if value else 0) for (name, value) in list(params.items())) return container.attach(**params) - -def split_buffer(reader, separator): - """ - Given a generator which yields strings and a separator string, - joins all input, splits on the separator and yields each chunk. - Requires that each input string is decodable as UTF-8. - """ - buffered = '' - - for data in reader: - lines = (buffered + data.decode('utf-8')).split(separator) - for line in lines[:-1]: - yield line + separator - if len(lines) > 1: - buffered = lines[-1] - - if len(buffered) > 0: - yield buffered diff --git a/fig/cli/utils.py b/fig/cli/utils.py index 2b0eb42d6..3116df98b 100644 --- a/fig/cli/utils.py +++ b/fig/cli/utils.py @@ -83,3 +83,28 @@ def mkdir(path, permissions=0o700): def docker_url(): return os.environ.get('DOCKER_HOST') + + +def split_buffer(reader, separator): + """ + Given a generator which yields strings and a separator string, + joins all input, splits on the separator and yields each chunk. + + Unlike string.split(), each chunk includes the trailing + separator, except for the last one if none was found on the end + of the input. + """ + buffered = str('') + separator = str(separator) + + for data in reader: + buffered += data + while True: + index = buffered.find(separator) + if index == -1: + break + yield buffered[:index+1] + buffered = buffered[index+1:] + + if len(buffered) > 0: + yield buffered diff --git a/tests/split_buffer_test.py b/tests/split_buffer_test.py new file mode 100644 index 000000000..b90463c07 --- /dev/null +++ b/tests/split_buffer_test.py @@ -0,0 +1,37 @@ +from __future__ import unicode_literals +from __future__ import absolute_import +from fig.cli.utils import split_buffer +from . import unittest + +class SplitBufferTest(unittest.TestCase): + def test_single_line_chunks(self): + def reader(): + yield "abc\n" + yield "def\n" + yield "ghi\n" + + self.assertEqual(list(split_buffer(reader(), '\n')), ["abc\n", "def\n", "ghi\n"]) + + def test_no_end_separator(self): + def reader(): + yield "abc\n" + yield "def\n" + yield "ghi" + + self.assertEqual(list(split_buffer(reader(), '\n')), ["abc\n", "def\n", "ghi"]) + + def test_multiple_line_chunk(self): + def reader(): + yield "abc\ndef\nghi" + + self.assertEqual(list(split_buffer(reader(), '\n')), ["abc\n", "def\n", "ghi"]) + + def test_chunked_line(self): + def reader(): + yield "a" + yield "b" + yield "c" + yield "\n" + yield "d" + + self.assertEqual(list(split_buffer(reader(), '\n')), ["abc\n", "d"])