From 4316b2fe203f142a6a89d94764a51159f9ccd408 Mon Sep 17 00:00:00 2001 From: Robert Bradshaw Date: Mon, 1 Apr 2024 16:49:58 -0700 Subject: [PATCH] Don't double-nest display data payloads. (#30812) --- sdks/python/apache_beam/transforms/display.py | 22 +++++++++---------- 1 file changed, 11 insertions(+), 11 deletions(-) diff --git a/sdks/python/apache_beam/transforms/display.py b/sdks/python/apache_beam/transforms/display.py index 0d1dd552413e..86bbf101f567 100644 --- a/sdks/python/apache_beam/transforms/display.py +++ b/sdks/python/apache_beam/transforms/display.py @@ -45,6 +45,7 @@ from datetime import timedelta from typing import TYPE_CHECKING from typing import List +from typing import Optional from typing import Union from apache_beam.portability import common_urns @@ -135,11 +136,7 @@ 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): - if isinstance(dd, beam_runner_api_pb2.DisplayData): - return dd - - display_data_dict = None + def create_payload(dd) -> Optional[beam_runner_api_pb2.LabelledPayload]: try: display_data_dict = dd.get_dict() except ValueError: @@ -186,12 +183,15 @@ def create_payload(dd): 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=dd_proto.SerializeToString())) + if isinstance(dd, beam_runner_api_pb2.DisplayData): + dd_protos.append(dd) + else: + dd_payload = create_payload(dd) + if dd_payload: + dd_protos.append( + beam_runner_api_pb2.DisplayData( + urn=common_urns.StandardDisplayData.DisplayData.LABELLED.urn, + payload=dd_payload.SerializeToString())) return dd_protos @classmethod