poe_data_tools_cli/commands/
dump_trees.rs

1use std::{
2    collections::HashMap,
3    fs::{File, create_dir_all},
4    io::BufWriter,
5    path::Path,
6};
7
8use anyhow::{Context, Result, ensure};
9use glob::{MatchOptions, Pattern};
10use poe_data_tools::{
11    Patch,
12    file_parsers::{
13        FileParser,
14        psg::{PSGParser, types::PSGFile},
15    },
16    fs::{FS, FileSystem},
17    tree::{
18        passive_info::{PassiveSkillInfo, load_passive_info},
19        psg::PassiveSkillGraph,
20    },
21};
22
23use crate::VERBOSE;
24
25fn process_file(
26    contents: &[u8],
27    output_path: &Path,
28    version: &Patch,
29    passive_info: &HashMap<u16, PassiveSkillInfo>,
30) -> Result<()> {
31    // Parse the PSG file
32    let psg_file = PSGParser {
33        version: version.major(),
34    }
35    .parse(contents)
36    .as_anyhow()
37    .context("Failed to parse passive skill tree")?;
38
39    // Add passive info - only nodes that are in the graph
40    let passive_info = {
41        let ids = psg_file
42            .groups
43            .iter()
44            .flat_map(|g| g.passives.iter().map(|p| p.id as u16));
45        ids.map(|id| (id, passive_info[&id].clone())).collect()
46    };
47
48    let passive_tree = {
49        let PSGFile {
50            version,
51            graph_type,
52            passives_per_orbit,
53            root_passives,
54            groups,
55        } = psg_file;
56
57        PassiveSkillGraph {
58            version,
59            graph_type,
60            passives_per_orbit,
61            root_passives,
62            groups,
63            passive_info,
64        }
65    };
66
67    // Write to file
68    create_dir_all(output_path.parent().context("No parent directory")?)
69        .context("Failed to create output dirs")?;
70
71    let f = File::create(output_path)
72        .with_context(|| format!("Failed to create file {:?}", output_path))?;
73    let f = BufWriter::new(f);
74
75    serde_json::to_writer_pretty(f, &passive_tree).context("Failed to serialise tree to JSON")
76}
77
78pub fn dump_trees(
79    fs: &mut FS,
80    patterns: &[Pattern],
81    output_folder: &Path,
82    version: &Patch,
83    cache_dir: &Path,
84) -> Result<()> {
85    for pattern in patterns {
86        ensure!(
87            pattern.as_str().ends_with(".psg"),
88            "Only .psg tree export is supported."
89        );
90    }
91
92    let filenames = fs
93        .list()
94        .filter(|filename| {
95            patterns.iter().any(|pattern| {
96                pattern.matches_with(
97                    filename,
98                    MatchOptions {
99                        require_literal_separator: true,
100                        ..Default::default()
101                    },
102                )
103            })
104        })
105        .collect::<Vec<_>>();
106
107    let passive_info = load_passive_info(fs, version, cache_dir)?
108        .into_iter()
109        .map(|p| (p.graph_passive_id, p))
110        .collect::<HashMap<_, _>>();
111
112    fs.batch_read(&filenames)
113        // Print and filter out errors
114        .filter_map(|(path, res)| match res {
115            Ok(b) => Some((path, b)),
116            Err(e) => {
117                log::error!("Failed to extract file: {:?}: {:?}", path, e);
118                None
119            }
120        })
121        // Attempt to read file contents
122        .map(|(filename, contents)| -> Result<_, anyhow::Error> {
123            // Convert the data table
124            let output_path = output_folder.join(filename.as_ref()).with_extension("json");
125            process_file(&contents, &output_path, version, &passive_info)
126                .with_context(|| format!("Failed to process file: {:?}", filename))?;
127
128            Ok(filename)
129        })
130        // Report results
131        .for_each(|result| match result {
132            Ok(filename) => log::info!("Extracted tree: {}", filename),
133            Err(e) => {
134                let error_message = if *VERBOSE.get().unwrap() {
135                    format!("{e:?}")
136                } else {
137                    format!("{e}")
138                };
139                log::error!("Failed to extract tree: {error_message}");
140            }
141        });
142
143    Ok(())
144}
poe_data_tools_cli/commands/dump_trees.rs

poe_data_tools_cli/commands/
dump_trees.rs
