ENH: missingdata: Add the NA mask members to PyArrayObject

10 files changed, 212 insertions, 90 deletions

diff --git a/numpy/core/SConscript b/numpy/core/SConscript
index 5bf62ea11..80a7d08de 100644
--- a/numpy/core/SConscript
+++ b/numpy/core/SConscript
@@ -452,6 +452,7 @@ if ENABLE_SEPARATE_COMPILATION:
         pjoin('src', 'multiarray', 'descriptor.c'),
         pjoin('src', 'multiarray', 'iterators.c'),
         pjoin('src', 'multiarray', 'mapping.c'),
+        pjoin('src', 'multiarray', 'na_mask.c'),
         pjoin('src', 'multiarray', 'number.c'),
         pjoin('src', 'multiarray', 'getset.c'),
         pjoin('src', 'multiarray', 'sequence.c'),
diff --git a/numpy/core/code_generators/genapi.py b/numpy/core/code_generators/genapi.py
index 29cbb271f..b083f6e77 100644
--- a/numpy/core/code_generators/genapi.py
+++ b/numpy/core/code_generators/genapi.py
@@ -52,6 +52,7 @@ API_FILES = [join('multiarray', 'methods.c'),
              join('multiarray', 'nditer_pywrap.c'),
              join('multiarray', 'nditer_templ.c.src'),
              join('multiarray', 'einsum.c.src'),
+             join('multiarray', 'na_mask.c'),
              join('umath', 'ufunc_object.c'),
              join('umath', 'ufunc_type_resolution.c'),
              join('umath', 'loops.c.src'),
diff --git a/numpy/core/code_generators/numpy_api.py b/numpy/core/code_generators/numpy_api.py
index a256d849a..682c8eb3d 100644
--- a/numpy/core/code_generators/numpy_api.py
+++ b/numpy/core/code_generators/numpy_api.py
@@ -318,7 +318,9 @@ multiarray_funcs_api = {
     # End 1.6 API
     'PyArray_MaskedCopyInto':               281,
     'PyArray_MaskedMoveInto':               282,
-    'PyArray_SetBaseObject':                      283,
+    'PyArray_SetBaseObject':                283,
+    'PyArray_HasNASupport':                 284,
+    'PyArray_ContainsNA':                   285,
 }
 
 ufunc_types_api = {
diff --git a/numpy/core/include/numpy/ndarraytypes.h b/numpy/core/include/numpy/ndarraytypes.h
index 94c45df88..85cc8a379 100644
--- a/numpy/core/include/numpy/ndarraytypes.h
+++ b/numpy/core/include/numpy/ndarraytypes.h
@@ -287,6 +287,40 @@ typedef enum {
     NPY_BUSDAY_RAISE
 } NPY_BUSDAY_ROLL;
 
+/*********************************************************************
+ * NumPy functions for dealing with masks, such as in masked iteration
+ *********************************************************************/
+
+typedef npy_uint8 npy_mask;
+#define NPY_MASK NPY_UINT8
+
+/*
+ * Bit 0 of the mask indicates whether a value is exposed
+ * or hidden. This is compatible with a 'where=' boolean
+ * mask, because NumPy booleans are 1 byte, and contain
+ * either the value 0 or 1.
+ */
+static NPY_INLINE npy_bool
+NpyMaskValue_IsExposed(npy_mask mask)
+{
+    return (mask & 0x01) != 0;
+}
+
+/*
+ * Bits 1 through 7 of the mask contain the payload.
+ */
+static NPY_INLINE npy_uint8
+NpyMaskValue_GetPayload(npy_mask mask)
+{
+    return ((npy_uint8)mask) >> 1;
+}
+
+static NPY_INLINE npy_mask
+NpyMaskValue_Create(npy_bool exposed, npy_uint8 payload)
+{
+    return (npy_mask)(exposed != 0) | (npy_mask)(payload << 1);
+}
+
 
 #define NPY_ERR(str) fprintf(stderr, #str); fflush(stderr);
 #define NPY_ERR2(str) fprintf(stderr, str); fflush(stderr);
@@ -583,32 +617,64 @@ typedef struct _arr_descr {
  */
 /* This struct will be moved to a private header in a future release */
 typedef struct tagPyArrayObject_fieldaccess {
-        PyObject_HEAD
-        char *data;             /* pointer to raw data buffer */
-        int nd;                 /* number of dimensions, also called ndim */
-        npy_intp *dimensions;   /* size in each dimension */
-        npy_intp *strides;      /*
-                                 * bytes to jump to get to the
-                                 * next element in each dimension
-                                 */
-        PyObject *base;         /*
-                                 * This object should be decref'd upon
-                                 * deletion of array
-                                 *
-                                 * For views it points to the original
-                                 * array
-                                 *
-                                 * For creation from buffer object it
-                                 * points to an object that shold be
-                                 * decref'd on deletion
-                                 *
-                                 * For UPDATEIFCOPY flag this is an
-                                 * array to-be-updated upon deletion
-                                 * of this one
-                                 */
-        PyArray_Descr *descr;   /* Pointer to type structure */
-        int flags;              /* Flags describing array -- see below */
-        PyObject *weakreflist;  /* For weakreferences */
+    PyObject_HEAD
+    /* Pointer to the raw data buffer */
+    char *data;
+    /* The number of dimensions, also called 'ndim' */
+    int nd;
+    /* The size in each dimension, also called 'shape' */
+    npy_intp *dimensions;
+    /*
+     * Number of bytes to jump to get to the
+     * next element in each dimension
+     */
+    npy_intp *strides;
+    /*
+     * This object is decref'd upon
+     * deletion of array. Except in the
+     * case of UPDATEIFCOPY which has
+     * special handling.
+     *
+     * For views it points to the original
+     * array, collapsed so no chains of
+     * views occur.
+     *
+     * For creation from buffer object it
+     * points to an object that shold be
+     * decref'd on deletion
+     *
+     * For UPDATEIFCOPY flag this is an
+     * array to-be-updated upon deletion
+     * of this one
+     */
+    PyObject *base;
+    /* Pointer to type structure */
+    PyArray_Descr *descr;
+    /* Flags describing array -- see below */
+    int flags;
+    /* For weak references */
+    PyObject *weakreflist;
+
+    /* New fields added as of NumPy 1.7 */
+
+    /*
+     * Descriptor for the mask dtype.
+     *   If no mask: NULL
+     *   If mask   : bool/uint8/structured dtype of mask dtypes
+     */
+    PyArray_Descr *maskna_descr;
+    /*
+     * Raw data buffer for mask. If the array has the flag
+     * NPY_ARRAY_OWNMASKNA enabled, it owns this memory and
+     * must call PyArray_free on it when destroyed.
+     */
+    npy_mask *maskna_data;
+    /*
+     * Just like dimensions and strides point into the same memory
+     * buffer, we now just make that buffer 3x the nd instead of 2x
+     * and use the same buffer.
+     */
+    npy_intp *maskna_strides;
 } PyArrayObject_fieldaccess;
 
 /*
@@ -688,7 +754,7 @@ typedef int (PyArray_FinalizeFunc)(PyArrayObject *, PyObject *);
  * the fastest in memory (strides array is reverse of C-contiguous
  * array)
  */
-#define NPY_ARRAY_F_CONTIGUOUS       0x0002
+#define NPY_ARRAY_F_CONTIGUOUS    0x0002
 
 /*
  * Note: all 0-d arrays are C_CONTIGUOUS and F_CONTIGUOUS. If a
@@ -699,7 +765,7 @@ typedef int (PyArray_FinalizeFunc)(PyArrayObject *, PyObject *);
  * If set, the array owns the data: it will be free'd when the array
  * is deleted.
  */
-#define NPY_ARRAY_OWNDATA       0x0004
+#define NPY_ARRAY_OWNDATA         0x0004
 
 /*
  * An array never has the next four set; they're only used as parameter
@@ -707,22 +773,22 @@ typedef int (PyArray_FinalizeFunc)(PyArrayObject *, PyObject *);
  */
 
 /* Cause a cast to occur regardless of whether or not it is safe. */
-#define NPY_ARRAY_FORCECAST     0x0010
+#define NPY_ARRAY_FORCECAST       0x0010
 
 /*
  * Always copy the array. Returned arrays are always CONTIGUOUS,
  * ALIGNED, and WRITEABLE.
  */
-#define NPY_ARRAY_ENSURECOPY    0x0020
+#define NPY_ARRAY_ENSURECOPY      0x0020
 
 /* Make sure the returned array is a base-class ndarray */
-#define NPY_ARRAY_ENSUREARRAY   0x0040
+#define NPY_ARRAY_ENSUREARRAY     0x0040
 
 /*
  * Make sure that the strides are in units of the element size Needed
  * for some operations with record-arrays.
  */
-#define NPY_ARRAY_ELEMENTSTRIDES 0x0080
+#define NPY_ARRAY_ELEMENTSTRIDES  0x0080
 
 /*
  * Array data is aligned on the appropiate memory address for the type
@@ -730,20 +796,27 @@ typedef int (PyArray_FinalizeFunc)(PyArrayObject *, PyObject *);
  * array of integers (4 bytes each) starts on a memory address that's
  * a multiple of 4)
  */
-#define NPY_ARRAY_ALIGNED       0x0100
+#define NPY_ARRAY_ALIGNED         0x0100
 
 /* Array data has the native endianness */
-#define NPY_ARRAY_NOTSWAPPED    0x0200
+#define NPY_ARRAY_NOTSWAPPED      0x0200
 
 /* Array data is writeable */
-#define NPY_ARRAY_WRITEABLE     0x0400
+#define NPY_ARRAY_WRITEABLE       0x0400
 
 /*
  * If this flag is set, then base contains a pointer to an array of
  * the same size that should be updated with the current contents of
  * this array when this array is deallocated
  */
-#define NPY_ARRAY_UPDATEIFCOPY  0x1000
+#define NPY_ARRAY_UPDATEIFCOPY    0x1000
+
+/*
+ * If this flag is set, then the array owns the memory for the
+ * missing values NA mask.
+ */
+#define NPY_ARRAY_OWNMASKNA       0x2000
+
 
 #define NPY_ARRAY_BEHAVED      (NPY_ARRAY_ALIGNED | \
                                 NPY_ARRAY_WRITEABLE)
@@ -772,7 +845,7 @@ typedef int (PyArray_FinalizeFunc)(PyArrayObject *, PyObject *);
                                 NPY_ARRAY_F_CONTIGUOUS | \
                                 NPY_ARRAY_ALIGNED)
 
-/* This flag is for the array interface */
+/* This flag is for the array interface, not PyArrayObject */
 #define NPY_ARR_HAS_DESCR  0x0800
 
 
@@ -845,9 +918,9 @@ typedef int (PyArray_FinalizeFunc)(PyArrayObject *, PyObject *);
 #define NPY_DISABLE_C_API
 #endif
 
-/*****************************
- * New iterator object
- *****************************/
+/**********************************
+ * The nditer object, added in 1.6
+ **********************************/
 
 /* The actual structure of the iterator is an internal detail */
 typedef struct NpyIter_InternalOnly NpyIter;
@@ -1278,8 +1351,6 @@ PyArrayNeighborhoodIter_Next2D(PyArrayNeighborhoodIterObject* iter);
 #define PyArray_FORTRAN_IF(m) ((PyArray_CHKFLAGS(m, NPY_ARRAY_F_CONTIGUOUS) ? \
                                NPY_ARRAY_F_CONTIGUOUS : 0))
 
