generate_id_constants.py

import os
import re
from pathlib import Path

import requests

HEADER = f'# DO NOT EDIT!\n# This file was automatically generated by "{os.path.basename(__file__)}"\n'
PREFIXES = ["Protoss", "Terran", "Zerg", "Neutral", "Effect"]


def clike_enum_parse(code):
    # remove comments etc.
    code = re.sub(r"#[^\n]*\n", "", code, flags=re.DOTALL)
    code = re.sub(r"//[^\n]+\n", "", code, flags=re.DOTALL)
    code = re.sub(r"/\*.*?\*/", "", code, flags=re.DOTALL)

    # normalize whitespace
    code = re.sub(r"\s+", " ", code)

    # parse enum blocks
    enums = {}
    for enum in re.findall(r"enum(?: class)? ([a-zA-Z_][a-zA-Z0-9_]*) {\s?(.+?)\s?}", code):
        name, body = enum
        body = {key: int(value) for key, value in re.findall(r"([a-zA-Z_][a-zA-Z0-9_]*) = (\d+),?\s?", body)}
        enums[name] = body

    return enums


def generate_python_code(enums):
    assert {"UNIT_TYPEID", "ABILITY_ID", "UPGRADE_ID", "BUFF_ID"} <= enums.keys()

    sc2dir = Path("sc2/")
    idsdir = sc2dir / "ids"
    idsdir.mkdir(exist_ok=True)

    with (idsdir / "__init__.py").open("w") as f:
        f.write("\n".join([HEADER, f"__all__ = {[n.lower() for n in enums.keys()] !r}\n"]))

    for name, body in enums.items():
        class_name = "".join(p.capitalize() for p in name.split("_")).replace("Typeid", "TypeId")

        code = [HEADER, "import enum", "", f"class {class_name}(enum.Enum):"]

        for key, value in sorted(body.items(), key=lambda p: p[1]):

            # Unprefixed alternative names for structures
            if any(key.startswith(p.upper() + "_") for p in PREFIXES):
                code.append(f"    {key.split('_',1)[1]} = {value}")

            code.append(f"    {key} = {value}")

        code += ["", f"for item in {class_name}:", f"    globals()[item.name] = item", ""]

        with (idsdir / name.lower()).with_suffix(".py").open("w") as f:
            f.write("\n".join(code))


if __name__ == "__main__":
    r = requests.get("https://raw.githubusercontent.com/Blizzard/s2client-api/master/include/sc2api/sc2_typeenums.h")
    code = generate_python_code(clike_enum_parse(r.text))