Implement initial complex number support (#224)

This patch adds minimal support for complex numbers.
Complex literals can be parsed and a very minimal
complex number type has been added to the runtime.
While minimal, the basic pieces are fleshed out enough
that producing follow-on patches to incrementally
develop `complex` should be easy.

compiler/expr_visitor.py

@@ -276,6 +276,8 @@ class ExprVisitor(ast.NodeVisitor):
         expr_str = expr_str + '.Neg()'
     elif isinstance(node.n, float):
       expr_str = 'NewFloat({})'.format(node.n)
+    elif isinstance(node.n, complex):
+      expr_str = 'NewComplex(complex({}, {}))'.format(node.n.real, node.n.imag)
     else:
       msg = 'number type not yet implemented: ' + type(node.n).__name__
       raise util.ParseError(node, msg)

compiler/expr_visitor_test.py

@@ -187,6 +187,7 @@ class ExprVisitorTest(unittest.TestCase):
   testNumFloatSciCap = _MakeLiteralTest(1E6)
   testNumFloatSciCapPlus = _MakeLiteralTest(1E+6)
   testNumFloatSciMinus = _MakeLiteralTest(1e-6)
+  testNumComplex = _MakeLiteralTest(3j)
 
   testSubscriptDictStr = _MakeExprTest('{"foo": 42}["foo"]')
   testSubscriptListInt = _MakeExprTest('[1, 2, 3][2]')

runtime/builtin_types.go

@@ -86,6 +86,7 @@ var builtinTypes = map[*Type]*builtinTypeInfo{
 	BoolType:                      {init: initBoolType, global: true},
 	BytesWarningType:              {global: true},
 	CodeType:                      {},
+	ComplexType:                   {init: initComplexType, global: true},
 	ClassMethodType:               {init: initClassMethodType, global: true},
 	DeprecationWarningType:        {global: true},
 	dictItemIteratorType:          {init: initDictItemIteratorType},

runtime/complex.go

+// Copyright 2016 Google Inc. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package grumpy
+
+import (
+	"fmt"
+	"reflect"
+	"strconv"
+)
+
+// ComplexType is the object representing the Python 'complex' type.
+var ComplexType = newBasisType("complex", reflect.TypeOf(Complex{}), toComplexUnsafe, ObjectType)
+
+// Complex represents Python 'complex' objects.
+type Complex struct {
+	Object
+	value complex128
+}
+
+// NewComplex returns a new Complex holding the given complex value.
+func NewComplex(value complex128) *Complex {
+	return &Complex{Object{typ: ComplexType}, value}
+}
+
+func toComplexUnsafe(o *Object) *Complex {
+	return (*Complex)(o.toPointer())
+}
+
+// ToObject upcasts c to an Object.
+func (c *Complex) ToObject() *Object {
+	return &c.Object
+}
+
+// Value returns the underlying complex value held by c.
+func (c *Complex) Value() complex128 {
+	return c.value
+}
+
+func complexRepr(f *Frame, o *Object) (*Object, *BaseException) {
+	c := toComplexUnsafe(o).Value()
+	rs, is := "", ""
+	pre, post := "", ""
+	sign := ""
+	if real(c) == 0.0 {
+		is = strconv.FormatFloat(imag(c), 'g', -1, 64)
+	} else {
+		pre = "("
+		rs = strconv.FormatFloat(real(c), 'g', -1, 64)
+		is = strconv.FormatFloat(imag(c), 'g', -1, 64)
+		if imag(c) >= 0.0 {
+			sign = "+"
+		}
+		post = ")"
+	}
+	return NewStr(fmt.Sprintf("%s%s%s%sj%s", pre, rs, sign, is, post)).ToObject(), nil
+}
+
+func initComplexType(dict map[string]*Object) {
+	ComplexType.slots.Repr = &unaryOpSlot{complexRepr}
+}

runtime/complex_test.go

+// Copyright 2016 Google Inc. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package grumpy
+
+import (
+	"testing"
+)
+
+func TestComplexRepr(t *testing.T) {
+	cases := []invokeTestCase{
+		{args: wrapArgs(complex(0.0, 0.0)), want: NewStr("0j").ToObject()},
+		{args: wrapArgs(complex(0.0, 1.0)), want: NewStr("1j").ToObject()},
+		{args: wrapArgs(complex(1.0, 2.0)), want: NewStr("(1+2j)").ToObject()},
+		{args: wrapArgs(complex(3.1, -4.2)), want: NewStr("(3.1-4.2j)").ToObject()},
+	}
+	for _, cas := range cases {
+		if err := runInvokeTestCase(wrapFuncForTest(Repr), &cas); err != "" {
+			t.Error(err)
+		}
+	}
+}

runtime/native.go

@@ -34,6 +34,8 @@ var (
 	// these kinds of values resolve directly to primitive Python types.
 	nativeTypes = map[reflect.Type]*Type{
 		reflect.TypeOf(bool(false)):     BoolType,
+		reflect.TypeOf(complex64(0)):    ComplexType,
+		reflect.TypeOf(complex128(0)):   ComplexType,
 		reflect.TypeOf(float32(0)):      FloatType,
 		reflect.TypeOf(float64(0)):      FloatType,
 		reflect.TypeOf(int(0)):          IntType,
@@ -312,6 +314,12 @@ func WrapNative(f *Frame, v reflect.Value) (*Object, *BaseException) {
 		// TODO: Make native bool subtypes singletons and add support
 		// for __new__ so we can use t.Call() here.
 		return (&Int{Object{typ: t}, i}).ToObject(), nil
+	case reflect.Complex64:
+	case reflect.Complex128:
+		c := v.Complex()
+		// TODO: Switch this over to calling the type when `complex.__new__`
+		// gets implemented.
+		return NewComplex(c).ToObject(), nil
 	case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Uint8, reflect.Uint16:
 		return t.Call(f, Args{NewInt(int(v.Int())).ToObject()}, nil)
 	// Handle potentially large ints separately in case of overflow.
@@ -391,6 +399,8 @@ func getNativeType(rtype reflect.Type) *Type {
 		// object.
 		base := nativeType
 		switch rtype.Kind() {
+		case reflect.Complex64, reflect.Complex128:
+			base = ComplexType
 		case reflect.Float32, reflect.Float64:
 			base = FloatType
 		case reflect.Func:

testing/complex_test.py

+# Copyright 2016 Google Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+assert repr(1j) == "1j"
+assert repr(complex()) == "0j"