Add support in native code

python · JukkaL · Apr 1, 2026 · Apr 1, 2026 · Apr 1, 2026 · Apr 1, 2026
commit 8ed64d6dce6d533853ad818e79c1f8b1e94a3a5d
diff --git a/mypyc/irbuild/expression.py b/mypyc/irbuild/expression.py
@@ -364,10 +364,22 @@ def transform_call_expr(builder: IRBuilder, expr: CallExpr) -> Value:
         ):
             item_type = builder.type_to_rtype(analyzed.types[0])
             vec_type = RVec(item_type)
-            if len(expr.args) == 0:
-                return vec_create(builder.builder, vec_type, 0, expr.line)
-            elif len(expr.args) == 1 and expr.arg_kinds == [ARG_POS]:
-                return translate_vec_create_from_iterable(builder, vec_type, expr.args[0])
+            cap = _get_vec_cap(builder, expr)
+            if len(expr.args) == 0 or (
+                len(expr.args) == 1 and expr.arg_kinds == [ARG_NAMED]
+            ):
+                # vec[T]() or vec[T](cap=N)
+                return vec_create(builder.builder, vec_type, 0, expr.line, cap=cap)
+            elif (
+                len(expr.args) == 1
+                and expr.arg_kinds == [ARG_POS]
+                or len(expr.args) == 2
+                and expr.arg_kinds == [ARG_POS, ARG_NAMED]
+            ):
+                # vec[T](items) or vec[T](items, cap=N)
+                return translate_vec_create_from_iterable(
+                    builder, vec_type, expr.args[0], cap=cap
+                )
         callee = analyzed.expr  # Unwrap type application
 
     if isinstance(callee, MemberExpr):
