Special cases for Float and Bool

Now accept normal Python floats or bools along with Float and Bool objects

Author: VirxEC

build.rs

@@ -26,6 +26,7 @@ struct PythonBindGenerator {
 
 impl PythonBindGenerator {
     const BASE_TYPES: [&'static str; 6] = ["bool", "i32", "u32", "f32", "String", "u8"];
+    const SPECIAL_BASE_TYPES: [&'static str; 2] = ["FloatT", "BoolT"];
 
     fn new(path: &Path) -> Option<Self> {
         // get the filename without the extension
@@ -952,6 +953,7 @@ impl PythonBindGenerator {
         }
 
         let mut signature_parts = Vec::new();
+        let mut needs_python = false;
 
         for variable_info in &self.types {
             let variable_name = &variable_info[0];
@@ -962,6 +964,15 @@ impl PythonBindGenerator {
             }
 
             if variable_type.starts_with("Option<") {
+                let inner_type = variable_type
+                    .trim_start_matches("Option<")
+                    .trim_start_matches("Box<")
+                    .trim_end_matches('>');
+
+                if Self::SPECIAL_BASE_TYPES.contains(&inner_type) {
+                    needs_python = true;
+                }
+
                 signature_parts.push(format!("{variable_name}=None"));
             } else if !Self::BASE_TYPES.contains(&variable_type.as_str())
                 && (variable_type.starts_with("Box<") || variable_type.ends_with('T'))
@@ -973,9 +984,12 @@ impl PythonBindGenerator {
         }
 
         self.write_string(format!("    #[pyo3(signature = ({}))]", signature_parts.join(", ")));
-
         self.write_str("    pub fn new(");
 
+        if needs_python {
+            self.write_str("        py: Python,");
+        }
+
         for variable_info in &self.types {
             let variable_name = &variable_info[0];
             let mut variable_type = variable_info[1].to_string();
@@ -1002,6 +1016,10 @@ impl PythonBindGenerator {
 
                 if Self::BASE_TYPES.contains(&inner_type) {
                     variable_type = format!("Option<{inner_type}>");
+                } else if inner_type == "Float" {
+                    variable_type = String::from("Option<crate::Floats>");
+                } else if inner_type == "Bool" {
+                    variable_type = String::from("Option<crate::Bools>");
                 } else {
                     variable_type = format!("Option<Py<super::{inner_type}>>");
                 }
@@ -1022,8 +1040,18 @@ impl PythonBindGenerator {
 
         for variable_info in &self.types {
             let variable_name = &variable_info[0];
+            let variable_type = variable_info[1]
+                .trim_start_matches("Option<")
+                .trim_start_matches("Box<")
+                .trim_end_matches('>');
 
-            self.file_contents.push(Cow::Owned(format!("            {variable_name},")));
+            if Self::SPECIAL_BASE_TYPES.contains(&variable_type) {
+                self.file_contents.push(Cow::Owned(format!(
+                    "            {variable_name}: {variable_name}.map(|x| x.into_gil(py)),"
+                )));
+            } else {
+                self.file_contents.push(Cow::Owned(format!("            {variable_name},")));
+            }
         }
 
         self.write_str("        }");
@@ -1407,6 +1435,10 @@ fn pyi_generator(type_data: &[(String, String, Vec<Vec<String>>)]) -> io::Result
                     "float"
                 } else if type_name == "String" {
                     "str"
+                } else if type_name == "Float" {
+                    "Float | float"
+                } else if type_name == "Bool" {
+                    "Bool | bool"
                 } else {
                     type_name
                 };

pytest.py

@@ -3,17 +3,6 @@
 from rlbot_flatbuffers import *
 
 if __name__ == "__main__":
-    color = Color(255, 0, 0)
-    print(repr(color))
-    print(color)
-    eval(repr(color))
-    print()
-
-    controller = ControllerState(throttle=1)
-    controller.boost = True
-
-    player_input = PlayerInput(0, controller)
-
     ready_message = ReadyMessage(True, wants_game_messages=True)
     print(repr(ready_message))
     print(ready_message)
@@ -45,18 +34,29 @@
 
     num_trials = 1_000_000
 
+    total_make_time = 0
     total_pack_time = 0
     total_unpack_time = 0
     for _ in range(num_trials):
         start = time_ns()
-        packed_bytes = player_input.pack()
+        desired_game_state = DesiredGameState(
+            DesiredBallState(DesiredPhysics()),
+            car_states=[DesiredCarState(boost_amount=Float(100))],
+            game_info_state=DesiredGameInfoState(game_speed=1, world_gravity_z=Float(-650), end_match=True),
+            console_commands=[ConsoleCommand("freeze")],
+        )
+        total_make_time += time_ns() - start
+
+        start = time_ns()
+        packed_bytes = desired_game_state.pack()
         total_pack_time += time_ns() - start
 
         start = time_ns()
-        PlayerInput.unpack(packed_bytes)
+        DesiredGameState.unpack(packed_bytes)
         total_unpack_time += time_ns() - start
 
-    print(f"Average time to pack: {total_pack_time / num_trials} ns")
-    print(f"Average time to unpack: {total_unpack_time / num_trials} ns")
+    print(f"Average time to make: {round(total_make_time / num_trials, 2)}ns")
+    print(f"Average time to pack: {round(total_pack_time / num_trials, 2)}ns")
+    print(f"Average time to unpack: {round(total_unpack_time / num_trials, 2)}ns")
 
-    print(f"Total time: {(total_pack_time + total_unpack_time) / 1000000 }ms")
+    print(f"Total time: {round((total_pack_time + total_unpack_time) / 1000000, 2)}ms")

src/lib.rs

@@ -61,6 +61,48 @@ pub const fn bool_to_str(b: bool) -> &'static str {
     }
 }
 
+#[derive(Debug, Clone, FromPyObject)]
+pub enum Floats {
+    Flat(Py<Float>),
+    Num(f32),
+}
+
+impl Default for Floats {
+    fn default() -> Self {
+        Floats::Flat(get_py_default())
+    }
+}
+
+impl FromGil<Floats> for Py<Float> {
+    fn from_gil(py: Python, floats: Floats) -> Self {
+        match floats {
+            Floats::Flat(float) => float,
+            Floats::Num(num) => Py::new(py, Float::new(num)).unwrap(),
+        }
+    }
+}
+
+#[derive(Debug, Clone, FromPyObject)]
+pub enum Bools {
+    Flat(Py<Bool>),
+    Num(bool),
+}
+
+impl Default for Bools {
+    fn default() -> Self {
+        Self::Flat(get_py_default())
+    }
+}
+
+impl FromGil<Bools> for Py<Bool> {
+    fn from_gil(py: Python, bools: Bools) -> Self {
+        match bools {
+            Bools::Flat(float) => float,
+            Bools::Num(num) => Py::new(py, Bool::new(num)).unwrap(),
+        }
+    }
+}
+
 macro_rules! pynamedmodule {
     (doc: $doc:literal, name: $name:tt, classes: [$($class_name:ident),*], vars: [$(($var_name:literal, $value:expr)),*]) => {
         #[doc = $doc]