<commit_before>import sys

from itertools import cycle

from .multiplexer import Multiplexer
from . import colors


class LogPrinter(object):
    def __init__(self, containers, attach_params=None):
        self.containers = containers
        self.attach_params = attach_params or {}
        self.generators = self._make_log_generators()

    def run(self):
        mux = Multiplexer(self.generators)
        for line in mux.loop():
            sys.stdout.write(line)

    def _make_log_generators(self):
        color_fns = cycle(colors.rainbow())
        generators = []

        for container in self.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):
        prefix = color_fn(container.name + " | ")
        return (prefix + line for line in self._readlines(self._attach(container)))

    def _attach(self, container):
        params = {
            'stdin': False,
            'stdout': True,
            'stderr': True,
            'logs': False,
            'stream': True,
        }
        params.update(self.attach_params)
        params = dict((name, 1 if value else 0) for (name, value) in params.items())
        return container.attach_socket(params=params)

    def _readlines(self, socket):
        for line in iter(socket.makefile().readline, b''):
            if not line.endswith('\n'):
                line += '\n'

            yield line

        socket.close()
<commit_msg>Attach with websocket and do manual line buffering<commit_after>import sys

from itertools import cycle

from .multiplexer import Multiplexer
from . import colors


class LogPrinter(object):
    def __init__(self, containers, attach_params=None):
        self.containers = containers
        self.attach_params = attach_params or {}
        self.generators = self._make_log_generators()

    def run(self):
        mux = Multiplexer(self.generators)
        for line in mux.loop():
            sys.stdout.write(line)

    def _make_log_generators(self):
        color_fns = cycle(colors.rainbow())
        generators = []

        for container in self.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):
        prefix = color_fn(container.name + " | ")
        websocket = self._attach(container)
        return (prefix + line for line in split_buffer(read_websocket(websocket), '\n'))

    def _attach(self, container):
        params = {
            'stdin': False,
            'stdout': True,
            'stderr': True,
            'logs': False,
            'stream': True,
        }
        params.update(self.attach_params)
        params = dict((name, 1 if value else 0) for (name, value) in params.items())
        return container.attach_socket(params=params, ws=True)

def read_websocket(websocket):
    while True:
        data = websocket.recv()
        if data:
            yield data
        else:
            break

def split_buffer(reader, separator):
    buffered = ''

    for data in reader:
        lines = (buffered + data).split(separator)
        for line in lines[:-1]:
            yield line + separator
        if len(lines) > 1:
            buffered = lines[-1]

    if len(buffered) > 0:
        yield buffered

{'fraction_non_alphanumeric': 0.10054897428488876, 'fraction_numerical': 0.0023114706732158337, 'mean_word_length': 2.2660377358490567, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2096586', 'n_tokens_mistral': 1012, 'n_tokens_neox': 970, 'n_words': 251}