jrmccluskey commented on issue #33189:
URL: https://github.com/apache/beam/issues/33189#issuecomment-2532768049
I actually came back to this while looking at another type checking issue
and realized that there's a better way to do this!
```py
import apache_beam as beam
from typing import TypeVar, Generic
from apache_beam.testing.util import assert_that, equal_to
from apache_beam.testing.test_pipeline import TestPipeline
T = TypeVar("T")
class ResultContainer(Generic[T]):
def __init__(self, payload: T) -> None:
self.payload = payload
class SomeDoFn(beam.DoFn):
def process(self, data: ResultContainer[int]):
yield data.payload + 1
def test_pardo():
with TestPipeline() as p:
output = (
p | beam.Create([ResultContainer(1)]).with_output_types(
ResultContainer[int])
| beam.ParDo(SomeDoFn()))
assert_that(
label="check result",
actual=output,
matcher=equal_to([2]),
)
if __name__ == '__main__':
test_pardo()
```
The `with_output_types` annotation addresses the same problem but avoids
forcing an extra DoFn to wrap the output. Hopefully this is a little more
helpful!
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
