#!/usr/bin/env python3 import os import subprocess import shutil from pathlib import Path import xml.etree.ElementTree as ET # Update paths to match actual project structure THRIFT_MODEL_DIR = Path("thrift_model") SERVICES_DIR = THRIFT_MODEL_DIR / "services" DATA_DIR = THRIFT_MODEL_DIR / "data" GEN_PY_DIR = THRIFT_MODEL_DIR / "gen_py" def get_version_from_pom(): """Parse version from pom.xml""" pom_path = THRIFT_MODEL_DIR / "pom.xml" tree = ET.parse(pom_path) root = tree.getroot() # XML namespaces ns = {'mvn': 'http://maven.apache.org/POM/4.0.0'} version = root.find('mvn:version', ns).text if version.endswith('-SNAPSHOT'): version = version.replace('-SNAPSHOT', '.dev0') return version def find_thrift_files(): """Find all .thrift files in the thrift_model directory""" data_files = list(DATA_DIR.glob("*.thrift")) service_files = list(SERVICES_DIR.glob("*.thrift")) # Process data files first (for dependencies), then service files return data_files + service_files def generate_python_code(thrift_files): """Generate Python code from Thrift files""" # First process data files (for dependencies) data_files = [f for f in thrift_files if f.parent == DATA_DIR] service_files = [f for f in thrift_files if f.parent == SERVICES_DIR] # Process in the right order: first data files, then service files ordered_files = data_files + service_files for thrift_file in ordered_files: print(f"Generating code for {thrift_file}...") try: subprocess.run([ "thrift", "--gen", "py", "-out", str(GEN_PY_DIR), str(thrift_file) ], check=True) print(f"Successfully generated code for {thrift_file}") except subprocess.CalledProcessError as e: print(f"Error generating code for {thrift_file}: {e}") raise def create_init_files(): """Create __init__.py files in all generated directories""" for root, dirs, files in os.walk(GEN_PY_DIR): path = Path(root) init_file = path / "__init__.py" if not init_file.exists(): print(f"Creating __init__.py in {path}") with open(init_file, 'w') as f: # For the top-level pangramia directory, we don't need special content if path.name == "pangramia": pass # For module directories, add the standard __all__ pattern if there are modules elif any(f.endswith('.py') and f != '__init__.py' for f in files): modules = [f[:-3] for f in files if f.endswith('.py') and f != '__init__.py'] if modules: f.write(f"__all__ = {repr(modules)}\n") # Ensure we have an __init__.py in the thrift_model directory thrift_model_init = THRIFT_MODEL_DIR / "__init__.py" if not thrift_model_init.exists(): print(f"Creating {thrift_model_init}") thrift_model_init.touch() def clean_gen_py(): """Clean the gen_py directory before generation""" if GEN_PY_DIR.exists(): print(f"Cleaning {GEN_PY_DIR}...") shutil.rmtree(GEN_PY_DIR) print(f"Cleaned {GEN_PY_DIR}") # Recreate the directory GEN_PY_DIR.mkdir(parents=True, exist_ok=True) def update_version_file(): """Update the version in __init__.py""" version = get_version_from_pom() print(f"Detected version from pom.xml: {version}") # Update the version in __init__.py init_path = Path("__init__.py") if init_path.exists(): with open(init_path, 'r') as f: content = f.read() # Replace the VERSION assignment if it exists if "VERSION = " in content: new_content = [] for line in content.splitlines(): if line.startswith("VERSION = "): new_content.append(f'VERSION = "{version}"') else: new_content.append(line) with open(init_path, 'w') as f: f.write('\n'.join(new_content)) print(f"Updated version in __init__.py to {version}") def main(): # Ensure directories exist SERVICES_DIR.mkdir(parents=True, exist_ok=True) DATA_DIR.mkdir(parents=True, exist_ok=True) # Clean existing generated code clean_gen_py() # Find all Thrift files thrift_files = find_thrift_files() if not thrift_files: print("No .thrift files found in thrift_model directory") return print(f"Found {len(thrift_files)} Thrift files to process") # Generate Python code generate_python_code(thrift_files) # Create __init__.py files create_init_files() # Update version file update_version_file() # Create a symbolic link to make the modules importable try: # Check if we're in the project root if not (Path.cwd() / "thrift_model").exists(): print("Warning: Not running from project root, symbolic link may not work correctly") # Create pangramia directory if it doesn't exist pangramia_dir = Path("pangramia") if not pangramia_dir.exists(): pangramia_dir.mkdir(parents=True, exist_ok=True) (pangramia_dir / "__init__.py").touch() print(f"Created {pangramia_dir} directory with __init__.py") # Create symbolic link from pangramia -> thrift_model/gen_py/pangramia link_path = Path("pangramia") # Link in the project root target_path = GEN_PY_DIR / "pangramia" # Ensure the target directory exists before creating the link if not target_path.exists(): print(f"Warning: Target directory {target_path} does not exist, cannot create symbolic link") else: # Remove existing link or directory at the destination if link_path.is_symlink(): print(f"Removing existing symbolic link: {link_path}") link_path.unlink() elif link_path.is_dir(): print(f"Removing existing directory: {link_path}") shutil.rmtree(link_path) elif link_path.exists(): # Handle case where it might be a file print(f"Removing existing file: {link_path}") link_path.unlink() # Create the new symbolic link try: # Use relative path for the link source for better portability relative_target = os.path.relpath(target_path, start=link_path.parent) os.symlink(relative_target, link_path, target_is_directory=True) print(f"Created symbolic link: {link_path} -> {relative_target}") except Exception as e: print(f"Error creating symbolic link: {e}") print("You may need to manually add the generated code to your Python path") # This else block corresponds to the `if not target_path.exists():` check further up # else: # print(f"Warning: Target directory {yt_target} does not exist, cannot create symbolic link") except Exception as e: print(f"An unexpected error occurred during symlink setup: {e}") # Optionally re-raise or handle more specifically print("\nThrift code generation completed successfully!") print(f"Generated Python code in {GEN_PY_DIR}") print(f"Current version: {get_version_from_pom()}") if __name__ == "__main__": main()