ENH: missingdata: Make comparisons with NA return NA, raise on 'if np.NA: ...'

5 files changed, 96 insertions, 10 deletions

diff --git a/numpy/core/code_generators/numpy_api.py b/numpy/core/code_generators/numpy_api.py
index 682c8eb3d..c8a7bdb01 100644
--- a/numpy/core/code_generators/numpy_api.py
+++ b/numpy/core/code_generators/numpy_api.py
@@ -68,6 +68,8 @@ multiarray_types_api = {
     'PyTimedeltaArrType_Type':          216,
     'PyHalfArrType_Type':               217,
     'NpyIter_Type':                     218,
+    # End 1.6 API
+    'NpyNA_Type':                       286,
 }
 
 #define NPY_NUMUSERTYPES (*(int *)PyArray_API[6])
diff --git a/numpy/core/numerictypes.py b/numpy/core/numerictypes.py
index 017898cf0..6dc82f706 100644
--- a/numpy/core/numerictypes.py
+++ b/numpy/core/numerictypes.py
@@ -86,11 +86,13 @@ __all__ = ['sctypeDict', 'sctypeNA', 'typeDict', 'typeNA', 'sctypes',
            'ScalarType', 'obj2sctype', 'cast', 'nbytes', 'sctype2char',
            'maximum_sctype', 'issctype', 'typecodes', 'find_common_type',
            'issubdtype', 'datetime_data','datetime_as_string',
-           'busday_offset', 'busday_count', 'is_busday', 'busdaycalendar']
+           'busday_offset', 'busday_count', 'is_busday', 'busdaycalendar',
+           'NA', 'NAType']
 
 from numpy.core.multiarray import typeinfo, ndarray, array, \
                       empty, dtype, datetime_data, datetime_as_string, \
-                      busday_offset, busday_count, is_busday, busdaycalendar
+                      busday_offset, busday_count, is_busday, busdaycalendar, \
+                      NA, NAType
 import types as _types
 import sys
 
diff --git a/numpy/core/src/multiarray/multiarraymodule.c b/numpy/core/src/multiarray/multiarraymodule.c
index 2c1061656..9491bc631 100644
--- a/numpy/core/src/multiarray/multiarraymodule.c
+++ b/numpy/core/src/multiarray/multiarraymodule.c
@@ -32,8 +32,6 @@
 
 NPY_NO_EXPORT int NPY_NUMUSERTYPES = 0;
 
-#define PyAO PyArrayObject
-
 /* Internal APIs */
 #include "arraytypes.h"
 #include "arrayobject.h"
@@ -50,6 +48,7 @@ NPY_NO_EXPORT int NPY_NUMUSERTYPES = 0;
 #include "datetime_strings.h"
 #include "datetime_busday.h"
 #include "datetime_busdaycal.h"
+#include "na_singleton.h"
 
 /* Only here for API compatibility */
 NPY_NO_EXPORT PyTypeObject PyBigArray_Type;
@@ -2685,7 +2684,7 @@ PyArray_Where(PyObject *condition, PyObject *x, PyObject *y)
         Py_DECREF(obj);
         return NULL;
     }
-    ret = PyArray_Choose((PyAO *)obj, tup, NULL, NPY_RAISE);
+    ret = PyArray_Choose((PyArrayObject *)obj, tup, NULL, NPY_RAISE);
     Py_DECREF(obj);
     Py_DECREF(tup);
     return ret;
