Improve Dataset#slice argument handling

scholarcy · Jan 6, 2018 · 67350a2 · 67350a2
1 parent 4e9291f
commit 67350a2
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/lib/wapiti/dataset.rb b/lib/wapiti/dataset.rb
@@ -85,8 +85,12 @@ def sample(n = 1, **options)
       Dataset.new sequences.sample(n, **options)
     end
 
-    def slice(*args)
-      Dataset.new [*sequences.slice(*args)]
+    def slice(start, length = 1)
+      if Range === start
+        Dataset.new sequences.slice(start)
+      else
+        Dataset.new sequences.slice(start, length)
+      end
     end
 
     def +(other)