ENH: added std::complex support to numpy.i

numpy.i now includes ready-made typemaps for std::complex<float>
and std::complex<double> .  Tests were added to testArray using
a newly defined ArrayZ class.

7 files changed, 333 insertions, 12 deletions

diff --git a/tools/swig/numpy.i b/tools/swig/numpy.i
index 217acd5bf..181625057 100644
--- a/tools/swig/numpy.i
+++ b/tools/swig/numpy.i
@@ -3071,15 +3071,13 @@
  *    %numpy_typemaps(long double, NPY_LONGDOUBLE, int)
  */
 
-/* ***************************************************************
- * Swig complains about a syntax error for the following macro
- * expansions:
- *
- *    %numpy_typemaps(complex float,  NPY_CFLOAT , int)
- *
- *    %numpy_typemaps(complex double, NPY_CDOUBLE, int)
- *
- *    %numpy_typemaps(complex long double, NPY_CLONGDOUBLE, int)
- */
+%#ifdef __cplusplus
+
+%include <std_complex.i>
+
+%numpy_typemaps(std::complex<float>,  NPY_CFLOAT , int)
+%numpy_typemaps(std::complex<double>, NPY_CDOUBLE, int)
+
+%#endif
 
 #endif /* SWIGPYTHON */
diff --git a/tools/swig/test/Array.i b/tools/swig/test/Array.i
index 6a8605eb6..64ae57eb0 100644
--- a/tools/swig/test/Array.i
+++ b/tools/swig/test/Array.i
@@ -6,6 +6,7 @@
 #define SWIG_FILE_WITH_INIT
 #include "Array1.h"
 #include "Array2.h"
+#include "ArrayZ.h"
 %}
 
 // Get the NumPy typemaps
@@ -51,6 +52,12 @@
 %apply (int* DIM1 , int* DIM2 , long** ARGOUTVIEW_ARRAY2)
       {(int* nrows, int* ncols, long** data             )};
 
+// Apply the 1D NumPy typemaps
+%apply (int DIM1  , std::complex<double>* INPLACE_ARRAY1)
+      {(int length, std::complex<double>* data          )};
+%apply (std::complex<double>** ARGOUTVIEW_ARRAY1, int* DIM1  )
+      {(std::complex<double>** data             , int* length)};
+
 // Array1 support
 %include "Array1.h"
 %extend Array1
@@ -100,3 +107,29 @@
     return self->asString();
   }
 }
