python ast 解析enum为C头文件

from _ast import Assign, ClassDef, FunctionDef, Import
import ast
from typing import Any

str="""
import enum

class MsgIdEnum(enum.Enum):
    MSG_ID_1 = (1, "消息1")
    MSG_ID_5 = (5, "消息5")
    MSG_ID_6 = (enum.auto(), "消息6")
    MSG_ID_7 = (enum.auto(), "消息7")
    MSG_ID_10 = (10, "消息10")
    MSG_ID_11 = (enum.auto(), "消息11")

    def __new__(cls, value, string_value):
        member = object.__new__(cls)
        member._value_ = value
        member.string_value = string_value
        return member
        
"""

class OStringStream():
    def __init__(self):
        self._buff=""
        self._len = 0
    
    @property
    def buff(self):
        return self._buff
        
    @property
    def len(self):
        return self._len
    
    def append(self, data):
        self._buff += data
        self._len += len(data)
        
class Visitor(ast.NodeVisitor):
    def __init__(self, file):
        self._file = file
        self._indent = 4
        self._data = OStringStream()
        self._data.append("""#ifndef MSG_ID_H
#define MSG_ID_H
""")
    def entry_indent(self, space=1):
        self._indent += space
        self._data.append(Visitor.backspace(self._indent))
        
    def exit_indent(self, space=1):
        self._indent -= space
        
    @staticmethod
    def backspace(num=1):
        return " "*num
    
    def line(self):
        self._data.append("\n")
    
    def visit(self, node):
        method = 'visit_' + node.__class__.__name__
        visitor = getattr(self, method, self.generic_visit)
        return visitor(node)
    
    def visit_FunctionDef(self, node: FunctionDef) -> Any:
        return True
    
    def visit_Assign(self, node: Assign) -> Any:
        is_auto_func = False
        msg_id_name = node.targets[0].id
        msg_id_value, doc = node.value.elts
        if isinstance(msg_id_value, ast.Call) and isinstance(msg_id_value.func, ast.Attribute) and msg_id_value.func.attr == 'auto':
            is_auto_func=True
        self.entry_indent(4)
        
        msg_id_value = "," if is_auto_func else f"={msg_id_value.value},"
        self._data.append(f"{msg_id_name}{msg_id_value:>10}{self.backspace(8)}//{self.backspace(4)}{doc.value}")
        self.line()
        self.exit_indent(4)
        return True
    
    def visit_Import(self, node: Import) -> Any:
        self._data.append(f"#include ")
        self.line()
        return True
    
    def visit_ClassDef(self, node: ClassDef) -> Any:
        if not node.name:
            print(f"Error:expect class name")
            return False
        self._data.append("enum %s%s" %(node.name, Visitor.backspace()))
        self.line()
        self._data.append("{")
        self.line()
        
        for stmt in node.body:
            res = self.visit(stmt)
            if not res:
                print(f"Error: stmt error")
                return False
        self._data.append("};")
        self.line()
        return True
    
    def generic_visit(self, node):
        return True
    
    def write_file(self):
        with open(self._file, "w") as f:
            f.write(self._data.buff)


if __name__ == "__main__":
    out = ast.parse(str)
    print(ast.dump(out, indent=4))
    for stmt in out.body:
        print(stmt)
        vs = Visitor("msg_id.h")
        vs.visit(stmt)
        vs._data.append("#endif")
        vs.write_file()




你可能感兴趣的:(python,ast)