-#define FORTRAN_IF PyArray_FORTRAN_IF
-
 #ifdef NPY_NO_DEPRECATED_API
 /*
  * Changing access macros into functions, to allow for future hiding
@@ -1364,17 +1435,14 @@ static NPY_INLINE PyObject *
 PyArray_GETITEM(PyArrayObject *arr, char *itemptr)
 {
     return ((PyArrayObject_fieldaccess *)arr)->descr->f->getitem(
-                                                itemptr,
-                                                arr);
+                                                        itemptr, arr);
 }
 
 static NPY_INLINE int
 PyArray_SETITEM(PyArrayObject *arr, char *itemptr, PyObject *v)
 {
     return ((PyArrayObject_fieldaccess *)arr)->descr->f->setitem(
-                                                v,
-                                                itemptr,
-                                                arr);
+                                                        v, itemptr, arr);
 }
 
 /* Same as PyArray_DATA */
@@ -1382,7 +1450,7 @@ PyArray_SETITEM(PyArrayObject *arr, char *itemptr, PyObject *v)
 
 #else
 
-/* Macros are deprecated as of NumPy 1.7. */
+/* These macros are deprecated as of NumPy 1.7. */
 #define PyArray_NDIM(obj) (((PyArrayObject_fieldaccess *)(obj))->nd)
 #define PyArray_BYTES(obj) ((char *)(((PyArrayObject_fieldaccess *)(obj))->data))
 #define PyArray_DATA(obj) ((void *)(((PyArrayObject_fieldaccess *)(obj))->data))
