diff --git a/py/complex.go b/py/complex.go index 8c1413ff..8f42a480 100644 --- a/py/complex.go +++ b/py/complex.go @@ -7,6 +7,7 @@ package py import ( + "fmt" "math" "math/cmplx" ) @@ -60,6 +61,14 @@ func convertToComplex(other Object) (Complex, bool) { return 0, false } +func (a Complex) M__str__() (Object, error) { + return String(fmt.Sprintf("(%g%+gj)", real(complex128(a)), imag(complex128(a)))), nil +} + +func (a Complex) M__repr__() (Object, error) { + return a.M__str__() +} + func (a Complex) M__neg__() (Object, error) { return -a, nil } @@ -286,6 +295,24 @@ func (a Complex) M__ge__(other Object) (Object, error) { return a.M__lt__(other) } +// Properties +func init() { + ComplexType.Dict["real"] = &Property{ + Fget: func(self Object) (Object, error) { + return Float(real(self.(Complex))), nil + }, + } + ComplexType.Dict["imag"] = &Property{ + Fget: func(self Object) (Object, error) { + return Float(imag(self.(Complex))), nil + }, + } + ComplexType.Dict["conjugate"] = MustNewMethod("conjugate", func(self Object) (Object, error) { + cnj := cmplx.Conj(complex128(self.(Complex))) + return Complex(cnj), nil + }, 0, "conjugate() -> Returns the complex conjugate.") +} + // Check interface is satisfied var _ floatArithmetic = Complex(complex(0, 0)) var _ richComparison = Complex(0) diff --git a/py/tests/complex.py b/py/tests/complex.py new file mode 100644 index 00000000..2673f988 --- /dev/null +++ b/py/tests/complex.py @@ -0,0 +1,34 @@ +# Copyright 2018 The go-python Authors. All rights reserved. +# Use of this source code is governed by a BSD-style +# license that can be found in the LICENSE file. + +from libtest import assertRaises + +doc="str" +assert str(3+4j) == "(3+4j)" + +doc="repr" +assert repr(3+4j) == "(3+4j)" + +doc="real" +assert (3+4j).real == 3.0 + +doc="imag" +assert (3+4j).imag == 4.0 + +doc="conjugate" +assert (3+4j).conjugate() == 3-4j + +doc="add" +assert (3+4j) + 2 == 5+4j +assert (3+4j) + 2j == 3+6j + +doc="sub" +assert (3+4j) - 1 == 2+4j +assert (3+4j) - 1j == 3+3j + +doc="mul" +assert (3+4j) * 2 == 6+8j +assert (3+4j) * 2j == -8+6j + +doc="finished"