py: add the 'add' method to the set class

kellrott · web-flow · commit 8e99b322adb3 · 2023-01-09T12:21:24.000+01:00
diff --git a/py/set.go b/py/set.go
@@ -46,6 +46,17 @@ func NewSetFromItems(items []Object) *Set {
 	return s
 }
 
+func init() {
+	SetType.Dict["add"] = MustNewMethod("add", func(self Object, args Tuple) (Object, error) {
+		setSelf := self.(*Set)
+		if len(args) != 1 {
+			return nil, ExceptionNewf(TypeError, "append() takes exactly one argument (%d given)", len(args))
+		}
+		setSelf.Add(args[0])
+		return NoneType{}, nil
+	}, 0, "add(value)")
+}
+
 // Add an item to the set
 func (s *Set) Add(item Object) {
 	s.items[item] = SetValue{}
diff --git a/py/tests/set.py b/py/tests/set.py
@@ -2,6 +2,8 @@
 # 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="__and__"
 a = {1, 2, 3}
 b = {2, 3, 4, 5}
@@ -81,6 +83,18 @@
 assert 4 in c
 assert 5 in c
 
+doc="add"
+a = set()
+a.add(1)
+a.add(2)
+a.add(3)
+assert len(a) == 3
+assert 1 in a
+assert 2 in a
+assert 3 in a
+assert 4 not in a
+assertRaises(TypeError, lambda: a.add())
+
 doc="__eq__, __ne__"
 a = set([1,2,3])
 assert a.__eq__(3) != True
