Test for wcs.wcs.cunit is added

astropy · Apr 17, 2019 · 2d2501b · 2d2501b
1 parent 5d8e6fe
commit 2d2501b
Show file tree

Hide file tree

Showing 3 changed files with 60 additions and 43 deletions.
diff --git a/CHANGES.rst b/CHANGES.rst
@@ -332,6 +332,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
@@ -4,8 +4,6 @@
 */
 
 #define NO_IMPORT_ARRAY
-#ifndef ASTROPY_WCS_API_H
-#define ASTROPY_WCS_API_H
 
 #include "astropy_wcs/pyutil.h"
 #include "astropy_wcs/str_list_proxy.h"
@@ -15,14 +13,15 @@
  ***************************************************************************/
 
 #define MAXSIZE 68
+#define ARRAYSIZE 72
 
 static PyTypeObject PyUnitListProxyType;
 
 typedef struct {
   PyObject_HEAD
   /*@null@*/ /*@shared@*/ PyObject* pyobject;
   Py_ssize_t size;
-  char (*array)[72];
+  char (*array)[ARRAYSIZE];
   PyObject* unit_class;
 } PyUnitListProxy;
 
@@ -96,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;
@@ -188,43 +187,35 @@ PyUnitListProxy_getitem(
   return result;
 }
 
-static PyObject *
-PyUnitListProxy_richcmp(PyUnitListProxy *self, PyUnitListProxy *other, int op)
-{int equal;
-  int status;
-
-
-  if ((op == Py_EQ || op == Py_NE) &&
-      PyObject_TypeCheck(other, &PyUnitListProxyType)) {
-
-
-    if (self->size!=other->size||self->unit_class!=other->unit_class||self->array!=other->array)
-       {status = 0;
-        equal = 0;}
-    else if (self == 0x0 || other == 0x0){
-       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;
-      }
-    } else {
-      return NULL;
-    }
+static PyObject*
+PyUnitListProxy_richcmp(
+	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;
   }
-
-  Py_INCREF(Py_NotImplemented);
-    return Py_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
 PyUnitListProxy_setitem(
@@ -338,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;
@@ -403,5 +394,3 @@ _setup_unit_list_proxy_type(
   return 0;
 }
 
-#endif /* ASTROPY_WCS_API_H */
-
diff --git a/astropy/wcs/tests/test_wcs.py b/astropy/wcs/tests/test_wcs.py
@@ -1183,3 +1183,26 @@ def test_footprint_contains():
 
     hasCoord = test_wcs.footprint_contains(SkyCoord(24,2,unit='deg'))
     assert hasCoord == False
+
+
+def test_cunit():
+    # Initializing WCS
+    w1 = wcs.WCS(naxis=2)
+    w2 = wcs.WCS(naxis=2)
+    w3 = wcs.WCS(naxis=2)
+    # Initializing the values of cunit
+    w1.wcs.cunit = ['deg', 'm/s']
+    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