diff --git a/mypyc/doc/list_operations.rst b/mypyc/doc/list_operations.rst index 5993c0a656bd..378568865501 100644 --- a/mypyc/doc/list_operations.rst +++ b/mypyc/doc/list_operations.rst @@ -32,6 +32,7 @@ Operators * ``lst[n]`` (get item by integer index) * ``lst[n:m]``, ``lst[n:]``, ``lst[:m]``, ``lst[:]`` (slicing) +* ``lst1 + lst2``, ``lst += iter`` * ``lst * n``, ``n * lst`` * ``obj in lst`` diff --git a/mypyc/doc/tuple_operations.rst b/mypyc/doc/tuple_operations.rst index fca9e63fc210..ed603fa9982d 100644 --- a/mypyc/doc/tuple_operations.rst +++ b/mypyc/doc/tuple_operations.rst @@ -21,6 +21,7 @@ Operators * ``tup[n]`` (integer index) * ``tup[n:m]``, ``tup[n:]``, ``tup[:m]`` (slicing) +* ``tup1 + tup2`` Statements ---------- diff --git a/mypyc/primitives/list_ops.py b/mypyc/primitives/list_ops.py index a453e568f00f..5cc8b3c0d1c6 100644 --- a/mypyc/primitives/list_ops.py +++ b/mypyc/primitives/list_ops.py @@ -271,6 +271,24 @@ error_kind=ERR_MAGIC, ) +# list + list +binary_op( + name="+", + arg_types=[list_rprimitive, list_rprimitive], + return_type=list_rprimitive, + c_function_name="PySequence_Concat", + error_kind=ERR_MAGIC, +) + +# list += list +binary_op( + name="+=", + arg_types=[list_rprimitive, object_rprimitive], + return_type=list_rprimitive, + c_function_name="PySequence_InPlaceConcat", + error_kind=ERR_MAGIC, +) + # list * int binary_op( name="*", diff --git a/mypyc/primitives/tuple_ops.py b/mypyc/primitives/tuple_ops.py index 0ea0243dc18b..f28d4ca5ec7a 100644 --- a/mypyc/primitives/tuple_ops.py +++ b/mypyc/primitives/tuple_ops.py @@ -15,7 +15,7 @@ object_rprimitive, tuple_rprimitive, ) -from mypyc.primitives.registry import custom_op, function_op, load_address_op, method_op +from mypyc.primitives.registry import binary_op, custom_op, function_op, load_address_op, method_op # Get the 'builtins.tuple' type object. load_address_op(name="builtins.tuple", type=object_rprimitive, src="PyTuple_Type") @@ -74,6 +74,15 @@ error_kind=ERR_MAGIC, ) +# tuple + tuple +binary_op( + name="+", + arg_types=[tuple_rprimitive, tuple_rprimitive], + return_type=tuple_rprimitive, + c_function_name="PySequence_Concat", + error_kind=ERR_MAGIC, +) + # tuple[begin:end] tuple_slice_op = custom_op( arg_types=[tuple_rprimitive, int_rprimitive, int_rprimitive], diff --git a/mypyc/test-data/fixtures/ir.py b/mypyc/test-data/fixtures/ir.py index b908b4c3fc1f..8a373a037e04 100644 --- a/mypyc/test-data/fixtures/ir.py +++ b/mypyc/test-data/fixtures/ir.py @@ -207,6 +207,10 @@ def __getitem__(self, i: slice) -> Tuple[T_co, ...]: pass def __len__(self) -> int: pass def __iter__(self) -> Iterator[T_co]: ... def __contains__(self, item: object) -> int: ... + @overload + def __add__(self, value: Tuple[T_co, ...], /) -> Tuple[T_co, ...]: ... + @overload + def __add__(self, value: Tuple[_T, ...], /) -> Tuple[T_co | _T, ...]: ... class function: pass @@ -223,7 +227,11 @@ def __rmul__(self, i: int) -> List[_T]: pass def __iter__(self) -> Iterator[_T]: pass def __len__(self) -> int: pass def __contains__(self, item: object) -> int: ... - def __add__(self, x: List[_T]) -> List[_T]: ... + @overload + def __add__(self, value: List[_T], /) -> List[_T]: ... + @overload + def __add__(self, value: List[_S], /) -> List[_S | _T]: ... + def __iadd__(self, value: Iterable[_T], /) -> List[_T]: ... # type: ignore[misc] def append(self, x: _T) -> None: pass def pop(self, i: int = -1) -> _T: pass def count(self, _T) -> int: pass diff --git a/mypyc/test-data/irbuild-lists.test b/mypyc/test-data/irbuild-lists.test index e2c656399821..b7ba1a783bb7 100644 --- a/mypyc/test-data/irbuild-lists.test +++ b/mypyc/test-data/irbuild-lists.test @@ -145,6 +145,32 @@ L0: x = r10 return 1 +[case testListAdd] +from typing import List +def f(a: List[int], b: List[int]) -> None: + c = a + b +[out] +def f(a, b): + a, b, r0, c :: list +L0: + r0 = PySequence_Concat(a, b) + c = r0 + return 1 + +[case testListIAdd] +from typing import List, Any +def f(a: List[int], b: Any) -> None: + a += b +[out] +def f(a, b): + a :: list + b :: object + r0 :: list +L0: + r0 = PySequence_InPlaceConcat(a, b) + a = r0 + return 1 + [case testListMultiply] from typing import List def f(a: List[int]) -> None: diff --git a/mypyc/test-data/irbuild-tuple.test b/mypyc/test-data/irbuild-tuple.test index abb180dde89b..e7280bb3b552 100644 --- a/mypyc/test-data/irbuild-tuple.test +++ b/mypyc/test-data/irbuild-tuple.test @@ -384,3 +384,37 @@ L3: L4: a = r1 return 1 + +[case testTupleAdd] +from typing import Tuple +def f(a: Tuple[int, ...], b: Tuple[int, ...]) -> None: + c = a + b + d = a + (1, 2) +def g(a: Tuple[int, int], b: Tuple[int, int]) -> None: + c = a + b +[out] +def f(a, b): + a, b, r0, c :: tuple + r1 :: tuple[int, int] + r2 :: object + r3, d :: tuple +L0: + r0 = PySequence_Concat(a, b) + c = r0 + r1 = (2, 4) + r2 = box(tuple[int, int], r1) + r3 = PySequence_Concat(a, r2) + d = r3 + return 1 +def g(a, b): + a, b :: tuple[int, int] + r0, r1 :: object + r2 :: tuple + r3, c :: tuple[int, int, int, int] +L0: + r0 = box(tuple[int, int], a) + r1 = box(tuple[int, int], b) + r2 = PySequence_Concat(r0, r1) + r3 = unbox(tuple[int, int, int, int], r2) + c = r3 + return 1 diff --git a/mypyc/test-data/run-lists.test b/mypyc/test-data/run-lists.test index 3b2721093e0f..84168f7254f5 100644 --- a/mypyc/test-data/run-lists.test +++ b/mypyc/test-data/run-lists.test @@ -267,6 +267,9 @@ print(g()) 7 [case testListOps] +from typing import Any, cast +from testutil import assertRaises + def test_slicing() -> None: # Use dummy adds to avoid constant folding zero = int() @@ -289,6 +292,27 @@ def test_slicing() -> None: assert s[long_int:] == [] assert s[-long_int:-1] == ["f", "o", "o", "b", "a"] +def in_place_add(l2: Any) -> list[Any]: + l1 = [1, 2] + l1 += l2 + return l1 + +def test_add() -> None: + res = [1, 2, 3, 4] + assert [1, 2] + [3, 4] == res + with assertRaises(TypeError, 'can only concatenate list (not "tuple") to list'): + assert [1, 2] + cast(Any, (3, 4)) == res + l1 = [1, 2] + id_l1 = id(l1) + l1 += [3, 4] + assert l1 == res + assert id_l1 == id(l1) + assert in_place_add([3, 4]) == res + assert in_place_add((3, 4)) == res + assert in_place_add({3, 4}) == res + assert in_place_add({3: "", 4: ""}) == res + assert in_place_add(range(3, 5)) == res + [case testOperatorInExpression] def tuple_in_int0(i: int) -> bool: diff --git a/mypyc/test-data/run-tuples.test b/mypyc/test-data/run-tuples.test index 1f1b0bc9eae7..afd3a956b871 100644 --- a/mypyc/test-data/run-tuples.test +++ b/mypyc/test-data/run-tuples.test @@ -146,7 +146,8 @@ assert Record.__annotations__ == { }, Record.__annotations__ [case testTupleOps] -from typing import Tuple, Final, List, Any, Optional +from typing import Tuple, Final, List, Any, Optional, cast +from testutil import assertRaises def f() -> Tuple[()]: return () @@ -254,3 +255,9 @@ TUPLE: Final[Tuple[str, ...]] = ('x', 'y') def test_final_boxed_tuple() -> None: t = TUPLE assert t == ('x', 'y') + +def test_add() -> None: + res = (1, 2, 3, 4) + assert (1, 2) + (3, 4) == res + with assertRaises(TypeError, 'can only concatenate tuple (not "list") to tuple'): + assert (1, 2) + cast(Any, [3, 4]) == res