From 83abbc0c796b201e2225be4f4b78edc02fbaee69 Mon Sep 17 00:00:00 2001 From: janasangeetha Date: Mon, 4 Nov 2024 14:08:40 +0530 Subject: [PATCH] Fix Ruff rule B027 --- tfx/tools/cli/handler/dag_runner_patcher.py | 22 ++++++++++++++++++++- 1 file changed, 21 insertions(+), 1 deletion(-) diff --git a/tfx/tools/cli/handler/dag_runner_patcher.py b/tfx/tools/cli/handler/dag_runner_patcher.py index 924c0799bf..a57ca3beed 100644 --- a/tfx/tools/cli/handler/dag_runner_patcher.py +++ b/tfx/tools/cli/handler/dag_runner_patcher.py @@ -24,7 +24,7 @@ from tfx.proto.orchestration import pipeline_pb2 -class DagRunnerPatcher(abc.ABC): +class ParentDagRunnerPatcher(abc.ABC): """Abstract base class for Patchers for various "DagRunner"s. These patcher classes "decorate" the `run` function of the DagRunners. @@ -56,11 +56,13 @@ def __init__(self, call_real_run=True): self._run_called = False self._call_real_run = call_real_run + @abc.abstractmethod def _before_run(self, runner: tfx_runner.TfxRunner, pipeline: Union[pipeline_pb2.Pipeline, tfx_pipeline.Pipeline], context: MutableMapping[str, Any]) -> None: pass + @abc.abstractmethod def _after_run(self, runner: tfx_runner.TfxRunner, pipeline: Union[pipeline_pb2.Pipeline, tfx_pipeline.Pipeline], context: MutableMapping[str, Any]) -> None: @@ -135,3 +137,21 @@ def wrapper(*args, **kwargs): return result return wrapper + +class DagRunnerPatcher(ParentDagRunnerPatcher): + """The child class for all abstract methods.""" + + def _before_run(self, runner: tfx_runner.TfxRunner, + pipeline: Union[pipeline_pb2.Pipeline, tfx_pipeline.Pipeline], + context: MutableMapping[str, Any]) -> None: + pass + + def _after_run(self, runner: tfx_runner.TfxRunner, + pipeline: Union[pipeline_pb2.Pipeline, tfx_pipeline.Pipeline], + context: MutableMapping[str, Any]) -> None: + pass + + def get_runner_class( + self + ) -> Union[Type[tfx_runner.TfxRunner], Type[portable_tfx_runner.TfxRunner]]: + raise NotImplementedError()