v0.2.2

CHANGELOG.md

@@ -1,5 +1,12 @@
 # Changelog
 
+## v0.2.2
+- Anonymization: output each format into its own directory
+- Support saving/loading anonymizer state (mapping of players <=> ID numbers)
+
+## v0.2.1
+- Fix a bug when parsing ELOs.
+
 ## v0.2.0
 - Ties are now properly anonymized.
 - ELO is now included in anonymized logs, but is rounded to the nearest 50 ELO.

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "psbattletools"
-version = "0.2.0"
+version = "0.2.2"
 edition = "2021"
 description = "Command-line program to manage Pokémon Showdown battle logs."
 license = "MIT"
@@ -17,6 +17,9 @@ lazy_static = "1.4.0"
 prettytable-rs = "0.8.0"
 rayon = "1.5.1"
 regex = "1.5.4"
+serde = "1.0.136"
+serde_derive = "1.0.136"
+serde_json = "1.0.79"
 structopt = "0.3.23"
 
 [dev-dependencies]

src/anonymize/anonymizer.rs

@@ -1,5 +1,6 @@
 // From https://github.com/AnnikaCodes/anonbattle/blob/main/src/anonymizer.rs
 
+use serde_derive::{Deserialize, Serialize};
 use std::collections::HashMap;
 use std::str::FromStr;
 use std::sync::Mutex;
@@ -14,6 +15,7 @@ lazy_static! {
 }
 
 /// Tracks players
+#[derive(Serialize, Deserialize)]
 struct SharedState {
     players: HashMap<String, String>,
     pub current_battle_number: u32,
@@ -40,6 +42,10 @@ impl SharedState {
             }
         }
     }
+
+    fn to_json(&self) -> serde_json::Result<String> {
+        serde_json::to_string(&self)
+    }
 }
 
 /// Anonymizes string JSON while tracking state
@@ -59,10 +65,19 @@ impl Anonymizer {
         }
     }
 
+    pub fn with_json(json: String, is_safe: bool, no_log: bool) -> Self {
+        let state: SharedState = serde_json::from_str(&json).unwrap();
+        Self {
+            state: Mutex::new(state),
+            is_safe,
+            no_log,
+        }
+    }
+
     /// Anonymizes a log.
     ///
-    /// Returns a tuple: (json, battle_number)
-    pub fn anonymize(&self, raw: &str) -> Result<(String, u32), BattleToolsError> {
+    /// Returns a tuple: (json, battle_number, format)
+    pub fn anonymize(&self, raw: &str) -> Result<(String, u32, String), BattleToolsError> {
         let json = json::parse(raw)?;
 
         let p1 = json["p1"]
@@ -254,7 +269,12 @@ impl Anonymizer {
             tracker.current_battle_number += 1;
             tracker.current_battle_number
         };
-        Ok((result, battle_number))
+        Ok((result, battle_number, json["format"].to_string()))
+    }
+
+    // TODO: actually save this & load it
+    pub fn get_state_json(&self) -> serde_json::Result<String> {
+        self.state.lock().unwrap().to_json()
     }
 }
 
