Add Slice function for range type (#83)

HyeockJinKim · corona10 · commit e427daf3db23 · 2019-09-27T00:23:15.000+09:00
* Add Slice function for range type Crate a new range object by calculating start, stop, and step when slice is entered as argument to the __getitem__ function of the range Fixes #77 * Add typecall for __index__ in Index function If __index__ is defined when class is used as index value, __index__ value is used. * Add tests for range
diff --git a/py/internal.go b/py/internal.go
@@ -87,10 +87,21 @@ func MakeGoInt64(a Object) (int64, error) {
 //
 // Will raise TypeError if Index can't be run on this object
 func Index(a Object) (Int, error) {
-	A, ok := a.(I__index__)
-	if ok {
+	if A, ok := a.(I__index__); ok {
 		return A.M__index__()
 	}
+
+	if A, ok, err := TypeCall0(a, "__index__"); ok {
+		if err != nil {
+			return 0, err
+		}
+		if res, ok := A.(Int); ok {
+			return res, nil
+		}
+
+		return 0, err
+	}
+
 	return 0, ExceptionNewf(TypeError, "unsupported operand type(s) for index: '%s'", a.Type().Name)
 }
 
diff --git a/py/range.go b/py/range.go
@@ -79,18 +79,18 @@ func RangeNew(metatype *Type, args Tuple, kwargs StringDict) (Object, error) {
 }
 
 func (r *Range) M__getitem__(key Object) (Object, error) {
+	if slice, ok := key.(*Slice); ok {
+		return computeRangeSlice(r, slice)
+	}
+
 	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
-	}
+	index = computeNegativeIndex(index, r.Length)
 
-	if index < 0 || index >= length {
-		return nil, ExceptionNewf(TypeError, "range object index out of range")
+	if index < 0 || index >= r.Length {
+		return nil, ExceptionNewf(IndexError, "range object index out of range")
 	}
 	result := computeItem(r, index)
 	return result, nil
@@ -160,6 +160,69 @@ func computeRangeLength(start, stop, step Int) Int {
 	return res
 }
 
+func computeNegativeIndex(index, length Int) Int {
+	if index < 0 {
+		index += length
+	}
+	return index
+}
+
+func computeBoundIndex(index, length Int) Int {
+	if index < 0 {
+		index = 0
+	} else if index > length {
+		index = length
+	}
+	return index
+}
+
+func computeRangeSlice(r *Range, s *Slice) (Object, error) {
+	start, err := Index(s.Start)
+	if err != nil {
+		start = 0
+	}
+	stop, err := Index(s.Stop)
+	if err != nil {
+		stop = r.Length
+	}
+
+	step, err := Index(s.Step)
+	if err != nil {
+		step = 1
+	}
+	if step == 0 {
+		return nil, ExceptionNewf(ValueError, "slice step cannot be zero")
+	}
+	start = computeNegativeIndex(start, r.Length)
+	stop = computeNegativeIndex(stop, r.Length)
+
+	start = computeBoundIndex(start, r.Length)
+	stop = computeBoundIndex(stop, r.Length)
+
+	startIndex := computeItem(r, start)
+	stopIndex := computeItem(r, stop)
+	stepIndex := step * r.Step
+
+	var sliceLength Int
+	if start < stop {
+		if stepIndex < 0 {
+			startIndex, stopIndex = stopIndex-1, startIndex-1
+		}
+	} else {
+		if stepIndex < 0 {
+			startIndex, stopIndex = stopIndex+1, startIndex+1
+		}
+	}
+	sliceLength = computeRangeLength(startIndex, stopIndex, stepIndex)
+
+	return &Range{
+		Start: startIndex,
+		Stop: stopIndex,
+		Step: stepIndex,
+		Length: sliceLength,
+	}, nil
+}
+
 // Check interface is satisfied
 var _ I__getitem__ = (*Range)(nil)
 var _ I__iter__ = (*Range)(nil)
diff --git a/py/tests/range.py b/py/tests/range.py
@@ -76,4 +76,44 @@
 assert str(range(0, 3)) == 'range(0, 3)'
 assert str(range(10, 3, -2)) == 'range(10, 3, -2)'
 
+doc="range_slice"
+a = range(10)
+assert a[::-1][0] == 9
+assert a[::-1][9] == 0
+assert a[0:3][0] == 0
+assert a[0:3][2] == 2
+assert a[-3:10][0] == 7
+assert a[-100:13][0] == 0
+assert a[-100:13][9] == 9
+
+try:
+    a[0:3][3]
+except IndexError:
+    pass
+else:
+    assert False, "IndexError not raised"
+try:
+    a[100:13][0]
+except IndexError:
+    pass
+else:
+    assert False, "IndexError not raised"
+try:
+    a[0:3:0]
+except ValueError:
+    pass
+else:
+    assert False, "ValueError not raised"
+
+doc="range_index"
+class Index:
+    def __index__(self):
+        return 1
+
+a = range(10)
+b = Index()
+assert a[b] == 1
+assert a[b:10] == a[1:10]
+assert a[10:b:-1] == a[10:1:-1]
+
 doc="finished"
