mirror of
https://github.com/ksyasuda/dotfiles.git
synced 2026-02-27 12:22:43 -08:00
339 lines
10 KiB
Python
339 lines
10 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Screenshot utilities with dual-mode support.
|
|
|
|
Provides unified screenshot handling with:
|
|
- File-based mode: Persistent artifacts for test documentation
|
|
- Inline base64 mode: Vision-based automation for agent analysis
|
|
- Size presets: Token optimization (full/half/quarter/thumb)
|
|
- Semantic naming: {appName}_{screenName}_{state}_{timestamp}.png
|
|
|
|
Supports resize operations via PIL (optional dependency).
|
|
|
|
Used by:
|
|
- test_recorder.py - Step-based screenshot recording
|
|
- app_state_capture.py - State snapshot captures
|
|
"""
|
|
|
|
import base64
|
|
import os
|
|
import subprocess
|
|
import sys
|
|
from datetime import datetime
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
# Try to import PIL for resizing, but make it optional
|
|
try:
|
|
from PIL import Image
|
|
|
|
HAS_PIL = True
|
|
except ImportError:
|
|
HAS_PIL = False
|
|
|
|
|
|
def generate_screenshot_name(
|
|
app_name: str | None = None,
|
|
screen_name: str | None = None,
|
|
state: str | None = None,
|
|
timestamp: str | None = None,
|
|
extension: str = "png",
|
|
) -> str:
|
|
"""Generate semantic screenshot filename.
|
|
|
|
Format: {appName}_{screenName}_{state}_{timestamp}.{ext}
|
|
Falls back to: screenshot_{timestamp}.{ext}
|
|
|
|
Args:
|
|
app_name: Application name (e.g., 'MyApp')
|
|
screen_name: Screen name (e.g., 'Login')
|
|
state: State description (e.g., 'Empty', 'Filled', 'Error')
|
|
timestamp: ISO timestamp (uses current time if None)
|
|
extension: File extension (default: 'png')
|
|
|
|
Returns:
|
|
Semantic filename ready for safe file creation
|
|
|
|
Example:
|
|
name = generate_screenshot_name('MyApp', 'Login', 'Empty')
|
|
# Returns: 'MyApp_Login_Empty_20251028-143052.png'
|
|
|
|
name = generate_screenshot_name()
|
|
# Returns: 'screenshot_20251028-143052.png'
|
|
"""
|
|
if timestamp is None:
|
|
timestamp = datetime.now().strftime("%Y%m%d-%H%M%S")
|
|
|
|
# Build semantic name
|
|
if app_name or screen_name or state:
|
|
parts = [app_name, screen_name, state]
|
|
parts = [p for p in parts if p] # Filter None/empty
|
|
name = "_".join(parts) + f"_{timestamp}"
|
|
else:
|
|
name = f"screenshot_{timestamp}"
|
|
|
|
return f"{name}.{extension}"
|
|
|
|
|
|
def get_size_preset(size: str = "half") -> tuple[float, float]:
|
|
"""Get scale factors for size preset.
|
|
|
|
Args:
|
|
size: 'full', 'half', 'quarter', 'thumb'
|
|
|
|
Returns:
|
|
Tuple of (scale_x, scale_y) for resizing
|
|
|
|
Example:
|
|
scale_x, scale_y = get_size_preset('half')
|
|
# Returns: (0.5, 0.5)
|
|
"""
|
|
presets = {
|
|
"full": (1.0, 1.0),
|
|
"half": (0.5, 0.5),
|
|
"quarter": (0.25, 0.25),
|
|
"thumb": (0.1, 0.1),
|
|
}
|
|
return presets.get(size, (0.5, 0.5))
|
|
|
|
|
|
def resize_screenshot(
|
|
input_path: str,
|
|
output_path: str | None = None,
|
|
size: str = "half",
|
|
quality: int = 85,
|
|
) -> tuple[str, int, int]:
|
|
"""Resize screenshot for token optimization.
|
|
|
|
Requires PIL (Pillow). Falls back gracefully without it.
|
|
|
|
Args:
|
|
input_path: Path to original screenshot
|
|
output_path: Output path (uses input_path if None)
|
|
size: 'full', 'half', 'quarter', 'thumb'
|
|
quality: JPEG quality (1-100, default: 85)
|
|
|
|
Returns:
|
|
Tuple of (output_path, width, height) of resized image
|
|
|
|
Raises:
|
|
FileNotFoundError: If input file doesn't exist
|
|
ValueError: If PIL not installed and size != 'full'
|
|
|
|
Example:
|
|
output, w, h = resize_screenshot(
|
|
'screenshot.png',
|
|
'screenshot_half.png',
|
|
'half'
|
|
)
|
|
print(f"Resized to {w}x{h}")
|
|
"""
|
|
input_file = Path(input_path)
|
|
if not input_file.exists():
|
|
raise FileNotFoundError(f"Screenshot not found: {input_path}")
|
|
|
|
# If full size, just copy
|
|
if size == "full":
|
|
if output_path:
|
|
import shutil
|
|
|
|
shutil.copy(input_path, output_path)
|
|
output_file = Path(output_path)
|
|
else:
|
|
output_file = input_file
|
|
|
|
# Get original dimensions
|
|
if HAS_PIL:
|
|
img = Image.open(str(output_file))
|
|
return (str(output_file), img.width, img.height)
|
|
return (str(output_file), 0, 0) # Dimensions unknown without PIL
|
|
|
|
# Need PIL to resize
|
|
if not HAS_PIL:
|
|
raise ValueError(
|
|
f"Size preset '{size}' requires PIL (Pillow). " "Install with: pip3 install pillow"
|
|
)
|
|
|
|
# Open original image
|
|
img = Image.open(str(input_file))
|
|
orig_w, orig_h = img.size
|
|
|
|
# Calculate new size
|
|
scale_x, scale_y = get_size_preset(size)
|
|
new_w = int(orig_w * scale_x)
|
|
new_h = int(orig_h * scale_y)
|
|
|
|
# Resize with high-quality resampling
|
|
resized = img.resize((new_w, new_h), Image.Resampling.LANCZOS)
|
|
|
|
# Determine output path
|
|
if output_path is None:
|
|
# Insert size marker before extension
|
|
stem = input_file.stem
|
|
suffix = input_file.suffix
|
|
output_path = str(input_file.parent / f"{stem}_{size}{suffix}")
|
|
|
|
# Save resized image
|
|
resized.save(output_path, quality=quality, optimize=True)
|
|
|
|
return (output_path, new_w, new_h)
|
|
|
|
|
|
def capture_screenshot(
|
|
udid: str,
|
|
output_path: str | None = None,
|
|
size: str = "half",
|
|
inline: bool = False,
|
|
app_name: str | None = None,
|
|
screen_name: str | None = None,
|
|
state: str | None = None,
|
|
) -> dict[str, Any]:
|
|
"""Capture screenshot with flexible output modes.
|
|
|
|
Supports both file-based (persistent artifacts) and inline base64 modes
|
|
(for vision-based automation).
|
|
|
|
Args:
|
|
udid: Device UDID
|
|
output_path: File path for file mode (generates semantic name if None)
|
|
size: 'full', 'half', 'quarter', 'thumb' (default: 'half')
|
|
inline: If True, returns base64 data instead of saving to file
|
|
app_name: App name for semantic naming
|
|
screen_name: Screen name for semantic naming
|
|
state: State description for semantic naming
|
|
|
|
Returns:
|
|
Dict with mode-specific fields:
|
|
|
|
File mode:
|
|
{
|
|
'mode': 'file',
|
|
'file_path': str,
|
|
'size_bytes': int,
|
|
'width': int,
|
|
'height': int,
|
|
'size_preset': str
|
|
}
|
|
|
|
Inline mode:
|
|
{
|
|
'mode': 'inline',
|
|
'base64_data': str,
|
|
'mime_type': 'image/png',
|
|
'width': int,
|
|
'height': int,
|
|
'size_preset': str
|
|
}
|
|
|
|
Example:
|
|
# File mode
|
|
result = capture_screenshot('ABC123', app_name='MyApp')
|
|
print(f"Saved to: {result['file_path']}")
|
|
|
|
# Inline mode
|
|
result = capture_screenshot('ABC123', inline=True, size='half')
|
|
print(f"Screenshot: {result['width']}x{result['height']}")
|
|
print(f"Base64: {result['base64_data'][:50]}...")
|
|
"""
|
|
try:
|
|
# Capture raw screenshot to temp file
|
|
temp_path = "/tmp/ios_simulator_screenshot.png"
|
|
cmd = ["xcrun", "simctl", "io", udid, "screenshot", temp_path]
|
|
|
|
subprocess.run(cmd, capture_output=True, text=True, check=True)
|
|
|
|
if inline:
|
|
# Inline mode: resize and convert to base64
|
|
# Resize if needed
|
|
if size != "full" and HAS_PIL:
|
|
resized_path, width, height = resize_screenshot(temp_path, size=size)
|
|
else:
|
|
resized_path = temp_path
|
|
# Get dimensions via PIL if available
|
|
if HAS_PIL:
|
|
img = Image.open(resized_path)
|
|
width, height = img.size
|
|
else:
|
|
width, height = 390, 844 # Fallback to common device size
|
|
|
|
# Read and encode as base64
|
|
with open(resized_path, "rb") as f:
|
|
base64_data = base64.b64encode(f.read()).decode("utf-8")
|
|
|
|
# Clean up temp files
|
|
Path(temp_path).unlink(missing_ok=True)
|
|
if resized_path != temp_path:
|
|
Path(resized_path).unlink(missing_ok=True)
|
|
|
|
return {
|
|
"mode": "inline",
|
|
"base64_data": base64_data,
|
|
"mime_type": "image/png",
|
|
"width": width,
|
|
"height": height,
|
|
"size_preset": size,
|
|
}
|
|
|
|
# File mode: save to output path with semantic naming
|
|
if output_path is None:
|
|
output_path = generate_screenshot_name(app_name, screen_name, state)
|
|
|
|
# Resize if needed
|
|
if size != "full" and HAS_PIL:
|
|
final_path, width, height = resize_screenshot(temp_path, output_path, size)
|
|
else:
|
|
# Just move temp to output
|
|
import shutil
|
|
|
|
shutil.move(temp_path, output_path)
|
|
final_path = output_path
|
|
|
|
# Get dimensions via PIL if available
|
|
if HAS_PIL:
|
|
img = Image.open(final_path)
|
|
width, height = img.size
|
|
else:
|
|
width, height = 390, 844 # Fallback
|
|
|
|
# Get file size
|
|
size_bytes = Path(final_path).stat().st_size
|
|
|
|
return {
|
|
"mode": "file",
|
|
"file_path": final_path,
|
|
"size_bytes": size_bytes,
|
|
"width": width,
|
|
"height": height,
|
|
"size_preset": size,
|
|
}
|
|
|
|
except subprocess.CalledProcessError as e:
|
|
raise RuntimeError(f"Failed to capture screenshot: {e.stderr}") from e
|
|
except Exception as e:
|
|
raise RuntimeError(f"Screenshot capture error: {e!s}") from e
|
|
|
|
|
|
def format_screenshot_result(result: dict[str, Any]) -> str:
|
|
"""Format screenshot result for human-readable output.
|
|
|
|
Args:
|
|
result: Result dictionary from capture_screenshot()
|
|
|
|
Returns:
|
|
Formatted string for printing
|
|
|
|
Example:
|
|
result = capture_screenshot('ABC123', inline=True)
|
|
print(format_screenshot_result(result))
|
|
"""
|
|
if result["mode"] == "file":
|
|
return (
|
|
f"Screenshot: {result['file_path']}\n"
|
|
f"Dimensions: {result['width']}x{result['height']}\n"
|
|
f"Size: {result['size_bytes']} bytes"
|
|
)
|
|
return (
|
|
f"Screenshot (inline): {result['width']}x{result['height']}\n"
|
|
f"Base64 length: {len(result['base64_data'])} chars"
|
|
)
|