Add --restore flag to apply command for clean package restoration

patch_package_py/cli.py

@@ -78,11 +78,18 @@ def cmd_apply(args):
     env_path = (
         Path(args.env_path) if args.env_path is not None else Path.cwd() / ".venv"
     )
-    site_packages_dir = find_site_packages(env_path)
 
     if not patches_dir.exists():
         return
 
+    try:
+        site_packages_dir = find_site_packages(env_path)
+    except FileNotFoundError:
+        logger.error(
+            f"Error: Could not determine site-packages directory for {env_path}",
+        )
+        sys.exit(1)
+
     if not site_packages_dir.exists():
         logger.error(
             f"Error: Site-packages directory {site_packages_dir} does not exist",
@@ -96,7 +103,11 @@ def cmd_apply(args):
         return
 
     for patch_file in patch_files:
-        apply_patch(patch_file, site_packages_dir)
+        apply_patch(
+            patch_file,
+            env_path,
+            restore=args.restore,
+        )
 
 
 def cli():
@@ -135,13 +146,18 @@ def cli():
     commit_parser.add_argument(
         "--skip-restore",
         action="store_true",
-        help="Skip reinstalling the target package before applying the new patch",
+        help="Skip restoring the clean package before applying the new patch",
     )
     commit_parser.set_defaults(func=cmd_commit)
 
     # apply command
     apply_parser = subparsers.add_parser("apply", help="Apply patches")
     apply_parser.add_argument("-e", "--env-path", help="Environment Path")
+    apply_parser.add_argument(
+        "--restore",
+        action="store_true",
+        help="Restore the clean package before applying each patch",
+    )
     apply_parser.set_defaults(func=cmd_apply)
 
     args = parser.parse_args()

patch_package_py/core.py

@@ -292,9 +292,8 @@ def commit_changes(
     if restore_target_package:
         restore_clean_package(package_name, version, target_env_path)
 
-    current_site_packages = find_site_packages(target_env_path)
     try:
-        apply_patch(patch_file_path, current_site_packages)
+        apply_patch(patch_file_path, target_env_path)
     except subprocess.CalledProcessError:
         msg = "Error: failed to apply the patch after creation."
         if restore_target_package:
@@ -313,7 +312,12 @@ def commit_changes(
     logger.info(f"Patch created and applied for {package_name}=={version}")
 
 
-def apply_patch(patch_file: Path, site_packages_dir: Path) -> None:
+def apply_patch(
+    patch_file: Path,
+    env_path: Path,
+    *,
+    restore: bool = False,
+) -> None:
     # Parse package name and version from patch file name
     patch_name = patch_file.stem  # Remove .patch extension
     if "+" not in patch_name:
@@ -324,6 +328,8 @@ def apply_patch(patch_file: Path, site_packages_dir: Path) -> None:
 
     package_name, version = patch_name.rsplit("+", 1)
 
+    site_packages_dir = find_site_packages(env_path)
+
     # Verify the package exists in site_packages_dir with matching version
     resolver = Resolver()
     result = resolver.resolve_in_site_packages(site_packages_dir, package_name)
@@ -339,6 +345,9 @@ def apply_patch(patch_file: Path, site_packages_dir: Path) -> None:
             f"Version mismatch: patch is for {package_name}=={version} but installed version is {installed_version}"
         )
 
+    if restore:
+        restore_clean_package(package_name, version, env_path)
+
     # First, check if the patch is already applied using dry-run
     try:
         subprocess.check_call(
@@ -355,21 +364,32 @@ def apply_patch(patch_file: Path, site_packages_dir: Path) -> None:
             stderr=subprocess.DEVNULL,
             stdout=subprocess.DEVNULL,
         )
-    except subprocess.CalledProcessError:
+    except subprocess.CalledProcessError as e:
+        if restore:
+            raise RuntimeError(
+                f"Failed to apply patch `{patch_name}` after restoring clean package."
+            ) from e
         logger.warning(
             f"Patch `{patch_name}` appears to be already applied, skipping...",
         )
         return
 
     # If dry-run succeeds, apply the patch for real
-    subprocess.check_call(
-        [
-            "patch",
-            "-p1",
-            "-N",
-            "--forward",
-            "-i",
-            str(patch_file.absolute()),
-        ],
-        cwd=site_packages_dir,
-    )
+    try:
+        subprocess.check_call(
+            [
+                "patch",
+                "-p1",
+                "-N",
+                "--forward",
+                "-i",
+                str(patch_file.absolute()),
+            ],
+            cwd=site_packages_dir,
+        )
+    except subprocess.CalledProcessError as e:
+        if restore:
+            raise RuntimeError(
+                f"Failed to apply patch `{patch_name}` after restoring clean package."
+            ) from e
+        raise

skills/patch-package-py/SKILL.md

@@ -16,7 +16,7 @@ The CLI has three commands:
 ```bash
 p12y patch <package> [-e <env-path>] [--amend]
 p12y commit <edit-path> [--skip-restore]
-p12y apply [-e <env-path>]
+p12y apply [-e <env-path>] [--restore]
 ```
 
 ## CLI availability and uv-based workflow
@@ -99,10 +99,9 @@ Use this sequence when guiding a user through patching a package:
    uv run p12y commit <edit-path>
    ```
 
-   `commit` writes the patch file, reinstalls the original package in the
+   `commit` writes the patch file, restores the clean package in the
    environment selected during `patch`, then applies the new patch. Use
-   `--skip-restore` when the target environment is already prepared for direct
-   patch application.
+   `--skip-restore` when the target environment is already in a clean state.
 
 6. Apply all patch files in `patches/` to the selected environment:
 
@@ -116,6 +115,12 @@ Use this sequence when guiding a user through patching a package:
    uv run p12y apply -e <env-path>
    ```
 
+   Use `--restore` to restore each package to its clean state before applying the patch. This is useful when the environment may already contain a previous version of the patch or manual edits:
+
+   ```bash
+   uv run p12y apply --restore
+   ```
+
 7. Run a small import verification check for the patched behavior.
 
 ## Custom environment paths
@@ -138,5 +143,5 @@ Run `uv run p12y commit <edit-path>` from the project root so the patch file lan
 - `Could not determine site-packages directory`: pass a virtual environment path that contains `Lib/site-packages` on Windows or `lib/python*/site-packages` on Unix-like systems.
 - `Version mismatch`: recreate the patch for the installed version, or install the package version named in the patch file.
 - `Invalid patch file name format`: use `patches/<package-name>+<version>.patch`.
-- `appears to be already applied`: the selected environment already contains the patch changes.
+- `appears to be already applied`: the selected environment already contains the patch changes. Run `uv run p12y apply --restore` to restore the clean package first and then re-apply.
 - `No changes detected`: edit files inside the path printed by `uv run p12y patch`, then rerun `uv run p12y commit <edit-path>`.

tests/test_core.py

@@ -125,10 +125,14 @@ def test_resolve_in_site_packages_not_found(self, tmp_path: Path):
 class TestApplyPatch:
     """Integration tests for the apply_patch workflow."""
 
-    def _setup_site_packages(self, tmp_path: Path, package_name: str, version: str):
-        """Helper to create a mock site-packages with a package installed."""
-        site_packages = tmp_path / "site-packages"
-        site_packages.mkdir()
+    def _setup_env(self, tmp_path: Path, package_name: str, version: str):
+        """Helper to create a mock venv with a package installed."""
+        env_path = tmp_path / ".venv"
+        if os.name == "nt":
+            site_packages = env_path / "Lib" / "site-packages"
+        else:
+            site_packages = env_path / "lib" / "python3.9" / "site-packages"
+        site_packages.mkdir(parents=True)
 
         # Create dist-info
         dist_info = (
@@ -146,40 +150,40 @@ def _setup_site_packages(self, tmp_path: Path, package_name: str, version: str):
         (pkg_dir / "__init__.py").write_text('__version__ = "1.0.0"\n')
         (pkg_dir / "core.py").write_text("def hello():\n    return 'hello'\n")
 
-        return site_packages
+        return env_path
 
     def test_apply_patch_invalid_name_format(self, tmp_path: Path, caplog):
         """Test that invalid patch file name is skipped."""
-        site_packages = self._setup_site_packages(tmp_path, "mypackage", "1.0.0")
+        env_path = self._setup_env(tmp_path, "mypackage", "1.0.0")
         patch_file = tmp_path / "invalid_name.patch"
         patch_file.write_text("some patch content")
 
-        apply_patch(patch_file, site_packages)
+        apply_patch(patch_file, env_path)
 
         assert "Invalid patch file name format" in caplog.text
 
     def test_apply_patch_package_not_found(self, tmp_path: Path, caplog):
         """Test that missing package is skipped."""
-        site_packages = self._setup_site_packages(tmp_path, "mypackage", "1.0.0")
+        env_path = self._setup_env(tmp_path, "mypackage", "1.0.0")
         patch_file = tmp_path / "otherpackage+1.0.0.patch"
         patch_file.write_text("some patch content")
 
-        apply_patch(patch_file, site_packages)
+        apply_patch(patch_file, env_path)
 
         assert "not found in site-packages" in caplog.text
 
     def test_apply_patch_version_mismatch(self, tmp_path: Path):
         """Test that version mismatch raises error."""
-        site_packages = self._setup_site_packages(tmp_path, "mypackage", "1.0.0")
+        env_path = self._setup_env(tmp_path, "mypackage", "1.0.0")
         patch_file = tmp_path / "mypackage+2.0.0.patch"
         patch_file.write_text("some patch content")
 
         with pytest.raises(ValueError, match="Version mismatch"):
-            apply_patch(patch_file, site_packages)
+            apply_patch(patch_file, env_path)
 
     def test_apply_patch_success(self, tmp_path: Path):
         """Test successful patch application."""
-        site_packages = self._setup_site_packages(tmp_path, "mypackage", "1.0.0")
+        env_path = self._setup_env(tmp_path, "mypackage", "1.0.0")
         patch_file = tmp_path / "mypackage+1.0.0.patch"
         patch_file.write_text(
             "--- a/mypackage/core.py\n"
@@ -191,25 +195,65 @@ def test_apply_patch_success(self, tmp_path: Path):
         )
 
         with patch("subprocess.check_call") as mock_check_call:
-            apply_patch(patch_file, site_packages)
+            apply_patch(patch_file, env_path)
 
             # Should be called twice: dry-run and actual apply
             assert mock_check_call.call_count == 2
 
     def test_apply_patch_already_applied(self, tmp_path: Path, caplog):
         """Test that already applied patch is skipped."""
-        site_packages = self._setup_site_packages(tmp_path, "mypackage", "1.0.0")
+        env_path = self._setup_env(tmp_path, "mypackage", "1.0.0")
         patch_file = tmp_path / "mypackage+1.0.0.patch"
         patch_file.write_text("some patch content")
 
         with patch("subprocess.check_call") as mock_check_call:
-            # Simulate dry-run failure (patch already applied)
             mock_check_call.side_effect = subprocess.CalledProcessError(1, "patch")
 
-            apply_patch(patch_file, site_packages)
+            apply_patch(patch_file, env_path)
 
             assert "already applied" in caplog.text
 
+    def test_apply_patch_restore_dry_run_failure_raises(self, tmp_path: Path):
+        """Test that a broken patch after restore raises RuntimeError (non-zero exit)."""
+        env_path = self._setup_env(tmp_path, "mypackage", "1.0.0")
+        patch_file = tmp_path / "mypackage+1.0.0.patch"
+        patch_file.write_text("some patch content")
+
+        def side_effect(cmd, *args, **kwargs):
+            if cmd[0] == "uv":
+                return None
+            raise subprocess.CalledProcessError(1, "patch")
+
+        with (
+            patch("subprocess.check_call", side_effect=side_effect),
+            pytest.raises(RuntimeError, match="Failed to apply patch"),
+        ):
+            apply_patch(patch_file, env_path, restore=True)
+
+    def test_apply_patch_restore_real_apply_failure_raises(self, tmp_path: Path):
+        """Test that a real apply failure after successful dry-run raises RuntimeError."""
+        env_path = self._setup_env(tmp_path, "mypackage", "1.0.0")
+        patch_file = tmp_path / "mypackage+1.0.0.patch"
+        patch_file.write_text("some patch content")
+
+        call_count = 0
+
+        def side_effect(cmd, *args, **kwargs):
+            nonlocal call_count
+            call_count += 1
+            if cmd[0] == "uv":
+                return None
+            # First patch call (dry-run) succeeds, second (real apply) fails
+            if call_count == 3:
+                raise subprocess.CalledProcessError(1, "patch")
+            return None
+
+        with (
+            patch("subprocess.check_call", side_effect=side_effect),
+            pytest.raises(RuntimeError, match="Failed to apply patch"),
+        ):
+            apply_patch(patch_file, env_path, restore=True)
+
 
 class TestCommitChanges:
     """Tests for creating patch files via commit_changes."""

tests/test_e2e.py

@@ -6,6 +6,7 @@
 
 from patch_package_py.core import (
     Resolver,
+    apply_patch,
     commit_changes,
     find_site_packages,
     prepare_patch_workspace,
@@ -31,9 +32,7 @@ def project(tmp_path, monkeypatch):
     monkeypatch.chdir(project_dir)
 
     target_env = project_dir / ".venv"
-    subprocess.check_call(
-        ["uv", "venv", str(target_env), "--python", sys.executable]
-    )
+    subprocess.check_call(["uv", "venv", str(target_env), "--python", sys.executable])
     subprocess.check_call(
         [
             "uv",
@@ -84,9 +83,7 @@ def test_new_workspace_has_existing_patch_applied(
         # Workspace 2 with amend: should carry over the patch
         ws2 = tmp_path / "ws2"
         monkeypatch.setattr(tempfile, "mkdtemp", _make_mock_mkdtemp(ws2))
-        prepare_patch_workspace(
-            module_path, PACKAGE, version, target_env, amend=True
-        )
+        prepare_patch_workspace(module_path, PACKAGE, version, target_env, amend=True)
 
         ws2_sp = find_site_packages(ws2 / "venv")
         assert "# patched by e2e test" in (ws2_sp / "six.py").read_text()
@@ -148,9 +145,7 @@ def test_amend_with_bad_patch_recovers_to_clean_state(
 
         ws = tmp_path / "ws"
         monkeypatch.setattr(tempfile, "mkdtemp", _make_mock_mkdtemp(ws))
-        prepare_patch_workspace(
-            module_path, PACKAGE, version, target_env, amend=True
-        )
+        prepare_patch_workspace(module_path, PACKAGE, version, target_env, amend=True)
 
         ws_sp = find_site_packages(ws / "venv")
         git_path = ws_sp.parent
@@ -176,3 +171,33 @@ def test_amend_with_bad_patch_recovers_to_clean_state(
             text=True,
         )
         assert ignored == ""
+
+
+class TestRestoreApplyE2E:
+    def test_restore_reinstalls_then_applies_patch(self, tmp_path, monkeypatch, project):
+        """--restore should restore the package to a clean state, then apply the patch."""
+        module_path = project["module_path"]
+        version = project["version"]
+        target_env = project["target_env"]
+
+        # Create a patch via the normal workflow
+        ws = tmp_path / "ws"
+        monkeypatch.setattr(tempfile, "mkdtemp", _make_mock_mkdtemp(ws))
+        prepare_patch_workspace(module_path, PACKAGE, version, target_env)
+
+        ws_sp = find_site_packages(ws / "venv")
+        six_py = ws_sp / "six.py"
+        original = six_py.read_text()
+        six_py.write_text(original + "\n# restore-e2e-marker\n")
+
+        commit_changes(PACKAGE, version, ws_sp, target_env)
+
+        patch_file = project["dir"] / "patches" / f"{PACKAGE}+{version}.patch"
+        assert patch_file.exists()
+
+        # The patch is now applied in target_env. Apply again with --restore
+        # should succeed (restore cleans the already-patched state).
+        apply_patch(patch_file, target_env, restore=True)
+
+        patched_content = (find_site_packages(target_env) / "six.py").read_text()
+        assert "# restore-e2e-marker" in patched_content