blob: b7ba2e0f5ce8421128936a3d3bdeee72278034bb [file] [log] [blame]
Guido van Rossume8769491992-08-13 12:14:11 +00001# Complex numbers
2
3
4from math import sqrt
5
6
Guido van Rossum7565b931993-12-17 14:23:52 +00007class complex:
Guido van Rossume8769491992-08-13 12:14:11 +00008
Guido van Rossum7565b931993-12-17 14:23:52 +00009 def __init__(self, re, im):
Guido van Rossume8769491992-08-13 12:14:11 +000010 self.re = float(re)
11 self.im = float(im)
Guido van Rossum7565b931993-12-17 14:23:52 +000012
13 def __coerce__(self, other):
14 if type(other) == type(self):
15 if other.__class__ == self.__class__:
16 return self, other
17 else:
18 raise TypeError, 'cannot coerce to complex'
19 else:
20 # The cast to float() may raise an exception!
21 return self, complex(float(other), 0.0)
Guido van Rossume8769491992-08-13 12:14:11 +000022
23 def __repr__(self):
24 return 'complex' + `self.re, self.im`
25
26 def __cmp__(a, b):
27 a = a.__abs__()
28 b = b.__abs__()
29 return (a > b) - (a < b)
30
31 def __float__(self):
32 if self.im:
33 raise ValueError, 'cannot convert complex to float'
34 return float(self.re)
35
36 def __long__(self):
37 return long(float(self))
38
39 def __int__(self):
40 return int(float(self))
41
42 def __abs__(self):
43 # XXX overflow?
44 return sqrt(self.re*self.re + self.im*self.im)
45
46 def __add__(a, b):
47 return complex(a.re + b.re, a.im + b.im)
48
49 def __sub__(a, b):
50 return complex(a.re - b.re, a.im - b.im)
51
52 def __mul__(a, b):
53 return complex(a.re*b.re - a.im*b.im, a.re*b.im + a.im*b.re)
54
55 def __div__(a, b):
56 q = (b.re*b.re + b.im*b.im)
57 re = (a.re*b.re + a.im*b.im) / q
58 im = (a.im*b.re - b.im*a.re) / q
59 return complex(re, im)
60
61 def __neg__(self):
62 return complex(-self.re, -self.im)
63
64
65def test():
66 a = complex(2, 0)
67 b = complex(3, 4)
Guido van Rossum7565b931993-12-17 14:23:52 +000068 print a
69 print b
70 print a+b
71 print a-b
72 print a*b
73 print a/b
74 print b+a
75 print b-a
76 print b*a
77 print b/a
Guido van Rossume8769491992-08-13 12:14:11 +000078 i = complex(0, 1)
79 print i, i*i, i*i*i, i*i*i*i
80 j = complex(1, 1)
81 print j, j*j, j*j*j, j*j*j*j
82 print abs(j), abs(j*j), abs(j*j*j), abs(j*j*j*j)
83 print i/-i
84
85test()