Avro consumer return schema

.gitignore

@@ -20,3 +20,6 @@ wheelhouse
 dl-*
 *.whl
 .pytest_cache
+.idea/
+.venv/
+

confluent_kafka/avro/serializer/message_serializer.py

@@ -59,6 +59,42 @@ def __exit__(self, *args):
         return False
 
 
+class HasSchemaMixin:
+    """
+    A mixing for decoded Avro record to make able to add schema attribute
+    """
+    def schema(self):
+        """
+        :return: Avro schema for used to decode this entity
+        :rtype: avro.schema.Schema
+        """
+        return self._schema
+
+
+def _wrap(value, schema):
+    """
+    Wraps a value into subclass with HasSchemaMixin
+    :param value: a decoded value
+    :param schema: corresponding Avro schema used to decode value
+    :return: An instance of a dynamically created class with schema fullname
+    """
+    if hasattr(schema, 'fullname'):
+        name = schema.fullname
+    elif hasattr(schema, 'namespace'):
+        name = "{namespace}.{name}".format(namespace=schema.namespace,
+                                           name=schema.name)
+    elif hasattr(schema, 'name'):
+        name = schema.name
+    else:
+        name = schema.type
+
+    new_class = type(str(name), (value.__class__, HasSchemaMixin), {})
+
+    wrapped = new_class(value)
+    wrapped._schema = schema
+    return wrapped
+
+
 class MessageSerializer(object):
     """
     A helper class that can serialize and deserialize messages
@@ -213,4 +249,7 @@ def decode_message(self, message):
             if magic != MAGIC_BYTE:
                 raise SerializerError("message does not start with magic byte")
             decoder_func = self._get_decoder_func(schema_id, payload)
-            return decoder_func(payload)
+            return _wrap(
+                decoder_func(payload),
+                self.registry_client.get_by_id(schema_id)
+            )

tests/avro/test_message_serializer.py

@@ -21,13 +21,14 @@
 #
 
 import struct
-
 import unittest
 
+from confluent_kafka import avro
+from confluent_kafka.avro.serializer.message_serializer import (
+    MessageSerializer, _wrap, HasSchemaMixin
+)
 from tests.avro import data_gen
-from confluent_kafka.avro.serializer.message_serializer import MessageSerializer
 from tests.avro.mock_schema_registry_client import MockSchemaRegistryClient
-from confluent_kafka import avro
 
 
 class TestMessageSerializer(unittest.TestCase):
@@ -82,3 +83,14 @@ def test_decode_none(self):
 
     def hash_func(self):
         return hash(str(self))
+
+    def test_schema_mixin_wrapper(self):
+        schema = avro.loads(data_gen.BASIC_SCHEMA)
+        for base_class in (int, float, dict, list):
+            val = base_class()
+            wrapped = _wrap(val, schema)
+            assert val == wrapped
+            assert isinstance(wrapped, base_class)
+            assert isinstance(wrapped, HasSchemaMixin)
+            assert wrapped.schema() is schema
+            assert wrapped.__class__.__name__ == 'python.test.basic.basic'

tox.ini

@@ -1,5 +1,5 @@
 [tox]
-envlist = flake8,py27,py34,py35,py36
+envlist = flake8,py27,py34,py35,py36,py37
 
 [testenv]
 setenv =
@@ -43,6 +43,11 @@ deps =
     {[base]deps}
     avro-python3
 
+[testenv:py37]
+deps =
+    {[base]deps}
+    avro-python3
+
 [testenv:flake8]
 deps = flake8
 commands = flake8