allow to retrieve partition info in mp consumer

3 files changed, 10 insertions, 8 deletions

diff --git a/kafka/consumer/base.py b/kafka/consumer/base.py
index 0800327..25c01a1 100644
--- a/kafka/consumer/base.py
+++ b/kafka/consumer/base.py
@@ -83,6 +83,14 @@ class Consumer(object):
         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 d03eb95..bd784cf 100644
--- a/kafka/consumer/multiprocess.py
+++ b/kafka/consumer/multiprocess.py
@@ -257,7 +257,8 @@ class MultiProcessConsumer(Consumer):
             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 733baa8..9b85f8c 100644
--- a/kafka/consumer/simple.py
+++ b/kafka/consumer/simple.py
@@ -131,7 +131,6 @@ class SimpleConsumer(Consumer):
                              (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 @@ class SimpleConsumer(Consumer):
             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