Implement parameter to control removal policy of exited containers

Author: qthibeault

src/multicosim/containers.py

@@ -22,6 +22,7 @@
 
 import abc
 import asyncio
+import contextlib
 import errno
 import logging
 import os
@@ -52,6 +53,7 @@
 from .simulations import Simulator as _Simulator
 
 Container: typing_extensions.TypeAlias = docker.models.containers.Container
+Remove = typing.Literal["always", "if_exit_success", "never"]
 
 
 @attrs.define()
@@ -427,10 +429,25 @@ async def send(self, component: ConnectedComponent[MsgT, DataT], msg: MsgT) -> D
         return retval.data
 
     @typing_extensions.override
-    def stop(self) -> None:
+    def stop(self, *, remove: Remove = "if_exit_success"):
         for child in self.children.values():
             child.container.stop()
 
+            match remove:
+                case "always":
+                    child.container.remove()
+                case "never":
+                    pass
+                case _:
+                    status = child.container.wait()
+
+                    # 0 -> ExitSucess, 137 -> SIGKILL, 143 -> SIGTERM
+                    if status["StatusCode"] in {0, 137, 143}:
+                        child.container.remove()
+
+        if remove is not "never":
+            self.context.network.remove()
+
 
 @attrs.define()
 class _Registration:
@@ -445,6 +462,9 @@ def start(self, context: Context) -> _ComponentSimulation:
         return _ComponentSimulation(self.component.start(context), self.dependencies)
 
 
+_SimulationGenerator: typing.TypeAlias = typing.Generator[Simulation, None, None]
+
+
 class Simulator(_Simulator[Context, Simulation]):
     """Simulator implementation using container-based components."""
 
@@ -499,6 +519,17 @@ def start(self) -> Simulation:
 
             raise e
 
+    @contextlib.contextmanager
+    @typing_extensions.override
+    def run(self, *, remove: Remove = "if_exit_success") -> _SimulationGenerator:
+        sim = self.start()
+
+        try:
+            yield sim
+        finally:
+            sim.stop(remove=remove)
+
+
 
 @attrs.frozen()
 class Success(typing.Generic[DataT]):