Correct return type hints with TaggedOutputs

user_725937 0
Viewed 185

I have a DoFn, where I return a regular and a tagged output:

class Foo(DoFn):
    UNPARSABLE_TAG = "unparseable_events"

    def process(
        self, element: Tuple[str, date]
    ) -> Iterable[Union[Event, str]]:
        ev, d = element
        if ev == "1":
            yield Event(ev)
        else:
            yield TaggedOutput(self.UNPARSABLE_TAG, ev)

Runtime type checking seems to expect a return hint that encompasses both regular and tagged output types.

While this is doable, when consuming both PCollections you then need to add a "decorator hack" to later get to correct types again:

    valid, invalid = ... | ParDo(Foo()).with_outputs(Foo.UNPARSABLE_TAG, main="main")
    valid | Map(_beamDecoratorHack)
    ...

def _beamDecoratorHack(
    ev: Union[Event, str]
) -> Event:
    assert isinstance(ev, Event)
    return ev

This seems sluggish. Is there a way of not needing the hack?

apache-beam python
2 Answers

If I understand Union correctly, it lets you explicitly select types from inside the union.

For example:

def func() -> Union[int, str]:
    return "123"


# Type safe (although clearly wrong)
a: int = func()

# Not type-safe (although passable for some downstream usages)
a: list = func()

So for your example, I think this should work:

    valid: PCollection[Event], invalid = ... | ParDo(Foo())
    valid | ...
user_51298810

The method process is annotated as returning an iterable Union[Event, str] - is this correct? From the code, I would assume that this should actually be Union[Event, TaggedOutput].

user_36341250
Related