uv-pep508: add routine for visiting all extras in a marker tree

We'll need this to be able to decode conflict markers back into
a series of inclusion/exclusion rules used during resolution.
This commit is contained in:
Andrew Gallant 2025-02-06 13:47:34 -05:00 committed by Andrew Gallant
parent 1af31d0cc4
commit b48e6b9862
1 changed files with 89 additions and 0 deletions

View File

@ -1162,6 +1162,49 @@ impl MarkerTree {
MarkerTree(INTERNER.lock().only_extras(self.0))
}
/// Calls the provided function on every `extra` in this tree.
///
/// The operator provided to the function is guaranteed to be
/// `MarkerOperator::Equal` or `MarkerOperator::NotEqual`.
pub fn visit_extras(self, mut f: impl FnMut(MarkerOperator, &ExtraName)) {
fn imp(tree: MarkerTree, f: &mut impl FnMut(MarkerOperator, &ExtraName)) {
match tree.kind() {
MarkerTreeKind::True | MarkerTreeKind::False => {}
MarkerTreeKind::Version(kind) => {
for (tree, _) in simplify::collect_edges(kind.edges()) {
imp(tree, f);
}
}
MarkerTreeKind::String(kind) => {
for (tree, _) in simplify::collect_edges(kind.children()) {
imp(tree, f);
}
}
MarkerTreeKind::In(kind) => {
for (_, tree) in kind.children() {
imp(tree, f);
}
}
MarkerTreeKind::Contains(kind) => {
for (_, tree) in kind.children() {
imp(tree, f);
}
}
MarkerTreeKind::Extra(kind) => {
if kind.low.is_false() {
f(MarkerOperator::Equal, kind.name().extra());
} else {
f(MarkerOperator::NotEqual, kind.name().extra());
}
for (_, tree) in kind.children() {
imp(tree, f);
}
}
}
}
imp(self, &mut f);
}
fn simplify_extras_with_impl(self, is_extra: &impl Fn(&ExtraName) -> bool) -> MarkerTree {
MarkerTree(INTERNER.lock().restrict(self.0, &|var| match var {
Variable::Extra(name) => is_extra(name.extra()).then_some(true),
@ -3211,4 +3254,50 @@ mod test {
m("os_name == 'Linux' or os_name != 'Linux'"),
);
}
#[test]
fn visit_extras() {
let collect = |m: MarkerTree| -> Vec<(&'static str, String)> {
let mut collected = vec![];
m.visit_extras(|op, extra| {
let op = match op {
MarkerOperator::Equal => "==",
MarkerOperator::NotEqual => "!=",
_ => unreachable!(),
};
collected.push((op, extra.as_str().to_string()));
});
collected
};
assert_eq!(collect(m("os_name == 'Linux'")), vec![]);
assert_eq!(
collect(m("extra == 'foo'")),
vec![("==", "foo".to_string())]
);
assert_eq!(
collect(m("extra == 'Linux' and extra == 'foo'")),
vec![("==", "foo".to_string()), ("==", "linux".to_string())]
);
assert_eq!(
collect(m("os_name == 'Linux' and extra == 'foo'")),
vec![("==", "foo".to_string())]
);
let marker = "
python_full_version >= '3.12'
and sys_platform == 'darwin'
and extra == 'extra-27-resolution-markers-for-days-cpu'
and extra != 'extra-27-resolution-markers-for-days-cu124'
";
assert_eq!(
collect(m(marker)),
vec![
("==", "extra-27-resolution-markers-for-days-cpu".to_string()),
(
"!=",
"extra-27-resolution-markers-for-days-cu124".to_string()
),
]
);
}
}