<commit_before>"""
Kafka reader
"""
import six
from exporters.readers.base_reader import BaseReader
from exporters.records.base_record import BaseRecord
from exporters.default_retries import retry_short
from exporters.utils import str_list


class KafkaScannerReader(BaseReader):
    """
    This reader retrieves items from kafka brokers.

        - batch_size (int)
            Number of items to be returned in each batch

        - brokers (list)
            List of brokers uris.

        - topic (str)
            Topic to read from.

        - partitions (list)
            Partitions to read from.

        - group (str)
            Reading group for kafka client.
    """

    # List of options to set up the reader
    supported_options = {
        'batch_size': {'type': six.integer_types, 'default': 10000},
        'brokers': {'type': str_list},
        'topic': {'type': six.string_types},
        'group': {'type': six.string_types},
        'partitions': {'type': str_list, 'default': None},
        'ssl_configs': {'type': dict, 'default': None},
    }

    def __init__(self, *args, **kwargs):
        from kafka_scanner import KafkaScanner
        super(KafkaScannerReader, self).__init__(*args, **kwargs)
        brokers = self.read_option('brokers')
        group = self.read_option('group')
        topic = self.read_option('topic')
        partitions = self.read_option('partitions')

        scanner = KafkaScanner(brokers, topic, group, partitions=partitions,
                               batchsize=self.read_option('batch_size'),
                               ssl_configs=self.read_option('ssl_configs'))

        self.batches = scanner.scan_topic_batches()

        if partitions:
            topic_str = '{} (partitions: {})'.format(topic, partitions)
        else:
            topic_str = topic
        self.logger.info('KafkaScannerReader has been initiated.'
                         'Topic: {}. Group: {}'.format(topic_str, group))

    @retry_short
    def get_from_kafka(self):
        return self.batches.next()<commit_msg>partition list is a list of integers, not strings<commit_after>"""
Kafka reader
"""
import six
from exporters.readers.base_reader import BaseReader
from exporters.records.base_record import BaseRecord
from exporters.default_retries import retry_short
from exporters.utils import str_list, int_list


class KafkaScannerReader(BaseReader):
    """
    This reader retrieves items from kafka brokers.

        - batch_size (int)
            Number of items to be returned in each batch

        - brokers (list)
            List of brokers uris.

        - topic (str)
            Topic to read from.

        - partitions (list)
            Partitions to read from.

        - group (str)
            Reading group for kafka client.
    """

    # List of options to set up the reader
    supported_options = {
        'batch_size': {'type': six.integer_types, 'default': 10000},
        'brokers': {'type': str_list},
        'topic': {'type': six.string_types},
        'group': {'type': six.string_types},
        'partitions': {'type': int_list, 'default': None},
        'ssl_configs': {'type': dict, 'default': None},
    }

    def __init__(self, *args, **kwargs):
        from kafka_scanner import KafkaScanner
        super(KafkaScannerReader, self).__init__(*args, **kwargs)
        brokers = self.read_option('brokers')
        group = self.read_option('group')
        topic = self.read_option('topic')
        partitions = self.read_option('partitions')

        scanner = KafkaScanner(brokers, topic, group, partitions=partitions,
                               batchsize=self.read_option('batch_size'),
                               ssl_configs=self.read_option('ssl_configs'))

        self.batches = scanner.scan_topic_batches()

        if partitions:
            topic_str = '{} (partitions: {})'.format(topic, partitions)
        else:
            topic_str = topic
        self.logger.info('KafkaScannerReader has been initiated.'
                         'Topic: {}. Group: {}'.format(topic_str, group))

    @retry_short
    def get_from_kafka(self):
        return self.batches.next()
{'fraction_non_alphanumeric': 0.11084452975047984, 'fraction_numerical': 0.0023992322456813818, 'mean_word_length': 2.3894308943089433, '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': '3004574', 'n_tokens_mistral': 1219, 'n_tokens_neox': 1157, 'n_words': 306}