complex: added __str__ and __repr__ plus missing properties

raff · ncw · commit d05bbccf75b8 · 2018-10-10T21:36:16.000+01:00
Added __str__ and __repr__ plus missing properties (real, imag, conjugate). Also added tests.
diff --git 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
@@ -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"
