diff --git a/src/restic/docker.py b/src/restic/docker.py index 493a945..ee58a85 100644 --- a/src/restic/docker.py +++ b/src/restic/docker.py @@ -12,31 +12,35 @@ def manage_containers(project: str, services: list[str]): def wrapper(*args, **kwargs): client = docker.from_env() - containers: list[Container] = [ - c - for c in client.containers.list() - if c.labels['com.docker.compose.project'] == project - and c.labels['com.docker.compose.service'] in services - ] try: + project_containers = ( + c + for c in client.containers.list() + if c.labels['com.docker.compose.project'] == project + ) + service_dict = { + c.labels['com.docker.compose.service']: c for c in project_containers + } + containers: list[Container] = [service_dict[s] for s in services] + except Exception as e: + raise e + else: # Start the containers for container in containers: container.stop() logger.info(f'Stopped container [blue]{container.name}[/]') - # Execute the wrapped function - result = func(*args, **kwargs) - - except Exception as e: - # Handle exceptions and ensure containers are stopped - raise e - else: - return result - finally: - # Stop the containers - for container in containers: - container.start() - logger.info(f'Started container [blue]{container.name}[/]') + try: + # Execute the wrapped function + return func(*args, **kwargs) + except Exception as e: + # Handle exceptions and ensure containers are stopped + raise e + finally: + # Stop the containers + for container in containers: + container.start() + logger.info(f'Started container [blue]{container.name}[/]') return wrapper