Commits

Brian Kearns  committed bb85257

avoid loops in numpy array creation

  • Participants
  • Parent commits 704d284

Comments (0)

Files changed (3)

File pypy/module/micronumpy/ctors.py

 from pypy.module.micronumpy import descriptor, loop, ufuncs
 from pypy.module.micronumpy.base import W_NDimArray, convert_to_array
 from pypy.module.micronumpy.converters import shape_converter
-from pypy.module.micronumpy.strides import find_shape_and_elems
 
 
 def build_scalar(space, w_dtype, w_state):
 @unwrap_spec(ndmin=int, copy=bool, subok=bool)
 def array(space, w_object, w_dtype=None, copy=True, w_order=None, subok=False,
           ndmin=0):
+    from pypy.module.micronumpy import strides
+
     # for anything that isn't already an array, try __array__ method first
     if not isinstance(w_object, W_NDimArray):
         w___array__ = space.lookup(w_object, "__array__")
         return w_ret
 
     # not an array or incorrect dtype
-    shape, elems_w = find_shape_and_elems(space, w_object, dtype)
+    shape, elems_w = strides.find_shape_and_elems(space, w_object, dtype)
     if dtype is None or (dtype.is_str_or_unicode() and dtype.elsize < 1):
-        for w_elem in elems_w:
-            if isinstance(w_elem, W_NDimArray) and w_elem.is_scalar():
-                w_elem = w_elem.get_scalar_value()
-            dtype = ufuncs.find_dtype_for_scalar(space, w_elem, dtype)
+        dtype = strides.find_dtype_for_seq(space, elems_w, dtype)
         if dtype is None:
             dtype = descriptor.get_dtype_cache(space).w_float64dtype
         elif dtype.is_str_or_unicode() and dtype.elsize < 1:
     if ndmin > len(shape):
         shape = [1] * (ndmin - len(shape)) + shape
     w_arr = W_NDimArray.from_shape(space, shape, dtype, order=order)
-    arr_iter = w_arr.create_iter()
-    for w_elem in elems_w:
-        arr_iter.setitem(dtype.coerce(space, w_elem))
-        arr_iter.next()
+    if len(elems_w) == 1:
+        w_arr.set_scalar_value(dtype.coerce(space, elems_w[0]))
+    else:
+        loop.assign(space, w_arr, elems_w)
     return w_arr
 
 

File pypy/module/micronumpy/loop.py

         arr_iter.setitem(box)
         arr_iter.next()
 
+def assign(space, arr, seq):
+    arr_iter = arr.create_iter()
+    arr_dtype = arr.get_dtype()
+    for item in seq:
+        arr_iter.setitem(arr_dtype.coerce(space, item))
+        arr_iter.next()
+
 where_driver = jit.JitDriver(name='numpy_where',
                              greens = ['shapelen', 'dtype', 'arr_dtype'],
                              reds = 'auto')

File pypy/module/micronumpy/strides.py

         return [], [w_iterable]
     if isinstance(w_iterable, W_NDimArray) and w_iterable.is_scalar():
         return [], [w_iterable]
+    return _find_shape_and_elems(space, w_iterable, is_rec_type)
+
+
+def _find_shape_and_elems(space, w_iterable, is_rec_type):
     shape = [space.len_w(w_iterable)]
     batch = space.listview(w_iterable)
     while True:
         batch = new_batch
 
 
+def find_dtype_for_seq(space, elems_w, dtype):
+    from pypy.module.micronumpy.ufuncs import find_dtype_for_scalar
+    if len(elems_w) == 1:
+        w_elem = elems_w[0]
+        if isinstance(w_elem, W_NDimArray) and w_elem.is_scalar():
+            w_elem = w_elem.get_scalar_value()
+        return find_dtype_for_scalar(space, w_elem, dtype)
+    return _find_dtype_for_seq(space, elems_w, dtype)
+
+
+def _find_dtype_for_seq(space, elems_w, dtype):
+    from pypy.module.micronumpy.ufuncs import find_dtype_for_scalar
+    for w_elem in elems_w:
+        if isinstance(w_elem, W_NDimArray) and w_elem.is_scalar():
+            w_elem = w_elem.get_scalar_value()
+        dtype = find_dtype_for_scalar(space, w_elem, dtype)
+    return dtype
+
+
 def to_coords(space, shape, size, order, w_item_or_slice):
     '''Returns a start coord, step, and length.
     '''