diff --git a/guppylang/std/array.py b/guppylang/std/array.py index 14e105d83..62169302b 100644 --- a/guppylang/std/array.py +++ b/guppylang/std/array.py @@ -9,6 +9,7 @@ from __future__ import annotations import builtins +from types import GeneratorType from typing import TYPE_CHECKING, Generic, TypeVar, no_type_check from guppylang.decorator import guppy @@ -71,6 +72,8 @@ def copy(self: array[T, n]) -> array[T, n]: ... def __new__(cls, *args: _T) -> builtins.list[_T]: # type: ignore[no-redef] # Runtime array constructor that is used for comptime. We return an actual list # in line with the comptime unpacking logic that turns arrays into lists. + if len(args) == 1 and isinstance(args[0], GeneratorType): + return list(args[0]) return [*args] diff --git a/tests/integration/tracing/test_array.py b/tests/integration/tracing/test_array.py index 4b6c5691d..809a6f28b 100644 --- a/tests/integration/tracing/test_array.py +++ b/tests/integration/tracing/test_array.py @@ -65,3 +65,12 @@ def main() -> int: validate(compiled) run_int_fn(compiled, 100) + +def test_comprehension(validate): + @guppy.comptime + def main() -> array[int, 5]: + xs = array(i for i in range(5)) + assert xs == [0, 1, 2, 3, 4] + return xs + + validate(guppy.compile(main))