generate_id_constants.py

from pathlib import Path
import re
import requests

HEADER = f"# DO NOT EDIT!\n# This file was automatically generated by \"{__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))