Fold pyclass enums into planus enums

Author: VirxEC

codegen/class_inject.rs

@@ -22,3 +22,16 @@ pub fn classes_to_lib_rs(mut class_names: Vec<&str>) -> io::Result<()> {
 
     Ok(())
 }
+
+pub fn add_pyclass_to_planus_enum(file: &mut String, name: &str) {
+    let find_str = format!("pub enum {name} {{");
+    let file_pos = file.find(&find_str).unwrap();
+    let end_pos = file_pos + find_str.len();
+    let start_pos = file[..file_pos].rfind("#[derive").unwrap();
+    file.replace_range(
+        start_pos..end_pos,
+        &format!(
+            "#[derive(Copy, Clone, Debug, Default, PartialEq, Eq, PartialOrd, Ord, Hash)]\n#[repr(u8)]\n#[::pyo3::pyclass(module = \"rlbot_flatbuffers\", from_py_object, frozen, hash, eq, eq_int)]\n{find_str}#[default]"
+        )
+    );
+}

codegen/enums.rs

@@ -95,67 +95,6 @@ impl<'a> EnumBindGenerator<'a> {
         }
     }
 
-    fn generate_definition(&mut self) {
-        write_str!(
-            self,
-            "#[derive(Debug, Default, Clone, PartialEq, Eq, Hash)]"
-        );
-        write_str!(
-            self,
-            "#[pyclass(module = \"rlbot_flatbuffers\", from_py_object, frozen, hash, eq, eq_int)]"
-        );
-        write_fmt!(self, "pub enum {} {{", self.name);
-        write_str!(self, "    #[default]");
-
-        for (var_num, var_info) in self.variants {
-            write_fmt!(self, "    {} = {var_num},", normalize_caps(&var_info.name));
-        }
-
-        write_str!(self, "}");
-        write_str!(self, "");
-    }
-
-    fn generate_from_flat_impls(&mut self) {
-        write_fmt!(self, "impl From<flat::{}> for {} {{", self.name, self.name);
-        write_fmt!(self, "    fn from(flat_t: flat::{}) -> Self {{", self.name);
-        write_str!(self, "        match flat_t {");
-
-        for var_info in self.variants.values() {
-            let var_name = normalize_caps(&var_info.name);
-            write_fmt!(
-                self,
-                "            flat::{}::{var_name} => Self::{var_name},",
-                self.name
-            );
-        }
-
-        write_str!(self, "        }");
-        write_str!(self, "    }");
-        write_str!(self, "}");
-        write_str!(self, "");
-    }
-
-    fn generate_to_flat_impls(&mut self) {
-        write_fmt!(self, "impl From<&{}> for flat::{} {{", self.name, self.name);
-        write_fmt!(self, "    fn from(py_type: &{}) -> Self {{", self.name);
-        write_str!(self, "        match py_type {");
-
-        for var_info in self.variants.values() {
-            let var_name = normalize_caps(&var_info.name);
-            write_fmt!(
-                self,
-                "            {}::{var_name} => Self::{var_name},",
-                self.name,
-            );
-        }
-
-        write_str!(self, "        }");
-        write_str!(self, "    }");
-        write_str!(self, "}");
-
-        write_str!(self, "");
-    }
-
     fn generate_new_method(&mut self) {
         write_str!(self, "    #[new]");
         assert!(u8::try_from(self.variants.len()).is_ok());
@@ -216,10 +155,9 @@ impl<'a> EnumBindGenerator<'a> {
             "use pyo3::{PyResult, exceptions::PyValueError, pyclass, pymethods};"
         );
         write_str!(self, "");
+        write_fmt!(self, "pub use flat::{};", self.name);
+        write_str!(self, "");
 
-        self.generate_definition();
-        self.generate_from_flat_impls();
-        self.generate_to_flat_impls();
         self.generate_py_methods();
 
         self.file_contents

codegen/main.rs

@@ -11,8 +11,8 @@ use std::{
 };
 
 use crate::{
-    enums::EnumBindGenerator, structs::StructBindGenerator, table::TableBindGenerator,
-    unions::UnionBindGenerator,
+    class_inject::add_pyclass_to_planus_enum, enums::EnumBindGenerator,
+    structs::StructBindGenerator, table::TableBindGenerator, unions::UnionBindGenerator,
 };
 
 mod class_inject;
@@ -274,6 +274,13 @@ fn main() -> eyre::Result<()> {
         .replace("::serde::Serialize,", "")
         .replace("::serde::Deserialize,", "");
 
+    for (_, full_type_name, item) in declarations.iter_declarations() {
+        if let DeclarationKind::Enum(_) = &item.kind {
+            let type_name = full_type_name.0.last().unwrap();
+            add_pyclass_to_planus_enum(&mut generated_planus, type_name);
+        }
+    }
+
     fs::write(OUT_FILE, format_string(&generated_planus)?.as_bytes())?;
 
     class_inject::classes_to_lib_rs(class_names)?;

codegen/table.rs

@@ -211,7 +211,7 @@ impl<'a> TableBindGenerator<'a> {
                         }
                     },
                     SimpleType::Enum(_) => {
-                        format!("flat_t.{field_name}.into()")
+                        format!("flat_t.{field_name}")
                     }
                 },
                 TypeKind::String => match field_info.assign_mode {
@@ -323,7 +323,7 @@ impl<'a> TableBindGenerator<'a> {
                         }
                     },
                     SimpleType::Enum(_) => {
-                        format!("(&py_type.{field_name}).into()")
+                        format!("py_type.{field_name}")
                     }
                 },
                 TypeKind::String => match field_info.assign_mode {