Flatten unions in struct constructors

Author: VirxEC

Cargo.lock

@@ -1,6 +1,6 @@
 [package]
 name = "rlbot-flatbuffers-py"
-version = "0.3.9"
+version = "0.4.0"
 edition = "2021"
 description = "A Python module implemented in Rust for serializing and deserializing RLBot's flatbuffers"
 repository = "https://github.com/VirxEC/rlbot_flatbuffers_py"

Cargo.toml

@@ -15,7 +15,10 @@ fn camel_to_snake_case(variable_name: &str) -> String {
     for c in variable_name.chars() {
         if c.is_uppercase() {
             if last_was_uppercase {
-                snake_case_parts.last_mut().unwrap().push(c.to_lowercase().next().unwrap());
+                snake_case_parts
+                    .last_mut()
+                    .unwrap()
+                    .push(c.to_lowercase().next().unwrap());
             } else {
                 snake_case_parts.push(c.to_lowercase().to_string());
             }
@@ -170,7 +173,11 @@ impl EnumBindGenerator {
 
         for variable_info in &self.types {
             let variable_name = variable_info.name.as_str();
-            write_fmt!(self, "            {} => Ok(Self::{variable_name}),", variable_info.raw_type);
+            write_fmt!(
+                self,
+                "            {} => Ok(Self::{variable_name}),",
+                variable_info.raw_type
+            );
         }
 
         if self.types.len() != usize::from(u8::MAX) {
@@ -192,7 +199,11 @@ impl EnumBindGenerator {
 
     fn generate_repr_method(&mut self) {
         write_str!(self, "    pub fn __repr__(&self) -> String {");
-        write_fmt!(self, "        format!(\"{}(value={{}})\", *self as u8)", self.struct_name);
+        write_fmt!(
+            self,
+            "        format!(\"{}(value={{}})\", *self as u8)",
+            self.struct_name
+        );
         write_str!(self, "    }");
     }
 
@@ -251,7 +262,12 @@ impl Generator for EnumBindGenerator {
     }
 
     fn generate_from_flat_impls(&mut self) {
-        write_fmt!(self, "impl From<flat::{}> for {} {{", self.struct_name, self.struct_name);
+        write_fmt!(
+            self,
+            "impl From<flat::{}> for {} {{",
+            self.struct_name,
+            self.struct_name
+        );
         write_fmt!(self, "    fn from(flat_t: flat::{}) -> Self {{", self.struct_name);
         write_str!(self, "        match flat_t {");
 
@@ -274,7 +290,12 @@ impl Generator for EnumBindGenerator {
     }
 
     fn generate_to_flat_impls(&mut self) {
-        write_fmt!(self, "impl From<&{}> for flat::{} {{", self.struct_name, self.struct_name);
+        write_fmt!(
+            self,
+            "impl From<&{}> for flat::{} {{",
+            self.struct_name,
+            self.struct_name
+        );
         write_fmt!(self, "    fn from(py_type: &{}) -> Self {{", self.struct_name);
         write_str!(self, "        match *py_type {");
 

codegen/enums.rs

@@ -58,6 +58,13 @@ impl PythonBindType {
         "GameMessageWrapper",
         "GameMessage",
     ];
+    pub const UNIONS: [&'static str; 5] = [
+        "PlayerClass",
+        "GameMessage",
+        "CollisionShape",
+        "RelativeAnchor",
+        "RenderType",
+    ];
 
     fn new(path: &Path) -> Option<Self> {
         // get the filename without the extension

codegen/main.rs

@@ -1,4 +1,5 @@
 use crate::{
+    generator::Generator,
     structs::{InnerOptionType, InnerVecType, RustType},
     PythonBindType,
 };
@@ -76,10 +77,16 @@ pub fn generator(type_data: &[PythonBindType]) -> io::Result<()> {
                 write_str!(file, "");
                 write_str!(file, "    def __init__(self, value: int = 0):");
                 write_str!(file, "        \"\"\"");
-                write_str!(file, "        :raises ValueError: If the `value` is not a valid enum value");
+                write_str!(
+                    file,
+                    "        :raises ValueError: If the `value` is not a valid enum value"
+                );
                 write_str!(file, "        \"\"\"\n");
                 write_str!(file, "    def __int__(self) -> int: ...");
-                write_fmt!(file, "    def __richcmp__(self, other: {type_name}, op: int) -> bool: ...");
+                write_fmt!(
+                    file,
+                    "    def __richcmp__(self, other: {type_name}, op: int) -> bool: ..."
+                );
             }
             PythonBindType::Struct(gen) => {
                 let mut python_types = Vec::new();
@@ -134,7 +141,7 @@ pub fn generator(type_data: &[PythonBindType]) -> io::Result<()> {
                             } else {
                                 type_name.as_str()
                             };
-                            
+
                             python_types.push(format!("Optional[{python_type}]"));
                         }
                         RustType::Box(inner_type) => {
@@ -145,6 +152,22 @@ pub fn generator(type_data: &[PythonBindType]) -> io::Result<()> {
                             python_types.push("str".to_string());
                             write_fmt!(file, "    {variable_name}: str");
                         }
+                        RustType::Union(type_name) => {
+                            write_fmt!(file, "    {variable_name}: {type_name}");
+
+                            // search for the union with the name `type_name` and get the types
+                            let union_types = type_data
+                                .iter()
+                                .find_map(|item| match item {
+                                    PythonBindType::Union(gen) if gen.struct_name() == type_name => {
+                                        Some(gen.types.iter().skip(1).map(|v| v.name.as_str()).collect::<Vec<_>>())
+                                    }
+                                    _ => None,
+                                })
+                                .unwrap();
+                            let types = union_types.join(" | ");
+                            python_types.push(format!("Optional[{types}]"));
+                        }
                         RustType::Custom(type_name) | RustType::Other(type_name) | RustType::Base(type_name) => {
                             python_types.push(type_name.to_string());
                             write_fmt!(file, "    {variable_name}: {type_name}");
@@ -168,14 +191,14 @@ pub fn generator(type_data: &[PythonBindType]) -> io::Result<()> {
                             "String" => Cow::Borrowed("\"\""),
                             "Vec<u8>" => Cow::Borrowed("b\"\""),
                             t => {
-                                if t.starts_with("Vec<") {
+                                if python_type.starts_with("Optional") || t.starts_with("Option<") {
+                                    Cow::Borrowed("None")
+                                } else if t.starts_with("Vec<") {
                                     Cow::Borrowed("[]")
                                 } else if t.starts_with("Box<") {
                                     let inner_type =
                                         t.trim_start_matches("Box<").trim_end_matches('>').trim_end_matches('T');
                                     Cow::Owned(format!("{inner_type}()"))
-                                } else if t.starts_with("Option<") {
-                                    Cow::Borrowed("None")
                                 } else {
                                     Cow::Owned(format!("{}()", t.trim_end_matches('T')))
                                 }