[red-knot] condense int literals (#11784)

Display `(Literal[1] | Literal[2])` as `Literal[1, 2]`, and `(Literal[1]
| Literal[2] | OtherType)` as `(Literal[1, 2] | OtherType)`.

Fixes #11782

---------

Co-authored-by: Alex Waygood <Alex.Waygood@Gmail.com>
This commit is contained in:
Carl Meyer 2024-06-06 16:30:40 -06:00 committed by GitHub
parent b2fc0df6db
commit cd101c83ae
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 60 additions and 20 deletions

View File

@ -746,16 +746,42 @@ pub(crate) struct UnionType {
impl UnionType {
fn display(&self, f: &mut std::fmt::Formatter<'_>, store: &TypeStore) -> std::fmt::Result {
f.write_str("(")?;
let (int_literals, other_types): (Vec<Type>, Vec<Type>) = self
.elements
.iter()
.copied()
.partition(|ty| matches!(ty, Type::IntLiteral(_)));
let mut first = true;
for ty in &self.elements {
if !int_literals.is_empty() {
f.write_str("Literal[")?;
let mut nums: Vec<i64> = int_literals
.into_iter()
.filter_map(|ty| {
if let Type::IntLiteral(n) = ty {
Some(n)
} else {
None
}
})
.collect();
nums.sort_unstable();
for num in nums {
if !first {
f.write_str(", ")?;
}
write!(f, "{num}")?;
first = false;
}
f.write_str("]")?;
}
for ty in other_types {
if !first {
f.write_str(" | ")?;
};
first = false;
write!(f, "{}", ty.display(store))?;
}
f.write_str(")")
Ok(())
}
}
@ -775,7 +801,6 @@ pub(crate) struct IntersectionType {
impl IntersectionType {
fn display(&self, f: &mut std::fmt::Formatter<'_>, store: &TypeStore) -> std::fmt::Result {
f.write_str("(")?;
let mut first = true;
for (neg, ty) in self
.positive
@ -792,7 +817,7 @@ impl IntersectionType {
};
write!(f, "{}", ty.display(store))?;
}
f.write_str(")")
Ok(())
}
}
@ -857,7 +882,7 @@ mod tests {
elems.into_iter().collect::<FxIndexSet<_>>()
);
let union = Type::Union(id);
assert_eq!(format!("{}", union.display(&store)), "(C1 | C2)");
assert_eq!(format!("{}", union.display(&store)), "C1 | C2");
}
#[test]
@ -880,9 +905,6 @@ mod tests {
neg.into_iter().collect::<FxIndexSet<_>>()
);
let intersection = Type::Intersection(id);
assert_eq!(
format!("{}", intersection.display(&store)),
"(C1 & C2 & ~C3)"
);
assert_eq!(format!("{}", intersection.display(&store)), "C1 & C2 & ~C3");
}
}

View File

@ -424,7 +424,7 @@ mod tests {
",
)?;
assert_public_type(&case, "a", "x", "(Literal[1] | Literal[2])")
assert_public_type(&case, "a", "x", "Literal[1, 2]")
}
#[test]
@ -450,7 +450,7 @@ mod tests {
",
)?;
assert_public_type(&case, "a", "x", "(Literal[2] | Literal[3])")
assert_public_type(&case, "a", "x", "Literal[2, 3]")
}
#[test]
@ -467,7 +467,7 @@ mod tests {
",
)?;
assert_public_type(&case, "a", "x", "(Unbound | Literal[1])")
assert_public_type(&case, "a", "x", "Literal[1] | Unbound")
}
#[test]
@ -492,7 +492,7 @@ mod tests {
",
)?;
assert_public_type(&case, "a", "x", "(Literal[3] | Literal[4] | Literal[5])")?;
assert_public_type(&case, "a", "x", "Literal[3, 4, 5]")?;
assert_public_type(&case, "a", "r", "Literal[2]")?;
assert_public_type(&case, "a", "s", "Literal[5]")
}
@ -515,7 +515,7 @@ mod tests {
",
)?;
assert_public_type(&case, "a", "x", "(Literal[2] | Literal[3] | Literal[4])")
assert_public_type(&case, "a", "x", "Literal[2, 3, 4]")
}
#[test]
@ -569,7 +569,7 @@ mod tests {
",
)?;
assert_public_type(&case, "a", "x", "(Literal[1] | Literal[2])")
assert_public_type(&case, "a", "x", "Literal[1, 2]")
}
#[test]
@ -587,9 +587,9 @@ mod tests {
",
)?;
assert_public_type(&case, "a", "x", "(Literal[1] | Literal[2])")?;
assert_public_type(&case, "a", "a", "(Literal[1] | Literal[0])")?;
assert_public_type(&case, "a", "b", "(Literal[0] | Literal[2])")
assert_public_type(&case, "a", "x", "Literal[1, 2]")?;
assert_public_type(&case, "a", "a", "Literal[0, 1]")?;
assert_public_type(&case, "a", "b", "Literal[0, 2]")
}
#[test]
@ -606,7 +606,25 @@ mod tests {
",
)?;
assert_public_type(&case, "a", "a", "(Literal[1] | Literal[2])")
assert_public_type(&case, "a", "a", "Literal[1, 2]")
}
#[test]
fn ifexpr_nested() -> anyhow::Result<()> {
let case = create_test()?;
write_to_path(
&case,
"a.py",
"
class C1: pass
class C2: pass
class C3: pass
x = C1 if flag else C2 if flag2 else C3
",
)?;
assert_public_type(&case, "a", "x", "Literal[C1] | Literal[C2] | Literal[C3]")
}
#[test]