Add type checking via mypy in CI

.github/workflows/python-app.yml

@@ -45,6 +45,10 @@ jobs:
           # exit-zero treats all errors as warnings.
           poetry run flake8 . --count --exit-zero --statistics
 
+      - name: Typecheck with mypy
+        run: |
+          poetry run mypy mockafka tests || true
+
       - name: Test with pytest
         run: |
           poetry run pytest --cov=./ --cov-report=xml

mockafka/admin_client.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from confluent_kafka.cimpl import NewTopic, NewPartitions
+from confluent_kafka.cimpl import NewTopic, NewPartitions  # type: ignore[import-untyped]
 
 from mockafka.cluster_metadata import ClusterMetadata
 from mockafka.kafka_store import KafkaStore

mockafka/aiokafka/aiokafka_admin_client.py

@@ -2,7 +2,7 @@
 
 from typing import Dict
 
-from aiokafka.admin import NewTopic, NewPartitions
+from aiokafka.admin import NewTopic, NewPartitions  # type: ignore[import-untyped]
 
 from mockafka.kafka_store import KafkaStore
 

mockafka/aiokafka/aiokafka_consumer.py

@@ -2,6 +2,7 @@
 
 import random
 from copy import deepcopy
+from typing import Any
 
 from mockafka.kafka_store import KafkaStore
 
@@ -34,18 +35,18 @@ class FakeAIOKafkaConsumer:
     - getmany(): Currently just calls getone().
     """
 
-    def __init__(self, *args, **kwargs):
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
         self.kafka = KafkaStore()
-        self.consumer_store = {}
+        self.consumer_store: dict[str, int] = {}
         self.subscribed_topic: list = []
 
-    async def start(self):
+    async def start(self) -> None:
         self.consumer_store = {}
-        self.subscribed_topic: list = []
+        self.subscribed_topic = []
 
-    async def stop(self):
+    async def stop(self) -> None:
         self.consumer_store = {}
-        self.subscribed_topic: list = []
+        self.subscribed_topic = []
 
     async def commit(self):
         for item in self.consumer_store:

mockafka/cluster_metadata.py

@@ -14,7 +14,7 @@ class ClusterMetadata(object):
     This class is typically not user instantiated.
     """
 
-    def __init__(self, topic: str = None):
+    def __init__(self, topic: str | None = None):
         self.kafka = KafkaStore()
         self.cluster_id = "test"
         self.controller_id = 1

mockafka/conumser.py

@@ -2,6 +2,7 @@
 
 import random
 from copy import deepcopy
+from typing import Any
 
 from mockafka.cluster_metadata import ClusterMetadata
 from mockafka.kafka_store import KafkaStore
@@ -40,7 +41,7 @@ class FakeConsumer(object):
     partitions) -> None: Incrementally unassign partitions (unsupported in mockafka).
     """
 
-    def __init__(self, *args, **kwargs):
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
         """
         Initialize the FakeConsumer.
 
@@ -49,7 +50,7 @@ def __init__(self, *args, **kwargs):
         - kwargs: Additional keyword arguments (unused).
         """
         self.kafka = KafkaStore()
-        self.consumer_store = {}
+        self.consumer_store: dict[str, int] = {}
         self.subscribed_topic: list = []
 
     def consume(self, num_messages=1, *args, **kwargs) -> list[Message]:
@@ -83,7 +84,7 @@ def close(self, *args, **kwargs):
         self.consumer_store = {}
         self.subscribed_topic = []
 
-    def commit(self, message: Message = None, *args, **kwargs):
+    def commit(self, message: Message | None = None, *args, **kwargs):
         """
         Commit offsets for consumed messages.
 

mockafka/decorators/asetup_kafka.py

@@ -2,12 +2,12 @@
 
 from functools import wraps
 
-from aiokafka.admin import NewTopic
+from aiokafka.admin import NewTopic  # type: ignore[import-untyped]
 
 from mockafka.aiokafka import FakeAIOKafkaAdmin
+from mockafka.decorators.typing import TopicConfig
 
