Improve Homedir detection

- Try passwd file
- Handle cases of empty strings

Cargo.lock

@@ -11,11 +11,18 @@ dependencies = [
  "winapi",
 ]
 
+[[package]]
+name = "libc"
+version = "0.2.180"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "bcc35a38544a891a5f7c865aca548a982ccb3b8650a5b06d0fd33a10283c56fc"
+
 [[package]]
 name = "prettyprompt"
-version = "0.2.0"
+version = "0.3.0"
 dependencies = [
  "ansi_term",
+ "libc",
 ]
 
 [[package]]

Cargo.toml

@@ -1,7 +1,23 @@
 [package]
 name = "prettyprompt"
-version = "0.2.0"
-edition = "2021"
+version = "0.3.0"
+edition = "2024"
 
 [dependencies]
 ansi_term = "0.12"
+libc = "0.2.180"
+
+[[bin]]
+name = "prettyprompt"
+path = "src/prettyprompt.rs"
+
+[[bin]]
+name = "prettyprompt-pwd"
+path = "src/prettyprompt_pwd.rs"
+
+[profile.release]
+opt-level = "z"
+lto = true
+codegen-units = 1
+panic = "abort"
+strip = true

README.md

@@ -35,7 +35,7 @@ A pretty shell prompt, written in rust.
 ## Current Limitations  
 
 - **Hard-Coded Configuration**: User customization is not available yet. 
-- **Exit Code Requirement**: Must pass the last command's exit code as a command-line argument.  
+- **Exit Code Requirement**: Must pass the last command’s exit code as a command-line argument.  
 
 ## Tested on  
 
@@ -59,7 +59,7 @@ cargo build --release
 **Step 2. Add to `$PATH`**  
   - Option 1. Move the binary to a directory in your `$PATH`. Eg:
 ```bash  
-sudo mv target/release/prettyprompt /usr/bin/
+sudo mv target/release/prettyprompt target/release/prettyprompt-pwd /usr/bin/ 
 ```  
   - Option 2. Add the directory containing the binary to `$PATH`  
    System-wide: `/etc/profile`  
@@ -96,13 +96,11 @@ end
 ```  
 
 ### `zsh`  
-Use the precmd function to set the `PROMPT` variable with the output of `prettyprompt $?` as its value.  
+Export the `PS1` variable with the output of `prettyprompt $?` as its value.  
 User-specific: `~/.zshrc`  
 System-wide: `/etc/zsh/zshrc`  
-```zsh  
-precmd() {  
-    PROMPT="$(prettyprompt $?)" # Notice the double quotes  
-}  
+```sh    
+export PS1='$(prettyprompt $?)' # Notice the single quotes
 ```  
 
 ### Other shells  

src/indicators/pwd.rs

@@ -1,4 +1,4 @@
-use crate::indicators::git::{repo_name, MAIN_COL};
+use crate::indicators::git::{MAIN_COL, repo_name};
 use ansi_term::ANSIGenericString;
 use ansi_term::Colour::RGB;
 use std::env::current_dir;
@@ -8,8 +8,19 @@ pub const PATH_COL: ansi_term::Colour = RGB(82, 82, 82);
 pub const REPO_COL: ansi_term::Colour = RGB(55, 120, 130);
 
 /// Find the current user's home directory.
