No op

PiperOrigin-RevId: 574448537

tfx/dsl/component/experimental/component_utils.py

@@ -13,28 +13,36 @@
 # limitations under the License.
 """Utils for TFX component types. Intended for internal usage only."""
 
-from typing import Any, Callable, Dict, Optional, Type
+from typing import Any, Callable, Mapping, Optional, Type
 
 from tfx import types
 from tfx.dsl.components.base import base_component
 from tfx.dsl.components.base import executor_spec as base_executor_spec
 from tfx.types import component_spec
 from tfx.types.system_executions import SystemExecution
 
+from google.protobuf import message
+
 
 def create_tfx_component_class(
     name: str,
     tfx_executor_spec: base_executor_spec.ExecutorSpec,
-    input_channel_parameters: Optional[Dict[
-        str, component_spec.ChannelParameter]] = None,
-    output_channel_parameters: Optional[Dict[
-        str, component_spec.ChannelParameter]] = None,
-    execution_parameters: Optional[Dict[
-        str, component_spec.ExecutionParameter]] = None,
+    input_channel_parameters: Optional[
+        Mapping[str, component_spec.ChannelParameter]
+    ] = None,
+    output_channel_parameters: Optional[
+        Mapping[str, component_spec.ChannelParameter]
+    ] = None,
+    execution_parameters: Optional[
+        Mapping[str, component_spec.ExecutionParameter]
+    ] = None,
     type_annotation: Optional[Type[SystemExecution]] = None,
-    default_init_args: Optional[Dict[str, Any]] = None,
+    default_init_args: Optional[Mapping[str, Any]] = None,
+    pre_execution_spec: Optional[message.Message] = None,
+    post_execution_spec: Optional[message.Message] = None,
     base_class: Type[
-        base_component.BaseComponent] = base_component.BaseComponent,
+        base_component.BaseComponent
+    ] = base_component.BaseComponent,
 ) -> Callable[..., base_component.BaseComponent]:
   """Creates a TFX component class dynamically."""
   tfx_component_spec_class = type(
@@ -79,6 +87,8 @@ def tfx_component_class_init(self, **kwargs):
       dict(
           SPEC_CLASS=tfx_component_spec_class,
           EXECUTOR_SPEC=tfx_executor_spec,
+          PRE_EXECUTABLE_SPEC=pre_execution_spec,
+          POST_EXECUTABLE_SPEC=post_execution_spec,
           __init__=tfx_component_class_init,
       ),
   )

tfx/dsl/components/base/base_component.py

@@ -46,6 +46,10 @@ class BaseComponent(base_node.BaseNode, abc.ABC):
     DRIVER_CLASS: a subclass of base_driver.BaseDriver as a custom driver for
       this component (optional, defaults to base_driver.BaseDriver). This is a
       class level value.
+    PRE_EXECUTABLE_SPEC: an optional PythonClassExecutableSpec of pre-execution
+      hook.
+    POST_EXECUTABLE_SPEC an optional PythonClassExecutableSpec of post-execution
+      hook.
     spec: an instance of `SPEC_CLASS`. See types.ComponentSpec for more details.
     platform_config: a protobuf message representing platform config for a
       component instance.
@@ -67,6 +71,9 @@ class level value.
   DRIVER_CLASS = base_driver.BaseDriver
   doc_controls.do_not_doc_in_subclasses(DRIVER_CLASS)
 
+  PRE_EXECUTABLE_SPEC = None
+  POST_EXECUTABLE_SPEC = None
+
   def __init__(
       self,
       spec: types.ComponentSpec,