MzTranslator/python/mz_core/reader.py

161 lines
4.2 KiB
Python

# -*- coding: utf-8 -*-
"""
Mz Translator Reader Module
Read Maya scene from .mz file (ZIP compressed, contains data.ma)
"""
import os
import zipfile
import tempfile
import io
import maya.OpenMaya as OpenMaya
def read_scene_from_mz(file_path, buffer=None):
"""
Read Maya scene from .mz file
Args:
file_path: Input .mz file path
buffer: Optional buffer for direct data reading
Returns:
str: data.ma content (if using buffer), or temp file path
Raises:
Exception: Raised when read fails
"""
if not os.path.exists(file_path):
raise FileNotFoundError(f"MZ file not found: {file_path}")
try:
# Open ZIP file
with zipfile.ZipFile(file_path, 'r') as zf:
# Check if contains data.ma
if 'data.ma' not in zf.namelist():
raise ValueError("Invalid .mz file: no data.ma found")
# Read data.ma content
ma_content = zf.read('data.ma')
if buffer is not None:
# Write to provided buffer
if hasattr(buffer, 'write'):
buffer.write(ma_content)
else:
raise ValueError("buffer must be a writable file-like object")
return buffer
else:
# Return content string (for stream reading)
return ma_content
except zipfile.BadZipFile:
raise ValueError("Invalid .mz file: not a valid ZIP file")
except Exception as e:
raise RuntimeError(f"Failed to read .mz file: {str(e)}")
def read_to_temp_file(file_path):
"""
Read from .mz file and write to temp file
Args:
file_path: Input .mz file path
Returns:
str: Temp data.ma file path
Note:
Need to call cleanup_temp_file after use
"""
if not os.path.exists(file_path):
raise FileNotFoundError(f"MZ file not found: {file_path}")
# Create temp directory
temp_dir = tempfile.mkdtemp(prefix='mz_import_')
temp_ma_path = os.path.normpath(os.path.join(temp_dir, 'data.ma'))
try:
# Open ZIP file
with zipfile.ZipFile(file_path, 'r') as zf:
# Check if contains data.ma
if 'data.ma' not in zf.namelist():
raise ValueError("Invalid .mz file: no data.ma found")
# Read and write to temp file
ma_content = zf.read('data.ma')
with open(temp_ma_path, 'wb') as f:
f.write(ma_content)
return temp_ma_path
except Exception as e:
# Cleanup temp files
_cleanup_temp_dir(temp_dir)
raise RuntimeError(f"Failed to extract .mz file: {str(e)}")
def read_to_buffer(buffer):
"""
Read .mz file data.ma content to buffer
Args:
buffer: Writable file object (must have write method)
Returns:
int: Number of bytes read
Note:
buffer should be an IOBase or similar object for Maya's file command
"""
# This is a wrapper for supporting stream reading
# Will be called in MPxFileTranslator.reader()
pass
def _cleanup_temp_dir(temp_dir):
"""
Cleanup temp directory
Args:
temp_dir: Temp directory path to delete
"""
import shutil
try:
if os.path.exists(temp_dir):
shutil.rmtree(temp_dir)
except Exception:
pass
def get_file_info(file_path):
"""
Get .mz file info
Args:
file_path: .mz file path
Returns:
dict: Dictionary containing file info
"""
if not os.path.exists(file_path):
raise FileNotFoundError(f"MZ file not found: {file_path}")
info = {
'file_size': os.path.getsize(file_path),
'files': []
}
try:
with zipfile.ZipFile(file_path, 'r') as zf:
for name in zf.namelist():
info['files'].append({
'name': name,
'size': zf.getinfo(name).file_size,
'compressed_size': zf.getinfo(name).compress_size
})
except Exception as e:
raise ValueError(f"Invalid .mz file: {str(e)}")
return info