-fn home_dir() -> String {
-    std::env::var("HOME").map_or_else(|_| String::new(), |path| path.to_string())
+fn home_dir() -> Option<String> {
+    // Try `passwd` file first
+    unsafe {
+        let pw = libc::getpwuid(libc::geteuid());
+        if !pw.is_null() {
+            let dir = std::ffi::CStr::from_ptr((*pw).pw_dir);
+            if let Ok(s) = dir.to_str() {
+                return Some(s.to_owned());
+            }
+        }
+    }
+    // Fallback: `$HOME` env var
+    std::env::var("HOME").ok()
 }
 
 /// Returns the full path of the current directory as a string.
@@ -27,7 +38,12 @@ fn remove_repo(pwd_path: &str, repo_path: &str) -> String {
 
 /// Replace the 'home directory' part of the path with a the '~' symbol
 fn replace_home(path: &str) -> String {
-    path.replacen(&home_dir(), "~", 1)
+    if let Some(home) = home_dir() {
+        if path.starts_with(&home) {
+            return path.replacen(&home, "~", 1);
+        }
+    }
+    path.to_owned()
 }
 
 fn short(path: &str, slash_limit: u8) -> String {
@@ -77,7 +93,7 @@ pub fn pwd(
                 format!(
                     "{}{}{}",
                     REPO_COL.paint(repo_name),     // Repo name
-                    MAIN_COL.paint(" \u{F02A2} "), // Seperator
+                    MAIN_COL.paint(" \u{F02A2} "), // Separator
                     PATH_COL.italic().paint(path)  // Sub-directory
                 )
                 .into()

src/indicators/ssh.rs

@@ -1,15 +1,17 @@
 use ansi_term::ANSIGenericString;
 use ansi_term::Colour::RGB;
-use std::env::var; // For environment variables
+use std::env::var_os; // For environment variables
 
 pub const SSH_SYMBOL: &str = "\u{276F}"; // SSH indicator symbol: "❯"
 pub const SSH_COL: ansi_term::Colour = RGB(255, 149, 0); // In SSH session
 pub const NORMIE_COL: ansi_term::Colour = RGB(82, 82, 82); // Non-SSH session
-const SSH_ENV_VARS: [&str; 2] = ["SSH_TTY", "SSH_CONNECTION"]; // Environment variables normally present in SSH sessions
+const SSH_ENV_VARS: [&str; 3] = ["SSH_TTY", "SSH_CONNECTION", "SSH_CLIENT"]; // Environment variables normally present in SSH sessions
 
 /// Checks if current session is an SSH session
 fn is_ssh_session() -> bool {
-    SSH_ENV_VARS.iter().any(|&var_name| var(var_name).is_ok()) // any() iterates through the iter and stops at first `true`. Equal to `||`(`or` condition)
+    SSH_ENV_VARS
+        .iter()
+        .any(|&var_name| var_os(var_name).is_some()) // any() iterates through the iter and stops at first `true`. Equal to `||`(`or` condition)
 }
 
 /// SSH shell indicator

src/indicators/user.rs

@@ -1,21 +1,20 @@
 use ansi_term::ANSIGenericString;
 use ansi_term::Colour::RGB;
-use std::env::var; // For environment variables
+
+use libc::geteuid;
 
 pub const USER_SYMBOL: &str = "\u{276F}"; // User indicator symbol: "❯"
 pub const ROOT_COL: ansi_term::Colour = RGB(255, 53, 94); // If the user is root
 pub const NORMIE_COL: ansi_term::Colour = RGB(0, 255, 180); // Regular user
-const ROOT_USER: &str = "root"; // Root username constant
 
-/// Username of current user
-pub fn username() -> String {
-    var("USER").unwrap_or_else(|_| "UnknownUser".to_string())
+// Check if the current user is root
+pub fn is_root() -> bool {
+    unsafe { geteuid() == 0 } // Uid for root is 0
 }
 
 /// Root user indicator
 pub fn indicator() -> ANSIGenericString<'static, str> {
-    let user = username();
-    if user == ROOT_USER {
+    if is_root() {
         ROOT_COL.paint(USER_SYMBOL)
     } else {
         NORMIE_COL.paint(USER_SYMBOL)

src/prettyprompt.rs

@@ -1,12 +1,11 @@
 mod indicators {
     pub mod error;
     pub mod git;
-    pub mod pwd;
     pub mod ssh;
     pub mod user;
 }
 
-use crate::indicators::{error, git, pwd, ssh, user};
+use crate::indicators::{error, git, ssh, user};
 use ansi_term::{ANSIGenericString, ANSIGenericStrings};
 
 // Add a component to the prompt if the condition is true.
@@ -31,13 +30,9 @@ fn main() {
     let indicate_ssh: bool = true;
     let indicate_err: bool = true;
     let indicate_git_branch: bool = true;
-    let show_pwd: bool = true;
-    let abbrev_home: bool = true;
-    let shorten_path: bool = true;
-    let replace_repo: bool = true;
 
     // Conditionally fetch Git-related info if required, or set to None
-    let git_info: Option<(String, String)> = if indicate_git_branch || replace_repo {
+    let git_info: Option<(String, String)> = if indicate_git_branch {
         git::info()
     } else {
         None
@@ -53,15 +48,6 @@ fn main() {
         error::indicator(&cmd_args)
     });
 
-    // Insert `pwd` at the beginning of the prompt
-    if show_pwd {
-        let repo_path = git_info.map(|info| info.0);
-        components.insert(
-            0,
-            pwd::pwd(abbrev_home, shorten_path, replace_repo, repo_path),
-        );
-    }
-
     // Finally, combine the parts into a single prompts string
     let prompt: ANSIGenericStrings<'_, str> = ANSIGenericStrings(&components[..]);
     // `print!()` prevents an extra newline, unlike `println!()`

src/prettyprompt_pwd.rs

@@ -0,0 +1,26 @@
+mod indicators {
+    pub mod git;
+    pub mod pwd;
+}
+
+use crate::indicators::{git, pwd};
+
+fn main() {
+    // Hard-coded configuration. This will be replaced by a configuration file in a future version
+    let indicate_git_branch: bool = true;
+    let abbrev_home: bool = true;
+    let shorten_path: bool = true;
+    let replace_repo: bool = true;
+
+    // Conditionally fetch Git-related info if required, or set to None
+    let git_info: Option<(String, String)> = if indicate_git_branch || replace_repo {
+        git::info()
+    } else {
+        None
+    };
+
+    let repo_path = git_info.map(|info| info.0);
+    let promt_pwd = pwd::pwd(abbrev_home, shorten_path, replace_repo, repo_path);
+
+    print!("{promt_pwd}"); // A trailing space for aesthetic formatting.
+}