Repository: nushell/nushell
File: crates/nu-path/src/pwd_per_drive.rs (Rust)

Code before:
    }
}

/// Helper to check if input path is relative path
/// with drive letter, it can be expanded with PWD-per-drive.
fn need_expand(path: &Path) -> bool {
    if let Some(path_str) = path.to_str() {
        let chars: Vec<char> = path_str.chars().collect();
        if chars.len() >= 2 {
            return chars[1] == ':' && (chars.len() == 2 || (chars[2] != '/' && chars[2] != '\\'));
        }
    }
    false
}

struct DriveToPwdMap {
    map: [Option<String>; 26], // Fixed-size array for A-Z
}

impl DriveToPwdMap {
    pub fn new() -> Self {
        // Initialize by current PWD-per-drive
        let mut map: [Option<String>; 26] = Default::default();
        for (drive_index, drive_letter) in ('A'..='Z').enumerate() {
            let env_var = Self::env_var_for_drive(drive_letter);
            if let Ok(env_pwd) = std::env::var(&env_var) {
                if env_pwd.len() > 3 {
                    map[drive_index] = Some(env_pwd);
                }
            }
        }
        Self { map }
    }

    pub fn env_var_for_drive(drive_letter: char) -> String {
        let drive_letter = drive_letter.to_ascii_uppercase();
        format!("={}:", drive_letter)
    }

    /// Collect PWD-per-drive as env vars (for child process)
    pub fn get_env_vars(&self, env: &mut HashMap<String, String>) {
        for (drive_index, drive_letter) in ('A'..='Z').enumerate() {
            if let Some(pwd) = self.map[drive_index].clone() {
                if pwd.len() > 3 {
                    let env_var_for_drive = Self::env_var_for_drive(drive_letter);
                    env.insert(env_var_for_drive, pwd);
                }
            }
        }
    }


Review comment (suggestion):
we typically don't allow `std::env::` since it's safer to use stack.get_env_var() and engine_state.get_env_var() and I don't think std::env::var is thread safe.

There are some instances where it's allowed but I'm not sure if this is one.

Code after:
                    env.insert(env_var_for_drive, pwd);
                }
            }
        }
    }

    /// Set the PWD for the drive letter in the absolute path.
    /// Return PathError for error.
    pub fn set_pwd(&mut self, path: &Path) -> Result<(), PathError> {
        if let (Some(drive_letter), Some(path_str)) =
            (Self::extract_drive_letter(path), path.to_str())
        {
            if drive_letter.is_ascii_alphabetic() {
                let drive_letter = drive_letter.to_ascii_uppercase();
                // Make sure saved drive letter is upper case
                let mut c = path_str.chars();
                match c.next() {
                    None => Err(PathError::InvalidDriveLetter),
                    Some(_) => {
                        let drive_index = drive_letter as usize - 'A' as usize;
                        let normalized_pwd = drive_letter.to_string() + c.as_str();
                        self.map[drive_index] = Some(normalized_pwd);
                        Ok(())
                    }
                }
            } else {
                Err(PathError::InvalidDriveLetter)
            }
        } else {
            Err(PathError::InvalidPath)
        }
    }

    /// Get the PWD for drive, if not yet, ask GetFullPathNameW() or omnipath,
    /// or else return default r"X:\".
    fn get_pwd(&self, drive_letter: char) -> Result<String, PathError> {
        if drive_letter.is_ascii_alphabetic() {
            let drive_letter = drive_letter.to_ascii_uppercase();
            let drive_index = drive_letter as usize - 'A' as usize;
            Ok(self.map[drive_index].clone().unwrap_or_else(|| {
                if let Some(sys_pwd) = get_full_path_name_w(&format!("{}:", drive_letter)) {
                    sys_pwd
                } else {
                    format!(r"{}:\", drive_letter)
                }
            }))
        } else {
            Err(PathError::InvalidDriveLetter)
        }
    }