@@ -55,6 +55,7 @@ pub struct StructBindGenerator {
5555 pub is_frozen : bool ,
5656 is_no_set : bool ,
5757 pub default_override : Option < ( & ' static str , & ' static str ) > ,
58+ freelist_size : usize ,
5859}
5960
6061macro_rules! write_str {
@@ -115,6 +116,17 @@ impl StructBindGenerator {
115116 }
116117 } ) ;
117118
119+ let freelist_size = PythonBindType :: FREELIST_TYPES
120+ . iter ( )
121+ . find_map ( |& ( name, size) | {
122+ if name == struct_name. as_str ( ) {
123+ Some ( size)
124+ } else {
125+ None
126+ }
127+ } )
128+ . unwrap_or_default ( ) ;
129+
118130 Some ( Self {
119131 filename,
120132 struct_name,
@@ -126,6 +138,7 @@ impl StructBindGenerator {
126138 is_frozen,
127139 is_no_set,
128140 default_override,
141+ freelist_size,
129142 } )
130143 }
131144
@@ -940,16 +953,20 @@ impl Generator for StructBindGenerator {
940953 }
941954
942955 fn generate_definition ( & mut self ) {
943- write_str ! (
944- self ,
945- if self . is_frozen {
946- "#[pyclass(module = \" rlbot_flatbuffers\" , subclass, get_all, frozen)]"
947- } else if self . types. is_empty( ) {
948- "#[pyclass(module = \" rlbot_flatbuffers\" , subclass, frozen)]"
949- } else {
950- "#[pyclass(module = \" rlbot_flatbuffers\" , subclass, get_all)]"
951- }
952- ) ;
956+ let freelist_str = if self . freelist_size > 0 {
957+ Cow :: Owned ( format ! ( ", freelist = {}" , self . freelist_size) )
958+ } else {
959+ Cow :: Borrowed ( "" )
960+ } ;
961+
962+ let pyclass_start_str = "#[pyclass(module = \" rlbot_flatbuffers\" , subclass, " ;
963+ if self . is_frozen {
964+ write_fmt ! ( self , "{pyclass_start_str}get_all, frozen{freelist_str})]" ) ;
965+ } else if self . types . is_empty ( ) {
966+ write_fmt ! ( self , "{pyclass_start_str}frozen{freelist_str})]" ) ;
967+ } else {
968+ write_fmt ! ( self , "{pyclass_start_str}get_all{freelist_str})]" ) ;
969+ }
953970
954971 if self . types . is_empty ( ) {
955972 write_str ! ( self , "#[derive(Default)]" ) ;
0 commit comments