Python Type Hints for Apache beam ValueProvider - apache-beam

Question

How does one use type hints for ValueProvider value types passed to PTransform and DoFn classes?
class MyPTransform(beam.PTransform):
def __init__(self, my_value_provider: ValueProvider):
# How do I enforce my_value_provider has value_type of str
self.my_value_provider = my_value_provider
I can make this a RuntimeValueProvider or StaticValueProvider and test this explicitly:
type(my_value_provider.type) == str
How do others do this? I didn't see anything here: https://beam.apache.org/documentation/sdks/python-type-safety

Answer 1

I don't think there is a way to enforce this with python's type checking, although you could always add your own runtime type check to potentially improve error messages.
Alternatively, you can avoid having to use ValueProvider at all by using Flex Templates instead.

Answer 2

I ended up creating a wrapper value provider like so, opinions/comments welcome.
from typing import Generic, TypeVar
from apache_beam.options.value_provider import ValueProvider, StaticValueProvider, RuntimeValueProvider
T = TypeVar('T')
class TypedValueProvider(ValueProvider, Generic[T]):
"""
Providers a wrapper around beam's ValueProvider to allow specifying value type hints
"""
def __init__(self, value_provider: ValueProvider, value_type: Type[T]):
self.value_provider = value_provider
self.value_type = value_type
assert value_provider.value_type == value_type
def is_accessible(self):
return self.value_provider.is_accessible()
def get(self):
return self.value_provider.get()
if __name__ == '__main__':
svp = StaticValueProvider(str, 'blah')
dvp: TypedValueProvider[str] = TypedValueProvider(svp, str)
print(isinstance(dvp, ValueProvider))
assert 'blah' == dvp.get()
rvp = RuntimeValueProvider('key', str, 'default')
RuntimeValueProvider.set_runtime_options({'key': 'value'})
dvp: TypedValueProvider[str] = TypedValueProvider(rvp, str)
print(isinstance(dvp, ValueProvider))
assert 'value' == dvp.get()