builtin: Implement builtin sum

corona10 · corona10 · commit 8d3d3ff113cc · 2018-09-03T23:03:08.000+09:00
diff --git a/builtin/builtin.go b/builtin/builtin.go
@@ -59,7 +59,7 @@ func init() {
 		py.MustNewMethod("round", builtin_round, 0, round_doc),
 		py.MustNewMethod("setattr", builtin_setattr, 0, setattr_doc),
 		// py.MustNewMethod("sorted", builtin_sorted, 0, sorted_doc),
-		// py.MustNewMethod("sum", builtin_sum, 0, sum_doc),
+		py.MustNewMethod("sum", builtin_sum, 0, sum_doc),
 		// py.MustNewMethod("vars", builtin_vars, 0, vars_doc),
 	}
 	globals := py.StringDict{
@@ -709,3 +709,51 @@ Update and return a dictionary containing the current scope's local variables.`
 const globals_doc = `globals() -> dictionary
 
 Return the dictionary containing the current scope's global variables.`
+
+const sum_doc = `sum($module, iterable, start=0, /)
+--
+Return the sum of a \'start\' value (default: 0) plus an iterable of numbers
+
+When the iterable is empty, return the start value.
+This function is intended specifically for use with numeric values and may
+reject non-numeric types.
+`
+
+func builtin_sum(self py.Object, args py.Tuple) (py.Object, error) {
+	var seq py.Object
+	var start py.Object
+	err := py.UnpackTuple(args, nil, "sum", 1, 2, &seq, &start)
+	if err != nil {
+		return nil, err
+	}
+	if start == nil {
+		start = py.Int(0)
+	} else {
+		switch start.(type) {
+		case py.Bytes:
+			return nil, py.ExceptionNewf(py.TypeError, "sum() can't sum bytes [use b''.join(seq) instead]")
+		case py.String:
+			return nil, py.ExceptionNewf(py.TypeError, "sum() can't sum strings [use ''.join(seq) instead]")
+		}
+	}
+
+	iter, err := py.Iter(seq)
+	if err != nil {
+		return nil, err
+	}
+
+	for {
+		item, err := py.Next(iter)
+		if err != nil {
+			if py.IsException(py.StopIteration, err) {
+				break
+			}
+			return nil, err
+		}
+		start, err = py.Add(start, item)
+		if err != nil {
+			return nil, err
+		}
+	}
+	return start, nil
+}
diff --git a/builtin/tests/builtin.py b/builtin/tests/builtin.py
@@ -159,6 +159,38 @@ class C: pass
 assert getattr(c, "potato") == "spud"
 assert c.potato == "spud"
 
+doc="sum"
+assert sum([1,2,3]) == 6
+assert sum([1,2,3], 3) == 9
+assert sum((1,2,3)) == 6
+assert sum((1,2,3), 3) == 9
+assert sum((1, 2.5, 3)) == 6.5
+assert sum((1, 2.5, 3), 3) == 9.5
+
+try:
+    sum([1,2,3], 'hi')
+except TypeError as e:
+    if e.args[0] != "sum() can't sum strings [use ''.join(seq) instead]":
+        raise
+    ok = True
+assert ok, "TypeError not raised"
+
+try:
+    sum([1,2,3], b'hi')
+except TypeError as e:
+    if e.args[0] != "sum() can't sum bytes [use b''.join(seq) instead]":
+        raise
+    ok = True
+assert ok, "TypeError not raised"
+
+try:
+    sum(['h', 'i'])
+except TypeError as e:
+    if e.args[0] != "unsupported operand type(s) for +: 'int' and 'string'":
+        raise
+    ok = True
+assert ok, "TypeError not raised"
+
 doc="__import__"
 lib = __import__("lib")
 assert lib.libfn() == 42
