Merge pull request RustPython#1381 from HyeockJinKim/range

Fixed calculation for range

Author: coolreader18

tests/snippets/builtin_range.py

@@ -19,6 +19,23 @@
 assert_raises(ValueError, lambda: range(10).index('foo'), 'not an int')
 assert_raises(ValueError, lambda: range(1, 10, 0), 'step is zero')
 
+# get tests
+assert range(10)[0] == 0
+assert range(10)[9] == 9
+assert range(10, 0, -1)[0] == 10
+assert range(10, 0, -1)[9] == 1
+assert_raises(IndexError, lambda: range(10)[10], 'out of bound')
+
+# slice tests
+assert range(10)[0:3] == range(3)
+assert range(10)[-5:9] == range(5, 9)
+assert range(10)[100:10] == range(10, 10)
+assert range(10)[-15:3] == range(0, 3)
+assert range(10, 100, 3)[4:1000:5] == range(22, 100, 15)
+assert range(10)[:] == range(10)
+assert range(10, 0, -2)[0:5:2] == range(10, 0, -4)
+assert range(10)[10:11] == range(10,10)
+
 # count tests
 assert range(10).count(2) == 1
 assert range(10).count(11) == 0
@@ -86,3 +103,5 @@
 assert_raises(IndexError, lambda: range(10)[-11], 'out of bound')
 assert range(10)[-2:4] == range(8, 4)
 assert range(10)[-6:-2] == range(4, 8)
+assert range(50, 0, -2)[-5] == 10
+assert range(50, 0, -2)[-5:3:5] == range(10, 44, -10)

vm/src/obj/objrange.rs

@@ -77,26 +77,32 @@ impl PyRange {
         let start = self.start.as_bigint();
         let stop = self.stop.as_bigint();
         let step = self.step.as_bigint();
+        let index = index.clone();
+        if self.is_empty() {
+            return None;
+        }
+
+        let length = if start < stop {
+            (stop - start - 1) / step + 1
+        } else {
+            (start - stop - 1) / (-step) + 1
+        };
 
-        let index = if *index < BigInt::zero() {
-            let index = stop + index;
-            if index < BigInt::zero() {
+        let index = if index < BigInt::zero() {
+            let new_index = &length + &index;
+            if new_index < BigInt::zero() {
                 return None;
             }
-            index
+            length + index
         } else {
-            index.clone()
+            if length <= index {
+                return None;
+            }
+            index
         };
 
-        let result = start + step * &index;
-
-        if (self.forward() && !self.is_empty() && result < *stop)
-            || (!self.forward() && !self.is_empty() && result > *stop)
-        {
-            Some(result)
-        } else {
-            None
-        }
+        let result = start + step * index;
+        Some(result)
     }
 }
 
@@ -293,17 +299,23 @@ impl PyRange {
     #[pymethod(name = "__getitem__")]
     fn getitem(&self, subscript: RangeIndex, vm: &VirtualMachine) -> PyResult {
         match subscript {
-            RangeIndex::Int(index) => {
-                if let Some(value) = self.get(index.as_bigint()) {
-                    Ok(PyInt::new(value).into_ref(vm).into_object())
-                } else {
-                    Err(vm.new_index_error("range object index out of range".to_string()))
-                }
-            }
             RangeIndex::Slice(slice) => {
+                let start = self.start.as_bigint();
+                let stop = self.stop.as_bigint();
+                let step = self.step.as_bigint();
+
                 let new_start = if let Some(int) = slice.start_index(vm)? {
-                    if let Some(i) = self.get(&int) {
+                    let int = &int;
+                    if let Some(i) = self.get(int) {
                         PyInt::new(i).into_ref(vm)
+                    } else if start < stop {
+                        if stop <= int {
+                            self.stop.clone()
+                        } else {
+                            self.start.clone()
+                        }
+                    } else if int < stop {
+                        self.stop.clone()
                     } else {
                         self.start.clone()
                     }
@@ -312,8 +324,17 @@ impl PyRange {
                 };
 
                 let new_end = if let Some(int) = slice.stop_index(vm)? {
-                    if let Some(i) = self.get(&int) {
+                    let int = &int;
+                    if let Some(i) = self.get(int) {
                         PyInt::new(i).into_ref(vm)
+                    } else if start < stop {
+                        if int < start {
+                            self.start.clone()
+                        } else {
+                            self.stop.clone()
+                        }
+                    } else if start < int {
+                        self.start.clone()
                     } else {
                         self.stop.clone()
                     }
@@ -322,7 +343,11 @@ impl PyRange {
                 };
 
                 let new_step = if let Some(int) = slice.step_index(vm)? {
-                    PyInt::new(int * self.step.as_bigint()).into_ref(vm)
+                    if step.is_zero() {
+                        return Err(vm.new_value_error("slice step cannot be zero".to_string()));
+                    } else {
+                        PyInt::new(int * self.step.as_bigint()).into_ref(vm)
+                    }
                 } else {
                     self.step.clone()
                 };
@@ -335,6 +360,13 @@ impl PyRange {
                 .into_ref(vm)
                 .into_object())
             }
+            RangeIndex::Int(index) => {
+                if let Some(value) = self.get(index.as_bigint()) {
+                    Ok(PyInt::new(value).into_ref(vm).into_object())
+                } else {
+                    Err(vm.new_index_error("range object index out of range".to_string()))
+                }
+            }
         }
     }