Allow to schedule a task using its function

Using the only the task name when scheduling a task is limiting as it
makes refactoring less obvious than using the function directly.

However having the possibility to use the name for scheduling is still
important when importing the function is not possible.

Author: NicolasLM

README.rst

@@ -44,7 +44,7 @@ Create task and schedule two jobs, one executed now and one later:
 
 
     # Schedule a job to be executed ASAP
-    spin.schedule('compute', 5, 3)
+    spin.schedule(compute, 5, 3)
 
     print('Starting workers, ^C to quit')
     spin.start_workers()

doc/index.rst

@@ -23,6 +23,24 @@ Quickstart
 
 .. literalinclude:: ../examples/quickstart.py
 
+The :class:`Engine` is the central part of Spinach, it allows to define tasks,
+schedule jobs to execute in the background and start background workers.
+:ref:`More details <engine>`.
+
+The Broker is the backend that background workers use to retrieve jobs to
+execute. Spinach provides two brokers: MemoryBroker for development and
+RedisBroker for production.
+
+The :meth:`Engine.task` decorator is used to register tasks. It requires at
+least a `name` to identify the task, but other options can be given to
+customize how the task behaves. :ref:`More details <tasks>`.
+
+Background jobs can then be scheduled by using either the task name or the task
+function::
+
+    spin.schedule('compute', 5, 3)  # identify a task by its name
+    spin.schedule(compute, 5, 3)    # identify a task by its function
+
 Getting started with spinach:
 
 .. toctree::

examples/queues.py

@@ -21,7 +21,7 @@ def slow():
     time.sleep(10)
 
 
-spin.schedule('slow')
-spin.schedule('fast')
+spin.schedule(slow)
+spin.schedule(fast)
 
 spin.start_workers(number=1, queue='high-priority', stop_when_queue_empty=True)

examples/quickstart.py

@@ -9,7 +9,7 @@ def compute(a, b):
 
 
 # Schedule a job to be executed ASAP
-spin.schedule('compute', 5, 3)
+spin.schedule(compute, 5, 3)
 
 print('Starting workers, ^C to quit')
 spin.start_workers()

spinach/contrib/spinachd/mail.py

@@ -6,7 +6,7 @@
 from django.core.mail import EmailMessage
 from django.core.mail.backends.base import BaseEmailBackend
 
-from .tasks import tasks
+from .tasks import tasks, send_emails
 
 
 class BackgroundEmailBackend(BaseEmailBackend):
@@ -17,7 +17,7 @@ def send_messages(self, messages):
             message.message()  # .message() triggers header validation
             msg_count += 1
         messages = serialize_email_messages(messages)
-        tasks.schedule('spinachd:send_emails', messages)
+        tasks.schedule(send_emails, messages)
 
         return msg_count
 

spinach/engine.py

@@ -2,7 +2,7 @@
 from logging import getLogger
 import threading
 
-from .task import Tasks, Batch
+from .task import Tasks, Batch, Schedulable
 from .utils import run_forever, handle_sigterm
 from .job import Job, JobStatus, advance_job_status
 from .brokers.base import Broker
@@ -58,31 +58,31 @@ def attach_tasks(self, tasks: Tasks):
         self._tasks.update(tasks)
         tasks._spin = self
 
-    def execute(self, task_name: str, *args, **kwargs):
-        return self._tasks.get(task_name).func(*args, **kwargs)
+    def execute(self, task: Schedulable, *args, **kwargs):
+        return self._tasks.get(task).func(*args, **kwargs)
 
-    def schedule(self, task_name: str, *args, **kwargs):
+    def schedule(self, task: Schedulable, *args, **kwargs):
         """Schedule a job to be executed as soon as possible.
 
-        :arg task_name: name of the task to execute
+        :arg task: the task or its name to execute in the background
         :arg args: args to be passed to the task function
         :arg kwargs: kwargs to be passed to the task function
         """
         at = datetime.now(timezone.utc)
-        return self.schedule_at(task_name, at, *args, **kwargs)
+        return self.schedule_at(task, at, *args, **kwargs)
 
