[mypyc] feat: extend get_expr_length to work with RTuple and TupleType

Author: BobTheBuidler

mypyc/irbuild/for_helpers.py

@@ -28,6 +28,9 @@
     TypeAlias,
     Var,
 )
+from mypy.types import (
+    TupleType,
+)
 from mypyc.ir.ops import (
     ERR_NEVER,
     BasicBlock,
@@ -1180,18 +1183,18 @@ def gen_cleanup(self) -> None:
             gen.gen_cleanup()
 
 
-def get_expr_length(expr: Expression) -> int | None:
+def get_expr_length(builder: IRBuilder, expr: Expression) -> int | None:
     if isinstance(expr, (StrExpr, BytesExpr)):
         return len(expr.value)
     elif isinstance(expr, (ListExpr, TupleExpr)):
         # if there are no star expressions, or we know the length of them,
         # we know the length of the expression
-        stars = [get_expr_length(i) for i in expr.items if isinstance(i, StarExpr)]
+        stars = [get_expr_length(builder, i) for i in expr.items if isinstance(i, StarExpr)]
         if None not in stars:
             other = sum(not isinstance(i, StarExpr) for i in expr.items)
             return other + sum(stars)  # type: ignore [arg-type]
     elif isinstance(expr, StarExpr):
-        return get_expr_length(expr.expr)
+        return get_expr_length(builder, expr.expr)
     elif (
         isinstance(expr, RefExpr)
         and isinstance(expr.node, Var)
@@ -1204,6 +1207,14 @@ def get_expr_length(expr: Expression) -> int | None:
     # performance boost and can be (sometimes) figured out pretty easily. set and dict
     # comps *can* be done as well but will need special logic to consider the possibility
     # of key conflicts. Range, enumerate, zip are all simple logic.
+
+    # we might still be able to get the length direcly from the type
+    expr_rtype = builder.node_type(expr)
+    if isinstance(expr_rtype, RTuple):
+        return len(expr_rtype.types)
+    expr_type = builder.types[expr]
+    if isinstance(expr_type, TupleType):
+        return len(expr_type.items)
     return None
 
 
@@ -1212,7 +1223,7 @@ def get_expr_length_value(
 ) -> Value:
     rtype = builder.node_type(expr)
     assert is_sequence_rprimitive(rtype), rtype
-    length = get_expr_length(expr)
+    length = get_expr_length(builder, expr)
     if length is None:
         # We cannot compute the length at compile time, so we will fetch it.
         return builder.builder.builtin_len(expr_reg, line, use_pyssize_t=use_pyssize_t)