diff --git a/aiojobs/_job.py b/aiojobs/_job.py index 5c3928d..85ccbc6 100644 --- a/aiojobs/_job.py +++ b/aiojobs/_job.py @@ -14,9 +14,15 @@ class Job(Generic[_T]): - def __init__(self, coro: Coroutine[object, object, _T], scheduler: Scheduler): + def __init__( + self, + coro: Coroutine[object, object, _T], + scheduler: Scheduler, + name: Optional[str] = None, + ): self._coro = coro self._scheduler: Optional[Scheduler] = scheduler + self._name = name loop = asyncio.get_running_loop() self._started = loop.create_future() @@ -50,6 +56,21 @@ def pending(self) -> bool: def closed(self) -> bool: return self._closed + def get_name(self) -> Optional[str]: + """Get the task name. + + See https://docs.python.org/3/library/asyncio-task.html#asyncio.Task.get_name. + Returns None if no name was set on the Job object and job has not yet started. + """ + if sys.version_info >= (3, 8) and self._task: + return self._task.get_name() + return self._name + + def set_name(self, name: str) -> None: + self._name = name + if sys.version_info >= (3, 8) and self._task is not None: + self._task.set_name(name) + async def _do_wait(self, timeout: Optional[float]) -> _T: async with async_timeout.timeout(timeout): # TODO: add a test for waiting for a pending coro @@ -118,7 +139,10 @@ async def _close(self, timeout: Optional[float]) -> None: def _start(self) -> None: assert self._task is None - self._task = asyncio.create_task(self._coro) + if sys.version_info >= (3, 8): + self._task = asyncio.create_task(self._coro, name=self._name) + else: + self._task = asyncio.create_task(self._coro) self._task.add_done_callback(self._done_callback) self._started.set_result(None) diff --git a/aiojobs/_scheduler.py b/aiojobs/_scheduler.py index 43b7269..d7d08ce 100644 --- a/aiojobs/_scheduler.py +++ b/aiojobs/_scheduler.py @@ -85,10 +85,12 @@ def pending_count(self) -> int: def closed(self) -> bool: return self._closed - async def spawn(self, coro: Coroutine[object, object, _T]) -> Job[_T]: + async def spawn( + self, coro: Coroutine[object, object, _T], name: Optional[str] = None + ) -> Job[_T]: if self._closed: raise RuntimeError("Scheduling a new job after closing") - job = Job(coro, self) + job = Job(coro, self, name=name) should_start = self._limit is None or self.active_count < self._limit if should_start: job._start() diff --git a/tests/test_job.py b/tests/test_job.py index 73d4a10..7013c7a 100644 --- a/tests/test_job.py +++ b/tests/test_job.py @@ -1,4 +1,5 @@ import asyncio +import sys from contextlib import suppress from typing import Awaitable, Callable, NoReturn from unittest import mock @@ -287,3 +288,39 @@ async def coro() -> NoReturn: await scheduler.spawn(coro()) await scheduler.close() handler.assert_called_once() + + +async def test_get_job_name(scheduler: Scheduler) -> None: + async def coro() -> None: + """Dummy function.""" + + job = await scheduler.spawn(coro(), name="test_job_name") + assert job.get_name() == "test_job_name" + if sys.version_info >= (3, 8): + assert job._task is not None + assert job._task.get_name() == "test_job_name" + + +async def test_get_default_job_name(scheduler: Scheduler) -> None: + async def coro() -> None: + """Dummy function.""" + + job = await scheduler.spawn(coro()) + if sys.version_info >= (3, 8): + job_name = job.get_name() + assert job_name is not None + assert job_name.startswith("Task-") + else: + assert job.get_name() is None + + +async def test_set_job_name(scheduler: Scheduler) -> None: + async def coro() -> None: + """Dummy function.""" + + job = await scheduler.spawn(coro(), name="original_name") + job.set_name("changed_name") + assert job.get_name() == "changed_name" + if sys.version_info >= (3, 8): + assert job._task is not None + assert job._task.get_name() == "changed_name"