GoLang 生成常量值map

在 golang 中有时候会需要用到类似其他语言中 enum 类型的通过值获取字段名称功能，通过以下脚本可以读取 go 文件中所有 全局 var/const 定义，并根据 值/iota值 生成 值:常量名 map。

gen_const_map.py

import io

import regex


class Const:
    const_name: str
    const_value: str

    def declare_str(self):
        return f'\t{self.const_value}: "{self.const_name}",\n'


class ConstGroup:
    const_type: str
    const_list: list[Const]
    map_keys: dict

    def __init__(self):
        self.map_keys = {}
        self.const_list = []

    def declare_str(self):
        s = f'var {self.const_type}Map = map[{self.const_type}]string{{\n'
        for c in self.const_list:
            s += c.declare_str()
        return s + "}\n"


const_map = {}
type_underlying = {}

with io.open("const.go", "rb") as f:
    const_file = f.read().decode("utf-8")

lines = const_file.split("\n")

cursor_type = ""
const_syntax = False
iota = -1

# 自定义常量的前缀
trim_prefix_reg = regex.compile("^(?P<prefix>Clz|SV|BT|TD|SP|FM|SGP|PassType|ShaderCompilerPlatform).+")

# extract const
for line in lines:
    params = [key.strip() for key in line.split(" ") if key != ""]
    if not params:
        continue
    if params[0] == "type":
        # type declare
        type_underlying[params[1]] = params[2]
    if params[0] == "const" and params[1] == "(":
        # const definition
        const_syntax = True
        continue
    elif params[0] == ")":
        const_syntax = False
        iota = -1
        cursor_type = ""
    if const_syntax:
        const_name = params[0]
        if len(params) == 4:
            # type declared
            const_type = params[1]
            const_value = params[3]
            if const_value == "iota":
                cursor_type = const_type
                iota = 0
                const_value = iota.__str__()
                iota += 1
        else:
            const_value = iota.__str__()
            const_type = cursor_type
            iota += 1
        if const_name == "_":
            continue
        if const_type in const_map:
            group = const_map[const_type]
        else:
            group = ConstGroup()
            group.const_type = const_type
            const_map[const_type] = group
        if const_value in group.map_keys:
            continue
        # trim prefix
        m = trim_prefix_reg.match(const_name)
        if m:
            prefix = m.group("prefix")
            const_name = const_name.removeprefix(prefix)
        # trim prefix type
        if const_name.startswith(const_type):
            const_name = const_name[len(const_type):]
        group.map_keys[const_value] = True  # avoid duplicated key
        const = Const()
        const.const_name = const_name
        const.const_value = const_value
        group.const_list.append(const)

print(f"ConstTypes: {type_underlying}")

# 修改包名
sb = "package unity\n\n"
sb += "// This is generated file. Do not modify.\n\n"

# write const
for k in const_map:
    group = const_map[k]
    sb += group.declare_str()
    sb += "\n"

with io.open("const_map.go", "wb") as f:
    f.write(sb.encode("utf-8"))

print("generated const_map.go")

使用

示例 const.go

type FileType int

const (
	FileTypeAssetsFile FileType = iota
	FileTypeBundleFile
	FileTypeWebFile
	FileTypeResourceFile
	FileTypeZipFile
)

生成文件 const_map.go

// This is generated file. Do not modify.

var FileTypeMap = map[FileType]string{
	0: "AssetsFile",
	1: "BundleFile",
	2: "WebFile",
	3: "ResourceFile",
	4: "ZipFile",
}