Add ability to reset to a lock-file based state (#83)

* Allow resetting from explicit file

* Provide CLI to reset to installer and migration lock file

* Create lock file on conda migrate

* Bump minimum conda version for get_package_records_from_explicit

* Update conda version in pyproject.toml

* Move reset file names to constants

* Increase number of print_explicit calls in migration tests

* Improve reset test typing

* Add test for a file-based reset

* Remove obsolete patches

* Add news

* Remove issue 55 from news

* Update pixi.lock

* Explicitly add conda-forge channel

* Explicitly set channels in test

* Revert "Explicitly set channels in test"

This reverts commit 4bb51da.

* Ensure consistent Python versions

* Remove ambiguity in channel environment variable names

* Auto-reset to file-based state

* Separate news items

* Rename --reset-to to --snapshot

* Use diff_for_unlink_link_precs instead of hand-crafted diffs

* Run macOS intel tests with conda-forge

* Update snapshot help text

* Refresh pixi.lock

* Update conda_self/reset.py

Co-authored-by: jaimergp <jaimergp@users.noreply.github.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Lint

* Apply suggestions from code review

Co-authored-by: Jannis Leidel <jannis@leidel.info>

* Replace state with snapshot

* Use less ambiguous snapshot variable names

---------

Co-authored-by: jaimergp <jaimergp@users.noreply.github.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Jannis Leidel <jannis@leidel.info>

Author: 4 people

.github/workflows/tests.yml

@@ -29,10 +29,10 @@ jobs:
         python-version: ["310", "311", "312", "313"]
         include:
           - os: macos-15-intel
-            channel: defaults
+            channel: conda-forge
             python-version: "310"
           - os: macos-latest
-            channel: conda-forge
+            channel: defaults
             python-version: "313"
 
     steps:
@@ -53,6 +53,8 @@ jobs:
           echo "channels: [${{ matrix.channel }}]" > .pixi/envs/test-py${{ matrix.python-version }}/.condarc
           pixi run --environment test-py${{ matrix.python-version }} python -m conda info
       - name: Run tests
+        env:
+          TEST_CONDA_CHANNEL: ${{ matrix.channel }}
         run: >
           pixi run
           --environment test-py${{ matrix.python-version }}

conda_self/cli/main_migrate.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from pathlib import Path
 from textwrap import dedent
 from typing import TYPE_CHECKING
 
@@ -85,6 +86,7 @@ def execute(args: argparse.Namespace) -> int:
     from conda.misc import clone_env
     from conda.reporters import confirm_yn
 
+    from ..constants import RESET_FILE_MIGRATE
     from ..query import permanent_dependencies
     from ..reset import reset
 
@@ -129,10 +131,10 @@ def execute(args: argparse.Namespace) -> int:
 
     # Take a snapshot of the current base environment by generating the explicit file.
     snapshot_filename = f"explicit.{datetime.now().strftime('%Y-%m-%d-%H-%M-%S')}.txt"
-    snapshot_dest = f"{src_prefix}/conda-meta/{snapshot_filename}"
+    snapshot_dest = Path(src_prefix, "conda-meta", snapshot_filename)
     if not context.quiet:
         print(f"Taking a snapshot of 'base' and saving it to '{snapshot_dest}'...")
-    with open(snapshot_dest, "w") as f:
+    with snapshot_dest.open(mode="w") as f:
         with redirect_stdout(f):
             print_explicit(src_prefix)
 
@@ -147,6 +149,12 @@ def execute(args: argparse.Namespace) -> int:
         print("Resetting 'base' environment...")
     reset(uninstallable_packages=uninstallable_packages)
 
+    # Save the state after migration to allow users to reset to post-migration state
+    migrate_state = Path(src_prefix, "conda-meta", RESET_FILE_MIGRATE)
+    with migrate_state.open(mode="w") as f:
+        with redirect_stdout(f):
+            print_explicit(src_prefix)
+
     # protect the base environment
     try:
         frozen_path = PrefixData(sys.prefix).prefix_path / PREFIX_FROZEN_FILE

conda_self/cli/main_reset.py

@@ -1,12 +1,34 @@
 from __future__ import annotations
 
+import sys
+from pathlib import Path
 from textwrap import dedent
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
     import argparse
+    from typing import TypedDict
+
+    class SnapshotData(TypedDict):
+        file_path: Path
+        snapshot_name: str
+
 
 HELP = "Reset 'base' environment to essential packages only."
+SNAPSHOT_HELP = dedent(
+    """
+    Snapshot to reset the `base` environment to.
+    `snapshot` removes all packages except for `conda`, its plugins,
+    and their dependencies.
+    `installer` resets the `base` environment to the snapshot provided
+    by the installer.
+    `migrate` resets the `base` environment to the snapshot after the last
+    `conda migrate` command run.
+
+    If not set, `conda self` will try to reset to the post-migration snapshot first,
+    then to the installer-provided, and finally to the current snapshot.
+    """
+).lstrip()
 
 WHAT_TO_EXPECT = dedent(
     """
@@ -20,38 +42,83 @@
     Reset the `base` environment to only the essential packages and plugins.
     """
 ).lstrip()
+SUCCESS_SNAPSHOT = dedent(
+    """
+    SUCCESS!
+    Reset the `base` environment to {snapshot_name} snapshot.
+    """
+).lstrip()
 
 
 def configure_parser(parser: argparse.ArgumentParser) -> None:
     from conda.cli.helpers import add_output_and_prompt_options
 
     parser.description = HELP
     add_output_and_prompt_options(parser)
+    parser.add_argument(
+        "--snapshot",
+        choices=("current", "installer", "migrate"),
+        help=SNAPSHOT_HELP,
+    )
     parser.set_defaults(func=execute)
 
 
 def execute(args: argparse.Namespace) -> int:
     from conda.base.context import context
     from conda.reporters import confirm_yn
 
+    from ..constants import RESET_FILE_INSTALLER, RESET_FILE_MIGRATE
     from ..query import permanent_dependencies
     from ..reset import reset
 
     if not context.quiet:
         print(WHAT_TO_EXPECT)
 
-    confirm_yn(
-        "Proceed with resetting your 'base' environment?[y/n]:\n",
-        default="no",
-        dry_run=context.dry_run,
-    )
+    reset_data: dict[str, SnapshotData] = {
+        "installer": {
+            "file_path": Path(sys.prefix, "conda-meta", RESET_FILE_INSTALLER),
+            "snapshot_name": "installer-provided",
+        },
+        "migrate": {
+            "file_path": Path(sys.prefix, "conda-meta", RESET_FILE_MIGRATE),
+            "snapshot_name": "post-migration",
+        },
+    }
+
+    reset_file: Path | None = None
+    snapshot_name = ""
+    if not args.snapshot:
+        for snapshot in ("migrate", "installer"):
+            snapshot_data = reset_data[snapshot]
+            if not snapshot_data["file_path"].exists():
+                continue
+            reset_file = snapshot_data["file_path"]
+            snapshot_name = snapshot_data["snapshot_name"]
+            break
+    elif args.snapshot in reset_data:
+        reset_file = reset_data[args.snapshot]["file_path"]
+        snapshot_name = reset_data[args.snapshot]["snapshot_name"]
+
+    if reset_file and not reset_file.exists():
+        raise FileNotFoundError(
+            f"Failed to reset to `{args.snapshot}`.\n"
+            f"Required file {reset_file} not found."
+        )
+
+    prompt = "Proceed with resetting your 'base' environment"
+    if snapshot_name:
+        prompt += f" to the {snapshot_name} snapshot"
+    confirm_yn(f"{prompt}?[y/n]:\n", default="no", dry_run=context.dry_run)
 
     if not context.quiet:
         print("Resetting 'base' environment...")
-    uninstallable_packages = permanent_dependencies()
-    reset(uninstallable_packages=uninstallable_packages)
+    uninstallable_packages = permanent_dependencies() if not reset_file else set()
+    reset(uninstallable_packages=uninstallable_packages, snapshot=reset_file)
 
     if not context.quiet:
-        print(SUCCESS)
+        if snapshot_name:
+            print(SUCCESS_SNAPSHOT.format(snapshot_name=snapshot_name))
+        else:
+            print(SUCCESS)
 
     return 0

conda_self/constants.py

@@ -1,3 +1,6 @@
 from typing import Final
 
 PERMANENT_PACKAGES: Final = ("conda", "conda-self")
+
+RESET_FILE_INSTALLER = "initial-state.explicit.txt"
+RESET_FILE_MIGRATE = "migrate-state.explicit.txt"

conda_self/reset.py

@@ -1,20 +1,50 @@
+from __future__ import annotations
+
 import sys
+from typing import TYPE_CHECKING
 
+from boltons.setutils import IndexedSet
 from conda.base.context import context
 from conda.core.link import PrefixSetup, UnlinkLinkTransaction
 from conda.core.prefix_data import PrefixData
+from conda.core.solve import diff_for_unlink_link_precs
+from conda.gateways.disk.read import yield_lines
+from conda.misc import get_package_records_from_explicit
+
+if TYPE_CHECKING:
+    from pathlib import Path
 
 
-def reset(prefix: str = sys.prefix, uninstallable_packages: set[str] = set()):
-    installed = sorted(PrefixData(prefix).iter_records(), key=lambda x: x.name)
-    packages_to_remove = [
-        pkg for pkg in installed if pkg.name not in uninstallable_packages
-    ]
+def reset(
+    prefix: str = sys.prefix,
+    uninstallable_packages: set[str] = set(),
+    snapshot: Path | None = None,
+):
+    if snapshot:
+        snapshot_content = list(yield_lines(snapshot))
+        packages_in_reset_env = IndexedSet(
+            get_package_records_from_explicit(snapshot_content)
+        )
+        packages_to_remove, packages_to_install = diff_for_unlink_link_precs(
+            prefix, packages_in_reset_env
+        )
+        if not packages_to_remove and not packages_to_install:
+            print(
+                "Nothing to do. "
+                "Packages in target environment match the selected snapshot."
+            )
+            return
+    else:
+        installed = sorted(PrefixData(prefix).iter_records(), key=lambda x: x.name)
+        packages_to_remove = tuple(
+            pkg for pkg in installed if pkg.name not in uninstallable_packages
+        )
+        packages_to_install = ()
 
     stp = PrefixSetup(
         target_prefix=prefix,
         unlink_precs=packages_to_remove,
-        link_precs=(),
+        link_precs=packages_to_install,
         remove_specs=(),
         update_specs=(),
         neutered_specs=(),

news/83-add-reset

@@ -0,0 +1,20 @@
+### Enhancements
+
+* Add option to reset protected environment to installer-provided or post-migration state. (#54 via #83)
+* Change `conda self reset` to try to reset to the post-migration state or the installer-provided state by default. (#55 via #83)
+
+### Bug fixes
+
+* <news item>
+
+### Deprecations
+
+* <news item>
+
+### Docs
+
+* <news item>
+
+### Other
+
+* <news item>