From 8360e373254339d65cb703d9f8510d89a7934d76 Mon Sep 17 00:00:00 2001
From: Ibraheem Ahmed <ibraheem@ibraheem.ca>
Date: Wed, 7 Aug 2024 21:51:40 -0400
Subject: [PATCH] simplify marker expressions when deserializing lockfile

---
 crates/uv-resolver/Cargo.toml          |  1 +
 crates/uv-resolver/src/lock.rs         | 48 ++++++++++++++++++++++++++
 crates/uv/src/commands/project/lock.rs |  2 +-
 3 files changed, 50 insertions(+), 1 deletion(-)

diff --git a/crates/uv-resolver/Cargo.toml b/crates/uv-resolver/Cargo.toml
index d67b9d24981db..ff775e2803a4b 100644
--- a/crates/uv-resolver/Cargo.toml
+++ b/crates/uv-resolver/Cargo.toml
@@ -56,6 +56,7 @@ textwrap = { workspace = true }
 thiserror = { workspace = true }
 tokio = { workspace = true }
 tokio-stream = { workspace = true }
+toml = { workspace = true }
 toml_edit = { workspace = true }
 tracing = { workspace = true }
 url = { workspace = true }
diff --git a/crates/uv-resolver/src/lock.rs b/crates/uv-resolver/src/lock.rs
index 463659707b7bb..9cb763d3e1a0f 100644
--- a/crates/uv-resolver/src/lock.rs
+++ b/crates/uv-resolver/src/lock.rs
@@ -10,6 +10,7 @@ use std::sync::Arc;
 use either::Either;
 use itertools::Itertools;
 use petgraph::visit::EdgeRef;
+use pubgrub::Range;
 use rustc_hash::{FxBuildHasher, FxHashMap, FxHashSet};
 use toml_edit::{value, Array, ArrayOfTables, InlineTable, Item, Table, Value};
 use url::Url;
@@ -77,6 +78,53 @@ pub struct Lock {
 }
 
 impl Lock {
+    /// Deserialize the [`Lock`] from a TOML string.
+    pub fn from_toml(s: &str) -> Result<Lock, toml::de::Error> {
+        let mut lock: Lock = toml::from_str(s)?;
+
+        // Simplify all marker expressions based on the requires-python bound.
+        //
+        // This is necessary to ensure the a `Lock` deserialized from a lockfile compares
+        // equally to a newly created `Lock`.
+        // TODO(ibraheem): we should only simplify python versions when serializing or ensure
+        // the requires-python bound is enforced on construction to avoid this step.
+        if let Some(requires_python) = &lock.requires_python {
+            let python_version = Range::from(requires_python.bound_major_minor().clone());
+            let python_full_version = Range::from(requires_python.bound().clone());
+
+            for package in &mut lock.packages {
+                for dep in &mut package.dependencies {
+                    if let Some(marker) = &mut dep.marker {
+                        *marker = marker.clone().simplify_python_versions(
+                            python_version.clone(),
+                            python_full_version.clone(),
+                        );
+                    }
+                }
+
+                for dep in package.optional_dependencies.values_mut().flatten() {
+                    if let Some(marker) = &mut dep.marker {
+                        *marker = marker.clone().simplify_python_versions(
+                            python_version.clone(),
+                            python_full_version.clone(),
+                        );
+                    }
+                }
+
+                for dep in package.dev_dependencies.values_mut().flatten() {
+                    if let Some(marker) = &mut dep.marker {
+                        *marker = marker.clone().simplify_python_versions(
+                            python_version.clone(),
+                            python_full_version.clone(),
+                        );
+                    }
+                }
+            }
+        }
+
+        Ok(lock)
+    }
+
     /// Initialize a [`Lock`] from a [`ResolutionGraph`].
     pub fn from_resolution_graph(graph: &ResolutionGraph) -> Result<Self, LockError> {
         let mut locked_dists = BTreeMap::new();
diff --git a/crates/uv/src/commands/project/lock.rs b/crates/uv/src/commands/project/lock.rs
index 9efcde1809bed..8d00e07a654ff 100644
--- a/crates/uv/src/commands/project/lock.rs
+++ b/crates/uv/src/commands/project/lock.rs
@@ -590,7 +590,7 @@ pub(crate) async fn commit(lock: &Lock, workspace: &Workspace) -> Result<(), Pro
 /// Returns `Ok(None)` if the lockfile does not exist.
 pub(crate) async fn read(workspace: &Workspace) -> Result<Option<Lock>, ProjectError> {
     match fs_err::tokio::read_to_string(&workspace.install_path().join("uv.lock")).await {
-        Ok(encoded) => match toml::from_str::<Lock>(&encoded) {
+        Ok(encoded) => match Lock::from_toml(&encoded) {
             Ok(lock) => Ok(Some(lock)),
             Err(err) => {
                 eprint!("Failed to parse lockfile; ignoring locked requirements: {err}");