diff --git a/conformance/third_party/conformance.exp b/conformance/third_party/conformance.exp index 100c600a1..ac7ccf4fa 100644 --- a/conformance/third_party/conformance.exp +++ b/conformance/third_party/conformance.exp @@ -7700,8 +7700,8 @@ { "code": -2, "column": 5, - "concise_description": "Augmented assignment produces a value of type `int`, which is not assignable to `Literal[3, 4, 5]`", - "description": "Augmented assignment produces a value of type `int`, which is not assignable to `Literal[3, 4, 5]`", + "concise_description": "Augmented assignment produces a value of type `Literal[6, 7, 8]`, which is not assignable to `Literal[3, 4, 5]`", + "description": "Augmented assignment produces a value of type `Literal[6, 7, 8]`, which is not assignable to `Literal[3, 4, 5]`", "line": 33, "name": "bad-assignment", "severity": "error", diff --git a/crates/pyrefly_types/src/lit_int.rs b/crates/pyrefly_types/src/lit_int.rs index 22458f5a8..148de193c 100644 --- a/crates/pyrefly_types/src/lit_int.rs +++ b/crates/pyrefly_types/src/lit_int.rs @@ -102,6 +102,23 @@ impl LitInt { } } + fn to_bigint(&self) -> BigInt { + match &self.0 { + LitIntInner::Small(x) => BigInt::from(*x), + LitIntInner::Big(x) => (**x).clone(), + } + } + + pub fn add(&self, other: &Self) -> Self { + match (&self.0, &other.0) { + (LitIntInner::Small(x), LitIntInner::Small(y)) => match x.checked_add(*y) { + Some(sum) => Self::new(sum), + None => Self::new_big(BigInt::from(*x) + BigInt::from(*y)), + }, + _ => Self::new_big(self.to_bigint() + other.to_bigint()), + } + } + pub fn negate(&self) -> Self { match &self.0 { LitIntInner::Small(x) => match x.checked_neg() { diff --git a/pyrefly/lib/alt/operators.rs b/pyrefly/lib/alt/operators.rs index 6663d5ef6..1b1511d18 100644 --- a/pyrefly/lib/alt/operators.rs +++ b/pyrefly/lib/alt/operators.rs @@ -273,6 +273,11 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> { && let Some(r) = self.untype_opt(rhs.clone(), x.right.range()) { Type::type_form(self.union(l, r)) + } else if x.op == Operator::Add + && let Type::Literal(Lit::Int(l)) = lhs + && let Type::Literal(Lit::Int(r)) = rhs + { + Type::Literal(Lit::Int(l.add(r))) } else if x.op == Operator::Add && ((*lhs == Type::LiteralString && rhs.is_literal_string()) || (*rhs == Type::LiteralString && lhs.is_literal_string())) @@ -324,6 +329,11 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> { && rhs.is_literal_string() { Type::LiteralString + } else if x.op == Operator::Add + && let Type::Literal(Lit::Int(l)) = lhs + && let Type::Literal(Lit::Int(r)) = rhs + { + Type::Literal(Lit::Int(l.add(r))) } else if x.op == Operator::Add && let Type::Tuple(ref l) = base && let Type::Tuple(r) = rhs diff --git a/pyrefly/lib/test/lsp/completion.rs b/pyrefly/lib/test/lsp/completion.rs index 320baa600..32d29ef1d 100644 --- a/pyrefly/lib/test/lsp/completion.rs +++ b/pyrefly/lib/test/lsp/completion.rs @@ -1844,7 +1844,7 @@ ff 4 | ff ^ Completion Results: -- (Variable) fff: int +- (Variable) fff: Literal[1, 2] "# .trim(), report.trim(), diff --git a/pyrefly/lib/test/lsp/hover_type.rs b/pyrefly/lib/test/lsp/hover_type.rs index 00b116034..882c78d0d 100644 --- a/pyrefly/lib/test/lsp/hover_type.rs +++ b/pyrefly/lib/test/lsp/hover_type.rs @@ -113,7 +113,7 @@ xy = xy + 1 # main.py 3 | xy = xy + 1 ^ -Hover Result: `int` +Hover Result: `Literal[6]` "# .trim(), report.trim(), diff --git a/pyrefly/lib/test/operators.rs b/pyrefly/lib/test/operators.rs index 15bf37c85..a5405901e 100644 --- a/pyrefly/lib/test/operators.rs +++ b/pyrefly/lib/test/operators.rs @@ -19,6 +19,25 @@ def f(a: int, b: int) -> None: "#, ); +testcase!( + test_literal_int_add, + r#" +from typing import Literal, assert_type + +a: Literal[1] = 1 +b: Literal[2] = 2 +c = a + b +assert_type(c, Literal[3]) +d: Literal[3] = c + +e = a + 5 +assert_type(e, Literal[6]) + +f = 5 + b +assert_type(f, Literal[7]) +"#, +); + testcase!( test_bounded_type_var_comparison, r#"