Start using classes in gen script to avoid globals

scripts/generate_messages.py

@@ -142,180 +142,6 @@ def name_to_snake(name_list: list[str]) -> str:
     name = '_'.join(name_list)
     return fix_name_conflict(name)
 
-def generate_msg(msg: NetMessageJson, game: Literal['game', 'system']) -> None:
-    name_snake = name_to_snake(msg['name'])
-    name_camel = name_to_camel(msg['name'])
-    dirname = os.path.dirname(__file__)
-    file_path= os.path.join(
-            dirname,
-            f'../twnet_parser/messages7/{game}/',
-            f'{name_snake}.py')
-    # if os.path.exists(file_path):
-    #     print(f"Warning: file already exists! {file_path}")
-    #     return
-    with open(file_path, 'w') as out_file:
-        print(f"Generating {file_path} ...")
-        out_file.write('# generated by scripts/generate_messages.py\n')
-        out_file.write('\n')
-        out_file.write('from twnet_parser.pretty_print import PrettyPrint\n')
-        if len(msg['members']) > 0:
-            out_file.write('from twnet_parser.packer import Unpacker\n')
-        out_file.write('from twnet_parser.chunk_header import ChunkHeader\n')
-        out_file.write(get_dependencies(msg))
-        out_file.write('\n')
-        out_file.write(f'class Msg{name_camel}(PrettyPrint):\n')
-        out_file.write('    def __init__(\n')
-        out_file.write('            self,\n')
-        args: list[str] = []
-        for member in msg['members']:
-            # {'name': ['message'], 'type': {'kind': 'string', 'disallow_cc': False}} 
-            ftype = 'int'
-            default = '-1'
-            if member['type']['kind'] == 'string':
-                ftype = 'str'
-                default = "'default'"
-            elif member['type']['kind'] == 'rest':
-                ftype = 'bytes'
-                default = "b'\\x00'"
-            elif member['type']['kind'] == 'sha256':
-                ftype = 'Annotated[bytes, 32]'
-                default = "bytes(32)"
-            elif member['type']['kind'] == 'data':
-                ftype = 'bytes'
-                default = "b'\\x00'"
-                if member['type']['size'] == 'specified_before':
-                    args.append('            data_size: Optional[int] = None')
-                else:
-                    raise ValueError(f"Error: unknown data size {member['type']}")
-            # {"name": ["mode"], "type": {"kind": "enum", "enum": ["chat"]}},
-            elif member['type']['kind'] == 'enum':
-                ftype = 'int'
-                default = '0'
-                # TODO: use ENUM_NAME_SOME_VALUE as default here
-            elif member['type']['kind'] in ('int32', 'tick'):
-                ftype = 'int'
-                default = '0'
-            elif member['type']['kind'] == 'boolean':
-                ftype = 'bool'
-                default = 'False'
-            elif member['type']['kind'] == 'tune_param':
-                ftype = 'float'
-                default = '0.0'
-            elif member['type']['kind'] == 'snapshot_object':
-                # TODO: think about snapshot_object
-                ftype = 'int'
-                default = '0'
-            elif member['type']['kind'] == 'array':
-                size: int = member['type']['count']
-                if size is None:
-                    print("Error: size is none for the following message")
-                    print(msg)
-                    exit(1)
-                arr_member: ArrayMemberTypeJson = member['type']['member_type']
-                if arr_member['kind'] == 'string':
-                    ftype = f'Annotated[list[str], {size}]'
-                    default = '[' + ', '.join(["''"] * size) + ']'
-                elif arr_member['kind'] == 'boolean':
-                    ftype = f'Annotated[list[bool], {size}]'
-                    default = '[' + ', '.join(["False"] * size) + ']'
-                elif arr_member['kind'] in ('int32', 'tick', 'enum'):
-                    ftype = f'Annotated[list[int], {size}]'
-                    default = '[' + ', '.join(["0"] * size) + ']'
-                else:
-                    raise ValueError( \
-                        f"Error: unknown array member type {member['type']}")
-                # Initializing lists with defaults
-                # And type annotation can get quite long
-                # So split it in two lines
-                default = f'\\\n                {default}'
-            elif member['type']['kind'] == 'flags': # TODO: think about flags
-                ftype = 'int'
-                default = '0'
-            elif member['type']['kind'] == 'optional':
-                if member['type']['inner']['kind'] == 'string':
-                    ftype = 'str'
-                    default = "''"
-                elif member['type']['inner']['kind'] in ('int32', 'tick'):
-                    ftype = 'int'
-                    default = '0'
-                else:
-                    raise ValueError(f"Error: unknown optional type {member['type']}")
-            else:
-                raise ValueError(f"Error: unknown type {member['type']}")
-            name = name_to_snake(member["name"])
-            manual_default = get_default(f"{game}.{name_snake}.{name}")
-            if manual_default:
-                default = manual_default
-            args.append(f'            {name}: {ftype} = {default}')
-        out_file.write(',\n'.join(args) + '\n')
-        out_file.write('    ) -> None:\n')
-        out_file.write(f"        self.message_name = '{name_snake}'\n")
-        sys: str = 'True' if game == 'system' else 'False'
-        out_file.write(f"        self.system_message = {sys}\n")
-        out_file.write("        self.header: ChunkHeader\n")
-        out_file.write('\n')
-        for member in msg['members']:
-            # {'name': ['message'], 'type': {'kind': 'string', 'disallow_cc': False}} 
-            ftype = 'int'
-            if member['type']['kind'] == 'string':
-                ftype = 'str'
-            elif member['type']['kind'] == 'rest':
-                ftype = 'bytes'
-            elif member['type']['kind'] == 'sha256':
-                ftype = 'Annotated[bytes, 32]'
-            elif member['type']['kind'] == 'data':
-                ftype = 'bytes'
-                if member['type']['size'] == 'specified_before':
-                    out_file.write("        " \
-                        "self.data_size: int = data_size if data_size else len(data)\n")
-                else:
-                    raise ValueError(f"Error: unknown data size {member['type']}")
-            # {"name": ["mode"], "type": {"kind": "enum", "enum": ["chat"]}},
-            elif member['type']['kind'] == 'enum':
-                ftype = 'int'
-            elif member['type']['kind'] in ('int32', 'tick'):
-                ftype = 'int'
-            elif member['type']['kind'] == 'boolean':
-                ftype = 'bool'
-            elif member['type']['kind'] == 'tune_param':
-                ftype = 'float'
-            elif member['type']['kind'] == 'snapshot_object':
-                # TODO: think about snapshot_object
-                ftype = 'int'
-            elif member['type']['kind'] == 'array':
-                # Array type annotations are so annoyingly long
-                # also there is a planned refactor
-                # https://gitlab.com/teeworlds-network/twnet_parser/-/issues/4
-                # so inherit type from constructor arguments
-                ftype = ''
-            elif member['type']['kind'] == 'flags': # TODO: think about flags
-                ftype = 'int'
-            elif member['type']['kind'] == 'optional':
-                if member['type']['inner']['kind'] == 'string':
-                    ftype = 'str'
-                elif member['type']['inner']['kind'] in ('int32', 'tick'):
-                    ftype = 'int'
-                else:
-                    raise ValueError(f"Error: unknown optional type {member['type']}")
-            else:
-                raise ValueError(f"Error: unknown type {member['type']}")
-            name = name_to_snake(member["name"])
-            if ftype != '':
-                ftype = f': {ftype}'
-            out_file.write(f"        self.{name}{ftype} = {name}\n")
-        out_file.write('\n')
-        out_file.write('    # first byte of data\n')
-        out_file.write('    # has to be the first byte of the message payload\n')
-        out_file.write('    # NOT the chunk header and NOT the message id\n')
-        out_file.write('    def unpack(self, data: bytes) -> bool:\n')
-        if len(msg['members']) > 0:
-            out_file.write('        unpacker = Unpacker(data)\n')
-        out_file.write(gen_unpack_members(msg))
-        out_file.write('        return True\n')
-        out_file.write('\n')
-        out_file.write('    def pack(self) -> bytes:\n')
-        out_file.write(gen_pack_return(msg))
-
 def gen_unpack_members(msg: NetMessageJson) -> str:
     res: str = ''
     for member in msg['members']:
