1use std::env;2use std::path::{Path, PathBuf};34use clap::{ArgMatches, Command, arg, crate_version};5use mdbook_driver::MDBook;6use mdbook_driver::errors::Result as Result3;7use mdbook_i18n_helpers::preprocessors::Gettext;8use mdbook_spec::Spec;9use mdbook_trpl::{Figure, Listing, Note};1011fn main() {12    let crate_version = concat!("v", crate_version!());13    let filter = tracing_subscriber::EnvFilter::builder()14        .with_env_var("MDBOOK_LOG")15        .with_default_directive(tracing_subscriber::filter::LevelFilter::INFO.into())16        .from_env_lossy();17    tracing_subscriber::fmt()18        .without_time()19        .with_ansi(std::io::IsTerminal::is_terminal(&std::io::stderr()))20        .with_writer(std::io::stderr)21        .with_env_filter(filter)22        .with_target(std::env::var_os("MDBOOK_LOG").is_some())23        .init();2425    // env_logger::Builder::from_env(env_logger::Env::default().default_filter_or("info")).init();26    let d_arg = arg!(-d --"dest-dir" <DEST_DIR>27"The output directory for your book\n(Defaults to ./book when omitted)")28    .required(false)29    .value_parser(clap::value_parser!(PathBuf));3031    let l_arg = arg!(-l --"lang" <LANGUAGE>32"The output language")33    .required(false)34    .value_parser(clap::value_parser!(String));3536    let root_arg = arg!(--"rust-root" <ROOT_DIR>37"Path to the root of the rust source tree")38    .required(false)39    .value_parser(clap::value_parser!(PathBuf));4041    let dir_arg = arg!([dir] "Root directory for the book\n\42                              (Defaults to the current directory when omitted)")43    .value_parser(clap::value_parser!(PathBuf));4445    // Note: we don't parse this into a `PathBuf` because it is comma separated46    // strings *and* we will ultimately pass it into `MDBook::test()`, which47    // accepts `Vec<&str>`. Although it is a bit annoying that `-l/--lang` and48    // `-L/--library-path` are so close, this is the same set of arguments we49    // would pass when invoking mdbook on the CLI, so making them match when50    // invoking rustbook makes for good consistency.51    let library_path_arg = arg!(52        -L --"library-path" <PATHS>53        "A comma-separated list of directories to add to the crate search\n\54        path when building tests"55    )56    .required(false)57    .value_parser(parse_library_paths);5859    let matches = Command::new("rustbook")60        .about("Build a book with mdBook")61        .author("Steve Klabnik <steve@steveklabnik.com>")62        .version(crate_version)63        .subcommand_required(true)64        .arg_required_else_help(true)65        .subcommand(66            Command::new("build")67                .about("Build the book from the markdown files")68                .arg(d_arg)69                .arg(l_arg)70                .arg(root_arg)71                .arg(&dir_arg),72        )73        .subcommand(74            Command::new("test")75                .about("Tests that a book's Rust code samples compile")76                .arg(dir_arg)77                .arg(library_path_arg),78        )79        .get_matches();8081    // Check which subcommand the user ran...82    match matches.subcommand() {83        Some(("build", sub_matches)) => {84            if let Err(e) = build(sub_matches) {85                handle_error(e);86            }87        }88        Some(("test", sub_matches)) => {89            if let Err(e) = test(sub_matches) {90                handle_error(e);91            }92        }93        _ => unreachable!(),94    };95}9697fn build(args: &ArgMatches) -> Result3<()> {98    let book_dir = get_book_dir(args);99    let dest_dir = args.get_one::<PathBuf>("dest-dir");100    let lang = args.get_one::<String>("lang");101    let rust_root = args.get_one::<PathBuf>("rust-root");102    let book = load_book(&book_dir, dest_dir, lang, rust_root.cloned())?;103    book.build()104}105106fn test(args: &ArgMatches) -> Result3<()> {107    let book_dir = get_book_dir(args);108    let mut book = load_book(&book_dir, None, None, None)?;109    let library_paths = args110        .try_get_one::<Vec<String>>("library-path")?111        .map(|v| v.iter().map(|s| s.as_str()).collect::<Vec<&str>>())112        .unwrap_or_default();113    book.test(library_paths)114}115116fn get_book_dir(args: &ArgMatches) -> PathBuf {117    if let Some(p) = args.get_one::<PathBuf>("dir") {118        // Check if path is relative from current dir, or absolute...119        if p.is_relative() { env::current_dir().unwrap().join(p) } else { p.to_path_buf() }120    } else {121        env::current_dir().unwrap()122    }123}124125fn load_book(126    book_dir: &Path,127    dest_dir: Option<&PathBuf>,128    lang: Option<&String>,129    rust_root: Option<PathBuf>,130) -> Result3<MDBook> {131    let mut book = MDBook::load(book_dir)?;132    book.config.set("output.html.input-404", "").unwrap();133    book.config.set("output.html.hash-files", true).unwrap();134135    if let Some(lang) = lang {136        let gettext = Gettext;137        book.with_preprocessor(gettext);138        book.config.set("book.language", lang).unwrap();139    }140141    // Set this to allow us to catch bugs in advance.142    book.config.build.create_missing = false;143144    if let Some(dest_dir) = dest_dir {145        book.config.build.build_dir = dest_dir.into();146    }147148    // NOTE: Replacing preprocessors using this technique causes error149    // messages to be displayed when the original preprocessor doesn't work150    // (but it otherwise succeeds).151    //152    // This should probably be fixed in mdbook to remove the existing153    // preprocessor, or this should modify the config and use154    // MDBook::load_with_config.155    if book.config.contains_key("preprocessor.trpl-note") {156        book.with_preprocessor(Note);157    }158159    if book.config.contains_key("preprocessor.trpl-listing") {160        book.with_preprocessor(Listing);161    }162163    if book.config.contains_key("preprocessor.trpl-figure") {164        book.with_preprocessor(Figure);165    }166167    if book.config.contains_key("preprocessor.spec") {168        book.with_preprocessor(Spec::new(rust_root)?);169    }170171    Ok(book)172}173174fn parse_library_paths(input: &str) -> Result<Vec<String>, String> {175    Ok(input.split(",").map(String::from).collect())176}177178fn handle_error(error: mdbook_driver::errors::Error) -> ! {179    eprintln!("Error: {}", error);180181    for cause in error.chain().skip(1) {182        eprintln!("\tCaused By: {}", cause);183    }184185    std::process::exit(101);186}