Skip to content

Commit 9a45508

Browse files
committed
feat: improve documentation and typing of async transformer
1 parent fc39a9d commit 9a45508

5 files changed

Lines changed: 26 additions & 10 deletions

File tree

src/athena_client/client/transformers/async_transformer.py

Lines changed: 19 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -6,21 +6,35 @@
66

77
from abc import ABC, abstractmethod
88
from collections.abc import AsyncIterator
9-
from typing import TypeVar
9+
from typing import Generic, TypeVar
1010

11+
TInput = TypeVar("TInput")
1112
T = TypeVar("T")
1213

1314

14-
class AsyncTransformer(ABC, AsyncIterator[T]):
15+
class AsyncTransformer(ABC, AsyncIterator[T], Generic[TInput, T]):
1516
"""Base class for image processing middleware."""
1617

17-
def __init__(self, source: AsyncIterator[bytes]) -> None:
18+
def __init__(self, source: AsyncIterator[TInput]) -> None:
1819
"""Initialize with source iterator."""
1920
self.source = source
2021

2122
@abstractmethod
22-
async def transform(self, data: bytes) -> T:
23-
"""Transform the image bytes."""
23+
async def transform(self, data: TInput) -> T:
24+
"""Asynchronously transform a single chunk of input data.
25+
26+
Args:
27+
data (bytes): The input data to be transformed, typically
28+
representing image bytes.
29+
30+
Returns:
31+
T: The transformed output, as defined by the subclass
32+
implementation.
33+
34+
Raises:
35+
NotImplementedError: If the method is not implemented by a subclass.
36+
37+
"""
2438
message = "Subclasses must implement this method"
2539
raise NotImplementedError(message)
2640

src/athena_client/client/transformers/brotli_compressor.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@
77
)
88

99

10-
class BrotliCompressor(AsyncTransformer[bytes]):
10+
class BrotliCompressor(AsyncTransformer[bytes, bytes]):
1111
"""Middleware for compressing bytes."""
1212

1313
async def transform(self, data: bytes) -> bytes:

src/athena_client/client/transformers/classification_input.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,9 @@
1212
)
1313

1414

15-
class ClassificationInputTransformer(AsyncTransformer[ClassificationInput]):
15+
class ClassificationInputTransformer(
16+
AsyncTransformer[bytes, ClassificationInput]
17+
):
1618
"""Transform image bytes into ClassifyRequests."""
1719

1820
def __init__(

src/athena_client/client/transformers/image_resizer.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@
1414
)
1515

1616

17-
class ImageResizer(AsyncTransformer[bytes]):
17+
class ImageResizer(AsyncTransformer[bytes, bytes]):
1818
"""Transform image bytes to ensure expected dimensions."""
1919

2020
def __init__(self, source: AsyncIterator[bytes]) -> None:

tests/transformers/test_async_transformer.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -11,15 +11,15 @@
1111
T = TypeVar("T")
1212

1313

14-
class DummyTransformer(AsyncTransformer[bytes]):
14+
class DummyTransformer(AsyncTransformer[bytes, bytes]):
1515
"""Test implementation of AsyncTransformer."""
1616

1717
async def transform(self, data: bytes) -> bytes:
1818
"""Simply return the input data."""
1919
return data
2020

2121

22-
class MockTransformer(AsyncTransformer[T]):
22+
class MockTransformer(AsyncTransformer[bytes, T]):
2323
"""Mock implementation for testing abstract base class."""
2424

2525
async def transform(self, data: bytes) -> T:

0 commit comments

Comments
 (0)