Add __match_args__ for structs

VirxEC · Jun 28, 2024 · db80588 · db80588
1 parent 0b9cb3b
commit db80588
Show file tree

Hide file tree

Showing 6 changed files with 88 additions and 13 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -1,13 +1,13 @@
 [package]
 name = "rlbot-flatbuffers-py"
-version = "0.4.0"
+version = "0.4.1"
 edition = "2021"
 description = "A Python module implemented in Rust for serializing and deserializing RLBot's flatbuffers"
 repository = "https://github.com/VirxEC/rlbot_flatbuffers_py"
 build = "codegen/main.rs"
 license = "MIT"
 readme = "README.md"
-exclude = [".github", "pytest.py", "rustfmt.toml", ".gitignore", ".gitmodules", "flatc_mac"]
+exclude = [".github", "pytest.py", "pybench.py", "rustfmt.toml", ".gitignore", ".gitmodules", "flatc_mac"]
 publish = false
 
 [lints.clippy]

diff --git a/codegen/enums.rs b/codegen/enums.rs
@@ -324,8 +324,6 @@ impl Generator for EnumBindGenerator {
         write_str!(self, "");
 
         self.generate_repr_method();
-        write_str!(self, "");
-
         write_str!(self, "}");
         write_str!(self, "");
     }

diff --git a/codegen/pyi.rs b/codegen/pyi.rs
@@ -174,6 +174,16 @@ pub fn generator(type_data: &[PythonBindType]) -> io::Result<()> {
                     }
                 }
 
+                if !gen.types.is_empty() {
+                    write_str!(file, "");
+                    write_str!(file, "    __match_args__ = (");
+
+                    for variable_info in &gen.types {
+                        write_fmt!(file, "        \"{}\",", variable_info.name);
+                    }
+                    write_str!(file, "    )");
+                }
+
                 if gen.types.is_empty() {
                     write_str!(file, "    def __init__(self): ...");
                 } else {

diff --git a/codegen/structs.rs b/codegen/structs.rs
@@ -1,5 +1,5 @@
 use crate::{generator::Generator, PythonBindType};
-use std::{borrow::Cow, fs, path::Path};
+use std::{borrow::Cow, fs, iter::repeat, path::Path};
 
 #[derive(Debug, PartialEq, Eq)]
 pub enum InnerVecType {
@@ -388,7 +388,10 @@ impl StructBindGenerator {
             return;
         }
 
-        write_str!(self, "    #[allow(unused_variables)]");
+        if !self.frozen_needs_py {
+            write_str!(self, "    #[allow(unused_variables)]");
+        }
+
         write_str!(self, "    pub fn __repr__(&self, py: Python) -> String {");
         write_str!(self, "        format!(");
 
@@ -496,6 +499,47 @@ impl StructBindGenerator {
         write_str!(self, "    }");
     }
 
+    fn generate_long_match_args(&mut self) {
+        write_str!(self, "    #[classattr]");
+        write_str!(
+            self,
+            "    fn __match_args__(py: Python) -> Bound<pyo3::types::PyTuple> {"
+        );
+        write_str!(self, "        pyo3::types::PyTuple::new_bound(py, [");
+
+        for variable_info in &self.types {
+            write_fmt!(self, "            \"{}\",", variable_info.name);
+        }
+
+        write_str!(self, "        ])");
+        write_str!(self, "    }");
+    }
+
+    fn generate_match_args(&mut self) {
+        if self.types.is_empty() {
+            return;
+        }
+
+        if self.types.len() > 12 {
+            self.generate_long_match_args();
+            return;
+        }
+
+        let sig_parts: Vec<_> = repeat("&'static str").take(self.types.len()).collect();
+        let sig = sig_parts.join(", ");
+
+        write_str!(self, "    #[classattr]");
+        write_fmt!(self, "    fn __match_args__() -> ({sig},) {{",);
+        write_str!(self, "        (");
+
+        for variable_info in &self.types {
+            write_fmt!(self, "            \"{}\",", variable_info.name);
+        }
+
+        write_str!(self, "        )");
+        write_str!(self, "    }");
+    }
+
     fn generate_pack_method(&mut self) {
         write_str!(
             self,
@@ -939,6 +983,9 @@ impl Generator for StructBindGenerator {
         self.generate_repr_method();
         write_str!(self, "");
 
+        self.generate_match_args();
+        write_str!(self, "");
+
         self.generate_pack_method();
         write_str!(self, "");
 

diff --git a/pytest.py b/pytest.py
@@ -49,8 +49,19 @@ def random_script_config():
     dgs = DesiredGameState(
         game_info_state=DesiredGameInfoState(game_speed=2, end_match=Bool())
     )
-    dgs.game_info_state.world_gravity_z = Float(-650)
-    dgs.game_info_state.end_match.val = True
+
+    match dgs.game_info_state:
+        case DesiredGameInfoState():
+            dgs.game_info_state.world_gravity_z = Float(-650)
+        case _:
+            assert False
+
+    match dgs.game_info_state.end_match:
+        case Bool(val):
+            dgs.game_info_state.end_match.val = not val
+        case _:
+            assert False
+
     dgs.console_commands = [ConsoleCommand("dump_items")]
     dgs.ball_state = DesiredBallState()
 
@@ -88,7 +99,14 @@ def random_script_config():
     print(comm.content.decode("utf-8"))
     print()
 
-    print(hash(AirState.Dodging))
+    air_state = AirState.Dodging
+    print(hash(air_state))
+
+    match air_state:
+        case AirState.Dodging:
+            pass
+        case _:
+            assert False
 
     try:
         AirState(8)
@@ -117,12 +135,14 @@ def random_script_config():
     renderPolyLine = RenderMessage(
         PolyLine3D(
             [Vector3() for _ in range(2048)],
-        )
+            Color(255),
+        ),
     )
 
     match renderPolyLine.variety.item:
-        case PolyLine3D():
-            assert len(renderPolyLine.variety.item.points) == 2048
+        case PolyLine3D(points, clr):
+            assert len(points) == 2048
+            assert clr.a == 255
         case _:
             assert False