feat: Add configuration to set how much AWS profile info is shown (#556)

1 files changed, 55 insertions, 36 deletions

diff --git a/src/modules/aws.rs b/src/modules/aws.rs
index 9225c0fe7..ceb36c85c 100644
--- a/src/modules/aws.rs
+++ b/src/modules/aws.rs
@@ -8,9 +8,12 @@ use dirs::home_dir;
 
 use super::{Context, Module, RootModuleConfig};
 
-use crate::configs::aws::AwsConfig;
+use crate::configs::aws::{AwsConfig, AwsItems};
 
-fn get_aws_region_from_config(aws_profile: &Option<String>) -> Option<String> {
+type Profile = String;
+type Region = String;
+
+fn get_aws_region_from_config(aws_profile: Option<&str>) -> Option<Region> {
     let config_location = env::var("AWS_CONFIG_FILE")
         .ok()
         .and_then(|path| PathBuf::from_str(&path).ok())
@@ -44,43 +47,38 @@ fn get_aws_region_from_config(aws_profile: &Option<String>) -> Option<String> {
     Some(region.to_string())
 }
 
-fn get_aws_region() -> Option<(String, String)> {
-    env::var("AWS_DEFAULT_REGION")
-        .ok()
-        .map(|region| (String::new(), region))
-        .or_else(|| {
-            let aws_profile = env::var("AWS_PROFILE").ok();
-            let aws_region = get_aws_region_from_config(&aws_profile);
-
-            if aws_profile.is_none() && aws_region.is_none() {
-                None
-            } else {
-                Some((
-                    aws_profile.unwrap_or_default(),
-                    aws_region.unwrap_or_default(),
-                ))
-            }
-        })
-        .or_else(|| {
-            env::var("AWS_REGION")
-                .ok()
-                .map(|region| (String::new(), region))
-        })
+fn get_aws_profile_and_region() -> (Option<Profile>, Option<Region>) {
+    match (
+        env::var("AWS_PROFILE").ok(),
+        env::var("AWS_REGION").ok(),
+        env::var("AWS_DEFAULT_REGION").ok(),
+    ) {
+        (Some(p), Some(_), Some(dr)) => (Some(p), Some(dr)),
+        (Some(p), Some(r), None) => (Some(p), Some(r)),
+        (None, Some(r), None) => (None, Some(r)),
+        (Some(p), None, Some(dr)) => (Some(p), Some(dr)),
+        (Some(ref p), None, None) => (Some(p.to_owned()), get_aws_region_from_config(Some(p))),
+        (None, None, Some(dr)) => (None, Some(dr)),
+        (None, Some(_), Some(dr)) => (None, Some(dr)),
+        (None, None, None) => (None, get_aws_region_from_config(None)),
+    }
+}
+
+fn get_aws_region() -> Option<Region> {
+    match (
+        env::var("AWS_REGION").ok(),
+        env::var("AWS_DEFAULT_REGION").ok(),
+    ) {
+        (Some(r), None) => Some(r),
+        (None, Some(dr)) => Some(dr),
+        (Some(_), Some(dr)) => Some(dr),
+        (None, None) => get_aws_region_from_config(None),
+    }
 }
 
 pub fn module<'a>(context: &'a Context) -> Option<Module<'a>> {
     const AWS_PREFIX: &str = "on ";
 
-    let (aws_profile, aws_region) = get_aws_region()?;
-    if aws_profile.is_empty() && aws_region.is_empty() {
-        return None;
-    }
-    let aws_region = if aws_profile.is_empty() || aws_region.is_empty() {
-        aws_region
-    } else {
-        format!("({})", aws_region)
-    };
-
     let mut module = context.new_module("aws");
     let config: AwsConfig = AwsConfig::try_load(module.config);
 
@@ -89,8 +87,29 @@ pub fn module<'a>(context: &'a Context) -> Option<Module<'a>> {
     module.get_prefix().set_value(AWS_PREFIX);
 
     module.create_segment("symbol", &config.symbol);
-    module.create_segment("profile", &config.profile.with_value(&aws_profile));
-    module.create_segment("region", &config.profile.with_value(&aws_region));
+    match config.displayed_items {
+        AwsItems::All => {
+            let (aws_profile, aws_region) = get_aws_profile_and_region();
+
+            let aws_segment = match (&aws_profile, &aws_region) {
+                (None, None) => return None,
+                (Some(p), Some(r)) => format!("{}({})", p, r),
+                (Some(p), None) => p.to_string(),
+                (None, Some(r)) => r.to_string(),
+            };
+            module.create_segment("all", &config.region.with_value(&aws_segment));
+        }
+        AwsItems::Profile => {
+            let aws_profile = env::var("AWS_PROFILE").ok()?;
+
+            module.create_segment("profile", &config.profile.with_value(&aws_profile));
+        }
+        AwsItems::Region => {
+            let aws_region = get_aws_region()?;
+
+            module.create_segment("region", &config.region.with_value(&aws_region));
+        }
+    };
 
     Some(module)
 }