fix shift

YieldsLabs · Jul 14, 2023 · c7be056 · c7be056
1 parent d22c204
commit c7be056
Showing 1 changed file with 12 additions and 1 deletion.
diff --git a/ta_lib/core/src/series.rs b/ta_lib/core/src/series.rs
@@ -44,7 +44,7 @@ impl<T: Clone> Series<T> {
     pub fn shift(&self, n: usize) -> Self {
         let data = repeat(None)
             .take(n)
-            .chain(self.data.iter().cloned().skip(n))
+            .chain(self.data.iter().cloned().take(self.len() - n))
             .collect();
 
         Self { data }
@@ -178,6 +178,17 @@ mod tests {
         assert_eq!(result, expected);
     }
 
+    #[test]
+    fn test_shift() {
+        let source = vec![1.0, 2.0, 3.0, 4.0, 5.0];
+        let n = 2;
+        let expected = vec![None, None, Some(1.0), Some(2.0), Some(3.0)];
+
+        let result = Series::from(&source).shift(n);
+
+        assert_eq!(result, expected);
+    }
+
     #[test]
     fn test_change() {
         let source = vec![