From 4b5809f0696bd69788a87d0304d8215e1bf8e241 Mon Sep 17 00:00:00 2001 From: Terry Yin Date: Tue, 14 Jan 2025 19:17:30 -0500 Subject: [PATCH] fix python callable --- dagfactory/dagbuilder.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/dagfactory/dagbuilder.py b/dagfactory/dagbuilder.py index a0f758b4..f738ef72 100644 --- a/dagfactory/dagbuilder.py +++ b/dagfactory/dagbuilder.py @@ -305,7 +305,9 @@ def make_task(operator: str, task_params: Dict[str, Any]) -> BaseOperator: raise Exception(f"Failed to import operator: {operator}") from err # pylint: disable=too-many-nested-blocks try: - if operator_obj in [PythonOperator, BranchPythonOperator, PythonSensor, ShortCircuitOperator]: + if (operator_obj in [PythonOperator, BranchPythonOperator, PythonSensor, ShortCircuitOperator] + or issubclass(operator_obj, PythonOperator) + or issubclass(operator_obj, PythonSensor)): if ( not task_params.get("python_callable") and not task_params.get("python_callable_name")