<commit_before>from __future__ import print_function
import zmq
import time
import ujson
import traceback

from oct.core.turrets_manager import TurretsManager
from oct.results.stats_handler import StatsHandler


class HightQuarter(object):
    """The main hight quarter that will receive informations from the turrets
    and send the start message

    :param int publish_port: the port for publishing information to turrets
    :param int rc_port: the result collector port for collecting results from the turrets
    :param StatsHandler stats_handler: the stats handler writer
    :param dict config: the configuration of the test
    """
    def __init__(self, publish_port, rc_port, output_dir, config):
        self.context = zmq.Context()
        self.poller = zmq.Poller()

        self.result_collector = self.context.socket(zmq.PULL)
        self.result_collector.set_hwm(0)
        self.result_collector.bind("tcp://*:{}".format(rc_port))

        self.stats_handler = StatsHandler(output_dir, config)

        self.poller.register(self.result_collector, zmq.POLLIN)

        self.turrets_manager = TurretsManager(publish_port)
        self.config = config
        self.started = False
        self.messages = 0

        # waiting for init sockets
        print("Warmup")
        time.sleep(1)

    def _process_socks(self, socks):
        if self.result_collector in socks:
            self.messages += 1
            data = self.result_collector.recv()
            if b'status' not in data:
                self.stats_handler.write_result(ujson.loads(data))
                self.messages += 1
            else:
                self.turrets_manager.process_message(ujson.loads(data))
        pass

    def _print_status(self, elapsed):
        display = 'turrets: {}, elapsed: {}  messages received: {}\r'
        print(display.format(len(self.turrets_manager.turrets), round(elapsed), self.messages,), end='')

    def _clean_queue(self):
        try:
            data = self.result_collector.recv(zmq.NOBLOCK)
            while data:
                data = self.result_collector.recv(zmq.NOBLOCK)
                if b'status' not in data:<commit_msg>add external publisher to hq and fix bug in counting messages<commit_after>from __future__ import print_function
import zmq
import time
import ujson
import traceback
from zmq.utils.strtypes import asbytes

from oct.core.turrets_manager import TurretsManager
from oct.results.stats_handler import StatsHandler


class HightQuarter(object):
    """The main hight quarter that will receive informations from the turrets
    and send the start message

    :param int publish_port: the port for publishing information to turrets
    :param int rc_port: the result collector port for collecting results from the turrets
    :param StatsHandler stats_handler: the stats handler writer
    :param dict config: the configuration of the test
    """
    def __init__(self, publish_port, rc_port, output_dir, config):
        self.context = zmq.Context()
        self.poller = zmq.Poller()
        self.topic = config.get('external_publisher_topic', 'oct-publish')

        self.result_collector = self.context.socket(zmq.PULL)
        self.result_collector.set_hwm(0)
        self.result_collector.bind("tcp://*:{}".format(rc_port))

        self.external_publisher = self.context.socket(zmq.PUB)
        self.external_publisher.bind("tcp://*:{}".format(config.get('external_publisher', 5002)))

        self.stats_handler = StatsHandler(output_dir, config)

        self.poller.register(self.result_collector, zmq.POLLIN)

        self.turrets_manager = TurretsManager(publish_port)
        self.config = config
        self.started = False
        self.messages = 0

        # waiting for init sockets
        print("Warmup")
        time.sleep(1)

    def _process_socks(self, socks):
        if self.result_collector in socks:
            data = self.result_collector.recv_string()
            if 'status' not in data:
                self.stats_handler.write_result(ujson.loads(data))
                self.external_publisher.send_multipart([asbytes(self.topic), asbytes(data)])
                self.messages += 1
            else:
                self.turrets_manager.process_message(ujson.loads(data))
        pass

    def _print_status(self, elapsed):
        display = 'turrets: {}, elapsed: {}  messages received: {}\r'
        print(display.format(len(self.turrets_manager.turrets), round(elapsed), self.messages,), end='')

    def _clean_queue(self):
        try:
            data = self.result_collector.recv(zmq.NOBLOCK)
            while data:
                data = self.result_collector.recv(zmq.NOBLOCK)
                if b'status' not in data:
{'fraction_non_alphanumeric': 0.09135226791013142, 'fraction_numerical': 0.002755404832556168, 'mean_word_length': 2.7964601769911503, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, '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': '4773563', 'n_tokens_mistral': 1398, 'n_tokens_neox': 1336, 'n_words': 336}