Unable to convert generator returning 1D arrays into a 2D array using fromiter/vstack about numba HOT 12 OPEN

$ python listtype.py 
True

TypeError: arrays to stack must be passed as a "sequence" type such as list or tuple.

from collections.abc import Sequence, Iterable
import numpy as np

def gen():
    for i in range(3):
        yield np.arange(i, i+3)

iterable = gen()
isinstance(iterable, Sequence)
# False
isinstance(iterable, Iterable)
# True

import numpy as np

def gen():
    """Generator function returning 1D-arrays."""
    for i in range(3):
        yield np.arange(i, i+3)

np.vstack(gen())
# TypeError: arrays to stack must be passed as a "sequence" type such as list or tuple.

@overload(np.vstack)
def impl_np_vstack(tup):
    if isinstance(tup, types.BaseTuple):
        def impl(tup):
            return _np_vstack(tup)
        return impl

import numpy as np
import numba as nb

@nb.njit
def gen():
    """Generator function returning 1D-arrays."""
    for i in range(3):
        yield np.arange(i, i+3)

@nb.njit
def gen_as_list():
    """Convert generator to list of 1D-arrays."""
    return nb.typed.List(gen())

@nb.njit
def gen_as_array_nb():
    """Convert generator to 2D array."""
    return np.asarray(gen_as_list())

def gen_as_array_np():
    """Convert generator to 2D array."""
    return np.asarray(gen_as_list())

print(gen_as_array_np())
# [[0 1 2]
#  [1 2 3]
#  [2 3 4]]

print(gen_as_array_nb())
# TypingError: No implementation of function Function(<built-in function asarray>) found for signature:
# asarray(ListType[array(int64, 1d, C)])

import numpy as np
import numba as nb
from numba import types
from numba.np.numpy_support import type_can_asarray

@nb.njit
def gen():
    """Generator function returning 1D-arrays."""
    for i in range(3):
        yield np.arange(i, i+3)

@nb.njit
def gen_as_list():
    """Convert generator to list of 1D-arrays."""
    return nb.typed.List(gen())

def gen_type_check():
    a = gen_as_list()
    # type checking as per @overload(np.asarray) in arraymath.py from line 4291
    if not type_can_asarray(a):
        return 1
    if isinstance(a, types.Array):
        return 2
    elif isinstance(a, (types.Sequence, types.Tuple)):
        return 3
    elif isinstance(a, (types.Number, types.Boolean)):
        return 4
    elif isinstance(a, types.containers.ListType):
        return 5
    elif isinstance(a, types.StringLiteral):
        return 6
    else:
        return 7

print(gen_type_check())
# 1

# def type_can_asarray(arr):
#     ok = (types.Array, types.Sequence, types.Tuple, types.StringLiteral,
#           types.Number, types.Boolean, types.containers.ListType)
#     return isinstance(arr, ok)

print(type_can_asarray(gen_as_list()))
# False

print(isinstance(gen_as_list(), types.containers.ListType))
# False

@nb.njit
def np_fromiter_impl(iter, dtype):
    arraylist = nb.typed.List(iter)
    size = len(arraylist)
    out = np.empty(size, dtype=dtype)
    for i in range(size):
        out[i] = arraylist[i]
    return out

dtype = np.dtype(('<i8', (3,)))

print(np_fromiter_impl.py_func(gen(), dtype=dtype))
# [[0 1 2]
#  [1 2 3]
#  [2 3 4]]

print(np_fromiter_impl(gen(), dtype=dtype))
# TypingError: non-precise type pyobject
# During: typing of argument at /tmp/ipykernel_1004317/1459046023.py (1)

@nb.njit
def make_dtype():
    return np.dtype(('<i8', (3,)))

print(make_dtype.py_func())
# ('<i8', (3,))

print(make_dtype())
# TypingError: No implementation of function Function(<class 'numpy.dtype'>) found for signature:
# dtype(Tuple(Literal[str](<i8), UniTuple(int64 x 1)))

print(isinstance(gen_as_list(), types.containers.ListType))
# False

import numpy as np
from numba import njit, typed

@njit
def give_list():
    return typed.List([1, 2, 3])

@njit
def give_arraylist():
    return typed.List([np.array([1, 2, 3]), np.array([1, 2, 3])])

@njit
def asarray(a):
    return np.asarray(a)

print(asarray(give_list()))
# [1 2 3]

print(asarray(give_arraylist()))
# TypingError: No implementation of function Function(<built-in function asarray>) found for signature:
# asarray(ListType[array(int64, 1d, C)])

import numpy as np
from numba import njit, types
from numba.extending import overload
from numba.core.errors import TypingError

@overload(np.fromiter)
def ovl_fromiter(iter, dtype):
    def np_fromiter_impl(iter, dtype):
        return np.asarray(list(iter), dtype)
    # Type check
    if not isinstance(iter, types.IterableType):
        raise TypingError("First argument must be an iterable.")
    else:
        return np_fromiter_impl

@njit
def gen_scalars():
    """Generator function returning scalars."""
    for i in range(3):
        yield i

@njit
def use_fromiter_scalars(dtype):
    return np.fromiter(gen_scalars(), dtype)

print(use_fromiter_scalars(np.int64))
# [0 1 2]

@overload(np.fromiter)
def ovl_fromiter(iter, dtype):
    def np_fromiter_impl(iter, dtype):
        return np.asarray(typed.List(iter), dtype)
    # Type check
    if not isinstance(iter, types.IterableType):
        raise TypingError("First argument must be an iterable.")
    else:
        return np_fromiter_impl

import numpy as np
from numba import njit, types
from numba.extending import overload
from numba.core.errors import TypingError

@overload(np.fromiter)
def ovl_fromiter(iter, dtype):
    def np_fromiter_impl(iter, dtype):
        return np.asarray(iter, dtype)
    # Type check
    if not isinstance(iter, types.IterableType):
        raise TypingError("First argument must be an iterable.")
    else:
        return np_fromiter_impl

@njit
def gen():
    """Generator function returning 1D-arrays."""
    for i in range(3):
        yield i

@njit
def use_fromiter(dtype):
    return np.fromiter(gen(), dtype)

print(use_fromiter(np.int64))
# [0 1 2]

print(np.asarray(gen.py_func()))
# <generator object gen at 0x7f7782e64880>