-    def schedule_at(self, task_name: str, at: datetime, *args, **kwargs):
+    def schedule_at(self, task: Schedulable, at: datetime, *args, **kwargs):
         """Schedule a job to be executed in the future.
 
-        :arg task_name: name of the task to execute
+        :arg task: the task or its name to execute in the background
         :arg at: date at which the job should start. It is advised to pass a
                  timezone aware datetime to lift any ambiguity. However if a
                  timezone naive datetime if given, it will be assumed to
                  contain UTC time.
         :arg args: args to be passed to the task function
         :arg kwargs: kwargs to be passed to the task function
         """
-        task = self._tasks.get(task_name)
+        task = self._tasks.get(task)
         job = Job(task.name, task.queue, at, task.max_retries, task_args=args,
                   task_kwargs=kwargs)
         return self._broker.enqueue_jobs([job])
@@ -96,8 +96,8 @@ def schedule_batch(self, batch: Batch):
         :arg batch: :class:`Batch` instance containing jobs to schedule
         """
         jobs = list()
-        for task_name, at, args, kwargs in batch.jobs_to_create:
-            task = self._tasks.get(task_name)
+        for task, at, args, kwargs in batch.jobs_to_create:
+            task = self._tasks.get(task)
             jobs.append(
                 Job(task.name, task.queue, at, task.max_retries,
                     task_args=args, task_kwargs=kwargs)

spinach/task.py

@@ -1,7 +1,7 @@
 from datetime import datetime, timezone, timedelta
 import functools
 import json
-from typing import Optional, Callable, List
+from typing import Optional, Callable, List, Union
 from numbers import Number
 
 from . import const, exc
@@ -29,6 +29,10 @@ def serialize(self):
             'periodicity': periodicity
         }, sort_keys=True)
 
+    @property
+    def task_name(self):
+        return self.name
+
     def __repr__(self):
         return 'Task({}, {}, {}, {}, {})'.format(
             self.func, self.name, self.queue, self.max_retries,
@@ -45,6 +49,9 @@ def __eq__(self, other):
         return True
 
 
+Schedulable = Union[str, Callable, Task]
+
+
 class Tasks:
     """Registry for tasks to be used by Spinach.
 
@@ -76,8 +83,12 @@ def names(self) -> List[str]:
     def tasks(self) -> dict:
         return self._tasks
 
-    def get(self, name: str):
-        task = self._tasks.get(name)
+    def get(self, name: Schedulable) -> Task:
+        try:
+            task_name = name.task_name
+        except AttributeError:
+            task_name = name
+        task = self._tasks.get(task_name)
         if task is not None:
             return task
 
