Script for enum modernization in Urho

convert_enums.py

import sys
import os
import re

re_enum = re.compile(r'\s*enum\s*(\w+)\s*(:.*)?\s*')
re_enum_value = re.compile(r'\s*(\w+)(?:\s*=\s*(.+))?,?(?:\s*\/\/.*)?\s*')
folders_blacklist = [
    # 'Urho3D/Audio',
    # 'Urho3D/Container',
    # 'Urho3D/Core',
    'Urho3D/CSharp',
    # 'Urho3D/Engine',
    # 'Urho3D/Glow',
    'Urho3D/Graphics',
    'Urho3D/IK',
    'Urho3D/Input',
    'Urho3D/IO',
    # 'Urho3D/Math',
    'Urho3D/Navigation',
    'Urho3D/Network',
    'Urho3D/Physics',
    'Urho3D/Resource',
    'Urho3D/RmlUI',
    'Urho3D/Scene',
    # 'Urho3D/Script',
    'Urho3D/SystemUI',
    'Urho3D/UI',
    'Urho3D/Urho2D',
]
enums_blacklist = [
    'LightVSVariation',
    'VertexLightVSVariation',
    'LightPSVariation',
    'DeferredLightVSVariation',
    'DeferredLightPSVariation',
    'LoopMode',
    'ShaderType',
    'Algorithm',
    'Feature',
    'CurveType',
    'Key',
    'Scancode',
]


def is_file_blocked(file_name):
    for blocked_folder in folders_blacklist:
        if blocked_folder in file_name:
            return True
    return False


def convert_name(value):
    return value.replace('_', ' ').title().replace(' ', '')


def sanitize_identifier(value):
    return value if value[0] < '0' or value[0] > '9' else '_' + value


def make_local_regex(value):
    return re.compile(r'\b%s\b' % value)


def make_global_regex(value):
    return re.compile(r'\b%s\b' % value)


def strip_common_prefix(values):
    common_prefix = os.path.commonprefix(values)
    return [sanitize_identifier(value[len(common_prefix):]) for value in values]


class EnumInfo:
    def __init__(self, file_name, enum_name):
        self.file_name = file_name
        self.enum_name = enum_name
        self.enum_values = []
        self.title_line = None
        self.begin_line = None
        self.end_line = None

    def __repr__(self):
        return self.enum_name

    def process(self):
        # Extract 'max' values
        max_enum_values = [enum_value for enum_value in self.enum_values if enum_value.startswith('MAX_')]
        max_enum_values_new = [convert_name(value) for value in max_enum_values]
        max_local_regexes = [make_local_regex(value) for value in max_enum_values]
        max_global_regexes = [make_global_regex(value) for value in max_enum_values]
        self.max_enum_values = list(zip(max_local_regexes, max_global_regexes, max_enum_values, max_enum_values_new))

        # Extract 'simple' values
        simple_enum_values = [enum_value for enum_value in self.enum_values if not enum_value.startswith('MAX_')]
        simple_enum_values_new = strip_common_prefix([convert_name(value) for value in simple_enum_values])
        simple_local_regexes = [make_local_regex(value) for value in simple_enum_values]
        simple_global_regexes = [make_global_regex(value) for value in simple_enum_values]
        self.simple_enum_values = list(zip(simple_local_regexes, simple_global_regexes, simple_enum_values, simple_enum_values_new))


def find_source_files(base_dir):
    file_names = []
    for subdir, _, files in os.walk(base_dir):
        for file in files:
            full_file_name = subdir + os.sep + file

            if full_file_name.endswith(".h") or full_file_name.endswith(".cpp") or full_file_name.endswith(".hpp"):
                file_names.append(full_file_name)
    return file_names


def find_enums(file_names):
    enums = []
    for file_name in file_names:
        with open(file_name, "r") as source_file:
            source_file_lines = [line for line in source_file]
            line_index = 0
            while line_index < len(source_file_lines):
                # Loop until find enum
                line_text = source_file_lines[line_index]
                line_index = line_index + 1
                enum_match = re_enum.fullmatch(line_text)
                if enum_match is None:
                    continue

                enum_info = EnumInfo(file_name, enum_match.group(1))
                enum_info.title_line = line_index - 1

                # Loop until find {
                while line_index < len(source_file_lines):
                    line_text = source_file_lines[line_index].strip()
                    line_index = line_index + 1
                    if line_text == '{':
                        enum_info.begin_line = line_index
                        break

                # Loop until find }
                while line_index < len(source_file_lines):
                    line_text = source_file_lines[line_index].strip()
                    line_index = line_index + 1
                    if len(line_text) > 0 and line_text[0] == '}':
                        enum_info.end_line = line_index - 1
                        break
                    value_match = re_enum_value.fullmatch(line_text)
                    if value_match is not None:
                        enum_info.enum_values.append(value_match.group(1))

                # Append if valid
                if enum_info.begin_line and enum_info.end_line:
                    enum_info.process()
                    enums.append(enum_info)
    return enums


