From 0f53e833ed655a68712553ea4766d0b24509f15b Mon Sep 17 00:00:00 2001
From: VirxEC <virx@virxcase.dev>
Date: Tue, 30 Jan 2024 22:30:20 -0500
Subject: [PATCH] Generate types in init methods

Update spec
---
 Cargo.lock         |  2 +-
 Cargo.toml         |  2 +-
 build.rs           | 35 ++++++++++++++++++++++++++---------
 flatbuffers-schema |  2 +-
 src/lib.rs         |  1 -
 5 files changed, 29 insertions(+), 13 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 8e626a6..0b0f3f9 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -314,7 +314,7 @@ dependencies = [
 
 [[package]]
 name = "rlbot-flatbuffers-py"
-version = "0.1.3"
+version = "0.1.4"
 dependencies = [
  "flatbuffers",
  "get-size",
diff --git a/Cargo.toml b/Cargo.toml
index 4d66ba8..a415141 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "rlbot-flatbuffers-py"
-version = "0.1.3"
+version = "0.1.4"
 edition = "2021"
 description = "A Python module implemented in Rust for serializing and deserializing RLBot's flatbuffers"
 repository = "https://github.com/VirxEC/rlbot-flatbuffers-py"
diff --git a/build.rs b/build.rs
index a4e748a..2690839 100644
--- a/build.rs
+++ b/build.rs
@@ -1086,8 +1086,10 @@ fn pyi_generator(type_data: &[(String, String, Vec<Vec<String>>)]) -> io::Result
             file_contents.push(Cow::Owned(format!("    item_type: {}Type", type_name)));
         }
 
+        let mut python_types = Vec::new();
+
         'outer: for variable_info in types {
-            let mut variable_name = &variable_info[0];
+            let mut variable_name = variable_info[0].as_str();
 
             if variable_name == "NONE" {
                 continue;
@@ -1108,6 +1110,7 @@ fn pyi_generator(type_data: &[(String, String, Vec<Vec<String>>)]) -> io::Result
 
             for (rust_type, python_type) in primitive_map {
                 if variable_type == rust_type {
+                    python_types.push(python_type.to_string());
                     file_contents.push(Cow::Owned(format!("    {variable_name}: {python_type}")));
                     continue 'outer;
                 }
@@ -1118,6 +1121,8 @@ fn pyi_generator(type_data: &[(String, String, Vec<Vec<String>>)]) -> io::Result
                     .trim_start_matches("Vec<")
                     .trim_end_matches('>')
                     .trim_end_matches('T');
+
+                python_types.push(format!("list[{type_name}]"));
                 file_contents.push(Cow::Owned(format!("    {variable_name}: list[{type_name}]")));
             } else if variable_type.starts_with("Option<") {
                 let type_name = variable_type
@@ -1126,24 +1131,32 @@ fn pyi_generator(type_data: &[(String, String, Vec<Vec<String>>)]) -> io::Result
                     .trim_end_matches('>')
                     .trim_end_matches('T');
 
-                if type_name == "bool" {
-                    file_contents.push(Cow::Owned(format!("    {variable_name}: Optional[bool]")));
+                let python_type = if type_name == "bool" {
+                    "bool"
                 } else if type_name == "i32" || type_name == "u32" {
-                    file_contents.push(Cow::Owned(format!("    {variable_name}: Optional[int]")));
+                    "int"
                 } else if type_name == "f32" {
-                    file_contents.push(Cow::Owned(format!("    {variable_name}: Optional[float]")));
+                    "float"
                 } else if type_name == "String" {
-                    file_contents.push(Cow::Owned(format!("    {variable_name}: Optional[str]")));
+                    "str"
                 } else {
-                    file_contents.push(Cow::Owned(format!("    {variable_name}: Optional[{type_name}]")));
-                }
+                    type_name
+                };
+
+                python_types.push(format!("Optional[{python_type}]"));
+                file_contents.push(Cow::Owned(format!("    {variable_name}: Optional[{python_type}]")));
             } else if variable_type.starts_with("Box<") && variable_type.ends_with("T>") {
                 let type_name = variable_type.trim_start_matches("Box<").trim_end_matches("T>");
+
+                python_types.push(type_name.to_string());
                 file_contents.push(Cow::Owned(format!("    {variable_name}: {type_name}")));
             } else if variable_type.ends_with('T') {
                 let type_name = variable_type.trim_end_matches('T');
+
+                python_types.push(type_name.to_string());
                 file_contents.push(Cow::Owned(format!("    {variable_name}: {type_name}")));
             } else {
+                python_types.push(variable_type.clone());
                 file_contents.push(Cow::Owned(format!("    {variable_name}: {variable_type}")));
             }
         }
@@ -1156,6 +1169,7 @@ fn pyi_generator(type_data: &[(String, String, Vec<Vec<String>>)]) -> io::Result
             file_contents.push(Cow::Borrowed("    def __init__("));
             file_contents.push(Cow::Borrowed("        self,"));
 
+            let mut i = 0;
             for variable_info in types {
                 if &variable_info[0] == "NONE" {
                     continue;
@@ -1187,7 +1201,10 @@ fn pyi_generator(type_data: &[(String, String, Vec<Vec<String>>)]) -> io::Result
                     }
                 };
 
-                file_contents.push(Cow::Owned(format!("        {variable_name}={default_value},")));
+                let python_type = &python_types[i];
+                file_contents.push(Cow::Owned(format!("        {variable_name}: {python_type}={default_value},")));
+
+                i += 1;
             }
 
             file_contents.push(Cow::Borrowed("    ): ..."));
diff --git a/flatbuffers-schema b/flatbuffers-schema
index 1468624..600d3a0 160000
--- a/flatbuffers-schema
+++ b/flatbuffers-schema
@@ -1 +1 @@
-Subproject commit 14686241bea255cd7fd332cc26519e89ca531bdf
+Subproject commit 600d3a04f97772214079755798a4c71d0f35b627
diff --git a/src/lib.rs b/src/lib.rs
index 4a31c2d..7d2f439 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -48,7 +48,6 @@ pynamedmodule! {
         GameStateType,
         GameSpeedOption,
         ScoreInfo,
-        BotSpawnIndex,
         MessagePacket,
         DemolishOption,
         Vector3Partial,