diff --git a/sync.py b/sync.py index b2f97d8..a5baf34 100644 --- a/sync.py +++ b/sync.py @@ -1,4 +1,6 @@ import asyncio +from types import CoroutineType +from typing import Any, Awaitable, Callable, Coroutine import aiofiles import yaml from pydantic import BaseModel @@ -39,11 +41,12 @@ async def docker_push(image: str) -> int | None: print(f"Pushing image {image}...") return await run_command(f"docker push {image}") -semaphore = asyncio.Semaphore(1) +semaphore = asyncio.Semaphore(5) -async def limited_task(task): +async def limited_task[T: Any](task: Callable[[], Coroutine[None, None, T]]) -> T: async with semaphore: - return await task + return await task() + async def main(): async with aiofiles.open('images.yaml', 'r') as file: @@ -59,19 +62,20 @@ async def main(): for image in config.images: if len(image.tags) > 0: for tag in image.tags: - async def task(): - await docker_pull(f"{image.source}:{tag}") - await docker_tag(f"{image.source}:{tag}", f"{image.target}:{tag}") - await docker_push(f"{image.target}:{tag}") - tasks.append(limited_task(task())) + async def task() -> int | None: + return (await docker_pull(f"{image.source}:{tag}")) \ + or (await docker_pull(f"{image.source}:{tag}")) \ + or (await docker_push(f"{image.target}:{tag}")) + tasks.append(limited_task(task)) else: - async def task(): - await docker_pull(image.source) - await docker_tag(image.source, image.target) - await docker_push(image.target) - tasks.append(limited_task(task())) + async def task() -> int | None: + return (await docker_pull(image.source)) \ + or (await docker_pull(image.source)) \ + or (await docker_push(image.target)) + tasks.append(limited_task(task)) results = await asyncio.gather(*tasks) + print(tasks) failed_tasks = 0 for result in results: if result is not None and result != 0: