casassg commented on a change in pull request #8962:
URL: https://github.com/apache/airflow/pull/8962#discussion_r431450538
##########
File path: airflow/operators/python.py
##########
@@ -145,6 +147,141 @@ 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:
+ # Check if we need to generate a new task_id
+ task_id = kwargs.get('task_id', None)
+ dag = kwargs.get('dag', None) or DagContext.get_current_dag()
+ if task_id and dag and task_id in dag.task_ids:
+ prefix = task_id.rsplit("__", 1)[0]
+ task_id = sorted(
+ filter(lambda x: x.startswith(prefix), dag.task_ids),
+ reverse=True
+ )[0]
+ num = int(task_id[-1] if '__' in task_id else '0') + 1
+ kwargs['task_id'] = f'{prefix}__{num}'
+
+ if not kwargs.get('do_xcom_push', True) and not multiple_outputs:
+ raise AirflowException('@task needs to have either
do_xcom_push=True or '
+ 'multiple_outputs=True.')
+ if not callable(python_callable):
+ raise AirflowException('`python_callable` param must be callable')
+ self._fail_if_method(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 _fail_if_method(python_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
Review comment:
With current approach you get both: unrolled and not unrolled. You
access each as needed. It's flexible and easier to use when you need several
outputs. Also it mimics other frameworks, so it should feel natural for the
user.
Resolving comment for now, if someone has anything else to discuss on the
topic, please reopen or start new thread.
----------------------------------------------------------------
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
