From 01774208457ebb1953be71d97f117c5263fb492f Mon Sep 17 00:00:00 2001 From: Dan Kotov Date: Thu, 29 Aug 2024 10:14:46 -0400 Subject: [PATCH] chore: add wanted prop states python bindings --- src/parser/src/first_pass/parser_settings.rs | 11 ++ src/python/Cargo.lock | 10 +- src/python/src/lib.rs | 188 ++++++++++++++++++- 3 files changed, 201 insertions(+), 8 deletions(-) diff --git a/src/parser/src/first_pass/parser_settings.rs b/src/parser/src/first_pass/parser_settings.rs index 47809432..d2a31ad9 100644 --- a/src/parser/src/first_pass/parser_settings.rs +++ b/src/parser/src/first_pass/parser_settings.rs @@ -169,6 +169,17 @@ pub fn rm_user_friendly_names(names: &Vec) -> Result, DemoPa Ok(real_names) } +pub fn rm_map_user_friendly_names(map: &AHashMap) -> Result, DemoParserError> { + let mut real_names_map: AHashMap = AHashMap::default(); + for (name, variant) in map { + match FRIENDLY_NAMES_MAPPING.get(&name) { + Some(real_name) => real_names_map.insert(real_name.to_string(), variant.clone()), + None => return Err(DemoParserError::UnknownPropName(name.to_string())), + }; + } + Ok(real_names_map) +} + pub fn create_mmap(path: String) -> Result { let file = match File::open(path) { Err(e) => return Err(DemoParserError::FileNotFound(format!("{}", e))), diff --git a/src/python/Cargo.lock b/src/python/Cargo.lock index 896894e0..8f99c620 100644 --- a/src/python/Cargo.lock +++ b/src/python/Cargo.lock @@ -374,7 +374,7 @@ dependencies = [ [[package]] name = "demoparser2" -version = "0.26.2" +version = "0.32.0" dependencies = [ "ahash", "csgoproto", @@ -1318,9 +1318,9 @@ dependencies = [ [[package]] name = "protobuf" -version = "3.3.0" +version = "3.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b65f4a8ec18723a734e5dc09c173e0abf9690432da5340285d536edcb4dac190" +checksum = "df67496db1a89596beaced1579212e9b7c53c22dca1d9745de00ead76573d514" dependencies = [ "bytes", "once_cell", @@ -1330,9 +1330,9 @@ dependencies = [ [[package]] name = "protobuf-support" -version = "3.3.0" +version = "3.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6872f4d4f4b98303239a2b5838f5bbbb77b01ffc892d627957f37a22d7cfe69c" +checksum = "70e2d30ab1878b2e72d1e2fc23ff5517799c9929e2cf81a8516f9f4dcf2b9cf3" dependencies = [ "thiserror", ] diff --git a/src/python/src/lib.rs b/src/python/src/lib.rs index a28d6de0..1fd623c6 100644 --- a/src/python/src/lib.rs +++ b/src/python/src/lib.rs @@ -2,6 +2,7 @@ use ahash::AHashMap; use itertools::Itertools; use memmap2::Mmap; use parser::first_pass::parser_settings::create_mmap; +use parser::first_pass::parser_settings::rm_map_user_friendly_names; use parser::first_pass::parser_settings::rm_user_friendly_names; use parser::first_pass::parser_settings::ParserInputs; use parser::first_pass::read_bits::DemoParserError; @@ -9,6 +10,7 @@ use parser::parse_demo::Parser; use parser::second_pass::game_events::EventField; use parser::second_pass::game_events::GameEvent; use parser::second_pass::parser_settings::create_huffman_lookup_table; +use parser::second_pass::variants::Sticker; use parser::second_pass::variants::VarVec; use parser::second_pass::variants::Variant; #[cfg(feature = "voice")] @@ -35,6 +37,61 @@ use std::sync::Arc; use pyo3::create_exception; create_exception!(DemoParser, Exception, pyo3::exceptions::PyException); +struct PyVariant(Variant); +struct PySticker(Sticker); + +impl<'source> FromPyObject<'source> for PyVariant { + fn extract_bound(obj: &pyo3::Bound<'source, PyAny>) -> PyResult { + if let Ok(val) = obj.extract::() { + Ok(PyVariant(Variant::Bool(val))) + } else if let Ok(val) = obj.extract::() { + Ok(PyVariant(Variant::U32(val))) + } else if let Ok(val) = obj.extract::() { + Ok(PyVariant(Variant::I32(val))) + } else if let Ok(val) = obj.extract::() { + Ok(PyVariant(Variant::I16(val))) + } else if let Ok(val) = obj.extract::() { + Ok(PyVariant(Variant::F32(val))) + } else if let Ok(val) = obj.extract::() { + Ok(PyVariant(Variant::U64(val))) + } else if let Ok(val) = obj.extract::() { + Ok(PyVariant(Variant::U8(val))) + } else if let Ok(val) = obj.extract::() { + Ok(PyVariant(Variant::String(val))) + } else if let Ok(val) = obj.extract::<[f32; 2]>() { + Ok(PyVariant(Variant::VecXY(val))) + } else if let Ok(val) = obj.extract::<[f32; 3]>() { + Ok(PyVariant(Variant::VecXYZ(val))) + } else if let Ok(val) = obj.extract::>() { + Ok(PyVariant(Variant::StringVec(val))) + } else if let Ok(val) = obj.extract::>() { + Ok(PyVariant(Variant::U32Vec(val))) + } else if let Ok(val) = obj.extract::>() { + Ok(PyVariant(Variant::U64Vec(val))) + } else if let Ok(val) = obj.extract::>() { + Ok(PyVariant(Variant::Stickers( + val.into_iter().map(|ps| ps.0).collect(), + ))) + } else { + Err(PyValueError::new_err("Unsupported type for Variant")) + } + } +} + +impl<'source> FromPyObject<'source> for PySticker { + fn extract_bound(obj: &pyo3::Bound<'source, PyAny>) -> PyResult { + let dict = obj.extract::<&PyDict>()?; + let sticker = Sticker { + name: dict.get_item("name").unwrap().extract::()?, + wear: dict.get_item("wear").unwrap().extract::()?, + id: dict.get_item("id").unwrap().extract::()?, + x: dict.get_item("x").unwrap().extract::()?, + y: dict.get_item("y").unwrap().extract::()?, + }; + Ok(PySticker(sticker)) + } +} + #[pymethods] impl DemoParser { #[new] @@ -60,6 +117,8 @@ impl DemoParser { wanted_players: vec![], wanted_player_props: vec![], wanted_other_props: vec![], + wanted_player_prop_states: AHashMap::default(), + wanted_other_prop_states: AHashMap::default(), wanted_events: vec![], parse_ents: false, wanted_ticks: vec![], @@ -87,6 +146,8 @@ impl DemoParser { wanted_players: vec![], wanted_player_props: vec![], wanted_other_props: vec![], + wanted_player_prop_states: AHashMap::default(), + wanted_other_prop_states: AHashMap::default(), wanted_events: vec!["all".to_string()], parse_ents: false, wanted_ticks: vec![], @@ -121,6 +182,8 @@ impl DemoParser { wanted_players: vec![], wanted_player_props: vec![], wanted_other_props: vec![], + wanted_player_prop_states: AHashMap::default(), + wanted_other_prop_states: AHashMap::default(), wanted_events: vec![], parse_ents: true, wanted_ticks: vec![], @@ -191,6 +254,8 @@ impl DemoParser { wanted_players: vec![], wanted_player_props: vec![], wanted_other_props: vec![], + wanted_player_prop_states: AHashMap::default(), + wanted_other_prop_states: AHashMap::default(), wanted_events: vec![], parse_ents: false, wanted_ticks: vec![], @@ -235,6 +300,8 @@ impl DemoParser { wanted_players: vec![], wanted_player_props: vec![], wanted_other_props: vec![], + wanted_player_prop_states: AHashMap::default(), + wanted_other_prop_states: AHashMap::default(), wanted_events: vec![], parse_ents: false, wanted_ticks: vec![], @@ -315,6 +382,8 @@ impl DemoParser { wanted_players: vec![], wanted_player_props: vec![], wanted_other_props: vec![], + wanted_player_prop_states: AHashMap::default(), + wanted_other_prop_states: AHashMap::default(), wanted_events: vec![], parse_ents: false, wanted_ticks: vec![], @@ -379,18 +448,32 @@ impl DemoParser { }) } - #[pyo3(signature = (event_name, *, player=None, other=None))] + #[pyo3(signature = (event_name, *, player=None, other=None, player_states=None, other_states=None))] pub fn parse_event( &self, py: Python<'_>, event_name: String, player: Option>, other: Option>, + player_states: Option>, + other_states: Option>, ) -> PyResult> { let wanted_player_props = player.unwrap_or_default(); let wanted_other_props = other.unwrap_or_default(); + let wanted_player_states = player_states + .unwrap_or_default() + .into_iter() + .map(|(k, v)| (k, v.0)) + .collect(); + let wanted_other_states = other_states + .unwrap_or_default() + .into_iter() + .map(|(k, v)| (k, v.0)) + .collect(); let real_player_props = rm_user_friendly_names(&wanted_player_props); let real_other_props = rm_user_friendly_names(&wanted_other_props); + let real_wanted_player_states = rm_map_user_friendly_names(&wanted_player_states); + let real_wanted_other_states = rm_map_user_friendly_names(&wanted_other_states); let real_player_props = match real_player_props { Ok(real_props) => real_props, @@ -400,6 +483,15 @@ impl DemoParser { Ok(real_props) => real_props, Err(e) => return Err(PyValueError::new_err(format!("{e}"))), }; + let real_wanted_player_states = match real_wanted_player_states { + Ok(real_wanted_player_states) => real_wanted_player_states, + Err(e) => return Err(Exception::new_err(format!("{e}"))), + }; + let real_wanted_other_states = match real_wanted_other_states { + Ok(real_wanted_other_states) => real_wanted_other_states, + Err(e) => return Err(Exception::new_err(format!("{e}"))), + }; + let mut real_name_to_og_name = AHashMap::default(); for (real_name, user_friendly_name) in real_player_props.iter().zip(&wanted_player_props) { real_name_to_og_name.insert(real_name.clone(), user_friendly_name.clone()); @@ -407,6 +499,18 @@ impl DemoParser { for (real_name, user_friendly_name) in real_other_props.iter().zip(&wanted_other_props) { real_name_to_og_name.insert(real_name.clone(), user_friendly_name.clone()); } + for (real_name, user_friendly_name) in real_wanted_player_states + .keys() + .zip(wanted_player_states.keys()) + { + real_name_to_og_name.insert(real_name.clone(), user_friendly_name.clone()); + } + for (real_name, user_friendly_name) in real_wanted_other_states + .keys() + .zip(wanted_other_states.keys()) + { + real_name_to_og_name.insert(real_name.clone(), user_friendly_name.clone()); + } let settings = ParserInputs { real_name_to_og_name, @@ -414,6 +518,8 @@ impl DemoParser { wanted_player_props: real_player_props, wanted_other_props: real_other_props, wanted_events: vec![event_name], + wanted_player_prop_states: real_wanted_player_states, + wanted_other_prop_states: real_wanted_other_states, parse_ents: true, wanted_ticks: vec![], parse_projectiles: false, @@ -435,18 +541,32 @@ impl DemoParser { Ok(event_series) } - #[pyo3(signature = (event_name, *, player=None, other=None))] + #[pyo3(signature = (event_name, *, player=None, other=None, player_states=None, other_states=None))] pub fn parse_events( &self, py: Python<'_>, event_name: Vec, player: Option>, other: Option>, + player_states: Option>, + other_states: Option>, ) -> PyResult> { let wanted_player_props = player.unwrap_or_default(); let wanted_other_props = other.unwrap_or_default(); + let wanted_player_states = player_states + .unwrap_or_default() + .into_iter() + .map(|(k, v)| (k, v.0)) + .collect(); + let wanted_other_states = other_states + .unwrap_or_default() + .into_iter() + .map(|(k, v)| (k, v.0)) + .collect(); let real_player_props = rm_user_friendly_names(&wanted_player_props); let real_other_props = rm_user_friendly_names(&wanted_other_props); + let real_wanted_player_states = rm_map_user_friendly_names(&wanted_player_states); + let real_wanted_other_states = rm_map_user_friendly_names(&wanted_other_states); let real_player_props = match real_player_props { Ok(real_props) => real_props, @@ -456,6 +576,15 @@ impl DemoParser { Ok(real_props) => real_props, Err(e) => return Err(PyValueError::new_err(format!("{e}"))), }; + let real_wanted_player_states = match real_wanted_player_states { + Ok(real_wanted_player_states) => real_wanted_player_states, + Err(e) => return Err(Exception::new_err(format!("{e}"))), + }; + let real_wanted_other_states = match real_wanted_other_states { + Ok(real_wanted_other_states) => real_wanted_other_states, + Err(e) => return Err(Exception::new_err(format!("{e}"))), + }; + let mut real_name_to_og_name = AHashMap::default(); for (real_name, user_friendly_name) in real_player_props.iter().zip(&wanted_player_props) { real_name_to_og_name.insert(real_name.clone(), user_friendly_name.clone()); @@ -463,6 +592,18 @@ impl DemoParser { for (real_name, user_friendly_name) in real_other_props.iter().zip(&wanted_other_props) { real_name_to_og_name.insert(real_name.clone(), user_friendly_name.clone()); } + for (real_name, user_friendly_name) in real_wanted_player_states + .keys() + .zip(wanted_player_states.keys()) + { + real_name_to_og_name.insert(real_name.clone(), user_friendly_name.clone()); + } + for (real_name, user_friendly_name) in real_wanted_other_states + .keys() + .zip(wanted_other_states.keys()) + { + real_name_to_og_name.insert(real_name.clone(), user_friendly_name.clone()); + } let settings = ParserInputs { real_name_to_og_name, @@ -470,6 +611,8 @@ impl DemoParser { wanted_player_props: real_player_props, wanted_other_props: real_other_props, wanted_events: event_name, + wanted_player_prop_states: real_wanted_player_states, + wanted_other_prop_states: real_wanted_other_states, parse_ents: true, wanted_ticks: vec![], parse_projectiles: false, @@ -496,6 +639,8 @@ impl DemoParser { wanted_players: vec![], wanted_player_props: vec![], wanted_other_props: vec![], + wanted_player_prop_states: AHashMap::default(), + wanted_other_prop_states: AHashMap::default(), wanted_events: vec![], wanted_ticks: vec![], real_name_to_og_name: AHashMap::default(), @@ -521,34 +666,71 @@ impl DemoParser { Ok(out_hm.to_object(py)) } - #[pyo3(signature = (wanted_props, *, players=None, ticks=None))] + #[pyo3(signature = (wanted_props, *, players=None, ticks=None, player_states=None, other_states=None))] pub fn parse_ticks( &self, py: Python, wanted_props: Vec, players: Option>, ticks: Option>, + player_states: Option>, + other_states: Option>, ) -> PyResult { let wanted_players = players.unwrap_or_default(); let wanted_ticks = ticks.unwrap_or_default(); + let wanted_player_states = player_states + .unwrap_or_default() + .into_iter() + .map(|(k, v)| (k, v.0)) + .collect(); + let wanted_other_states = other_states + .unwrap_or_default() + .into_iter() + .map(|(k, v)| (k, v.0)) + .collect(); let real_props = rm_user_friendly_names(&wanted_props); + let real_wanted_player_states = rm_map_user_friendly_names(&wanted_player_states); + let real_wanted_other_states = rm_map_user_friendly_names(&wanted_other_states); let real_props = match real_props { Ok(real_props) => real_props, Err(e) => return Err(Exception::new_err(format!("{e}"))), }; + let real_wanted_player_states = match real_wanted_player_states { + Ok(real_wanted_player_states) => real_wanted_player_states, + Err(e) => return Err(Exception::new_err(format!("{e}"))), + }; + let real_wanted_other_states = match real_wanted_other_states { + Ok(real_wanted_other_states) => real_wanted_other_states, + Err(e) => return Err(Exception::new_err(format!("{e}"))), + }; let arc_huf = Arc::new(&self.huf); let mut real_name_to_og_name = AHashMap::default(); for (real_name, user_friendly_name) in real_props.iter().zip(&wanted_props) { real_name_to_og_name.insert(real_name.clone(), user_friendly_name.clone()); } + for (real_name, user_friendly_name) in real_wanted_player_states + .keys() + .zip(wanted_player_states.keys()) + { + real_name_to_og_name.insert(real_name.clone(), user_friendly_name.clone()); + } + for (real_name, user_friendly_name) in real_wanted_other_states + .keys() + .zip(wanted_other_states.keys()) + { + real_name_to_og_name.insert(real_name.clone(), user_friendly_name.clone()); + } + let settings = ParserInputs { real_name_to_og_name, wanted_players, wanted_player_props: real_props, wanted_other_props: vec![], wanted_events: vec![], + wanted_player_prop_states: real_wanted_player_states, + wanted_other_prop_states: real_wanted_other_states, parse_ents: true, wanted_ticks, parse_projectiles: false,