Skip to content

Commit b39bd4e

Browse files
committed
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
1 parent eb115a9 commit b39bd4e

File tree

1 file changed

+64
-6
lines changed

1 file changed

+64
-6
lines changed

py/range.go

Lines changed: 64 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -79,17 +79,17 @@ func RangeNew(metatype *Type, args Tuple, kwargs StringDict) (Object, error) {
7979
}
8080

8181
func (r *Range) M__getitem__(key Object) (Object, error) {
82+
if slice, ok := key.(*Slice); ok {
83+
return computeRangeSlice(r, slice)
84+
}
85+
8286
index, err := Index(key)
8387
if err != nil {
8488
return nil, err
8589
}
86-
// TODO(corona10): Support slice case
87-
length := computeRangeLength(r.Start, r.Stop, r.Step)
88-
if index < 0 {
89-
index += length
90-
}
90+
index = computeNegativeIndex(index, r.Length)
9191

92-
if index < 0 || index >= length {
92+
if index < 0 || index >= r.Length {
9393
return nil, ExceptionNewf(TypeError, "range object index out of range")
9494
}
9595
result := computeItem(r, index)
@@ -152,6 +152,64 @@ func computeRangeLength(start, stop, step Int) Int {
152152
return res
153153
}
154154

155+
func computeNegativeIndex(index, length Int) Int {
156+
if index < 0 {
157+
index += length
158+
}
159+
return index
160+
}
161+
162+
func computeBoundIndex(index, length Int) Int {
163+
if index < 0 {
164+
index = 0
165+
} else if index >= length {
166+
index = length
167+
}
168+
return index
169+
}
170+
171+
func computeRangeSlice(r *Range, s *Slice) (Object, error) {
172+
start, err := Index(s.Start)
173+
if err != nil {
174+
start = 0
175+
}
176+
stop, err := Index(s.Stop)
177+
if err != nil {
178+
stop = r.Length
179+
}
180+
181+
step, err := Index(s.Step)
182+
if err != nil {
183+
step = 1
184+
}
185+
if step == 0 {
186+
return nil, ExceptionNewf(ValueError, "slice step cannot be zero")
187+
}
188+
start = computeNegativeIndex(start, r.Length)
189+
stop = computeNegativeIndex(stop, r.Length)
190+
191+
start = computeBoundIndex(start, r.Length)
192+
stop = computeBoundIndex(stop, r.Length)
193+
194+
startIndex := computeItem(r, start)
195+
stopIndex := computeItem(r, stop)
196+
stepIndex := step * r.Step
197+
198+
var sliceLength Int
199+
if startIndex < stopIndex {
200+
sliceLength = (stopIndex - startIndex - 1) / stepIndex + 1
201+
} else {
202+
sliceLength = (stopIndex - startIndex + 1) / stepIndex + 1
203+
}
204+
205+
return &Range{
206+
Start: startIndex,
207+
Stop: stopIndex,
208+
Step: stepIndex,
209+
Length: sliceLength,
210+
}, nil
211+
}
212+
155213
// Check interface is satisfied
156214
var _ I__getitem__ = (*Range)(nil)
157215
var _ I__iter__ = (*Range)(nil)

0 commit comments

Comments
 (0)