Add Slice function for range type

HyeockJinKim · HyeockJinKim · commit b39bd4e82a31 · 2019-09-16T12:25:00.000+09:00
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
diff --git a/py/range.go b/py/range.go
@@ -79,17 +79,17 @@ 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 {
+	if index < 0 || index >= r.Length {
 		return nil, ExceptionNewf(TypeError, "range object index out of range")
 	}
 	result := computeItem(r, index)
@@ -152,6 +152,64 @@ 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 startIndex < stopIndex {
+		sliceLength = (stopIndex - startIndex - 1) / stepIndex + 1
+	} else {
+		sliceLength = (stopIndex - startIndex + 1) / stepIndex + 1
+	}
+
+	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)
