From 43164ad096cf3dc9f252d0c6e4377ff025bb01f0 Mon Sep 17 00:00:00 2001
From: Martin Olveyra <molveyra@gmail.com>
Date: Mon, 29 Jun 2015 13:28:12 -0200
Subject: [PATCH] allow to retrieve partition info in mp consumer

---
 kafka/consumer/base.py         | 8 ++++++++
 kafka/consumer/multiprocess.py | 3 ++-
 kafka/consumer/simple.py       | 7 -------
 3 files changed, 10 insertions(+), 8 deletions(-)

diff --git a/kafka/consumer/base.py b/kafka/consumer/base.py
index 08003270f..25c01a150 100644
--- a/kafka/consumer/base.py
+++ b/kafka/consumer/base.py
@@ -83,6 +83,14 @@ def cleanup(obj):
         self._cleanup_func = cleanup
         atexit.register(cleanup, self)
 
+        self.partition_info = False     # Do not return partition info in msgs
+
+    def provide_partition_info(self):
+        """
+        Indicates that partition info must be returned by the consumer
+        """
+        self.partition_info = True
+
     def fetch_last_known_offsets(self, partitions=None):
         if self.group is None:
             raise ValueError('KafkaClient.group must not be None')
diff --git a/kafka/consumer/multiprocess.py b/kafka/consumer/multiprocess.py
index d03eb95c5..bd784cf3b 100644
--- a/kafka/consumer/multiprocess.py
+++ b/kafka/consumer/multiprocess.py
@@ -257,7 +257,8 @@ def get_messages(self, count=1, block=True, timeout=10):
             except Empty:
                 break
 
-            messages.append(message)
+            _msg = (partition, message) if self.partition_info else message
+            messages.append(_msg)
             new_offsets[partition] = message.offset + 1
             count -= 1
             if timeout is not None:
diff --git a/kafka/consumer/simple.py b/kafka/consumer/simple.py
index 733baa8e7..9b85f8c3e 100644
--- a/kafka/consumer/simple.py
+++ b/kafka/consumer/simple.py
@@ -131,7 +131,6 @@ def __init__(self, client, group, topic, auto_commit=True, partitions=None,
                              (buffer_size, max_buffer_size))
         self.buffer_size = buffer_size
         self.max_buffer_size = max_buffer_size
-        self.partition_info = False     # Do not return partition info in msgs
         self.fetch_max_wait_time = FETCH_MAX_WAIT_TIME
         self.fetch_min_bytes = fetch_size_bytes
         self.fetch_offsets = self.offsets.copy()
@@ -182,12 +181,6 @@ def reset_partition_offset(self, partition):
             self.fetch_offsets[partition] = resp.offsets[0]
             return resp.offsets[0]
 
-    def provide_partition_info(self):
-        """
-        Indicates that partition info must be returned by the consumer
-        """
-        self.partition_info = True
-
     def seek(self, offset, whence=None, partition=None):
         """
         Alter the current offset in the consumer, similar to fseek