Merge util

aiokafka/coordinator/consumer.py

@@ -5,13 +5,13 @@
 import time
 
 from kafka.future import Future
-from kafka.util import WeakMethod
 
 import aiokafka.errors as Errors
 from aiokafka.metrics import AnonMeasurable
 from aiokafka.metrics.stats import Avg, Count, Max, Rate
 from aiokafka.protocol.commit import OffsetCommitRequest, OffsetFetchRequest
 from aiokafka.structs import OffsetAndMetadata, TopicPartition
+from aiokafka.util import WeakMethod
 
 from .base import BaseCoordinator, Generation
 from .assignors.range import RangePartitionAssignor

aiokafka/protocol/message.py

@@ -1,5 +1,6 @@
 import io
 import time
+from binascii import crc32
 
 from aiokafka.codec import (
     has_gzip,
@@ -12,23 +13,24 @@
     lz4_decode,
     lz4_decode_old_kafka,
 )
+from aiokafka.util import WeakMethod
+
 from .frame import KafkaBytes
 from .struct import Struct
-from .types import Int8, Int32, Int64, Bytes, Schema, AbstractType
-from kafka.util import crc32, WeakMethod
+from .types import Int8, Int32, UInt32, Int64, Bytes, Schema, AbstractType
 
 
 class Message(Struct):
     SCHEMAS = [
         Schema(
-            ("crc", Int32),
+            ("crc", UInt32),
             ("magic", Int8),
             ("attributes", Int8),
             ("key", Bytes),
             ("value", Bytes),
         ),
         Schema(
-            ("crc", Int32),
+            ("crc", UInt32),
             ("magic", Int8),
             ("attributes", Int8),
             ("timestamp", Int64),

aiokafka/protocol/struct.py

@@ -1,10 +1,10 @@
 from io import BytesIO
 
+from aiokafka.util import WeakMethod
+
 from .abstract import AbstractType
 from .types import Schema
 
-from kafka.util import WeakMethod
-
 
 class Struct(AbstractType):
     SCHEMA = Schema()

aiokafka/protocol/types.py

@@ -64,6 +64,19 @@ def decode(cls, data):
         return _unpack(cls._unpack, data.read(4))
 
 
+class UInt32(AbstractType):
+    _pack = struct.Struct(">I").pack
+    _unpack = struct.Struct(">I").unpack
+
+    @classmethod
+    def encode(cls, value):
+        return _pack(cls._pack, value)
+
+    @classmethod
+    def decode(cls, data):
+        return _unpack(cls._unpack, data.read(4))
+
+
 class Int64(AbstractType):
     _pack = struct.Struct(">q").pack
     _unpack = struct.Struct(">q").unpack

aiokafka/util.py

@@ -1,6 +1,8 @@
 import asyncio
 import os
+import weakref
 from asyncio import AbstractEventLoop
+from types import MethodType
 from typing import Any, Awaitable, Coroutine, Dict, Tuple, TypeVar, Union, cast
 
 import async_timeout
@@ -91,3 +93,38 @@ def get_running_loop() -> asyncio.AbstractEventLoop:
 
 INTEGER_MAX_VALUE = 2**31 - 1
 INTEGER_MIN_VALUE = -(2**31)
+
+
+class WeakMethod(object):
+    """
+    Callable that weakly references a method and the object it is bound to. It
+    is based on https://stackoverflow.com/a/24287465.
+
+    Arguments:
+
+        object_dot_method: A bound instance method (i.e. 'object.method').
+    """
+
+    def __init__(self, object_dot_method: MethodType) -> None:
+        self.target = weakref.ref(object_dot_method.__self__)
+        self._target_id = id(self.target())
+        self.method = weakref.ref(object_dot_method.__func__)
+        self._method_id = id(self.method())
+
+    def __call__(self, *args: Any, **kwargs: Any) -> Any:
+        """
+        Calls the method on target with args and kwargs.
+        """
+        method = self.method()
+        assert method is not None
+        return method(self.target(), *args, **kwargs)
+
+    def __hash__(self) -> int:
+        return hash(self.target) ^ hash(self.method)
+
+    def __eq__(self, other: Any) -> bool:
+        if not isinstance(other, WeakMethod):
+            return False
+        return (
+            self._target_id == other._target_id and self._method_id == other._method_id
+        )