feat: implement multi-stage build for better layer caching

Refactor generate_containerfile() to use multi-stage Docker/Podman builds
with conditional stages that only exist when needed:
- Stage 1 (init): Always present - boost marker, upgrade, install deps
- Stage 2 (pre-hooks): Only if pre_init_hooks provided
- Stage 3 (packages): Only if additional_packages provided
- Stage 4 (runner): Only if init_hooks provided

Each stage inherits from the previous existing stage, enabling efficient
layer caching - only changed stages need to rebuild.

Author: xz-dev

src/distrobox_plus/utils/builder.py

@@ -69,7 +69,16 @@ def generate_containerfile(
     init_hooks: str = "",
     pre_init_hooks: str = "",
 ) -> str:
-    """Generate Containerfile content for a boosted image.
+    """Generate multi-stage Containerfile content for a boosted image.
+
+    Creates a multi-stage build with conditional stages:
+    - Stage 1: init - Always present, installs base dependencies
+    - Stage 2: pre-hooks - Only if pre_init_hooks provided
+    - Stage 3: packages - Only if additional_packages provided
+    - Stage 4: runner - Only if init_hooks provided
+
+    Each stage inherits from the previous existing stage, enabling
+    efficient layer caching by Docker/Podman.
 
     Args:
         base_image: Base image to build from