def find_and_replace(file_name, enums):
    with open(file_name, 'r') as file:
        filedata = file.read()

    for enum in enums:
        for _, regex, _, new_value in enum.max_enum_values:
            filedata = regex.sub(f'{new_value}', filedata)
        for _, regex, _, new_value in enum.simple_enum_values:
            filedata = regex.sub(f'{enum.enum_name}::{new_value}', filedata)

    with open(file_name, 'w') as file:
        file.write(filedata)


def main():
    if len(sys.argv) != 2:
        print('Command line syntax: gene.py path/to/source')
        return

    source_folder = sys.argv[1]
    urho3d_folder = source_folder + '/Urho3D/'

    # Scan for Urho3D source files
    print(f'Scanning {urho3d_folder}...')
    urho3d_files = find_source_files(urho3d_folder)
    urho3d_files = [file_name for file_name in urho3d_files if not is_file_blocked(file_name)]
    print(f'{len(urho3d_files)} files found!')

    # Scan for enums
    enums = find_enums(urho3d_files)
    enums = [enum for enum in enums if enum.enum_name not in enums_blacklist]

    # Group enums by files
    enums_by_file = {}
    for enum in enums:
        if enum.file_name not in enums_by_file:
            enums_by_file[enum.file_name] = []
        enums_by_file[enum.file_name].append(enum)

    print(f'{len(enums_by_file)} files with enums found!')

    # Patch enums in source files
    print('Patching enum declarations...')
    for file_name in enums_by_file:
        # Read source file
        with open(file_name, "r") as source_file:
            source_file_lines = [line for line in source_file]

        # Sort enums from end to begin
        enums_in_file = sorted(enums_by_file[file_name], key=lambda enum: enum.begin_line, reverse=True)

        for enum in enums_in_file:
            # Patch title
            source_file_lines[enum.title_line] = source_file_lines[enum.title_line].replace('enum ', 'enum class ')

            # Patch enum values
            for line_index in range(enum.begin_line, enum.end_line):
                line = source_file_lines[line_index]
                for regex, _, _, new_value in enum.max_enum_values + enum.simple_enum_values:
                    line = regex.sub(new_value, line)
                source_file_lines[line_index] = line

        # Write source file back
        with open(file_name, "w") as dest_file:
            dest_file.writelines(source_file_lines)
    print('Done')

    # Find and replace in all files
    all_files = [file_name for file_name in find_source_files(source_folder) if 'ThirdParty' not in file_name]
    print(f'Patching enum usage in {len(all_files)} files...')
    for file_name in all_files:
        find_and_replace(file_name, enums)
    print('Done')

    # Patch enums in source files
    print('Spawning extra declarations...')
    for file_name in enums_by_file:
        # Read source file
        with open(file_name, "r") as source_file:
            source_file_lines = [line for line in source_file]

        # Sort enums from end to begin
        enums_in_file = sorted(enums_by_file[file_name], key=lambda enum: enum.begin_line, reverse=True)

        for enum in enums_in_file:
            # Spawn max values
            new_lines = []
            if len(enum.max_enum_values) > 0:
                new_lines.append('\n')
            for _, _, _, new_value in enum.max_enum_values:
                new_lines.append(f'static const auto {new_value} = static_cast<unsigned>({enum.enum_name}::{new_value});\n')

            # Spawn legacy enums
            new_lines.append('\n')
            new_lines.append('#if 1 // #ifdef URHO3D_LEGACY_ENUMS\n')
            for _, _, old_value, new_value in enum.simple_enum_values + enum.max_enum_values:
                new_lines.append(f'static const {enum.enum_name} {old_value} = {enum.enum_name}::{new_value};\n')
            new_lines.append('#endif\n')
            new_lines.append('\n')

            source_file_lines = source_file_lines[:enum.end_line + 1] + new_lines + source_file_lines[enum.end_line + 1:]

        # Write source file back
        with open(file_name, "w") as dest_file:
            dest_file.writelines(source_file_lines)
    print('Done')



main()