@@ -110,6 +121,10 @@ def task(self, func: Optional[Callable]=None, name: Optional[str]=None,
         self.add(func, name=name, queue=queue, max_retries=max_retries,
                  periodicity=periodicity)
 
+        # Add an attribute to the function to be able to conveniently use it as
+        # spin.schedule(function) instead of spin.schedule('task_name')
+        func.task_name = name
+
         return func
 
     def add(self, func: Callable, name: Optional[str]=None,
@@ -161,23 +176,23 @@ def _require_attached_tasks(self):
                 'a Spinach Engine.'
             )
 
-    def schedule(self, task_name: str, *args, **kwargs):
+    def schedule(self, task: Schedulable, *args, **kwargs):
         """Schedule a job to be executed as soon as possible.
 
-        :arg task_name: name of the task to execute in the background
+        :arg task: the task or its name to execute in the background
         :arg args: args to be passed to the task function
         :arg kwargs: kwargs to be passed to the task function
 
         This method can only be used once tasks have been attached to a
         Spinach :class:`Engine`.
         """
         self._require_attached_tasks()
-        self._spin.schedule(task_name, *args, **kwargs)
+        self._spin.schedule(task, *args, **kwargs)
 
-    def schedule_at(self, task_name: str, at: datetime, *args, **kwargs):
+    def schedule_at(self, task: Schedulable, at: datetime, *args, **kwargs):
         """Schedule a job to be executed in the future.
 
-        :arg task_name: name of the task to execute in the background
+        :arg task: the task or its name to execute in the background
         :arg at: Date at which the job should start. It is advised to pass a
                  timezone aware datetime to lift any ambiguity. However if a
                  timezone naive datetime if given, it will be assumed to
@@ -189,7 +204,7 @@ def schedule_at(self, task_name: str, at: datetime, *args, **kwargs):
         Spinach :class:`Engine`.
         """
         self._require_attached_tasks()
-        self._spin.schedule_at(task_name, at, *args, **kwargs)
+        self._spin.schedule_at(task, at, *args, **kwargs)
 
     def schedule_batch(self, batch: 'Batch'):
         """Schedule many jobs at once.
@@ -225,28 +240,28 @@ class Batch:
     def __init__(self):
         self.jobs_to_create = list()
 
-    def schedule(self, task_name: str, *args, **kwargs):
+    def schedule(self, task: Schedulable, *args, **kwargs):
         """Add a job to be executed ASAP to the batch.
 
-        :arg task_name: name of the task to execute in the background
+        :arg task: the task or its name to execute in the background
         :arg args: args to be passed to the task function
         :arg kwargs: kwargs to be passed to the task function
         """
         at = datetime.now(timezone.utc)
-        self.schedule_at(task_name, at, *args, **kwargs)
+        self.schedule_at(task, at, *args, **kwargs)
 
-    def schedule_at(self, task_name: str, at: datetime, *args, **kwargs):
+    def schedule_at(self, task: Schedulable, at: datetime, *args, **kwargs):
         """Add a job to be executed in the future to the batch.
 
-        :arg task_name: name of the task to execute in the background
+        :arg task: the task or its name to execute in the background
         :arg at: Date at which the job should start. It is advised to pass a
                  timezone aware datetime to lift any ambiguity. However if a
                  timezone naive datetime if given, it will be assumed to
                  contain UTC time.
         :arg args: args to be passed to the task function
         :arg kwargs: kwargs to be passed to the task function
         """
-        self.jobs_to_create.append((task_name, at, args, kwargs))
+        self.jobs_to_create.append((task, at, args, kwargs))
 
 
 class RetryException(Exception):

tests/test_task.py

@@ -104,11 +104,13 @@ def bar():
         pass
 
     assert 'foo' in str(tasks.tasks['foo'].func)
+    assert foo.task_name == 'foo'
     assert tasks.tasks['foo'].name == 'foo'
     assert tasks.tasks['foo'].queue == 'tasks_queue'
     assert tasks.tasks['foo'].max_retries == const.DEFAULT_MAX_RETRIES
 
     assert 'bar' in str(tasks.tasks['bar'].func)
+    assert bar.task_name == 'bar'
     assert tasks.tasks['bar'].name == 'bar'
     assert tasks.tasks['bar'].queue == 'task_queue'
     assert tasks.tasks['bar'].max_retries == 20
@@ -145,18 +147,53 @@ def test_tasks_names():
     assert sorted(tasks.names) == ['bar', 'foo']
 
 
-def test_tasks_get():
+def test_tasks_get_by_name():
     tasks = Tasks()
-    with pytest.raises(exc.UnknownTask):
-        tasks.get('foo')
-
     tasks.add(print, 'foo')
+
     r = tasks.get('foo')
     assert isinstance(r, Task)
     assert r.name == 'foo'
     assert r.func == print
 
 
+def test_tasks_get_by_function():
+    tasks = Tasks()
+
+    @tasks.task(name='foo')
+    def foo():
+        pass
+
+    r = tasks.get(foo)
+    assert isinstance(r, Task)
+    assert r.name == 'foo'
+    assert r.func == foo
+
+
+def test_tasks_get_by_task_object(task):
+    tasks = Tasks()
+    tasks._tasks[task.name] = task
+
+    r = tasks.get(task)
+    assert isinstance(r, Task)
+    assert r.name == task.name
+    assert r.func == task.func
+
+
+def test_tasks_get_by_unknown_or_wrong_object():
+    tasks = Tasks()
+    with pytest.raises(exc.UnknownTask):
+        tasks.get('foo')
+    with pytest.raises(exc.UnknownTask):
+        tasks.get(None)
+    with pytest.raises(exc.UnknownTask):
+        tasks.get(object())
+    with pytest.raises(exc.UnknownTask):
+        tasks.get(b'foo')
+    with pytest.raises(exc.UnknownTask):
+        tasks.get(RuntimeError)
+
+
 def test_tasks_scheduling(task):
     tasks = Tasks()
     tasks.add(print, 'write_to_stdout', queue='foo_queue')