admin

Author: fangnx

src/confluent_kafka/_model/__init__.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import List, Optional
+from typing import List, Optional, Any
 from enum import Enum
 from .. import cimpl
 from ..cimpl import TopicPartition
@@ -62,7 +62,7 @@ class ConsumerGroupTopicPartitions:
         List of topic partitions information.
     """
 
-    def __init__(self, group_id: str, topic_partitions: Optional[List[TopicPartition]] = None) -> None:
+    def __init__(self, group_id: str, topic_partitions: Optional[List['cimpl.TopicPartition']] = None) -> None:
         self.group_id = group_id
         self.topic_partitions = topic_partitions
 
@@ -91,8 +91,8 @@ class ConsumerGroupState(Enum):
     #: Consumer Group is empty.
     EMPTY = cimpl.CONSUMER_GROUP_STATE_EMPTY
 
-    def __lt__(self, other: object) -> bool:
-        if not isinstance(other, ConsumerGroupState):
+    def __lt__(self, other: 'ConsumerGroupState') -> Any:
+        if self.__class__ != other.__class__:
             return NotImplemented
         return self.value < other.value
 
@@ -111,8 +111,8 @@ class ConsumerGroupType(Enum):
     #: Classic Type
     CLASSIC = cimpl.CONSUMER_GROUP_TYPE_CLASSIC
 
-    def __lt__(self, other: object) -> bool:
-        if not isinstance(other, ConsumerGroupType):
+    def __lt__(self, other: 'ConsumerGroupType') -> Any:
+        if self.__class__ != other.__class__:
             return NotImplemented
         return self.value < other.value
 
@@ -167,8 +167,8 @@ class IsolationLevel(Enum):
     READ_UNCOMMITTED = cimpl.ISOLATION_LEVEL_READ_UNCOMMITTED  #: Receive all the offsets.
     READ_COMMITTED = cimpl.ISOLATION_LEVEL_READ_COMMITTED  #: Skip offsets belonging to an aborted transaction.
 
-    def __lt__(self, other: object) -> bool:
-        if not isinstance(other, IsolationLevel):
+    def __lt__(self, other: 'IsolationLevel') -> Any:
+        if self.__class__ != other.__class__:
             return NotImplemented
         return self.value < other.value
 
@@ -186,7 +186,7 @@ class ElectionType(Enum):
     #: Unclean election
     UNCLEAN = cimpl.ELECTION_TYPE_UNCLEAN
 
-    def __lt__(self, other: object) -> bool:
-        if not isinstance(other, ElectionType):
+    def __lt__(self, other: 'ElectionType') -> Any:
+        if self.__class__ != other.__class__:
             return NotImplemented
         return self.value < other.value

src/confluent_kafka/_util/conversion_util.py

@@ -12,13 +12,16 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Union, Type
+from typing import Union, Type, TypeVar
 from enum import Enum
 
+# Generic type for enum conversion
+E = TypeVar('E', bound=Enum)
+
 
 class ConversionUtil:
     @staticmethod
-    def convert_to_enum(val: Union[str, int, Enum], enum_clazz: Type[Enum]) -> Enum:
+    def convert_to_enum(val: Union[str, int, E], enum_clazz: Type[E]) -> E:
         if type(enum_clazz) is not type(Enum):
             raise TypeError("'enum_clazz' must be of type Enum")
 

src/confluent_kafka/admin/_cluster.py

@@ -41,7 +41,7 @@ def __init__(self, controller: Node, nodes: List[Node], cluster_id: Optional[str
         self.cluster_id = cluster_id
         self.controller = controller
         self.nodes = nodes
-        self.authorized_operations = None
+        self.authorized_operations: Optional[List[AclOperation]] = None
         if authorized_operations:
             self.authorized_operations = []
             for op in authorized_operations:

src/confluent_kafka/admin/_resource.py

@@ -28,8 +28,8 @@ class ResourceType(Enum):
     BROKER = _cimpl.RESOURCE_BROKER  #: Broker resource. Resource name is broker id.
     TRANSACTIONAL_ID = _cimpl.RESOURCE_TRANSACTIONAL_ID  #: Transactional ID resource.
 
-    def __lt__(self, other: object) -> Any:
-        if not isinstance(other, ResourceType):
+    def __lt__(self, other: 'ResourceType') -> Any:
+        if self.__class__ != other.__class__:
             return NotImplemented
         return self.value < other.value
 
@@ -44,7 +44,7 @@ class ResourcePatternType(Enum):
     LITERAL = _cimpl.RESOURCE_PATTERN_LITERAL  #: Literal: A literal resource name
     PREFIXED = _cimpl.RESOURCE_PATTERN_PREFIXED  #: Prefixed: A prefixed resource name
 
-    def __lt__(self, other: object) -> Any:
-        if not isinstance(other, ResourcePatternType):
+    def __lt__(self, other: 'ResourcePatternType') -> Any:
+        if self.__class__ != other.__class__:
             return NotImplemented
         return self.value < other.value

src/confluent_kafka/admin/_scram.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import List, Optional
+from typing import List, Optional, Any
 from enum import Enum
 
 from .. import cimpl
@@ -26,8 +26,8 @@ class ScramMechanism(Enum):
     SCRAM_SHA_256 = cimpl.SCRAM_MECHANISM_SHA_256  #: SCRAM-SHA-256 mechanism
     SCRAM_SHA_512 = cimpl.SCRAM_MECHANISM_SHA_512  #: SCRAM-SHA-512 mechanism
 
-    def __lt__(self, other: object) -> bool:
-        if not isinstance(other, ScramMechanism):
+    def __lt__(self, other: 'ScramMechanism') -> Any:
+        if self.__class__ != other.__class__:
             return NotImplemented
         return self.value < other.value
 

src/confluent_kafka/admin/_topic.py

@@ -46,7 +46,7 @@ def __init__(self, name: str, topic_id: Uuid, is_internal: bool,
         self.topic_id = topic_id
         self.is_internal = is_internal
         self.partitions = partitions
-        self.authorized_operations = None
+        self.authorized_operations: Optional[List[AclOperation]] = None
         if authorized_operations:
             self.authorized_operations = []
             for op in authorized_operations:

src/confluent_kafka/deserializing_consumer.py

@@ -16,7 +16,7 @@
 # limitations under the License.
 #
 
-from typing import Any, Optional, Callable, List
+from typing import Optional, List
 
 from confluent_kafka.cimpl import Consumer as _ConsumerImpl, Message
 from .error import (ConsumeError,
@@ -80,7 +80,7 @@ def __init__(self, conf: ConfigDict) -> None:
 
         super(DeserializingConsumer, self).__init__(conf_copy)
 
-    def poll(self, timeout: float = -1) -> Optional['Message']:
+    def poll(self, timeout: float = -1) -> Optional[Message]:
         """
         Consume messages and calls callbacks.
 

