py: Implement range M__getitem__

corona10 · ncw · commit 9985b4959687 · 2019-01-03T18:07:19.000Z
diff --git a/py/range.go b/py/range.go
@@ -78,6 +78,24 @@ func RangeNew(metatype *Type, args Tuple, kwargs StringDict) (Object, error) {
 	}, nil
 }
 
+func (r *Range) M__getitem__(key Object) (Object, error) {
+	index, err := Index(key)
+	if err != nil {
+		return nil, err
+	}
+	// TODO(corona10): Support slice case
+	length := computeRangeLength(r.Start, r.Stop, r.Step)
+	if index < 0 {
+		index += length
+	}
+
+	if index < 0 || index >= length {
+		return nil, ExceptionNewf(TypeError, "range object index out of range")
+	}
+	result := computeItem(r, index)
+	return result, nil
+}
+
 // Make a range iterator from a range
 func (r *Range) M__iter__() (Object, error) {
 	return &RangeIterator{
@@ -109,6 +127,12 @@ func (it *RangeIterator) M__next__() (Object, error) {
 	return r, nil
 }
 
+func computeItem(r *Range, item Int) Int {
+	incr := item * r.Step
+	res := r.Start + incr
+	return res
+}
+
 func computeRangeLength(start, stop, step Int) Int {
 	var lo, hi Int
 	if step > 0 {
@@ -129,5 +153,6 @@ func computeRangeLength(start, stop, step Int) Int {
 }
 
 // Check interface is satisfied
+var _ I__getitem__ = (*Range)(nil)
 var _ I__iter__ = (*Range)(nil)
 var _ I_iterator = (*RangeIterator)(nil)
diff --git a/py/tests/range.py b/py/tests/range.py
@@ -18,4 +18,24 @@
 assert len(a) == 100
 assert len(b) == 100
 
+doc="range_get_item"
+a = range(3)
+assert a[2] == 2
+assert a[1] == 1
+assert a[0] == 0
+assert a[-1] == 2
+assert a[-2] == 1
+assert a[-3] == 0
+
+b = range(0, 10, 2)
+assert b[4] == 8
+assert b[3] == 6
+assert b[2] == 4
+assert b[1] == 2
+assert b[0] == 0
+assert b[-4] == 2
+assert b[-3] == 4
+assert b[-2] == 6
+assert b[-1] == 8
+
 doc="finished"
