images: Added --iterative option. Closes 227

1 files changed, 37 insertions, 7 deletions

diff --git a/pywal/image.py b/pywal/image.py
index 1b982ce..d23e82b 100644
--- a/pywal/image.py
+++ b/pywal/image.py
@@ -4,6 +4,7 @@ Get the image file.
 import logging
 import os
 import random
+import re
 import sys
 
 from .settings import CACHE_DIR
@@ -11,15 +12,20 @@ from . import util
 from . import wallpaper
 
 
-def get_random_image(img_dir):
-    """Pick a random image file from a directory."""
+def get_image_dir(img_dir):
+    """Get all images in a directory."""
     current_wall = wallpaper.get()
     current_wall = os.path.basename(current_wall)
 
     file_types = (".png", ".jpg", ".jpeg", ".jpe", ".gif")
 
-    images = [img for img in os.scandir(img_dir)
-              if img.name.lower().endswith(file_types)]
+    return [img.name for img in os.scandir(img_dir)
+            if img.name.lower().endswith(file_types)], current_wall
+
+
+def get_random_image(img_dir):
+    """Pick a random image file from a directory."""
+    images, current_wall = get_image_dir(img_dir)
 
     if len(images) > 2 and current_wall in images:
         images.remove(current_wall)
@@ -29,17 +35,41 @@ def get_random_image(img_dir):
         sys.exit(1)
 
     random.shuffle(images)
+    return os.path.join(img_dir, images[0])
+
 
-    return os.path.join(img_dir, images[0].name)
+def get_next_image(img_dir):
+    """Get the next image in a dir."""
+    images, current_wall = get_image_dir(img_dir)
+    images.sort(key=lambda img: [int(x) if x.isdigit() else x
+                                 for x in re.split('([0-9]+)', img)])
 
+    try:
+        next_index = images.index(current_wall) + 1
 
-def get(img, cache_dir=CACHE_DIR):
+    except ValueError:
+        next_index = 0
+
+    try:
+        image = images[next_index]
+
+    except IndexError:
+        image = images[0]
+
+    return os.path.join(img_dir, image)
+
+
+def get(img, cache_dir=CACHE_DIR, iterative=False):
     """Validate image input."""
     if os.path.isfile(img):
         wal_img = img
 
     elif os.path.isdir(img):
-        wal_img = get_random_image(img)
+        if iterative:
+            wal_img = get_next_image(img)
+
+        else:
+            wal_img = get_random_image(img)
 
     else:
         logging.error("No valid image file found.")