<commit_before># -*- coding: utf-8 -*-

from bisect import bisect, bisect_left, insort
from collections import Counter
from hashlib import md5
from sys import version_info


class HashRing(object):
    """Implement a ketama compatible consistent hashing ring."""

    def __init__(self, nodes=[], replicas=4, vnodes=40, compat=True):
        """Create a new HashRing.

        :param nodes: nodes used to create the continuum (see doc for format).
        :param replicas: number of replicas per node (=4 if compat).
        :param vnodes: default number of vnodes per node.
        :param compat: use a ketama compatible hash calculation.
        """
        self._default_vnodes = vnodes
        self._distribution = Counter()
        self._keys = []
        self._nodes = {}
        self._replicas = 4 if compat else replicas
        self._ring = {}

        self._configure_hashi(compat)
        self._configure_nodes(nodes)
        self._create_ring()

    def _configure_hashi(self, compat):
        """Use a ketama compatible hash in compatibility mode (default).

        There is a big performance gap in the hash calculation between
        the ketama C binding and its pure python counterpart.

        Python 3 is doing way better than python 2 thanks to its
        native bytes/int representation.

        Quick benchmark, for 1 million generated ketama compatible keys:
            - python_ketama C binding: 0.8427069187164307 ms
            - python 2: 5.462762832641602 ms
            - python 3: 3.570068597793579 ms
            - pypy: 1.6146340370178223 ms

        When using python 2 and ketama compatibility is not important, you
        can get a better hashing speed using the other provided hashing.

        Quick benchmark, for 1 million generated hash keys:
            - python 2: 3.7595579624176025 ms
            - python 3: 3.268343687057495 ms
            - pypy: 1.9193649291992188 ms
        """
        if compat:
            if version_info >= (3, ):
                self._listbytes = lambda x: x
            self.hashi = self._hashi_ketama
        else:
            self.hashi = self._hashi_faster

    def _configure_nodes(self, nodes):
        """Parse and set up the given nodes.

        :param nodes: nodes used to create the continuum (see doc for format).
        """
        if isinstance(nodes, str):
            nodes = [nodes]
        elif not isinstance(nodes, (dict, list)):
            raise ValueError(
                'nodes configuration should be a list or a dict,'
                ' got {}'.format(type(nodes)))

        for node in nodes:
            conf = {
                'hostname': node,
                'instance': None,
                'nodename': node,
                'port': None,
                'vnodes': self._default_vnodes,
                'weight': 1
            }
            nodename = node
            if isinstance(nodes, dict):
                node_conf = nodes[node]
                if isinstance(node_conf, int):
                    conf['weight'] = node_conf
                elif isinstance(node_conf, dict):
                    for k, v in node_conf.items():
                        if k in conf:
                            conf[k] = v
                else:
                    raise ValueError(
                        'node configuration should be a dict or an int,'
                        ' got {}'.format(type(node_conf)))
            self._nodes[nodename] = conf

    def _create_ring(self):
        """Generate a ketama compatible continuum/ring.
        """
        _weight_sum = 0
        for conf in self._nodes.values():
            _weight_sum += conf['weight']
        self._weight_sum = _weight_sum

        _distribution = Counter()
        _keys = []
        _ring = {}
        for nodename in self._nodes:
            for h in self._hashi_weight_generator(nodename):
                _ring[h] = nodename
                insort(_keys, h)
                _distribution[nodename] += 1
        self._distribution = _distribution
        self._keys = _keys
        self._ring = _ring

    def _get(self, key, what):
        """Generic getter magic method.

        The node with the nearest but not less hash value is returned.

        :param key: the key to look for.
        :param what: the information to look for in, allowed values:
            - instance (default): associated node instance
            - nodename: node name
            - pos: index of the given key in the ring
            - tuple: ketama compatible (pos, name) tuple
            - weight: node weight
        """
        if not self._ring:
            return None

        pos = self._get_pos(key)
        if what == 'pos':
            return pos

        nodename = self._ring[self._keys[pos]]
        if what in ['hostname', 'instance', 'port', 'weight']:
            return self._nodes[nodename][what]
        elif what == 'dict':
            return self._nodes[nodename]
        elif what == 'nodename':
            return nodename<commit_msg>implement a weight_fn parameter allowing users to pass their own weight calculation function to the hash ring<commit_after># -*- coding: utf-8 -*-

from bisect import bisect, bisect_left, insort
from collections import Counter
from hashlib import md5
from sys import version_info