@@ -1429,6 +1497,34 @@ PyArray_CLEARFLAGS(PyArrayObject *arr, int flags)
     ((PyArrayObject_fieldaccess *)arr)->flags &= ~flags;
 }
 
+/* Access to the missing values NA mask, added in 1.7 */
+
+static NPY_INLINE PyArray_Descr *
+PyArray_MASKNA_DESCR(PyArrayObject *arr)
+{
+    return ((PyArrayObject_fieldaccess *)arr)->maskna_descr;
+}
+
+static NPY_INLINE npy_mask *
+PyArray_MASKNA_DATA(PyArrayObject *arr)
+{
+    return ((PyArrayObject_fieldaccess *)arr)->maskna_data;
+}
+
+/* For the corresponding DIMS, use PyArray_DIMS(arr) */
+static NPY_INLINE npy_intp *
+PyArray_MASKNA_STRIDES(PyArrayObject *arr)
+{
+    return ((PyArrayObject_fieldaccess *)arr)->maskna_strides;
+}
+
+static NPY_INLINE npy_bool
+PyArray_HASMASKNA(PyArrayObject *arr)
+{
+    return ((PyArrayObject_fieldaccess *)arr)->maskna_data != NULL;
+}
+
+
 #define PyTypeNum_ISBOOL(type) ((type) == NPY_BOOL)
 
 #define PyTypeNum_ISUNSIGNED(type) (((type) == NPY_UBYTE) ||   \
