Skip to content

Commit

Permalink
[red-knot] Literal[True,False] normalized to builtins.bool (#13178)
Browse files Browse the repository at this point in the history
The `UnionBuilder` builds `builtins.bool` when handed `Literal[True]`
and `Literal[False]`.

Caveat: If the builtins module is unfindable somehow, the builder falls
back to the union type of these two literals.

First task from #12694

---------

Co-authored-by: Carl Meyer <[email protected]>
  • Loading branch information
dylwil3 and carljm authored Sep 1, 2024
1 parent d3b6e8f commit 52d8847
Showing 1 changed file with 72 additions and 5 deletions.
77 changes: 72 additions & 5 deletions crates/red_knot_python_semantic/src/types/builder.rs
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,8 @@
use crate::types::{IntersectionType, Type, UnionType};
use crate::{Db, FxOrderSet};

use super::builtins_symbol_ty_by_name;

pub(crate) struct UnionBuilder<'db> {
elements: FxOrderSet<Type<'db>>,
db: &'db dyn Db,
Expand Down Expand Up @@ -56,7 +58,24 @@ impl<'db> UnionBuilder<'db> {
self
}

pub(crate) fn build(self) -> Type<'db> {
/// Performs the following normalizations:
/// - Replaces `Literal[True,False]` with `bool`.
/// - TODO For enums `E` with members `X1`,...,`Xn`, replaces
/// `Literal[E.X1,...,E.Xn]` with `E`.
fn simplify(&mut self) {
if self
.elements
.is_superset(&[Type::BooleanLiteral(true), Type::BooleanLiteral(false)].into())
{
let bool_ty = builtins_symbol_ty_by_name(self.db, "bool");
self.elements.remove(&Type::BooleanLiteral(true));
self.elements.remove(&Type::BooleanLiteral(false));
self.elements.insert(bool_ty);
}
}

pub(crate) fn build(mut self) -> Type<'db> {
self.simplify();
match self.elements.len() {
0 => Type::Never,
1 => self.elements[0],
Expand Down Expand Up @@ -247,17 +266,38 @@ impl<'db> InnerIntersectionBuilder<'db> {
mod tests {
use super::{IntersectionBuilder, IntersectionType, Type, UnionBuilder, UnionType};
use crate::db::tests::TestDb;

fn setup_db() -> TestDb {
TestDb::new()
}
use crate::program::{Program, SearchPathSettings};
use crate::python_version::PythonVersion;
use crate::types::builtins_symbol_ty_by_name;
use crate::ProgramSettings;
use ruff_db::system::{DbWithTestSystem, SystemPathBuf};

impl<'db> UnionType<'db> {
fn elements_vec(self, db: &'db TestDb) -> Vec<Type<'db>> {
self.elements(db).into_iter().collect()
}
}

fn setup_db() -> TestDb {
let db = TestDb::new();

let src_root = SystemPathBuf::from("/src");
db.memory_file_system()
.create_directory_all(&src_root)
.unwrap();

Program::from_settings(
&db,
&ProgramSettings {
target_version: PythonVersion::default(),
search_paths: SearchPathSettings::new(src_root),
},
)
.expect("Valid search path settings");

db
}

#[test]
fn build_union() {
let db = setup_db();
Expand Down Expand Up @@ -296,6 +336,33 @@ mod tests {
assert_eq!(ty, t0);
}

#[test]
fn build_union_bool() {
let db = setup_db();
let bool_ty = builtins_symbol_ty_by_name(&db, "bool");

let t0 = Type::BooleanLiteral(true);
let t1 = Type::BooleanLiteral(true);
let t2 = Type::BooleanLiteral(false);
let t3 = Type::IntLiteral(17);

let Type::Union(union) = UnionBuilder::new(&db).add(t0).add(t1).add(t3).build() else {
panic!("expected a union");
};
assert_eq!(union.elements_vec(&db), &[t0, t3]);
let Type::Union(union) = UnionBuilder::new(&db)
.add(t0)
.add(t1)
.add(t2)
.add(t3)
.build()
else {
panic!("expected a union");
};

assert_eq!(union.elements_vec(&db), &[t3, bool_ty]);
}

#[test]
fn build_union_flatten() {
let db = setup_db();
Expand Down

0 comments on commit 52d8847

Please sign in to comment.