How to improve this snippet?

[minimum]

The below snippet detect the venv name of the Python virtual environment created by .direnv and update the config of Pyright language server.

E.g. If a /Users/user/git/rust/project/.direnv/python-3.10.6/bin/python executable is found, it update the config /Users/user/git/rust/project/pyrightconfig.json with value

{
  "venv": "python-3.10.6",
  "venvPath": ".direnv"
}

The Python version of this is here, with no nearly error handling.

I'd like to have proper error handling in the Rust snippet, and simplify the code if possible.

Below is the Rust snippet.

use std::{
    fs::File,
    path::{Component, PathBuf},
    process::Command,
};

use serde_json::{json, Value};

fn main() -> Result<(), anyhow::Error> {
    let output = String::from_utf8(Command::new("which").arg("python").output()?.stdout)?;
    let python = output.trim();
    if python.is_empty() {
        anyhow::bail!("failed to determine Python executable");
    }

    let python = PathBuf::from(python);
    if python.components().rev().nth(3) != Some(Component::Normal(".direnv".as_ref())) {
        anyhow::bail!(".direnv not detected");
    }

    if let Some(project_root) = python.ancestors().nth(4) {
        if let Some(venv) = python.ancestors().nth(2) {
            if let Some(venv) = venv.file_name() {
                if let Some(venv) = venv.to_str() {
                    let config_file = project_root.join("pyrightconfig.json");
                    let mut config: Value;
                    if config_file.exists() {
                        config = serde_json::from_reader(File::open(config_file.clone())?)?;
                    } else {
                        config = json!({});
                    }

                    config["venv"] = Value::String(venv.to_string());
                    config["venvPath"] = Value::String(".direnv".to_string());

                    serde_json::to_writer_pretty(File::create(config_file)?, &config)?;
                } else {
                    anyhow::bail!("invalid venv name: {:?}", venv);
                }
            } else {
                anyhow::bail!("invalid venv name: {:?}", venv.file_name());
            }
        } else {
            anyhow::bail!("failed to detect venv name")
        }
    } else {
        anyhow::bail!("failed to detect project root");
    }

    Ok(())
}

[Dushistov]

I would be use anyhow::ensure! instead of if { return Err, like this:

anyhow::ensure!(!python.is_empty(), "failed to determine Python executable");

and also nested if looks ugly,
you can use Option::ok_or_else in this case:

 let project_root = python
        .ancestors()
        .nth(4)
        .ok_or_else(|| anyhow::anyhow!("failed to detect project root"))?;
    let venv = python
        .ancestors()
        .nth(2)
        .ok_or_else(|| anyhow::anyhow!("failed to detect venv name"))?;

also you can assign config only once:

let mut config: Value = if config_file.exists() {
    serde_json::from_reader(File::open(config_file.clone())?)?
} else {
   json!({})
};

And no need to clone PathBuf, you can just use reference here: serde_json::from_reader(File::open(&config_file)?)?

[vague]

pub mod pyrightconfig {
    use crate::Result;
    use std::path::{Path, PathBuf};

    pub const JSON_FN: &str = "pyrightconfig.json";
    pub const DIRENV: &str = ".direnv";

    pub fn create(path: &Path) -> Result<()> {
        anyhow::ensure!(
            path.to_str().map(|s| s.contains(DIRENV)).unwrap_or(false),
            "{path:?} doesn't contain {DIRENV}"
        );

        let (path, py) = pyrightconfig_path_and_venv(path);
        if let Some(venv) = py.as_deref() {
            write_pyrightconfig(&path, venv)?;
            Ok(())
        } else {
            anyhow::bail!("the version of py venv is not found in the given path {path:?}")
        }
    }

    fn pyrightconfig_path_and_venv(path: &Path) -> (PathBuf, Option<String>) {
        let mut comp = path.components();
        let mut prefix: PathBuf = comp
            .by_ref()
            .take_while(|p| p.as_os_str() != DIRENV)
            .collect();
        prefix.push(JSON_FN);

        // python version
        let py = comp
            .next()
            .and_then(|c| c.as_os_str().to_str().map(String::from));

        (prefix, py)
    }

    fn write_pyrightconfig(path: &Path, venv: &str) -> Result<()> {
        use std::{fs::File, io::Write};

        let mut json: serde_json::Value = File::open(path)
            .ok()
            .and_then(|file| serde_json::from_reader(file).ok())
            .unwrap_or_default();
        json["venv"] = venv.to_owned().into();
        json["venvPath"] = DIRENV.to_owned().into();
        write!(std::fs::File::create(path)?, "{json}")?;
        Ok(())
    }
}

Full code with tests: Rust Playground

[nerditation]

you can use the anyhow::Context extension trait to replace the bail!() statements under nested if let on Result and Option types:

    use anyhow::Context;
    //...
    let project_root = python.ancestors().nth(4).context("failed to detect project root")?;
    let venv = python.ancestors().nth(2).context("failed to detect venv name")?;
    let venv = venv.file_name().with_context(|| format!("invalid venv name: {:?}", venv.file_name()))?;
    let venv = venv.to_str().with_context(|| format!("invalid venv name: {:?}", venv))?;
    //...

you can also chain those venv conversions, but you lose the ability to format error messages using the intermidate values.

    let venv = python
        .ancestors()
        .nth(2)
        .context("failed to detect venv name")?
        .file_name()
        .context("invalid venv name")?
        .to_str()
        .context("invalid venv name")?;

alternatively, you can use ok_or() (lazily: ok_or_else()) to turn Option<T> into anyhow::Result<T> without using anyhow::Context, but it's more verbose.

    let project_root = python.ancestors().nth(4).ok_or(anyhow!("failed to detect project root"))?;

or if you somehow prefer bail!(), let else is a thing now:

    let Some(project_root) = python.ancestors().nth(4) else {
        bail!("failed to detect project root");
    };
    let Some(venv) = python.ancestors().nth(2) else {
        bail!("failed to detect venv name");
    };
    //...

[minimum]

Oh, I add that clone because the borrow checker tells me the variable is moved when I try to use it again to write back to the file later. Didn't notice that I can use reference here. Thank you!

[minimum]

Thank you for so many methods to flatten the if statements. The code looks pretty better now!

[minimum]

Thank you for the complete example! I did forget that dividing the code to small functions can help in this case!