Handled more options at bytearray init

Author: windom

grumpy-runtime-src/runtime/bytearray.go

@@ -98,12 +98,52 @@ func byteArrayGT(f *Frame, v, w *Object) (*Object, *BaseException) {
 }
 
 func byteArrayInit(f *Frame, o *Object, args Args, _ KWArgs) (*Object, *BaseException) {
-	if raised := checkFunctionArgs(f, "__init__", args, IntType); raised != nil {
+	argc := len(args)
+	if argc == 0 {
+		return None, nil
+	}
+	expectedTypes := []*Type{ObjectType}
+	if raised := checkFunctionArgs(f, "__init__", args, expectedTypes...); raised != nil {
 		return nil, raised
 	}
 	a := toByteArrayUnsafe(o)
+	if args[0].isInstance(IntType) {
+		a.mutex.Lock()
+		a.value = make([]byte, toIntUnsafe(args[0]).Value())
+		a.mutex.Unlock()
+		return None, nil
+	}
+	// else it must be an iterable
+	value := []byte{}
+	iter, raised := Iter(f, args[0])
+	if raised != nil {
+		return nil, raised
+	}
+	item, raised := Next(f, iter)
+	for ; raised == nil; item, raised = Next(f, iter) {
+		switch {
+		case item.isInstance(StrType):
+			sval := toStrUnsafe(item).Value()
+			if len(sval) != 1 {
+				return nil, f.RaiseType(ValueErrorType, "string must be of size 1")
+			}
+			value = append(value, sval[0])
+		case item.isInstance(IntType):
+			ival := toIntUnsafe(item).Value()
+			if ival < 0 || ival >= 256 {
+				return nil, f.RaiseType(ValueErrorType, "byte must be in range(0, 256)")
+			}
+			value = append(value, byte(ival))
+		default:
+			return nil, f.RaiseType(TypeErrorType, "an integer or string of size 1 is required")
+		}
+	}
+	if !raised.isInstance(StopIterationType) {
+		return nil, raised
+	}
+	f.RestoreExc(nil, nil)
 	a.mutex.Lock()
-	a.value = make([]byte, toIntUnsafe(args[0]).Value())
+	a.value = value
 	a.mutex.Unlock()
 	return None, nil
 }

grumpy-runtime-src/runtime/bytearray_test.go

@@ -68,8 +68,18 @@ func TestByteArrayGetItem(t *testing.T) {
 
 func TestByteArrayInit(t *testing.T) {
 	cases := []invokeTestCase{
+		{args: wrapArgs(), want: newTestByteArray("").ToObject()},
 		{args: wrapArgs(3), want: newTestByteArray("\x00\x00\x00").ToObject()},
-		{args: wrapArgs(newObject(ObjectType)), wantExc: mustCreateException(TypeErrorType, `'__init__' requires a 'int' object but received a "object"`)},
+		{args: wrapArgs([]int{3, 2, 1}), want: newTestByteArray("\x03\x02\x01").ToObject()},
+		{args: wrapArgs("abc"), want: newTestByteArray("abc").ToObject()},
+		{args: wrapArgs([]string{"a", "b", "c"}), want: newTestByteArray("abc").ToObject()},
+		{args: wrapArgs(newTestXRange(3)), want: newTestByteArray("\x00\x01\x02").ToObject()},
+		{args: wrapArgs(newTestRange(3)), want: newTestByteArray("\x00\x01\x02").ToObject()},
+		{args: wrapArgs(newObject(ObjectType)), wantExc: mustCreateException(TypeErrorType, `'object' object is not iterable`)},
+		{args: wrapArgs([]int{-1}), wantExc: mustCreateException(ValueErrorType, `byte must be in range(0, 256)`)},
+		{args: wrapArgs([]int{256}), wantExc: mustCreateException(ValueErrorType, `byte must be in range(0, 256)`)},
+		{args: wrapArgs([]string{"ab"}), wantExc: mustCreateException(ValueErrorType, `string must be of size 1`)},
+		{args: wrapArgs([]interface{}{5, []interface{}{}}), wantExc: mustCreateException(TypeErrorType, `an integer or string of size 1 is required`)},
 	}
 	for _, cas := range cases {
 		if err := runInvokeTestCase(ByteArrayType.ToObject(), &cas); err != "" {