✨ Add stream output for pruning containers (#566)

python_on_whales/components/container/cli_wrapper.py

@@ -1159,11 +1159,35 @@ def pause(self, containers: Union[ValidContainer, List[ValidContainer]]) -> None
 
         run(full_cmd)
 
-    def prune(self, filters: Dict[str, str] = {}) -> None:
+    @overload
+    def prune(
+        self,
+        filters: Dict[str, str] = {},
+        stream_logs: Literal[True] = ...,
+    ) -> Iterable[Tuple[str, bytes]]:
+        ...
+
+    @overload
+    def prune(
+        self,
+        filters: Dict[str, str] = {},
+        stream_logs: Literal[False] = ...,
+    ) -> None:
+        ...
+
+    def prune(
+        self,
+        filters: Dict[str, str] = {},
+        stream_logs: bool = False,
+    ):
         """Remove containers that are not running.
 
         Parameters:
             filters: Filters as strings or list of strings
+            stream_logs: If `True` this function will return an iterator of strings.
+                You can then read the logs as they arrive. If `False` (the default value), then
+                the function returns `None`, but when it returns, then the prune operation has already been
+                done.
         """
         if isinstance(filter, list):
             raise TypeError(
@@ -1173,6 +1197,8 @@ def prune(self, filters: Dict[str, str] = {}) -> None:
             )
         full_cmd = self.docker_cmd + ["container", "prune", "--force"]
         full_cmd.add_args_list("--filter", format_dict_for_cli(filters))
+        if stream_logs:
+            return stream_stdout_and_stderr(full_cmd)
         run(full_cmd)
 
     def rename(self, container: ValidContainer, new_name: str) -> None:

tests/python_on_whales/components/test_container.py

@@ -924,6 +924,86 @@ def test_prune(ctr_client: DockerClient):
     assert container not in ctr_client.container.list(all=True)
 
 
+@pytest.mark.parametrize("ctr_client", ["docker", "podman"], indirect=True)
+def test_prune_streaming(ctr_client: DockerClient):
+    for container in ctr_client.container.list(filters={"name": "test-container"}):
+        ctr_client.container.remove(container, force=True)
+    container = ctr_client.container.create("busybox")
+    assert container in ctr_client.container.list(all=True)
+
+    # container not pruned because it is not old enough
+    # podman does not provide logs when not pruned
+    logs = list(ctr_client.container.prune(filters={"until": "100h"}, stream_logs=True))
+    assert container in ctr_client.container.list(all=True)
+
+    if ctr_client.client_config.client_type == "docker":
+        assert len(logs) >= 1
+    logs_as_big_binary = b""
+    for log_type, log_value in logs:
+        if ctr_client.client_config.client_type == "docker":
+            assert log_type in ("stdout", "stderr")
+            logs_as_big_binary += log_value
+
+    if ctr_client.client_config.client_type == "docker":
+        assert b"Total reclaimed space:" in logs_as_big_binary
+
+    # container not pruned because it is does not have label "dne"
+    # podman does not provide logs when not pruned
+    logs = list(ctr_client.container.prune(filters={"label": "dne"}, stream_logs=True))
+    assert container in ctr_client.container.list(all=True)
+
+    if ctr_client.client_config.client_type == "docker":
+        assert len(logs) >= 1
+    logs_as_big_binary = b""
+    for log_type, log_value in logs:
+        if ctr_client.client_config.client_type == "docker":
+            assert log_type in ("stdout", "stderr")
+            logs_as_big_binary += log_value
+
+    if ctr_client.client_config.client_type == "docker":
+        assert b"Total reclaimed space:" in logs_as_big_binary
+
+    # container not pruned because it is not old enough and does not have label "dne"
+    # podman does not provide logs when not pruned
+    logs = list(
+        ctr_client.container.prune(
+            filters={"until": "100h", "label": "dne"}, stream_logs=True
+        )
+    )
+    assert container in ctr_client.container.list(all=True)
+
+    if ctr_client.client_config.client_type == "docker":
+        assert len(logs) >= 1
+    logs_as_big_binary = b""
+    for log_type, log_value in logs:
+        if ctr_client.client_config.client_type == "docker":
+            assert log_type in ("stdout", "stderr")
+        logs_as_big_binary += log_value
+
+    if ctr_client.client_config.client_type == "docker":
+        assert b"Total reclaimed space:" in logs_as_big_binary
+
+    # container pruned
+    # podman does provide logs when pruned
+    logs = list(ctr_client.container.prune(stream_logs=True))
+    assert container not in ctr_client.container.list(all=True)
+
+    if ctr_client.client_config.client_type == "docker":
+        len(logs) >= 3
+
+    if ctr_client.client_config.client_type == "podman":
+        len(logs) >= 1
+
+    logs_as_big_binary = b""
+    for log_type, log_value in logs:
+        if ctr_client.client_config.client_type == "docker":
+            assert log_type in ("stdout", "stderr")
+        logs_as_big_binary += log_value
+
+    if ctr_client.client_config.client_type == "docker":
+        assert b"Total reclaimed space:" in logs_as_big_binary
+
+
 @pytest.mark.parametrize("ctr_client", ["docker", "podman"], indirect=True)
 def test_run_detached_interactive(ctr_client: DockerClient):
     with ctr_client.run("ubuntu", interactive=True, detach=True, tty=False) as c: