feat: improve documentation and typing of async transformer

Author: snus-kin

src/athena_client/client/transformers/async_transformer.py

@@ -6,25 +6,39 @@
 
 from abc import ABC, abstractmethod
 from collections.abc import AsyncIterator
-from typing import TypeVar
+from typing import Generic, TypeVar
 
-T = TypeVar("T")
+TInput = TypeVar("TInput")
+TOutput = TypeVar("TOutput")
 
 
-class AsyncTransformer(ABC, AsyncIterator[T]):
+class AsyncTransformer(ABC, AsyncIterator[TOutput], Generic[TInput, TOutput]):
     """Base class for image processing middleware."""
 
-    def __init__(self, source: AsyncIterator[bytes]) -> None:
+    def __init__(self, source: AsyncIterator[TInput]) -> None:
         """Initialize with source iterator."""
         self.source = source
 
     @abstractmethod
-    async def transform(self, data: bytes) -> T:
-        """Transform the image bytes."""
+    async def transform(self, data: TInput) -> TOutput:
+        """Asynchronously transform a single chunk of input data.
+
+        Args:
+            data (bytes): The input data to be transformed, typically
+            representing image bytes.
+
+        Returns:
+            T: The transformed output, as defined by the subclass
+                implementation.
+
+        Raises:
+            NotImplementedError: If the method is not implemented by a subclass.
+
+        """
         message = "Subclasses must implement this method"
         raise NotImplementedError(message)
 
-    async def __anext__(self) -> T:
+    async def __anext__(self) -> TOutput:
         """Get next transformed bytes."""
         data = await anext(self.source)
         return await self.transform(data)

src/athena_client/client/transformers/brotli_compressor.py

@@ -7,7 +7,7 @@
 )
 
 
-class BrotliCompressor(AsyncTransformer[bytes]):
+class BrotliCompressor(AsyncTransformer[bytes, bytes]):
     """Middleware for compressing bytes."""
 
     async def transform(self, data: bytes) -> bytes:

src/athena_client/client/transformers/classification_input.py

@@ -12,7 +12,9 @@
 )
 
 
-class ClassificationInputTransformer(AsyncTransformer[ClassificationInput]):
+class ClassificationInputTransformer(
+    AsyncTransformer[bytes, ClassificationInput]
+):
     """Transform image bytes into ClassifyRequests."""
 
     def __init__(

src/athena_client/client/transformers/image_resizer.py

@@ -14,7 +14,7 @@
 )
 
 
-class ImageResizer(AsyncTransformer[bytes]):
+class ImageResizer(AsyncTransformer[bytes, bytes]):
     """Transform image bytes to ensure expected dimensions."""
 
     def __init__(self, source: AsyncIterator[bytes]) -> None:

tests/transformers/test_async_transformer.py

@@ -11,15 +11,15 @@
 T = TypeVar("T")
 
 
-class DummyTransformer(AsyncTransformer[bytes]):
+class DummyTransformer(AsyncTransformer[bytes, bytes]):
     """Test implementation of AsyncTransformer."""
 
     async def transform(self, data: bytes) -> bytes:
         """Simply return the input data."""
         return data
 
 
-class MockTransformer(AsyncTransformer[T]):
+class MockTransformer(AsyncTransformer[bytes, T]):
     """Mock implementation for testing abstract base class."""
 
     async def transform(self, data: bytes) -> T: