Add display_data to AppliedPTransforms (#35380)

* define display_data as a appliedptransform attribute similar to annotatoins

* Adjust for transforms that define their display data at expand time

* extend base case unit test

* remove unused diplaydatacontainer class

Author: hjtran

sdks/python/apache_beam/pipeline.py

@@ -1235,6 +1235,8 @@ def __init__(
 
     self.annotations = annotations
 
+    self.display_data = {}
+
   @property
   def inputs(self):
     return tuple(self.main_inputs.values())
@@ -1438,6 +1440,11 @@ def transform_to_runner_api(
         (transform_urn not in Pipeline.runner_implemented_transforms())):
       environment_id = context.get_environment_id_for_resource_hints(
           self.resource_hints)
+    if self.transform is not None:
+      display_data = DisplayData.create_from(
+          self.transform, extra_items=self.display_data)
+    else:
+      display_data = None
 
     return beam_runner_api_pb2.PTransform(
         unique_name=self.full_label,
@@ -1457,8 +1464,7 @@ def transform_to_runner_api(
         environment_id=environment_id,
         annotations=self.annotations,
         # TODO(https://github.com/apache/beam/issues/18012): Add display_data.
-        display_data=DisplayData.create_from(self.transform).to_proto()
-        if self.transform else None)
+        display_data=display_data.to_proto() if display_data else None)
 
   @staticmethod
   def from_runner_api(

sdks/python/apache_beam/transforms/display.py

@@ -226,7 +226,7 @@ def create_from_options(cls, pipeline_options):
     return cls(pipeline_options._get_display_data_namespace(), items)
 
   @classmethod
-  def create_from(cls, has_display_data):
+  def create_from(cls, has_display_data, extra_items=None):
     """ Creates :class:`~apache_beam.transforms.display.DisplayData` from a
     :class:`HasDisplayData` instance.
 
@@ -243,9 +243,11 @@ def create_from(cls, has_display_data):
       raise ValueError(
           'Element of class {}.{} does not subclass HasDisplayData'.format(
               has_display_data.__module__, has_display_data.__class__.__name__))
+    if extra_items is None:
+      extra_items = {}
     return cls(
         has_display_data._get_display_data_namespace(),
-        has_display_data.display_data())
+        dict(**has_display_data.display_data(), **extra_items))
 
 
 class DisplayDataItem(object):

sdks/python/apache_beam/transforms/display_test.py

@@ -217,6 +217,14 @@ def display_data(self):
     ]
 
     hc.assert_that(dd.items, hc.has_items(*expected_items))
+    expected_items.append(
+        DisplayDataItemMatcher(
+            key='extra_key', value='extra_value', namespace=nspace))
+    hc.assert_that(
+        DisplayData.create_from(fn, extra_items={
+            'extra_key': 'extra_value'
+        }).items,
+        hc.has_items(*expected_items))
 
   def test_drop_if_none(self):
     class MyDoFn(beam.DoFn):