feat: Implement AWS region aliases (#646)

5 files changed, 53 insertions, 6 deletions

diff --git a/docs/config/README.md b/docs/config/README.md
index 9ca73a8ae..a460562b3 100644
--- a/docs/config/README.md
+++ b/docs/config/README.md
@@ -128,11 +128,12 @@ The `aws` module shows the current AWS region and profile. This is based on
 ### Options
 
 | Variable          | Default         | Description                                                                 |
-| ----------------- | --------------- | --------------------------------------------------------------------------- |
-| `symbol`          | `"☁️ "`         | The symbol used before displaying the current AWS profile.                  |
+| ----------------- | --------------- | ----------------------------------------------------------------------------|
+| `symbol`          | `"☁️  "`         | The symbol used before displaying the current AWS profile.                  |
+| `displayed_items` | `all`           | Choose which item to display. Possible values: [`all`, `profile`, `region`] |
+| `region_aliases`  |                 | Table of region aliases to display in addition to the AWS name.             |
 | `style`           | `"bold yellow"` | The style for the module.                                                   |
 | `disabled`        | `false`         | Disables the `AWS` module.                                                  |
-| `displayed_items` | `all`           | Choose which item to display. Possible values: [`all`, `profile`, `region`] |
 
 ### Example
 
@@ -143,6 +144,9 @@ The `aws` module shows the current AWS region and profile. This is based on
 style = "bold blue"
 symbol = "🅰 "
 displayed_items = "region"
+[aws.region_aliases]
+ap-southeast-2 = "au"
+us-east-1 = "va"
 ```
 
 ## Battery
diff --git a/src/config.rs b/src/config.rs
index c2c8c0a45..d3f868f16 100644
--- a/src/config.rs
+++ b/src/config.rs
@@ -3,6 +3,7 @@ use crate::utils;
 use ansi_term::{Color, Style};
 
 use std::clone::Clone;
+use std::collections::HashMap;
 use std::marker::Sized;
 
 use dirs::home_dir;
@@ -127,6 +128,22 @@ where
     }
 }
 
+impl<'a, T, S: ::std::hash::BuildHasher + Default> ModuleConfig<'a> for HashMap<String, T, S>
+where
+    T: ModuleConfig<'a>,
+    S: Clone,
+{
+    fn from_config(config: &'a Value) -> Option<Self> {
+        let mut hm = HashMap::default();
+
+        for (x, y) in config.as_table()?.iter() {
+            hm.insert(x.clone(), T::from_config(y)?);
+        }
+
+        Some(hm)
+    }
+}
+
 impl<'a, T> ModuleConfig<'a> for Option<T>
 where
     T: ModuleConfig<'a> + Sized,
diff --git a/src/configs/aws.rs b/src/configs/aws.rs
index 54dca3055..228558a24 100644
--- a/src/configs/aws.rs
+++ b/src/configs/aws.rs
@@ -1,4 +1,5 @@
 use crate::config::{ModuleConfig, RootModuleConfig, SegmentConfig};
+use std::collections::HashMap;
 
 use ansi_term::{Color, Style};
 use starship_module_config_derive::ModuleConfig;
@@ -18,6 +19,7 @@ pub struct AwsConfig<'a> {
     pub style: Style,
     pub disabled: bool,
     pub displayed_items: AwsItems,
+    pub region_aliases: HashMap<String, &'a str>,
 }
 
 impl<'a> RootModuleConfig<'a> for AwsConfig<'a> {
@@ -29,6 +31,7 @@ impl<'a> RootModuleConfig<'a> for AwsConfig<'a> {
             style: Color::Yellow.bold(),
             disabled: false,
             displayed_items: AwsItems::All,
+            region_aliases: HashMap::new(),
         }
     }
 }
diff --git a/src/modules/aws.rs b/src/modules/aws.rs
index ceb36c85c..c6fc8e92a 100644
--- a/src/modules/aws.rs
+++ b/src/modules/aws.rs
@@ -1,3 +1,4 @@
+use std::collections::HashMap;
 use std::env;
 use std::fs::File;
 use std::io::{BufRead, BufReader};
@@ -76,6 +77,13 @@ fn get_aws_region() -> Option<Region> {
     }
 }
 
+fn alias_region(region: &str, aliases: &HashMap<String, &str>) -> String {
+    match aliases.get(region) {
+        None => region.to_string(),
+        Some(alias) => alias.to_string(),
+    }
+}
+
 pub fn module<'a>(context: &'a Context) -> Option<Module<'a>> {
     const AWS_PREFIX: &str = "on ";
 
@@ -93,9 +101,9 @@ pub fn module<'a>(context: &'a Context) -> Option<Module<'a>> {
 
             let aws_segment = match (&aws_profile, &aws_region) {
                 (None, None) => return None,
-                (Some(p), Some(r)) => format!("{}({})", p, r),
+                (Some(p), Some(r)) => format!("{}({})", p, alias_region(r, &config.region_aliases)),
                 (Some(p), None) => p.to_string(),
-                (None, Some(r)) => r.to_string(),
+                (None, Some(r)) => alias_region(r, &config.region_aliases),
             };
             module.create_segment("all", &config.region.with_value(&aws_segment));
         }
@@ -105,7 +113,7 @@ pub fn module<'a>(context: &'a Context) -> Option<Module<'a>> {
             module.create_segment("profile", &config.profile.with_value(&aws_profile));
         }
         AwsItems::Region => {
-            let aws_region = get_aws_region()?;
+            let aws_region = alias_region(&get_aws_region()?, &config.region_aliases);
 
             module.create_segment("region", &config.region.with_value(&aws_region));
         }
diff --git a/tests/testsuite/aws.rs b/tests/testsuite/aws.rs
index 766d16d14..cb31c4a77 100644
--- a/tests/testsuite/aws.rs
+++ b/tests/testsuite/aws.rs
@@ -29,6 +29,21 @@ fn region_set() -> io::Result<()> {
 }
 
 #[test]
+fn region_set_with_alias() -> io::Result<()> {
+    let output = common::render_module("aws")
+        .env("AWS_REGION", "ap-southeast-2")
+        .use_config(toml::toml! {
+            [aws.region_aliases]
+            ap-southeast-2 = "au"
+        })
+        .output()?;
+    let expected = format!("on {} ", Color::Yellow.bold().paint("☁️  au"));
+    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("AWS_REGION", "ap-northeast-2")