@@ -310,7 +330,7 @@ mod unit_tests {
     #[test]
     pub fn anonymization() {
         let anonymizer = Anonymizer::new(true, false);
-        let (json, _) = anonymizer.anonymize(&SAMPLE_JSON).unwrap();
+        let (json, _, _) = anonymizer.anonymize(&SAMPLE_JSON).unwrap();
         assert_ne!(json, *SAMPLE_JSON);
 
         for term in ["00:00:01", "Annika", "annika", "Rust Haters", "rusthaters"] {
@@ -338,7 +358,7 @@ mod unit_tests {
     #[test]
     pub fn tie() {
         let anonymizer = Anonymizer::new(true, false);
-        let (json, _) = anonymizer.anonymize(&TIE_WINNERSTRING_JSON).unwrap();
+        let (json, _, _) = anonymizer.anonymize(&TIE_WINNERSTRING_JSON).unwrap();
         assert_eq!(
             gjson::get(&json, "winner").to_string(),
             "".to_string(),
@@ -350,8 +370,8 @@ mod unit_tests {
     pub fn no_log() {
         let anonymizer_logs = Anonymizer::new(false, false);
         let anonymizer_no_logs = Anonymizer::new(false, true);
-        let (logs_json, _) = anonymizer_logs.anonymize(&SAMPLE_JSON).unwrap();
-        let (no_logs_json, _) = anonymizer_no_logs.anonymize(&SAMPLE_JSON).unwrap();
+        let (logs_json, _, _) = anonymizer_logs.anonymize(&SAMPLE_JSON).unwrap();
+        let (no_logs_json, _, _) = anonymizer_no_logs.anonymize(&SAMPLE_JSON).unwrap();
 
         for should_be_arr in ["log", "inputLog"] {
             assert!(

src/anonymize/mod.rs

@@ -4,7 +4,7 @@ mod anonymizer;
 use std::path::PathBuf;
 
 use crate::{directory::LogParser, BattleToolsError};
-use anonymizer::Anonymizer;
+pub use anonymizer::Anonymizer;
 pub struct AnonymizingDirectoryParser {
     anonymizer: Anonymizer,
     output_directory: PathBuf,
@@ -17,6 +17,17 @@ impl AnonymizingDirectoryParser {
             output_directory,
         }
     }
+
+    pub fn with_anonymizer(anonymizer: Anonymizer, output_directory: PathBuf) -> Self {
+        Self {
+            anonymizer,
+            output_directory,
+        }
+    }
+
+    pub fn get_state_json(&self) -> serde_json::Result<String> {
+        self.anonymizer.get_state_json()
+    }
 }
 
 impl LogParser<()> for AnonymizingDirectoryParser {
@@ -25,8 +36,10 @@ impl LogParser<()> for AnonymizingDirectoryParser {
         raw_json: String,
         _: &std::path::Path,
     ) -> Result<(), BattleToolsError> {
-        let (json, battle_num) = self.anonymizer.anonymize(&raw_json)?;
+        let (json, battle_num, directory) = self.anonymizer.anonymize(&raw_json)?;
         let mut out_file = self.output_directory.clone();
+        out_file.push(directory);
+        std::fs::create_dir_all(&out_file)?;
         out_file.push(format!("{}.log.json", battle_num));
         std::fs::write(out_file, json)?;
         Ok(())

src/main.rs

@@ -110,6 +110,16 @@ struct Options {
         help = "The maximum number of threads to use for concurrent processing"
     )]
     threads: Option<usize>,
+    #[structopt(
+        long = "save-state-to",
+        help = "Save the state of the anonymizer to this file"
+    )]
+    save_state_to: Option<PathBuf>,
+    #[structopt(
+        long = "load-state-from",
+        help = "Load the state of the anonymizer from this file"
+    )]
+    load_state_from: Option<PathBuf>,
 }
 
 #[derive(Debug)]
@@ -121,7 +131,9 @@ pub enum BattleToolsError {
     InvalidLog(String),
     PathConversion(String),
     IncompleteAnonymization(String),
+    SerializationError(serde_json::Error),
 }
+
 impl From<std::io::Error> for BattleToolsError {
     fn from(error: std::io::Error) -> Self {
         BattleToolsError::IOError(error)
@@ -147,6 +159,11 @@ impl From<rayon::ThreadPoolBuildError> for BattleToolsError {
         BattleToolsError::ThreadPoolError(error)
     }
 }
+impl From<serde_json::Error> for BattleToolsError {
+    fn from(error: serde_json::Error) -> Self {
+        BattleToolsError::SerializationError(error)
+    }
+}
 
 fn main() -> Result<(), BattleToolsError> {
     let options = Options::from_args();
@@ -200,8 +217,21 @@ fn main() -> Result<(), BattleToolsError> {
         } => {
             // create dir if needed
             fs::create_dir_all(&output_dir)?;
-            let mut anonymizer = AnonymizingDirectoryParser::new(is_safe, no_log, output_dir);
+
+            let mut anonymizer = if let Some(load_path) = options.load_state_from {
+                let json = fs::read_to_string(load_path)?;
+                let anonymizer = anonymize::Anonymizer::with_json(json, is_safe, no_log);
+                AnonymizingDirectoryParser::with_anonymizer(anonymizer, output_dir)
+            } else {
+                AnonymizingDirectoryParser::new(is_safe, no_log, output_dir)
+            };
+
             anonymizer.handle_directories(directories, options.exclude)?;
+
+            if let Some(save_state_path) = options.save_state_to {
+                let json = anonymizer.get_state_json()?;
+                fs::write(save_state_path, json)?;
+            }
         }
     }
 

tests/anonymize.rs

@@ -24,12 +24,12 @@ fn test_no_terms_and_identicality() {
     assert!(output.status.success(), "command failed");
 
     let mut out_file_1 = out_dir.clone();
-    out_file_1.push("1.log.json");
+    out_file_1.push("gen8randombattle/1.log.json");
     let output_1 = std::fs::read_to_string(&out_file_1)
         .unwrap_or_else(|_| panic!("Couldn't read output file {:?}", out_file_1));
 
     let mut out_file_999 = out_dir;
-    out_file_999.push("999.log.json");
+    out_file_999.push("gen8randombattle/999.log.json");
     let output_999 = std::fs::read_to_string(out_file_999).expect("Couldn't read output file");
 
     for term in ["00:00:01", "Annika", "annika", "Rust Haters", "rusthaters"] {