Merge pull request #393 from LoLei/270-recursive-i-argument

Add support for #270:  Recursive -i argument

2 files changed, 41 insertions, 10 deletions

diff --git a/pywal/__main__.py b/pywal/__main__.py
index 2e50ca9..de992b7 100644
--- a/pywal/__main__.py
+++ b/pywal/__main__.py
@@ -53,6 +53,11 @@ def get_args():
                           "flag is used: Go through the images in order "
                           "instead of shuffled.")
 
+    arg.add_argument("--recursive", action="store_true",
+                     help="When pywal is given a directory as input and this "
+                          "flag is used: Search for images recursively in "
+                          "subdirectories instead of the root only.")
+
     arg.add_argument("--saturate", metavar="0.0-1.0",
                      help="Set the color saturation.")
 
@@ -159,7 +164,8 @@ def parse_args(parser):
         util.Color.alpha_num = args.a
 
     if args.i:
-        image_file = image.get(args.i, iterative=args.iterative)
+        image_file = image.get(args.i, iterative=args.iterative,
+                               recursive=args.recursive)
         colors_plain = colors.get(image_file, args.l, args.backend,
                                   sat=args.saturate)
 
diff --git a/pywal/image.py b/pywal/image.py
index 897f443..6ef74fc 100644
--- a/pywal/image.py
+++ b/pywal/image.py
@@ -12,6 +12,24 @@ from . import util
 from . import wallpaper
 
 
+def get_image_dir_recursive(img_dir):
+    """Get all images in a directory recursively."""
+    current_wall = wallpaper.get()
+    current_wall = os.path.basename(current_wall)
+
+    file_types = (".png", ".jpg", ".jpeg", ".jpe", ".gif")
+
+    images = []
+    for path, _, files in os.walk(img_dir):
+        for name in files:
+            if name.lower().endswith(file_types):
+                if name.endswith(current_wall):
+                    current_wall = os.path.join(path, name)
+                images.append(os.path.join(path, name))
+
+    return images, current_wall
+
+
 def get_image_dir(img_dir):
     """Get all images in a directory."""
     current_wall = wallpaper.get()
@@ -23,9 +41,12 @@ def get_image_dir(img_dir):
             if img.name.lower().endswith(file_types)], current_wall
 
 
-def get_random_image(img_dir):
+def get_random_image(img_dir, recursive):
     """Pick a random image file from a directory."""
-    images, current_wall = get_image_dir(img_dir)
+    if recursive:
+        images, current_wall = get_image_dir_recursive(img_dir)
+    else:
+        images, current_wall = get_image_dir(img_dir)
 
     if len(images) > 2 and current_wall in images:
         images.remove(current_wall)
@@ -35,12 +56,16 @@ 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 if not recursive else "", images[0])
 
 
-def get_next_image(img_dir):
+def get_next_image(img_dir, recursive):
     """Get the next image in a dir."""
-    images, current_wall = get_image_dir(img_dir)
+    if recursive:
+        images, current_wall = get_image_dir_recursive(img_dir)
+    else:
+        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)])
 
@@ -56,20 +81,20 @@ def get_next_image(img_dir):
     except IndexError:
         image = images[0]
 
-    return os.path.join(img_dir, image)
+    return os.path.join(img_dir if not recursive else "", image)
 
 
-def get(img, cache_dir=CACHE_DIR, iterative=False):
+def get(img, cache_dir=CACHE_DIR, iterative=False, recursive=False):
     """Validate image input."""
     if os.path.isfile(img):
         wal_img = img
 
     elif os.path.isdir(img):
         if iterative:
-            wal_img = get_next_image(img)
+            wal_img = get_next_image(img, recursive)
 
         else:
-            wal_img = get_random_image(img)
+            wal_img = get_random_image(img, recursive)
 
     else:
         logging.error("No valid image file found.")