Some suggested changes

astropy · Apr 10, 2019 · b80fce8 · b80fce8
1 parent cb7972c
commit b80fce8
Show file tree

Hide file tree

Showing 3 changed files with 42 additions and 46 deletions.
diff --git a/CHANGES.rst b/CHANGES.rst
@@ -313,6 +313,11 @@ astropy.visualization
 astropy.wcs
 ^^^^^^^^^^^
 
+- Added a ``PyUnitListProxy_richcmp`` method in ``UnitListProxy`` class to enable 
+  ``WCS.wcs.cunit`` equality testing. It helps to check whether the two instances of 
+  ``WCS.wcs.cunit`` are equal or not by comparing the data members of 
+  ``UnitListProxy`` class [#8480]
+
 Other Changes and Additions
 ---------------------------
 

diff --git a/astropy/wcs/src/unit_list_proxy.c b/astropy/wcs/src/unit_list_proxy.c
@@ -21,7 +21,7 @@ typedef struct {
   PyObject_HEAD
   /*@null@*/ /*@shared@*/ PyObject* pyobject;
   Py_ssize_t size;
-  char (*array)[72];
+  char (*array)[ARRAYSIZE];
   PyObject* unit_class;
 } PyUnitListProxy;
 
@@ -95,7 +95,7 @@ PyUnitListProxy_clear(
 PyUnitListProxy_New(
     /*@shared@*/ PyObject* owner,
     Py_ssize_t size,
-    char (*array)[72]) {
+    char (*array)[ARRAYSIZE]) {
 
   PyUnitListProxy* self = NULL;
   PyObject *units_module;
@@ -187,51 +187,34 @@ PyUnitListProxy_getitem(
   return result;
 }
 
-static PyObject *
+static PyObject*
 PyUnitListProxy_richcmp(
-    PyObject *a,
-    PyObject *b,
-    int op) {
-
-  int equal;
-  int status;
-  PyUnitListProxy *ax,*bx;
-  if ((op == Py_EQ || op == Py_NE) &&
-      PyObject_TypeCheck(b, &PyUnitListProxyType) &&
-      PyObject_TypeCheck(a, &PyUnitListProxyType)) {
-      ax = (PyUnitListProxy *)a;
-      bx = (PyUnitListProxy *)b;
-    if (!PyObject_RichCompareBool(ax->unit_class, bx->unit_class,Py_EQ)||
-        ax->size != bx->size||
-        strncmp(ax->array, bx->array, ARRAYSIZE) ) {
-        status = 0;
-        equal = 0;
-    } else if (ax == NULL || bx == NULL ||
-    	       PyObject_RichCompareBool(ax->unit_class, bx->unit_class,Py_EQ) == -1) {
-            status = 1;
-    } else {
-       status = 0;
-       equal = 1;
-    }
-
-    if (status == 0) {
-      if (op == Py_NE) {
-        equal = !equal;
-      }
-      if (equal) {
-        Py_RETURN_TRUE;
-      } else {
-        Py_RETURN_FALSE;
-      }
-    } 
-    if (status == 1){
-    	PyErr_SetString(PyExc_TypeError,
-                        "Some Internal Errors occured during comparison or elements are NULL");
-    	return NULL;
-    }
+	PyObject *a,
+	PyObject *b,
+	int op){
+  PyUnitListProxy *lhs, *rhs;
+  assert(a != NULL && b != NULL);
+  if (!PyObject_TypeCheck(a, &PyUnitListProxyType) ||
+      !PyObject_TypeCheck(b, &PyUnitListProxyType)) {
+    Py_RETURN_NOTIMPLEMENTED;
   }
+  if (op != Py_EQ && op != Py_NE) {
     Py_RETURN_NOTIMPLEMENTED;
-
+  }
+  lhs = (PyUnitListProxy *)a;
+  rhs = (PyUnitListProxy *)b;
+  int equal = PyObject_RichCompareBool(lhs->unit_class, rhs->unit_class, Py_EQ);
+  if (equal == -1) {
+    return NULL;  // Exception will be set because the rich-compare failed
+  }
+  equal = equal == 1 && !strncmp(lhs->array, rhs->array, ARRAYSIZE) && lhs->size == rhs->size;
+  if ((op == Py_EQ && equal == 1) ||
+      (op == Py_NE && equal == 0)) {
+    Py_RETURN_TRUE;
+  } 
+  else {
+    Py_RETURN_FALSE;
+  }
 }
 
 static int
@@ -346,7 +329,7 @@ set_unit_list(
     const char* propname,
     PyObject* value,
     Py_ssize_t len,
-    char (*dest)[72]) {
+    char (*dest)[ARRAYSIZE]) {
 
   PyObject*  unit  = NULL;
   PyObject*  proxy = NULL;

diff --git a/astropy/wcs/tests/test_wcs.py b/astropy/wcs/tests/test_wcs.py
@@ -1186,7 +1186,7 @@ def test_footprint_contains():
 
 
 def test_cunit():
-    # Initializinig WCS
+    # Initializing WCS
     w1 = wcs.WCS(naxis=2)
     w2 = wcs.WCS(naxis=2)
     w3 = wcs.WCS(naxis=2)
@@ -1195,6 +1195,14 @@ def test_cunit():
     w2.wcs.cunit = ['km/h', 'km/h']
     w3.wcs.cunit = ['deg', 'm/s']
 
+    # Equality checking a cunit with itself
     assert w1.wcs.cunit == w1.wcs.cunit
+    # Equality checking of two different cunit object having same values
     assert w1.wcs.cunit == w3.wcs.cunit
+    # Inequality checking of two different cunit object having different values
     assert w1.wcs.cunit != w2.wcs.cunit
+    # Inequality checking of cunit with a list of literals
+    assert w1.wcs.cunit != [1, 2, 3]
+    # Comparison is not implemented TypeError will raise
+    with pytest.raises(TypeError):
+        w1.wcs.cunit < w2.wcs.cunit