[pypy-svn] pypy cmath: (lac, arigo)

[arigo]

Author: Armin Rigo <arigo at tunes.org>
Branch: cmath
Changeset: r40799:3e320a25e3d9
Date: 2011-01-17 18:11 +0100
http://bitbucket.org/pypy/pypy/changeset/3e320a25e3d9/

Log:	(lac, arigo)

	Add space.unpackcomplex().

diff --git a/pypy/objspace/std/test/test_complexobject.py b/pypy/objspace/std/test/test_complexobject.py
--- a/pypy/objspace/std/test/test_complexobject.py
+++ b/pypy/objspace/std/test/test_complexobject.py
@@ -39,7 +39,16 @@
         test_cparse('.e+5', '.e+5', '0.0')
         test_cparse('(1+2j)', '1', '2')
         test_cparse('(1-6j)', '1', '-6')
-        
+
+    def test_unpackcomplex(self):
+        space = self.space
+        w_z = W_ComplexObject(2.0, 3.5)
+        assert space.unpackcomplex(w_z) == (2.0, 3.5)
+        space.raises_w(space.w_TypeError, space.unpackcomplex, space.w_None)
+        w_f = space.newfloat(42.5)
+        assert space.unpackcomplex(w_f) == (42.5, 0.0)
+        w_l = space.wrap(-42L)
+        assert space.unpackcomplex(w_l) == (-42.0, 0.0)
 
     def test_pow(self):
         def _pow((r1, i1), (r2, i2)):

diff --git a/pypy/objspace/std/objspace.py b/pypy/objspace/std/objspace.py
--- a/pypy/objspace/std/objspace.py
+++ b/pypy/objspace/std/objspace.py
@@ -264,6 +264,20 @@
     def newcomplex(self, realval, imagval):
         return W_ComplexObject(realval, imagval)
 
+    def unpackcomplex(self, w_complex):
+        if isinstance(w_complex, W_ComplexObject):
+            return (w_complex.realval, w_complex.imagval)
+        else:
+            try:
+                floatval = self.float_w(w_complex)
+            except OperationError, e:
+                if not e.match(self, self.w_TypeError):
+                    raise
+                raise operationerrfmt(self.w_TypeError,
+                                      "complex number expected, got '%s'",
+                                      self.type(w_complex).getname(self))
+            return (floatval, 0.0)
+
     def newlong(self, val): # val is an int
         return W_LongObject.fromint(self, val)