class HashRing(object):
    """Implement a ketama compatible consistent hashing ring."""

    def __init__(self,
                 nodes=[],
                 replicas=4,
                 vnodes=40,
                 compat=True,
                 weight_fn=None):
        """Create a new HashRing.

        :param nodes: nodes used to create the continuum (see doc for format).
        :param replicas: number of replicas per node (=4 if compat).
        :param vnodes: default number of vnodes per node.
        :param compat: use a ketama compatible hash calculation.
        :param weight_fn: use this function to calculate the node's weight.
        """
        self._default_vnodes = vnodes
        self._distribution = Counter()
        self._keys = []
        self._nodes = {}
        self._replicas = 4 if compat else replicas
        self._ring = {}

        if weight_fn and not hasattr(weight_fn, '__call__'):
            raise TypeError('weight_fn should be a callable function')
        self._weight_fn = weight_fn

        self._configure_hashi(compat)
        self._configure_nodes(nodes)
        self._create_ring()

    def _configure_hashi(self, compat):
        """Use a ketama compatible hash in compatibility mode (default).

        There is a big performance gap in the hash calculation between
        the ketama C binding and its pure python counterpart.

        Python 3 is doing way better than python 2 thanks to its
        native bytes/int representation.

        Quick benchmark, for 1 million generated ketama compatible keys:
            - python_ketama C binding: 0.8427069187164307 ms
            - python 2: 5.462762832641602 ms
            - python 3: 3.570068597793579 ms
            - pypy: 1.6146340370178223 ms

        When using python 2 and ketama compatibility is not important, you
        can get a better hashing speed using the other provided hashing.

        Quick benchmark, for 1 million generated hash keys:
            - python 2: 3.7595579624176025 ms
            - python 3: 3.268343687057495 ms
            - pypy: 1.9193649291992188 ms
        """
        if compat:
            if version_info >= (3, ):
                self._listbytes = lambda x: x
            self.hashi = self._hashi_ketama
        else:
            self.hashi = self._hashi_faster

    def _configure_nodes(self, nodes):
        """Parse and set up the given nodes.

        :param nodes: nodes used to create the continuum (see doc for format).
        """
        if isinstance(nodes, str):
            nodes = [nodes]
        elif not isinstance(nodes, (dict, list)):
            raise ValueError(
                'nodes configuration should be a list or a dict,'
                ' got {}'.format(type(nodes)))

        for node in nodes:
            conf = {
                'hostname': node,
                'instance': None,
                'nodename': node,
                'port': None,
                'vnodes': self._default_vnodes,
                'weight': 1
            }
            nodename = node
            if isinstance(nodes, dict):
                node_conf = nodes[node]
                if isinstance(node_conf, int):
                    conf['weight'] = node_conf
                elif isinstance(node_conf, dict):
                    for k, v in node_conf.items():
                        if k in conf:
                            conf[k] = v
                else:
                    raise ValueError(
                        'node configuration should be a dict or an int,'
                        ' got {}'.format(type(node_conf)))
            if self._weight_fn:
                conf['weight'] = self._weight_fn(**conf)
            self._nodes[nodename] = conf

    def _create_ring(self):
        """Generate a ketama compatible continuum/ring.
        """
        _weight_sum = 0
        for nodename, conf in self._nodes.items():
            if self._weight_fn:
                self._nodes[nodename]['weight'] = self._weight_fn(**conf)
            _weight_sum += self._nodes[nodename]['weight']
        self._weight_sum = _weight_sum

        _distribution = Counter()
        _keys = []
        _ring = {}
        for nodename in self._nodes:
            for h in self._hashi_weight_generator(nodename):
                _ring[h] = nodename
                insort(_keys, h)
                _distribution[nodename] += 1
        self._distribution = _distribution
        self._keys = _keys
        self._ring = _ring

    regenerate = _create_ring

    def _get(self, key, what):
        """Generic getter magic method.

        The node with the nearest but not less hash value is returned.

        :param key: the key to look for.
        :param what: the information to look for in, allowed values:
            - instance (default): associated node instance
            - nodename: node name
            - pos: index of the given key in the ring
            - tuple: ketama compatible (pos, name) tuple
            - weight: node weight
        """
        if not self._ring:
            return None

        pos = self._get_pos(key)
        if what == 'pos':
            return pos

        nodename = self._ring[self._keys[pos]]
        if what in ['hostname', 'instance', 'port', 'weight']:
            return self._nodes[nodename][what]
        elif what == 'dict':
            return self._nodes[nodename]
        elif what == 'nodename':
            return nodename
{'fraction_non_alphanumeric': 0.08139318308809834, 'fraction_numerical': 0.02533060160178804, 'mean_word_length': 1.7002765903947699, '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': '6465887', 'n_tokens_mistral': 3015, 'n_tokens_neox': 2707, 'n_words': 976}