1 // Copyright 2012-2016 Tinyarray authors.
2 //
3 // This file is part of Tinyarray. It is subject to the license terms in the
4 // file LICENSE.rst found in the top-level directory of this distribution and
5 // at https://gitlab.kwant-project.org/kwant/tinyarray/blob/master/LICENSE.rst.
6 // A list of Tinyarray authors can be found in the README.rst file at the
7 // top-level directory of this distribution and at
8 // https://gitlab.kwant-project.org/kwant/tinyarray.
9
10 #include <Python.h>
11 #include <cstddef>
12 #include <sstream>
13 #include <limits>
14 #include <algorithm>
15 #include "array.hh"
16 #include "arithmetic.hh"
17 #include "functions.hh"
18 #include "conversion.hh"
19 #include "version.hh"
20
21 template <>
22 const char *Array<long>::pyformat = "l";
23 template <>
24 const char *Array<double>::pyformat = "d";
25 template <>
26 const char *Array<Complex>::pyformat = "Zd";
27
28 const char *dtype_names[] = {"int", "float", "complex"};
29
30 const char *format_names[] = {
31 "int32 little endian", "int32 big endian",
32 "int64 little endian", "int64 big endian",
33 "float64 little endian", "float64 big endian",
34 "complex128 little endian", "complex128 big endian",
35 "unknown"};
36
37 Format format_by_dtype[NONE];
38
39 namespace {
40
is_big_endian()41 bool is_big_endian()
42 {
43 union {
44 unsigned int i;
45 char c[4];
46 } bint = {0x0a0b0c0d};
47
48 if (bint.c[0] == 0x0a) {
49 assert(bint.c[1] == 0x0b);
50 assert(bint.c[2] == 0x0c);
51 assert(bint.c[3] == 0x0d);
52 return true;
53 } else {
54 assert(bint.c[0] == 0x0d);
55 assert(bint.c[1] == 0x0c);
56 assert(bint.c[2] == 0x0b);
57 assert(bint.c[3] == 0x0a);
58 return false;
59 }
60 }
61
62 PyObject *int_str, *long_str, *float_str, *complex_str, *index_str,
63 *reconstruct;
64
dtype_of_scalar(PyObject * obj)65 Dtype dtype_of_scalar(PyObject *obj)
66 {
67 if (PyComplex_Check(obj)) return COMPLEX;
68 if (PyFloat_Check(obj)) return DOUBLE;
69 if (PyInt_Check(obj)) return LONG;
70 // TODO: The following line should be removed for Python 3.
71 if (PyLong_Check(obj)) return LONG;
72 if (PyObject_HasAttr(obj, index_str)) return LONG;
73
74 if (PyObject_HasAttr(obj, complex_str)) return COMPLEX;
75 if (PyObject_HasAttr(obj, float_str)) return DOUBLE;
76 if (PyObject_HasAttr(obj, int_str)) return LONG;
77 // TODO: The following line should be removed for Python 3.
78 if (PyObject_HasAttr(obj, long_str)) return LONG;
79
80 return NONE;
81 }
82
dtype_of_buffer(Py_buffer * view)83 Dtype dtype_of_buffer(Py_buffer *view)
84 {
85 const char *fmt = view->format;
86 Dtype dtype = NONE;
87
88 // Currently, we only understand native endianness and alignment.
89 if (*fmt == '@') fmt++;
90
91 if (strchr("cbB?hHiIlLqQnN", *fmt)) {
92 dtype = LONG;
93 fmt++;
94 } else if (strchr("fdg", *fmt)) {
95 dtype = DOUBLE;
96 fmt++;
97 } else if (*fmt == 'Z') {
98 fmt++;
99 if (strchr("fdg", *fmt)) {
100 dtype = COMPLEX;
101 }
102 fmt++;
103 }
104
105 // Right now, no composite data structures are supported; if we found a
106 // single supported data type, we should be at the end of the string.
107 if (*fmt != '\0') return NONE;
108
109 return dtype;
110 }
111
112 template<typename O, typename I>
convert_array(PyObject * in_,int ndim,size_t * shape)113 PyObject *convert_array(PyObject *in_, int ndim, size_t *shape)
114 {
115 assert(Array<I>::check_exact(in_)); Array<I> *in = (Array<I>*)in_;
116 size_t size;
117 if (ndim == -1) {
118 assert(shape == 0);
119 in->ndim_shape(&ndim, &shape);
120 } else {
121 #ifndef NDEBUG
122 int in_ndim;
123 size_t *in_shape;
124 in->ndim_shape(&in_ndim, &in_shape);
125 assert(shape);
126 assert(calc_size(ndim, shape) == calc_size(in_ndim, in_shape));
127 #endif
128 }
129 Array<O> *out = Array<O>::make(ndim, shape, &size);
130 I *src = in->data();
131 O *dest = out->data();
132 for (size_t i = 0; i < size; ++i) dest[i] = src[i];
133 return (PyObject*)out;
134 }
135
136 typedef PyObject *Convert_array(PyObject*, int, size_t*);
137
138 Convert_array *convert_array_dtable[][3] = {
139 {convert_array<long, long>,
140 convert_array<long, double>,
141 0},
142 {convert_array<double, long>,
143 convert_array<double, double>,
144 0},
145 {convert_array<Complex, long>,
146 convert_array<Complex, double>,
147 convert_array<Complex, Complex>}
148 };
149
convert_array(Dtype dtype_out,PyObject * in,Dtype dtype_in,int ndim=-1,size_t * shape=0)150 PyObject *convert_array(Dtype dtype_out, PyObject *in, Dtype dtype_in,
151 int ndim = -1, size_t *shape = 0)
152 {
153 if (dtype_in == NONE)
154 dtype_in = get_dtype(in);
155 assert(get_dtype(in) == get_dtype(in));
156 Convert_array *func = convert_array_dtable[int(dtype_out)][int(dtype_in)];
157 if (!func) {
158 PyErr_Format(PyExc_TypeError, "Cannot convert %s to %s.",
159 dtype_names[int(dtype_in)], dtype_names[int(dtype_out)]);
160 return 0;
161 }
162 return func(in, ndim, shape);
163 }
164
165 const char *seq_err_msg =
166 "A sequence does not support sequence protocol - "
167 "this is probably due to a bug in numpy for 0-d arrays.";
168
169 // This function determines the shape of an array-like sequence (of
170 // sequences...) given to it as first parameter. `dtype_guess' is the dtype of
171 // the first element of the sequence.
172 //
173 // All four arguments after the first one are written to. `shape' and `seqs'
174 // must have space for at least `max_ndim' elements.
175 //
176 // After successful execution `seqs' will contain `ndim' new references
177 // returned by PySequence_Fast.
examine_sequence(PyObject * arraylike,int * ndim,size_t * shape,PyObject ** seqs,Dtype * dtype_guess)178 int examine_sequence(PyObject *arraylike, int *ndim, size_t *shape,
179 PyObject **seqs, Dtype *dtype_guess)
180 {
181 PyObject *p = arraylike;
182 int d = -1;
183 assert(PySequence_Check(p));
184 for (bool is_sequence = true; ; is_sequence = PySequence_Check(p)) {
185 if (is_sequence) {
186 ++d;
187 if (d == ptrdiff_t(max_ndim)) {
188 // Strings are, in a way, infinitely nested sequences because
189 // the first element of a string is a again a string.
190 if (PyString_Check(p))
191 PyErr_SetString(PyExc_TypeError, "Expecting a number.");
192 else
193 PyErr_SetString(PyExc_ValueError, "Too many dimensions.");
194 --d;
195 goto fail;
196 }
197 } else {
198 // We are in the innermost sequence. Determine the dtype if
199 // requested.
200 if (dtype_guess) {
201 *dtype_guess = dtype_of_scalar(p);
202 if (*dtype_guess == NONE) {
203 PyErr_SetString(PyExc_TypeError, "Expecting a number.");
204 goto fail;
205 }
206 }
207 break;
208 }
209
210 // See https://github.com/numpy/numpy/issues/652.
211 seqs[d] = PySequence_Fast(p, seq_err_msg);
212 if (!seqs[d]) {--d; goto fail;}
213
214 if ((shape[d] = PySequence_Fast_GET_SIZE(seqs[d]))) {
215 p = *PySequence_Fast_ITEMS(seqs[d]);
216 } else {
217 // We are in the innermost sequence which is empty.
218 if (dtype_guess) *dtype_guess = NONE;
219 break;
220 }
221 }
222 *ndim = d + 1;
223 return 0;
224
225 fail:
226 for (; d >= 0; --d) Py_DECREF(seqs[d]);
227 return -1;
228 }
229
230 // This function is designed to be run after examine_sequence. It takes care
231 // of releasing the references passed to it in seqs.
232 template <typename T>
readin_seqs(PyObject ** seqs,T * dest,int ndim,const size_t * shape,bool exact)233 int readin_seqs(PyObject **seqs, T *dest, int ndim, const size_t *shape,
234 bool exact)
235 {
236 assert(ndim > 0);
237
238 // seqs is the stack of sequences being processed, all returned by
239 // PySequence_Fast. ps[d] and es[d] are the begin and end of the elements
240 // of seqs[d - 1].
241 PyObject **ps[max_ndim], **es[max_ndim];
242 es[0] = ps[0] = 0;
243
244 for (int d = 1; d < ndim; ++d) {
245 PyObject **p = PySequence_Fast_ITEMS(seqs[d - 1]);
246 ps[d] = p + 1;
247 es[d] = p + shape[d - 1];
248 }
249
250 int d = ndim - 1;
251 size_t len = shape[d];
252 PyObject **p = PySequence_Fast_ITEMS(seqs[d]), **e = p + len;
253 while (true) {
254 if (len && PySequence_Check(p[0])) {
255 if (d + 1 == ndim) {
256 PyErr_SetString(PyExc_ValueError,
257 "Input has irregular nesting depth.");
258 goto fail;
259 }
260 ++d;
261 ps[d] = p;
262 es[d] = e;
263 } else {
264 // Read-in a leaf sequence.
265 while (p < e) {
266 T value;
267 if (exact)
268 value = number_from_pyobject_exact<T>(*p++);
269 else
270 value = number_from_pyobject<T>(*p++);
271 if (value == T(-1) && PyErr_Occurred()) goto fail;
272 *dest++ = value;
273 }
274 Py_DECREF(seqs[d]);
275
276 while (ps[d] == es[d]) {
277 if (d == 0) {
278 // Success!
279 return 0;
280 }
281 --d;
282 Py_DECREF(seqs[d]);
283 }
284 if (!PySequence_Check(*ps[d])) {
285 --d;
286 PyErr_SetString(PyExc_ValueError,
287 "Input has irregular nesting depth.");
288 goto fail;
289 }
290 }
291
292 // See https://github.com/numpy/numpy/issues/652.
293 seqs[d] = PySequence_Fast(*ps[d]++, seq_err_msg);
294 if (!seqs[d]) {--d; goto fail;}
295 len = PySequence_Fast_GET_SIZE(seqs[d]);
296
297 // Verify that the length of the current sequence agrees with the
298 // shape.
299 if (len != shape[d]) {
300 PyErr_SetString(PyExc_ValueError,
301 "Input has irregular shape.");
302 goto fail;
303 }
304
305 p = PySequence_Fast_ITEMS(seqs[d]);
306 e = p + len;
307 }
308
309 fail:
310 while (true) {
311 Py_DECREF(seqs[d]);
312 if (d == 0) break;
313 --d;
314 }
315 return -1;
316 }
317
318 template <typename T>
readin_seqs_into_new(PyObject ** seqs,int ndim_in,int ndim_out,const size_t * shape_out,bool exact)319 PyObject *readin_seqs_into_new(PyObject **seqs, int ndim_in, int ndim_out,
320 const size_t *shape_out, bool exact)
321 {
322 Array<T> *result = Array<T>::make(ndim_out, shape_out);
323 assert(ndim_out >= ndim_in);
324 #ifndef NDEBUG
325 for (int d = 0, e = ndim_out - ndim_in; d < e; ++d)
326 assert(shape_out[d] == 1);
327 #endif
328 if (result == 0) return 0;
329 if (readin_seqs<T>(seqs, result->data(), ndim_in,
330 shape_out + ndim_out - ndim_in, exact)
331 == -1) {
332 Py_DECREF(result);
333 return 0;
334 }
335 return (PyObject*)result;
336 }
337
338 PyObject *(*readin_seqs_into_new_dtable[])(
339 PyObject**, int, int, const size_t*, bool) =
340 DTYPE_DISPATCH(readin_seqs_into_new);
341
342 template <typename T>
readin_scalar_into_new(PyObject * in,bool exact,int ndim=0)343 PyObject *readin_scalar_into_new(PyObject *in, bool exact, int ndim = 0)
344 {
345 T value;
346 if (exact)
347 value = number_from_pyobject_exact<T>(in);
348 else
349 value = number_from_pyobject<T>(in);
350
351 if (value == T(-1) && PyErr_Occurred()) return 0;
352
353 Array<T> *result = Array<T>::make(ndim, 1);
354 *result->data() = value;
355
356 size_t *shape;
357 result->ndim_shape(0, &shape);
358 for (int d = 0; d < ndim; ++d) shape[d] = 1;
359
360 return (PyObject*)result;
361 }
362
363 PyObject *(*readin_scalar_into_new_dtable[])(PyObject*, bool, int) =
364 DTYPE_DISPATCH(readin_scalar_into_new);
365
examine_buffer(PyObject * in,Py_buffer * view,Dtype * dtype)366 int examine_buffer(PyObject *in, Py_buffer *view, Dtype *dtype)
367 {
368 if (!PyObject_CheckBuffer(in)) return -1;
369 Dtype dt = NONE;
370 memset(view, 0, sizeof(Py_buffer));
371
372 // I don't know if the following makes much sense: I try to get the buffer
373 // using less and less demanding flags. NumPy does the same.
374 if (PyObject_GetBuffer(in, view, PyBUF_ND | PyBUF_FORMAT) == 0)
375 dt = dtype_of_buffer(view);
376 else if (PyObject_GetBuffer(in, view, PyBUF_STRIDES | PyBUF_FORMAT) == 0)
377 dt = dtype_of_buffer(view);
378 else if (PyObject_GetBuffer(in, view, PyBUF_FULL_RO) == 0)
379 dt = dtype_of_buffer(view);
380 PyErr_Clear();
381
382 // Check if the buffer can actually be converted into one of our
383 // formats.
384 if (dt == NONE) return -1;
385
386 if (dtype) *dtype = dt;
387 return 0;
388 }
389
390 template<typename T>
391 T (*get_buffer_converter_complex(const char *fmt))(const void *);
392
393 template<>
get_buffer_converter_complex(const char *)394 long (*get_buffer_converter_complex(const char *))(const void *)
395 {
396 // Complex can only be cast to complex.
397 PyErr_Format(PyExc_TypeError, "Complex cannot be cast to int.");
398
399 return 0;
400 }
401
402 template<>
get_buffer_converter_complex(const char *)403 double (*get_buffer_converter_complex(const char *))(const void *)
404 {
405 // complex can only be cast to complex
406 PyErr_Format(PyExc_TypeError, "Complex cannot be cast to float.");
407
408 return 0;
409 }
410
411 template<>
get_buffer_converter_complex(const char * fmt)412 Complex (*get_buffer_converter_complex(const char *fmt))(const void *)
413 {
414 switch(*(fmt + 1)){
415 case 'f':
416 return number_from_ptr<Complex, std::complex<float> >;
417 case 'd':
418 return number_from_ptr<Complex, std::complex<double> >;
419 case 'g':
420 return number_from_ptr<Complex, std::complex<long double> >;
421 }
422
423 return 0;
424 }
425
426 template<typename T>
get_buffer_converter(Py_buffer * view)427 T (*get_buffer_converter(Py_buffer *view))(const void *)
428 {
429 // currently, we only understand native endianness and alignment
430 char *fmt = view->format;
431
432 if (*fmt == '@') {
433 fmt++;
434 }
435
436 switch(*fmt) {
437 case 'c':
438 return number_from_ptr<T, char>;
439 case 'b':
440 return number_from_ptr<T, signed char>;
441 case 'B':
442 return number_from_ptr<T, unsigned char>;
443 case '?':
444 return number_from_ptr<T, bool>;
445 case 'h':
446 return number_from_ptr<T, short>;
447 case 'H':
448 return number_from_ptr<T, unsigned short>;
449 case 'i':
450 return number_from_ptr<T, int>;
451 case 'I':
452 return number_from_ptr<T, unsigned int>;
453 case 'l':
454 return number_from_ptr<T, long>;
455 case 'L':
456 return number_from_ptr<T, unsigned long>;
457 case 'q':
458 return number_from_ptr<T, long long>;
459 case 'Q':
460 return number_from_ptr<T, unsigned long long>;
461 case 'n':
462 return number_from_ptr<T, ssize_t>;
463 case 'N':
464 return number_from_ptr<T, size_t>;
465 case 'f':
466 return number_from_ptr<T, float>;
467 case 'd':
468 return number_from_ptr<T, double>;
469 case 'g':
470 return number_from_ptr<T, long double>;
471 case 'Z':
472 return get_buffer_converter_complex<T>(fmt);
473 }
474
475 return 0;
476 }
477
478 template<typename T>
readin_buffer(T * dest,Py_buffer * view)479 int readin_buffer(T *dest, Py_buffer *view)
480 {
481 if (view->len == 0) return 0;
482
483 T (*number_from_ptr)(const void *) = get_buffer_converter<T>(view);
484 if (!number_from_ptr) return -1;
485
486 if (view->ndim == 0) {
487 *dest = (*number_from_ptr)(view->buf);
488 if (PyErr_Occurred()) return -1;
489 else return 0;
490 }
491
492 Py_ssize_t indices[max_ndim];
493 for (int i = 0; i < view->ndim; i++) {
494 indices[i] = 0;
495 }
496
497 if (view->suboffsets) {
498 while(indices[0] < view->shape[0]) {
499 char *pointer = (char*)view->buf;
500 for (int i = 0; i < view->ndim; i++) {
501 pointer += view->strides[i] * indices[i];
502 if (view->suboffsets[i] >=0 ) {
503 pointer = *((char**)pointer) + view->suboffsets[i];
504 }
505 }
506
507 *dest++ = (*number_from_ptr)(pointer);
508 if (PyErr_Occurred()) return -1;
509
510 indices[view->ndim-1]++;
511
512 for (int i = view->ndim - 1; i > 0; i--) {
513 if (indices[i] < view->shape[i]) break;
514 assert(indices[i] == view->shape[i]);
515 indices[i-1]++;
516 indices[i] = 0;
517 }
518 }
519 } else if (view->strides) {
520 char *ptr = (char *)view->buf;
521
522 while(indices[0] < view->shape[0]) {
523 *dest++ = (*number_from_ptr)(ptr);
524 if (PyErr_Occurred()) return -1;
525
526 indices[view->ndim-1]++;
527 ptr += view->strides[view->ndim-1];
528
529 for (int i = view->ndim - 1; i > 0; i--) {
530 if (indices[i] < view->shape[i]) break;
531 assert(indices[i] == view->shape[i]);
532 indices[i-1]++;
533 ptr += view->strides[i-1];
534 indices[i] = 0;
535 ptr -= view->strides[i] * view->shape[i];
536 }
537 }
538 } else {
539 // Must be C-contiguous.
540 char *end = (char *)view->buf + view->len;
541 char *p = (char *)view->buf;
542 while(p < end) {
543 *dest++ = (*number_from_ptr)(p);
544 if (PyErr_Occurred()) return -1;
545
546 p += view->itemsize;
547 }
548 }
549
550 return 0;
551 }
552
553 template <typename T>
make_and_readin_buffer(Py_buffer * view,int ndim_out,const size_t * shape_out)554 PyObject *make_and_readin_buffer(Py_buffer *view, int ndim_out,
555 const size_t *shape_out)
556 {
557 Array<T> *result = Array<T>::make(ndim_out, shape_out);
558 assert(ndim_out >= view->ndim);
559 #ifndef NDEBUG
560 for (int d = 0, e = ndim_out - view->ndim; d < e; ++d)
561 assert(shape_out[d] == 1);
562 #endif
563 if (result == 0) return 0;
564 if (readin_buffer<T>(result->data(), view) == -1) {
565 Py_DECREF(result);
566 return 0;
567 }
568 return (PyObject*)result;
569 }
570
571 PyObject *(*make_and_readin_buffer_dtable[])(
572 Py_buffer *, int, const size_t*) = DTYPE_DISPATCH(make_and_readin_buffer);
573
574 template <typename T>
to_pystring(Array<T> * self,PyObject * to_str (PyObject *),const char * header,const char * trailer,const char * indent,const char * separator)575 PyObject *to_pystring(Array<T> *self, PyObject* to_str(PyObject *),
576 const char *header, const char *trailer,
577 const char *indent, const char *separator)
578 {
579 int ndim;
580 size_t *shape;
581 self->ndim_shape(&ndim, &shape);
582
583 std::ostringstream o;
584 o << header;
585
586 const T *p = self->data();
587 if (ndim > 0) {
588 int d = 0;
589 size_t i[max_ndim];
590 i[0] = shape[0];
591
592 o << '[';
593 while (true) {
594 if (i[d]) {
595 --i[d];
596 if (d < ndim - 1) {
597 o << '[';
598 ++d;
599 i[d] = shape[d];
600 } else {
601 PyObject *num = pyobject_from_number(*p++);
602 PyObject *str = to_str(num);
603 o << PyString_AsString(str);
604 Py_DECREF(str);
605 Py_DECREF(num);
606 if (i[d] > 0) o << separator << ' ';
607 }
608 } else {
609 o << ']';
610 if (d == 0) break;
611 --d;
612 if (i[d]) {
613 o << separator << "\n " << indent;
614 for (int i = 0; i < d; ++i) o << ' ';
615 }
616 }
617 }
618 } else {
619 PyObject *num = pyobject_from_number(*p);
620 PyObject *str = to_str(num);
621 o << PyString_AsString(str);
622 Py_DECREF(str);
623 Py_DECREF(num);
624 }
625 o << trailer;
626
627 return PyString_FromString(o.str().c_str());
628 }
629
630 template <typename T>
repr(PyObject * obj)631 PyObject *repr(PyObject *obj)
632 {
633 Array<T> *self = reinterpret_cast<Array<T> *>(obj);
634 return to_pystring(self, PyObject_Repr, "array(", ")", " ", ",");
635 }
636
637 template <typename T>
str(PyObject * obj)638 PyObject *str(PyObject *obj)
639 {
640 Array<T> *self = reinterpret_cast<Array<T> *>(obj);
641 return to_pystring(self, PyObject_Str, "", "", "", "");
642 }
643
len(Array_base * self)644 Py_ssize_t len(Array_base *self)
645 {
646 int ndim;
647 size_t *shape;
648 self->ndim_shape(&ndim, &shape);
649 if (ndim == 0) {
650 PyErr_SetString(PyExc_TypeError, "len() of unsized object.");
651 return -1;
652 }
653 return shape[0];
654 }
655
index_from_key(int ndim,const size_t * shape,PyObject * key)656 Py_ssize_t index_from_key(int ndim, const size_t *shape, PyObject *key)
657 {
658 long indices[max_ndim];
659 int res = load_index_seq_as_long(key, indices, max_ndim);
660 if (res == -1) {
661 PyErr_SetString(PyExc_IndexError, "Invalid index.");
662 return -1;
663 }
664 if (res != ndim) {
665 PyErr_SetString(PyExc_IndexError, "Number of indices "
666 "must be equal to number of dimensions.");
667 return -1;
668 }
669
670 int d = 0;
671 Py_ssize_t s = shape[0];
672 Py_ssize_t index = indices[0];
673 if (index < 0) index += s;
674 if (index < 0 || index >= s) goto out_of_range;
675 for (d = 1; d < ndim; ++d) {
676 s = shape[d];
677 Py_ssize_t i = indices[d];
678 if (i < 0) i += s;
679 if (i < 0 || i >= s) goto out_of_range;
680 index *= s;
681 index += i;
682 }
683 return index;
684
685 out_of_range:
686 PyErr_Format(PyExc_IndexError, "Index %ld out of range "
687 "(-%lu <= index < %lu) in dimension %d.",
688 indices[d], (unsigned long)s, (unsigned long)s, d);
689 return -1;
690 }
691
692 template <typename T>
getitem(PyObject * obj,PyObject * key)693 PyObject *getitem(PyObject *obj, PyObject *key)
694 {
695 Array<T> *self = reinterpret_cast<Array<T> *>(obj);
696
697 if (PySlice_Check(key)) {
698 PyErr_SetString(PyExc_NotImplementedError,
699 "Slices are not implemented.");
700 return 0;
701 } else {
702 int ndim;
703 size_t *shape;
704 self->ndim_shape(&ndim, &shape);
705 T *data = self->data();
706 Py_ssize_t index = index_from_key(ndim, shape, key);
707 if (index == -1) return 0;
708 return pyobject_from_number(data[index]);
709 }
710 }
711
712 template <typename T>
seq_getitem(PyObject * obj,Py_ssize_t index)713 PyObject *seq_getitem(PyObject *obj, Py_ssize_t index)
714 {
715 int ndim;
716 size_t *shape;
717 Array<T> *self = reinterpret_cast<Array<T> *>(obj);
718 self->ndim_shape(&ndim, &shape);
719 assert(ndim != 0);
720
721 if (index < 0) index += shape[0];
722 if (size_t(index) >= shape[0]) {
723 PyErr_SetString(PyExc_IndexError, "Invalid index.");
724 return 0;
725 }
726
727 T *src = self->data();
728 if (ndim == 1) {
729 assert(index >= 0);
730 assert(size_t(index) < shape[0]);
731 return pyobject_from_number(src[index]);
732 }
733
734 assert(ndim > 1);
735 size_t item_size;
736 Array<T> *result = Array<T>::make(ndim - 1, shape + 1, &item_size);
737 if (!result) return 0;
738 src += index * item_size;
739 T *dest = result->data();
740 for (size_t i = 0; i < item_size; ++i) dest[i] = src[i];
741 return (PyObject*)result;
742 }
743
744 template <typename T>
getbuffer(PyObject * obj,Py_buffer * view,int flags)745 int getbuffer(PyObject *obj, Py_buffer *view, int flags)
746 {
747 int ndim;
748 size_t *shape, size;
749 Array<T> *self = reinterpret_cast<Array<T> *>(obj);
750
751 assert(view);
752 if ((flags & PyBUF_F_CONTIGUOUS) == PyBUF_F_CONTIGUOUS) {
753 PyErr_SetString(PyExc_BufferError,
754 "Tinyarrays are not Fortran contiguous.");
755 goto fail;
756 }
757 // The following has been commented out as a workaround for Cython's
758 // inability to deal with read-only buffers.
759 //
760 // if ((flags & PyBUF_WRITEABLE) == PyBUF_WRITEABLE) {
761 // PyErr_SetString(PyExc_BufferError, "Tinyarrays are not writeable");
762 // goto fail;
763 // }
764
765 self->ndim_shape(&ndim, &shape);
766 size = calc_size(ndim, shape);
767
768 view->buf = self->data();
769 view->itemsize = sizeof(T);
770 view->len = size * view->itemsize;
771 view->readonly = 1;
772 if ((flags & PyBUF_FORMAT) == PyBUF_FORMAT)
773 view->format = (char*)Array<T>::pyformat;
774 else
775 view->format = 0;
776 if ((flags & PyBUF_ND) == PyBUF_ND) {
777 view->ndim = ndim;
778 view->shape = (Py_ssize_t*)shape;
779 // From the documentation it's not clear whether it is allowed not to
780 // set strides to NULL (for C continuous arrays), but it works, so we
781 // do it. However, there is a bug in current numpy
782 // (http://projects.scipy.org/numpy/ticket/2197) which requires strides
783 // if view->len == 0. Because we don't have proper strides, we just
784 // set strides to shape. This dirty trick seems to work well -- no one
785 // looks at strides when len == 0.
786 if (size != 0)
787 view->strides = 0;
788 else
789 view->strides = (Py_ssize_t*)shape;
790 } else {
791 view->ndim = 0;
792 view->shape = 0;
793 view->strides = 0;
794 }
795 view->internal = 0;
796 view->suboffsets = 0;
797
798 // Success.
799 Py_INCREF(self);
800 view->obj = (PyObject*)self;
801
802 return 0;
803
804 fail:
805 view->obj = 0;
806 return -1;
807 }
808
809 // Given the same input, these hash functions return the same result as those
810 // in Python. As tinyarrays compare equal to equivalent tuples it is important
811 // for the hashes to agree. If not, there will be problems with dictionaries.
812
813 #if PY_MAJOR_VERSION >= 3
814
815 // the only documentation for this is in the Python sourcecode
816 const Py_hash_t HASH_IMAG = _PyHASH_IMAG;
817
hash(void * inst,long x)818 Py_hash_t hash(void *inst, long x)
819 {
820 // For integers the hash is just the integer itself modulo _PyHASH_MODULUS
821 // except for the singular case of -1.
822 // define 'sign' of the correct width to avoid overflow
823 Py_hash_t sign = x < 0 ? -1 : 1;
824 Py_hash_t result = sign * ((sign * x) % _PyHASH_MODULUS);
825 return result == -1 ? -2 : result;
826 }
827
828 #else
829
830 /* In Python 2 hashes were long integers, as indicated by
831 https://github.com/python/cpython/blob/2.7/Include/object.h#L314
832 */
833 typedef long Py_hash_t;
834 typedef unsigned long Py_uhash_t;
835 const Py_hash_t HASH_IMAG = 1000003L;
836
hash(void * inst,long x)837 Py_hash_t hash(void *inst, long x)
838 {
839 return x != -1 ? x : -2;
840 }
841
842 #endif
843
844 // We used to have our own implementation of this, but the extra function
845 // call is quite negligible compared to the execution time of the function.
846 #if PY_MAJOR_VERSION == 3 && PY_MINOR_VERSION >= 10
hash(void * inst,double x)847 Py_hash_t hash(void *inst, double x)
848 {
849 return _Py_HashDouble((PyObject *)inst, x);
850 #else
851 Py_hash_t hash(void *, double x)
852 {
853 return _Py_HashDouble(x);
854 #endif
855 }
856
857 Py_hash_t hash(void *inst, Complex x)
858 {
859 // x.imag == 0 => hash(x.imag) == 0 => hash(x) == hash(x.real)
860 return hash(inst, x.real()) + HASH_IMAG * hash(inst, x.imag());
861 }
862
863 // The following routine calculates the hash of a multi-dimensional array. The
864 // hash is equal to that of an arrangement of nested tuples equivalent to the
865 // array.
866 //
867 // It exists in two versions because Python's tuplehash has changed in Python
868 // 3.8 with the following motivation: "The hash function for tuples is now
869 // based on xxHash which gives better collision results on (formerly)
870 // pathological cases. Additionally, on 64-bit systems it improves tuple hashes
871 // in general."
872
873 #if (PY_MAJOR_VERSION < 3 || PY_MINOR_VERSION < 8) && PY_MAJOR_VERSION < 4
874
875 // Version for Python < 3.8
876 template <typename T>
877 Py_hash_t hash(PyObject *obj)
878 {
879 int ndim;
880 size_t *shape;
881 Array<T> *self = reinterpret_cast<Array<T> *>(obj);
882 self->ndim_shape(&ndim, &shape);
883 T *p = self->data();
884 if (ndim == 0) return hash(p, *p);
885
886 const Py_uhash_t mult_init = 1000003, r_init = 0x345678;
887 const Py_uhash_t mul_addend = 82520, r_addend = 97531;
888 Py_ssize_t i[max_ndim];
889 Py_uhash_t mult[max_ndim], r[max_ndim];
890 --ndim; // For convenience.
891 int d = 0;
892 i[0] = shape[0];
893 mult[0] = mult_init;
894 r[0] = r_init;
895 while (true) {
896 if (i[d]) {
897 --i[d];
898 if (d == ndim) {
899 // Innermost loop body.
900 r[d] = (r[d] ^ hash(p, *p)) * mult[d];
901 p++;
902 mult[d] += mul_addend + 2 * i[d];
903 } else {
904 // Entering a loop.
905 ++d;
906 i[d] = shape[d];
907 mult[d] = mult_init;
908 r[d] = r_init;
909 }
910 } else {
911 // Exiting a loop.
912 if (d == 0) {
913 // Exiting the outermost loop.
914 Py_uhash_t r_next = r[0] + r_addend;
915 return r_next == Py_uhash_t(-1) ? -2 : r_next;
916 }
917 --d;
918 Py_uhash_t r_next = r[d+1] + r_addend;
919 r_next = r_next == Py_uhash_t(-1) ? -2 : r_next;
920 r[d] = (r[d] ^ r_next) * mult[d];
921 mult[d] += mul_addend + 2 * i[d];
922 }
923 }
924 }
925
926 #else
927
928 #if SIZEOF_PY_UHASH_T > 4
929
930 const Py_uhash_t _hash_init = 2870177450012600261U;
931
932 inline void _hash_inner_loop(Py_uhash_t &acc, Py_uhash_t lane)
933 {
934 acc += lane * 14029467366897019727U;
935 acc = ((acc << 31) | (acc >> 33)); // Rotate left 31 bits.
936 acc *= 11400714785074694791U;
937 }
938
939 #else
940
941 const Py_uhash_t _hash_init = 374761393U;
942
943 inline void _hash_inner_loop(Py_uhash_t &acc, Py_uhash_t lane)
944 {
945 acc += lane * 2246822519U;
946 acc = ((acc << 13) | (acc >> 19)); // Rotate left 13 bits.
947 acc *= 2654435761U;
948 }
949
950 #endif
951
952 inline Py_uhash_t _hash_loop_end(Py_uhash_t acc, Py_uhash_t len)
953 {
954 acc += len ^ (_hash_init ^ 3527539UL);
955 if (acc == Py_uhash_t(-1)) return 1546275796;
956 return acc;
957 }
958
959 // Version for Python >= 3.8
960 template <typename T>
961 Py_hash_t hash(PyObject *obj)
962 {
963 int ndim;
964 size_t *shape;
965 Array<T> *self = reinterpret_cast<Array<T> *>(obj);
966 self->ndim_shape(&ndim, &shape);
967 T *p = self->data();
968 if (ndim == 0) return hash(p, *p);
969
970 Py_ssize_t i[max_ndim];
971 Py_uhash_t acc[max_ndim];
972 --ndim; // For convenience.
973 int d = 0;
974 i[0] = shape[0];
975 acc[0] = _hash_init;
976 // The following is equivalent to 'ndim' (the original value) nested loops.
977 while (true) {
978 if (i[d]) {
979 --i[d];
980 if (d == ndim) {
981 _hash_inner_loop(acc[d], hash(p, *p));
982 p++;
983 } else {
984 ++d;
985 i[d] = shape[d];
986 acc[d] = _hash_init;
987 }
988 } else {
989 if (d == 0) return _hash_loop_end(acc[0], shape[0]);
990 --d;
991 _hash_inner_loop(acc[d], _hash_loop_end(acc[d+1], shape[d+1]));
992 }
993 }
994 }
995
996 #endif
997
998 template <typename T>
999 bool compare_scalar(const int op, const T a, const T b) {
1000 switch(op){
1001 case Py_EQ: return a == b;
1002 case Py_NE: return a != b;
1003 case Py_LE: return a <= b;
1004 case Py_GE: return a >= b;
1005 case Py_LT: return a < b;
1006 case Py_GT: return a > b;
1007 default:
1008 assert(false); // If we get here something is very wrong.
1009 return false; // Stop the compiler complaining.
1010 }
1011 }
1012
1013
1014 template <>
1015 bool compare_scalar<Complex>(const int op, const Complex a, const Complex b) {
1016 switch(op){
1017 case Py_EQ: return a == b;
1018 case Py_NE: return a != b;
1019 // This function is never called in a context where
1020 // the following code path is run -- fall through.
1021 case Py_LE:
1022 case Py_GT:
1023 case Py_LT:
1024 case Py_GE:
1025 default:
1026 assert(false);
1027 return false; // Stop the compiler complaining.
1028 }
1029 }
1030
1031
1032 template <typename T>
1033 bool compare_data(int op, PyObject *a_, PyObject *b_, size_t size)
1034 {
1035 assert(Array<T>::check_exact(a_)); Array<T> *a = (Array<T>*)a_;
1036 assert(Array<T>::check_exact(b_)); Array<T> *b = (Array<T>*)b_;
1037 const T *data_a = a->data();
1038 const T *data_b = b->data();
1039 // Sequences are ordered the same as their first differing elements, see:
1040 // https://docs.python.org/3/reference/expressions.html#value-comparisons
1041 // comparison for "multidimensional" sequences is identical to comparing
1042 // the flattened sequences when they have the same shape (the present case).
1043 size_t i = 0;
1044 for (; i < size; ++i)
1045 if (data_a[i] != data_b[i]) break;
1046 // Any of these operations should return true when objects are equal.
1047 if (i == size) return ((op == Py_EQ) || (op == Py_LE) || (op == Py_GE));
1048 // encapsulate this into a function to handle the COMPLEX case
1049 return compare_scalar<T>(op, data_a[i], data_b[i]);
1050 }
1051
1052 bool (*compare_data_dtable[])(int, PyObject*, PyObject*, size_t) =
1053 DTYPE_DISPATCH(compare_data);
1054
1055 #if PY_MAJOR_VERSION < 3
1056 #define PY2_RAISE_TYPEERROR(msg)\
1057 PyErr_SetString(PyExc_TypeError, msg);\
1058 return 0
1059 #else
1060 #define PY2_RAISE_TYPEERROR(msg)
1061 #endif // PY_MAJOR_VERSION < 3
1062
1063 PyObject *richcompare(PyObject *a, PyObject *b, int op)
1064 {
1065 PyObject *result;
1066 const bool equality_comparison = (op == Py_EQ || op == Py_NE);
1067
1068 // Short circuit when we are comparing the same object.
1069 bool equal = (a == b);
1070 if (equal) {
1071 // Any of these operations should return true when objects are equal
1072 equal = (op == Py_EQ) || (op == Py_GE) || (op == Py_LE);
1073 result = equal ? Py_True : Py_False;
1074 goto done;
1075 }
1076
1077 Dtype dtype;
1078 if (coerce_to_arrays(&a, &b, &dtype) < 0) return 0;
1079
1080 // Obviate the need for `compare_scalar<Complex` to
1081 // handle the case of an undefined comparison.
1082 if (dtype == COMPLEX && !equality_comparison) {
1083 PY2_RAISE_TYPEERROR("unorderable type: complex()");
1084 result = Py_NotImplemented;
1085 goto decref_then_done;
1086 }
1087
1088 int ndim_a, ndim_b;
1089 size_t *shape_a, *shape_b;
1090 reinterpret_cast<Array_base*>(a)->ndim_shape(&ndim_a, &shape_a);
1091 reinterpret_cast<Array_base*>(b)->ndim_shape(&ndim_b, &shape_b);
1092
1093 // TODO: Enable array comparisons between arrays of differing dimensions.
1094 if (ndim_a != ndim_b) {
1095 if (equality_comparison) {
1096 goto equality_then_done;
1097 } else {
1098 PY2_RAISE_TYPEERROR("Unorderable type: only arrays with"
1099 "the same shape can be ordered.");
1100 result = Py_NotImplemented;
1101 goto decref_then_done;
1102 }
1103 }
1104 for (int d = 0; d < ndim_a; ++d) {
1105 if (shape_a[d] != shape_b[d]) {
1106 if (equality_comparison) {
1107 goto equality_then_done;
1108 } else {
1109 PY2_RAISE_TYPEERROR("Unorderable type: only arrays with"
1110 "the same shape can be ordered.");
1111 result = Py_NotImplemented;
1112 goto decref_then_done;
1113 }
1114 }
1115 }
1116
1117 // Actually compare the data.
1118 equal = compare_data_dtable[int(dtype)](op, a, b, calc_size(ndim_a, shape_a));
1119 result = equal ? Py_True : Py_False;
1120 goto decref_then_done;
1121
1122 // Non error-path exit points from this function
1123 equality_then_done:
1124 result = ((op == Py_EQ) == equal) ? Py_True : Py_False;
1125 decref_then_done:
1126 Py_DECREF(a);
1127 Py_DECREF(b);
1128 done:
1129 Py_INCREF(result);
1130 return result;
1131 }
1132
1133 PyObject *get_dtype_py(PyObject *self, void *)
1134 {
1135 static PyObject *dtypes[] = {
1136 (PyObject*)&PyInt_Type,
1137 (PyObject*)&PyFloat_Type,
1138 (PyObject*)&PyComplex_Type
1139 };
1140 int dtype = int(get_dtype(self));
1141 assert(dtype < int(NONE));
1142 Py_INCREF(dtypes[dtype]);
1143 return dtypes[dtype];
1144 }
1145
1146 PyObject *get_ndim(Array_base *self, void *)
1147 {
1148 int ndim;
1149 self->ndim_shape(&ndim, 0);
1150 return PyLong_FromLong(ndim);
1151 }
1152
1153 PyObject *get_size(Array_base *self, void *)
1154 {
1155 int ndim;
1156 size_t *shape;
1157 self->ndim_shape(&ndim, &shape);
1158 return PyLong_FromSize_t(calc_size(ndim, shape));
1159 }
1160
1161 PyObject *get_shape(Array_base *self, void *)
1162 {
1163 int ndim;
1164 size_t *shape;
1165 self->ndim_shape(&ndim, &shape);
1166 size_t result_shape = ndim;
1167 Array<long> *result = Array<long>::make(1, &result_shape);
1168 if (!result) return 0;
1169 long *data = result->data();
1170 for (int d = 0; d < ndim; ++d) data[d] = shape[d];
1171 return (PyObject*)result;
1172 }
1173
1174 PyGetSetDef getset[] = {
1175 {(char*)"dtype", get_dtype_py, 0, 0, 0},
1176 {(char*)"ndim", (getter)get_ndim, 0, 0, 0},
1177 {(char*)"size", (getter)get_size, 0, 0, 0},
1178 {(char*)"shape", (getter)get_shape, 0, 0, 0},
1179 {0, 0, 0, 0, 0} // Sentinel
1180 };
1181
1182 // **************** Iterator ****************
1183
1184 template <typename T>
1185 class Array_iter {
1186 public:
1187 static Array_iter *make(Array<T> *array);
1188 static void dealloc(Array_iter<T> *self);
1189 static PyObject *next(Array_iter<T> *self);
1190 static PyObject *len(Array_iter<T> *self);
1191 private:
1192 static PyMethodDef methods[];
1193 static PyTypeObject pytype;
1194 static const char *pyname;
1195 PyObject ob_base;
1196 size_t index;
1197 Array<T> *array; // Set to 0 when iterator is exhausted.
1198 };
1199
1200 template <>
1201 const char *Array_iter<long>::pyname = "tinyarray.ndarrayiter_int";
1202 template <>
1203 const char *Array_iter<double>::pyname = "tinyarray.ndarrayiter_float";
1204 template <>
1205 const char *Array_iter<Complex>::pyname = "tinyarray.ndarrayiter_complex";
1206
1207 template <typename T>
1208 Array_iter<T> *Array_iter<T>::make(Array<T> *array)
1209 {
1210 int ndim;
1211 assert(Array<T>::check_exact((PyObject*)array));
1212 array->ndim_shape(&ndim, 0);
1213 if (ndim == 0) {
1214 PyErr_SetString(PyExc_TypeError, "Iteration over a 0-d array.");
1215 return 0;
1216 }
1217 Array_iter<T> *ret = PyObject_New(Array_iter<T>, &Array_iter<T>::pytype);
1218 if (ret == 0) return 0;
1219 ret->index = 0;
1220 Py_INCREF(array);
1221 ret->array = array;
1222 return ret;
1223 }
1224
1225 template <typename T>
1226 void Array_iter<T>::dealloc(Array_iter<T> *self)
1227 {
1228 // We use Py_XDECREF as array is already decref'ed when the iterator gets
1229 // exhausted.
1230 Py_XDECREF(self->array);
1231 PyObject_Del(self);
1232 }
1233
1234 template <typename T>
1235 PyObject *Array_iter<T>::next(Array_iter<T> *self)
1236 {
1237 Array<T> *array = self->array;
1238 if (array == 0) return 0;
1239 int ndim;
1240 size_t *shape;
1241 array->ndim_shape(&ndim, &shape);
1242 assert(ndim != 0);
1243
1244 if (self->index == shape[0]) {
1245 // End of iteration.
1246 Py_DECREF(array);
1247 self->array = 0;
1248 return 0;
1249 }
1250
1251 T *src = array->data();
1252
1253 if (ndim == 1) {
1254 assert(size_t(self->index) < shape[0]);
1255 return pyobject_from_number(src[self->index++]);
1256 }
1257
1258 assert(ndim > 1);
1259 size_t item_size;
1260 Array<T> *result = Array<T>::make(ndim - 1, shape + 1, &item_size);
1261 if (!result) return 0;
1262 src += item_size * self->index++;
1263 T *dest = result->data();
1264 for (size_t i = 0; i < item_size; ++i) dest[i] = src[i];
1265 return (PyObject*)result;
1266 }
1267
1268 template <typename T>
1269 PyObject *Array_iter<T>::len(Array_iter<T> *self)
1270 {
1271 Py_ssize_t len = 0;
1272 if (self->array) {
1273 #ifndef NDEBUG
1274 int ndim;
1275 self->array->ndim_shape(&ndim, 0);
1276 assert(ndim != 0);
1277 #endif
1278 size_t *shape;
1279 self->array->ndim_shape(0, &shape);
1280 len = shape[0] - self->index;
1281 }
1282 return PyInt_FromSsize_t(len);
1283 }
1284
1285 template <typename T>
1286 PyMethodDef Array_iter<T>::methods[] = {
1287 {"__length_hint__", (PyCFunction)Array_iter<T>::len, METH_NOARGS,
1288 "Private method returning an estimate of len(list(it))."},
1289 {0, 0} // Sentinel
1290 };
1291
1292 template <typename T>
1293 PyTypeObject Array_iter<T>::pytype = {
1294 PyVarObject_HEAD_INIT(&PyType_Type, 0)
1295 pyname, // tp_name
1296 sizeof(Array_iter<T>), // tp_basicsize
1297 0, // tp_itemsize
1298 // methods
1299 (destructor)Array_iter<T>::dealloc, // tp_dealloc
1300 0, // tp_print
1301 0, // tp_getattr
1302 0, // tp_setattr
1303 0, // tp_compare
1304 0, // tp_repr
1305 0, // tp_as_number
1306 0, // tp_as_sequence
1307 0, // tp_as_mapping
1308 0, // tp_hash
1309 0, // tp_call
1310 0, // tp_str
1311 PyObject_GenericGetAttr, // tp_getattro
1312 0, // tp_setattro
1313 0, // tp_as_buffer
1314 Py_TPFLAGS_DEFAULT, // tp_flags
1315 0, // tp_doc
1316 0, // tp_traverse
1317 0, // tp_clear
1318 0, // tp_richcompare
1319 0, // tp_weaklistoffset
1320 PyObject_SelfIter, // tp_iter
1321 (iternextfunc)Array_iter<T>::next, // tp_iternext
1322 Array_iter<T>::methods // tp_methods
1323 };
1324
1325 // The following explicit instantiations are necessary for GCC 4.6 but not for
1326 // GCC 4.7. I don't know why.
1327
1328 template PyObject *repr<long>(PyObject*);
1329 template PyObject *repr<double>(PyObject*);
1330 template PyObject *repr<Complex>(PyObject*);
1331
1332 template PyObject *str<long>(PyObject*);
1333 template PyObject *str<double>(PyObject*);
1334 template PyObject *str<Complex>(PyObject*);
1335
1336 template Py_hash_t hash<long>(PyObject*);
1337 template Py_hash_t hash<double>(PyObject*);
1338 template Py_hash_t hash<Complex>(PyObject*);
1339
1340 template int getbuffer<long>(PyObject*, Py_buffer*, int);
1341 template int getbuffer<double>(PyObject*, Py_buffer*, int);
1342 template int getbuffer<Complex>(PyObject*, Py_buffer*, int);
1343
1344 template PyObject *getitem<long>(PyObject*, PyObject*);
1345 template PyObject *getitem<double>(PyObject*, PyObject*);
1346 template PyObject *getitem<Complex>(PyObject*, PyObject*);
1347
1348 template PyObject *seq_getitem<long>(PyObject*, Py_ssize_t);
1349 template PyObject *seq_getitem<double>(PyObject*, Py_ssize_t);
1350 template PyObject *seq_getitem<Complex>(PyObject*, Py_ssize_t);
1351
1352 } // Anonymous namespace
1353
1354 // **************** Public interface ****************
1355
1356 PyDoc_STRVAR(tinyarray_doc,
1357 "Arrays of numbers for Python, optimized for small sizes\n\n\
1358 \n\
1359 The tinyarray module provides multi-dimensional arrays of numbers, similar to\n\
1360 NumPy, but with the following differences:\n\
1361 \n\
1362 * Optimized for small sizes: Compared to NumPy, common operations on small\n\
1363 arrays (e.g. length 3) are up to 35 times faster, and 3 times less memory is\n\
1364 used to store them.\n\
1365 \n\
1366 * Arrays are immutable and hashable, and can be thus used as dictionary keys.\n\
1367 \n\
1368 * The tinyarray module provides only the functionality that is deemed essential\n\
1369 with small arrays. For example, there exists a fast tinyarray.dot function,\n\
1370 but there is no fancy indexing or slicing.\n\
1371 \n\
1372 The module's interface is a basic subset of that of NumPy and hence should be\n\
1373 familiar to many Python programmers. All functions are simplified versions of\n\
1374 their NumPy counterparts.\n\
1375 \n\
1376 For example, arrays can be created with:\n\
1377 \n\
1378 tinyarray.array(arraylike, [dtype])\n\
1379 \n\
1380 where arraylike can be a number, a sequence (of sequences, ...) of numbers,\n\
1381 a NumPy or tinyarray array, or another object supporting the buffer protocol.\n\
1382 The dtype parameter is optional and can only take the values int, float, and\n\
1383 complex. Note that dtype is a positional argument and cannot be used as a\n\
1384 keyword argument. Arrays can be also created with the functions identity,\n\
1385 zeros, and ones.\n\
1386 \n\
1387 Tinyarrays support iteration and indexing (currently without slicing), as well\n\
1388 as vectorized elementwise arithmetics. A small number of operations like dot,\n\
1389 floor, and transpose are provided. Printing works as well as pickling.\n\
1390 Whenever an operation is missing from Tinyarray, NumPy can be used directly,\n\
1391 e.g.: numpy.linalg.det(my_tinyarray).");
1392
1393 extern "C"
MOD_INIT_FUNC(tinyarray)1394 MOD_INIT_FUNC(tinyarray)
1395 {
1396 // Determine storage formats.
1397 bool be = is_big_endian();
1398 if (std::numeric_limits<double>::is_iec559 &&
1399 sizeof(double) == 8) {
1400 format_by_dtype[int(DOUBLE)] = Format(FLOAT64_LE + be);
1401 format_by_dtype[int(COMPLEX)] = Format(COMPLEX128_LE + be);
1402 } else {
1403 format_by_dtype[int(DOUBLE)] = UNKNOWN;
1404 format_by_dtype[int(COMPLEX)] = UNKNOWN;
1405 }
1406 if (sizeof(long) == 8)
1407 format_by_dtype[int(LONG)] = Format(INT64_LE + be);
1408 else if (sizeof(long) == 4)
1409 format_by_dtype[int(LONG)] = Format(INT32_LE + be);
1410 else
1411 format_by_dtype[int(LONG)] = UNKNOWN;
1412
1413 if (PyType_Ready(&Array<long>::pytype) < 0) return MOD_ERROR_VAL;
1414 if (PyType_Ready(&Array<double>::pytype) < 0) return MOD_ERROR_VAL;
1415 if (PyType_Ready(&Array<Complex>::pytype) < 0) return MOD_ERROR_VAL;
1416
1417 PyObject *m;
1418 MOD_DEF(m, "tinyarray", functions, tinyarray_doc);
1419
1420 reconstruct = PyObject_GetAttrString(m, "_reconstruct");
1421
1422 Py_INCREF(&Array<long>::pytype);
1423 Py_INCREF(&Array<double>::pytype);
1424 Py_INCREF(&Array<Complex>::pytype);
1425
1426 PyModule_AddObject(m, "__version__", PyString_FromString(VERSION));
1427
1428 PyObject *all = PyList_New(0);
1429 for (const PyMethodDef *f = functions; f->ml_name; ++f) {
1430 if (f->ml_name[0] == '_') continue;
1431 PyObject *f_py = PyObject_GetAttrString(m, f->ml_name);
1432 PyList_Append(all, PyObject_GetAttrString(f_py, "__name__"));
1433 Py_DECREF(f_py);
1434 }
1435 PyModule_AddObject(m, "__all__", all);
1436
1437 PyModule_AddObject(m, "ndarray_int",
1438 (PyObject *)&Array<long>::pytype);
1439 PyModule_AddObject(m, "ndarray_float",
1440 (PyObject *)&Array<double>::pytype);
1441 PyModule_AddObject(m, "ndarray_complex",
1442 (PyObject *)&Array<Complex>::pytype);
1443
1444 // export information on the sizes of different dtypes in bytes
1445 PyObject *dtype_size = PyDict_New();
1446 PyDict_SetItem(dtype_size,
1447 (PyObject*)&PyInt_Type,
1448 PyInt_FromSize_t(sizeof(long)));
1449 PyDict_SetItem(dtype_size,
1450 (PyObject*)&PyFloat_Type,
1451 PyInt_FromSize_t(sizeof(double)));
1452 PyDict_SetItem(dtype_size,
1453 (PyObject*)&PyComplex_Type,
1454 PyInt_FromSize_t(sizeof(Complex)));
1455 PyModule_AddObject(m, "dtype_size", dtype_size);
1456
1457 // We never release these references but this is not a problem. The Python
1458 // interpreter does the same, see try_complex_special_method in
1459 // complexobject.c
1460 int_str = PyString_InternFromString("__int__");
1461 if (int_str == 0) return MOD_ERROR_VAL;
1462 long_str = PyString_InternFromString("__long__");
1463 if (long_str == 0) return MOD_ERROR_VAL;
1464 float_str = PyString_InternFromString("__float__");
1465 if (float_str == 0) return MOD_ERROR_VAL;
1466 complex_str = PyString_InternFromString("__complex__");
1467 if (complex_str == 0) return MOD_ERROR_VAL;
1468 index_str = PyString_InternFromString("__index__");
1469 if (complex_str == 0) return MOD_ERROR_VAL;
1470
1471 return MOD_SUCCESS_VAL(m);
1472 }
1473
load_index_seq_as_long(PyObject * obj,long * out,int maxlen)1474 int load_index_seq_as_long(PyObject *obj, long *out, int maxlen)
1475 {
1476 assert(maxlen >= 1);
1477 int len;
1478 if (PySequence_Check(obj)) {
1479 // get a new reference -- don't forget to DECREF on all codepaths
1480 obj = PySequence_Fast(obj, "Bug in tinyarray, load_index_seq_as_long");
1481 if (!obj) return -1;
1482 Py_ssize_t long_len = PySequence_Fast_GET_SIZE(obj);
1483 if (long_len > maxlen) {
1484 PyErr_Format(PyExc_ValueError, "Sequence too long."
1485 " Maximum length is %d.", maxlen);
1486 goto fail;
1487 }
1488 len = static_cast<int>(long_len);
1489 for (PyObject **p = PySequence_Fast_ITEMS(obj), **e = p + len; p < e;
1490 ++p, ++out) {
1491 PyObject *index = PyNumber_Index(*p);
1492 if (index == 0) goto fail;
1493 *out = PyInt_AsLong(index);
1494 Py_DECREF(index);
1495 if (*out == -1 && PyErr_Occurred()) goto fail;
1496 }
1497 Py_DECREF(obj);
1498 } else {
1499 len = 1;
1500 *out = PyInt_AsLong(obj);
1501 if (*out == -1 && PyErr_Occurred()) return -1;
1502 }
1503 return len;
1504
1505 fail:
1506 Py_DECREF(obj);
1507 return -1;
1508 }
1509
load_index_seq_as_ulong(PyObject * obj,unsigned long * uout,int maxlen,const char * errmsg)1510 int load_index_seq_as_ulong(PyObject *obj, unsigned long *uout,
1511 int maxlen, const char *errmsg)
1512 {
1513 long *out = reinterpret_cast<long*>(uout);
1514 int len = load_index_seq_as_long(obj, out, maxlen);
1515 if (len == -1) return -1;
1516 for (int i = 0; i < len; ++i)
1517 if (out[i] < 0) {
1518 if (errmsg == 0)
1519 errmsg = "Sequence may not contain negative values.";
1520 PyErr_SetString(PyExc_ValueError, errmsg);
1521 return -1;
1522 }
1523 return len;
1524 }
1525
1526 // If *dtype == NONE the simplest fitting dtype (at least dtype_min)
1527 // will be used and written back to *dtype. Any other value of *dtype requests
1528 // an array of the given dtype.
array_from_arraylike(PyObject * in,Dtype * dtype,Dtype dtype_min,bool as_matrix)1529 PyObject *array_from_arraylike(PyObject *in, Dtype *dtype, Dtype dtype_min,
1530 bool as_matrix)
1531 {
1532 Dtype dtype_in = get_dtype(in), dt = *dtype;
1533 int ndim;
1534 size_t shape[max_ndim];
1535 PyObject *seqs[max_ndim], *result;
1536 if (dtype_in != NONE) {
1537 // `in` is already an array.
1538 if (dt == NONE)
1539 dt = Dtype(std::max(int(dtype_in), int(dtype_min)));
1540 if (as_matrix) {
1541 size_t *in_shape;
1542 reinterpret_cast<Array_base*>(in)->ndim_shape(&ndim, &in_shape);
1543 if (ndim == 2) {
1544 if (dt == dtype_in)
1545 Py_INCREF(result = in);
1546 else
1547 result = convert_array(dt, in, dtype_in);
1548 } else if (ndim < 2) {
1549 shape[1] = (ndim == 0) ? 1 : in_shape[0];
1550 shape[0] = 1;
1551 result = convert_array(dt, in, dtype_in, 2, shape);
1552 } else {
1553 PyErr_SetString(PyExc_ValueError,
1554 "Matrix must be 2-dimensional.");
1555 result = 0;
1556 }
1557 } else {
1558 if (dt == dtype_in)
1559 Py_INCREF(result = in);
1560 else
1561 result = convert_array(dt, in, dtype_in);
1562 }
1563
1564 *dtype = dt;
1565 return result;
1566 } else if(PySequence_Check(in)) {
1567 // `in` is not an array, but is a sequence.
1568
1569 bool find_type = (dt == NONE);
1570
1571 // Try if buffer interface is supported.
1572 Py_buffer view;
1573 if (examine_buffer(in, &view, find_type ? &dt : 0) == 0) {
1574 if (find_type && int(dt) < int(dtype_min)) dt = dtype_min;
1575 for (int i = 0; i < view.ndim; i++)
1576 shape[i] = view.shape[i];
1577 if (as_matrix && view.ndim != 2) {
1578 if (view.ndim > 2) {
1579 PyErr_SetString(PyExc_ValueError,
1580 "Matrix must be 2-dimensional.");
1581 return 0;
1582 }
1583 shape[1] = (view.ndim == 0) ? 1 : shape[0];
1584 shape[0] = 1;
1585 }
1586 result = make_and_readin_buffer_dtable[int(dt)](
1587 &view, (as_matrix ? 2 : view.ndim), shape);
1588 PyBuffer_Release(&view);
1589
1590 *dtype = dt;
1591 return result;
1592 }
1593
1594 if (examine_sequence(in, &ndim, shape, seqs,
1595 find_type ? &dt : 0) == 0) {
1596 if (as_matrix && ndim != 2) {
1597 if (ndim > 2) {
1598 for (int d = 0; d < ndim; ++d) Py_DECREF(seqs[d]);
1599 PyErr_SetString(PyExc_ValueError,
1600 "Matrix must be 2-dimensional.");
1601 return 0;
1602 }
1603 shape[1] = (ndim == 0) ? 1 : shape[0];
1604 shape[0] = 1;
1605 }
1606 if (find_type) {
1607 PyObject *seqs_copy[max_ndim];
1608 for (int d = 0; d < ndim; ++d)
1609 Py_INCREF(seqs_copy[d] = seqs[d]);
1610 if (dt == NONE) {
1611 assert(shape[(as_matrix ? 2 : ndim) - 1] == 0);
1612 dt = default_dtype;
1613 }
1614 if (int(dt) < int(dtype_min)) dt = dtype_min;
1615 while (true) {
1616 result = readin_seqs_into_new_dtable[int(dt)](
1617 seqs, ndim, (as_matrix ? 2 : ndim), shape, true);
1618 if (result) break;
1619 dt = Dtype(int(dt) + 1);
1620 if (dt == NONE) {
1621 result = 0;
1622 break;
1623 }
1624 PyErr_Clear();
1625 for (int d = 0; d < ndim; ++d)
1626 Py_INCREF(seqs[d] = seqs_copy[d]);
1627 }
1628 for (int d = 0; d < ndim; ++d) Py_DECREF(seqs_copy[d]);
1629 } else {
1630 // A specific dtype has been requested.
1631 result = readin_seqs_into_new_dtable[int(dt)](
1632 seqs, ndim, (as_matrix ? 2 : ndim), shape, false);
1633 }
1634
1635 *dtype = dt;
1636 return result;
1637 }
1638 } else {
1639 // `in` is a scalar, or an invalid input.
1640
1641 dtype_in = dtype_of_scalar(in);
1642 bool find_type = (dt == NONE);
1643
1644 if (dtype_in == NONE) {
1645 PyErr_SetString(PyExc_TypeError,
1646 "Expecting an arraylike object or a scalar.");
1647 return 0;
1648 }
1649
1650 if (find_type) {
1651 dt = Dtype(std::max(int(dtype_in), int(dtype_min)));
1652 result = readin_scalar_into_new_dtable[int(dt)](
1653 in, true, (as_matrix ? 2 : 0));
1654 } else {
1655 result = readin_scalar_into_new_dtable[int(dt)](
1656 in, false, (as_matrix ? 2 : 0));
1657 }
1658
1659 *dtype = dt;
1660 return result;
1661 }
1662
1663 return 0;
1664 }
1665
coerce_to_arrays(PyObject ** a_,PyObject ** b_,Dtype * coerced_dtype)1666 int coerce_to_arrays(PyObject **a_, PyObject **b_, Dtype *coerced_dtype)
1667 {
1668 PyObject *a = *a_, *b = *b_;
1669
1670 // Make sure a and b are tinyarrays.
1671 Dtype dtype_a = NONE, dtype_b = NONE, dtype;
1672 a = array_from_arraylike(a, &dtype_a);
1673 if (!a) return -1;
1674 b = array_from_arraylike(b, &dtype_b, dtype_a);
1675 if (!b) {
1676 Py_DECREF(a);
1677 return -1;
1678 }
1679
1680 // Promote to a common dtype.
1681 dtype = Dtype(std::max(int(dtype_a), int(dtype_b)));
1682 if (dtype_a != dtype) {
1683 PyObject *temp = convert_array(dtype, a, dtype_a);
1684 if (temp == 0) goto fail;
1685 Py_DECREF(a);
1686 a = temp;
1687 } else if (dtype_b != dtype) {
1688 PyObject *temp = convert_array(dtype, b, dtype_b);
1689 if (temp == 0) goto fail;
1690 Py_DECREF(b);
1691 b = temp;
1692 }
1693
1694 // Success
1695 *a_ = a; *b_ = b; *coerced_dtype = dtype;
1696 return 0;
1697
1698 fail:
1699 Py_DECREF(a);
1700 Py_DECREF(b);
1701 return -1;
1702 }
1703
1704 template <typename T>
transpose(PyObject * in_,PyObject *)1705 PyObject *transpose(PyObject *in_, PyObject *)
1706 {
1707 assert(Array<T>::check_exact(in_)); Array<T> *in = (Array<T>*)in_;
1708
1709 int ndim;
1710 ptrdiff_t hops[max_ndim];
1711 size_t *shape_in, shape_out[max_ndim], stride = 1;
1712 in->ndim_shape(&ndim, &shape_in);
1713 if (ndim == 0) {
1714 Py_INCREF(in_);
1715 return in_;
1716 }
1717 for (int id = ndim - 1, od = 0; id >= 0; --id, ++od) {
1718 size_t ext = shape_in[id];
1719 shape_out[od] = ext;
1720 hops[od] = stride;
1721 stride *= ext;
1722 }
1723 for (int d = 1; d < ndim; ++d) hops[d - 1] -= hops[d] * shape_out[d];
1724 Array<T> *out = Array<T>::make(ndim, shape_out);
1725 if (!out) return 0;
1726 T *src = in->data(), *dest = out->data();
1727
1728 int d = 0;
1729 size_t i[max_ndim];
1730 --ndim;
1731 i[0] = shape_out[0];
1732 while (true) {
1733 if (i[d]) {
1734 --i[d];
1735 if (d == ndim) {
1736 *dest++ = *src;
1737 src += hops[d];
1738 } else {
1739 ++d;
1740 i[d] = shape_out[d];
1741 }
1742 } else {
1743 if (d == 0) return (PyObject*)out;
1744 --d;
1745 src += hops[d];
1746 }
1747 }
1748 }
1749
1750 template <typename T>
conjugate(PyObject * in_,PyObject *)1751 PyObject *conjugate(PyObject *in_, PyObject *)
1752 {
1753 return apply_unary_ufunc<Conjugate<T> >(in_);
1754 }
1755
1756 template <typename T>
size_of(PyObject * in_,PyObject *)1757 PyObject *size_of(PyObject *in_, PyObject *)
1758 {
1759 assert(Array<T>::check_exact(in_));
1760 return PyInt_FromSsize_t(((const Array<T>*) in_)->object_size());
1761 }
1762
1763 template <typename T>
make(int ndim,size_t size)1764 Array<T> *Array<T>::make(int ndim, size_t size)
1765 {
1766 Py_ssize_t ob_size = size;
1767 assert(ndim != 0 || size == 1);
1768 if (ndim > 1)
1769 ob_size += (ndim * sizeof(size_t) + sizeof(T) - 1) / sizeof(T);
1770 Array *result = PyObject_NewVar(Array<T>, &Array<T>::pytype, ob_size);
1771 if (!result) return 0;
1772 if (ndim > 1)
1773 result->ob_base.ob_size = -ndim;
1774 else if (ndim == 0)
1775 result->ob_base.ob_size = -1;
1776 return result;
1777 }
1778
1779 template <typename T>
make(int ndim,const size_t * shape,size_t * sizep)1780 Array<T> *Array<T>::make(int ndim, const size_t *shape, size_t *sizep)
1781 {
1782 // Check shape and calculate size, the total number of elements.
1783 size_t size = 1;
1784 // `reserve' allows to detect overflow.
1785 size_t reserve = PY_SSIZE_T_MAX;
1786 for (int d = 0; d < ndim; ++d) {
1787 size_t elem = shape[d];
1788 if (elem > reserve) {
1789 PyErr_SetString(PyExc_ValueError, "Array would be too big.");
1790 return 0;
1791 }
1792 size *= elem;
1793 if (elem) reserve /= elem;
1794 }
1795
1796 Array *result = Array<T>::make(ndim, size);
1797 if (!result) return 0;
1798 size_t *result_shape;
1799 result->ndim_shape(0, &result_shape);
1800 for (int d = 0; d < ndim; ++d) result_shape[d] = shape[d];
1801
1802 if (sizep) *sizep = size;
1803 return result;
1804 }
1805
1806
1807 template <typename T>
reduce(PyObject * self_,PyObject *)1808 PyObject *reduce(PyObject *self_, PyObject*)
1809 {
1810 assert(Array<T>::check_exact(self_)); Array<T> *self = (Array<T>*)self_;
1811 PyObject *result = PyTuple_New(2);
1812 if (!result) return 0;
1813
1814 int ndim;
1815 size_t *shape;
1816 self->ndim_shape(&ndim, &shape);
1817 size_t size_in_bytes = calc_size(ndim, shape) * sizeof(T);
1818
1819 PyObject *pyshape = PyTuple_New(ndim);
1820 for (int i=0; i < ndim; ++i)
1821 PyTuple_SET_ITEM(pyshape, i, PyInt_FromSize_t(shape[i]));
1822 PyObject *format = PyInt_FromLong(format_by_dtype[int(get_dtype(self_))]);
1823 PyObject *data = PyBytes_FromStringAndSize((char*)self->data(),
1824 size_in_bytes);
1825
1826 // PyTuple_SET_ITEM steals references, so we need to INCREF
1827 Py_INCREF(reconstruct);
1828 PyTuple_SET_ITEM(result, 0, reconstruct);
1829 // Py_BuildValue does not steal references, so we need to DECREF
1830 PyTuple_SET_ITEM(result, 1, Py_BuildValue("(OOO)", pyshape, format, data));
1831 Py_DECREF(pyshape);
1832 Py_DECREF(format);
1833 Py_DECREF(data);
1834
1835 return result;
1836 }
1837
1838 // **************** Type object ****************
1839
1840 template <>
1841 const char *Array<long>::pyname = "tinyarray.ndarray_int";
1842 template <>
1843 const char *Array<double>::pyname = "tinyarray.ndarray_float";
1844 template <>
1845 const char *Array<Complex>::pyname = "tinyarray.ndarray_complex";
1846
1847 template <typename T>
1848 PySequenceMethods Array<T>::as_sequence = {
1849 (lenfunc)len, // sq_length
1850 0, // sq_concat
1851 0, // sq_repeat
1852 seq_getitem<T> // sq_item
1853 };
1854
1855 template <typename T>
1856 PyMappingMethods Array<T>::as_mapping = {
1857 (lenfunc)len, // mp_length
1858 getitem<T> // mp_subscript
1859 };
1860
1861 #if PY_MAJOR_VERSION >= 3
1862 template <typename T>
1863 PyBufferProcs Array<T>::as_buffer = {
1864 getbuffer<T>, // bf_getbuffer
1865 0, // bf_releasebuffer
1866 };
1867 #else
1868 template <typename T>
1869 PyBufferProcs Array<T>::as_buffer = {
1870 // We only support the new buffer protocol.
1871 0, // bf_getreadbuffer
1872 0, // bf_getwritebuffer
1873 0, // bf_getsegcount
1874 0, // bf_getcharbuffer
1875 getbuffer<T> // bf_getbuffer
1876 };
1877 #endif
1878
1879 template <typename T>
1880 PyMethodDef Array<T>::methods[] = {
1881 {"transpose", (PyCFunction)transpose<T>, METH_NOARGS},
1882 {"conjugate", (PyCFunction)conjugate<T>, METH_NOARGS},
1883 {"__sizeof__", (PyCFunction)size_of<T>, METH_NOARGS},
1884 {"__reduce__", (PyCFunction)reduce<T>, METH_NOARGS},
1885 {0, 0} // Sentinel
1886 };
1887
1888 #if PY_MAJOR_VERSION >= 3 // don't need flags for buffers or checking types
1889 static unsigned long _tp_flags = Py_TPFLAGS_DEFAULT;
1890 #else
1891 static long _tp_flags = Py_TPFLAGS_DEFAULT |
1892 Py_TPFLAGS_HAVE_NEWBUFFER |
1893 Py_TPFLAGS_CHECKTYPES;
1894 #endif
1895
1896 template <typename T>
1897 PyTypeObject Array<T>::pytype = {
1898 PyVarObject_HEAD_INIT(&PyType_Type, 0)
1899 pyname,
1900 sizeof(Array_base), // tp_basicsize
1901 sizeof(T), // tp_itemsize
1902 (destructor)PyObject_Del, // tp_dealloc
1903 0, // tp_print
1904 0, // tp_getattr
1905 0, // tp_setattr
1906 0, // tp_compare (tp_reserved in Python 3.x)
1907 repr<T>, // tp_repr
1908 &as_number, // tp_as_number
1909 &as_sequence, // tp_as_sequence
1910 &as_mapping, // tp_as_mapping
1911 hash<T>, // tp_hash
1912 0, // tp_call
1913 str<T>, // tp_str
1914 PyObject_GenericGetAttr, // tp_getattro
1915 0, // tp_setattro
1916 &as_buffer, // tp_as_buffer
1917 _tp_flags, // tp_flags
1918 0, // tp_doc
1919 0, // tp_traverse
1920 0, // tp_clear
1921 (richcmpfunc)richcompare, // tp_richcompare
1922 0, // tp_weaklistoffset
1923 (getiterfunc)Array_iter<T>::make, // tp_iter
1924 0, // tp_iternext
1925 methods, // tp_methods
1926 0, // tp_members
1927 getset, // tp_getset
1928 0, // tp_base
1929 0, // tp_dict
1930 0, // tp_descr_get
1931 0, // tp_descr_set
1932 0, // tp_dictoffset
1933 0, // tp_init
1934 0, // tp_alloc
1935 0, // tp_new
1936 PyObject_Del // tp_free
1937 };
1938
1939 // **************** Explicit instantiations ****************
1940 template class Array<long>;
1941 template class Array<double>;
1942 template class Array<Complex>;
1943
1944 template PyObject *transpose<long>(PyObject*, PyObject*);
1945 template PyObject *transpose<double>(PyObject*, PyObject*);
1946 template PyObject *transpose<Complex>(PyObject*, PyObject*);
1947
1948 template PyObject *conjugate<long>(PyObject*, PyObject*);
1949 template PyObject *conjugate<double>(PyObject*, PyObject*);
1950 template PyObject *conjugate<Complex>(PyObject*, PyObject*);
1951
1952 template PyObject *size_of<long>(PyObject*, PyObject*);
1953 template PyObject *size_of<double>(PyObject*, PyObject*);
1954 template PyObject *size_of<Complex>(PyObject*, PyObject*);
1955
1956 template PyObject *reduce<long>(PyObject*, PyObject*);
1957 template PyObject *reduce<double>(PyObject*, PyObject*);
1958 template PyObject *reduce<Complex>(PyObject*, PyObject*);
1959