drgn/scripts/gen_tests_elf_py.py

#!/usr/bin/env python3
# Copyright (c) Meta Platforms, Inc. and affiliates.
# SPDX-License-Identifier: LGPL-2.1-or-later

import argparse
from pathlib import Path
import re
import sys


def main() -> None:
    argparse.ArgumentParser(
        description="Generate tests/elf.py from libdrgn/include/elf.h"
    ).parse_args()

    contents = Path("libdrgn/include/elf.h").read_text()
    contents = re.sub(r"/\*.*?\*/", "", contents, flags=re.DOTALL)
    contents = re.sub(r"\\\n", "", contents)

    enums = {
        name: []
        for name in (
            "ET",
            "PT",
            "SHF",
            "SHN",
            "SHT",
            "STB",
            "STT",
            "STV",
        )
    }
    for match in re.finditer(
        r"^\s*#\s*define\s+(?P<enum>"
        + "|".join(enums)
        + r")_(?P<name>\w+)\s+(?:(?P<value>0x[0-9a-fA-F]+|[0-9]+)|(?:\(\s*1U?\s*<<\s*(?P<bitshift>[0-9]+)\s*\)))",
        contents,
        re.MULTILINE,
    ):
        enum = match.group("enum")
        name = match.group("name")
        if match.group("value"):
            value = int(match.group("value"), 0)
        else:
            value = 1 << int(match.group("bitshift"), 10)
        enums[enum].append((name, value))

    f = sys.stdout
    f.write(
        """\
# Copyright (c) Meta Platforms, Inc. and affiliates.
# SPDX-License-Identifier: LGPL-2.1-or-later
# Generated by scripts/gen_tests_elf_py.py.

import enum
from typing import Text
"""
    )
    for type_name, constants in enums.items():
        assert constants
        f.write(f"\n\nclass {type_name}(enum.IntEnum):\n")
        for name, value in constants:
            f.write(f"    {name} = 0x{value:X}\n")
        f.write(
            f"""
    @classmethod
    def str(cls, value: int) -> Text:
        try:
            return f"{type_name}_{{cls(value).name}}"
        except ValueError:
            return hex(value)
"""
        )


if __name__ == "__main__":
    main()