pystorm/examples/basic_operations.py

#!/usr/bin/env python3
"""
Example: Basic MPQ Archive Operations
Demonstrates how to work with MPQ archives using PyStorm
"""

from pystorm import MPQArchive, MPQ_CREATE_ARCHIVE_V2, MPQ_FILE_COMPRESS
import sys


def basic_operations_example():
    """Demonstrate basic MPQ archive operations"""
    print("=== Basic MPQ Archive Operations ===\n")
    
    # Create a new archive
    print("1. Creating a new archive...")
    try:
        with MPQArchive("test_archive.mpq", flags=MPQ_CREATE_ARCHIVE_V2) as archive:
            print("   ✓ Archive created successfully")
            
            # Add a text file
            print("\n2. Adding a text file...")
            test_content = b"Hello, MPQ World! This is a test file."
            with open("test_file.txt", "wb") as f:
                f.write(test_content)
            
            archive.add_file("test_file.txt", "internal/test.txt", 
                           flags=MPQ_FILE_COMPRESS)
            print("   ✓ File added successfully")
            
            # Flush changes
            archive.flush()
            print("   ✓ Changes flushed to disk")
    except Exception as e:
        print(f"   ✗ Error: {e}")
        return
    
    # Open and read from the archive
    print("\n3. Opening the archive and reading the file...")
    try:
        with MPQArchive("test_archive.mpq") as archive:
            # Check if file exists
            if archive.has_file("internal/test.txt"):
                print("   ✓ File exists in archive")
                
                # Open and read the file
                with archive.open_file("internal/test.txt") as mpq_file:
                    content = mpq_file.read()
                    print(f"   ✓ File content: {content.decode('utf-8')}")
            else:
                print("   ✗ File not found in archive")
    except Exception as e:
        print(f"   ✗ Error: {e}")
        return
    
    # List files in archive
    print("\n4. Listing all files in archive...")
    try:
        with MPQArchive("test_archive.mpq") as archive:
            files = archive.find_files("*")
            print(f"   Found {len(files)} file(s):")
            for file_info in files:
                print(f"     - {file_info['name']}")
                print(f"       Size: {file_info['size']} bytes")
                print(f"       Compressed: {file_info['compressed_size']} bytes")
    except Exception as e:
        print(f"   ✗ Error: {e}")
        return
    
    # Extract file
    print("\n5. Extracting file from archive...")
    try:
        with MPQArchive("test_archive.mpq") as archive:
            archive.extract_file("internal/test.txt", "extracted_test.txt")
            print("   ✓ File extracted successfully")
            
            # Verify extraction
            with open("extracted_test.txt", "rb") as f:
                extracted_content = f.read()
            print(f"   ✓ Extracted content: {extracted_content.decode('utf-8')}")
    except Exception as e:
        print(f"   ✗ Error: {e}")
        return
    
    # Clean up
    print("\n6. Cleaning up test files...")
    import os
    for file in ["test_file.txt", "extracted_test.txt", "test_archive.mpq"]:
        try:
            if os.path.exists(file):
                os.remove(file)
            print(f"   ✓ Removed {file}")
        except Exception as e:
            print(f"   ✗ Error removing {file}: {e}")
    
    print("\n=== Example completed successfully! ===")


def main():
    """Main entry point"""
    try:
        basic_operations_example()
    except KeyboardInterrupt:
        print("\n\nInterrupted by user")
        sys.exit(1)
    except Exception as e:
        print(f"\n\nUnexpected error: {e}")
        sys.exit(1)


if __name__ == "__main__":
    main()