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

[arigo]

Author: Armin Rigo <arigo at tunes.org>
Branch: cmath
Changeset: r40809:5dad3217caee
Date: 2011-01-17 19:07 +0100
http://bitbucket.org/pypy/pypy/changeset/5dad3217caee/

Log:	(lac, arigo)

	Fix rect() and polar() to accept or return a tuple of two floats
	instead of a complex.

diff --git a/pypy/module/cmath/__init__.py b/pypy/module/cmath/__init__.py
--- a/pypy/module/cmath/__init__.py
+++ b/pypy/module/cmath/__init__.py
@@ -33,5 +33,9 @@
     appleveldefs = {
     }
 
-    interpleveldefs = dict([(name, 'interp_cmath.wrapped_' + name)
-                            for name in names_and_docstrings])
+    interpleveldefs = {
+        'pi': 'space.wrap(interp_cmath.pi)',
+        'e':  'space.wrap(interp_cmath.e)',
+    }
+    interpleveldefs.update(dict([(name, 'interp_cmath.wrapped_' + name)
+                                 for name in names_and_docstrings]))

diff --git a/pypy/module/cmath/test/test_cmath.py b/pypy/module/cmath/test/test_cmath.py
--- a/pypy/module/cmath/test/test_cmath.py
+++ b/pypy/module/cmath/test/test_cmath.py
@@ -45,6 +45,22 @@
         z = cmath.log(100j, 10j)
         assert abs(z - (1.6824165174565446-0.46553647994440367j)) < 1e-10
 
+    def test_pi_e(self):
+        import cmath, math
+        assert cmath.pi == math.pi
+        assert cmath.e == math.e
+
+    def test_rect(self):
+        import cmath
+        z = cmath.rect(2.0, cmath.pi/2)
+        assert abs(z - 2j) < 1e-10
+
+    def test_polar(self):
+        import cmath
+        r, phi = cmath.polar(2j)
+        assert r == 2
+        assert abs(phi - cmath.pi/2) < 1e-10
+
 
 def parse_testfile(fname):
     """Parse a file with test values

diff --git a/pypy/module/cmath/interp_cmath.py b/pypy/module/cmath/interp_cmath.py
--- a/pypy/module/cmath/interp_cmath.py
+++ b/pypy/module/cmath/interp_cmath.py
@@ -1,5 +1,6 @@
 import math
 from math import fabs
+from pypy.rlib.objectmodel import specialize
 from pypy.rlib.rarithmetic import copysign, asinh, log1p, isinf, isnan
 from pypy.interpreter.gateway import ObjSpace, W_Root, NoneNotWrapped
 from pypy.module.cmath import Module, names_and_docstrings
@@ -21,18 +22,27 @@
 from pypy.module.cmath.special_value import tanh_special_values
 from pypy.module.cmath.special_value import rect_special_values
 
+pi = math.pi
+e  = math.e
+
+
+ at specialize.arg(0)
+def call_c_func(c_func, x, y):
+    try:
+        resx, resy = c_func(x, y)
+    except ValueError:
+        raise OperationError(space.w_ValueError,
+                             space.wrap("math domain error"))
+    except OverflowError:
+        raise OperationError(space.w_OverflowError,
+                             space.wrap("math range error"))
+    return resx, resy
+
 
 def unaryfn(c_func):
     def wrapper(space, w_z):
         x, y = space.unpackcomplex(w_z)
-        try:
-            resx, resy = c_func(x, y)
-        except ValueError:
-            raise OperationError(space.w_ValueError,
-                                 space.wrap("math domain error"))
-        except OverflowError:
-            raise OperationError(space.w_OverflowError,
-                                 space.wrap("math range error"))
+        resx, resy = call_c_func(c_func, x, y)
         return space.newcomplex(resx, resy)
     #
     name = c_func.func_name
@@ -283,6 +293,8 @@
         return space.truediv(w_logz, w_logbase)
     else:
         return w_logz
+wrapped_log.unwrap_spec = [ObjSpace, W_Root, W_Root]
+wrapped_log.func_doc = _inner_wrapped_log.func_doc
 
 
 @unaryfn
@@ -452,7 +464,6 @@
     return sy, -sx
 
 
- at unaryfn
 def c_rect(r, phi):
     if not isfinite(r) or not isfinite(phi):
         # if r is +/-infinity and phi is finite but nonzero then
@@ -479,6 +490,14 @@
     imag = r * math.sin(phi)
     return real, imag
 
+def wrapped_rect(space, w_x, w_y):
+    x = space.float_w(w_x)
+    y = space.float_w(w_y)
+    resx, resy = call_c_func(c_rect, x, y)
+    return space.newcomplex(resx, resy)
+wrapped_rect.unwrap_spec = [ObjSpace, W_Root, W_Root]
+wrapped_rect.func_doc = names_and_docstrings['rect']
+
 
 def c_atan2(x, y):
     # Windows screws up atan2 for inf and nan, and alpha Tru64 5.1 doesn't
@@ -505,8 +524,14 @@
     return math.atan2(y, x)
 
 
- at unaryfn
 def c_polar(x, y):
     phi = c_atan2(x, y)
     r = math.hypot(x, y)
     return r, phi
+
+def wrapped_polar(space, w_z):
+    x, y = space.unpackcomplex(w_z)
+    resx, resy = call_c_func(c_polar, x, y)
+    return space.newtuple([space.newfloat(resx), space.newfloat(resy)])
+wrapped_polar.unwrap_spec = [ObjSpace, W_Root]
+wrapped_polar.func_doc = names_and_docstrings['polar']