def to_proto(self):
# type: (...) -> List[beam_runner_api_pb2.DisplayData]
"""Returns a List of Beam proto representation of Display data."""
def create_payload(dd):
display_data_dict = None
try:
display_data_dict = dd.get_dict()
except ValueError:
# Skip if the display data is invalid.
return None
# We use 'label' or 'key' properties to populate the 'label' attribute of
# 'LabelledPayload'. 'label' is a better choice since it's expected to be
# more human readable but some transforms, sources, etc. may not set a
# 'label' property when configuring DisplayData.
label = (
display_data_dict['label']
if 'label' in display_data_dict else display_data_dict['key'])
value = display_data_dict['value']
if isinstance(value, str):
return beam_runner_api_pb2.LabelledPayload(
label=label,
string_value=value,
key=display_data_dict['key'],
namespace=display_data_dict.get('namespace', ''))
elif isinstance(value, bool):
return beam_runner_api_pb2.LabelledPayload(
label=label,
bool_value=value,
key=display_data_dict['key'],
namespace=display_data_dict.get('namespace', ''))
elif isinstance(value, int):
return beam_runner_api_pb2.LabelledPayload(
label=label,
int_value=value,
key=display_data_dict['key'],
namespace=display_data_dict.get('namespace', ''))
elif isinstance(value, (float, complex)):
return beam_runner_api_pb2.LabelledPayload(
label=label,
double_value=value,
key=display_data_dict['key'],
namespace=display_data_dict.get('namespace', ''))
else:
raise ValueError(
'Unsupported type %s for value of display data %s' %
(type(value), label))
dd_protos = []
for dd in self.items:
dd_proto = create_payload(dd)
if dd_proto:
dd_protos.append(
beam_runner_api_pb2.DisplayData(
urn=common_urns.StandardDisplayData.DisplayData.LABELLED.urn,
payload=create_payload(dd).SerializeToString()))
return dd_protos
def to_proto(self):
# type: (...) -> List[beam_runner_api_pb2.DisplayData]
"""Returns a List of Beam proto representation of Display data."""
def create_payload(dd):
display_data_dict = None
try:
display_data_dict = dd.get_dict()
except ValueError:
# Skip if the display data is invalid.
return None
if 'value' not in display_data_dict or 'label' not in display_data_dict:
return None
label = display_data_dict['label']
value = display_data_dict['value']
if isinstance(value, str):
return beam_runner_api_pb2.LabelledPayload(
label=label, string_value=value)
elif isinstance(value, bool):
return beam_runner_api_pb2.LabelledPayload(
label=label, bool_value=value)
elif isinstance(value, (int, float, complex)):
return beam_runner_api_pb2.LabelledPayload(
label=label, double_value=value)
else:
raise ValueError(
'Unsupported type %s for value of display data %s' %
(type(value), label))
dd_protos = []
for dd in self.items:
dd_proto = create_payload(dd)
if dd_proto:
dd_protos.append(
beam_runner_api_pb2.DisplayData(
urn=common_urns.StandardDisplayData.DisplayData.LABELLED.urn,
payload=create_payload(dd).SerializeToString()))
return dd_protos
def test_display_data(self):
class MyParentTransform(beam.PTransform):
def expand(self, p):
self.p = p
return p | beam.Create([None])
def display_data(self): # type: () -> dict
parent_dd = super(MyParentTransform, self).display_data()
parent_dd['p_dd_string'] = DisplayDataItem(
'p_dd_string_value', label='p_dd_string_label')
parent_dd['p_dd_bool'] = DisplayDataItem(True, label='p_dd_bool_label')
parent_dd['p_dd_int'] = DisplayDataItem(1, label='p_dd_int_label')
return parent_dd
class MyPTransform(MyParentTransform):
def expand(self, p):
self.p = p
return p | beam.Create([None])
def display_data(self): # type: () -> dict
parent_dd = super(MyPTransform, self).display_data()
parent_dd['dd_string'] = DisplayDataItem(
'dd_string_value', label='dd_string_label')
parent_dd['dd_bool'] = DisplayDataItem(False, label='dd_bool_label')
parent_dd['dd_int'] = DisplayDataItem(1.1, label='dd_int_label')
return parent_dd
p = beam.Pipeline()
p | MyPTransform() # pylint: disable=expression-not-assigned
from apache_beam.portability.api import beam_runner_api_pb2
proto_pipeline = Pipeline.to_runner_api(p, use_fake_coders=True)
my_transform, = [
transform
for transform in proto_pipeline.components.transforms.values()
if transform.unique_name == 'MyPTransform'
]
self.assertIsNotNone(my_transform)
self.assertListEqual(
list(my_transform.display_data),
[
beam_runner_api_pb2.DisplayData(
urn=common_urns.StandardDisplayData.DisplayData.LABELLED.urn,
payload=beam_runner_api_pb2.LabelledPayload(
label='p_dd_string_label',
string_value='p_dd_string_value').SerializeToString()),
beam_runner_api_pb2.DisplayData(
urn=common_urns.StandardDisplayData.DisplayData.LABELLED.urn,
payload=beam_runner_api_pb2.LabelledPayload(
label='p_dd_bool_label',
bool_value=True).SerializeToString()),
beam_runner_api_pb2.DisplayData(
urn=common_urns.StandardDisplayData.DisplayData.LABELLED.urn,
payload=beam_runner_api_pb2.LabelledPayload(
label='p_dd_int_label',
double_value=1).SerializeToString()),
beam_runner_api_pb2.DisplayData(
urn=common_urns.StandardDisplayData.DisplayData.LABELLED.urn,
payload=beam_runner_api_pb2.LabelledPayload(
label='dd_string_label',
string_value='dd_string_value').SerializeToString()),
beam_runner_api_pb2.DisplayData(
urn=common_urns.StandardDisplayData.DisplayData.LABELLED.urn,
payload=beam_runner_api_pb2.LabelledPayload(
label='dd_bool_label',
bool_value=False).SerializeToString()),
beam_runner_api_pb2.DisplayData(
urn=common_urns.StandardDisplayData.DisplayData.LABELLED.urn,
payload=beam_runner_api_pb2.LabelledPayload(
label='dd_int_label',
double_value=1.1).SerializeToString()),
])