turbaszek commented on a change in pull request #8962:
URL: https://github.com/apache/airflow/pull/8962#discussion_r432401785
##########
File path: airflow/operators/python.py
##########
@@ -145,6 +149,131 @@ def execute_callable(self):
return self.python_callable(*self.op_args, **self.op_kwargs)
+class _PythonFunctionalOperator(BaseOperator):
+ """
+ Wraps a Python callable and captures args/kwargs when called for execution.
+
+ :param python_callable: A reference to an object that is callable
+ :type python_callable: python callable
+ :param op_kwargs: a dictionary of keyword arguments that will get unpacked
+ in your function
+ :type op_kwargs: dict (templated)
+ :param op_args: a list of positional arguments that will get unpacked when
+ calling your callable
+ :type op_args: list (templated)
+ :param multiple_outputs: if set, function return value will be
+ unrolled to multiple XCom values. Dict will unroll to xcom values with
keys as keys.
+ Defaults to False.
+ :type multiple_outputs: bool
+ """
+
+ template_fields = ('op_args', 'op_kwargs')
+ ui_color = '#ffefeb'
Review comment:
```suggestion
ui_color = PythonOperator. ui_color
```
To keep it consistent. WDYT?
##########
File path: airflow/operators/python.py
##########
@@ -145,6 +149,131 @@ def execute_callable(self):
return self.python_callable(*self.op_args, **self.op_kwargs)
+class _PythonFunctionalOperator(BaseOperator):
+ """
+ Wraps a Python callable and captures args/kwargs when called for execution.
+
+ :param python_callable: A reference to an object that is callable
+ :type python_callable: python callable
+ :param op_kwargs: a dictionary of keyword arguments that will get unpacked
+ in your function
+ :type op_kwargs: dict (templated)
+ :param op_args: a list of positional arguments that will get unpacked when
+ calling your callable
+ :type op_args: list (templated)
+ :param multiple_outputs: if set, function return value will be
+ unrolled to multiple XCom values. Dict will unroll to xcom values with
keys as keys.
+ Defaults to False.
+ :type multiple_outputs: bool
+ """
+
+ template_fields = ('op_args', 'op_kwargs')
+ ui_color = '#ffefeb'
+
+ # since we won't mutate the arguments, we should just do the shallow copy
+ # there are some cases we can't deepcopy the objects(e.g protobuf).
Review comment:
```suggestion
# there are some cases we can't deepcopy the objects (e.g protobuf).
```
##########
File path: airflow/operators/python.py
##########
@@ -145,6 +149,131 @@ def execute_callable(self):
return self.python_callable(*self.op_args, **self.op_kwargs)
+class _PythonFunctionalOperator(BaseOperator):
+ """
+ Wraps a Python callable and captures args/kwargs when called for execution.
+
+ :param python_callable: A reference to an object that is callable
+ :type python_callable: python callable
+ :param op_kwargs: a dictionary of keyword arguments that will get unpacked
+ in your function
+ :type op_kwargs: dict (templated)
+ :param op_args: a list of positional arguments that will get unpacked when
+ calling your callable
+ :type op_args: list (templated)
+ :param multiple_outputs: if set, function return value will be
+ unrolled to multiple XCom values. Dict will unroll to xcom values with
keys as keys.
+ Defaults to False.
+ :type multiple_outputs: bool
+ """
+
+ template_fields = ('op_args', 'op_kwargs')
+ ui_color = '#ffefeb'
+
+ # since we won't mutate the arguments, we should just do the shallow copy
+ # there are some cases we can't deepcopy the objects(e.g protobuf).
+ shallow_copy_attrs = ('python_callable',)
+
+ @apply_defaults
+ def __init__(
+ self,
+ python_callable: Callable,
+ op_args: Tuple[Any],
+ op_kwargs: Dict[str, Any],
+ multiple_outputs: bool = False,
+ **kwargs
+ ) -> None:
+ kwargs['task_id'] = self._get_unique_task_id(kwargs['task_id'],
kwargs.get('dag', None))
+ super().__init__(**kwargs)
+ self.python_callable = python_callable
+
+ # Check that arguments can be binded
+ signature(python_callable).bind(*op_args, **op_kwargs)
+ self.multiple_outputs = multiple_outputs
+ self.op_args = op_args
+ self.op_kwargs = op_kwargs
+
+ @staticmethod
+ def _get_unique_task_id(task_id: str, dag: Optional[DAG]) -> str:
+ dag = dag or DagContext.get_current_dag()
+ if not dag or task_id not in dag.task_ids:
+ return task_id
+ core = re.split(r'__\d+$', task_id)[0]
+ suffixes = sorted(
+ [int(re.split(r'^.+__', task_id)[1])
+ for task_id in dag.task_ids
+ if re.match(rf'^{core}__\d+$', task_id)]
+ )
+ if not suffixes:
+ return f'{core}__1'
+ return f'{core}__{suffixes[-1] + 1}'
+
+ @staticmethod
+ def validate_python_callable(python_callable):
+ """Validate that python callable can be wrapped by operator.
+ Raises exception if invalid.
+
+ :param python_callable: Python object to be validated
+ :raises: TypeError, AirflowException
+ """
+ if not callable(python_callable):
+ raise TypeError('`python_callable` param must be callable')
+ if 'self' in signature(python_callable).parameters.keys():
+ raise AirflowException('@task does not support methods')
Review comment:
This is tricky one... and I would be in favor of using
`inspect.ismethod` instead of argument checking. Otherwise this will rise error:
```python
@task
def troll_airflow(self):
return f"Haha it doesn't work"
```
##########
File path: airflow/operators/python.py
##########
@@ -145,6 +149,131 @@ def execute_callable(self):
return self.python_callable(*self.op_args, **self.op_kwargs)
+class _PythonFunctionalOperator(BaseOperator):
+ """
+ Wraps a Python callable and captures args/kwargs when called for execution.
+
+ :param python_callable: A reference to an object that is callable
+ :type python_callable: python callable
+ :param op_kwargs: a dictionary of keyword arguments that will get unpacked
+ in your function
+ :type op_kwargs: dict (templated)
+ :param op_args: a list of positional arguments that will get unpacked when
+ calling your callable
+ :type op_args: list (templated)
+ :param multiple_outputs: if set, function return value will be
+ unrolled to multiple XCom values. Dict will unroll to xcom values with
keys as keys.
+ Defaults to False.
+ :type multiple_outputs: bool
+ """
+
+ template_fields = ('op_args', 'op_kwargs')
+ ui_color = '#ffefeb'
+
+ # since we won't mutate the arguments, we should just do the shallow copy
+ # there are some cases we can't deepcopy the objects(e.g protobuf).
+ shallow_copy_attrs = ('python_callable',)
+
+ @apply_defaults
+ def __init__(
+ self,
+ python_callable: Callable,
+ op_args: Tuple[Any],
+ op_kwargs: Dict[str, Any],
+ multiple_outputs: bool = False,
+ **kwargs
+ ) -> None:
+ kwargs['task_id'] = self._get_unique_task_id(kwargs['task_id'],
kwargs.get('dag', None))
+ super().__init__(**kwargs)
+ self.python_callable = python_callable
+
+ # Check that arguments can be binded
+ signature(python_callable).bind(*op_args, **op_kwargs)
+ self.multiple_outputs = multiple_outputs
+ self.op_args = op_args
+ self.op_kwargs = op_kwargs
+
+ @staticmethod
+ def _get_unique_task_id(task_id: str, dag: Optional[DAG]) -> str:
+ dag = dag or DagContext.get_current_dag()
+ if not dag or task_id not in dag.task_ids:
+ return task_id
+ core = re.split(r'__\d+$', task_id)[0]
+ suffixes = sorted(
+ [int(re.split(r'^.+__', task_id)[1])
+ for task_id in dag.task_ids
+ if re.match(rf'^{core}__\d+$', task_id)]
+ )
+ if not suffixes:
+ return f'{core}__1'
+ return f'{core}__{suffixes[-1] + 1}'
+
+ @staticmethod
+ def validate_python_callable(python_callable):
+ """Validate that python callable can be wrapped by operator.
+ Raises exception if invalid.
+
+ :param python_callable: Python object to be validated
+ :raises: TypeError, AirflowException
+ """
+ if not callable(python_callable):
+ raise TypeError('`python_callable` param must be callable')
+ if 'self' in signature(python_callable).parameters.keys():
+ raise AirflowException('@task does not support methods')
+
+ def execute(self, context: Dict):
+ return_value = self.python_callable(*self.op_args, **self.op_kwargs)
+ self.log.info("Done. Returned value was: %s", return_value)
Review comment:
Not sure we want to log this. It can be small or big output. At least
use debug
##########
File path: airflow/operators/python.py
##########
@@ -145,6 +148,142 @@ def execute_callable(self):
return self.python_callable(*self.op_args, **self.op_kwargs)
+class _PythonFunctionalOperator(BaseOperator):
+ """
+ Wraps a Python callable and captures args/kwargs when called for execution.
+
+ :param python_callable: A reference to an object that is callable
+ :type python_callable: python callable
+ :param multiple_outputs: if set, function return value will be
+ unrolled to multiple XCom values. List/Tuples will unroll to xcom
values
+ with index as key. Dict will unroll to xcom values with keys as keys.
+ Defaults to False.
+ :type multiple_outputs: bool
+ """
+
+ template_fields = ('_op_args', '_op_kwargs')
+ ui_color = '#ffefeb'
+
+ # since we won't mutate the arguments, we should just do the shallow copy
+ # there are some cases we can't deepcopy the objects(e.g protobuf).
+ shallow_copy_attrs = ('python_callable',)
+
+ @apply_defaults
+ def __init__(
+ self,
+ python_callable: Callable,
+ multiple_outputs: bool = False,
+ *args,
+ **kwargs
+ ) -> None:
+ dag = kwargs.get('dag', None) or DagContext.get_current_dag()
+ kwargs['task_id'] = self._get_unique_task_id(kwargs['task_id'], dag)
+ self._validate_python_callable(python_callable)
+ super().__init__(*args, **kwargs)
+ self.python_callable = python_callable
+ self.multiple_outputs = multiple_outputs
+ self._kwargs = kwargs
+ self._op_args: List[Any] = []
+ self._called = False
+ self._op_kwargs: Dict[str, Any] = {}
+
+ @staticmethod
+ def _get_unique_task_id(task_id, dag):
+ if not dag or task_id not in dag.task_ids:
+ return task_id
+ core = re.split(r'__\d+$', task_id)[0]
+ suffixes = sorted(
+ [int(re.split(r'^.+__', task_id)[1])
+ for task_id in dag.task_ids
+ if re.match(rf'^{core}__\d+$', task_id)]
+ )
+ if not suffixes:
+ return f'{core}__1'
+ return f'{core}__{suffixes[-1] + 1}'
+
+ @staticmethod
+ def _validate_python_callable(python_callable):
+ if not callable(python_callable):
+ raise AirflowException('`python_callable` param must be callable')
+ if 'self' in signature(python_callable).parameters.keys():
+ raise AirflowException('@task does not support methods')
+
+ def __call__(self, *args, **kwargs):
+ # If args/kwargs are set, then operator has been called. Raise
exception
+ if self._called:
+ raise AirflowException('@task decorated functions can only be
called once. If you need to reuse '
+ 'it several times in a DAG, use the `copy`
method.')
+
+ # If we have no DAG, reinitialize class to capture DAGContext and DAG
default args.
+ if not self.has_dag():
+ self.__init__(python_callable=self.python_callable,
+ multiple_outputs=self.multiple_outputs,
+ **self._kwargs)
+
+ # Capture args/kwargs
+ self._op_args = args
+ self._op_kwargs = kwargs
+ self._called = True
+ return XComArg(self)
+
+ def copy(self, task_id: Optional[str] = None, **kwargs):
+ """
+ Create a copy of the task, allow to overwrite ctor kwargs if needed.
+
+ If alias is created a new DAGContext, apply defaults and set new DAG
as the operator DAG.
+
+ :param task_id: Task id for the new operator
+ :type task_id: Optional[str]
+ """
+ if task_id:
+ self._kwargs['task_id'] = task_id
+ return _PythonFunctionalOperator(
+ python_callable=self.python_callable,
+ multiple_outputs=self.multiple_outputs,
+ **{**kwargs, **self._kwargs}
+ )
+
+ def execute(self, context: Dict):
+ return_value = self.python_callable(*self._op_args, **self._op_kwargs)
+ self.log.info("Done. Returned value was: %s", return_value)
+ if not self.multiple_outputs:
+ return return_value
+ if isinstance(return_value, dict):
+ for key, value in return_value.items():
+ self.xcom_push(context, str(key), value)
+ elif isinstance(return_value, (list, tuple)):
+ for key, value in enumerate(return_value):
+ self.xcom_push(context, str(key), value)
+ return return_value
+
+
+def task(python_callable: Optional[Callable] = None, **kwargs):
+ """
+ Python operator decorator. Wraps a function into an Airflow operator.
+ Accepts kwargs for operator kwarg. Will try to wrap operator into DAG at
declaration or
+ on function invocation. Use alias to reuse function in the DAG.
+
+ :param python_callable: Function to decorate
+ :type python_callable: Optional[Callable]
+ :param multiple_outputs: if set, function return value will be
+ unrolled to multiple XCom values. List/Tuples will unroll to xcom
values
+ with index as key. Dict will unroll to xcom values with keys as keys.
+ Defaults to False.
+ :type multiple_outputs: bool
+
+ """
+ def wrapper(f):
+ """Python wrapper to generate PythonFunctionalOperator out of simple
python functions.
+ Used for Airflow functional interface
+ """
+ return _PythonFunctionalOperator(python_callable=f,
task_id=f.__name__, **kwargs)
+ if callable(python_callable):
+ return wrapper(python_callable)
+ elif python_callable is not None:
+ raise AirflowException('No args allowed while using @task, use kwargs
instead')
+ return wrapper
Review comment:
Is there any particular reason why we don't do simply:
```python
def task(*args, **kwargs):
"""
Python operator decorator. Wraps a function into an Airflow operator.
Accepts kwargs for operator kwarg. Will try to wrap operator into DAG at
declaration or
on function invocation. Use alias to reuse function in the DAG.
"""
if args:
raise AirflowException("No args allowed")
def wrapper(f):
"""
Python wrapper to generate PythonFunctionalOperator out of simple
python functions.
Used for Airflow functional interface
"""
_PythonFunctionalOperator.validate_python_callable(f)
@functools.wraps(f)
def factory(*args, **f_kwargs):
op = _PythonFunctionalOperator(
python_callable=f,
task_id=f.__name__,
op_args=args,
op_kwargs=f_kwargs,
**kwargs
)
return XComArg(op)
return factory
return wrapper
```
##########
File path: airflow/operators/python.py
##########
@@ -145,6 +149,131 @@ def execute_callable(self):
return self.python_callable(*self.op_args, **self.op_kwargs)
+class _PythonFunctionalOperator(BaseOperator):
+ """
+ Wraps a Python callable and captures args/kwargs when called for execution.
+
+ :param python_callable: A reference to an object that is callable
+ :type python_callable: python callable
+ :param op_kwargs: a dictionary of keyword arguments that will get unpacked
+ in your function
+ :type op_kwargs: dict (templated)
+ :param op_args: a list of positional arguments that will get unpacked when
+ calling your callable
+ :type op_args: list (templated)
+ :param multiple_outputs: if set, function return value will be
+ unrolled to multiple XCom values. Dict will unroll to xcom values with
keys as keys.
+ Defaults to False.
+ :type multiple_outputs: bool
+ """
+
+ template_fields = ('op_args', 'op_kwargs')
+ ui_color = '#ffefeb'
+
+ # since we won't mutate the arguments, we should just do the shallow copy
+ # there are some cases we can't deepcopy the objects(e.g protobuf).
+ shallow_copy_attrs = ('python_callable',)
+
+ @apply_defaults
+ def __init__(
+ self,
+ python_callable: Callable,
+ op_args: Tuple[Any],
+ op_kwargs: Dict[str, Any],
+ multiple_outputs: bool = False,
+ **kwargs
+ ) -> None:
+ kwargs['task_id'] = self._get_unique_task_id(kwargs['task_id'],
kwargs.get('dag', None))
+ super().__init__(**kwargs)
Review comment:
```suggestion
super().__init__(**kwargs)
self.task_id = self._get_unique_task_id(kwargs['task_id'],
kwargs.get('dag', None))
```
##########
File path: airflow/operators/python.py
##########
@@ -145,6 +149,131 @@ def execute_callable(self):
return self.python_callable(*self.op_args, **self.op_kwargs)
+class _PythonFunctionalOperator(BaseOperator):
+ """
+ Wraps a Python callable and captures args/kwargs when called for execution.
+
+ :param python_callable: A reference to an object that is callable
+ :type python_callable: python callable
+ :param op_kwargs: a dictionary of keyword arguments that will get unpacked
+ in your function
+ :type op_kwargs: dict (templated)
+ :param op_args: a list of positional arguments that will get unpacked when
+ calling your callable
+ :type op_args: list (templated)
+ :param multiple_outputs: if set, function return value will be
+ unrolled to multiple XCom values. Dict will unroll to xcom values with
keys as keys.
+ Defaults to False.
+ :type multiple_outputs: bool
+ """
+
+ template_fields = ('op_args', 'op_kwargs')
+ ui_color = '#ffefeb'
+
+ # since we won't mutate the arguments, we should just do the shallow copy
+ # there are some cases we can't deepcopy the objects(e.g protobuf).
+ shallow_copy_attrs = ('python_callable',)
+
+ @apply_defaults
+ def __init__(
+ self,
+ python_callable: Callable,
+ op_args: Tuple[Any],
+ op_kwargs: Dict[str, Any],
+ multiple_outputs: bool = False,
+ **kwargs
+ ) -> None:
+ kwargs['task_id'] = self._get_unique_task_id(kwargs['task_id'],
kwargs.get('dag', None))
+ super().__init__(**kwargs)
+ self.python_callable = python_callable
+
+ # Check that arguments can be binded
+ signature(python_callable).bind(*op_args, **op_kwargs)
+ self.multiple_outputs = multiple_outputs
+ self.op_args = op_args
+ self.op_kwargs = op_kwargs
+
+ @staticmethod
+ def _get_unique_task_id(task_id: str, dag: Optional[DAG]) -> str:
Review comment:
```suggestion
def _get_unique_task_id(task_id: str, dag: Optional[DAG] = None) -> str:
```
##########
File path: airflow/operators/python.py
##########
@@ -145,6 +149,131 @@ def execute_callable(self):
return self.python_callable(*self.op_args, **self.op_kwargs)
+class _PythonFunctionalOperator(BaseOperator):
+ """
+ Wraps a Python callable and captures args/kwargs when called for execution.
+
+ :param python_callable: A reference to an object that is callable
+ :type python_callable: python callable
+ :param op_kwargs: a dictionary of keyword arguments that will get unpacked
+ in your function
+ :type op_kwargs: dict (templated)
+ :param op_args: a list of positional arguments that will get unpacked when
+ calling your callable
+ :type op_args: list (templated)
+ :param multiple_outputs: if set, function return value will be
+ unrolled to multiple XCom values. Dict will unroll to xcom values with
keys as keys.
+ Defaults to False.
+ :type multiple_outputs: bool
+ """
+
+ template_fields = ('op_args', 'op_kwargs')
+ ui_color = '#ffefeb'
+
+ # since we won't mutate the arguments, we should just do the shallow copy
+ # there are some cases we can't deepcopy the objects(e.g protobuf).
+ shallow_copy_attrs = ('python_callable',)
+
+ @apply_defaults
+ def __init__(
+ self,
+ python_callable: Callable,
+ op_args: Tuple[Any],
+ op_kwargs: Dict[str, Any],
+ multiple_outputs: bool = False,
+ **kwargs
+ ) -> None:
+ kwargs['task_id'] = self._get_unique_task_id(kwargs['task_id'],
kwargs.get('dag', None))
+ super().__init__(**kwargs)
+ self.python_callable = python_callable
+
+ # Check that arguments can be binded
+ signature(python_callable).bind(*op_args, **op_kwargs)
+ self.multiple_outputs = multiple_outputs
+ self.op_args = op_args
+ self.op_kwargs = op_kwargs
+
+ @staticmethod
+ def _get_unique_task_id(task_id: str, dag: Optional[DAG]) -> str:
+ dag = dag or DagContext.get_current_dag()
+ if not dag or task_id not in dag.task_ids:
+ return task_id
+ core = re.split(r'__\d+$', task_id)[0]
+ suffixes = sorted(
+ [int(re.split(r'^.+__', task_id)[1])
+ for task_id in dag.task_ids
+ if re.match(rf'^{core}__\d+$', task_id)]
+ )
+ if not suffixes:
+ return f'{core}__1'
+ return f'{core}__{suffixes[-1] + 1}'
Review comment:
Would you mind adding a comment about how the auto generated id looks
like? It would make it easier to understand what we are doing here.
##########
File path: airflow/operators/python.py
##########
@@ -145,6 +149,131 @@ def execute_callable(self):
return self.python_callable(*self.op_args, **self.op_kwargs)
+class _PythonFunctionalOperator(BaseOperator):
+ """
+ Wraps a Python callable and captures args/kwargs when called for execution.
+
+ :param python_callable: A reference to an object that is callable
+ :type python_callable: python callable
+ :param op_kwargs: a dictionary of keyword arguments that will get unpacked
+ in your function
+ :type op_kwargs: dict (templated)
+ :param op_args: a list of positional arguments that will get unpacked when
+ calling your callable
+ :type op_args: list (templated)
+ :param multiple_outputs: if set, function return value will be
+ unrolled to multiple XCom values. Dict will unroll to xcom values with
keys as keys.
+ Defaults to False.
+ :type multiple_outputs: bool
+ """
+
+ template_fields = ('op_args', 'op_kwargs')
+ ui_color = '#ffefeb'
+
+ # since we won't mutate the arguments, we should just do the shallow copy
+ # there are some cases we can't deepcopy the objects(e.g protobuf).
+ shallow_copy_attrs = ('python_callable',)
+
+ @apply_defaults
+ def __init__(
+ self,
+ python_callable: Callable,
+ op_args: Tuple[Any],
+ op_kwargs: Dict[str, Any],
+ multiple_outputs: bool = False,
+ **kwargs
+ ) -> None:
+ kwargs['task_id'] = self._get_unique_task_id(kwargs['task_id'],
kwargs.get('dag', None))
+ super().__init__(**kwargs)
+ self.python_callable = python_callable
+
+ # Check that arguments can be binded
+ signature(python_callable).bind(*op_args, **op_kwargs)
+ self.multiple_outputs = multiple_outputs
+ self.op_args = op_args
+ self.op_kwargs = op_kwargs
+
+ @staticmethod
+ def _get_unique_task_id(task_id: str, dag: Optional[DAG]) -> str:
+ dag = dag or DagContext.get_current_dag()
+ if not dag or task_id not in dag.task_ids:
+ return task_id
+ core = re.split(r'__\d+$', task_id)[0]
+ suffixes = sorted(
+ [int(re.split(r'^.+__', task_id)[1])
+ for task_id in dag.task_ids
+ if re.match(rf'^{core}__\d+$', task_id)]
+ )
+ if not suffixes:
+ return f'{core}__1'
+ return f'{core}__{suffixes[-1] + 1}'
+
+ @staticmethod
+ def validate_python_callable(python_callable):
+ """Validate that python callable can be wrapped by operator.
+ Raises exception if invalid.
+
+ :param python_callable: Python object to be validated
+ :raises: TypeError, AirflowException
+ """
+ if not callable(python_callable):
+ raise TypeError('`python_callable` param must be callable')
+ if 'self' in signature(python_callable).parameters.keys():
+ raise AirflowException('@task does not support methods')
+
+ def execute(self, context: Dict):
+ return_value = self.python_callable(*self.op_args, **self.op_kwargs)
+ self.log.info("Done. Returned value was: %s", return_value)
+ if not self.multiple_outputs:
+ return return_value
+ if isinstance(return_value, dict):
+ for key in return_value.keys():
+ if not isinstance(key, str):
+ raise AirflowException('Returned dictionary keys must be
strings when using '
+ f'multiple_outputs, found {key}
({type(key)}) instead')
+ for key, value in return_value.items():
+ self.xcom_push(context, key, value)
+ else:
+ self.log.info(f'Returned output was type {type(return_value)}
expected dictionary '
+ 'for multiple_outputs')
Review comment:
As a user I would be confused: does it work or not? We should fail here
I think.
##########
File path: airflow/operators/python.py
##########
@@ -145,6 +149,131 @@ def execute_callable(self):
return self.python_callable(*self.op_args, **self.op_kwargs)
+class _PythonFunctionalOperator(BaseOperator):
+ """
+ Wraps a Python callable and captures args/kwargs when called for execution.
+
+ :param python_callable: A reference to an object that is callable
+ :type python_callable: python callable
+ :param op_kwargs: a dictionary of keyword arguments that will get unpacked
+ in your function
+ :type op_kwargs: dict (templated)
+ :param op_args: a list of positional arguments that will get unpacked when
+ calling your callable
+ :type op_args: list (templated)
+ :param multiple_outputs: if set, function return value will be
+ unrolled to multiple XCom values. Dict will unroll to xcom values with
keys as keys.
+ Defaults to False.
+ :type multiple_outputs: bool
+ """
+
+ template_fields = ('op_args', 'op_kwargs')
+ ui_color = '#ffefeb'
+
+ # since we won't mutate the arguments, we should just do the shallow copy
+ # there are some cases we can't deepcopy the objects(e.g protobuf).
+ shallow_copy_attrs = ('python_callable',)
+
+ @apply_defaults
+ def __init__(
+ self,
+ python_callable: Callable,
+ op_args: Tuple[Any],
+ op_kwargs: Dict[str, Any],
+ multiple_outputs: bool = False,
+ **kwargs
+ ) -> None:
+ kwargs['task_id'] = self._get_unique_task_id(kwargs['task_id'],
kwargs.get('dag', None))
+ super().__init__(**kwargs)
+ self.python_callable = python_callable
+
+ # Check that arguments can be binded
+ signature(python_callable).bind(*op_args, **op_kwargs)
+ self.multiple_outputs = multiple_outputs
+ self.op_args = op_args
+ self.op_kwargs = op_kwargs
+
+ @staticmethod
+ def _get_unique_task_id(task_id: str, dag: Optional[DAG]) -> str:
+ dag = dag or DagContext.get_current_dag()
+ if not dag or task_id not in dag.task_ids:
+ return task_id
+ core = re.split(r'__\d+$', task_id)[0]
+ suffixes = sorted(
+ [int(re.split(r'^.+__', task_id)[1])
+ for task_id in dag.task_ids
+ if re.match(rf'^{core}__\d+$', task_id)]
+ )
+ if not suffixes:
+ return f'{core}__1'
+ return f'{core}__{suffixes[-1] + 1}'
+
+ @staticmethod
+ def validate_python_callable(python_callable):
+ """Validate that python callable can be wrapped by operator.
Review comment:
```suggestion
"""
Validate that python callable can be wrapped by operator.
```
##########
File path: airflow/operators/python.py
##########
@@ -145,6 +149,131 @@ def execute_callable(self):
return self.python_callable(*self.op_args, **self.op_kwargs)
+class _PythonFunctionalOperator(BaseOperator):
+ """
+ Wraps a Python callable and captures args/kwargs when called for execution.
+
+ :param python_callable: A reference to an object that is callable
+ :type python_callable: python callable
+ :param op_kwargs: a dictionary of keyword arguments that will get unpacked
+ in your function
+ :type op_kwargs: dict (templated)
+ :param op_args: a list of positional arguments that will get unpacked when
+ calling your callable
+ :type op_args: list (templated)
+ :param multiple_outputs: if set, function return value will be
+ unrolled to multiple XCom values. Dict will unroll to xcom values with
keys as keys.
+ Defaults to False.
+ :type multiple_outputs: bool
+ """
+
+ template_fields = ('op_args', 'op_kwargs')
+ ui_color = '#ffefeb'
+
+ # since we won't mutate the arguments, we should just do the shallow copy
+ # there are some cases we can't deepcopy the objects(e.g protobuf).
+ shallow_copy_attrs = ('python_callable',)
+
+ @apply_defaults
+ def __init__(
+ self,
+ python_callable: Callable,
+ op_args: Tuple[Any],
+ op_kwargs: Dict[str, Any],
+ multiple_outputs: bool = False,
+ **kwargs
+ ) -> None:
+ kwargs['task_id'] = self._get_unique_task_id(kwargs['task_id'],
kwargs.get('dag', None))
+ super().__init__(**kwargs)
+ self.python_callable = python_callable
+
+ # Check that arguments can be binded
+ signature(python_callable).bind(*op_args, **op_kwargs)
+ self.multiple_outputs = multiple_outputs
+ self.op_args = op_args
+ self.op_kwargs = op_kwargs
+
+ @staticmethod
+ def _get_unique_task_id(task_id: str, dag: Optional[DAG]) -> str:
+ dag = dag or DagContext.get_current_dag()
+ if not dag or task_id not in dag.task_ids:
+ return task_id
+ core = re.split(r'__\d+$', task_id)[0]
+ suffixes = sorted(
+ [int(re.split(r'^.+__', task_id)[1])
+ for task_id in dag.task_ids
+ if re.match(rf'^{core}__\d+$', task_id)]
+ )
+ if not suffixes:
+ return f'{core}__1'
+ return f'{core}__{suffixes[-1] + 1}'
+
+ @staticmethod
+ def validate_python_callable(python_callable):
+ """Validate that python callable can be wrapped by operator.
+ Raises exception if invalid.
+
+ :param python_callable: Python object to be validated
+ :raises: TypeError, AirflowException
+ """
+ if not callable(python_callable):
+ raise TypeError('`python_callable` param must be callable')
+ if 'self' in signature(python_callable).parameters.keys():
+ raise AirflowException('@task does not support methods')
+
+ def execute(self, context: Dict):
+ return_value = self.python_callable(*self.op_args, **self.op_kwargs)
+ self.log.info("Done. Returned value was: %s", return_value)
+ if not self.multiple_outputs:
+ return return_value
+ if isinstance(return_value, dict):
+ for key in return_value.keys():
+ if not isinstance(key, str):
+ raise AirflowException('Returned dictionary keys must be
strings when using '
+ f'multiple_outputs, found {key}
({type(key)}) instead')
+ for key, value in return_value.items():
+ self.xcom_push(context, key, value)
+ else:
+ self.log.info(f'Returned output was type {type(return_value)}
expected dictionary '
+ 'for multiple_outputs')
+ return return_value
+
+
+def task(python_callable: Optional[Callable] = None, **kwargs):
+ """
+ Python operator decorator. Wraps a function into an Airflow operator.
+ Accepts kwargs for operator kwarg. Will try to wrap operator into DAG at
declaration or
+ on function invocation. Use alias to reuse function in the DAG.
Review comment:
Is this still valid?
##########
File path: airflow/operators/python.py
##########
@@ -145,6 +148,140 @@ def execute_callable(self):
return self.python_callable(*self.op_args, **self.op_kwargs)
+class _PythonFunctionalOperator(BaseOperator):
+ """
+ Wraps a Python callable and captures args/kwargs when called for execution.
+
+ :param python_callable: A reference to an object that is callable
+ :type python_callable: python callable
+ :param multiple_outputs: if set, function return value will be
+ unrolled to multiple XCom values. Dict will unroll to xcom values with
keys as keys.
+ Defaults to False.
+ :type multiple_outputs: bool
+ """
+
+ template_fields = ('_op_args', '_op_kwargs')
+ ui_color = '#ffefeb'
+
+ # since we won't mutate the arguments, we should just do the shallow copy
+ # there are some cases we can't deepcopy the objects(e.g protobuf).
+ shallow_copy_attrs = ('python_callable',)
+
+ @apply_defaults
+ def __init__(
+ self,
+ python_callable: Callable,
+ multiple_outputs: bool = False,
+ **kwargs
+ ) -> None:
+ self._validate_python_callable(python_callable)
+ super().__init__(**kwargs)
+ self.python_callable = python_callable
+ self.multiple_outputs = multiple_outputs
+ self._kwargs = kwargs
+ self._op_args: List[Any] = []
+ self._called = False
+ self._op_kwargs: Dict[str, Any] = {}
+
+ @staticmethod
+ def _get_unique_task_id(task_id: str, dag: Optional[DAG]) -> str:
+ dag = dag or DagContext.get_current_dag()
+ if not dag or task_id not in dag.task_ids:
+ return task_id
+ core = re.split(r'__\d+$', task_id)[0]
+ suffixes = sorted(
+ [int(re.split(r'^.+__', task_id)[1])
+ for task_id in dag.task_ids
+ if re.match(rf'^{core}__\d+$', task_id)]
+ )
+ if not suffixes:
+ return f'{core}__1'
+ return f'{core}__{suffixes[-1] + 1}'
Review comment:
We can use f-strings
##########
File path: airflow/operators/python.py
##########
@@ -145,6 +149,131 @@ def execute_callable(self):
return self.python_callable(*self.op_args, **self.op_kwargs)
+class _PythonFunctionalOperator(BaseOperator):
+ """
+ Wraps a Python callable and captures args/kwargs when called for execution.
+
+ :param python_callable: A reference to an object that is callable
+ :type python_callable: python callable
+ :param op_kwargs: a dictionary of keyword arguments that will get unpacked
+ in your function
+ :type op_kwargs: dict (templated)
+ :param op_args: a list of positional arguments that will get unpacked when
+ calling your callable
+ :type op_args: list (templated)
+ :param multiple_outputs: if set, function return value will be
+ unrolled to multiple XCom values. Dict will unroll to xcom values with
keys as keys.
+ Defaults to False.
+ :type multiple_outputs: bool
+ """
+
+ template_fields = ('op_args', 'op_kwargs')
+ ui_color = '#ffefeb'
+
+ # since we won't mutate the arguments, we should just do the shallow copy
+ # there are some cases we can't deepcopy the objects(e.g protobuf).
+ shallow_copy_attrs = ('python_callable',)
+
+ @apply_defaults
+ def __init__(
+ self,
+ python_callable: Callable,
+ op_args: Tuple[Any],
+ op_kwargs: Dict[str, Any],
+ multiple_outputs: bool = False,
+ **kwargs
+ ) -> None:
+ kwargs['task_id'] = self._get_unique_task_id(kwargs['task_id'],
kwargs.get('dag', None))
+ super().__init__(**kwargs)
+ self.python_callable = python_callable
+
+ # Check that arguments can be binded
+ signature(python_callable).bind(*op_args, **op_kwargs)
+ self.multiple_outputs = multiple_outputs
+ self.op_args = op_args
+ self.op_kwargs = op_kwargs
+
+ @staticmethod
+ def _get_unique_task_id(task_id: str, dag: Optional[DAG]) -> str:
+ dag = dag or DagContext.get_current_dag()
+ if not dag or task_id not in dag.task_ids:
+ return task_id
+ core = re.split(r'__\d+$', task_id)[0]
+ suffixes = sorted(
+ [int(re.split(r'^.+__', task_id)[1])
+ for task_id in dag.task_ids
+ if re.match(rf'^{core}__\d+$', task_id)]
+ )
+ if not suffixes:
+ return f'{core}__1'
+ return f'{core}__{suffixes[-1] + 1}'
+
+ @staticmethod
+ def validate_python_callable(python_callable):
+ """Validate that python callable can be wrapped by operator.
+ Raises exception if invalid.
+
+ :param python_callable: Python object to be validated
+ :raises: TypeError, AirflowException
+ """
+ if not callable(python_callable):
+ raise TypeError('`python_callable` param must be callable')
+ if 'self' in signature(python_callable).parameters.keys():
+ raise AirflowException('@task does not support methods')
+
+ def execute(self, context: Dict):
+ return_value = self.python_callable(*self.op_args, **self.op_kwargs)
+ self.log.info("Done. Returned value was: %s", return_value)
+ if not self.multiple_outputs:
+ return return_value
+ if isinstance(return_value, dict):
+ for key in return_value.keys():
+ if not isinstance(key, str):
+ raise AirflowException('Returned dictionary keys must be
strings when using '
+ f'multiple_outputs, found {key}
({type(key)}) instead')
+ for key, value in return_value.items():
+ self.xcom_push(context, key, value)
+ else:
+ self.log.info(f'Returned output was type {type(return_value)}
expected dictionary '
+ 'for multiple_outputs')
+ return return_value
+
+
+def task(python_callable: Optional[Callable] = None, **kwargs):
+ """
+ Python operator decorator. Wraps a function into an Airflow operator.
+ Accepts kwargs for operator kwarg. Will try to wrap operator into DAG at
declaration or
+ on function invocation. Use alias to reuse function in the DAG.
+
+ :param python_callable: Function to decorate
+ :type python_callable: Optional[Callable]
+ :param multiple_outputs: if set, function return value will be
+ unrolled to multiple XCom values. List/Tuples will unroll to xcom
values
+ with index as key. Dict will unroll to xcom values with keys as keys.
+ Defaults to False.
+ :type multiple_outputs: bool
Review comment:
This is an unexpected parameter. There's no such param in `task`
signature.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
