Skip to content

Commit

Permalink
Add unary inference for integer and boolean literals (#13559)
Browse files Browse the repository at this point in the history
## Summary

Just trying to familiarize myself with the general patterns, testing,
etc.

Part of #12701.
  • Loading branch information
charliermarsh authored Sep 30, 2024
1 parent 5f4b282 commit d86b73e
Showing 1 changed file with 65 additions and 0 deletions.
65 changes: 65 additions & 0 deletions crates/red_knot_python_semantic/src/types/infer.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2209,8 +2209,16 @@ impl<'db> TypeInferenceBuilder<'db> {
} = unary;

match (op, self.infer_expression(operand)) {
(UnaryOp::UAdd, Type::IntLiteral(value)) => Type::IntLiteral(value),
(UnaryOp::USub, Type::IntLiteral(value)) => Type::IntLiteral(-value),
(UnaryOp::Invert, Type::IntLiteral(value)) => Type::IntLiteral(!value),

(UnaryOp::UAdd, Type::BooleanLiteral(bool)) => Type::IntLiteral(i64::from(bool)),
(UnaryOp::USub, Type::BooleanLiteral(bool)) => Type::IntLiteral(-i64::from(bool)),
(UnaryOp::Invert, Type::BooleanLiteral(bool)) => Type::IntLiteral(!i64::from(bool)),

(UnaryOp::Not, ty) => ty.bool(self.db).negate().into_type(self.db),

_ => Type::Unknown, // TODO other unary op types
}
}
Expand Down Expand Up @@ -6571,4 +6579,61 @@ mod tests {
assert_public_ty(&db, "/src/a.py", "c", "bool");
Ok(())
}

#[test]
fn unary_add() -> anyhow::Result<()> {
let mut db = setup_db();

db.write_dedented(
"/src/a.py",
"
a = +0
b = +1
c = +True
",
)?;

assert_public_ty(&db, "/src/a.py", "a", "Literal[0]");
assert_public_ty(&db, "/src/a.py", "b", "Literal[1]");
assert_public_ty(&db, "/src/a.py", "c", "Literal[1]");
Ok(())
}

#[test]
fn unary_sub() -> anyhow::Result<()> {
let mut db = setup_db();

db.write_dedented(
"/src/a.py",
"
a = -0
b = -1
c = -True
",
)?;

assert_public_ty(&db, "/src/a.py", "a", "Literal[0]");
assert_public_ty(&db, "/src/a.py", "b", "Literal[-1]");
assert_public_ty(&db, "/src/a.py", "c", "Literal[-1]");
Ok(())
}

#[test]
fn unary_invert() -> anyhow::Result<()> {
let mut db = setup_db();

db.write_dedented(
"/src/a.py",
"
a = ~0
b = ~1
c = ~True
",
)?;

assert_public_ty(&db, "/src/a.py", "a", "Literal[-1]");
assert_public_ty(&db, "/src/a.py", "b", "Literal[-2]");
assert_public_ty(&db, "/src/a.py", "c", "Literal[-2]");
Ok(())
}
}

0 comments on commit d86b73e

Please sign in to comment.