BA-translator/generate_flatbuffer_folders.py

#!/usr/bin/env python3
"""
FlatBuffers Python Code Generator

Generates Python modules from .fbs schema files for BA-translator.
"""

import os
import sys
import subprocess
import shutil
import argparse
from pathlib import Path


def log_info(message, verbose=False):
    """Print info message if verbose mode is enabled."""
    if verbose:
        print(f"[INFO] {message}")
        
def log_error(message):
    """Print error message."""
    print(f"[ERROR] {message}")

def check_flatc(flatc_path='flatc'):
    """Check if flatc compiler is available."""
    try:
        result = subprocess.run([flatc_path, '--version'], capture_output=True, text=True)
        return result.returncode == 0
    except FileNotFoundError:
        return False
        
def get_namespace(fbs_file):
    """Extract namespace from .fbs file."""
    try:
        with open(fbs_file, 'r', encoding='utf-8') as f:
            for line in f:
                line = line.strip()
                if line.startswith('namespace '):
                    return line.replace('namespace ', '').replace(';', '').strip()
    except Exception:
        pass
    return 'FlatData'
        
def create_init_file(directory):
    """Create __init__.py file in directory."""
    init_file = Path(directory) / '__init__.py'
    if not init_file.exists():
        init_file.write_text('# Generated FlatBuffers Python package\n')
                
def generate_python_modules(fbs_files, flatc_path='flatc', clean=False, verbose=False):
    """Generate Python modules from .fbs files."""
    if not fbs_files:
        log_error("No .fbs files provided")
        return False
        
    if not check_flatc(flatc_path):
        log_error(f"flatc compiler not found. Please install FlatBuffers.")
        return False
        
    script_dir = Path(__file__).parent.absolute()
    
    # Clean existing files if requested
    if clean:
        for dirname in ['FlatData', 'MX']:
            dir_path = script_dir / dirname
            if dir_path.exists():
                log_info(f"Cleaning {dir_path}", verbose)
                shutil.rmtree(dir_path)
                
    generated_dirs = set()
    
    for fbs_file in fbs_files:
        if not Path(fbs_file).exists():
            log_error(f"Schema file not found: {fbs_file}")
            return False
            
        log_info(f"Processing {fbs_file}", verbose)
        
        # Generate Python code using flatc
        cmd = [
            flatc_path,
            '--python',
            '--gen-object-api',
            '-o', str(script_dir),
            fbs_file
        ]
        
        try:
            result = subprocess.run(cmd, capture_output=True, text=True)
            
            if result.returncode != 0:
                log_error(f"flatc failed for {fbs_file}")
                if result.stderr:
                    print(f"stderr: {result.stderr}")
                return False
                
            # Get namespace and add to generated dirs
            namespace = get_namespace(fbs_file)
            generated_dirs.add(script_dir / namespace)
            log_info(f"Generated modules for {fbs_file}", verbose)
            
        except Exception as e:
            log_error(f"Exception running flatc: {e}")
            return False
            
    # Create __init__.py files
    for gen_dir in generated_dirs:
        if gen_dir.exists():
            create_init_file(gen_dir)
            log_info(f"Created __init__.py in {gen_dir}", verbose)
            
    return True

def find_fbs_files(directory='.'):
    """Find all .fbs files in directory."""
    fbs_files = []
    for root, _, files in os.walk(directory):
        for file in files:
            if file.endswith('.fbs'):
                fbs_files.append(os.path.join(root, file))
    return fbs_files


def main():
    """Main entry point."""
    parser = argparse.ArgumentParser(
        description='Generate Python modules from FlatBuffers schema files')
    
    parser.add_argument('fbs_files', nargs='*', help='.fbs schema files to process')
    parser.add_argument('--auto', action='store_true', help='Auto-find all .fbs files')
    parser.add_argument('--flatc-path', default='flatc', help='Path to flatc compiler')
    parser.add_argument('--clean', action='store_true', help='Clean existing files first')
    parser.add_argument('--verbose', '-v', action='store_true', help='Verbose output')
    
    args = parser.parse_args()
    
    # Get .fbs files
    if args.auto:
        fbs_files = find_fbs_files()
        if not fbs_files:
            log_error("No .fbs files found")
            sys.exit(1)
        log_info(f"Found {len(fbs_files)} .fbs files", args.verbose)
    elif args.fbs_files:
        fbs_files = args.fbs_files
    else:
        parser.print_help()
        sys.exit(1)
        
    # Generate modules
    success = generate_python_modules(
        fbs_files=fbs_files,
        flatc_path=args.flatc_path,
        clean=args.clean,
        verbose=args.verbose
    )
    
    if success:
        print("[SUCCESS] Python modules generated successfully")
    else:
        sys.exit(1)


if __name__ == '__main__':
    main()