Guido van Rossum | e876949 | 1992-08-13 12:14:11 +0000 | [diff] [blame] | 1 | # Complex numbers |
| 2 | |
| 3 | |
| 4 | from math import sqrt |
| 5 | |
| 6 | |
Guido van Rossum | 7565b93 | 1993-12-17 14:23:52 +0000 | [diff] [blame] | 7 | class complex: |
Guido van Rossum | e876949 | 1992-08-13 12:14:11 +0000 | [diff] [blame] | 8 | |
Guido van Rossum | 7565b93 | 1993-12-17 14:23:52 +0000 | [diff] [blame] | 9 | def __init__(self, re, im): |
Guido van Rossum | e876949 | 1992-08-13 12:14:11 +0000 | [diff] [blame] | 10 | self.re = float(re) |
| 11 | self.im = float(im) |
Guido van Rossum | 7565b93 | 1993-12-17 14:23:52 +0000 | [diff] [blame] | 12 | |
| 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 Rossum | e876949 | 1992-08-13 12:14:11 +0000 | [diff] [blame] | 22 | |
| 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 | |
| 65 | def test(): |
| 66 | a = complex(2, 0) |
| 67 | b = complex(3, 4) |
Guido van Rossum | 7565b93 | 1993-12-17 14:23:52 +0000 | [diff] [blame] | 68 | 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 Rossum | e876949 | 1992-08-13 12:14:11 +0000 | [diff] [blame] | 78 | 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 | |
| 85 | test() |