Improve concurrency

dataservice/clients.py

@@ -195,6 +195,8 @@ def _get_context_kwargs(self, request: Request) -> dict[str, Any]:
         context_kwargs = {}
         if request.proxy:
             context_kwargs["proxy"] = {"server": request.proxy.url}
+        if request.headers:
+            context_kwargs["extra_http_headers"] = request.headers
         if self.config is not None and self.config.device:
             context_kwargs.update(self.config.device)
         return context_kwargs

dataservice/worker.py

@@ -351,8 +351,14 @@ async def fetch(self) -> None:
             while self.has_jobs():
                 logger.debug(f"Work queue size: {self._work_queue.qsize()}")
                 logger.debug(f"Data queue size: {self._data_queue.qsize()}")
-                item = self._work_queue.get_nowait()
-                tasks = [task async for task in self._iter_callbacks(item)]
+                items = []
+                for _ in range(
+                    min(self.config.max_concurrency, self._work_queue.qsize())
+                ):
+                    items.append(self._work_queue.get_nowait())
+                tasks = [
+                    task for item in items async for task in self._iter_callbacks(item)
+                ]
                 await asyncio.gather(*tasks)
                 if self.config.cache.use:
                     await cache.write_periodically(self.config.cache.write_interval)

noxfile.py

@@ -8,4 +8,11 @@
 @nox_poetry.session
 def tests(session):
     session.install(".")
-    session.run("pytest", "tests/", "-v", "--cov=dataservice/", external=True)
+    session.run(
+        "pytest",
+        "tests/",
+        "-v",
+        "--cov=dataservice/",
+        "--ignore=tests/integration/test_pw_extra_http_headers.py",
+        external=True,
+    )

pyproject.toml

@@ -3,7 +3,7 @@ name = "python-dataservice"
 packages = [
     { include = "dataservice"},
 ]
-version = "0.11.7"
+version = "0.11.8"
 description = "Lightweight async data gathering for Python"
 authors = ["NomadMonad <[email protected]>"]
 readme = "README.rst"

tests/integration/test_pw_extra_http_headers.py

@@ -0,0 +1,45 @@
+import json
+
+import pytest
+
+from dataservice.clients import PlaywrightClient
+from dataservice.models import Request, Response
+from dataservice.service import DataService
+
+
+@pytest.fixture
+def client():
+    return PlaywrightClient()
+
+
+@pytest.fixture
+def user_agent():
+    return "Mozilla/5.0 (Linux; Android 6.0; Nexus 5 Build/MRA58N) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/130.0.0.0 Mobile Safari/537.36"
+
+
+@pytest.fixture
+def start_requests(client, user_agent):
+    return [
+        Request(
+            url="https://httpbin.org/headers",
+            callback=parse_headers,
+            client=client,
+            headers={"User-Agent": user_agent},
+        )
+    ]
+
+
+@pytest.fixture
+def data_service(start_requests):
+    return DataService(requests=start_requests)
+
+
+def parse_headers(response: Response):
+    headers = json.loads(response.html.find("pre").text)
+    return headers
+
+
+def test_headers(data_service, user_agent):
+    data = tuple(data_service)
+    assert len(data) == 1
+    assert data[0]["headers"]["User-Agent"] == user_agent