@@ -561,8 +573,16 @@ def translate_super_method_call(builder: IRBuilder, expr: CallExpr, callee: Supe
     return builder.builder.call(decl, arg_values, arg_kinds, arg_names, expr.line)
 
 
+def _get_vec_cap(builder: IRBuilder, expr: CallExpr) -> Value | None:
+    """Extract the 'cap' keyword argument value from a vec() call, or None."""
+    for i, (kind, name) in enumerate(zip(expr.arg_kinds, expr.arg_names)):
+        if kind == ARG_NAMED and name == "cap":
+            return builder.accept(expr.args[i])
+    return None
+
+
 def translate_vec_create_from_iterable(
-    builder: IRBuilder, vec_type: RVec, arg: Expression
+    builder: IRBuilder, vec_type: RVec, arg: Expression, *, cap: Value | None = None
 ) -> Value:
     line = arg.line
     item_type = vec_type.item_type
@@ -581,28 +601,33 @@ def translate_vec_create_from_iterable(
         if is_int64_rprimitive(other_type) or is_int_rprimitive(other_type):
             length = builder.accept(other)
             init = builder.accept(lst.items[0])
-            return vec_create_initialized(builder.builder, vec_type, length, init, line)
+            return vec_create_initialized(builder.builder, vec_type, length, init, line, cap=cap)
         assert False, other_type
     if isinstance(arg, ListExpr):
         items = []
         for item in arg.items:
             value = builder.accept(item)
             items.append(builder.coerce(value, item_type, line))
-        return vec_create_from_values(builder.builder, vec_type, items, line)
+        return vec_create_from_values(builder.builder, vec_type, items, line, cap=cap)
     if isinstance(arg, ListComprehension):
-        return translate_vec_comprehension(builder, vec_type, arg.generator)
-    return vec_from_iterable(builder, vec_type, arg, line)
+        return translate_vec_comprehension(builder, vec_type, arg.generator, cap=cap)
+    return vec_from_iterable(builder, vec_type, arg, line, cap=cap)
 
 
 def vec_from_iterable(
-    builder: IRBuilder, vec_type: RVec, iterable: Expression, line: int
+    builder: IRBuilder,
+    vec_type: RVec,
+    iterable: Expression,
+    line: int,
+    *,
+    cap: Value | None = None,
 ) -> Value:
     """Construct a vec from an arbitrary iterable."""
     # Translate it as a vec comprehension vec[t]([<name> for <name> in
     # iterable]). This way we can use various special casing supported
     # by for loops and comprehensions.
     vec = Register(vec_type)
-    builder.assign(vec, vec_create(builder.builder, vec_type, 0, line), line)
+    builder.assign(vec, vec_create(builder.builder, vec_type, 0, line, cap=cap), line)
     name = f"___tmp_{line}"
     var = Var(name)
     reg = builder.add_local(var, vec_type.item_type)

diff --git a/mypyc/irbuild/for_helpers.py b/mypyc/irbuild/for_helpers.py
@@ -363,7 +363,9 @@ def gen_inner_stmts() -> None:
     return builder.read(set_ops, gen.line)
 
 
-def translate_vec_comprehension(builder: IRBuilder, vec_type: RVec, gen: GeneratorExpr) -> Value:
+def translate_vec_comprehension(
+    builder: IRBuilder, vec_type: RVec, gen: GeneratorExpr, *, cap: Value | None = None
+) -> Value:
     def set_item(x: Value, y: Value, z: Value, line: int) -> None:
         vec_init_item_unsafe(builder.builder, x, y, z, line)
 
@@ -372,15 +374,15 @@ def set_item(x: Value, y: Value, z: Value, line: int) -> None:
         builder,
         gen,
         empty_op_llbuilder=lambda length, line: vec_create(
-            builder.builder, vec_type, length, line
+            builder.builder, vec_type, length, line, cap=cap
         ),
         set_item_op=set_item,
     )
     if val is not None:
         return val
 
     vec = Register(vec_type)
-    builder.assign(vec, vec_create(builder.builder, vec_type, 0, gen.line), gen.line)
+    builder.assign(vec, vec_create(builder.builder, vec_type, 0, gen.line, cap=cap), gen.line)
     loop_params = list(zip(gen.indices, gen.sequences, gen.condlists, gen.is_async))
 
     def gen_inner_stmts() -> None:

diff --git a/mypyc/irbuild/vec.py b/mypyc/irbuild/vec.py
@@ -74,17 +74,28 @@ def as_platform_int(builder: LowLevelIRBuilder, v: Value, line: int) -> Value:
     return builder.coerce(v, c_pyssize_t_rprimitive, line)
 
 
-def vec_create(builder: LowLevelIRBuilder, vtype: RVec, length: int | Value, line: int) -> Value:
+def vec_create(
+    builder: LowLevelIRBuilder,
+    vtype: RVec,
+    length: int | Value,
+    line: int,
+    *,
+    cap: Value | None = None,
+) -> Value:
     if isinstance(length, int):
         length = Integer(length, c_pyssize_t_rprimitive)
     length = as_platform_int(builder, length, line)
+    if cap is not None:
+        cap = as_platform_int(builder, cap, line)
+    else:
+        cap = length
 
     item_type = vtype.item_type
     api_name = vec_api_by_item_type.get(item_type)
     if api_name is not None:
         call = CallC(
             f"{api_name}.alloc",
-            [length, length],
+            [length, cap],
             vtype,
             False,
             False,
@@ -110,7 +121,7 @@ def vec_create(builder: LowLevelIRBuilder, vtype: RVec, length: int | Value, lin
         if depth == 0:
             call = CallC(
                 "VecTApi.alloc",
-                [length, length, typeval],
+                [length, cap, typeval],
                 vtype,
                 False,
                 False,
@@ -121,7 +132,7 @@ def vec_create(builder: LowLevelIRBuilder, vtype: RVec, length: int | Value, lin
         else:
             call = CallC(
                 "VecNestedApi.alloc",
-                [length, length, typeval, Integer(depth, int32_rprimitive)],
+                [length, cap, typeval, Integer(depth, int32_rprimitive)],
                 vtype,
                 False,
                 False,
@@ -134,7 +145,13 @@ def vec_create(builder: LowLevelIRBuilder, vtype: RVec, length: int | Value, lin
 
 
 def vec_create_initialized(
-    builder: LowLevelIRBuilder, vtype: RVec, length: int | Value, init: Value, line: int
+    builder: LowLevelIRBuilder,
+    vtype: RVec,
+    length: int | Value,
+    init: Value,
+    line: int,
+    *,
+    cap: Value | None = None,
 ) -> Value:
     """Create vec with items initialized to the given value."""
     if isinstance(length, int):
@@ -143,7 +160,7 @@ def vec_create_initialized(
 
     item_type = vtype.item_type
     init = builder.coerce(init, item_type, line)
-    vec = vec_create(builder, vtype, length, line)
+    vec = vec_create(builder, vtype, length, line, cap=cap)
 
     items_start = vec_items(builder, vec)
     step = step_size(item_type)
@@ -160,9 +177,14 @@ def vec_create_initialized(
 
 
 def vec_create_from_values(
-    builder: LowLevelIRBuilder, vtype: RVec, values: list[Value], line: int
+    builder: LowLevelIRBuilder,
+    vtype: RVec,
+    values: list[Value],
+    line: int,
+    *,
+    cap: Value | None = None,
 ) -> Value:
-    vec = vec_create(builder, vtype, len(values), line)
+    vec = vec_create(builder, vtype, len(values), line, cap=cap)
     ptr = vec_items(builder, vec)
     item_type = vtype.item_type
     step = step_size(item_type)

diff --git a/mypyc/test-data/irbuild-basic.test b/mypyc/test-data/irbuild-basic.test
@@ -3815,3 +3815,12 @@ def get_length():
 L0:
     r0 = 'https://w3id.org/cwl/salad#'
     return 54
+
+[case testXXX]
+X = tuple[str, ...]
+
+def f() -> tuple[str, ...]:
+    return X(str(x) for x in range(5))
+
+def f2() -> tuple[str, ...]:
+    return tuple(str(x) for x in range(5))
diff --git a/mypyc/test-data/irbuild-vec-i64.test b/mypyc/test-data/irbuild-vec-i64.test
@@ -820,3 +820,56 @@ L1:
     r2 = r1.len
 L2:
     return 1
+
+[case testVecI64CreateWithCap]
+from librt.vecs import vec
+from mypy_extensions import i64
+
+def cap_only() -> vec[i64]:
+    return vec[i64](cap=5)
+
+def cap_zero() -> vec[i64]:
+    return vec[i64](cap=0)
+
+def cap_variable(n: i64) -> vec[i64]:
+    return vec[i64](cap=n)
+[out]
+def cap_only():
+    r0 :: vec[i64]
+L0:
+    r0 = VecI64Api.alloc(0, 5)
+    return r0
+def cap_zero():
+    r0 :: vec[i64]
+L0:
+    r0 = VecI64Api.alloc(0, 0)
+    return r0
+def cap_variable(n):
+    n :: i64
+    r0 :: vec[i64]
+L0:
+    r0 = VecI64Api.alloc(0, n)
+    return r0
+
+[case testVecI64CreateFromListWithCap]
+from librt.vecs import vec
+from mypy_extensions import i64
+
+def list_with_cap() -> vec[i64]:
+    return vec[i64]([1, 2], cap=5)
+[out]
+def list_with_cap():
+    r0 :: vec[i64]
+    r1 :: object
+    r2, r3, r4 :: ptr
+L0:
+    r0 = VecI64Api.alloc(2, 5)
+    r1 = r0.buf
+    r2 = get_element_ptr r1 items :: VecI64BufObject
+    set_mem r2, 1 :: i64*
+    r3 = r2 + 8
+    set_mem r3, 2 :: i64*
+    r4 = r3 + 8
+    keep_alive r0
+    return r0
+
diff --git a/mypyc/test-data/irbuild-vec-nested.test b/mypyc/test-data/irbuild-vec-nested.test
@@ -309,3 +309,41 @@ L5:
     r10 = load_mem r9 :: vec[vec[str]]*
     keep_alive v
     return r10
+
+[case testVecNestedCreateWithCap]
+from librt.vecs import vec
+from mypy_extensions import i64
+
+def cap_str() -> vec[vec[str]]:
+    return vec[vec[str]](cap=5)
+
+def cap_i64() -> vec[vec[i64]]:
+    return vec[vec[i64]](cap=5)
+
+def cap_variable(n: i64) -> vec[vec[str]]:
+    return vec[vec[str]](cap=n)
+[out]
+def cap_str():
+    r0 :: object
+    r1 :: ptr
+    r2 :: vec[vec[str]]
+L0:
+    r0 = load_address PyUnicode_Type
+    r1 = r0
+    r2 = VecNestedApi.alloc(0, 5, r1, 1)
+    return r2
+def cap_i64():
+    r0 :: vec[vec[i64]]
+L0:
+    r0 = VecNestedApi.alloc(0, 5, 2, 1)
+    return r0
+def cap_variable(n):
+    n :: i64
+    r0 :: object
+    r1 :: ptr
+    r2 :: vec[vec[str]]
+L0:
+    r0 = load_address PyUnicode_Type
+    r1 = r0
+    r2 = VecNestedApi.alloc(0, n, r1, 1)
+    return r2
diff --git a/mypyc/test-data/irbuild-vec-t.test b/mypyc/test-data/irbuild-vec-t.test
@@ -359,3 +359,75 @@ def f() -> None:
 main:5: error: Invalid item type for "vec"
 main:6: error: Invalid item type for "vec"
 main:7: error: Invalid item type for "vec"
+
+[case testVecTCreateWithCap]
+from librt.vecs import vec
+from mypy_extensions import i64
+
+def cap_only() -> vec[str]:
+    return vec[str](cap=5)
+
+def cap_zero() -> vec[str]:
+    return vec[str](cap=0)
+
+def cap_variable(n: i64) -> vec[str]:
+    return vec[str](cap=n)
+[out]
+def cap_only():
+    r0 :: object
+    r1 :: ptr
+    r2 :: vec[str]
+L0:
+    r0 = load_address PyUnicode_Type
+    r1 = r0
+    r2 = VecTApi.alloc(0, 5, r1)
+    return r2
+def cap_zero():
+    r0 :: object
+    r1 :: ptr
+    r2 :: vec[str]
+L0:
+    r0 = load_address PyUnicode_Type
+    r1 = r0
+    r2 = VecTApi.alloc(0, 0, r1)
+    return r2
+def cap_variable(n):
+    n :: i64
+    r0 :: object
+    r1 :: ptr
+    r2 :: vec[str]
+L0:
+    r0 = load_address PyUnicode_Type
+    r1 = r0
+    r2 = VecTApi.alloc(0, n, r1)
+    return r2
+
+[case testVecTCreateFromListWithCap]
+from librt.vecs import vec
+from mypy_extensions import i64
+
+def list_with_cap() -> vec[str]:
+    return vec[str](['a', 'b'], cap=5)
+[out]
+def list_with_cap():
+    r0, r1 :: str
+    r2 :: object
+    r3 :: ptr
+    r4 :: vec[str]
+    r5 :: object
+    r6, r7, r8 :: ptr
+L0:
+    r0 = 'a'
+    r1 = 'b'
+    r2 = load_address PyUnicode_Type
+    r3 = r2
+    r4 = VecTApi.alloc(2, 5, r3)
+    r5 = r4.buf
+    r6 = get_element_ptr r5 items :: VecTBufObject
+    set_mem r6, r0 :: builtins.str*
+    r7 = r6 + 8
+    set_mem r7, r1 :: builtins.str*
+    r8 = r7 + 8
+    keep_alive r4
+    return r4
+