Mark arbitrary function as tasks with @pytask.mark.task. (#200)

Author: tobiasraabe

docs/source/changes.rst

@@ -14,9 +14,12 @@ all releases are available on `PyPI <https://pypi.org/project/pytask>`_ and
 - :gh:`193` adds more figures to the documentation.
 - :gh:`194` updates the ``README.rst``.
 - :gh:`196` references the two new cookiecutters for projects and plugins.
-- :gh:`198` fixes the documentation of ``@pytask.mark.skipif``. (Closes :gh:`195`)
+- :gh:`198` fixes the documentation of :func:`@pytask.mark.skipif
+  <_pytask.skipping.skipif>`. (Closes :gh:`195`)
 - :gh:`199` extends the error message when paths are ambiguous on case-insensitive file
   systems.
+- :gh:`200` implements the :func:`@pytask.mark.task <_pytask.task.task>` decorator to
+  mark functions as tasks regardless whether they are prefixed with ``task_`` or not.
 
 
 0.1.5 - 2022-01-10

docs/source/tutorials/how_to_write_a_task.rst

@@ -4,8 +4,11 @@ How to write a task
 Starting from the project structure in the :doc:`previous tutorial
 <how_to_set_up_a_project>`, this tutorial teaches you how to write your first task.
 
-The task will be defined in ``src/my_project/task_data_preparation.py`` and it will
-generate artificial data which will be stored in ``bld/data.pkl``. We will call the
+By default, pytask will look for tasks in modules whose name is prefixed with ``task_``.
+Tasks are functions in these modules whose name also starts with ``task_``.
+
+Our first task will be defined in ``src/my_project/task_data_preparation.py`` and it
+will generate artificial data which will be stored in ``bld/data.pkl``. We will call the
 function in the module :func:`task_create_random_data`.
 
 .. code-block::
@@ -62,10 +65,35 @@ subsequent directories.
 
 .. image:: /_static/images/how-to-write-a-task.png
 
-.. important::
 
-    By default, pytask assumes that tasks are functions and both, the function name and
-    the module name, must be prefixed with ``task_``.
+Customize task names
+--------------------
+
+Use the :func:`@pytask.mark.task <_pytask.task_utils.task>` decorator to mark a function
+as a task regardless of its function name. You can optionally pass a new name for the
+task. Otherwise, the function name is used.
+
+.. code-block:: python
+
+    # The id will be '.../task_data_preparation.py::create_random_data'
+
+
+    @pytask.mark.task
+    def create_random_data():
+        ...
+
+
+    # The id will be '.../task_data_preparation.py::create_data'
+
+
+    @pytask.mark.task(name="create_data")
+    def create_random_data():
+        ...
+
+
+Customize task module names
+---------------------------
 
-    Use the configuration value :confval:`task_files` if you prefer a different naming
-    scheme for the task modules.
+Use the configuration value :confval:`task_files` if you prefer a different naming
+scheme for the task modules. By default, it is set to ``task_*.py``. You can specify one
+or multiple patterns to collect tasks from other files.

src/_pytask/cli.py

@@ -66,6 +66,7 @@ def pytask_add_hooks(pm: pluggy.PluginManager) -> None:
     from _pytask import profile
     from _pytask import resolve_dependencies
     from _pytask import skipping
+    from _pytask import task
 
     pm.register(build)
     pm.register(capture)
@@ -86,6 +87,7 @@ def pytask_add_hooks(pm: pluggy.PluginManager) -> None:
     pm.register(profile)
     pm.register(resolve_dependencies)
     pm.register(skipping)
+    pm.register(task)
 
 
 @click.group(

src/_pytask/mark/structures.py

@@ -182,7 +182,7 @@ def store_mark(obj: Callable[..., Any], mark: Mark) -> None:
     assert isinstance(mark, Mark), mark
     # Always reassign name to avoid updating pytaskmark in a reference that was only
     # borrowed.
-    obj.pytaskmark = get_unpacked_marks(obj) + [mark]  # type: ignore
+    obj.pytaskmark = get_unpacked_marks(obj) + [mark]  # type: ignore[attr-defined]
 
 
 class MarkGenerator:

src/_pytask/mark_utils.py

@@ -5,6 +5,7 @@
 """
 from typing import Any
 from typing import List
+from typing import Tuple
 from typing import TYPE_CHECKING
 
 
@@ -30,3 +31,12 @@ def get_marks_from_obj(obj: Any, marker_name: str) -> "List[Mark]":
 def has_marker(obj: Any, marker_name: str) -> bool:
     """Determine whether a task function has a certain marker."""
     return any(marker.name == marker_name for marker in getattr(obj, "pytaskmark", []))
+
+
+def remove_markers_from_func(obj: Any, marker_name: str) -> Tuple[Any, List["Mark"]]:
+    """Remove parametrize markers from the object."""
+    markers = [i for i in getattr(obj, "pytaskmark", []) if i.name == marker_name]
+    others = [i for i in getattr(obj, "pytaskmark", []) if i.name != marker_name]
+    obj.pytaskmark = others
+
+    return obj, markers

src/_pytask/parametrize.py

@@ -18,8 +18,11 @@
 from _pytask.console import TASK_ICON
 from _pytask.mark import Mark
 from _pytask.mark import MARK_GEN as mark  # noqa: N811
+from _pytask.mark_utils import has_marker
+from _pytask.mark_utils import remove_markers_from_func
 from _pytask.nodes import find_duplicates
 from _pytask.session import Session
+from _pytask.task_utils import parse_task_marker
 
 
 def parametrize(
@@ -93,14 +96,19 @@ def pytask_parametrize_task(
 
     """
     if callable(obj):
-        obj, markers = _remove_parametrize_markers_from_func(obj)
+        obj, markers = remove_markers_from_func(obj, "parametrize")
 
         if len(markers) > 1:
             raise NotImplementedError(
                 "Multiple parametrizations are currently not implemented since it is "
                 "not possible to define products for tasks from a Cartesian product."
             )
 
+        if has_marker(obj, "task"):
+            parsed_name = parse_task_marker(obj)
+            if parsed_name is not None:
+                name = parsed_name
+
         base_arg_names, arg_names, arg_values = _parse_parametrize_markers(
             markers, name
         )
@@ -119,8 +127,9 @@ def pytask_parametrize_task(
 
             # Copy function and attributes to allow in-place changes.
             func = _copy_func(obj)  # type: ignore
-            func.pytaskmark = copy.deepcopy(obj.pytaskmark)  # type: ignore
-
+            func.pytaskmark = copy.deepcopy(  # type: ignore[attr-defined]
+                obj.pytaskmark  # type: ignore[attr-defined]
+            )
             # Convert parametrized dependencies and products to decorator.
             session.hook.pytask_parametrize_kwarg_to_marker(obj=func, kwargs=kwargs)
 
@@ -148,15 +157,6 @@ def pytask_parametrize_task(
         return names_and_functions
 
 
-def _remove_parametrize_markers_from_func(obj: Any) -> Tuple[Any, List[Mark]]:
-    """Remove parametrize markers from the object."""
-    parametrize_markers = [i for i in obj.pytaskmark if i.name == "parametrize"]
-    others = [i for i in obj.pytaskmark if i.name != "parametrize"]
-    obj.pytaskmark = others
-
-    return obj, parametrize_markers
-
-
 def _parse_parametrize_marker(
     marker: Mark, name: str
 ) -> Tuple[Tuple[str, ...], List[Tuple[str, ...]], List[Tuple[Any, ...]]]:

src/_pytask/task.py

@@ -0,0 +1,37 @@
+from pathlib import Path
+from typing import Any
+from typing import Dict
+from typing import Optional
+
+from _pytask.config import hookimpl
+from _pytask.mark_utils import has_marker
+from _pytask.nodes import PythonFunctionTask
+from _pytask.session import Session
+from _pytask.task_utils import parse_task_marker
+
+
+@hookimpl
+def pytask_parse_config(config: Dict[str, Any]) -> None:
+    config["markers"]["task"] = "Mark a function as a task regardless of its name."
+
+
+@hookimpl
+def pytask_collect_task(
+    session: Session, path: Path, name: str, obj: Any
+) -> Optional[PythonFunctionTask]:
+    """Collect a task which is a function.
+
+    There is some discussion on how to detect functions in this `thread
+    <https://stackoverflow.com/q/624926/7523785>`_. :class:`types.FunctionType` does not
+    detect built-ins which is not possible anyway.
+
+    """
+    if has_marker(obj, "task") and callable(obj):
+        parsed_name = parse_task_marker(obj)
+        if parsed_name is not None:
+            name = parsed_name
+        return PythonFunctionTask.from_path_name_function_session(
+            path, name, obj, session
+        )
+    else:
+        return None

src/_pytask/task_utils.py

@@ -0,0 +1,27 @@
+from typing import Any
+from typing import Callable
+from typing import Optional
+
+from _pytask.mark import Mark
+from _pytask.mark_utils import remove_markers_from_func
+
+
+def task(name: Optional[str] = None) -> str:
+    return name
+
+
+def parse_task_marker(obj: Callable[..., Any]) -> str:
+    obj, task_markers = remove_markers_from_func(obj, "task")
+
+    if len(task_markers) != 1:
+        raise ValueError(
+            "The @pytask.mark.task decorator cannot be applied more than once to a "
+            "single task."
+        )
+    task_marker = task_markers[0]
+
+    name = task(*task_marker.args, **task_marker.kwargs)
+
+    obj.pytaskmark.append(Mark("task", (), {}))  # type: ignore[attr-defined]
+
+    return name

tests/test_task.py

@@ -0,0 +1,68 @@
+import textwrap
+
+import pytest
+from _pytask.nodes import create_task_name
+from _pytask.outcomes import ExitCode
+from pytask import main
+
+
+@pytest.mark.parametrize("func_name", ["task_example", "func"])
+@pytest.mark.parametrize("task_name", ["the_only_task", None])
+def test_task_with_task_decorator(tmp_path, func_name, task_name):
+    task_decorator_input = f"{task_name!r}" if task_name else task_name
+    source = f"""
+    import pytask
+
+    @pytask.mark.task({task_decorator_input})
+    @pytask.mark.produces("out.txt")
+    def {func_name}(produces):
+        produces.write_text("Hello. It's me.")
+    """
+    tmp_path.joinpath("task_module.py").write_text(textwrap.dedent(source))
+
+    session = main({"paths": tmp_path})
+
+    assert session.exit_code == ExitCode.OK
+
+    if task_name:
+        assert session.tasks[0].name == create_task_name(
+            tmp_path.joinpath("task_module.py"), task_name
+        )
+    else:
+        assert session.tasks[0].name == create_task_name(
+            tmp_path.joinpath("task_module.py"), func_name
+        )
+
+
+@pytest.mark.parametrize("func_name", ["task_example", "func"])
+@pytest.mark.parametrize("task_name", ["the_only_task", None])
+def test_task_with_task_decorator_with_parametrize(tmp_path, func_name, task_name):
+    task_decorator_input = f"{task_name!r}" if task_name else task_name
+    source = f"""
+    import pytask
+
+    @pytask.mark.task({task_decorator_input})
+    @pytask.mark.parametrize("produces", ["out_1.txt", "out_2.txt"])
+    def {func_name}(produces):
+        produces.write_text("Hello. It's me.")
+    """
+    tmp_path.joinpath("task_module.py").write_text(textwrap.dedent(source))
+
+    session = main({"paths": tmp_path})
+
+    assert session.exit_code == ExitCode.OK
+
+    if task_name:
+        assert session.tasks[0].name == create_task_name(
+            tmp_path.joinpath("task_module.py"), f"{task_name}[out_1.txt]"
+        )
+        assert session.tasks[1].name == create_task_name(
+            tmp_path.joinpath("task_module.py"), f"{task_name}[out_2.txt]"
+        )
+    else:
+        assert session.tasks[0].name == create_task_name(
+            tmp_path.joinpath("task_module.py"), f"{func_name}[out_1.txt]"
+        )
+        assert session.tasks[1].name == create_task_name(
+            tmp_path.joinpath("task_module.py"), f"{func_name}[out_2.txt]"
+        )