@@ -1592,40 +1688,6 @@ struct NpyAuxData_tag {
 #define NPY_AUXDATA_CLONE(auxdata) \
     ((auxdata)->clone(auxdata))
 
-/*********************************************************************
- * NumPy functions for dealing with masks, such as in masked iteration
- *********************************************************************/
-
-typedef npy_uint8 npy_mask;
-#define NPY_MASK NPY_UINT8
-
-/*
- * Bit 0 of the mask indicates whether a value is exposed
- * or hidden. This is compatible with a 'where=' boolean
- * mask, because NumPy booleans are 1 byte, and contain
- * either the value 0 or 1.
- */
-static NPY_INLINE npy_bool
-NpyMask_IsExposed(npy_mask mask)
-{
-    return (mask & 0x01) != 0;
-}
-
-/*
- * Bits 1 through 7 of the mask contain the payload.
- */
-static NPY_INLINE npy_uint8
-NpyMask_GetPayload(npy_mask mask)
-{
-    return ((npy_uint8)mask) >> 1;
-}
-
-static NPY_INLINE npy_mask
-NpyMask_Create(npy_bool exposed, npy_uint8 payload)
-{
-    return (npy_mask)(exposed != 0) | (npy_mask)(payload << 1);
-}
-
 /*
  * This is the form of the struct that's returned pointed by the
  * PyCObject attribute of an array __array_struct__. See
diff --git a/numpy/core/include/numpy/npy_deprecated_api.h b/numpy/core/include/numpy/npy_deprecated_api.h
index 413d24d4e..a268f504a 100644
--- a/numpy/core/include/numpy/npy_deprecated_api.h
+++ b/numpy/core/include/numpy/npy_deprecated_api.h
@@ -89,4 +89,10 @@
  */
 #define fortran fortran_
 
+/*
+ * Deprecated as of NumPy 1.7, as it is a namespace-polluting
+ * macro.
+ */
+#define FORTRAN_IF PyArray_FORTRAN_IF
+
 #endif
diff --git a/numpy/core/setup.py b/numpy/core/setup.py
index 6b5b0df3b..e2cdf35e1 100644
--- a/numpy/core/setup.py
+++ b/numpy/core/setup.py
@@ -773,6 +773,7 @@ def configuration(parent_package='',top_path=None):
             join('src', 'multiarray', 'mapping.c'),
             join('src', 'multiarray', 'methods.c'),
             join('src', 'multiarray', 'multiarraymodule.c'),
+            join('src', 'multiarray', 'na_mask.c'),
             join('src', 'multiarray', 'nditer_templ.c.src'),
             join('src', 'multiarray', 'nditer_api.c'),
             join('src', 'multiarray', 'nditer_constr.c'),
diff --git a/numpy/core/src/multiarray/dtype_transfer.c b/numpy/core/src/multiarray/dtype_transfer.c
index ce688efd5..045e6c4b5 100644
--- a/numpy/core/src/multiarray/dtype_transfer.c
+++ b/numpy/core/src/multiarray/dtype_transfer.c
@@ -3077,7 +3077,7 @@ void _strided_masked_wrapper_decsrcref_transfer_function(
     while (N > 0) {
         /* Skip masked values, still calling decsrcref for move_references */
         subloopsize = 0;
-        while (subloopsize < N && !NpyMask_IsExposed(*mask)) {
+        while (subloopsize < N && !NpyMaskValue_IsExposed(*mask)) {
             ++subloopsize;
             mask += mask_stride;
         }
@@ -3088,7 +3088,7 @@ void _strided_masked_wrapper_decsrcref_transfer_function(
         N -= subloopsize;
         /* Process unmasked values */
         subloopsize = 0;
-        while (subloopsize < N && NpyMask_IsExposed(*mask)) {
+        while (subloopsize < N && NpyMaskValue_IsExposed(*mask)) {
             ++subloopsize;
             mask += mask_stride;
         }
@@ -3120,7 +3120,7 @@ void _strided_masked_wrapper_transfer_function(
     while (N > 0) {
         /* Skip masked values */
         subloopsize = 0;
-        while (subloopsize < N && !NpyMask_IsExposed(*mask)) {
+        while (subloopsize < N && !NpyMaskValue_IsExposed(*mask)) {
             ++subloopsize;
             mask += mask_stride;
         }
@@ -3129,7 +3129,7 @@ void _strided_masked_wrapper_transfer_function(
         N -= subloopsize;
         /* Process unmasked values */
         subloopsize = 0;
-        while (subloopsize < N && NpyMask_IsExposed(*mask)) {
+        while (subloopsize < N && NpyMaskValue_IsExposed(*mask)) {
             ++subloopsize;
             mask += mask_stride;
         }
diff --git a/numpy/core/src/multiarray/getset.c b/numpy/core/src/multiarray/getset.c
index e9e052683..6503d8f12 100644
--- a/numpy/core/src/multiarray/getset.c
+++ b/numpy/core/src/multiarray/getset.c
@@ -775,7 +775,7 @@ array_flat_set(PyArrayObject *self, PyObject *val)
     typecode = PyArray_DESCR(self);
     Py_INCREF(typecode);
     arr = (PyArrayObject *)PyArray_FromAny(val, typecode,
-                          0, 0, NPY_ARRAY_FORCECAST | FORTRAN_IF(self), NULL);
+                  0, 0, NPY_ARRAY_FORCECAST | PyArray_FORTRAN_IF(self), NULL);
     if (arr == NULL) {
         return -1;
     }
diff --git a/numpy/core/src/multiarray/na_mask.c b/numpy/core/src/multiarray/na_mask.c
new file mode 100644
index 000000000..ce8dd1e4b
--- /dev/null
+++ b/numpy/core/src/multiarray/na_mask.c
@@ -0,0 +1,49 @@
+/*
+ * This file implements missing value NA mask support for the NumPy array.
+ *
+ * Written by Mark Wiebe (mwwiebe@gmail.com)
+ * Copyright (c) 2011 by Enthought, Inc.
+ *
+ * See LICENSE.txt for the license.
+ */
+
+#define PY_SSIZE_T_CLEAN
+#include <Python.h>
+
+#define NPY_NO_DEPRECATED_API
+#define _MULTIARRAYMODULE
+#include <numpy/arrayobject.h>
+
+#include "npy_config.h"
+#include "numpy/npy_3kcompat.h"
+
+/*NUMPY_API
+ *
+ * Returns true if the array has an NA mask. When
+ * NA dtypes are implemented, will also return true
+ * if the array's dtype has NA support.
+ */
+NPY_NO_EXPORT npy_bool
+PyArray_HasNASupport(PyArrayObject *arr)
+{
+    return PyArray_HASMASKNA(arr);
+}
+
+/*NUMPY_API
+ *
+ * Returns false if the array has no NA support. Returns
+ * true if the array has NA support AND there is an
+ * NA anywhere in the array.
+ */
+NPY_NO_EXPORT npy_bool
+PyArray_ContainsNA(PyArrayObject *arr)
+{
+    /* Need NA support to contain NA */
+    if (!PyArray_HasNASupport(arr)) {
+        return 0;
+    }
+
+    /* TODO: Loop through NA mask */
+
+    return 0;
+}
diff --git a/numpy/core/src/umath/ufunc_type_resolution.c b/numpy/core/src/umath/ufunc_type_resolution.c
index fb7352070..df3bb19bd 100644
--- a/numpy/core/src/umath/ufunc_type_resolution.c
+++ b/numpy/core/src/umath/ufunc_type_resolution.c
@@ -1408,7 +1408,7 @@ unmasked_ufunc_loop_as_masked(
         /* Skip masked values */
         subloopsize = 0;
         while (subloopsize < loopsize &&
-                        !NpyMask_IsExposed(*(npy_mask *)mask)) {
+                        !NpyMaskValue_IsExposed(*(npy_mask *)mask)) {
             ++subloopsize;
             mask += mask_stride;
         }
@@ -1422,7 +1422,7 @@ unmasked_ufunc_loop_as_masked(
          */
         subloopsize = 0;
         while (subloopsize < loopsize &&
-                        NpyMask_IsExposed(*(npy_mask *)mask)) {
+                        NpyMaskValue_IsExposed(*(npy_mask *)mask)) {
             ++subloopsize;
             mask += mask_stride;
         }