@@ -80,63 +89,60 @@ def generate_containerfile(
     Returns:
         Containerfile content as string
     """
-    lines = [
-        f"FROM {base_image}",
-        "",
-        "# Marker for boosted image - distrobox-init will skip package setup",
-        "RUN touch /.distrobox-boost",
-        "",
-    ]
-
-    # Add upgrade command
-    upgrade_cmd = generate_upgrade_cmd()
+    lines: list[str] = []
+    current_stage = "init"
+
+    # Stage 1: init (always present)
     lines.extend(
         [
+            f"FROM {base_image} AS init",
+            "",
+            "# Marker for boosted image - distrobox-init will skip package setup",
+            "RUN touch /.distrobox-boost",
+            "",
             "# Upgrade existing packages",
-            f"RUN {upgrade_cmd}",
+            f"RUN {generate_upgrade_cmd()}",
+            "",
+            "# Install distrobox dependencies (conditional per-package check)",
+            f"RUN {generate_install_cmd()}",
             "",
         ]
     )
 
-    # Add pre-init hooks if specified
+    # Stage 2: pre-hooks (conditional)
     if pre_init_hooks:
-        hooks_cmd = generate_hooks_cmd(pre_init_hooks)
         lines.extend(
             [
+                f"FROM {current_stage} AS pre-hooks",
+                "",
                 "# Pre-init hooks",
-                f"RUN {hooks_cmd}",
+                f"RUN {generate_hooks_cmd(pre_init_hooks)}",
                 "",
             ]
         )
+        current_stage = "pre-hooks"
 
-    # Add distrobox package installation
-    install_cmd = generate_install_cmd()
-    lines.extend(
-        [
-            "# Install distrobox dependencies (conditional per-package check)",
-            f"RUN {install_cmd}",
-            "",
-        ]
-    )
-
-    # Add additional packages if specified
+    # Stage 3: packages (conditional)
     if additional_packages:
-        additional_cmd = generate_additional_packages_cmd(additional_packages)
         lines.extend(
             [
+                f"FROM {current_stage} AS packages",
+                "",
                 "# Install additional packages",
-                f"RUN {additional_cmd}",
+                f"RUN {generate_additional_packages_cmd(additional_packages)}",
                 "",
             ]
         )
+        current_stage = "packages"
 
-    # Add init hooks if specified
+    # Stage 4: runner (conditional)
     if init_hooks:
-        hooks_cmd = generate_hooks_cmd(init_hooks)
         lines.extend(
             [
+                f"FROM {current_stage} AS runner",
+                "",
                 "# Init hooks",
-                f"RUN {hooks_cmd}",
+                f"RUN {generate_hooks_cmd(init_hooks)}",
                 "",
             ]
         )

tests/unit/test_builder.py

@@ -81,10 +81,10 @@ def test_hash_length(self):
 class TestGenerateContainerfile:
     """Tests for generate_containerfile function."""
 
-    def test_starts_with_from(self):
-        """Test that Containerfile starts with FROM."""
+    def test_starts_with_from_and_init_stage(self):
+        """Test that Containerfile starts with FROM and init stage."""
         result = generate_containerfile("alpine:latest")
-        assert result.startswith("FROM alpine:latest")
+        assert result.startswith("FROM alpine:latest AS init")
 
     def test_contains_boost_marker(self):
         """Test that Containerfile creates boost marker."""
@@ -136,7 +136,7 @@ def test_has_pre_init_hooks_with_arg(self):
         assert "echo hello" in result
 
     def test_order_of_sections(self):
-        """Test that sections appear in correct order."""
+        """Test that sections appear in correct order with multi-stage build."""
         result = generate_containerfile(
             "alpine:latest",
             additional_packages="git",
@@ -146,44 +146,173 @@ def test_order_of_sections(self):
         lines = result.split("\n")
 
         # Find positions
+        init_stage_pos = None
         boost_pos = None
         upgrade_pos = None
-        pre_hook_pos = None
         install_pos = None
-        additional_pos = None
-        init_hook_pos = None
+        pre_hooks_stage_pos = None
+        packages_stage_pos = None
+        runner_stage_pos = None
 
         for i, line in enumerate(lines):
+            if "FROM alpine:latest AS init" in line:
+                init_stage_pos = i
             if "/.distrobox-boost" in line:
                 boost_pos = i
             if "# Upgrade" in line:
                 upgrade_pos = i
-            if "# Pre-init hooks" in line:
-                pre_hook_pos = i
             if "# Install distrobox dependencies" in line:
                 install_pos = i
-            if "# Install additional packages" in line:
-                additional_pos = i
-            if "# Init hooks" in line:
-                init_hook_pos = i
-
-        # Verify order
+            if "FROM init AS pre-hooks" in line:
+                pre_hooks_stage_pos = i
+            if "FROM pre-hooks AS packages" in line:
+                packages_stage_pos = i
+            if "FROM packages AS runner" in line:
+                runner_stage_pos = i
+
+        # Verify all stages present
+        assert init_stage_pos is not None
         assert boost_pos is not None
         assert upgrade_pos is not None
-        assert pre_hook_pos is not None
         assert install_pos is not None
-        assert additional_pos is not None
-        assert init_hook_pos is not None
+        assert pre_hooks_stage_pos is not None
+        assert packages_stage_pos is not None
+        assert runner_stage_pos is not None
 
+        # Verify order within init stage
+        assert init_stage_pos < boost_pos
         assert boost_pos < upgrade_pos
-        assert upgrade_pos < pre_hook_pos
-        assert pre_hook_pos < install_pos
-        assert install_pos < additional_pos
-        assert additional_pos < init_hook_pos
+        assert upgrade_pos < install_pos
+
+        # Verify stage order
+        assert install_pos < pre_hooks_stage_pos
+        assert pre_hooks_stage_pos < packages_stage_pos
+        assert packages_stage_pos < runner_stage_pos
 
     def test_returns_multiline_string(self):
         """Test that result is a valid multiline Containerfile."""
         result = generate_containerfile("alpine:latest")
         lines = result.split("\n")
         assert len(lines) > 5  # Should have multiple lines
         assert lines[0].startswith("FROM")
+
+
+class TestMultiStageBuild:
+    """Tests for multi-stage build structure."""
+
+    def test_multi_stage_all_options(self):
+        """Test that all 4 stages are present when all options provided."""
+        result = generate_containerfile(
+            "alpine:latest",
+            additional_packages="git vim",
+            init_hooks="echo init",
+            pre_init_hooks="echo pre",
+        )
+
+        assert "FROM alpine:latest AS init" in result
+        assert "FROM init AS pre-hooks" in result
+        assert "FROM pre-hooks AS packages" in result
+        assert "FROM packages AS runner" in result
+
+    def test_multi_stage_skip_pre_hooks(self):
+        """Test that pre-hooks stage is skipped when no pre_init_hooks."""
+        result = generate_containerfile(
+            "alpine:latest",
+            additional_packages="git",
+            init_hooks="echo init",
+        )
+
+        assert "FROM alpine:latest AS init" in result
+        assert "AS pre-hooks" not in result
+        # packages should inherit from init directly
+        assert "FROM init AS packages" in result
+        assert "FROM packages AS runner" in result
+
+    def test_multi_stage_skip_packages(self):
+        """Test that packages stage is skipped when no additional_packages."""
+        result = generate_containerfile(
+            "alpine:latest",
+            init_hooks="echo init",
+            pre_init_hooks="echo pre",
+        )
+
+        assert "FROM alpine:latest AS init" in result
+        assert "FROM init AS pre-hooks" in result
+        assert "AS packages" not in result
+        # runner should inherit from pre-hooks directly
+        assert "FROM pre-hooks AS runner" in result
+
+    def test_multi_stage_only_packages(self):
+        """Test with only additional_packages."""
+        result = generate_containerfile(
+            "alpine:latest",
+            additional_packages="git vim",
+        )
+
+        assert "FROM alpine:latest AS init" in result
+        assert "AS pre-hooks" not in result
+        assert "FROM init AS packages" in result
+        assert "AS runner" not in result
+
+    def test_multi_stage_base_only(self):
+        """Test that only init stage exists when no optional args."""
+        result = generate_containerfile("alpine:latest")
+
+        assert "FROM alpine:latest AS init" in result
+        assert "AS pre-hooks" not in result
+        assert "AS packages" not in result
+        assert "AS runner" not in result
+        # Verify base content still present
+        assert "touch /.distrobox-boost" in result
+        assert "# Upgrade existing packages" in result
+        assert "# Install distrobox dependencies" in result
+
+    def test_stage_chain_init_to_runner(self):
+        """Test correct FROM chain: init -> runner (no middle stages)."""
+        result = generate_containerfile(
+            "alpine:latest",
+            init_hooks="echo init",
+        )
+
+        assert "FROM alpine:latest AS init" in result
+        assert "FROM init AS runner" in result
+        assert "AS pre-hooks" not in result
+        assert "AS packages" not in result
+
+    def test_stage_chain_with_pre_hooks_only(self):
+        """Test correct FROM chain: init -> pre-hooks (no packages/runner)."""
+        result = generate_containerfile(
+            "alpine:latest",
+            pre_init_hooks="echo pre",
+        )
+
+        assert "FROM alpine:latest AS init" in result
+        assert "FROM init AS pre-hooks" in result
+        assert "AS packages" not in result
+        assert "AS runner" not in result
+
+    def test_stage_chain_packages_and_runner(self):
+        """Test correct FROM chain when skipping pre-hooks."""
+        result = generate_containerfile(
+            "alpine:latest",
+            additional_packages="git",
+            init_hooks="echo init",
+        )
+
+        # packages should come from init (not pre-hooks)
+        assert "FROM init AS packages" in result
+        # runner should come from packages
+        assert "FROM packages AS runner" in result
+
+    def test_stage_chain_pre_hooks_and_runner(self):
+        """Test correct FROM chain when skipping packages."""
+        result = generate_containerfile(
+            "alpine:latest",
+            init_hooks="echo init",
+            pre_init_hooks="echo pre",
+        )
+
+        # pre-hooks from init
+        assert "FROM init AS pre-hooks" in result
+        # runner should come from pre-hooks (not packages)
+        assert "FROM pre-hooks AS runner" in result