import argparse import sys import json from pathlib import Path # Add parent directories to the Python path to find utils and parsers source_path = Path(__file__).parent.parent sys.path.append(str(source_path / 'utils')) sys.path.append(str(source_path / 'parsers')) from file_utils import find_files_by_extension from json_utils import write_json from config_utils import load_config from trs_processor import TRSSceneProcessor def generate_guid_mappers(input_dir, output_dir, indent=None, shrink=False, ignored_folders=None): """ Finds all .meta files and generates JSON files mapping GUIDs to asset paths. This is a necessary setup step to allow the TRS processor to find prefab files. """ assets_dir = input_dir / "Assets" if not assets_dir.is_dir(): print(f"Error: 'Assets' directory not found in '{input_dir}'", file=sys.stderr) return None print("--> Finding all .meta files for GUID mapping...") meta_files = find_files_by_extension(str(assets_dir), '.meta', ignored_folders=ignored_folders, project_root=input_dir) print(f"--> Found {len(meta_files)} .meta files.") guid_map = {} print("--> Parsing .meta files and mapping GUIDs to asset paths...") for meta_file_path_str in meta_files: meta_file_path = Path(meta_file_path_str) asset_file_path = Path(meta_file_path_str.rsplit('.meta', 1)[0]) if not asset_file_path.is_file(): continue guid = None try: with open(meta_file_path, 'r', encoding='utf-8') as f: for line in f: if line.strip().startswith('guid:'): guid = line.strip().split(':')[1].strip() break except Exception as e: print(f"Warning: Could not read or parse guid from {meta_file_path}. {e}", file=sys.stderr) continue if guid: # Use the full, resolved path for the guid_map value guid_map[guid] = asset_file_path.resolve().as_posix() print("--> GUID mapping complete.") return guid_map def main(): """ Main function to run the TRS-only data extraction for scenes. """ parser = argparse.ArgumentParser( description="Generates a simplified JSON representation of scene hierarchies with only TRS data." ) parser.add_argument( "--input", type=str, required=True, help="The root directory of the target Unity project." ) parser.add_argument( "--output", type=str, required=True, help="The directory where the generated JSON files will be saved." ) args = parser.parse_args() # --- Load Configuration --- config = load_config() ignored_folders = config.get('ignored_folders', []) shrink_json = config.get('shrink_json', False) indent_level = config.get('indentation_level', 4) input_dir = Path(args.input).resolve() output_dir = Path(args.output).resolve() if not input_dir.is_dir(): print(f"Error: Input path '{input_dir}' is not a valid directory.", file=sys.stderr) sys.exit(1) # --- Setup Output Directory --- try: output_dir.mkdir(parents=True, exist_ok=True) print(f"Output will be saved to: {output_dir}") except OSError as e: print(f"Error: Could not create output directory '{output_dir}'. {e}", file=sys.stderr) sys.exit(1) assets_dir = input_dir / "Assets" if not assets_dir.is_dir(): print(f"Warning: 'Assets' directory not found in '{input_dir}'. Skipping all processing.", file=sys.stderr) return print("\n--- Running Scene TRS Extraction ---") # --- Task 1: Generate GUID Map (for prefab lookups) --- print("\n[1/2] Generating GUID Mappers...") guid_map = generate_guid_mappers( input_dir, output_dir, # Not saving mappers, just using the map in memory ignored_folders=ignored_folders ) if guid_map is None: print("Error: Failed to generate GUID map. Aborting.", file=sys.stderr) sys.exit(1) print("--> GUID Mapper generation complete.") # --- Task 2: Find and Process Scene Files --- print("\n[2/2] Parsing Scene Hierarchies for TRS data...") scene_files = find_files_by_extension(str(assets_dir), '.unity', ignored_folders=ignored_folders, project_root=input_dir) if not scene_files: print("--> No scene files found to process.") return print(f"--> Found {len(scene_files)} scene file(s) to process.") total_files = len(scene_files) for i, file_path_str in enumerate(scene_files): file_path = Path(file_path_str) output_json_path = output_dir / f"{file_path.stem}.json" try: print(f"\n--- Processing {file_path.name} ({i+1}/{total_files}) ---") processor = TRSSceneProcessor(guid_map) print(f" -> Loading and parsing file...") if not processor.load_documents(file_path): print(f"Warning: Could not load or parse {file_path.name}. Skipping.", file=sys.stderr) continue print(f" -> Building hierarchy and extracting TRS data...") hierarchy = processor.process() # The hierarchy is the list of root objects, which will be the top-level JSON array. write_json(hierarchy, output_json_path, indent=indent_level, shrink=shrink_json) print(f"--> Successfully processed TRS hierarchy for {file_path.name} -> {output_json_path}") except Exception as e: print(f"Error processing TRS hierarchy for {file_path.name}: {e}", file=sys.stderr) print("\nScene TRS extraction complete.") if __name__ == "__main__": main()