"""A module for implementation of the Splitter engine."""

from splitio.engine.evaluator import CONTROL

from splitio.engine.hashfns import get_hash_fn

class Splitter(object):

"""Class responsible for choosing the right partition."""

def get_treatment(self, key, seed, partitions, algo):

"""

Return the appropriate treatment or CONTROL if no partitions are found.

:param key: The key for which to determine the treatment

:type key: str

:param seed: The feature seed

:type seed: int

:param partitions: The condition partitions

:type partitions: list

:return: The treatment

:rtype: str

"""

if not partitions:

return CONTROL

if len(partitions) == 1 and partitions[0].size == 100:

return partitions[0].treatment

return self.get_treatment_for_bucket(

self.get_bucket(key, seed, algo),

partitions

)

@staticmethod

def get_bucket(key, seed, algo):

"""

Get the bucket for a key hash.

:param key_hash: The hash for a key

:type key_hash: int

:return: The bucked for a hash

:rtype: int

"""

hashfn = get_hash_fn(algo)

key_hash = hashfn(key, seed)

return abs(key_hash) % 100 + 1

@staticmethod

def get_treatment_for_bucket(bucket, partitions):

"""

Get the treatment for a given bucket and partitions.

:param bucket: The bucket number generated by get_bucket

:type bucket: int

:param partitions: The condition partitions

:type partitions: list

:return: The treatment

:rtype: str

"""

covered_buckets = 0

for partition in partitions:

covered_buckets += partition.size

if covered_buckets >= bucket:

return partition.treatment

return CONTROL