Fix uv sources unwind on dotted/sub-table spellings (#544, #524)

crates/socket-patch-cli/tests/e2e_vendor_pypi_build.rs

@@ -828,6 +828,124 @@ fn uv_vendor_fresh_checkout_frozen_offline_and_revert() {
     );
 }
 
+/// Vendor then revert six on a REAL uv project whose existing sources use
+/// `sources_spelling`; the unwind must be silent, byte-identical, and leave
+/// a pair `uv lock --check` accepts (#544 dotted keys, #524 sub-tables).
+fn uv_sources_spelling_round_trip(tag: &str, sources_spelling: &str) {
+    let Some((uv, python)) = capstone_uv(tag) else {
+        return;
+    };
+    bake_leak_guards();
+    let tmp = tempfile::tempdir().unwrap();
+    let proj = tmp.path().join("proj");
+    std::fs::create_dir_all(&proj).unwrap();
+    let cache = tmp.path().join("uv-cache");
+    let mut cache_env: Vec<(&str, &str)> = vec![("UV_CACHE_DIR", cache.to_str().unwrap())];
+    if let Some(py) = python.as_deref() {
+        cache_env.push(("UV_PYTHON", py));
+    }
+    if !setup_uv_six_project(&uv, &proj, &cache_env, tag) {
+        return;
+    }
+    // Re-lock with a second, user-authored source in the spelling under
+    // test (a direct wheel URL, so no local build backend is needed).
+    let idna = "https://files.pythonhosted.org/packages/e5/3e/741d8c82801c347547f8a2a06aa57dbb1992be9e948df2ea0eda2c8b79e8/idna-3.7-py3-none-any.whl";
+    std::fs::write(
+        proj.join("pyproject.toml"),
+        format!(
+            "[project]\nname = \"vendor-capstone\"\nversion = \"0.1.0\"\nrequires-python = \">=3.9\"\ndependencies = [\"six==1.16.0\", \"idna\"]\n\n{}",
+            sources_spelling.replace("{IDNA}", idna)
+        ),
+    )
+    .unwrap();
+    let relock = tool(&uv, &proj, &["lock", "-q"], &cache_env);
+    if !relock.status.success() {
+        println!(
+            "SKIP e2e_vendor_pypi_build({tag}): `uv lock` with the idna source failed:\n{}",
+            String::from_utf8_lossy(&relock.stderr)
+        );
+        return;
+    }
+    assert_tool_ok(&tool(&uv, &proj, &["sync", "-q"], &cache_env), "uv sync");
+    let installed_six = site_packages(&proj.join(".venv")).join("six.py");
+    stage_patch(&proj, &installed_six);
+    let pyproject_before = std::fs::read(proj.join("pyproject.toml")).unwrap();
+    let uvlock_before = std::fs::read(proj.join("uv.lock")).unwrap();
+
+    let (code, stdout, stderr) = run_vendored(&VendorDriver::VendorOffline, &proj);
+    assert_eq!(
+        code, 0,
+        "vendor failed.\nstdout:\n{stdout}\nstderr:\n{stderr}"
+    );
+    assert_vendored_applied(&parse_envelope(&stdout));
+    let check = tool(&uv, &proj, &["lock", "--check"], &cache_env);
+    assert_tool_ok(&check, "`uv lock --check` on the wired pair");
+
+    let (code, stdout, stderr) = run_socket(
+        &proj,
+        &[
+            "vendor",
+            "--revert",
+            "--json",
+            "--cwd",
+            proj.to_str().unwrap(),
+        ],
+    );
+    assert_eq!(
+        code, 0,
+        "revert failed.\nstdout:\n{stdout}\nstderr:\n{stderr}"
+    );
+    let renv = parse_envelope(&stdout);
+    assert_eq!(renv["status"], "success", "revert envelope: {renv}");
+    assert!(
+        !stdout.contains("vendor_lock_entry_drifted"),
+        "socket-patch's own sources line is not drift: {renv}"
+    );
+    assert_eq!(
+        String::from_utf8(std::fs::read(proj.join("pyproject.toml")).unwrap()).unwrap(),
+        String::from_utf8(pyproject_before).unwrap(),
+        "revert must restore pyproject.toml byte-identical"
+    );
+    assert_eq!(
+        std::fs::read(proj.join("uv.lock")).unwrap(),
+        uvlock_before,
+        "revert must restore uv.lock byte-identical"
+    );
+    let check = tool(&uv, &proj, &["lock", "--check"], &cache_env);
+    assert_tool_ok(&check, "`uv lock --check` after the revert");
+    assert!(
+        !proj.join(".socket/vendor").exists(),
+        ".socket/vendor must be fully removed after revert"
+    );
+}
+
+#[test]
+#[serial_test::serial]
+fn uv_vendor_revert_dotted_sources_key() {
+    uv_sources_spelling_round_trip(
+        "uv-dotted-sources",
+        "[tool.uv]\nsources.idna = { url = \"{IDNA}\" }\n",
+    );
+}
+
+#[test]
+#[serial_test::serial]
+fn uv_vendor_revert_dotted_sources_url_key() {
+    uv_sources_spelling_round_trip(
+        "uv-dotted-sources-url",
+        "[tool.uv]\nsources.idna.url = \"{IDNA}\"\n",
+    );
+}
+
+#[test]
+#[serial_test::serial]
+fn uv_vendor_revert_sub_table_sources() {
+    uv_sources_spelling_round_trip(
+        "uv-sub-table-sources",
+        "[tool.uv.sources.idna]\nurl = \"{IDNA}\"\n",
+    );
+}
+
 /// `get <uuid> --mode vendored` twin of the uv capstone above (v3.6): the
 /// SAME vendor engine and wiring, driven through get's uuid path — exempt
 /// from installed narrowing, so only the mocked `view/{uuid}` route is