+
+// ArrayZ support
+%include "ArrayZ.h"
+%extend ArrayZ
+{
+  void __setitem__(int i, std::complex<double> v)
+  {
+    self->operator[](i) = v;
+  }
+
+  std::complex<double> __getitem__(int i)
+  {
+    return self->operator[](i);
+  }
+
+  int __len__()
+  {
+    return self->length();
+  }
+
+  std::string __str__()
+  {
+    return self->asString();
+  }
+}
+
diff --git a/tools/swig/test/ArrayZ.cxx b/tools/swig/test/ArrayZ.cxx
new file mode 100644
index 000000000..4756c6981
--- /dev/null
+++ b/tools/swig/test/ArrayZ.cxx
@@ -0,0 +1,131 @@
+#include "ArrayZ.h"
+#include <iostream>
+#include <sstream>
+
+// Default/length/array constructor
+ArrayZ::ArrayZ(int length, std::complex<double>* data) :
+  _ownData(false), _length(0), _buffer(0)
+{
+  resize(length, data);
+}
+
+// Copy constructor
+ArrayZ::ArrayZ(const ArrayZ & source) :
+  _length(source._length)
+{
+  allocateMemory();
+  *this = source;
+}
+
+// Destructor
+ArrayZ::~ArrayZ()
+{
+  deallocateMemory();
+}
+
+// Assignment operator
+ArrayZ & ArrayZ::operator=(const ArrayZ & source)
+{
+  int len = _length < source._length ? _length : source._length;
+  for (int i=0;  i < len; ++i)
+  {
+    (*this)[i] = source[i];
+  }
+  return *this;
+}
+
+// Equals operator
+bool ArrayZ::operator==(const ArrayZ & other) const
+{
+  if (_length != other._length) return false;
+  for (int i=0; i < _length; ++i)
+  {
+    if ((*this)[i] != other[i]) return false;
+  }
+  return true;
+}
+
+// Length accessor
+int ArrayZ::length() const
+{
+  return _length;
+}
+
+// Resize array
+void ArrayZ::resize(int length, std::complex<double>* data)
+{
+  if (length < 0) throw std::invalid_argument("ArrayZ length less than 0");
+  if (length == _length) return;
+  deallocateMemory();
+  _length = length;
+  if (!data)
+  {
+    allocateMemory();
+  }
+  else
+  {
+    _ownData = false;
+    _buffer  = data;
+  }
+}
+
+// Set item accessor
+std::complex<double> & ArrayZ::operator[](int i)
+{
+  if (i < 0 || i >= _length) throw std::out_of_range("ArrayZ index out of range");
+  return _buffer[i];
+}
+
+// Get item accessor
+const std::complex<double> & ArrayZ::operator[](int i) const
+{
+  if (i < 0 || i >= _length) throw std::out_of_range("ArrayZ index out of range");
+  return _buffer[i];
+}
+
+// String output
+std::string ArrayZ::asString() const
+{
+  std::stringstream result;
+  result << "[";
+  for (int i=0; i < _length; ++i)
+  {
+    result << " " << _buffer[i];
+    if (i < _length-1) result << ",";
+  }
+  result << " ]";
+  return result.str();
+}
+
+// Get view
+void ArrayZ::view(std::complex<double>** data, int* length) const
+{
+  *data   = _buffer;
+  *length = _length;
+}
+
+// Private methods
+ void ArrayZ::allocateMemory()
+ {
+   if (_length == 0)
+   {
+     _ownData = false;
+     _buffer  = 0;
+   }
+   else
+   {
+     _ownData = true;
+     _buffer = new std::complex<double>[_length];
+   }
+ }
+
+ void ArrayZ::deallocateMemory()
+ {
+   if (_ownData && _length && _buffer)
+   {
+     delete [] _buffer;
+   }
+   _ownData = false;
+   _length  = 0;
+   _buffer  = 0;
+ }
diff --git a/tools/swig/test/ArrayZ.h b/tools/swig/test/ArrayZ.h
new file mode 100644
index 000000000..75abae0b2
--- /dev/null
+++ b/tools/swig/test/ArrayZ.h
@@ -0,0 +1,56 @@
+#ifndef ARRAYZ_H
+#define ARRAYZ_H
+
+#include <stdexcept>
+#include <string>
+#include <complex>
+
+class ArrayZ
+{
+public:
+
+  // Default/length/array constructor
+  ArrayZ(int length = 0, std::complex<double>* data = 0);
+
+  // Copy constructor
+  ArrayZ(const ArrayZ & source);
+
+  // Destructor
+  ~ArrayZ();
+
+  // Assignment operator
+  ArrayZ & operator=(const ArrayZ & source);
+
+  // Equals operator
+  bool operator==(const ArrayZ & other) const;
+
+  // Length accessor
+  int length() const;
+
+  // Resize array
+  void resize(int length, std::complex<double>* data = 0);
+
+  // Set item accessor
+  std::complex<double> & operator[](int i);
+
+  // Get item accessor
+  const std::complex<double> & operator[](int i) const;
+
+  // String output
+  std::string asString() const;
+
+  // Get view
+  void view(std::complex<double>** data, int* length) const;
+
+private:
+  // Members
+  bool _ownData;
+  int _length;
+  std::complex<double> * _buffer;
+
+  // Methods
+  void allocateMemory();
+  void deallocateMemory();
+};
+
+#endif
diff --git a/tools/swig/test/Makefile b/tools/swig/test/Makefile
index 5360b1ced..5632e7ad0 100644
--- a/tools/swig/test/Makefile
+++ b/tools/swig/test/Makefile
@@ -5,7 +5,8 @@ PROXIES    = $(INTERFACES:.i=.py      )
 
 # Default target: build the tests
 .PHONY : all
-all: $(WRAPPERS) Array1.cxx Array1.h Farray.cxx Farray.h Vector.cxx Vector.h \
+all: $(WRAPPERS) Array1.cxx Array1.h Array2.cxx Array2.h ArrayZ.cxx ArrayZ.h \
+    Farray.cxx Farray.h Vector.cxx Vector.h \
 	Matrix.cxx Matrix.h Tensor.cxx Tensor.h Fortran.h Fortran.cxx
 	./setup.py build_ext -i
 
diff --git a/tools/swig/test/setup.py b/tools/swig/test/setup.py
index e39114f91..81df1b8ed 100755
--- a/tools/swig/test/setup.py
+++ b/tools/swig/test/setup.py
@@ -15,7 +15,8 @@ numpy_include = numpy.get_include()
 _Array = Extension("_Array",
                    ["Array_wrap.cxx",
                     "Array1.cxx",
-                    "Array2.cxx"],
+                    "Array2.cxx",
+                    "ArrayZ.cxx"],
                    include_dirs = [numpy_include],
                    )
 
