Implement dict.setdefault (#365)

Author: m4ns0ur

runtime/dict.go

@@ -326,7 +326,7 @@ func (d *Dict) incVersion() {
 // DelItem removes the entry associated with key from d. It returns true if an
 // item was removed, or false if it did not exist in d.
 func (d *Dict) DelItem(f *Frame, key *Object) (bool, *BaseException) {
-	originValue, raised := d.putItem(f, key, nil)
+	originValue, raised := d.putItem(f, key, nil, true)
 	if raised != nil {
 		return false, raised
 	}
@@ -365,7 +365,7 @@ func (d *Dict) GetItemString(f *Frame, key string) (*Object, *BaseException) {
 // Pop looks up key in d, returning and removing the associalted value if exist,
 // or nil if key is not present in d.
 func (d *Dict) Pop(f *Frame, key *Object) (*Object, *BaseException) {
-	return d.putItem(f, key, nil)
+	return d.putItem(f, key, nil, true)
 }
 
 // Keys returns a list containing all the keys in d.
@@ -390,7 +390,7 @@ func (d *Dict) Len() int {
 
 // putItem associates value with key in d, returning the old associated value if
 // the key was added, or nil if it was not already present in d.
-func (d *Dict) putItem(f *Frame, key, value *Object) (*Object, *BaseException) {
+func (d *Dict) putItem(f *Frame, key, value *Object, overwrite bool) (*Object, *BaseException) {
 	hash, raised := Hash(f, key)
 	if raised != nil {
 		return nil, raised
@@ -413,7 +413,7 @@ func (d *Dict) putItem(f *Frame, key, value *Object) (*Object, *BaseException) {
 					d.table.incUsed(-1)
 					d.incVersion()
 				}
-			} else {
+			} else if overwrite || entry == nil {
 				newEntry := &dictEntry{hash.Value(), key, value}
 				if newTable, ok := t.writeEntry(f, index, newEntry); ok {
 					if newTable != nil {
@@ -435,7 +435,7 @@ func (d *Dict) putItem(f *Frame, key, value *Object) (*Object, *BaseException) {
 
 // SetItem associates value with key in d.
 func (d *Dict) SetItem(f *Frame, key, value *Object) *BaseException {
-	_, raised := d.putItem(f, key, value)
+	_, raised := d.putItem(f, key, value, true)
 	return raised
 }
 
@@ -755,6 +755,36 @@ func dictRepr(f *Frame, o *Object) (*Object, *BaseException) {
 	return NewStr(buf.String()).ToObject(), nil
 }
 
+func dictSetDefault(f *Frame, args Args, _ KWArgs) (*Object, *BaseException) {
+	argc := len(args)
+	if argc == 1 {
+		return nil, f.RaiseType(TypeErrorType, "setdefault expected at least 1 arguments, got 0")
+	}
+	if argc > 3 {
+		return nil, f.RaiseType(TypeErrorType, fmt.Sprintf("setdefault expected at most 2 arguments, got %v", argc-1))
+	}
+	expectedTypes := []*Type{DictType, ObjectType, ObjectType}
+	if argc == 2 {
+		expectedTypes = expectedTypes[:2]
+	}
+	if raised := checkMethodArgs(f, "setdefault", args, expectedTypes...); raised != nil {
+		return nil, raised
+	}
+	d := toDictUnsafe(args[0])
+	key := args[1]
+	var value *Object
+	if argc > 2 {
+		value = args[2]
+	} else {
+		value = None
+	}
+	originValue, raised := d.putItem(f, key, value, false)
+	if originValue != nil {
+		return originValue, raised
+	}
+	return value, raised
+}
+
 func dictSetItem(f *Frame, o, key, value *Object) *BaseException {
 	return toDictUnsafe(o).SetItem(f, key, value)
 }
@@ -803,6 +833,7 @@ func initDictType(dict map[string]*Object) {
 	dict["itervalues"] = newBuiltinFunction("itervalues", dictIterValues).ToObject()
 	dict["keys"] = newBuiltinFunction("keys", dictKeys).ToObject()
 	dict["pop"] = newBuiltinFunction("pop", dictPop).ToObject()
+	dict["setdefault"] = newBuiltinFunction("setdefault", dictSetDefault).ToObject()
 	dict["update"] = newBuiltinFunction("update", dictUpdate).ToObject()
 	dict["values"] = newBuiltinFunction("values", dictValues).ToObject()
 	DictType.slots.Contains = &binaryOpSlot{dictContains}

runtime/dict_test.go

@@ -615,6 +615,30 @@ func TestDictNewRaises(t *testing.T) {
 	}
 }
 
+func TestDictSetDefault(t *testing.T) {
+	setDefaultMethod := mustNotRaise(GetAttr(NewRootFrame(), DictType.ToObject(), NewStr("setdefault"), nil))
+	setDefault := newBuiltinFunction("TestDictSetDefault", func(f *Frame, args Args, kwargs KWArgs) (*Object, *BaseException) {
+		i, raised := setDefaultMethod.Call(f, args, kwargs)
+		if raised != nil {
+			return nil, raised
+		}
+		return NewTuple(i, args[0]).ToObject(), nil
+	}).ToObject()
+	cases := []invokeTestCase{
+		{args: wrapArgs(NewDict(), "foo"), want: newTestTuple(None, newTestDict("foo", None)).ToObject()},
+		{args: wrapArgs(NewDict(), "foo", 42), want: newTestTuple(42, newTestDict("foo", 42)).ToObject()},
+		{args: wrapArgs(newTestDict("foo", 42), "foo"), want: newTestTuple(42, newTestDict("foo", 42)).ToObject()},
+		{args: wrapArgs(newTestDict("foo", 42), "foo", 43), want: newTestTuple(42, newTestDict("foo", 42)).ToObject()},
+		{args: wrapArgs(NewDict()), wantExc: mustCreateException(TypeErrorType, "setdefault expected at least 1 arguments, got 0")},
+		{args: wrapArgs(NewDict(), "foo", "bar", "baz"), wantExc: mustCreateException(TypeErrorType, "setdefault expected at most 2 arguments, got 3")},
+	}
+	for _, cas := range cases {
+		if err := runInvokeTestCase(setDefault, &cas); err != "" {
+			t.Error(err)
+		}
+	}
+}
+
 func TestDictSetItem(t *testing.T) {
 	setItem := newBuiltinFunction("TestDictSetItem", func(f *Frame, args Args, _ KWArgs) (*Object, *BaseException) {
 		if raised := checkFunctionArgs(f, "TestDictSetItem", args, DictType, ObjectType, ObjectType); raised != nil {

runtime/set.go

@@ -81,7 +81,7 @@ func toSetUnsafe(o *Object) *Set {
 
 // Add inserts key into s. If key already exists then does nothing.
 func (s *Set) Add(f *Frame, key *Object) (bool, *BaseException) {
-	origin, raised := s.dict.putItem(f, key, None)
+	origin, raised := s.dict.putItem(f, key, None, true)
 	if raised != nil {
 		return false, raised
 	}

third_party/stdlib/test/test_dict.py

@@ -288,7 +288,6 @@ def test_get(self):
         self.assertRaises(TypeError, d.get)
         self.assertRaises(TypeError, d.get, None, None, None)
 
-    @unittest.expectedFailure
     def test_setdefault(self):
         # dict.setdefault()
         d = {}
@@ -316,7 +315,6 @@ def __hash__(self):
         x.fail = True
         self.assertRaises(Exc, d.setdefault, x, [])
 
-    @unittest.expectedFailure
     def test_setdefault_atomic(self):
         # Issue #13521: setdefault() calls __hash__ and __eq__ only once.
         class Hashed(object):