Add support for literal addition

Author: cdce8p

mypy/checkexpr.py

@@ -3445,6 +3445,9 @@ def visit_op_expr(self, e: OpExpr) -> Type:
                             items=proper_left_type.items + [UnpackType(mapped)]
                         )
 
+        if e.op == "+" and (result := self.literal_expression_addition(e, left_type)):
+            return result
+
         use_reverse: UseReverse = USE_REVERSE_DEFAULT
         if e.op == "|":
             if is_named_instance(proper_left_type, "builtins.dict"):
@@ -3505,6 +3508,57 @@ def visit_op_expr(self, e: OpExpr) -> Type:
         else:
             raise RuntimeError(f"Unknown operator {e.op}")
 
+    def literal_value_from_expr(
+        self, expr: Expression, typ: Type | None = None
+    ) -> tuple[list[str | int], str] | None:
+        if isinstance(expr, StrExpr):
+            return [expr.value], "builtins.str"
+        if isinstance(expr, IntExpr):
+            return [expr.value], "builtins.int"
+        if isinstance(expr, BytesExpr):
+            return [expr.value], "builtins.bytes"
+
+        typ = typ or self.accept(expr)
+        ptype = get_proper_type(typ)
+
+        if isinstance(ptype, LiteralType) and not isinstance(ptype.value, (bool, float)):
+            return [ptype.value], ptype.fallback.type.fullname
+
+        if isinstance(ptype, UnionType):
+            fallback: str | None = None
+            values: list[str | int] = []
+            for item in ptype.items:
+                pitem = get_proper_type(item)
+                if not isinstance(pitem, LiteralType) or isinstance(pitem.value, (float, bool)):
+                    break
+                if fallback is None:
+                    fallback = pitem.fallback.type.fullname
+                if fallback != pitem.fallback.type.fullname:
+                    break
+                values.append(pitem.value)
+            else:
+                assert fallback is not None
+                return values, fallback
+        return None
+
+    def literal_expression_addition(self, e: OpExpr, left_type: Type) -> Type | None:
+        """Check if literal values can be combined with addition."""
+        assert e.op == "+"
+        if not (lvalue := self.literal_value_from_expr(e.left, left_type)):
+            return None
+        if not (rvalue := self.literal_value_from_expr(e.right)) or lvalue[1] != rvalue[1]:
+            return None
+
+        values: list[int | str] = sorted(
+            {
+                val[0] + val[1]  # type: ignore[operator]
+                for val in itertools.product(lvalue[0], rvalue[0])
+            }
+        )
+        if len(values) == 1:
+            return LiteralType(values[0], self.named_type(lvalue[1]))
+        return UnionType([LiteralType(val, self.named_type(lvalue[1])) for val in values])
+
     def visit_comparison_expr(self, e: ComparisonExpr) -> Type:
         """Type check a comparison expression.
 

test-data/unit/check-literal.test

@@ -1416,19 +1416,19 @@ c: Literal[4]
 d: Literal['foo']
 e: str
 
-reveal_type(a + a)      # N: Revealed type is "builtins.int"
+reveal_type(a + a)      # N: Revealed type is "Literal[6]"
 reveal_type(a + b)      # N: Revealed type is "builtins.int"
 reveal_type(b + a)      # N: Revealed type is "builtins.int"
-reveal_type(a + 1)      # N: Revealed type is "builtins.int"
-reveal_type(1 + a)      # N: Revealed type is "builtins.int"
-reveal_type(a + c)      # N: Revealed type is "builtins.int"
-reveal_type(c + a)      # N: Revealed type is "builtins.int"
+reveal_type(a + 1)      # N: Revealed type is "Literal[4]"
+reveal_type(1 + a)      # N: Revealed type is "Literal[4]"
+reveal_type(a + c)      # N: Revealed type is "Literal[7]"
+reveal_type(c + a)      # N: Revealed type is "Literal[7]"
 
-reveal_type(d + d)      # N: Revealed type is "builtins.str"
+reveal_type(d + d)      # N: Revealed type is "Literal['foofoo']"
 reveal_type(d + e)      # N: Revealed type is "builtins.str"
 reveal_type(e + d)      # N: Revealed type is "builtins.str"
-reveal_type(d + 'foo')  # N: Revealed type is "builtins.str"
-reveal_type('foo' + d)  # N: Revealed type is "builtins.str"
+reveal_type(d + 'bar')  # N: Revealed type is "Literal['foobar']"
+reveal_type('bar' + d)  # N: Revealed type is "Literal['barfoo']"
 
 reveal_type(a.__add__(b))  # N: Revealed type is "builtins.int"
 reveal_type(b.__add__(a))  # N: Revealed type is "builtins.int"
@@ -2960,3 +2960,87 @@ class C(B[Literal["word"]]):
 reveal_type(C().collection)  # N: Revealed type is "builtins.list[Literal['word']]"
 reveal_type(C().word)  # N: Revealed type is "Literal['word']"
 [builtins fixtures/tuple.pyi]
+
+[case testLiteralAddition]
+from typing import Union
+from typing_extensions import Literal
+
+str_a: Literal["a"]
+str_b: Literal["b"]
+str_union_1: Literal["a", "b"]
+str_union_2: Literal["c", "d"]
+s: str
+int_1: Literal[1]
+int_2: Literal[2]
+int_union_1: Literal[1, 2]
+int_union_2: Literal[3, 4]
+i: int
+bytes_a: Literal[b"a"]
+bytes_b: Literal[b"b"]
+bytes_union_1: Literal[b"a", b"b"]
+bytes_union_2: Literal[b"c", b"d"]
+b: bytes
+
+misc_union: Literal["a", 1]
+
+reveal_type(str_a + str_b)  # N: Revealed type is "Literal['ab']"
+reveal_type(str_a + "b")    # N: Revealed type is "Literal['ab']"
+reveal_type("a" + str_b)    # N: Revealed type is "Literal['ab']"
+reveal_type(str_union_1 + "b")    # N: Revealed type is "Union[Literal['ab'], Literal['bb']]"
+reveal_type(str_union_1 + str_b)  # N: Revealed type is "Union[Literal['ab'], Literal['bb']]"
+reveal_type("a" + str_union_1)    # N: Revealed type is "Union[Literal['aa'], Literal['ab']]"
+reveal_type(str_a + str_union_1)  # N: Revealed type is "Union[Literal['aa'], Literal['ab']]"
+reveal_type(str_union_1 + str_union_2)  # N: Revealed type is "Union[Literal['ac'], Literal['ad'], Literal['bc'], Literal['bd']]"
+reveal_type(str_a + s)  # N: Revealed type is "builtins.str"
+reveal_type(s + str_a)  # N: Revealed type is "builtins.str"
+reveal_type(str_union_1 + s)  # N: Revealed type is "builtins.str"
+reveal_type(s + str_union_1)  # N: Revealed type is "builtins.str"
+
+reveal_type(int_1 + int_2)  # N: Revealed type is "Literal[3]"
+reveal_type(int_1 + 1)      # N: Revealed type is "Literal[2]"
+reveal_type(1 + int_1)      # N: Revealed type is "Literal[2]"
+reveal_type(int_union_1 + 1)      # N: Revealed type is "Union[Literal[2], Literal[3]]"
+reveal_type(int_union_1 + int_1)  # N: Revealed type is "Union[Literal[2], Literal[3]]"
+reveal_type(1 + int_union_1)      # N: Revealed type is "Union[Literal[2], Literal[3]]"
+reveal_type(int_1 + int_union_1)  # N: Revealed type is "Union[Literal[2], Literal[3]]"
+reveal_type(int_union_1 + int_union_2)  # N: Revealed type is "Union[Literal[4], Literal[5], Literal[6]]"
+reveal_type(int_1 + i)  # N: Revealed type is "builtins.int"
+reveal_type(i + int_1)  # N: Revealed type is "builtins.int"
+reveal_type(int_union_1 + i)  # N: Revealed type is "builtins.int"
+reveal_type(i + int_union_1)  # N: Revealed type is "builtins.int"
+
+reveal_type(bytes_a + bytes_b)  # N: Revealed type is "Literal[b'ab']"
+reveal_type(bytes_a + b"b")    # N: Revealed type is "Literal[b'ab']"
+reveal_type(b"a" + bytes_b)    # N: Revealed type is "Literal[b'ab']"
+reveal_type(bytes_union_1 + b"b")    # N: Revealed type is "Union[Literal[b'ab'], Literal[b'bb']]"
+reveal_type(bytes_union_1 + bytes_b)  # N: Revealed type is "Union[Literal[b'ab'], Literal[b'bb']]"
+reveal_type(b"a" + bytes_union_1)    # N: Revealed type is "Union[Literal[b'aa'], Literal[b'ab']]"
+reveal_type(bytes_a + bytes_union_1)  # N: Revealed type is "Union[Literal[b'aa'], Literal[b'ab']]"
+reveal_type(bytes_union_1 + bytes_union_2)  # N: Revealed type is "Union[Literal[b'ac'], Literal[b'ad'], Literal[b'bc'], Literal[b'bd']]"
+reveal_type(bytes_a + b)  # N: Revealed type is "builtins.bytes"
+reveal_type(b + bytes_a)  # N: Revealed type is "builtins.bytes"
+reveal_type(bytes_union_1 + b)  # N: Revealed type is "builtins.bytes"
+reveal_type(b + bytes_union_1)  # N: Revealed type is "builtins.bytes"
+
+reveal_type(misc_union + "a")  # N: Revealed type is "Union[builtins.str, builtins.int]" \
+    # E: Unsupported operand types for + ("Literal[1]" and "str") \
+    # N: Left operand is of type "Literal['a', 1]"
+reveal_type("a" + misc_union)  # E: Unsupported operand types for + ("str" and "Literal[1]") \
+    # N: Right operand is of type "Literal['a', 1]" \
+    # N: Revealed type is "builtins.str"
+[builtins fixtures/primitives.pyi]
+
+[case testLiteralAdditionTypedDict]
+from typing import TypedDict
+from typing_extensions import Literal
+
+class LookupDict(TypedDict):
+    top_var: str
+    bottom_var: str
+    var: str
+
+def func(d: LookupDict, pos: Literal["top_", "bottom_", ""]) -> str:
+    return d[pos + "var"]
+
+[builtins fixtures/dict.pyi]
+[typing fixtures/typing-typeddict.pyi]

test-data/unit/cmdline.test

@@ -912,8 +912,8 @@ test_between(1	+	1)
 tabs.py:2: error: Incompatible return value type (got "None", expected "str")
             return None
                    ^~~~
-tabs.py:4: error: Argument 1 to "test_between" has incompatible type "int";
-expected "str"
+tabs.py:4: error: Argument 1 to "test_between" has incompatible type
+"Literal[2]"; expected "str"
     test_between(1  +       1)
                  ^~~~~~~~~~~~
 

test-data/unit/fixtures/primitives.pyi

@@ -32,6 +32,7 @@ class str(Sequence[str]):
     def __getitem__(self, item: int) -> str: pass
     def format(self, *args: object, **kwargs: object) -> str: pass
 class bytes(Sequence[int]):
+    def __add__(self, x: bytes) -> bytes: pass
     def __iter__(self) -> Iterator[int]: pass
     def __contains__(self, other: object) -> bool: pass
     def __getitem__(self, item: int) -> int: pass

test-data/unit/typexport-basic.test

@@ -142,7 +142,7 @@ class str: pass
 class list: pass
 class dict: pass
 [out]
-OpExpr(3) : builtins.int
+OpExpr(3) : Literal[3]
 OpExpr(4) : builtins.float
 OpExpr(5) : builtins.float
 OpExpr(6) : builtins.float