{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Python-tasks\n",
    "\n",
    "Python tasks are Python functions that are parameterised in a separate step before\n",
    "they are executed or added to a workflow.\n",
    "\n",
    "## Define decorator\n",
    "\n",
    "The simplest way to define a Python task is to decorate a function with `pydra.compose.python.define`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pydra.compose import python\n",
    "\n",
    "\n",
    "# Note that we use PascalCase because the object returned by the decorator is actually a class\n",
    "@python.define\n",
    "def MyFirstTask(a, b):\n",
    "    \"\"\"Sample function for testing\"\"\"\n",
    "    return a + b"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The resulting task-definition class can be then parameterized (instantiated), and\n",
    "executed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Instantiate the task, setting all parameters\n",
    "my_first_task = MyFirstTask(a=1, b=2.0)\n",
    "\n",
    "# Execute the task\n",
    "outputs = my_first_task()\n",
    "\n",
    "print(outputs.out)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "By default, the name of the output field for a function with only one output is `out`. To\n",
    "name this something else, or in the case where there are multiple output fields, the `outputs`\n",
    "argument can be provided to `python.define`\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@python.define(outputs=[\"c\", \"d\"])\n",
    "def NamedOutputTask(a, b):\n",
    "    \"\"\"Sample function for testing\"\"\"\n",
    "    return a + b, a - b\n",
    "\n",
    "\n",
    "named_output_task = NamedOutputTask(a=2, b=1)\n",
    "\n",
    "outputs = named_output_task()\n",
    "\n",
    "print(outputs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The input and output field attributes automatically extracted from the function, explicit\n",
    "attributes can be augmented"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@python.define(\n",
    "    inputs={\"a\": python.arg(allowed_values=[1, 2, 3]), \"b\": python.arg(default=10.0)},\n",
    "    outputs={\n",
    "        \"c\": python.out(type=float, help=\"the sum of the inputs\"),\n",
    "        \"d\": python.out(type=float, help=\"the difference of the inputs\"),\n",
    "    },\n",
    ")\n",
    "def AugmentedTask(a, b):\n",
    "    \"\"\"Sample function for testing\"\"\"\n",
    "    return a + b, a - b"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Type annotations\n",
    "\n",
    "If provided, type annotations are included in the task, and are checked at\n",
    "the time of parameterisation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pydra.compose import python\n",
    "\n",
    "\n",
    "@python.define\n",
    "def MyTypedTask(a: int, b: float) -> float:\n",
    "    \"\"\"Sample function for testing\"\"\"\n",
    "    return a + b\n",
    "\n",
    "\n",
    "try:\n",
    "    # 1.5 is not an integer so this should raise a TypeError\n",
    "    my_typed_task = MyTypedTask(a=1.5, b=2.0)\n",
    "except TypeError as e:\n",
    "    print(f\"Type error caught: {e}\")\n",
    "else:\n",
    "    assert False, \"Expected a TypeError\"\n",
    "\n",
    "# While 2 is an integer, it can be implicitly coerced to a float\n",
    "my_typed_task = MyTypedTask(a=1, b=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "## Docstring parsing\n",
    "\n",
    "Instead of explicitly providing help strings and output names in `inputs` and `outputs`\n",
    "arguments, if the function describes the its inputs and/or outputs in the doc string, \n",
    "in either reST, Google or NumpyDoc style, then they will be extracted and included in the\n",
    "input or output fields\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pydra.utils import print_help\n",
    "\n",
    "\n",
    "@python.define(outputs=[\"c\", \"d\"])\n",
    "def DocStrExample(a: int, b: float) -> tuple[float, float]:\n",
    "    \"\"\"Example python task with help strings pulled from doc-string\n",
    "\n",
    "    Args:\n",
    "        a: First input\n",
    "            to be inputted\n",
    "        b: Second input\n",
    "\n",
    "    Returns:\n",
    "        c: Sum of a and b\n",
    "        d: Product of a and b\n",
    "    \"\"\"\n",
    "    return a + b, a * b\n",
    "\n",
    "\n",
    "print_help(DocStrExample)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Wrapping external functions\n",
    "\n",
    "Like all decorators, `python.define` is just a function, so can also be used to convert\n",
    "a function that is defined separately into a Python task."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "NumpyCorrelate = python.define(np.correlate)\n",
    "\n",
    "numpy_correlate = NumpyCorrelate(a=[1, 2, 3], v=[0, 1, 0.5])\n",
    "\n",
    "outputs = numpy_correlate()\n",
    "\n",
    "print(outputs.out)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Like with decorated functions, input and output fields can be explicitly augmented via\n",
    "the `inputs` and `outputs` arguments"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "NumpyCorrelate = python.define(np.correlate, outputs=[\"correlation\"])\n",
    "\n",
    "numpy_correlate = NumpyCorrelate(a=[1, 2, 3], v=[0, 1, 0.5])\n",
    "\n",
    "outputs = numpy_correlate()\n",
    "\n",
    "print(outputs.correlation)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "wf13",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}