Docker Buildx Pruning Logs Streaming (#578)

python_on_whales/components/buildx/cli_wrapper.py

@@ -4,7 +4,18 @@
 import tempfile
 from enum import Enum
 from pathlib import Path
-from typing import Any, Dict, Iterator, List, Literal, Optional, Union
+from typing import (
+    Any,
+    Dict,
+    Iterable,
+    Iterator,
+    List,
+    Literal,
+    Optional,
+    Tuple,
+    Union,
+    overload,
+)
 
 import python_on_whales.components.image.cli_wrapper
 from python_on_whales.client_config import (
@@ -514,16 +525,42 @@ def list(self) -> List[Builder]:
             Builder(self.client_config, x, is_immutable_id=True) for x in builders_names
         ]
 
-    def prune(self, all: bool = False, filters: Dict[str, str] = {}) -> None:
+    @overload
+    def prune(
+        self,
+        all: bool = False,
+        filters: Dict[str, str] = {},
+        stream_logs: Literal[True] = ...,
+    ) -> Iterable[Tuple[str, bytes]]:
+        ...
+
+    @overload
+    def prune(
+        self,
+        all: bool = False,
+        filters: Dict[str, str] = {},
+        stream_logs: Literal[False] = ...,
+    ) -> None:
+        ...
+
+    def prune(
+        self, all: bool = False, filters: Dict[str, str] = {}, stream_logs: bool = False
+    ):
         """Remove build cache on the current builder.
 
         Parameters:
             all: Remove all cache, not just dangling layers
             filters: Filters to use, for example `filters=dict(until="24h")`
+            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.
         """
         full_cmd = self.docker_cmd + ["buildx", "prune", "--force"]
         full_cmd.add_flag("--all", all)
         full_cmd.add_args_list("--filter", format_dict_for_cli(filters))
+        if stream_logs:
+            return stream_buildx_logs(full_cmd)
         run(full_cmd)
 
     def remove(self, builder: Union[Builder, str]) -> None:

tests/python_on_whales/components/buildx/test_buildx_cli_wrapper.py

@@ -46,6 +46,11 @@ def with_container_driver():
     docker.buildx.use(current_builder)
 
 
+@pytest.fixture
+def prune_all():
+    docker.buildx.prune(all=True)
+
+
 @pytest.fixture
 def with_oci_layout_compliant_dir(tmp_path):
     (tmp_path / "Dockerfile").write_text(dockerfile_content1)
@@ -644,8 +649,55 @@ def test_bake_stream_logs(monkeypatch):
     assert output[-1].startswith("#")
 
 
-def test_prune():
-    docker.buildx.prune(filters=dict(until="3m"))
+@pytest.mark.usefixtures("with_docker_driver")
+@pytest.mark.usefixtures("prune_all")
+def test_prune_all_empty():
+    logs = docker.buildx.prune(all=True, stream_logs=True)
+    logs = list(logs)
+    assert len(logs) == 1  # nothing reclaimable
+    assert logs[0] == "Total:\t0B\n"
+
+
+@pytest.mark.usefixtures("with_docker_driver")
+def test_prune(tmp_path):
+    (tmp_path / "Dockerfile").write_text(dockerfile_content1)
+
+    with docker.buildx.create(driver_options=dict(network="host"), use=True):
+        docker.buildx.build(
+            tmp_path, cache_from=dict(type="local", src=tmp_path / "cache")
+        )
+        logs = list(docker.buildx.prune(all=True, stream_logs=True))
+        table_first_row_categories = logs[0]
+        table_total_freed_space = logs[-1]
+        random_row_with_reclaimable = logs[1]
+        assert (
+            table_first_row_categories
+            == "ID\t\t\t\t\t\tRECLAIMABLE\tSIZE\t\tLAST ACCESSED\n"
+        )
+        assert "Total:" in table_total_freed_space
+        assert "true" in random_row_with_reclaimable
+        assert len(logs) >= 2  # table header plus at least one freed file
+
+
+@pytest.mark.usefixtures("with_docker_driver")
+def test_prune_cache_mode_max(tmp_path):
+    (tmp_path / "Dockerfile").write_text(dockerfile_content1)
+
+    with docker.buildx.create(driver_options=dict(network="host"), use=True):
+        docker.buildx.build(
+            tmp_path, cache_to=dict(type="local", dest=tmp_path / "cache", mode="max")
+        )
+        logs = list(docker.buildx.prune(all=True, stream_logs=True))
+        table_first_row_categories = logs[0]
+        table_total_freed_space = logs[-1]
+        random_row_with_reclaimable = logs[1]
+        assert (
+            table_first_row_categories
+            == "ID\t\t\t\t\t\tRECLAIMABLE\tSIZE\t\tLAST ACCESSED\n"
+        )
+        assert "Total:" in table_total_freed_space
+        assert "true" in random_row_with_reclaimable
+        assert len(logs) >= 2  # table header plus at least one freed file
 
 
 def test_list():