diff --git a/tools/swig/test/testArray.py b/tools/swig/test/testArray.py
index d986de3b3..2ccca19b4 100755
--- a/tools/swig/test/testArray.py
+++ b/tools/swig/test/testArray.py
@@ -269,12 +269,113 @@ class Array2TestCase(unittest.TestCase):
 
 ######################################################################
 
+class ArrayZTestCase(unittest.TestCase):
+
+    def setUp(self):
+        self.length = 5
+        self.array3 = Array.ArrayZ(self.length)
+
+    def testConstructor0(self):
+        "Test ArrayZ default constructor"
+        a = Array.ArrayZ()
+        self.failUnless(isinstance(a, Array.ArrayZ))
+        self.failUnless(len(a) == 0)
+
+    def testConstructor1(self):
+        "Test ArrayZ length constructor"
+        self.failUnless(isinstance(self.array3, Array.ArrayZ))
+
+    def testConstructor2(self):
+        "Test ArrayZ array constructor"
+        na = np.arange(self.length, dtype=np.complex128)
+        aa = Array.ArrayZ(na)
+        self.failUnless(isinstance(aa, Array.ArrayZ))
+
+    def testConstructor3(self):
+        "Test ArrayZ copy constructor"
+        for i in range(self.array3.length()): self.array3[i] = complex(i,-i)
+        arrayCopy = Array.ArrayZ(self.array3)
+        self.failUnless(arrayCopy == self.array3)
+
+    def testConstructorBad(self):
+        "Test ArrayZ length constructor, negative"
+        self.assertRaises(ValueError, Array.ArrayZ, -4)
+
+    def testLength(self):
+        "Test ArrayZ length method"
+        self.failUnless(self.array3.length() == self.length)
+
+    def testLen(self):
+        "Test ArrayZ __len__ method"
+        self.failUnless(len(self.array3) == self.length)
+
+    def testResize0(self):
+        "Test ArrayZ resize method, length"
+        newLen = 2 * self.length
+        self.array3.resize(newLen)
+        self.failUnless(len(self.array3) == newLen)
+
+    def testResize1(self):
+        "Test ArrayZ resize method, array"
+        a = np.zeros((2*self.length,), dtype=np.complex128)
+        self.array3.resize(a)
+        self.failUnless(len(self.array3) == a.size)
+
+    def testResizeBad(self):
+        "Test ArrayZ resize method, negative length"
+        self.assertRaises(ValueError, self.array3.resize, -5)
+
+    def testSetGet(self):
+        "Test ArrayZ __setitem__, __getitem__ methods"
+        n = self.length
+        for i in range(n):
+            self.array3[i] = i*i
+        for i in range(n):
+            self.failUnless(self.array3[i] == i*i)
+
+    def testSetBad1(self):
+        "Test ArrayZ __setitem__ method, negative index"
+        self.assertRaises(IndexError, self.array3.__setitem__, -1, 0)
+
+    def testSetBad2(self):
+        "Test ArrayZ __setitem__ method, out-of-range index"
+        self.assertRaises(IndexError, self.array3.__setitem__, self.length+1, 0)
+
+    def testGetBad1(self):
+        "Test ArrayZ __getitem__ method, negative index"
+        self.assertRaises(IndexError, self.array3.__getitem__, -1)
+
+    def testGetBad2(self):
+        "Test ArrayZ __getitem__ method, out-of-range index"
+        self.assertRaises(IndexError, self.array3.__getitem__, self.length+1)
+
+    def testAsString(self):
+        "Test ArrayZ asString method"
+        for i in range(self.array3.length()): self.array3[i] = complex(i+1,-i-1)
+        self.failUnless(self.array3.asString() == "[ (1,-1), (2,-2), (3,-3), (4,-4), (5,-5) ]")
+
+    def testStr(self):
+        "Test ArrayZ __str__ method"
+        for i in range(self.array3.length()): self.array3[i] = complex(i-2,(i-2)*2)
+        self.failUnless(str(self.array3) == "[ (-2,-4), (-1,-2), (0,0), (1,2), (2,4) ]")
+
+    def testView(self):
+        "Test ArrayZ view method"
+        for i in range(self.array3.length()): self.array3[i] = complex(i+1,i+2)
+        a = self.array3.view()
+        self.failUnless(isinstance(a, np.ndarray))
+        self.failUnless(len(a) == self.length)
+        self.failUnless((a == [1+2j, 2+3j, 3+4j, 4+5j, 5+6j]).all())
+
+######################################################################
+
 if __name__ == "__main__":
 
     # Build the test suite
     suite = unittest.TestSuite()
     suite.addTest(unittest.makeSuite(Array1TestCase))
     suite.addTest(unittest.makeSuite(Array2TestCase))
+    suite.addTest(unittest.makeSuite(ArrayZTestCase))
 
     # Execute the test suite
     print("Testing Classes of Module Array")