yt-dlp-dags/ytdlp-ops-auth/generate-thrift.py
2025-04-03 13:25:49 +03:00

195 lines
7.5 KiB
Python
Executable File

#!/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()