src/confluent_kafka/error.py

@@ -35,7 +35,8 @@ class _KafkaClientError(KafkaException):
         by the broker.
     """
 
-    def __init__(self, kafka_error: KafkaError, exception: Optional[Exception] = None, kafka_message: Optional[Message] = None) -> None:
+    def __init__(self, kafka_error: KafkaError, exception: Optional[Exception] = None,
+                 kafka_message: Optional[Message] = None) -> None:
         super(_KafkaClientError, self).__init__(kafka_error)
         self.exception = exception
         self.kafka_message = kafka_message
@@ -67,7 +68,8 @@ class ConsumeError(_KafkaClientError):
 
     """
 
-    def __init__(self, kafka_error: KafkaError, exception: Optional[Exception] = None, kafka_message: Optional[Message] = None) -> None:
+    def __init__(self, kafka_error: KafkaError, exception: Optional[Exception] = None,
+                 kafka_message: Optional[Message] = None) -> None:
         super(ConsumeError, self).__init__(kafka_error, exception, kafka_message)
 
 

src/confluent_kafka/serializing_producer.py

@@ -16,7 +16,7 @@
 # limitations under the License.
 #
 
-from typing import Any, Optional, Callable
+from typing import Any, Optional
 
 from confluent_kafka.cimpl import Producer as _ProducerImpl
 from .serialization import (MessageField,