Merge structs

Author: ods

aiokafka/admin/client.py

@@ -16,11 +16,11 @@
     AlterConfigsRequest,
     ListGroupsRequest,
     ApiVersionRequest_v0)
-from kafka.structs import TopicPartition, OffsetAndMetadata
 
 from aiokafka import __version__
-from aiokafka.errors import IncompatibleBrokerVersion, for_code
 from aiokafka.client import AIOKafkaClient
+from aiokafka.errors import IncompatibleBrokerVersion, for_code
+from aiokafka.structs import TopicPartition, OffsetAndMetadata
 
 from .config_resource import ConfigResourceType, ConfigResource
 from .new_topic import NewTopic

aiokafka/coordinator/assignors/roundrobin.py

@@ -2,13 +2,12 @@
 import itertools
 import logging
 
-from kafka.structs import TopicPartition
-
 from aiokafka.coordinator.assignors.abstract import AbstractPartitionAssignor
 from aiokafka.coordinator.protocol import (
     ConsumerProtocolMemberMetadata,
     ConsumerProtocolMemberAssignment,
 )
+from aiokafka.structs import TopicPartition
 
 log = logging.getLogger(__name__)
 

aiokafka/coordinator/assignors/sticky/sticky_assignor.py

@@ -4,7 +4,6 @@
 
 from kafka.protocol.struct import Struct
 from kafka.protocol.types import String, Array, Int32
-from kafka.structs import TopicPartition
 
 from aiokafka.coordinator.assignors.abstract import AbstractPartitionAssignor
 from aiokafka.coordinator.assignors.sticky.partition_movements import PartitionMovements
@@ -14,6 +13,7 @@
     ConsumerProtocolMemberAssignment,
 )
 from aiokafka.coordinator.protocol import Schema
+from aiokafka.structs import TopicPartition
 
 log = logging.getLogger(__name__)
 

aiokafka/coordinator/consumer.py

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

aiokafka/coordinator/protocol.py

@@ -1,6 +1,7 @@
 from kafka.protocol.struct import Struct
 from kafka.protocol.types import Array, Bytes, Int16, Int32, Schema, String
-from kafka.structs import TopicPartition
+
+from aiokafka.structs import TopicPartition
 
 
 class ConsumerProtocolMemberMetadata(Struct):

aiokafka/structs.py

@@ -1,12 +1,7 @@
 from dataclasses import dataclass
-from typing import Generic, NamedTuple, Optional, Sequence, Tuple, TypeVar
+from typing import Generic, List, NamedTuple, Optional, Sequence, Tuple, TypeVar
 
-from kafka.structs import (
-    BrokerMetadata,
-    OffsetAndMetadata,
-    PartitionMetadata,
-    TopicPartition,
-)
+from aiokafka.errors import KafkaError
 
 
 __all__ = [
@@ -19,6 +14,74 @@
 ]
 
 
+class TopicPartition(NamedTuple):
+    """A topic and partition tuple"""
+
+    topic: str
+    "A topic name"
+
+    partition: int
+    "A partition id"
+
+
+class BrokerMetadata(NamedTuple):
+    """A Kafka broker metadata used by admin tools"""
+
+    nodeId: int
+    "The Kafka broker id"
+
+    host: str
+    "The Kafka broker hostname"
+
+    port: int
+    "The Kafka broker port"
+
+    rack: Optional[str]
+    """The rack of the broker, which is used to in rack aware partition
+    assignment for fault tolerance.
+    Examples: `RACK1`, `us-east-1d`. Default: None
+    """
+
+
+class PartitionMetadata(NamedTuple):
+    """A topic partition metadata describing the state in the MetadataResponse"""
+
+    topic: str
+    "The topic name of the partition this metadata relates to"
+
+    partition: int
+    "The id of the partition this metadata relates to"
+
+    leader: int
+    "The id of the broker that is the leader for the partition"
+
+    replicas: List[int]
+    "The ids of all brokers that contain replicas of the partition"
+    isr: List[int]
+    "The ids of all brokers that contain in-sync replicas of the partition"
+
+    error: Optional[KafkaError]
+    "A KafkaError object associated with the request for this partition metadata"
+
+
+class OffsetAndMetadata(NamedTuple):
+    """The Kafka offset commit API
+
+    The Kafka offset commit API allows users to provide additional metadata
+    (in the form of a string) when an offset is committed. This can be useful
+    (for example) to store information about which node made the commit,
+    what time the commit was made, etc.
+    """
+
+    offset: int
+    "The offset to be committed"
+
+    metadata: str
+    "Non-null metadata"
+
+    # TODO add leaderEpoch:
+
+
 class RecordMetadata(NamedTuple):
     """Returned when a :class:`~.AIOKafkaProducer` sends a message"""
 

docs/api.rst

@@ -149,7 +149,7 @@ Structs
 
 .. automodule:: aiokafka.structs
 
-.. autoclass:: kafka.structs.TopicPartition
+.. autoclass:: aiokafka.structs.TopicPartition
     :members:
 
 .. autoclass:: aiokafka.structs.RecordMetadata

examples/ssl_consume_produce.py

@@ -1,7 +1,7 @@
 import asyncio
 from aiokafka import AIOKafkaProducer, AIOKafkaConsumer
 from aiokafka.helpers import create_ssl_context
-from kafka.structs import TopicPartition
+from aiokafka.structs import TopicPartition
 
 context = create_ssl_context(
     cafile="./ca-cert",  # CA used to sign certificate.

kafka/__init__.py

@@ -19,7 +19,6 @@ def emit(self, record):
 
 
 from kafka.serializer import Serializer, Deserializer
-from kafka.structs import TopicPartition, OffsetAndMetadata
 
 
 __all__ = [

kafka/scram.py

@@ -6,13 +6,13 @@
 
 import pytest
 
-from kafka.structs import TopicPartition
 from aiokafka.coordinator.assignors.range import RangePartitionAssignor
 from aiokafka.coordinator.assignors.roundrobin import RoundRobinPartitionAssignor
 from aiokafka.coordinator.assignors.sticky.sticky_assignor import (
     StickyPartitionAssignor,
 )
 from aiokafka.coordinator.protocol import ConsumerProtocolMemberAssignment
+from aiokafka.structs import TopicPartition
 
 
 @pytest.fixture(autouse=True)

kafka/structs.py

@@ -1,6 +1,5 @@
-from kafka.structs import TopicPartition
-
 from aiokafka.coordinator.assignors.sticky.partition_movements import PartitionMovements
+from aiokafka.structs import TopicPartition
 
 
 def test_empty_movements_are_sticky():

tests/coordinator/test_assignors.py

@@ -3,8 +3,6 @@
 import unittest
 from unittest import mock
 
-from kafka.structs import TopicPartition
-
 from aiokafka.cluster import ClusterMetadata
 from aiokafka.errors import (
     KafkaTimeoutError, NotLeaderForPartitionError, LeaderNotAvailableError
@@ -13,6 +11,7 @@
 from aiokafka.producer.message_accumulator import (
     MessageAccumulator, MessageBatch, BatchBuilder
 )
+from aiokafka.structs import TopicPartition
 
 from ._testutil import run_until_complete