From c01f8a7f4805a09917d2d88480895a7b195af567 Mon Sep 17 00:00:00 2001
From: Roey Darwish Dror <roey.ghost@gmail.com>
Date: Wed, 22 Jul 2020 06:19:08 +0300
Subject: [PATCH] Pull zsh themes (fix #471) (#472)

---
 src/steps/zsh.rs | 44 +++++++++++++++++++++++++++-----------------
 1 file changed, 27 insertions(+), 17 deletions(-)

diff --git a/src/steps/zsh.rs b/src/steps/zsh.rs
index debc90e5..15f938b2 100644
--- a/src/steps/zsh.rs
+++ b/src/steps/zsh.rs
@@ -7,9 +7,9 @@ use anyhow::Result;
 use directories::BaseDirs;
 use log::debug;
 use std::env;
-use std::fs;
 use std::path::{Path, PathBuf};
 use std::process::Command;
+use walkdir::WalkDir;
 
 pub fn run_zr(base_dirs: &BaseDirs, run_type: RunType) -> Result<()> {
     let zsh = require("zsh")?;
@@ -89,26 +89,36 @@ pub fn run_oh_my_zsh(ctx: &ExecutionContext) -> Result<()> {
 
     print_separator("oh-my-zsh");
 
-    let mut custom_dir = PathBuf::from(
-        env::var::<_>("ZSH_CUSTOM")
-            .or_else(|_| Command::new("zsh").args(&["-c", "echo -n $ZSH_CUSTOM"]).check_output())?,
-    );
-    custom_dir.push("plugins");
+    let custom_dir = env::var::<_>("ZSH_CUSTOM")
+        .or_else(|_| {
+            Command::new("zsh")
+                .args(&["-c", "test $ZSH_CUSTOM && echo -n $ZSH_CUSTOM"])
+                .check_output()
+        })
+        .map(PathBuf::from)
+        .unwrap_or_else(|e| {
+            let default_path = oh_my_zsh.join("custom");
+            debug!(
+                "Running zsh returned {}. Using default path: {}",
+                e,
+                default_path.display()
+            );
+            default_path
+        });
 
     debug!("oh-my-zsh custom dir: {}", custom_dir.display());
 
-    if let Ok(custom_plugins_dir) = fs::read_dir(custom_dir) {
-        let mut custom_plugins = Repositories::new(ctx.git());
+    let mut custom_repos = Repositories::new(ctx.git());
 
-        for entry in custom_plugins_dir {
-            let entry = entry?;
-            custom_plugins.insert_if_repo(entry.path());
-        }
-        custom_plugins.remove(&oh_my_zsh.to_string_lossy());
-        if !custom_plugins.is_empty() {
-            println!("Pulling custom plugins");
-            ctx.git().multi_pull(&custom_plugins, ctx)?;
-        }
+    for entry in WalkDir::new(custom_dir).max_depth(2) {
+        let entry = entry?;
+        custom_repos.insert_if_repo(entry.path());
+    }
+
+    custom_repos.remove(&oh_my_zsh.to_string_lossy());
+    if !custom_repos.is_empty() {
+        println!("Pulling custom plugins and themes");
+        ctx.git().multi_pull(&custom_repos, ctx)?;
     }
 
     ctx.run_type()