feat: use a semaphore to limit how many tasks make requests

1 files changed, 5 insertions, 4 deletions

diff --git a/main.py b/main.py
index 11dae63..e1c6273 100644
--- a/main.py
+++ b/main.py
@@ -33,9 +33,9 @@ def get_urls(filename: str, limit: int | None = None) -> list[str]:
         return urls
 
 
-async def fetch(session: aiohttp.ClientSession, url: str) -> str:
+async def fetch(session: aiohttp.ClientSession, url: str, semaphore: asyncio.Semaphore) -> str:
     try:
-        async with session.get(url) as response:
+        async with semaphore, session.get(url) as response:
             print(".", end="", flush=True)
             if response.ok:
                 return "😁"
@@ -47,7 +47,7 @@ async def fetch(session: aiohttp.ClientSession, url: str) -> str:
         return "😴"
 
 
-async def ping(urls: list[str], max_concurrency=None) -> list[str]:
+async def ping(urls: list[str], max_concurrency: int = 1) -> list[str]:
     """Make a GET request to members of URLS.
 
     If MAX_CONCURRENCY is None, browse every site at once.
@@ -62,7 +62,8 @@ async def ping(urls: list[str], max_concurrency=None) -> list[str]:
         aiohttp.ClientSession(max_field_size=8190 * 2, timeout=aiohttp.ClientTimeout(5)) as session,
         asyncio.TaskGroup() as tg,
     ):
-        tasks = [tg.create_task(fetch(session, url)) for url in urls]
+        semaphore = asyncio.Semaphore(max_concurrency)
+        tasks = [tg.create_task(fetch(session, url, semaphore)) for url in urls]
 
     return [t.result() for t in tasks]