.. DO NOT EDIT.
.. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY.
.. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE:
.. "auto/core/flyte_basics/decorating_tasks.py"
.. LINE NUMBERS ARE GIVEN BELOW.

.. only:: html

    .. note::
        :class: sphx-glr-download-link-note

        Click :ref:`here <sphx_glr_download_auto_core_flyte_basics_decorating_tasks.py>`
        to download the full example code

.. rst-class:: sphx-glr-example-title

.. _sphx_glr_auto_core_flyte_basics_decorating_tasks.py:


Decorating Tasks
----------------

A simple way of modifying the behavior of tasks is by using decorators to wrap your task functions.

In order to make sure that your decorated function contains all the type annotation and docstring
information that Flyte needs, you'll need to use the built-in :py:func:`~functools.wraps` decorator.

.. GENERATED FROM PYTHON SOURCE LINES 11-20

.. code-block:: default


    import logging
    from functools import partial, wraps

    from flytekit import task, workflow

    logger = logging.getLogger(__file__)



.. GENERATED FROM PYTHON SOURCE LINES 21-25

Using a Single Decorator
^^^^^^^^^^^^^^^^^^^^^^^^

Here we define decorator that logs the input and output information of a decorated task.

.. GENERATED FROM PYTHON SOURCE LINES 26-39

.. code-block:: default



    def log_io(fn):
        @wraps(fn)
        def wrapper(*args, **kwargs):
            logger.info(f"task {fn.__name__} called with args: {args}, kwargs: {kwargs}")
            out = fn(*args, **kwargs)
            logger.info(f"task {fn.__name__} output: {out}")
            return out

        return wrapper



.. GENERATED FROM PYTHON SOURCE LINES 40-44

Next, we define a task called ``t1`` which is decorated with ``log_io``.

.. note::
   The order of invoking the decorators is important. ``@task`` should always be the outer-most decorator.

.. GENERATED FROM PYTHON SOURCE LINES 44-52

.. code-block:: default



    @task
    @log_io
    def t1(x: int) -> int:
        return x + 1



.. GENERATED FROM PYTHON SOURCE LINES 53-63

Stacking Multiple Decorators
^^^^^^^^^^^^^^^^^^^^^^^^^^^^

You can also stack multiple decorators on top of each other as long as ``@task`` is the outer-most decorator.

Here we define a decorator that validates that the output of the decorated function is a positive number before
returning it, raising a ``ValueError`` if it violates this assumption.

.. note::
   The ``validate_output`` output uses :py:func:`~functools.partial` to implement parameterized decorators.

.. GENERATED FROM PYTHON SOURCE LINES 64-82

.. code-block:: default



    def validate_output(fn=None, *, floor=0):
        @wraps(fn)
        def wrapper(*args, **kwargs):
            out = fn(*args, **kwargs)
            if out <= floor:
                raise ValueError(
                    f"output of task {fn.__name__} must be a positive number, found {out}"
                )
            return out

        if fn is None:
            return partial(validate_output, floor=floor)

        return wrapper



.. GENERATED FROM PYTHON SOURCE LINES 83-84

Now let's define a function that uses both the logging and validator decorators:

.. GENERATED FROM PYTHON SOURCE LINES 84-93

.. code-block:: default



    @task
    @log_io
    @validate_output(floor=10)
    def t2(x: int) -> int:
        return x + 10



.. GENERATED FROM PYTHON SOURCE LINES 94-95

Finally, we compose a workflow that calls ``t1`` and ``t2``.

.. GENERATED FROM PYTHON SOURCE LINES 95-106

.. code-block:: default



    @workflow
    def wf(x: int) -> int:
        return t2(x=t1(x=x))


    if __name__ == "__main__":
        print(f"Running wf(x=10) {wf(x=10)}")



.. GENERATED FROM PYTHON SOURCE LINES 107-111

In this example, you learned how to modify the behavior of tasks via function decorators using the built-in
:py:func:`~functools.wraps` decorator pattern. To learn more about how to extend Flyte at a deeper level, for
example creating custom types, custom tasks, or backend plugins,
see :ref:`Extending Flyte <sphx_glr_auto_core_extend_flyte>`.


.. rst-class:: sphx-glr-timing

   **Total running time of the script:** ( 0 minutes  0.000 seconds)


.. _sphx_glr_download_auto_core_flyte_basics_decorating_tasks.py:


.. only :: html

 .. container:: sphx-glr-footer
    :class: sphx-glr-footer-example



  .. container:: sphx-glr-download sphx-glr-download-python

     :download:`Download Python source code: decorating_tasks.py <decorating_tasks.py>`



  .. container:: sphx-glr-download sphx-glr-download-jupyter

     :download:`Download Jupyter notebook: decorating_tasks.ipynb <decorating_tasks.ipynb>`


.. only:: html

 .. rst-class:: sphx-glr-signature

    `Gallery generated by Sphinx-Gallery <https://sphinx-gallery.github.io>`_