crates/socket-patch-core/src/patch/redirect/upstream/uv.rs

@@ -1127,6 +1127,12 @@ fn restore_metadata(meta: &mut Metadata, hit: &Hit, ctx: &Ctx<'_>) -> bool {
             uv.remove("sources");
         }
     }
+    // Adding a key to a header-less parent (one only implied by
+    // `[tool.uv.sources.<pkg>]` sub-tables) printed its header; once just
+    // sub-tables remain, the header is the rewrite's own bytes (#524).
+    if let Some(Item::Table(sources)) = uv.get_mut("sources") {
+        hide_header_over_sub_tables(sources);
+    }
     let mut removed = false;
     if let Some(index) = pushed {
         if let Some(overrides) = uv
@@ -1143,12 +1149,29 @@ fn restore_metadata(meta: &mut Metadata, hit: &Hit, ctx: &Ctx<'_>) -> bool {
     if uv.is_empty() {
         tool.remove("uv");
     }
+    if let Some(Item::Table(uv)) = tool.get_mut("uv") {
+        hide_header_over_sub_tables(uv);
+    }
     if tool.is_empty() {
         meta.doc.remove("tool");
     }
     removed
 }
 
+/// Make a standard table implicit again when it holds only sub-tables, so
+/// it renders as the `[parent.<sub>]` headers alone — the spelling it had
+/// before the hosted rewrite added (and restore removed) a key under it.
+fn hide_header_over_sub_tables(table: &mut toml_edit::Table) {
+    if !table.is_dotted()
+        && !table.is_empty()
+        && table
+            .iter()
+            .all(|(_, item)| matches!(item, Item::Table(t) if !t.is_dotted()))
+    {
+        table.set_implicit(true);
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -1188,6 +1211,88 @@ mod tests {
         assert!(spec_clauses("x @ https://h/x.whl").is_err());
     }
 
+    const HOSTED_SIX: &str = "https://patch.socket.dev/patch/pypi/six/1.16.0/g/e828efa5-5c6d-43f3-9909-03f5ac232b98/six-1.16.0-py2.py3-none-any.whl";
+
+    /// Hosted rewrite of `six` into `original`, then `restore_metadata`:
+    /// the pyproject must come back byte-identically (#524).
+    fn assert_metadata_round_trips(original: &str) {
+        use crate::utils::python_lock::ArtifactSource;
+        let rewritten = crate::utils::python_script::rewrite_project_metadata(
+            original,
+            "six",
+            "1.16.0",
+            ArtifactSource::Url(HOSTED_SIX),
+        )
+        .unwrap()
+        .expect("the rewrite adds a source");
+        assert!(rewritten.contains(HOSTED_SIX), "{rewritten}");
+        let mut meta = Metadata {
+            rel: "pyproject.toml".into(),
+            text: rewritten.clone(),
+            script: false,
+            doc: rewritten.parse().unwrap(),
+        };
+        let hit = Hit {
+            index: 0,
+            uuid: "e828efa5-5c6d-43f3-9909-03f5ac232b98".into(),
+            name: "six".into(),
+            version: "1.16.0".into(),
+        };
+        let client = super::super::UpstreamClient::new(true);
+        let ctx = Ctx {
+            client: &client,
+            origins: &[],
+            bun_lockb: false,
+        };
+        restore_metadata(&mut meta, &hit, &ctx);
+        assert_eq!(
+            meta.render().unwrap(),
+            original,
+            "rewritten was:\n{rewritten}"
+        );
+    }
+
+    const SUB_TABLE_SOURCES: &str = "[tool.uv.sources.idna]\nurl = \"https://files.pythonhosted.org/packages/e5/3e/idna-3.7-py3-none-any.whl\"\n";
+
+    #[test]
+    fn restore_drops_sources_header_made_explicit_over_sub_tables() {
+        assert_metadata_round_trips(&format!(
+            "[project]\nname = \"app\"\nversion = \"0.1.0\"\ndependencies = [\"six==1.16.0\", \"idna==3.7\"]\n\n{SUB_TABLE_SOURCES}"
+        ));
+    }
+
+    #[test]
+    fn restore_drops_headers_made_explicit_over_sub_tables_transitive() {
+        // six is transitive, so the rewrite also adds an override under the
+        // header-less `[tool.uv]` parent.
+        assert_metadata_round_trips(&format!(
+            "[project]\nname = \"app\"\nversion = \"0.1.0\"\ndependencies = [\"idna==3.7\"]\n\n{SUB_TABLE_SOURCES}"
+        ));
+    }
+
+    #[test]
+    fn restore_drops_sources_header_made_explicit_over_sub_tables_crlf() {
+        assert_metadata_round_trips(
+            &format!(
+                "[project]\nname = \"app\"\nversion = \"0.1.0\"\ndependencies = [\"six==1.16.0\", \"idna==3.7\"]\n\n{SUB_TABLE_SOURCES}"
+            )
+            .replace('\n', "\r\n"),
+        );
+    }
+
+    #[test]
+    fn restore_keeps_user_sources_spellings() {
+        let head = "[project]\nname = \"app\"\nversion = \"0.1.0\"\ndependencies = [\"six==1.16.0\", \"idna==3.7\"]\n\n";
+        for sources in [
+            "[tool.uv.sources]\nidna = { url = \"https://h/idna-3.7-py3-none-any.whl\" }\n",
+            "[tool.uv]\nsources.idna = { url = \"https://h/idna-3.7-py3-none-any.whl\" }\n",
+            "[tool.uv]\nsources.idna.url = \"https://h/idna-3.7-py3-none-any.whl\"\n",
+            "[tool.uv]\ndev-dependencies = []\n\n[tool.uv.sources.idna]\nurl = \"https://h/idna-3.7-py3-none-any.whl\"\n",
+        ] {
+            assert_metadata_round_trips(&format!("{head}{sources}"));
+        }
+    }
+
     #[test]
     fn spec_style_joins() {
         let clauses = vec![">=1".to_string(), "<2".to_string()];