Ability to pass request headers in task

hikka-io · May 19, 2024 · e167566 · e167566
1 parent 97c4e11
commit e167566
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 1 deletion.
diff --git a/kibun/__init__.py b/kibun/__init__.py
@@ -7,7 +7,8 @@ async def kibun_loop(tasks, semaphore_number=1):
     semaphore = asyncio.Semaphore(semaphore_number)
     step = 0
 
-    logger.info(f"Got {len(tasks)} tasks")
+    if len(tasks) > 0:
+        logger.info(f"Got {len(tasks)} tasks")
 
     while len(tasks) > 0:
         # Shuffle tasks to make them less predictable

diff --git a/kibun/backend/aio/request.py b/kibun/backend/aio/request.py
@@ -20,6 +20,7 @@ async def request_aio(session, endpoint, task):
             params=task.params,
             method=task.method,
             proxy=task.proxy,
+            headers=task.headers,
             error_markers=task.error_markers,
         )
 

diff --git a/kibun/task.py b/kibun/task.py
@@ -30,6 +30,7 @@ class Task:
     delete_kwargs: dict
     max_fails: int | None
     error_markers: list
+    headers: dict | None
 
 
 @dataclass
@@ -67,6 +68,7 @@ def create_task(
     delete_kwargs={},
     max_fails=10,
     error_markers=constants.DEFAULT_MARKERS,
+    headers=None,
 ):
     if constants.NETWORK_ERROR not in done_status:
         done_status.append(constants.NETWORK_ERROR)
@@ -101,6 +103,7 @@ def create_task(
         delete_kwargs,
         max_fails,
         error_markers,
+        headers,
     )