@@ -3851,6 +3850,9 @@ PyMODINIT_FUNC initmultiarray(void) {
     if (PyType_Ready(&NpyIter_Type) < 0) {
         return RETVAL;
     }
+    if (PyType_Ready(&NpyNA_Type) < 0) {
+        return RETVAL;
+    }
 
     PyArrayDescr_Type.tp_hash = PyArray_DescrHash;
     if (PyType_Ready(&PyArrayDescr_Type) < 0) {
@@ -3946,6 +3948,12 @@ PyMODINIT_FUNC initmultiarray(void) {
     Py_INCREF(&NpyBusDayCalendar_Type);
     PyDict_SetItemString(d, "busdaycalendar",
                             (PyObject *)&NpyBusDayCalendar_Type);
+    /* NA Type */
+    PyDict_SetItemString(d, "NAType", (PyObject *)&NpyNA_Type);
+    Py_INCREF(&NpyNA_Type);
+    /* NA  Singleton */
+    PyDict_SetItemString(d, "NA", (PyObject *)&_Npy_NASingleton);
+    Py_INCREF(&_Npy_NASingleton);
 
     set_flaginfo(d);
 
diff --git a/numpy/core/src/multiarray/na_singleton.c b/numpy/core/src/multiarray/na_singleton.c
index 73ac755ad..7f866ec19 100644
--- a/numpy/core/src/multiarray/na_singleton.c
+++ b/numpy/core/src/multiarray/na_singleton.c
@@ -71,6 +71,26 @@ na_init(NpyNA_fieldaccess *self, PyObject *args, PyObject *kwds)
     return 0;
 }
 
+/*
+ * The call function proxies to the na_init function to handle
+ * the payload and dtype parameters.
+ */
+static PyObject *
+na_call(PyObject *NPY_UNUSED(self), PyObject *args, PyObject *kwds)
+{
+    NpyNA_fieldaccess *ret;
+
+    ret = (NpyNA_fieldaccess *)na_new(&NpyNA_Type, NULL, NULL);
+    if (ret != NULL) {
+        if (na_init(ret, args, kwds) < 0) {
+            Py_DECREF(ret);
+            return NULL;
+        }
+    }
+
+    return (PyObject *)ret;
+}
+
 static void
 na_dealloc(NpyNA_fieldaccess *self)
 {
@@ -106,6 +126,31 @@ na_repr(NpyNA_fieldaccess *self)
     }
 }
 
+/*
+ * The str function is the same as repr, except it throws away
+ * the dtype. It is always either "NA" or "NA(payload)".
+ */
+static PyObject *
+na_str(NpyNA_fieldaccess *self)
+{
+    if (self->payload == 255) {
+        return PyUString_FromString("NA");
+    }
+    else {
+        return PyUString_FromFormat("NA(%d)", (int)self->payload);
+    }
+}
+
+/*
+ * Any comparison with NA produces an NA.
+ */
+static PyObject *
+na_richcompare(NpyNA_fieldaccess *self, PyObject *other, int cmp_op)
+{
+    Py_INCREF(Npy_NA);
+    return Npy_NA;
+}
+
 static PyObject *
 na_payload_get(NpyNA_fieldaccess *self)
 {
@@ -214,6 +259,33 @@ static PyGetSetDef na_getsets[] = {
     {NULL, NULL, NULL, NULL, NULL}
 };
 
+/* Using NA in an if statement is always an error */
+static int
+na_nonzero(PyObject *NPY_UNUSED(self))
+{
+    PyErr_SetString(PyExc_ValueError,
+            "numpy.NA represents an unknown missing value, "
+            "so its truth value cannot be determined");
+    return -1;
+}
+
+NPY_NO_EXPORT PyNumberMethods na_as_number = {
+    0,                                          /*nb_add*/
+    0,                                          /*nb_subtract*/
+    0,                                          /*nb_multiply*/
+#if defined(NPY_PY3K)
+#else
+    0,                                          /*nb_divide*/
+#endif
+    0,                                          /*nb_remainder*/
+    0,                                          /*nb_divmod*/
+    0,                                          /*nb_power*/
+    0,                                          /*nb_neg*/
+    0,                                          /*nb_pos*/
+    0,                                          /*(unaryfunc)array_abs,*/
+    (inquiry)na_nonzero,                        /*nb_nonzero*/
+};
+
 NPY_NO_EXPORT PyTypeObject NpyNA_Type = {
 #if defined(NPY_PY3K)
     PyVarObject_HEAD_INIT(NULL, 0)
@@ -235,12 +307,12 @@ NPY_NO_EXPORT PyTypeObject NpyNA_Type = {
     0,                                          /* tp_compare */
 #endif
     (reprfunc)na_repr,                          /* tp_repr */
-    0,                                          /* tp_as_number */
+    &na_as_number,                              /* tp_as_number */
     0,                                          /* tp_as_sequence */
     0,                                          /* tp_as_mapping */
     0,                                          /* tp_hash */
-    0,                                          /* tp_call */
-    0,                                          /* tp_str */
+    (ternaryfunc)na_call,                       /* tp_call */
+    (reprfunc)na_str,                           /* tp_str */
     0,                                          /* tp_getattro */
     0,                                          /* tp_setattro */
     0,                                          /* tp_as_buffer */
@@ -248,7 +320,7 @@ NPY_NO_EXPORT PyTypeObject NpyNA_Type = {
     0,                                          /* tp_doc */
     0,                                          /* tp_traverse */
     0,                                          /* tp_clear */
-    0,                                          /* tp_richcompare */
+    (richcmpfunc)na_richcompare,                /* tp_richcompare */
     0,                                          /* tp_weaklistoffset */
     0,                                          /* tp_iter */
     0,                                          /* tp_iternext */
@@ -278,7 +350,7 @@ NPY_NO_EXPORT PyTypeObject NpyNA_Type = {
 
 NPY_NO_EXPORT NpyNA_fieldaccess _Npy_NASingleton = {
     PyObject_HEAD_INIT(&NpyNA_Type)
-    0,    /* payload */
+    255,  /* payload (255 means no payload) */
     NULL, /* dtype */
     1     /* is_singleton */
 };
diff --git a/numpy/core/src/multiarray/na_singleton.h b/numpy/core/src/multiarray/na_singleton.h
index d9dd4e9d3..a0a36430a 100644
--- a/numpy/core/src/multiarray/na_singleton.h
+++ b/numpy/core/src/multiarray/na_singleton.h
@@ -14,4 +14,6 @@ typedef struct {
 
 NPY_NO_EXPORT NpyNA_fieldaccess _Npy_NASingleton;
 
+#define Npy_NA ((PyObject *)&_Npy_NASingleton)
+
 #endif