From fed1341e22967d04a8ce26b386be5aade741d926 Mon Sep 17 00:00:00 2001
From: Zhenhui Xie <xiezh0831@yahoo.co.jp>
Date: Mon, 28 Oct 2019 21:41:16 +0800
Subject: [PATCH] feat: Add an option to limit the duration of starship
 directory scanning (#589)

---
 docs/config/README.md        |  3 +++
 src/config.rs                | 17 +++++++++++++++++
 src/configs/starship_root.rs |  2 ++
 src/context.rs               | 11 +++++++++++
 4 files changed, 33 insertions(+)

diff --git a/docs/config/README.md b/docs/config/README.md
index 326c6155196a..897e8697e0e3 100644
--- a/docs/config/README.md
+++ b/docs/config/README.md
@@ -65,6 +65,7 @@ This is the list of prompt-wide configuration options.
 | -------------- | ----------------------------- | ------------------------------------------------------ |
 | `add_newline`  | `true`                        | Add a new line before the start of the prompt.         |
 | `prompt_order` | [link](#default-prompt-order) | Configure the order in which the prompt module occurs. |
+| `scan_timeout` | `30`                          | Timeout for starship to scan files (in milliseconds).  |
 
 ### Example
 
@@ -75,6 +76,8 @@ This is the list of prompt-wide configuration options.
 add_newline = false
 # Overwrite a default_prompt_order and  use custom prompt_order
 prompt_order=["rust","line_break","package","line_break","character"]
+# Wait 10 milliseconds for starship to check files under the current directory.
+scan_timeout = 10
 ```
 
 ### Default Prompt Order
diff --git a/src/config.rs b/src/config.rs
index eac1d6ba515a..92d4add0a331 100644
--- a/src/config.rs
+++ b/src/config.rs
@@ -75,6 +75,23 @@ impl<'a> ModuleConfig<'a> for i64 {
     }
 }
 
+impl<'a> ModuleConfig<'a> for u64 {
+    fn from_config(config: &Value) -> Option<Self> {
+        match config {
+            Value::Integer(value) => {
+                // Converting i64 to u64
+                if *value > 0 {
+                    Some(*value as u64)
+                } else {
+                    None
+                }
+            }
+            Value::String(value) => value.parse::<u64>().ok(),
+            _ => None,
+        }
+    }
+}
+
 impl<'a> ModuleConfig<'a> for f64 {
     fn from_config(config: &Value) -> Option<Self> {
         config.as_float()
diff --git a/src/configs/starship_root.rs b/src/configs/starship_root.rs
index bc0da2399bfa..3873e96d14a4 100644
--- a/src/configs/starship_root.rs
+++ b/src/configs/starship_root.rs
@@ -6,6 +6,7 @@ use starship_module_config_derive::ModuleConfig;
 pub struct StarshipRootConfig<'a> {
     pub add_newline: bool,
     pub prompt_order: Vec<&'a str>,
+    pub scan_timeout: u64,
 }
 
 impl<'a> RootModuleConfig<'a> for StarshipRootConfig<'a> {
@@ -47,6 +48,7 @@ impl<'a> RootModuleConfig<'a> for StarshipRootConfig<'a> {
                 "time",
                 "character",
             ],
+            scan_timeout: 30,
         }
     }
 }
diff --git a/src/context.rs b/src/context.rs
index f1ab2997784a..b306a78d2c40 100644
--- a/src/context.rs
+++ b/src/context.rs
@@ -10,6 +10,7 @@ use std::ffi::OsStr;
 use std::fs;
 use std::path::{Path, PathBuf};
 use std::string::String;
+use std::time::{Duration, SystemTime};
 
 /// Context contains data or common methods that may be used by multiple modules.
 /// The data contained within Context will be relevant to this particular rendering
@@ -137,13 +138,23 @@ impl<'a> Context<'a> {
     }
 
     pub fn get_dir_files(&self) -> Result<&Vec<PathBuf>, std::io::Error> {
+        let start_time = SystemTime::now();
+        let scan_timeout = Duration::from_millis(self.config.get_root_config().scan_timeout);
+
         self.dir_files
             .get_or_try_init(|| -> Result<Vec<PathBuf>, std::io::Error> {
                 let dir_files = fs::read_dir(&self.current_dir)?
+                    .take_while(|_item| {
+                        SystemTime::now().duration_since(start_time).unwrap() < scan_timeout
+                    })
                     .filter_map(Result::ok)
                     .map(|entry| entry.path())
                     .collect::<Vec<PathBuf>>();
 
+                log::trace!(
+                    "Building a vector of directory files took {:?}",
+                    SystemTime::now().duration_since(start_time).unwrap()
+                );
                 Ok(dir_files)
             })
     }