-
-def asetup_kafka(topics: [dict[str, str]], clean: bool = False):
+def asetup_kafka(topics: list[TopicConfig], clean: bool = False):
     """
     asetup_kafka is a decorator for setting up mock Kafka topics using a FakeAIOKafkaAdminClient.
 

mockafka/decorators/bulk_producer.py

@@ -1,15 +1,17 @@
 from __future__ import annotations
 
 from functools import wraps
+
 from mockafka import FakeProducer
+from mockafka.decorators.typing import MessageDict
 
 
-def bulk_produce(list_of_messages: list[dict[str, str]]):
+def bulk_produce(list_of_messages: list[MessageDict]):
     """
     A decorator for bulk-producing messages using a FakeProducer.
 
     Parameters:
-    - list_of_messages (list[dict[str, str]]): A list of dictionaries containing message details.
+    - list_of_messages (list[dict]): A list of dictionaries containing message details.
 
     Each dictionary should have the following optional keys:
     - 'value': The value of the message.

mockafka/decorators/setup_kafka.py

@@ -1,15 +1,17 @@
 from __future__ import annotations
 
 from functools import wraps
+
 from mockafka.admin_client import FakeAdminClientImpl, NewTopic
+from mockafka.decorators.typing import TopicConfig
 
 
-def setup_kafka(topics: [dict[str, str]], clean: bool = False):
+def setup_kafka(topics: list[TopicConfig], clean: bool = False):
     """
     A decorator for setting up Mockafka with specified topics using a FakeAdminClient.
 
     Parameters:
-    - topics (list[dict[str, str]]): A list of dictionaries containing topic details.
+    - topics (list[dict]): A list of dictionaries containing topic details.
         Each dictionary should have the keys 'topic' and 'partition'.
     - clean (bool): Option to have a clean Kafka (remove existing topics) or not.
 

mockafka/decorators/typing.py

@@ -0,0 +1,18 @@
+from __future__ import annotations
+
+from typing import TypedDict
+from typing_extensions import NotRequired
+
+
+class MessageDict(TypedDict):
+    value: NotRequired[str]
+    key: NotRequired[str]
+    topic: NotRequired[str]
+    partition: NotRequired[int]
+    timestamp: NotRequired[int]
+    headers: NotRequired[dict]
+
+
+class TopicConfig(TypedDict):
+    topic: str
+    partition: int

mockafka/kafka_store.py

@@ -18,7 +18,7 @@
 
 from __future__ import annotations
 
-from confluent_kafka import KafkaException
+from confluent_kafka import KafkaException  # type: ignore[import-untyped]
 from .message import Message
 from copy import deepcopy
 
@@ -30,7 +30,7 @@
 
 
 class SingletonMeta(type):
-    _instances = {}
+    _instances: dict[type[SingletonMeta], SingletonMeta] = {}
 
     def __call__(cls, *args, **kwargs):
         if cls not in cls._instances or "clean" in kwargs.keys():

mockafka/message.py

@@ -1,12 +1,12 @@
 from __future__ import annotations
 
-from typing import Optional
+from typing import Optional, Any
 
-from confluent_kafka import KafkaError
+from confluent_kafka import KafkaError  # type: ignore[import-untyped]
 
 
 class Message:
-    def __init__(self, *args, **kwargs):
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
         self._headers: Optional[dict] = kwargs.get("headers", None)
         self._key: Optional[str] = kwargs.get("key", None)
         self._value: Optional[str] = kwargs.get("value", None)

mockafka/producer.py

@@ -8,7 +8,7 @@
 
 
 class FakeProducer(object):
-    def __init__(self, config: dict = None):
+    def __init__(self, config: dict | None = None):
         self.kafka = KafkaStore()
 
     def produce(self, topic, value=None, *args, **kwargs):

pyproject.toml

@@ -16,6 +16,7 @@ confluent-kafka = ">= 1.9.2"
 aiokafka = ">=0.10,<0.12"
 pytest-cov = ">=4.1,<6.0"
 pytest-asyncio = "^0.23.5"
+typing-extensions = "^4.12.2"
 
 [tool.poetry.group.dev.dependencies]
 flake8 = "*"

setup.cfg

@@ -2,3 +2,7 @@
 max-complexity = 10
 # The GitHub editor is 127 chars wide
 max-line-length = 127
+
+
+[mypy]
+enable_error_code = ignore-without-code

tests/test_aiokafka/test_aiokafka_producer.py

@@ -3,7 +3,7 @@
 from unittest import IsolatedAsyncioTestCase
 
 import pytest
-from aiokafka.admin import NewTopic
+from aiokafka.admin import NewTopic  # type: ignore[import-untyped]
 
 from mockafka import Message
 from mockafka.aiokafka.aiokafka_admin_client import FakeAIOKafkaAdmin
@@ -18,6 +18,10 @@ def setUp(self) -> None:
         self.producer = FakeAIOKafkaProducer()
         self.admin_client = FakeAIOKafkaAdmin()
 
+        self.topic = "test1"
+        self.key = "test_key"
+        self.value = "test_value"
+
     async def _create_mock_topic(self):
         await self.admin_client.create_topics(
             new_topics=[
@@ -27,18 +31,6 @@ async def _create_mock_topic(self):
             ]
         )
 
-    @pytest.fixture(autouse=True)
-    def topic(self):
-        self.topic = "test1"
-
-    @pytest.fixture(autouse=True)
-    def key(self):
-        self.key = "test_key"
-
-    @pytest.fixture(autouse=True)
-    def value(self):
-        self.value = "test_value"
-
     async def test_produce_failed_topic_not_exist(self):
         with pytest.raises(KafkaException):
             await self.producer.send(
@@ -69,7 +61,7 @@ async def test_produce_fail_for_none_partition(self):
                 partition=None,
             )
 
-    async def test_produce_once(self):
+    async def test_produce_once(self) -> None:
         await self._create_mock_topic()
         await self.producer.send(
             headers={},
@@ -88,7 +80,7 @@ async def test_produce_once(self):
         self.assertEqual(message.error(), None)
         self.assertEqual(message.latency(), None)
 
-    async def test_send_and_wait(self):
+    async def test_send_and_wait(self) -> None:
         await self._create_mock_topic()
 
         await self.producer.start()

tests/test_aiokafka/test_async_decorators.py

@@ -4,7 +4,7 @@
 from unittest import IsolatedAsyncioTestCase
 
 import pytest
-from aiokafka.admin import NewTopic
+from aiokafka.admin import NewTopic  # type: ignore[import-untyped]
 
 from mockafka import Message
 from mockafka.aiokafka import (
@@ -93,7 +93,7 @@ async def test_produce_with_kafka_setup_decorator(self):
     @aproduce(topic="test_topic", partition=5, key="test_", value="test_value1")
     @aproduce(topic="test_topic", partition=5, key="test_", value="test_value1")
     @aconsume(topics=["test_topic"])
-    async def test_consumer_decorator(self, message: Message = None):
+    async def test_consumer_decorator(self, message: Message | None = None):
         if message is None:
             return
 

tests/test_aiokafka/test_fake_aiokafka_admin_client.py

@@ -3,7 +3,7 @@
 from unittest import IsolatedAsyncioTestCase
 
 import pytest
-from aiokafka.admin import NewTopic, NewPartitions
+from aiokafka.admin import NewTopic, NewPartitions  # type: ignore[import-untyped]
 
 from mockafka.aiokafka.aiokafka_admin_client import FakeAIOKafkaAdmin
 from mockafka.kafka_store import KafkaStore

tests/test_async_mockafka.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import pytest
-from aiokafka.admin import NewTopic
+from aiokafka.admin import NewTopic  # type: ignore[import-untyped]
 
 from mockafka import aproduce, asetup_kafka, aconsume
 from mockafka.aiokafka import (

tests/test_docrators.py

@@ -3,10 +3,11 @@
 from mockafka import FakeConsumer, produce, bulk_produce, setup_kafka, Message
 from mockafka.admin_client import FakeAdminClientImpl, NewTopic
 from mockafka.producer import FakeProducer
+from mockafka.decorators.typing import MessageDict
 from mockafka.decorators.consumer import consume
 from unittest import TestCase
 
-sample_for_bulk_produce = [
+sample_for_bulk_produce: list[MessageDict] = [
     {
         "key": "test_key",
         "value": "test_value",
@@ -98,7 +99,7 @@ def test_produce_with_kafka_setup_decorator(self):
     @produce(topic="test_topic", partition=5, key="test_", value="test_value1")
     @produce(topic="test_topic", partition=5, key="test_", value="test_value1")
     @consume(topics=["test_topic"])
-    def test_consumer_decorator(self, message: Message = None):
+    def test_consumer_decorator(self, message: Message | None = None):
         if message is None:
             return
 

tests/test_kafka_store.py

@@ -2,7 +2,7 @@
 
 from unittest import TestCase
 
-from parameterized import parameterized
+from parameterized import parameterized  # type: ignore[import-untyped]
 
 from mockafka.kafka_store import KafkaStore, KafkaException, Message