feat: Show AWS region in aws module (starship#482)

Author: Thomas Lee

docs/config/README.md

@@ -114,8 +114,9 @@ prompt_order = [
 
 ## AWS
 
-The `aws` module shows the current AWS profile. This is based on the
-`AWS_PROFILE` env var.
+The `aws` module shows the current AWS region and profile. This is based on
+`AWS_REGION`, `AWS_DEFAULT_REGION`, and `AWS_PROFILE` env var with
+`~/.aws/config` file.
 
 ### Options
 

src/configs/aws.rs

@@ -7,6 +7,7 @@ use starship_module_config_derive::ModuleConfig;
 pub struct AwsConfig<'a> {
     pub symbol: SegmentConfig<'a>,
     pub profile: SegmentConfig<'a>,
+    pub region: SegmentConfig<'a>,
     pub style: Style,
     pub disabled: bool,
 }
@@ -16,6 +17,7 @@ impl<'a> RootModuleConfig<'a> for AwsConfig<'a> {
         AwsConfig {
             symbol: SegmentConfig::new("☁️  "),
             profile: SegmentConfig::default(),
+            region: SegmentConfig::default(),
             style: Color::Yellow.bold(),
             disabled: false,
         }

src/modules/aws.rs

@@ -1,16 +1,85 @@
 use std::env;
+use std::fs::File;
+use std::io::{BufRead, BufReader};
+use std::path::PathBuf;
+use std::str::FromStr;
+
+use dirs::home_dir;
 
 use super::{Context, Module, RootModuleConfig};
 
 use crate::configs::aws::AwsConfig;
 
+fn get_aws_region_from_config(aws_profile: &Option<String>) -> Option<String> {
+    let config_location = env::var("AWS_CONFIG_FILE")
+        .ok()
+        .and_then(|path| PathBuf::from_str(&path).ok())
+        .or_else(|| {
+            let mut home = home_dir()?;
+            home.push(".aws/config");
+            Some(home)
+        })?;
+
+    let file = File::open(&config_location).ok()?;
+    let reader = BufReader::new(file);
+    let lines = reader.lines().filter_map(Result::ok);
+
+    let region_line = if let Some(ref aws_profile) = aws_profile {
+        lines
+            .skip_while(|line| line != &format!("[profile {}]", aws_profile))
+            .skip(1)
+            .take_while(|line| !line.starts_with('['))
+            .find(|line| line.starts_with("region"))
+    } else {
+        lines
+            .skip_while(|line| line != "[default]")
+            .skip(1)
+            .take_while(|line| !line.starts_with('['))
+            .find(|line| line.starts_with("region"))
+    }?;
+
+    let region = region_line.split('=').nth(1)?;
+    let region = region.trim();
+
+    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))
+        })
+}
+
 pub fn module<'a>(context: &'a Context) -> Option<Module<'a>> {
     const AWS_PREFIX: &str = "on ";
 
-    let aws_profile = env::var("AWS_PROFILE").ok()?;
-    if aws_profile.is_empty() {
+    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);
@@ -21,6 +90,7 @@ pub fn module<'a>(context: &'a Context) -> Option<Module<'a>> {
 
     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));
 
     Some(module)
 }

tests/testsuite/aws.rs

@@ -1,17 +1,44 @@
+use std::fs::File;
+use std::io::{self, Write};
+
 use ansi_term::Color;
-use std::io;
 
 use crate::common;
 
 #[test]
-fn no_profile_set() -> io::Result<()> {
+fn no_region_set() -> io::Result<()> {
     let output = common::render_module("aws").env_clear().output()?;
     let expected = "";
     let actual = String::from_utf8(output.stdout).unwrap();
     assert_eq!(expected, actual);
     Ok(())
 }
 
+#[test]
+fn region_set() -> io::Result<()> {
+    let output = common::render_module("aws")
+        .env_clear()
+        .env("AWS_REGION", "ap-northeast-2")
+        .output()?;
+    let expected = format!("on {} ", Color::Yellow.bold().paint("☁️  ap-northeast-2"));
+    let actual = String::from_utf8(output.stdout).unwrap();
+    assert_eq!(expected, actual);
+    Ok(())
+}
+
+#[test]
+fn default_region_set() -> io::Result<()> {
+    let output = common::render_module("aws")
+        .env_clear()
+        .env("AWS_REGION", "ap-northeast-2")
+        .env("AWS_DEFAULT_REGION", "ap-northeast-1")
+        .output()?;
+    let expected = format!("on {} ", Color::Yellow.bold().paint("☁️  ap-northeast-1"));
+    let actual = String::from_utf8(output.stdout).unwrap();
+    assert_eq!(expected, actual);
+    Ok(())
+}
+
 #[test]
 fn profile_set() -> io::Result<()> {
     let output = common::render_module("aws")
@@ -23,3 +50,59 @@ fn profile_set() -> io::Result<()> {
     assert_eq!(expected, actual);
     Ok(())
 }
+
+#[test]
+fn default_profile_set() -> io::Result<()> {
+    let dir = common::new_tempdir()?;
+    let config_path = dir.path().join("config");
+    let mut file = File::create(&config_path)?;
+
+    file.write_all(
+        "[default]
+region = us-east-1
+
+[profile astronauts]
+region = us-east-2
+"
+        .as_bytes(),
+    )?;
+
+    let output = common::render_module("aws")
+        .env_clear()
+        .env("AWS_CONFIG_FILE", config_path.to_string_lossy().as_ref())
+        .output()?;
+    let expected = format!("on {} ", Color::Yellow.bold().paint("☁️  us-east-1"));
+    let actual = String::from_utf8(output.stdout).unwrap();
+    assert_eq!(expected, actual);
+    Ok(())
+}
+
+#[test]
+fn profile_and_config_set() -> io::Result<()> {
+    let dir = common::new_tempdir()?;
+    let config_path = dir.path().join("config");
+    let mut file = File::create(&config_path)?;
+
+    file.write_all(
+        "[default]
+region = us-east-1
+
+[profile astronauts]
+region = us-east-2
+"
+        .as_bytes(),
+    )?;
+
+    let output = common::render_module("aws")
+        .env_clear()
+        .env("AWS_CONFIG_FILE", config_path.to_string_lossy().as_ref())
+        .env("AWS_PROFILE", "astronauts")
+        .output()?;
+    let expected = format!(
+        "on {} ",
+        Color::Yellow.bold().paint("☁️  astronauts(us-east-2)")
+    );
+    let actual = String::from_utf8(output.stdout).unwrap();
+    assert_eq!(expected, actual);
+    Ok(())
+}