From fc15a24ab6c77f6dba27a58c51439c77cb3d0366 Mon Sep 17 00:00:00 2001
From: John Lancaster <32917998+jsl12@users.noreply.github.com>
Date: Mon, 27 May 2024 18:33:41 -0500
Subject: [PATCH] reworked service mapping

---
 src/restic/docker.py | 42 +++++++++++++++++++++++-------------------
 1 file changed, 23 insertions(+), 19 deletions(-)

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