from __future__ import absolute_import
from __future__ import unicode_literals

import logging
from collections import OrderedDict


def assignment_to_plan(assignment):
    """Convert an assignment to the format used by Kafka to describe a reassignment plan"""
    return {
        'version': 1,
        'partitions':
        [{'topic': t_p[0],
          'partition': t_p[1],
          'replicas': replica
          } for t_p, replica in assignment.iteritems()]
    }


class KafkaInterface(object):
    """This class acts as an interface to interact with Kafka."""

    def __init__(self, script_path=None):
        self._kafka_script_path = script_path
        self.log = logging.getLogger(self.__class__.__name__)

    def get_cluster_assignment(self, zk):
        """Generate the reassignment plan for given zookeeper
        configuration.
        """
        plan = zk.get_cluster_plan()
        assignment = {}
        for elem in plan['partitions']:
            assignment[
                (elem['topic'], elem['partition'])
            ] = elem['replicas']
        # assignment map created in sorted order for deterministic solution
        return OrderedDict(sorted(assignment.items(), key=lambda t: t[0]))

    def execute_plan(self, zk, plan):
        """Execute the proposed plan.

        Execute the given proposed plan over given
        zookeeper configuration

        Arguments:
        plan:   Proposed plan in json format
        """
        return zk.execute_plan(plan)

@@ -1,137 +1,4 @@
-from __future__ import absolute_import%0Afrom __future__ import unicode_literals%0A%0Aimport logging%0Afrom collections import OrderedDict%0A%0A%0A
 def 
@@ -88,16 +88,20 @@
 Kafka to
+%0A   
  describ
@@ -121,16 +121,22 @@
 ent plan
+.%0A    
 %22%22%22%0A    
@@ -342,1028 +342,4 @@
   %7D%0A
-%0A%0Aclass KafkaInterface(object):%0A    %22%22%22This class acts as an interface to interact with Kafka.%22%22%22%0A%0A    def __init__(self, script_path=None):%0A        self._kafka_script_path = script_path%0A        self.log = logging.getLogger(self.__class__.__name__)%0A%0A    def get_cluster_assignment(self, zk):%0A        %22%22%22Generate the reassignment plan for given zookeeper%0A        configuration.%0A        %22%22%22%0A        plan = zk.get_cluster_plan()%0A        assignment = %7B%7D%0A        for elem in plan%5B'partitions'%5D:%0A            assignment%5B%0A                (elem%5B'topic'%5D, elem%5B'partition'%5D)%0A            %5D = elem%5B'replicas'%5D%0A        # assignment map created in sorted order for deterministic solution%0A        return OrderedDict(sorted(assignment.items(), key=lambda t: t%5B0%5D))%0A%0A    def execute_plan(self, zk, plan):%0A        %22%22%22Execute the proposed plan.%0A%0A        Execute the given proposed plan over given%0A        zookeeper configuration%0A%0A        Arguments:%0A        plan:   Proposed plan in json format%0A        %22%22%22%0A        return zk.execute_plan(plan)%0A

Remove KafkaInterface