summaryrefslogtreecommitdiff
path: root/kazoo/hosts.py
blob: 2c21f89d64364a22c3388f4d69701f7dd5727df4 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
import random


class RandomHostIterator:
    """ An iterator that returns a randomly selected host.  A host is
        guaranteed to not be selected twice unless there is only one
        host in the collection.
    """

    def __init__(self, hosts):
        self.last = 0
        self.hosts = hosts

    def __iter__(self):
        hostslist = self.hosts[:]
        random.shuffle(hostslist)
        for host in hostslist:
            yield host

    def __len__(self):
        return len(self.hosts)


def collect_hosts(hosts):
    """ Collect a set of hosts and an optional chroot from a string.
    """
    host_ports, chroot = hosts.partition("/")[::2]
    chroot = "/" + chroot if chroot else None

    result = []
    for host_port in host_ports.split(","):
        host, port = host_port.partition(":")[::2]
        port = int(port.strip()) if port else 2181
        result.append((host.strip(), port))
    return (RandomHostIterator(result), chroot)