@@ -580,21 +406,217 @@ def gen_enum_file7(enums: list[GameEnumJson]):
         print(f"Generating {file_path} ...")
         out_file.write(enum_code)
 
-def generate(spec: str) -> None:
+class CodeGenerator():
-    print(f"generating classes from {spec} ...")
+    def generate_msg(
-    with open(spec) as spec_io:
+            self,
-        spec_data: SpecJson = json.load(spec_io)
+            msg: NetMessageJson,
-        # for msg in [spec_data['game_messages'][1]]:
+            game: Literal['game', 'system']
-        game_enums: list[GameEnumJson] = spec_data['game_enumerations']
+    ) -> None:
-        game_messages: list[NetMessageJson] = spec_data['game_messages']
+        name_snake = name_to_snake(msg['name'])
-        system_messages: list[NetMessageJson] = spec_data['system_messages']
+        name_camel = name_to_camel(msg['name'])
-        gen_enum_file7(game_enums)
+        dirname = os.path.dirname(__file__)
-        gen_match_file7('game', game_messages)
+        file_path= os.path.join(
-        gen_match_file7('system', system_messages)
+                dirname,
-        for msg in game_messages:
+                f'../twnet_parser/messages7/{game}/',
-            generate_msg(msg, 'game')
+                f'{name_snake}.py')
-        for msg in system_messages:
+        # if os.path.exists(file_path):
-            generate_msg(msg, 'system')
+        #     print(f"Warning: file already exists! {file_path}")
+        #     return
+        with open(file_path, 'w') as out_file:
+            print(f"Generating {file_path} ...")
+            out_file.write('# generated by scripts/generate_messages.py\n')
+            out_file.write('\n')
+            out_file.write('from twnet_parser.pretty_print import PrettyPrint\n')
+            if len(msg['members']) > 0:
+                out_file.write('from twnet_parser.packer import Unpacker\n')
+            out_file.write('from twnet_parser.chunk_header import ChunkHeader\n')
+            out_file.write(get_dependencies(msg))
+            out_file.write('\n')
+            out_file.write(f'class Msg{name_camel}(PrettyPrint):\n')
+            out_file.write('    def __init__(\n')
+            out_file.write('            self,\n')
+            args: list[str] = []
+            for member in msg['members']:
+                # {
+                #   'name': ['message'],
+                #   'type': {
+                #     'kind': 'string',
+                #     'disallow_cc': False
+                #   }
+                # } 
+                ftype = 'int'
+                default = '-1'
+                if member['type']['kind'] == 'string':
+                    ftype = 'str'
+                    default = "'default'"
+                elif member['type']['kind'] == 'rest':
+                    ftype = 'bytes'
+                    default = "b'\\x00'"
+                elif member['type']['kind'] == 'sha256':
+                    ftype = 'Annotated[bytes, 32]'
+                    default = "bytes(32)"
+                elif member['type']['kind'] == 'data':
+                    ftype = 'bytes'
+                    default = "b'\\x00'"
+                    if member['type']['size'] == 'specified_before':
+                        args.append('            data_size: Optional[int] = None')
+                    else:
+                        raise ValueError(f"Error: unknown data size {member['type']}")
+                # {"name": ["mode"], "type": {"kind": "enum", "enum": ["chat"]}},
+                elif member['type']['kind'] == 'enum':
+                    ftype = 'int'
+                    default = '0'
+                    # TODO: use ENUM_NAME_SOME_VALUE as default here
+                elif member['type']['kind'] in ('int32', 'tick'):
+                    ftype = 'int'
+                    default = '0'
+                elif member['type']['kind'] == 'boolean':
+                    ftype = 'bool'
+                    default = 'False'
+                elif member['type']['kind'] == 'tune_param':
+                    ftype = 'float'
+                    default = '0.0'
+                elif member['type']['kind'] == 'snapshot_object':
+                    # TODO: think about snapshot_object
+                    ftype = 'int'
+                    default = '0'
+                elif member['type']['kind'] == 'array':
+                    size: int = member['type']['count']
+                    if size is None:
+                        print("Error: size is none for the following message")
+                        print(msg)
+                        exit(1)
+                    arr_member: ArrayMemberTypeJson = member['type']['member_type']
+                    if arr_member['kind'] == 'string':
+                        ftype = f'Annotated[list[str], {size}]'
+                        default = '[' + ', '.join(["''"] * size) + ']'
+                    elif arr_member['kind'] == 'boolean':
+                        ftype = f'Annotated[list[bool], {size}]'
+                        default = '[' + ', '.join(["False"] * size) + ']'
+                    elif arr_member['kind'] in ('int32', 'tick', 'enum'):
+                        ftype = f'Annotated[list[int], {size}]'
+                        default = '[' + ', '.join(["0"] * size) + ']'
+                    else:
+                        raise ValueError( \
+                            f"Error: unknown array member type {member['type']}")
+                    # Initializing lists with defaults
+                    # And type annotation can get quite long
+                    # So split it in two lines
+                    default = f'\\\n                {default}'
+                elif member['type']['kind'] == 'flags': # TODO: think about flags
+                    ftype = 'int'
+                    default = '0'
+                elif member['type']['kind'] == 'optional':
+                    if member['type']['inner']['kind'] == 'string':
+                        ftype = 'str'
+                        default = "''"
+                    elif member['type']['inner']['kind'] in ('int32', 'tick'):
+                        ftype = 'int'
+                        default = '0'
+                    else:
+                        raise \
+                            ValueError( \
+                            f"Error: unknown optional type {member['type']}")
+                else:
+                    raise ValueError(f"Error: unknown type {member['type']}")
+                name = name_to_snake(member["name"])
+                manual_default = get_default(f"{game}.{name_snake}.{name}")
+                if manual_default:
+                    default = manual_default
+                args.append(f'            {name}: {ftype} = {default}')
+            out_file.write(',\n'.join(args) + '\n')
+            out_file.write('    ) -> None:\n')
+            out_file.write(f"        self.message_name = '{name_snake}'\n")
+            sys: str = 'True' if game == 'system' else 'False'
+            out_file.write(f"        self.system_message = {sys}\n")
+            out_file.write("        self.header: ChunkHeader\n")
+            out_file.write('\n')
+            for member in msg['members']:
+                # {
+                #   'name': ['message'],
+                #   'type': {
+                #     'kind': 'string',
+                #     'disallow_cc': False
+                #   }
+                # } 
+                ftype = 'int'
+                if member['type']['kind'] == 'string':
+                    ftype = 'str'
+                elif member['type']['kind'] == 'rest':
+                    ftype = 'bytes'
+                elif member['type']['kind'] == 'sha256':
+                    ftype = 'Annotated[bytes, 32]'
+                elif member['type']['kind'] == 'data':
+                    ftype = 'bytes'
+                    if member['type']['size'] == 'specified_before':
+                        out_file.write("        " \
+                            "self.data_size: int =" \
+                            " data_size if data_size else len(data)\n")
+                    else:
+                        raise ValueError(f"Error: unknown data size {member['type']}")
+                # {"name": ["mode"], "type": {"kind": "enum", "enum": ["chat"]}},
+                elif member['type']['kind'] == 'enum':
+                    ftype = 'int'
+                elif member['type']['kind'] in ('int32', 'tick'):
+                    ftype = 'int'
+                elif member['type']['kind'] == 'boolean':
+                    ftype = 'bool'
+                elif member['type']['kind'] == 'tune_param':
+                    ftype = 'float'
+                elif member['type']['kind'] == 'snapshot_object':
+                    # TODO: think about snapshot_object
+                    ftype = 'int'
+                elif member['type']['kind'] == 'array':
+                    # Array type annotations are so annoyingly long
+                    # also there is a planned refactor
+                    # https://gitlab.com/teeworlds-network/twnet_parser/-/issues/4
+                    # so inherit type from constructor arguments
+                    ftype = ''
+                elif member['type']['kind'] == 'flags': # TODO: think about flags
+                    ftype = 'int'
+                elif member['type']['kind'] == 'optional':
+                    if member['type']['inner']['kind'] == 'string':
+                        ftype = 'str'
+                    elif member['type']['inner']['kind'] in ('int32', 'tick'):
+                        ftype = 'int'
+                    else:
+                        raise \
+                            ValueError( \
+                            f"Error: unknown optional type {member['type']}")
+                else:
+                    raise ValueError(f"Error: unknown type {member['type']}")
+                name = name_to_snake(member["name"])
+                if ftype != '':
+                    ftype = f': {ftype}'
+                out_file.write(f"        self.{name}{ftype} = {name}\n")
+            out_file.write('\n')
+            out_file.write('    # first byte of data\n')
+            out_file.write('    # has to be the first byte of the message payload\n')
+            out_file.write('    # NOT the chunk header and NOT the message id\n')
+            out_file.write('    def unpack(self, data: bytes) -> bool:\n')
+            if len(msg['members']) > 0:
+                out_file.write('        unpacker = Unpacker(data)\n')
+            out_file.write(gen_unpack_members(msg))
+            out_file.write('        return True\n')
+            out_file.write('\n')
+            out_file.write('    def pack(self) -> bytes:\n')
+            out_file.write(gen_pack_return(msg))
+
+    def generate(self, spec: str) -> None:
+        print(f"generating classes from {spec} ...")
+        with open(spec) as spec_io:
+            spec_data: SpecJson = json.load(spec_io)
+            # for msg in [spec_data['game_messages'][1]]:
+            game_enums: list[GameEnumJson] = spec_data['game_enumerations']
+            game_messages: list[NetMessageJson] = spec_data['game_messages']
+            system_messages: list[NetMessageJson] = spec_data['system_messages']
+            gen_enum_file7(game_enums)
+            gen_match_file7('game', game_messages)
+            gen_match_file7('system', system_messages)
+            for msg in game_messages:
+                self.generate_msg(msg, 'game')
+            for msg in system_messages:
+                self.generate_msg(msg, 'system')
 
 def main() -> None:
     dirname = os.path.dirname(__file__)
@@ -602,7 +624,8 @@ def main() -> None:
             dirname,
             '../../libtw2/gamenet/generate/spec/teeworlds-0.7.5.json')
     if os.path.exists(spec_07):
-        generate(spec_07)
+        generator = CodeGenerator()
+        generator.generate(spec_07)
     else:
         print(f"Error: file not found {spec_07}")
         print("        try running these commands")