[pypy-commit] pypy numpy-complex: Initial start. It leaks memory like a seive. We need lltype.malloc(flavor="stack") I think.

[alex_gaynor]

Author: Alex Gaynor <alex.gaynor at gmail.com>
Branch: numpy-complex
Changeset: r48065:a973209a2d26
Date: 2011-10-14 15:28 -0400
http://bitbucket.org/pypy/pypy/changeset/a973209a2d26/

Log:	Initial start. It leaks memory like a seive. We need
	lltype.malloc(flavor="stack") I think.

diff --git a/pypy/module/micronumpy/interp_dtype.py b/pypy/module/micronumpy/interp_dtype.py
--- a/pypy/module/micronumpy/interp_dtype.py
+++ b/pypy/module/micronumpy/interp_dtype.py
@@ -18,6 +18,7 @@
 SIGNEDLTR = "i"
 BOOLLTR = "b"
 FLOATINGLTR = "f"
+COMPLEXLTR = "c"
 
 class W_Dtype(Wrappable):
     def __init__(self, space):
@@ -118,9 +119,19 @@
     W_LowLevelDtype.name = name
     W_LowLevelDtype.aliases = aliases
     W_LowLevelDtype.applevel_types = applevel_types
-    W_LowLevelDtype.num_bytes = rffi.sizeof(T)
-    if expected_size is not None:
-        assert W_LowLevelDtype.num_bytes == expected_size
+
+    if hasattr(T, "_arrayfld"):
+        # Hack for structs (read: complex), rffi.sizeof(Struct) returns a
+        # symbolic, which can't be compared
+        assert T._arrayfld is None
+        primitive_size = sum(rffi.sizeof(VT) for VT in T._flds.itervalues())
+    else:
+        primitive_size = rffi.sizeof(T)
+    if expected_size is None:
+        expected_size = primitive_size
+    assert expected_size == primitive_size
+    W_LowLevelDtype.num_bytes = expected_size
+
     return W_LowLevelDtype
 
 
@@ -465,12 +476,55 @@
 class W_Float64Dtype(FloatArithmeticDtype, W_Float64Dtype):
     pass
 
+ComplexDouble = lltype.Struct("Complex Double",
+    ("real", lltype.Float),
+    ("imag", lltype.Float),
+)
+W_Complex128Dtype = create_low_level_dtype(
+    num = 15, kind = COMPLEXLTR, name = "complex128",
+    aliases = [],
+    applevel_types = ["complex"],
+    T = ComplexDouble,
+    valtype = ComplexDouble,
+    expected_size = 16,
+)
+class W_Complex128Dtype(W_Complex128Dtype):
+    def _create_complex_struct(self, real, imag):
+        # Really we want to stack alloc it, doesn't seem possible though.
+        # Unforutnately a raw alloc means we'll leak memory (no gc) and it
+        # can't be a GcStruct because the array storage isn't a GcArray.
+        # So basically we need stack allocs.
+        c = lltype.malloc(ComplexDouble, flavor="raw")
+        c.real = real
+        c.imag = imag
+        return c
+
+    @specialize.argtype(1)
+    def adapt_val(self, val):
+        if hasattr(val, "_T"):
+            assert val._T is ComplexDouble
+        else:
+            val = self._create_complex_struct(rffi.cast(lltype.Float, val), 0.0)
+        return self.box(val)
+
+    def unwrap(self, space, w_item):
+        real, imag = space.unpackcomplex(w_item)
+        return self.adapt_val(self._create_complex_struct(real, imag))
+
+    def setitem(self, storage, i, item):
+        val = self.unbox(item)
+        # You can't set a full struct, gotta do it one field at a time.
+        self.unerase(storage)[i].real = val.real
+        self.unerase(storage)[i].imag = val.imag
+
+
 ALL_DTYPES = [
     W_BoolDtype,
     W_Int8Dtype, W_UInt8Dtype, W_Int16Dtype, W_UInt16Dtype,
     W_Int32Dtype, W_UInt32Dtype, W_LongDtype, W_ULongDtype,
     W_Int64Dtype, W_UInt64Dtype,
     W_Float32Dtype, W_Float64Dtype,
+    W_Complex128Dtype,
 ]
 
 dtypes_by_alias = unrolling_iterable([
@@ -496,6 +550,7 @@
     __str__ = interp2app(W_Dtype.descr_str),
 
     num = interp_attrproperty("num", cls=W_Dtype),
+    name = interp_attrproperty("name", cls=W_Dtype),
     kind = interp_attrproperty("kind", cls=W_Dtype),
     itemsize = interp_attrproperty("num_bytes", cls=W_Dtype),
     shape = GetSetProperty(W_Dtype.descr_get_shape),
diff --git a/pypy/module/micronumpy/interp_numarray.py b/pypy/module/micronumpy/interp_numarray.py
--- a/pypy/module/micronumpy/interp_numarray.py
+++ b/pypy/module/micronumpy/interp_numarray.py
@@ -38,8 +38,6 @@
             w_dtype = None
             for w_item in l:
                 w_dtype = interp_ufuncs.find_dtype_for_scalar(space, w_item, w_dtype)
-                if w_dtype is space.fromcache(interp_dtype.W_Float64Dtype):
-                    break
             if w_dtype is None:
                 w_dtype = space.w_None
 
diff --git a/pypy/module/micronumpy/interp_ufuncs.py b/pypy/module/micronumpy/interp_ufuncs.py
--- a/pypy/module/micronumpy/interp_ufuncs.py
+++ b/pypy/module/micronumpy/interp_ufuncs.py
@@ -241,6 +241,7 @@
     bool_dtype = space.fromcache(interp_dtype.W_BoolDtype)
     long_dtype = space.fromcache(interp_dtype.W_LongDtype)
     int64_dtype = space.fromcache(interp_dtype.W_Int64Dtype)
+    complex128_dtype = space.fromcache(interp_dtype.W_Complex128Dtype)
 
     if space.is_w(w_type, space.w_bool):
         if current_guess is None or current_guess is bool_dtype:
@@ -256,6 +257,8 @@
             current_guess is long_dtype or current_guess is int64_dtype):
             return int64_dtype
         return current_guess
+    elif space.is_w(w_type, space.w_complex):
+        return complex128_dtype
     return space.fromcache(interp_dtype.W_Float64Dtype)
 
 
diff --git a/pypy/module/micronumpy/test/test_dtypes.py b/pypy/module/micronumpy/test/test_dtypes.py
--- a/pypy/module/micronumpy/test/test_dtypes.py
+++ b/pypy/module/micronumpy/test/test_dtypes.py
@@ -161,3 +161,12 @@
 
         # You can't subclass dtype
         raises(TypeError, type, "Foo", (dtype,), {})
+
+    def test_complex_dtype(self):
+        from numpy import dtype
+
+        d = dtype(complex)
+        assert d.kind == "c"
+        assert d.name == "complex128"
+        assert d.num == 15
+        assert d.itemsize == 16
\ No newline at end of file
diff --git a/pypy/module/micronumpy/test/test_numarray.py b/pypy/module/micronumpy/test/test_numarray.py
--- a/pypy/module/micronumpy/test/test_numarray.py
+++ b/pypy/module/micronumpy/test/test_numarray.py
@@ -558,6 +558,8 @@
         assert array([1.2, True]).dtype is dtype(float)
         assert array([1.2, 5]).dtype is dtype(float)
         assert array([]).dtype is dtype(float)
+        assert array([1 + 2j]).dtype is dtype(complex)
+        assert array([1.0, 1 + 2j]).dtype is dtype(complex)
 
     def test_comparison(self):
         import operator