m-morozov-kw commented on issue #33189:
URL: https://github.com/apache/beam/issues/33189#issuecomment-2535328586
> I actually came back to this while looking at another type checking issue
and realized that there's a better way to do this!
>
> ```python
> 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!
Oh, that is cool. That code looks more idiomatic in terms of beam code.
Thank you
--
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: github-unsubscr...@beam.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
