Compare commits
20
Commits
de6771ac86
...
74e887db04
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
74e887db04 | ||
|
|
0d97094092 | ||
|
|
6eb0ab7d5a | ||
|
|
d572ca199f | ||
|
|
1939ef0c0d | ||
|
|
e12feae399 | ||
|
|
48fe3f025e | ||
|
|
cc5cb24b59 | ||
|
|
6a1ce2ad18 | ||
|
|
f4764699e8 | ||
|
|
9a3b78105e | ||
|
|
9e2e098928 | ||
|
|
aadb47a955 | ||
|
|
9e2659d3f1 | ||
|
|
b913274f87 | ||
|
|
f3526b60e5 | ||
|
|
4e61ace0df | ||
|
|
eb6b73e392 | ||
|
|
4bc948d77d | ||
|
|
2bbf6b36d4 |
No files matched your search
@@ -0,0 +1,35 @@
|
||||
- Every script is self contained.
|
||||
|
||||
- Do not be chatty, do not summarize the code afterwards.
|
||||
- Use `#!/usr/bin/env -S uv run --script` as the shebang line.
|
||||
- Add dependencies under the shebang line as
|
||||
|
||||
```
|
||||
# /// script
|
||||
# dependencies = [list of external dependencies as json array]
|
||||
# ///
|
||||
```
|
||||
|
||||
and keep them updated as the code changes.
|
||||
|
||||
- Keep comments to minimum
|
||||
- Use `uv` for package management.
|
||||
- Use logging instead of print statements. Use `logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s: %(message)s")` for configuration.
|
||||
- Use argparse for command-line arguments, keep argument parsing in a func called `parse_args() -> argparse.Namespace`, which gets called from `main`. Use `argparse.ArgumentDefaultsHelpFormatter` for the formatter class.
|
||||
- Use `pathlib.Path` for file operations. When the CLI args include a path, set its `type=Path`
|
||||
- Place imports at the top of the file.
|
||||
- Use httpx for HTTP requests. Do not use async.
|
||||
- Always add type annotations to function signatures, prefer built-in types, e.g. `list[str]` over `typing` module types when possible.
|
||||
- Prefer kwargs calling style for function calls, e.g. `func(arg1=value1, arg2=value2)` instead of `func(value1, value2)`.
|
||||
- Use `Path.cwd()` for the current working directory.
|
||||
- Prefer `os.getenv` over `os.environ.get` for environment variables.
|
||||
- Use `subprocess.run` for running shell commands, prefer `check=True` to raise an error on failure.
|
||||
- Use `tempfile.TemporaryDirectory()` for temporary directories.
|
||||
- Do not shorten parameter names unnecessarily, unless it's an idiomatic abbreviation.
|
||||
- Use f-strings for string formatting.
|
||||
- When calling an external command, use the long form of the arguments, e.g. `--output` instead of `-o`.
|
||||
|
||||
## Testing
|
||||
|
||||
- Use pytest for tests and name test files `<script_name>_test.py`.
|
||||
- Make test scripts self contained like `docker_refresh_test.py`: use the same uv shebang, declare pytest and all required external dependencies in inline script metadata, and run `pytest.main(args=[str(Path(__file__).resolve()), *sys.argv[1:]])` via `sys.exit` under the `__main__` guard so pytest CLI arguments are forwarded.
|
||||
+20
-4
@@ -261,9 +261,23 @@
|
||||
}
|
||||
#model-selector {
|
||||
left: 1rem;
|
||||
padding: 0;
|
||||
}
|
||||
#model-selector > summary {
|
||||
padding: 0.5rem;
|
||||
}
|
||||
#model-selector > .model-opt__label {
|
||||
margin: 0 0.5rem;
|
||||
}
|
||||
#model-selector > .model-opt__label:last-child {
|
||||
margin-bottom: 0.5rem;
|
||||
}
|
||||
#history-panel {
|
||||
right: 1rem;
|
||||
padding: 0;
|
||||
}
|
||||
#history-panel > summary {
|
||||
padding: 0.5rem;
|
||||
}
|
||||
#model-selector:is(:hover, [open]),
|
||||
#history-panel:is(:hover, [open]) {
|
||||
@@ -272,6 +286,7 @@
|
||||
.history-list {
|
||||
min-width: 220px;
|
||||
max-width: 320px;
|
||||
margin: 0 0.5rem 0.5rem;
|
||||
padding-top: 0.5rem;
|
||||
background: var(--bg-primary);
|
||||
border: 1px solid var(--border-color);
|
||||
@@ -411,7 +426,7 @@
|
||||
<button class="submit" :class="isBusy && 'busy'" @click="isBusy ? abort() : handleSubmit()" x-text="isBusy ? 'Cancel' : 'Send'"></button>
|
||||
|
||||
<!-- Model Selector -->
|
||||
<details id="model-selector" x-show="!attachment">
|
||||
<details id="model-selector" x-show="!attachment" @click.outside="$el.open = false" @keydown.escape.window="$el.open = false">
|
||||
<summary x-text="'Model: ' + activeModelName"></summary>
|
||||
<template x-for="m in models" :key="m.name">
|
||||
<label class="model-opt__label" style="display: block">
|
||||
@@ -422,7 +437,7 @@
|
||||
</details>
|
||||
|
||||
<!-- History -->
|
||||
<details id="history-panel" x-ref="historyPanel" x-show="!attachment">
|
||||
<details id="history-panel" x-ref="historyPanel" x-show="!attachment" @click.outside="$el.open = false" @keydown.escape.window="$el.open = false">
|
||||
<summary>History</summary>
|
||||
<div class="history-list">
|
||||
<template x-if="historyItems.length === 0">
|
||||
@@ -511,7 +526,7 @@
|
||||
pendingImages: [],
|
||||
attachment: initialAttachment,
|
||||
promptText: initialPromptText,
|
||||
activeModelName: "gemini-3.5-flash-lite",
|
||||
activeModelName: "gemini-3.8-flash",
|
||||
historyItems: [],
|
||||
isBusy: false,
|
||||
abortController: null,
|
||||
@@ -519,8 +534,9 @@
|
||||
sessionId: null,
|
||||
|
||||
models: [
|
||||
{ name: "gpt-5.6-luna", chat: chatWithOpenAI, key: window.env.OPENAI_API_KEY },
|
||||
{ name: "gpt-6-luna", chat: chatWithOpenAI, key: window.env.OPENAI_API_KEY },
|
||||
{ name: "gemini-3.5-flash-lite", chat: chatWithGemini, key: window.env.GEMINI_API_KEY },
|
||||
{ name: "gemini-3.8-flash", chat: chatWithGemini, key: window.env.GEMINI_API_KEY },
|
||||
{ name: "claude-haiku-4-5-20251001", chat: chatWithClaude, key: window.env.ANTHROPIC_API_KEY },
|
||||
],
|
||||
|
||||
|
||||
Executable
+282
@@ -0,0 +1,282 @@
|
||||
#!/usr/bin/env -S uv run --script
|
||||
# /// script
|
||||
# requires-python = ">=3.9"
|
||||
# dependencies = [
|
||||
# "bottle",
|
||||
# "httpx",
|
||||
# "pycryptodomex",
|
||||
# ]
|
||||
# ///
|
||||
"""Small local Deezer search and download server.
|
||||
|
||||
Examples:
|
||||
./deezer_cli.py --arl "$DEEZER_ARL"
|
||||
DEEZER_ARL=... ./deezer_cli.py --port 8080
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import hashlib
|
||||
import html
|
||||
import html.parser
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import threading
|
||||
import webbrowser
|
||||
from binascii import a2b_hex, b2a_hex
|
||||
|
||||
from bottle import Bottle, request, response
|
||||
import httpx
|
||||
from Cryptodome.Cipher import AES, Blowfish
|
||||
|
||||
|
||||
API_URL = "https://api.deezer.com"
|
||||
WEB_URL = "https://www.deezer.com"
|
||||
MEDIA_URL = "https://media.deezer.com/v1/get_url"
|
||||
HEADERS = {
|
||||
"Origin": WEB_URL,
|
||||
"Accept-Encoding": "utf-8",
|
||||
"Referer": f"{WEB_URL}/login",
|
||||
}
|
||||
|
||||
|
||||
class ScriptExtractor(html.parser.HTMLParser):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.scripts = []
|
||||
self.in_script = False
|
||||
|
||||
def handle_starttag(self, tag, attrs):
|
||||
self.in_script = tag.lower() == "script"
|
||||
|
||||
def handle_data(self, data):
|
||||
if self.in_script:
|
||||
self.scripts.append(data)
|
||||
|
||||
def handle_endtag(self, tag):
|
||||
if tag.lower() == "script":
|
||||
self.in_script = False
|
||||
|
||||
|
||||
def api_get(client, path, **params):
|
||||
response = client.get(f"{API_URL}/{path.lstrip('/')}", params=params)
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
if "error" in data:
|
||||
raise RuntimeError(data["error"].get("message", "Deezer API request failed."))
|
||||
return data
|
||||
|
||||
|
||||
def search_tracks(client, query):
|
||||
return api_get(client, "search/track", q=query, limit=25).get("data", [])
|
||||
|
||||
|
||||
def track_metadata(client, track_id):
|
||||
data = api_get(client, f"track/{track_id}")
|
||||
if not data.get("readable"):
|
||||
raise RuntimeError(f"Track {track_id} is unavailable or not playable.")
|
||||
return data
|
||||
|
||||
|
||||
def website_song_data(client, track_id):
|
||||
response = client.get(f"{WEB_URL}/us/track/{track_id}")
|
||||
response.raise_for_status()
|
||||
parser = ScriptExtractor()
|
||||
parser.feed(response.text)
|
||||
for script in parser.scripts:
|
||||
match = re.search(r'{"DATA":.*', script)
|
||||
if match:
|
||||
state = json.loads(match.group(0))
|
||||
return state["DATA"]
|
||||
raise RuntimeError("Could not find Deezer track data in the web page.")
|
||||
|
||||
|
||||
def md5hex(data):
|
||||
return b2a_hex(hashlib.md5(data).digest())
|
||||
|
||||
|
||||
def encrypted_url(song_id, md5_origin, media_version, quality):
|
||||
joined = b"\xa4".join([
|
||||
md5_origin.encode(),
|
||||
str(quality).encode(),
|
||||
str(song_id).encode(),
|
||||
str(media_version).encode(),
|
||||
])
|
||||
data = b"\xa4".join([md5hex(joined), joined]) + b"\xa4"
|
||||
data += b"\0" * (-len(data) % 16)
|
||||
cipher = AES.new(b"jo6aey6haid2Teih", AES.MODE_ECB)
|
||||
return b2a_hex(cipher.encrypt(data)).decode()
|
||||
|
||||
|
||||
def decrypt_stream(data, song_id):
|
||||
song_id_md5 = md5hex(str(song_id).encode())
|
||||
key = bytes(song_id_md5[i] ^ song_id_md5[i + 16] ^ b"g4el58wc0zvf9na1"[i]
|
||||
for i in range(16))
|
||||
cipher = Blowfish.new(key, Blowfish.MODE_CBC, iv=a2b_hex("0001020304050607"))
|
||||
output = bytearray()
|
||||
for block_number, start in enumerate(range(0, len(data), 2048)):
|
||||
block = data[start:start + 2048]
|
||||
if block_number % 3 == 0 and len(block) == 2048:
|
||||
block = cipher.decrypt(block)
|
||||
output.extend(block)
|
||||
return bytes(output)
|
||||
|
||||
|
||||
def safe_filename(value):
|
||||
value = str(value or "")
|
||||
value = re.sub(r'[<>:"/\\|?*\x00-\x1f]', "-", value)
|
||||
return value.rstrip(" .") or "unknown"
|
||||
|
||||
|
||||
def audio_data(client, token, track_id):
|
||||
metadata = track_metadata(client, track_id)
|
||||
song = website_song_data(client, track_id)
|
||||
|
||||
choices = (
|
||||
("MP3_320", 3, "FILESIZE_MP3_320"),
|
||||
("MP3_256", 5, "FILESIZE_MP3_256"),
|
||||
("MP3_128", 1, "FILESIZE_MP3_128"),
|
||||
)
|
||||
song_format, quality, _ = choices[-1]
|
||||
for choice in choices:
|
||||
if int(song.get(choice[2]) or 0) > 0:
|
||||
song_format, quality, _ = choice
|
||||
break
|
||||
|
||||
response = client.post(
|
||||
MEDIA_URL,
|
||||
headers=HEADERS,
|
||||
json={
|
||||
"license_token": token,
|
||||
"media": [{"type": "FULL", "formats": [{"cipher": "BF_CBC_STRIPE", "format": song_format}]}],
|
||||
"track_tokens": [song["TRACK_TOKEN"]],
|
||||
},
|
||||
)
|
||||
response.raise_for_status()
|
||||
media_items = response.json().get("data") or []
|
||||
media = (media_items[0].get("media") or []) if media_items else []
|
||||
sources = (media[0].get("sources") or []) if media else []
|
||||
if sources:
|
||||
media_url = sources[0]["url"]
|
||||
else:
|
||||
media_url = (
|
||||
f"https://e-cdns-proxy-{song['MD5_ORIGIN'][0]}.dzcdn.net/mobile/1/"
|
||||
f"{encrypted_url(song['SNG_ID'], song['MD5_ORIGIN'], song['MEDIA_VERSION'], 1)}"
|
||||
)
|
||||
|
||||
media_response = client.get(media_url)
|
||||
media_response.raise_for_status()
|
||||
filename = f"{safe_filename(metadata.get('artist', {}).get('name'))} - {safe_filename(metadata.get('title'))}.mp3"
|
||||
return filename, decrypt_stream(media_response.content, song["SNG_ID"])
|
||||
|
||||
|
||||
def render_page(query="", results=(), error=""):
|
||||
rows = "".join(
|
||||
"<tr>"
|
||||
f"<td class=art><img src=\"{html.escape(result.get('album', {}).get('cover_medium', ''), quote=True)}\" alt=\"\" loading=lazy></td>"
|
||||
f"<td>{html.escape(result.get('title', ''))}</td>"
|
||||
f"<td>{html.escape(result.get('artist', {}).get('name', ''))}</td>"
|
||||
f"<td>{html.escape(result.get('album', {}).get('title', ''))}</td>"
|
||||
f"<td><a class=button href=/download/{result['id']}>Download</a></td>"
|
||||
"</tr>"
|
||||
for result in results
|
||||
)
|
||||
table = (
|
||||
"<table><thead><tr><th></th><th>Title</th><th>Artist</th><th>Album</th><th></th></tr></thead>"
|
||||
f"<tbody>{rows}</tbody></table>"
|
||||
if results else ("<p class=muted>No results.</p>" if query and not error else "")
|
||||
)
|
||||
error_box = f"<p class=error>{html.escape(error)}</p>" if error else ""
|
||||
return f"""<!doctype html>
|
||||
<html lang=en>
|
||||
<head>
|
||||
<meta charset=utf-8>
|
||||
<meta name=viewport content="width=device-width,initial-scale=1">
|
||||
<title>Deezer Search</title>
|
||||
<style>
|
||||
body {{ max-width: 900px; margin: 40px auto; padding: 0 16px; font: 15px system-ui, sans-serif; color: #222; }}
|
||||
form {{ display: flex; gap: 8px; margin-bottom: 24px; }}
|
||||
input {{ flex: 1; padding: 9px; font: inherit; }}
|
||||
button, .button {{ border: 0; border-radius: 4px; padding: 9px 13px; background: #a238ff; color: white; text-decoration: none; cursor: pointer; }}
|
||||
table {{ width: 100%; border-collapse: collapse; }}
|
||||
th, td {{ border-bottom: 1px solid #ddd; padding: 14px 10px; text-align: left; }}
|
||||
th:last-child, td:last-child {{ text-align: right; }}
|
||||
.art {{ width: 56px; padding-right: 0; }}
|
||||
.art img {{ display: block; width: 56px; height: 56px; object-fit: cover; border-radius: 4px; background: #eee; }}
|
||||
.error {{ color: #b00020; }}
|
||||
.muted {{ color: #777; }}
|
||||
</style>
|
||||
</head>
|
||||
<body>
|
||||
<h1>Deezer Search</h1>
|
||||
<form method=get>
|
||||
<input name=q value="{html.escape(query, quote=True)}" placeholder="Artist or track" autofocus>
|
||||
<button type=submit>Search</button>
|
||||
</form>
|
||||
{error_box}
|
||||
{table}
|
||||
</body>
|
||||
</html>"""
|
||||
|
||||
|
||||
def create_app(client, token):
|
||||
app = Bottle()
|
||||
|
||||
@app.get("/")
|
||||
def index():
|
||||
query = request.query.q.strip()
|
||||
if not query:
|
||||
return render_page()
|
||||
try:
|
||||
return render_page(query, search_tracks(client, query))
|
||||
except (httpx.HTTPError, KeyError, ValueError, RuntimeError) as exc:
|
||||
return render_page(query, error=str(exc))
|
||||
|
||||
@app.get("/download/<track_id:re:[0-9]+>")
|
||||
def download_route(track_id):
|
||||
try:
|
||||
filename, data = audio_data(client, token, track_id)
|
||||
except (httpx.HTTPError, KeyError, IndexError, ValueError, RuntimeError) as exc:
|
||||
response.status = 502
|
||||
return f"Download failed: {html.escape(str(exc))}"
|
||||
response.content_type = "audio/mpeg"
|
||||
response.headers["Content-Disposition"] = f'attachment; filename="{filename}"'
|
||||
response.headers["Content-Length"] = str(len(data))
|
||||
return data
|
||||
|
||||
return app
|
||||
|
||||
|
||||
def main(argv=None):
|
||||
parser = argparse.ArgumentParser(description=__doc__.split("\n", 1)[0])
|
||||
parser.add_argument("--arl", default=os.environ.get("DEEZER_ARL"),
|
||||
help="Deezer ARL cookie (defaults to DEEZER_ARL).")
|
||||
parser.add_argument("--host", default="127.0.0.1", help="Bind address (default: 127.0.0.1).")
|
||||
parser.add_argument("--port", default=8080, type=int, help="Bind port (default: 8080).")
|
||||
args = parser.parse_args(argv)
|
||||
if not args.arl:
|
||||
parser.error("--arl is required unless DEEZER_ARL is set.")
|
||||
|
||||
try:
|
||||
client = httpx.Client(headers=HEADERS, cookies={"arl": args.arl, "comeback": "1"},
|
||||
follow_redirects=True, timeout=30)
|
||||
user_data = client.post(
|
||||
f"{WEB_URL}/ajax/gw-light.php",
|
||||
params={"api_version": "1.0", "api_token": "null", "input": "3",
|
||||
"method": "deezer.getUserData"},
|
||||
)
|
||||
user_data.raise_for_status()
|
||||
token = user_data.json()["results"]["USER"]["OPTIONS"]["license_token"]
|
||||
except (httpx.HTTPError, KeyError, IndexError, ValueError, RuntimeError, json.JSONDecodeError) as exc:
|
||||
parser.error(str(exc))
|
||||
|
||||
app = create_app(client, token)
|
||||
browser_host = "127.0.0.1" if args.host in ("0.0.0.0", "::") else args.host
|
||||
browser_url = f"http://{browser_host}:{args.port}/"
|
||||
print(f"Serving {browser_url}")
|
||||
threading.Timer(0.5, webbrowser.open, args=(browser_url,)).start()
|
||||
app.run(host=args.host, port=args.port)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Executable
+204
@@ -0,0 +1,204 @@
|
||||
#!/usr/bin/env -S uv run --script
|
||||
# /// script
|
||||
# dependencies = []
|
||||
# ///
|
||||
|
||||
import argparse
|
||||
from dataclasses import dataclass, field
|
||||
import json
|
||||
import logging
|
||||
from pathlib import Path
|
||||
import re
|
||||
import shlex
|
||||
import subprocess
|
||||
import sys
|
||||
from typing import Any
|
||||
|
||||
|
||||
@dataclass
|
||||
class Project:
|
||||
name: str
|
||||
directory: Path
|
||||
files: tuple[Path, ...]
|
||||
environment_files: tuple[Path, ...]
|
||||
services: set[str] = field(default_factory=set)
|
||||
running_images: dict[str, set[str]] = field(default_factory=dict)
|
||||
|
||||
def command(self) -> list[str]:
|
||||
command = [
|
||||
"docker", "compose", "--project-name", self.name,
|
||||
"--project-directory", str(self.directory), "--profile", "*",
|
||||
]
|
||||
for path in self.files:
|
||||
command.extend(["--file", str(path)])
|
||||
for path in self.environment_files:
|
||||
command.extend(["--env-file", str(path)])
|
||||
return command
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Pull and redeploy running Compose services with unpinned images.",
|
||||
formatter_class=argparse.ArgumentDefaultsHelpFormatter,
|
||||
epilog="Digests and tags starting with a version number (e.g. v1.2.3 or 16-alpine) "
|
||||
"are skipped. Run on the Compose host with the original files and environment available.",
|
||||
)
|
||||
parser.add_argument("projects", nargs="*", help="Only refresh these Compose projects; omit to refresh all eligible projects")
|
||||
parser.add_argument("--dry-run", action="store_true", help="Discover services and log commands without changing deployments")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def is_pinned(image: str) -> bool:
|
||||
if "@" in image:
|
||||
return True
|
||||
final_component = image.rsplit("/", 1)[-1]
|
||||
if ":" not in final_component:
|
||||
return False
|
||||
tag = final_component.rsplit(":", 1)[-1]
|
||||
return re.match(r"^v?\d+(?:[._-]|$)", tag) is not None
|
||||
|
||||
|
||||
def inspect_containers() -> list[dict[str, Any]]:
|
||||
result = subprocess.run(
|
||||
args=["docker", "container", "ls", "--quiet"],
|
||||
check=True, capture_output=True, text=True, timeout=60,
|
||||
)
|
||||
identifiers = result.stdout.split()
|
||||
containers = []
|
||||
for offset in range(0, len(identifiers), 100):
|
||||
result = subprocess.run(
|
||||
args=["docker", "container", "inspect", *identifiers[offset:offset + 100]],
|
||||
check=True, capture_output=True, text=True, timeout=60,
|
||||
)
|
||||
value = json.loads(result.stdout)
|
||||
if not isinstance(value, list):
|
||||
raise ValueError("Docker returned an unexpected response")
|
||||
containers.extend(value)
|
||||
return containers
|
||||
|
||||
|
||||
def discover_projects(containers: list[dict[str, Any]], selected_projects: set[str]) -> list[Project]:
|
||||
projects: dict[str, Project] = {}
|
||||
for container in containers:
|
||||
if not container.get("State", {}).get("Running"):
|
||||
continue
|
||||
config = container.get("Config") or {}
|
||||
labels = config.get("Labels") or {}
|
||||
name = labels.get("com.docker.compose.project")
|
||||
service = labels.get("com.docker.compose.service")
|
||||
if selected_projects and name not in selected_projects:
|
||||
continue
|
||||
if labels.get("com.docker.compose.oneoff", "").lower() == "true":
|
||||
continue
|
||||
image = config.get("Image") or ""
|
||||
if not image:
|
||||
logging.info("Skipping %s/%s: container image reference is missing", name, service)
|
||||
continue
|
||||
if is_pinned(image=image):
|
||||
logging.warning("Skipping %s/%s: deployed image %s is pinned", name, service, image)
|
||||
continue
|
||||
if not name or not service:
|
||||
logging.warning("Skipping %s: unpinned image %s is not a Compose service", container.get("Name"), image)
|
||||
continue
|
||||
directory = labels.get("com.docker.compose.project.working_dir")
|
||||
files = labels.get("com.docker.compose.project.config_files")
|
||||
if not directory or not files:
|
||||
raise ValueError(f"Project {name}: missing Compose working directory or config file labels")
|
||||
working_directory = Path(directory)
|
||||
config_files = tuple(working_directory / path for path in files.split(",") if path)
|
||||
environment_files = tuple(
|
||||
working_directory / path
|
||||
for path in labels.get("com.docker.compose.project.environment_file", "").split(",") if path
|
||||
)
|
||||
if not config_files:
|
||||
raise ValueError(f"Project {name}: no Compose config files")
|
||||
project = Project(name=name, directory=working_directory, files=config_files, environment_files=environment_files)
|
||||
existing = projects.setdefault(name, project)
|
||||
if (existing.directory, existing.files, existing.environment_files) != (project.directory, project.files, project.environment_files):
|
||||
raise ValueError(f"Project {name}: containers have inconsistent Compose file labels")
|
||||
existing.services.add(service)
|
||||
existing.running_images.setdefault(service, set()).add(container.get("Image") or "")
|
||||
return sorted(projects.values(), key=lambda project: project.name)
|
||||
|
||||
|
||||
def refresh_project(project: Project, dry_run: bool) -> None:
|
||||
if not project.directory.is_dir():
|
||||
raise ValueError(f"Project {project.name}: directory is unavailable: {project.directory}; run on the Compose host")
|
||||
for path in (*project.files, *project.environment_files):
|
||||
if not path.is_file():
|
||||
raise ValueError(f"Project {project.name}: file is unavailable: {path}")
|
||||
command = project.command()
|
||||
result = subprocess.run(
|
||||
args=[*command, "config", "--format", "json"], cwd=project.directory,
|
||||
check=True, capture_output=True, text=True, timeout=60,
|
||||
)
|
||||
configuration = json.loads(result.stdout)
|
||||
configured_services = configuration.get("services", {})
|
||||
if any(service.get("build") is not None for service in configured_services.values()):
|
||||
logging.info("Skipping project %s: contains locally built services", project.name)
|
||||
return
|
||||
services = []
|
||||
for name in sorted(project.services):
|
||||
service = configured_services.get(name)
|
||||
if service is None:
|
||||
raise ValueError(f"Project {project.name}: service {name} is absent from current Compose config")
|
||||
image = service.get("image") or ""
|
||||
if not image or is_pinned(image=image):
|
||||
logging.warning("Skipping %s/%s: current Compose image %s is missing or pinned", project.name, name, image)
|
||||
else:
|
||||
logging.info("Refreshing %s/%s (%s)", project.name, name, image)
|
||||
services.append(name)
|
||||
if not services:
|
||||
return
|
||||
pull_command = [*command, "pull", "--policy", "always", *services]
|
||||
logging.info("%s", shlex.join(pull_command))
|
||||
if not dry_run:
|
||||
subprocess.run(args=pull_command, cwd=project.directory, check=True)
|
||||
changed_services = []
|
||||
for name in services:
|
||||
result = subprocess.run(
|
||||
args=["docker", "image", "inspect", "--format", "{{.Id}}", configured_services[name]["image"]],
|
||||
check=True, capture_output=True, text=True, timeout=60,
|
||||
)
|
||||
image_id = result.stdout.strip()
|
||||
if not image_id:
|
||||
raise ValueError(f"Project {project.name}: image ID is missing for service {name}")
|
||||
if project.running_images.get(name) == {image_id}:
|
||||
logging.info("Skipping %s/%s: running image is up to date", project.name, name)
|
||||
else:
|
||||
changed_services.append(name)
|
||||
services = changed_services
|
||||
if not services:
|
||||
return
|
||||
up_command = [*command, "up", "--detach", "--no-deps", "--no-build", "--pull", "never", *services]
|
||||
logging.info("%s", shlex.join(up_command))
|
||||
if not dry_run:
|
||||
subprocess.run(args=up_command, cwd=project.directory, check=True)
|
||||
|
||||
|
||||
def main() -> int:
|
||||
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s: %(message)s")
|
||||
arguments = parse_args()
|
||||
try:
|
||||
selected_projects = set(arguments.projects)
|
||||
projects = discover_projects(containers=inspect_containers(), selected_projects=selected_projects)
|
||||
if not projects:
|
||||
if selected_projects:
|
||||
logging.info("No eligible services to refresh in selected projects: %s", ", ".join(sorted(selected_projects)))
|
||||
else:
|
||||
logging.info("No eligible services to refresh")
|
||||
failed = False
|
||||
for project in projects:
|
||||
try:
|
||||
refresh_project(project=project, dry_run=arguments.dry_run)
|
||||
except (OSError, ValueError, subprocess.SubprocessError) as error:
|
||||
logging.error("Project %s: %s", project.name, getattr(error, "stderr", None) or error)
|
||||
failed = True
|
||||
return int(failed)
|
||||
except (OSError, ValueError, subprocess.SubprocessError) as error:
|
||||
logging.error("%s", getattr(error, "stderr", None) or error)
|
||||
return 1
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
Executable
+189
@@ -0,0 +1,189 @@
|
||||
#!/usr/bin/env -S uv run --script
|
||||
# /// script
|
||||
# dependencies = ["pytest"]
|
||||
# ///
|
||||
|
||||
import json
|
||||
import logging
|
||||
from pathlib import Path
|
||||
import subprocess
|
||||
import sys
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
import docker_refresh as script
|
||||
|
||||
|
||||
@pytest.mark.parametrize("image, expected", [
|
||||
("nginx", False), ("nginx:latest", False), ("app:stable", False),
|
||||
("app:main", False), ("app:nightly", False), ("app:alpine", False),
|
||||
("localhost:5000/team/app", False),
|
||||
("localhost:5000/team/app:latest", False),
|
||||
("nginx:1.27.3", True), ("postgres:16-alpine", True),
|
||||
("app:v2.0.1", True), ("app:2026-10-03", True),
|
||||
("app@sha256:abc", True), ("app:latest@sha256:abc", True),
|
||||
])
|
||||
def test_image_pinning(image: str, expected: bool) -> None:
|
||||
assert script.is_pinned(image=image) == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize("selected, expected", [
|
||||
(set(), ["other", "web"]),
|
||||
({"web"}, ["web"]),
|
||||
({"web", "other"}, ["other", "web"]),
|
||||
({"other", "missing"}, ["other"]),
|
||||
({"missing"}, []),
|
||||
])
|
||||
def test_discovery_filters_and_groups_replicas(selected: set[str], expected: list[str]) -> None:
|
||||
containers = []
|
||||
for image, running, oneoff, project in [
|
||||
("app:latest", True, "False", "web"),
|
||||
("app:latest", True, "False", "web"),
|
||||
("db:16", True, "False", "web"),
|
||||
("stopped:latest", False, "False", "web"),
|
||||
("oneoff:latest", True, "True", "web"),
|
||||
("app:latest", True, "False", "other"),
|
||||
]:
|
||||
containers.append({
|
||||
"Image": "sha256:running",
|
||||
"State": {"Running": running},
|
||||
"Config": {"Image": image, "Labels": {
|
||||
"com.docker.compose.project": project,
|
||||
"com.docker.compose.service": image.split(":")[0],
|
||||
"com.docker.compose.oneoff": oneoff,
|
||||
"com.docker.compose.project.working_dir": "/srv/web",
|
||||
"com.docker.compose.project.config_files": "compose.yml,override.yml",
|
||||
"com.docker.compose.project.environment_file": "/srv/web/custom.env",
|
||||
}},
|
||||
})
|
||||
projects = script.discover_projects(containers=containers, selected_projects=selected)
|
||||
assert [project.name for project in projects] == expected
|
||||
for project in projects:
|
||||
assert project.services == {"app"}
|
||||
assert project.running_images == {"app": {"sha256:running"}}
|
||||
assert project.files == (Path("/srv/web/compose.yml"), Path("/srv/web/override.yml"))
|
||||
assert "--env-file" in project.command()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("names", [[], ["web"], ["web", "api"]])
|
||||
def test_positional_project_arguments(names: list[str]) -> None:
|
||||
with patch("sys.argv", ["docker_refresh.py", "--dry-run", *names]):
|
||||
arguments = script.parse_args()
|
||||
assert arguments.projects == names
|
||||
assert arguments.dry_run
|
||||
|
||||
|
||||
def test_selected_pinned_project_reports_skip_reason(caplog: pytest.LogCaptureFixture) -> None:
|
||||
container = {
|
||||
"State": {"Running": True},
|
||||
"Config": {
|
||||
"Image": "ghcr.io/abdusco/uno:v1.0.0",
|
||||
"Labels": {
|
||||
"com.docker.compose.project": "uno",
|
||||
"com.docker.compose.service": "uno",
|
||||
},
|
||||
},
|
||||
}
|
||||
with patch("sys.argv", ["docker_refresh.py", "uno"]):
|
||||
with patch.object(script, "inspect_containers", return_value=[container]):
|
||||
with caplog.at_level(logging.INFO):
|
||||
assert script.main() == 0
|
||||
assert caplog.record_tuples == [
|
||||
("root", logging.WARNING, "Skipping uno/uno: deployed image ghcr.io/abdusco/uno:v1.0.0 is pinned"),
|
||||
("root", logging.INFO, "No eligible services to refresh in selected projects: uno"),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dry_run", [True, False])
|
||||
def test_refresh_commands_and_current_config_filtering(tmp_path: Path, dry_run: bool) -> None:
|
||||
compose = tmp_path / "compose.yml"
|
||||
compose.touch()
|
||||
project = script.Project(name="web", directory=tmp_path, files=(compose,), environment_files=(), services={"app", "db"})
|
||||
response = subprocess.CompletedProcess(args=[], returncode=0, stdout=json.dumps({
|
||||
"services": {
|
||||
"app": {"image": "app:latest"},
|
||||
"db": {"image": "db:16"},
|
||||
},
|
||||
}))
|
||||
image_response = subprocess.CompletedProcess(args=[], returncode=0, stdout="sha256:new\n")
|
||||
with patch.object(script.subprocess, "run", side_effect=[response, response, image_response, response]) as run:
|
||||
script.refresh_project(project=project, dry_run=dry_run)
|
||||
assert run.call_count == (1 if dry_run else 4)
|
||||
if not dry_run:
|
||||
assert run.call_args_list[1].kwargs["args"] == [*project.command(), "pull", "--policy", "always", "app"]
|
||||
assert run.call_args_list[2].kwargs["args"] == ["docker", "image", "inspect", "--format", "{{.Id}}", "app:latest"]
|
||||
assert run.call_args_list[3].kwargs["args"] == [*project.command(), "up", "--detach", "--no-deps", "--no-build", "--pull", "never", "app"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("running_images, expected_services", [
|
||||
({"app": {"sha256:current"}, "worker": {"sha256:current"}}, []),
|
||||
({"app": {"sha256:old"}, "worker": {"sha256:current"}}, ["app"]),
|
||||
({"app": {"sha256:current", "sha256:old"}, "worker": {"sha256:current"}}, ["app"]),
|
||||
({"app": {"sha256:old"}, "worker": {"sha256:old"}}, ["app", "worker"]),
|
||||
])
|
||||
def test_only_changed_services_are_redeployed(
|
||||
tmp_path: Path, running_images: dict[str, set[str]], expected_services: list[str],
|
||||
) -> None:
|
||||
compose = tmp_path / "compose.yml"
|
||||
compose.touch()
|
||||
project = script.Project(
|
||||
name="web", directory=tmp_path, files=(compose,), environment_files=(),
|
||||
services={"app", "worker"}, running_images=running_images,
|
||||
)
|
||||
configuration = subprocess.CompletedProcess(args=[], returncode=0, stdout=json.dumps({
|
||||
"services": {"app": {"image": "app:latest"}, "worker": {"image": "worker:latest"}},
|
||||
}))
|
||||
image_response = subprocess.CompletedProcess(args=[], returncode=0, stdout="sha256:current\n")
|
||||
with patch.object(script.subprocess, "run", side_effect=[
|
||||
configuration, configuration, image_response, image_response, configuration,
|
||||
]) as run:
|
||||
script.refresh_project(project=project, dry_run=False)
|
||||
assert run.call_args_list[1].kwargs["args"] == [*project.command(), "pull", "--policy", "always", "app", "worker"]
|
||||
assert run.call_count == (5 if expected_services else 4)
|
||||
if expected_services:
|
||||
assert run.call_args.kwargs["args"] == [
|
||||
*project.command(), "up", "--detach", "--no-deps", "--no-build", "--pull", "never", *expected_services,
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("build", [".", {"context": "."}])
|
||||
def test_project_with_any_built_service_is_skipped(tmp_path: Path, build: str | dict[str, str]) -> None:
|
||||
compose = tmp_path / "compose.yml"
|
||||
compose.touch()
|
||||
project = script.Project(name="web", directory=tmp_path, files=(compose,), environment_files=(), services={"app"})
|
||||
response = subprocess.CompletedProcess(args=[], returncode=0, stdout=json.dumps({
|
||||
"services": {
|
||||
"app": {"image": "app:latest"},
|
||||
"local": {"build": build},
|
||||
},
|
||||
}))
|
||||
with patch.object(script.subprocess, "run", return_value=response) as run:
|
||||
script.refresh_project(project=project, dry_run=False)
|
||||
assert run.call_count == 1
|
||||
assert run.call_args.kwargs["args"] == [*project.command(), "config", "--format", "json"]
|
||||
|
||||
|
||||
def test_failed_pull_does_not_redeploy(tmp_path: Path) -> None:
|
||||
compose = tmp_path / "compose.yml"
|
||||
compose.touch()
|
||||
project = script.Project(name="web", directory=tmp_path, files=(compose,), environment_files=(), services={"app"})
|
||||
response = subprocess.CompletedProcess(args=[], returncode=0, stdout='{"services":{"app":{"image":"app:latest"}}}')
|
||||
with patch.object(script.subprocess, "run", side_effect=[
|
||||
response, subprocess.CalledProcessError(returncode=1, cmd="pull"),
|
||||
]) as run:
|
||||
with pytest.raises(subprocess.CalledProcessError):
|
||||
script.refresh_project(project=project, dry_run=False)
|
||||
assert run.call_count == 2
|
||||
|
||||
|
||||
def test_missing_files_fail_before_docker_commands(tmp_path: Path) -> None:
|
||||
project = script.Project(name="web", directory=tmp_path, files=(tmp_path / "missing.yml",), environment_files=(), services={"app"})
|
||||
with patch.object(script.subprocess, "run") as run:
|
||||
with pytest.raises(ValueError):
|
||||
script.refresh_project(project=project, dry_run=False)
|
||||
run.assert_not_called()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main(args=[str(Path(__file__).resolve()), *sys.argv[1:]]))
|
||||
Executable
+245
@@ -0,0 +1,245 @@
|
||||
#!/usr/bin/env python3
|
||||
# /// script
|
||||
# dependencies = []
|
||||
# ///
|
||||
|
||||
import argparse
|
||||
import ipaddress
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
from pathlib import Path
|
||||
import re
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
from typing import Any
|
||||
|
||||
|
||||
class HelpFormatter(argparse.ArgumentDefaultsHelpFormatter, argparse.RawDescriptionHelpFormatter):
|
||||
pass
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Generate Caddy reverse proxies from Docker container labels.",
|
||||
formatter_class=HelpFormatter,
|
||||
epilog="""Compose service labels:
|
||||
labels:
|
||||
caddy.domain: ascii.app.abdus.dev
|
||||
caddy.port: "3000"
|
||||
caddy.basic_auth: "true"
|
||||
|
||||
caddy.port selects the internal TCP port, not the published host port.
|
||||
Omit it when exactly one TCP port is published. Only running labeled
|
||||
containers generate routes. Potential Docker port clashes warn on stderr.
|
||||
caddy.basic_auth accepts true or false (default). When true, the route
|
||||
imports the (auth) snippet defined in your handwritten Caddy config.
|
||||
|
||||
Local usage:
|
||||
python3 docker_to_caddy.py
|
||||
python3 docker_to_caddy.py --output docker.generated.caddy
|
||||
|
||||
Remote Docker over SSH:
|
||||
env -u DOCKER_CONTEXT DOCKER_HOST=ssh://abdus@your-server docker ps
|
||||
env -u DOCKER_CONTEXT DOCKER_HOST=ssh://abdus@your-server \\
|
||||
python3 docker_to_caddy.py --output docker.generated.caddy
|
||||
|
||||
The SSH user needs Docker socket access on the server. These commands do
|
||||
not change your selected Docker context (for example, OrbStack). Output is
|
||||
written locally; loopback upstreams refer to the server where Caddy runs.
|
||||
|
||||
Caddy setup:
|
||||
Add a top-level import docker.generated.caddy to your handwritten config.
|
||||
Remove migrated proxy blocks, then validate and reload Caddy manually.
|
||||
The script does not configure DNS or reload Caddy.
|
||||
""",
|
||||
)
|
||||
parser.add_argument("--output", type=Path, help="Output file; otherwise write to stdout")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def docker_json(arguments: list[str]) -> list[Any]:
|
||||
result = subprocess.run(
|
||||
args=["docker", *arguments], check=True, capture_output=True, text=True,
|
||||
timeout=60,
|
||||
)
|
||||
value = json.loads(result.stdout)
|
||||
if not isinstance(value, list):
|
||||
raise ValueError("Docker returned an unexpected response")
|
||||
return value
|
||||
|
||||
|
||||
def inspect_containers() -> list[dict[str, Any]]:
|
||||
result = subprocess.run(
|
||||
args=["docker", "container", "ls", "--all", "--quiet"],
|
||||
check=True, capture_output=True, text=True, timeout=60,
|
||||
)
|
||||
identifiers = result.stdout.split()
|
||||
containers = []
|
||||
for offset in range(0, len(identifiers), 100):
|
||||
containers.extend(docker_json(arguments=[
|
||||
"container", "inspect", *identifiers[offset:offset + 100],
|
||||
]))
|
||||
return containers
|
||||
|
||||
|
||||
def container_name(container: dict[str, Any]) -> str:
|
||||
return str(container.get("Name") or container["Id"]).lstrip("/")
|
||||
|
||||
|
||||
def port_number(value: str) -> int:
|
||||
if not re.fullmatch(r"[0-9]+", value) or not 1 <= int(value) <= 65535:
|
||||
raise ValueError(f"Invalid port: {value!r}")
|
||||
return int(value)
|
||||
|
||||
|
||||
def normalize_domain(value: str) -> str:
|
||||
domain = value.lower().rstrip(".")
|
||||
labels = domain.split(".")
|
||||
if len(domain) > 253 or len(labels) < 2 or any(
|
||||
not re.fullmatch(r"[a-z0-9](?:[a-z0-9-]{0,61}[a-z0-9])?", label)
|
||||
for label in labels
|
||||
):
|
||||
raise ValueError(f"Invalid domain: {value!r}; use a single DNS hostname")
|
||||
return domain
|
||||
|
||||
|
||||
def upstream(binding: dict[str, str]) -> tuple[int, str, int]:
|
||||
address = ipaddress.ip_address(binding.get("HostIp") or "0.0.0.0")
|
||||
if address.is_unspecified:
|
||||
address = ipaddress.ip_address("127.0.0.1" if address.version == 4 else "::1")
|
||||
port = port_number(value=binding["HostPort"])
|
||||
return address.version, str(address), port
|
||||
|
||||
|
||||
def render_config(containers: list[dict[str, Any]]) -> str:
|
||||
routes: dict[str, tuple[str, str, bool]] = {}
|
||||
for container in containers:
|
||||
labels = container.get("Config", {}).get("Labels") or {}
|
||||
if "caddy.domain" not in labels or not container.get("State", {}).get("Running"):
|
||||
continue
|
||||
name = container_name(container=container)
|
||||
try:
|
||||
domain = normalize_domain(value=labels["caddy.domain"])
|
||||
auth_label = labels.get("caddy.basic_auth", "false")
|
||||
if not isinstance(auth_label, str) or auth_label.lower() not in ("true", "false"):
|
||||
raise ValueError("caddy.basic_auth must be true or false")
|
||||
basic_auth = auth_label.lower() == "true"
|
||||
if container.get("HostConfig", {}).get("NetworkMode") == "host":
|
||||
raise ValueError("Host networking is unsupported; publish a TCP port instead")
|
||||
ports = {
|
||||
key: [binding for binding in bindings or [] if binding.get("HostPort")]
|
||||
for key, bindings in (container.get("NetworkSettings", {}).get("Ports") or {}).items()
|
||||
if key.endswith("/tcp")
|
||||
}
|
||||
ports = {key: bindings for key, bindings in ports.items() if bindings}
|
||||
if "caddy.port" in labels:
|
||||
key = f"{port_number(value=labels['caddy.port'])}/tcp"
|
||||
if key not in ports:
|
||||
raise ValueError(f"Selected container port {key} has no published host port")
|
||||
elif len(ports) == 1:
|
||||
key = next(iter(ports))
|
||||
else:
|
||||
raise ValueError("Expected one published TCP port; set caddy.port to select one")
|
||||
version, address, port = min(upstream(binding=binding) for binding in ports[key])
|
||||
target = f"[{address}]:{port}" if version == 6 else f"{address}:{port}"
|
||||
if domain in routes:
|
||||
raise ValueError(f"Domain {domain} is also claimed by {routes[domain][0]}")
|
||||
routes[domain] = (name, target, basic_auth)
|
||||
except (ValueError, KeyError, TypeError) as error:
|
||||
raise ValueError(f"Container {name}: {error}") from error
|
||||
blocks = ["# Generated by docker_to_caddy.py; do not edit.\n"]
|
||||
for domain, (name, target, basic_auth) in sorted(routes.items()):
|
||||
safe_name = re.sub(r"[^a-zA-Z0-9_.-]", "_", name)
|
||||
auth_import = " import auth\n" if basic_auth else ""
|
||||
blocks.append(
|
||||
f"# Container: {safe_name}\n{domain} {{\n{auth_import} reverse_proxy {target}\n}}\n"
|
||||
)
|
||||
return "\n".join(blocks)
|
||||
|
||||
|
||||
def addresses_overlap(first: str, second: str) -> bool:
|
||||
left = ipaddress.ip_address(first or "0.0.0.0")
|
||||
right = ipaddress.ip_address(second or "0.0.0.0")
|
||||
if left.version != right.version:
|
||||
return (left.version == 6 and left.is_unspecified) or (
|
||||
right.version == 6 and right.is_unspecified
|
||||
)
|
||||
return left == right or left.is_unspecified or right.is_unspecified
|
||||
|
||||
|
||||
def warn_port_clashes(containers: list[dict[str, Any]]) -> None:
|
||||
bindings: list[tuple[str, str, str, int, str]] = []
|
||||
for container in containers:
|
||||
name = container_name(container=container)
|
||||
state = container.get("State", {}).get("Status", "unknown")
|
||||
configured = container.get("HostConfig", {}).get("PortBindings") or {}
|
||||
actual = container.get("NetworkSettings", {}).get("Ports") or {}
|
||||
seen: set[tuple[str, int, str]] = set()
|
||||
for source in (configured, actual):
|
||||
for container_port, entries in source.items():
|
||||
protocol = container_port.rsplit("/", 1)[-1]
|
||||
for binding in entries or []:
|
||||
value = binding.get("HostPort")
|
||||
if not value or value == "0":
|
||||
continue
|
||||
try:
|
||||
port = port_number(value=value)
|
||||
address = str(ipaddress.ip_address(binding.get("HostIp") or "0.0.0.0"))
|
||||
except ValueError as error:
|
||||
logging.warning("Cannot check binding for %s: %s", name, error)
|
||||
continue
|
||||
identity = (address, port, protocol)
|
||||
if identity not in seen:
|
||||
seen.add(identity)
|
||||
bindings.append((name, state, address, port, protocol))
|
||||
warned: set[tuple[str, str, int, str]] = set()
|
||||
for index, first in enumerate(bindings):
|
||||
for second in bindings[index + 1:]:
|
||||
if first[0] == second[0] or first[3:] != second[3:]:
|
||||
continue
|
||||
if not addresses_overlap(first=first[2], second=second[2]):
|
||||
continue
|
||||
identity = (*sorted((first[0], second[0])), first[3], first[4])
|
||||
if identity in warned:
|
||||
continue
|
||||
warned.add(identity)
|
||||
logging.warning(
|
||||
"Potential port clash: %s (%s) [%s]:%s/%s and %s (%s) [%s]:%s/%s",
|
||||
*first, *second,
|
||||
)
|
||||
|
||||
|
||||
def write_config(config: str, output: Path) -> None:
|
||||
output = output.absolute()
|
||||
with tempfile.TemporaryDirectory(dir=output.parent, prefix=".docker-to-caddy-") as directory:
|
||||
temporary = Path(directory) / "config"
|
||||
temporary.write_text(data=config, encoding="utf-8")
|
||||
if output.exists():
|
||||
temporary.chmod(mode=output.stat().st_mode & 0o777)
|
||||
os.replace(src=temporary, dst=output)
|
||||
|
||||
|
||||
def main() -> int:
|
||||
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s: %(message)s")
|
||||
arguments = parse_args()
|
||||
try:
|
||||
containers = inspect_containers()
|
||||
warn_port_clashes(containers=containers)
|
||||
config = render_config(containers=containers)
|
||||
if arguments.output is None:
|
||||
sys.stdout.write(config)
|
||||
else:
|
||||
write_config(config=config, output=arguments.output)
|
||||
logging.info("Wrote %s", arguments.output)
|
||||
return 0
|
||||
except subprocess.CalledProcessError as error:
|
||||
logging.error("Docker command failed: %s", (error.stderr or str(error)).strip())
|
||||
except (OSError, ValueError, subprocess.TimeoutExpired) as error:
|
||||
logging.error("%s", error)
|
||||
return 1
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
@@ -0,0 +1,168 @@
|
||||
#!/usr/bin/env -S uv run --script
|
||||
# /// script
|
||||
# dependencies = ["pytest"]
|
||||
# ///
|
||||
|
||||
import copy
|
||||
from pathlib import Path
|
||||
import subprocess
|
||||
import sys
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
import docker_to_caddy as script
|
||||
|
||||
|
||||
def container(name: str = "app", domain: str = "ascii.app.abdus.dev") -> dict:
|
||||
return {
|
||||
"Name": f"/{name}",
|
||||
"Config": {"Labels": {"caddy.domain": domain}},
|
||||
"State": {"Running": True, "Status": "running"},
|
||||
"HostConfig": {"PortBindings": {
|
||||
"3000/tcp": [{"HostIp": "127.0.0.1", "HostPort": "4510"}],
|
||||
}},
|
||||
"NetworkSettings": {"Ports": {
|
||||
"3000/tcp": [{"HostIp": "127.0.0.1", "HostPort": "4510"}],
|
||||
}},
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("address, expected", [
|
||||
("127.0.0.1", "127.0.0.1:4510"),
|
||||
("0.0.0.0", "127.0.0.1:4510"),
|
||||
("::", "[::1]:4510"),
|
||||
("::1", "[::1]:4510"),
|
||||
("192.168.1.2", "192.168.1.2:4510"),
|
||||
])
|
||||
def test_bindings(address: str, expected: str) -> None:
|
||||
app = container()
|
||||
app["NetworkSettings"]["Ports"]["3000/tcp"][0]["HostIp"] = address
|
||||
assert f"reverse_proxy {expected}" in script.render_config(containers=[app])
|
||||
|
||||
|
||||
@pytest.mark.parametrize("selector, succeeds", [(None, False), ("3000", True), ("9000", False)])
|
||||
def test_port_selection(selector: str | None, succeeds: bool) -> None:
|
||||
app = container()
|
||||
app["NetworkSettings"]["Ports"]["4000/tcp"] = [
|
||||
{"HostIp": "::1", "HostPort": "4600"},
|
||||
]
|
||||
if selector is not None:
|
||||
app["Config"]["Labels"]["caddy.port"] = selector
|
||||
if succeeds:
|
||||
assert "127.0.0.1:4510" in script.render_config(containers=[app])
|
||||
else:
|
||||
with pytest.raises(ValueError):
|
||||
script.render_config(containers=[app])
|
||||
|
||||
|
||||
@pytest.mark.parametrize("domain", ["", "bad host.dev", "example.dev {", "*.abdus.dev", "https://abdus.dev"])
|
||||
def test_invalid_domains(domain: str) -> None:
|
||||
with pytest.raises(ValueError):
|
||||
script.render_config(containers=[container(domain=domain)])
|
||||
|
||||
|
||||
def test_duplicate_domains() -> None:
|
||||
with pytest.raises(ValueError):
|
||||
script.render_config(containers=[container(), container(name="other")])
|
||||
|
||||
|
||||
def test_missing_published_ports() -> None:
|
||||
app = container()
|
||||
app["NetworkSettings"]["Ports"] = {}
|
||||
with pytest.raises(ValueError):
|
||||
script.render_config(containers=[app])
|
||||
|
||||
|
||||
def test_host_networking() -> None:
|
||||
app = container()
|
||||
app["HostConfig"]["NetworkMode"] = "host"
|
||||
with pytest.raises(ValueError):
|
||||
script.render_config(containers=[app])
|
||||
|
||||
|
||||
def test_stopped_unlabeled_and_empty() -> None:
|
||||
stopped = container()
|
||||
stopped["State"] = {"Running": False, "Status": "exited"}
|
||||
unlabeled = container()
|
||||
unlabeled["Config"]["Labels"] = None
|
||||
assert script.render_config(containers=[stopped, unlabeled]) == script.render_config(containers=[])
|
||||
|
||||
|
||||
def test_deterministic_and_ipv4_preference() -> None:
|
||||
first = container(name="a", domain="a.abdus.dev")
|
||||
second = container(name="z", domain="z.abdus.dev")
|
||||
first["NetworkSettings"]["Ports"]["3000/tcp"].append(
|
||||
{"HostIp": "::", "HostPort": "4510"},
|
||||
)
|
||||
forward = script.render_config(containers=[first, second])
|
||||
assert forward == script.render_config(containers=[second, first])
|
||||
assert "reverse_proxy 127.0.0.1:4510" in forward
|
||||
assert forward.index("a.abdus.dev") < forward.index("z.abdus.dev")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("left, right, protocol, port, expected", [
|
||||
("0.0.0.0", "127.0.0.1", "tcp", "4510", True),
|
||||
("127.0.0.1", "127.0.0.2", "tcp", "4510", False),
|
||||
("::", "127.0.0.1", "tcp", "4510", True),
|
||||
("127.0.0.1", "127.0.0.1", "udp", "4510", False),
|
||||
("127.0.0.1", "127.0.0.1", "tcp", "0", False),
|
||||
])
|
||||
def test_clashes(left: str, right: str, protocol: str, port: str, expected: bool) -> None:
|
||||
first = container()
|
||||
first["HostConfig"]["PortBindings"]["3000/tcp"][0]["HostIp"] = left
|
||||
first["NetworkSettings"]["Ports"] = copy.deepcopy(first["HostConfig"]["PortBindings"])
|
||||
second = container(name="stopped")
|
||||
second["Config"]["Labels"] = None
|
||||
second["State"] = {"Running": False, "Status": "exited"}
|
||||
second["NetworkSettings"]["Ports"] = {}
|
||||
second["HostConfig"]["PortBindings"] = {
|
||||
f"3000/{protocol}": [{"HostIp": right, "HostPort": port}],
|
||||
}
|
||||
with patch.object(script.logging, "warning") as warning:
|
||||
script.warn_port_clashes(containers=[first, second])
|
||||
assert warning.call_count == int(expected)
|
||||
|
||||
|
||||
def test_output_and_failure_preservation(tmp_path: Path) -> None:
|
||||
output = tmp_path / "generated.caddy"
|
||||
output.write_text(data="previous", encoding="utf-8")
|
||||
output.chmod(mode=0o640)
|
||||
with patch("sys.argv", ["docker_to_caddy.py", "--output", str(output)]):
|
||||
with patch.object(script, "inspect_containers", return_value=[container(), container()]):
|
||||
assert script.main() == 1
|
||||
assert output.read_text() == "previous"
|
||||
with patch.object(script, "inspect_containers", side_effect=subprocess.CalledProcessError(
|
||||
returncode=1, cmd="docker", stderr="permission denied",
|
||||
)):
|
||||
assert script.main() == 1
|
||||
assert output.read_text() == "previous"
|
||||
script.write_config(config="replacement", output=output)
|
||||
assert output.read_text() == "replacement"
|
||||
assert output.stat().st_mode & 0o777 == 0o640
|
||||
|
||||
|
||||
def test_stdout(capsys: pytest.CaptureFixture[str]) -> None:
|
||||
with patch("sys.argv", ["docker_to_caddy.py"]):
|
||||
with patch.object(script, "inspect_containers", return_value=[container()]):
|
||||
assert script.main() == 0
|
||||
assert capsys.readouterr().out == script.render_config(containers=[container()])
|
||||
|
||||
|
||||
def test_docker_discovery() -> None:
|
||||
responses = [
|
||||
subprocess.CompletedProcess(args=[], returncode=0, stdout="id1\nid2\n"),
|
||||
subprocess.CompletedProcess(args=[], returncode=0, stdout='[{"Name":"/one"},{"Name":"/two"}]'),
|
||||
]
|
||||
with patch.object(script.subprocess, "run", side_effect=responses) as run:
|
||||
assert script.inspect_containers() == [{"Name": "/one"}, {"Name": "/two"}]
|
||||
assert run.call_args_list[1].kwargs["args"] == ["docker", "container", "inspect", "id1", "id2"]
|
||||
with patch.object(script.subprocess, "run", return_value=subprocess.CompletedProcess(
|
||||
args=[], returncode=0, stdout="",
|
||||
)) as run:
|
||||
assert script.inspect_containers() == []
|
||||
assert run.call_count == 1
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main(args=[str(Path(__file__).resolve()), *sys.argv[1:]]))
|
||||
+35
-11
@@ -9,8 +9,8 @@ import os
|
||||
import re
|
||||
import struct
|
||||
import subprocess
|
||||
import sys
|
||||
import urllib.parse
|
||||
import zipfile
|
||||
from functools import partial
|
||||
from pathlib import Path
|
||||
from typing import Callable
|
||||
@@ -391,6 +391,9 @@ def filehost_cmd(urls: tuple[str, ...], name: str | None, cwd: Path, token: str)
|
||||
|
||||
_extract_and_process(output_dir)
|
||||
|
||||
def is_filehost(url: str) -> bool:
|
||||
kws = ["rg.to", "rapidgator.net", "trbt.cc"]
|
||||
return any(kw in url for kw in kws)
|
||||
|
||||
@cli.command("download")
|
||||
@click.argument("urls", nargs=-1, required=True)
|
||||
@@ -401,7 +404,7 @@ def filehost_cmd(urls: tuple[str, ...], name: str | None, cwd: Path, token: str)
|
||||
def download_cmd(
|
||||
ctx: click.Context, urls: tuple[str, ...], name: str | None, cwd: Path, token: str
|
||||
) -> None:
|
||||
filehost_urls = [u for u in urls if "rg.to" in u or "rapidgator.net" in u]
|
||||
filehost_urls = [u for u in urls if is_filehost(u)]
|
||||
gallery_urls = [u for u in urls if u not in set(filehost_urls)]
|
||||
if filehost_urls:
|
||||
ctx.invoke(filehost_cmd, urls=tuple(filehost_urls), name=name, cwd=cwd, token=token)
|
||||
@@ -410,37 +413,58 @@ def download_cmd(
|
||||
|
||||
|
||||
@cli.command("process")
|
||||
@click.argument("dirs", nargs=-1, required=True)
|
||||
def process_cmd(dirs: tuple[str, ...]) -> None:
|
||||
for d in dirs:
|
||||
p = Path(d)
|
||||
@click.argument("paths", nargs=-1, required=True, type=click.Path(path_type=Path))
|
||||
def process_cmd(paths: tuple[Path, ...]) -> None:
|
||||
"""Process directories and ZIP files, extracting ZIPs into sibling directories."""
|
||||
for p in paths:
|
||||
archive_path: Path | None = None
|
||||
if not p.exists():
|
||||
print(f"skipping {p}: does not exist", file=sys.stderr)
|
||||
logging.warning(f"skipping {p}: does not exist")
|
||||
continue
|
||||
if p.is_file() and p.suffix.lower() == ".zip":
|
||||
output_dir = p.with_suffix("")
|
||||
if output_dir.exists() and not output_dir.is_dir():
|
||||
logging.warning(f"skipping {p}: {output_dir} is not a directory")
|
||||
continue
|
||||
try:
|
||||
with zipfile.ZipFile(file=p) as archive:
|
||||
logging.info(f"extracting {p} → {output_dir}")
|
||||
archive.extractall(path=output_dir)
|
||||
except zipfile.BadZipFile:
|
||||
logging.warning(f"skipping {p}: invalid ZIP file")
|
||||
continue
|
||||
archive_path = p
|
||||
p = output_dir
|
||||
if not p.is_dir():
|
||||
logging.warning(f"skipping {p}: expected a directory or ZIP file")
|
||||
continue
|
||||
if _needs_resize(p):
|
||||
resized_dir = p / "resized"
|
||||
resized_dir.mkdir(exist_ok=True)
|
||||
originals = [f for f in p.rglob("*") if is_image(f)]
|
||||
print(f"resizing {len(originals)} image(s) in {p} → {resized_dir}")
|
||||
logging.info(f"resizing {len(originals)} image(s) in {p} → {resized_dir}")
|
||||
for img in originals:
|
||||
run_command(
|
||||
[
|
||||
"imgz",
|
||||
"resize",
|
||||
"-o",
|
||||
"--output",
|
||||
str(resized_dir / f"{img.stem}.resized.jpg"),
|
||||
str(img),
|
||||
]
|
||||
)
|
||||
if originals:
|
||||
print(f"trashing {len(originals)} original(s)")
|
||||
logging.info(f"trashing {len(originals)} original(s)")
|
||||
move_to_trash(originals)
|
||||
else:
|
||||
print(f"skipping {p}: already within size limits")
|
||||
logging.info(f"skipping {p}: already within size limits")
|
||||
|
||||
rename_children(p, flatten=True)
|
||||
remove_empty_dirs(p)
|
||||
if archive_path is not None:
|
||||
move_to_trash(file_paths=[archive_path])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s: %(message)s")
|
||||
cli()
|
||||
@@ -0,0 +1,86 @@
|
||||
#!/usr/bin/env -S uv run --script
|
||||
# /// script
|
||||
# requires-python = ">=3.11"
|
||||
# dependencies = ["pytest", "httpx", "pillow", "click"]
|
||||
# ///
|
||||
import sys
|
||||
import zipfile
|
||||
from pathlib import Path
|
||||
from unittest.mock import Mock
|
||||
|
||||
import pytest
|
||||
from click.testing import CliRunner
|
||||
from PIL import Image
|
||||
|
||||
import download_imageset
|
||||
|
||||
|
||||
@pytest.mark.parametrize("extension", [".zip", ".ZIP"])
|
||||
def test_process_zip_and_directory(
|
||||
tmp_path: Path, monkeypatch: pytest.MonkeyPatch, extension: str
|
||||
) -> None:
|
||||
directory = tmp_path / "existing"
|
||||
nested = directory / "nested"
|
||||
nested.mkdir(parents=True)
|
||||
image_path = nested / "photo.png"
|
||||
Image.new(mode="RGB", size=(10, 10)).save(fp=image_path)
|
||||
archive_path = tmp_path / f"gallery{extension}"
|
||||
with zipfile.ZipFile(file=archive_path, mode="w") as archive:
|
||||
archive.write(filename=image_path, arcname="album/photo.png")
|
||||
trash = Mock()
|
||||
monkeypatch.setattr(download_imageset, "move_to_trash", trash)
|
||||
|
||||
result = CliRunner().invoke(
|
||||
cli=download_imageset.cli,
|
||||
args=["process", str(archive_path), str(directory)],
|
||||
)
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert (tmp_path / "gallery" / "gallery__0000.png").is_file()
|
||||
assert (directory / "existing__0000.png").is_file()
|
||||
assert not (tmp_path / "gallery" / "album").exists()
|
||||
assert not nested.exists()
|
||||
trash.assert_called_once_with(file_paths=[archive_path])
|
||||
|
||||
|
||||
@pytest.mark.parametrize("filename", ["missing", "unsupported.txt", "broken.zip"])
|
||||
def test_process_skips_invalid_inputs(
|
||||
tmp_path: Path, monkeypatch: pytest.MonkeyPatch, filename: str
|
||||
) -> None:
|
||||
source = tmp_path / filename
|
||||
if filename != "missing":
|
||||
source.write_text(data="not an archive")
|
||||
trash = Mock()
|
||||
monkeypatch.setattr(download_imageset, "move_to_trash", trash)
|
||||
|
||||
result = CliRunner().invoke(
|
||||
cli=download_imageset.cli, args=["process", str(source)]
|
||||
)
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
trash.assert_not_called()
|
||||
if filename != "missing":
|
||||
assert source.read_text() == "not an archive"
|
||||
assert not (tmp_path / "broken").exists()
|
||||
|
||||
|
||||
def test_process_keeps_zip_when_processing_fails(
|
||||
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
archive_path = tmp_path / "gallery.zip"
|
||||
with zipfile.ZipFile(file=archive_path, mode="w") as archive:
|
||||
archive.writestr(zinfo_or_arcname="photo.png", data=b"invalid image")
|
||||
trash = Mock()
|
||||
monkeypatch.setattr(download_imageset, "move_to_trash", trash)
|
||||
|
||||
result = CliRunner().invoke(
|
||||
cli=download_imageset.cli, args=["process", str(archive_path)]
|
||||
)
|
||||
|
||||
assert result.exit_code != 0
|
||||
assert archive_path.is_file()
|
||||
trash.assert_not_called()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main(args=[str(Path(__file__).resolve()), *sys.argv[1:]]))
|
||||
@@ -215,6 +215,49 @@ def parse_release(
|
||||
actors = []
|
||||
title = None
|
||||
|
||||
# Prefer known actor names, checking longer names first so a short
|
||||
# entry cannot consume the beginning of a longer one.
|
||||
known_names = sorted(
|
||||
{name.strip() for name in known_actors if name.strip()},
|
||||
key=lambda name: (-len(name), name.casefold()),
|
||||
)
|
||||
words = remaining.split()
|
||||
|
||||
def match_known_name(start: int) -> tuple[str, int] | None:
|
||||
for name in known_names:
|
||||
name_words = name.split()
|
||||
end = start + len(name_words)
|
||||
if [word.casefold() for word in words[start:end]] == [
|
||||
word.casefold() for word in name_words
|
||||
]:
|
||||
return name, end
|
||||
return None
|
||||
|
||||
# Episode codes precede actor names in some sources, for example
|
||||
# "E1109 Uma Jolie". Keep the code as the title and match actors after it.
|
||||
actor_start = 0
|
||||
if words and words[0].startswith("E") and words[0][1:].isdigit():
|
||||
title = words[0]
|
||||
actor_start = 1
|
||||
|
||||
cursor = actor_start
|
||||
while known_names and (match := match_known_name(cursor)):
|
||||
actor, cursor = match
|
||||
actors.append(actor)
|
||||
if cursor < len(words) and words[cursor].casefold() == "and":
|
||||
cursor += 1
|
||||
|
||||
if actors:
|
||||
title_words = words[cursor:]
|
||||
if title_words:
|
||||
title = " ".join(([title] if title else []) + title_words)
|
||||
return Release(
|
||||
studio=studio,
|
||||
actors=sorted(set(actors)),
|
||||
title=title,
|
||||
released_at=parsed_date.date if parsed_date else None,
|
||||
)
|
||||
|
||||
# PornFidelity.E1109.Uma.Jolie
|
||||
match remaining.split(" "):
|
||||
case [episode, *names] if episode.startswith("E") and episode[1:].isdigit():
|
||||
|
||||
Executable
+124
@@ -0,0 +1,124 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Rename files according to filename mappings returned by OpenRouter."""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
from urllib.error import HTTPError, URLError
|
||||
from urllib.request import Request, urlopen
|
||||
|
||||
# Fill in the instructions that should guide the filename changes.
|
||||
SYSTEM_PROMPT = """
|
||||
Your task is to rename files. You will be given a JSON array of filenames.
|
||||
You must return a JSON array of objects, each with "find" and "replace" fields, where "find" is the original filename and "replace" is the new filename.
|
||||
Only include files that need to be renamed; do not include files that should remain unchanged.
|
||||
Ensure that the new filenames are valid and do not contain any directory paths or invalid characters.
|
||||
|
||||
The output format is "$actors -- @$studio -- $title -- $date"
|
||||
- $actors are separated by commas, and should be the names of the actors in the video.
|
||||
- $studio is the name of the studio that produced the video. Use "Unknown" if the studio is not known.
|
||||
- Not all files have title or studio or date, leave those bits out if they are not known.
|
||||
- $date must be in YYYY-MM-DD format. PRT releases are usually in YY.MM.DD format, so you will need to convert that to YYYY-MM-DD.
|
||||
|
||||
Examples:
|
||||
- 15.01.10_-_Aspen_Ora_-_Cock_Crazy_-_x265.mp4 -> Aspen Ora -- @Unknown -- Cock Crazy -- 2015-01-10.mp4
|
||||
- [onlyfans.com] Lilijunex - Huge Dildo Machine Slams Lilijunex [2026-06-28 -> Lilijunex -- @OnlyFans -- Huge Dildo Machine Slams Lilijunex -- 2026-06-28.mp4
|
||||
"""
|
||||
OPENROUTER_URL = "https://openrouter.ai/api/v1/chat/completions"
|
||||
|
||||
|
||||
def send_request(payload: dict[str, Any], api_key: str) -> dict[str, Any]:
|
||||
request = Request(
|
||||
OPENROUTER_URL,
|
||||
data=json.dumps(payload).encode("utf-8"),
|
||||
headers={
|
||||
"Authorization": f"Bearer {api_key}",
|
||||
"Content-Type": "application/json",
|
||||
},
|
||||
method="POST",
|
||||
)
|
||||
|
||||
try:
|
||||
with urlopen(request, timeout=60) as response:
|
||||
result = json.loads(response.read().decode("utf-8"))
|
||||
except HTTPError as error:
|
||||
detail = error.read().decode("utf-8", errors="replace")
|
||||
raise RuntimeError(f"OpenRouter returned HTTP {error.code}: {detail}") from error
|
||||
except URLError as error:
|
||||
raise RuntimeError(f"Could not reach OpenRouter: {error.reason}") from error
|
||||
return result
|
||||
|
||||
|
||||
def get_renames(filenames: list[str], api_key: str, model: str) -> dict[str, str]:
|
||||
payload = {
|
||||
"model": model,
|
||||
"messages": [
|
||||
{"role": "system", "content": SYSTEM_PROMPT},
|
||||
{"role": "user", "content": json.dumps(filenames)},
|
||||
],
|
||||
}
|
||||
result = send_request(payload, api_key)
|
||||
|
||||
try:
|
||||
content = result["choices"][0]["message"]["content"]
|
||||
mappings = json.loads(content)
|
||||
except (KeyError, IndexError, TypeError, json.JSONDecodeError) as error:
|
||||
raise RuntimeError("OpenRouter response did not contain a JSON mapping array.") from error
|
||||
|
||||
if not isinstance(mappings, list):
|
||||
raise RuntimeError("OpenRouter response must be a JSON array.")
|
||||
for mapping in mappings:
|
||||
if not isinstance(mapping, dict) or not isinstance(mapping.get("find"), str) or not isinstance(mapping.get("replace"), str):
|
||||
raise RuntimeError('Each mapping must have string "find" and "replace" fields.')
|
||||
return {mapping["find"]: mapping["replace"] for mapping in mappings}
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description="Ask OpenRouter for filename changes and rename the supplied files.")
|
||||
parser.add_argument("paths", nargs="+", type=Path, help="one or more file paths")
|
||||
parser.add_argument(
|
||||
"--openrouter-api-key",
|
||||
default=os.environ.get("OPENROUTER_API_KEY"),
|
||||
help="OpenRouter API key (defaults to OPENROUTER_API_KEY)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--openrouter-model",
|
||||
default=os.environ.get("OPENROUTER_MODEL", "~deepseek/deepseek-flash-latest"),
|
||||
help="OpenRouter model (defaults to OPENROUTER_MODEL)",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
for path in args.paths:
|
||||
if not path.exists():
|
||||
parser.error(f"file does not exist: {path}")
|
||||
|
||||
filenames = [path.name for path in args.paths]
|
||||
by_name = get_renames(filenames, args.openrouter_api_key, args.openrouter_model)
|
||||
|
||||
planned = []
|
||||
for path in args.paths:
|
||||
replacement = by_name.get(path.name)
|
||||
if replacement is not None and replacement != path.name:
|
||||
if not replacement or replacement in {".", ".."} or "/" in replacement or "\\" in replacement:
|
||||
raise RuntimeError(f"Invalid replacement filename: {replacement!r}")
|
||||
planned.append((path, path.with_name(replacement)))
|
||||
|
||||
destinations = [destination for _, destination in planned]
|
||||
if len(set(destinations)) != len(destinations):
|
||||
raise RuntimeError("Multiple files would be renamed to the same path.")
|
||||
for source, destination in planned:
|
||||
if destination.exists() and destination != source:
|
||||
raise RuntimeError(f"Destination already exists: {destination}")
|
||||
|
||||
for source, destination in planned:
|
||||
source.rename(destination)
|
||||
print(f"{source} -> {destination}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
try:
|
||||
main()
|
||||
except RuntimeError as error:
|
||||
raise SystemExit(str(error))
|
||||
+18
-4
@@ -6,14 +6,27 @@
|
||||
import argparse
|
||||
import logging
|
||||
import subprocess
|
||||
from pathlib import Path
|
||||
from threading import Thread
|
||||
from typing import Callable
|
||||
from pathlib import Path
|
||||
|
||||
GIGAPIXEL_EXE = Path(
|
||||
"/Applications/Topaz Gigapixel AI.app/Contents/MacOS/Topaz Gigapixel AI"
|
||||
)
|
||||
GIGAPIXEL_EXE = Path("/Applications/Topaz Gigapixel.app/Contents/MacOS/Topaz Gigapixel")
|
||||
MAX_LENGTH = 4000
|
||||
FACE_RECOVERY_ARGS = [
|
||||
"--sharpen",
|
||||
"1",
|
||||
"--denoise",
|
||||
"4",
|
||||
"--compression",
|
||||
"1",
|
||||
"--fr",
|
||||
"60",
|
||||
"--frv",
|
||||
"2",
|
||||
"--frc",
|
||||
"0",
|
||||
"--gamma-correction",
|
||||
]
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
@@ -91,6 +104,7 @@ def upscale_images(
|
||||
output_format,
|
||||
"--suffix",
|
||||
"-upscaled",
|
||||
*FACE_RECOVERY_ARGS,
|
||||
]
|
||||
|
||||
if output_format == "jpg":
|
||||
|
||||
Executable
+112
@@ -0,0 +1,112 @@
|
||||
#!/usr/bin/env -S uv run --script
|
||||
# /// script
|
||||
# dependencies = ["httpx"]
|
||||
# ///
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import sys
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from dataclasses import asdict, dataclass
|
||||
from pathlib import Path
|
||||
|
||||
import httpx
|
||||
|
||||
|
||||
@dataclass
|
||||
class UploadedImage:
|
||||
url: str
|
||||
thumb_url: str
|
||||
delete_url: str
|
||||
|
||||
|
||||
def upload_image_to_hamsterimg(api_key: str, image_paths: list[Path]) -> list[UploadedImage]:
|
||||
if not image_paths:
|
||||
return []
|
||||
if not api_key:
|
||||
raise ValueError("HamsterImg API key is required")
|
||||
for image_path in image_paths:
|
||||
if not image_path.is_file():
|
||||
raise FileNotFoundError(f"Image file not found: {image_path}")
|
||||
with httpx.Client(headers={"X-API-Key": api_key}, timeout=60) as client:
|
||||
|
||||
def upload(image_path: Path) -> UploadedImage:
|
||||
with image_path.open(mode="rb") as image_file:
|
||||
response = client.post(
|
||||
url="https://hamsterimg.net/api/1/upload",
|
||||
data={"format": "json"},
|
||||
files={"source": (image_path.name, image_file)},
|
||||
)
|
||||
try:
|
||||
result = response.json()
|
||||
except ValueError:
|
||||
result = {}
|
||||
if response.is_error or not isinstance(result, dict) or result.get("status_code") != 200:
|
||||
error = result.get("error", {}) if isinstance(result, dict) else {}
|
||||
message = error.get("message") if isinstance(error, dict) else str(error)
|
||||
status = result.get("status_txt") if isinstance(result, dict) else None
|
||||
reason = message or status or response.reason_phrase or "Unknown error"
|
||||
raise RuntimeError(f"HamsterImg upload failed for {image_path}: {reason}")
|
||||
try:
|
||||
image = result["image"]
|
||||
url = image["url"]
|
||||
thumb_url = image["thumb"]["url"]
|
||||
delete_url = image["delete_url"]
|
||||
except (KeyError, TypeError) as error:
|
||||
raise ValueError(f"HamsterImg returned incomplete image details for {image_path}") from error
|
||||
if not all(isinstance(value, str) and value for value in (url, thumb_url, delete_url)):
|
||||
raise ValueError(f"HamsterImg returned missing image URLs for {image_path}")
|
||||
return UploadedImage(url=url, thumb_url=thumb_url, delete_url=delete_url)
|
||||
|
||||
with ThreadPoolExecutor(max_workers=min(8, len(image_paths))) as executor:
|
||||
futures = [executor.submit(upload, image_path=image_path) for image_path in image_paths]
|
||||
return [future.result() for future in futures]
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Upload images to HamsterImg and print their URLs.",
|
||||
formatter_class=argparse.ArgumentDefaultsHelpFormatter,
|
||||
)
|
||||
parser.add_argument("image_paths", nargs="+", type=Path, help="Image files to upload")
|
||||
parser.add_argument("--api-key", help="HamsterImg API key; also reads HAMSTERIMG_API_KEY")
|
||||
parser.add_argument(
|
||||
"--format",
|
||||
dest="output_format",
|
||||
choices=["text", "json", "bbcode"],
|
||||
default="text",
|
||||
help="Output format; BBCode shows thumbnails linked to full images",
|
||||
)
|
||||
arguments = parser.parse_args()
|
||||
arguments.api_key = arguments.api_key or os.getenv("HAMSTERIMG_API_KEY")
|
||||
if not arguments.api_key:
|
||||
parser.error("Provide --api-key or set HAMSTERIMG_API_KEY")
|
||||
for image_path in arguments.image_paths:
|
||||
if not image_path.is_file():
|
||||
parser.error(f"Image file not found: {image_path}")
|
||||
return arguments
|
||||
|
||||
|
||||
def main() -> int:
|
||||
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s: %(message)s")
|
||||
arguments = parse_args()
|
||||
try:
|
||||
images = upload_image_to_hamsterimg(api_key=arguments.api_key, image_paths=arguments.image_paths)
|
||||
except (OSError, ValueError, RuntimeError, httpx.HTTPError) as error:
|
||||
logging.error(f"{error}")
|
||||
return 1
|
||||
if arguments.output_format == "json":
|
||||
print(json.dumps([asdict(obj=image) for image in images], indent=2, ensure_ascii=False))
|
||||
return 0
|
||||
for image in images:
|
||||
if arguments.output_format == "bbcode":
|
||||
print(f"[url={image.url}][img]{image.thumb_url}[/img][/url]")
|
||||
else:
|
||||
print(image.url)
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
@@ -0,0 +1,195 @@
|
||||
#!/usr/bin/env -S uv run --script
|
||||
# /// script
|
||||
# dependencies = ["httpx", "pytest"]
|
||||
# ///
|
||||
|
||||
import json
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from threading import Barrier
|
||||
from unittest.mock import patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
import hamsterimg as script
|
||||
|
||||
|
||||
def test_hamsterimg_uploads_in_parallel_and_preserves_order(tmp_path: Path) -> None:
|
||||
image_paths = [tmp_path / "first.jpg", tmp_path / "second.png"]
|
||||
for image_path in image_paths:
|
||||
image_path.write_bytes(image_path.name.encode())
|
||||
barrier = Barrier(parties=2)
|
||||
|
||||
def respond(request: httpx.Request) -> httpx.Response:
|
||||
assert str(request.url) == "https://hamsterimg.net/api/1/upload"
|
||||
assert request.method == "POST"
|
||||
assert request.headers["X-API-Key"] == "test-key"
|
||||
assert request.headers["Content-Type"].startswith("multipart/form-data;")
|
||||
body = request.read()
|
||||
assert b'name="source"' in body
|
||||
assert b'name="format"' in body
|
||||
filename = "first.jpg" if b'filename="first.jpg"' in body else "second.png"
|
||||
barrier.wait(timeout=5)
|
||||
return httpx.Response(status_code=200, json={
|
||||
"status_code": 200,
|
||||
"image": {
|
||||
"url": f"https://images.example/{filename}",
|
||||
"thumb": {"url": f"https://images.example/thumb/{filename}"},
|
||||
"delete_url": f"https://hamsterimg.net/{filename}/delete/token",
|
||||
},
|
||||
})
|
||||
|
||||
client = httpx.Client(transport=httpx.MockTransport(handler=respond), headers={"X-API-Key": "test-key"})
|
||||
with patch.object(script.httpx, "Client", return_value=client) as factory:
|
||||
images = script.upload_image_to_hamsterimg(api_key="test-key", image_paths=image_paths)
|
||||
factory.assert_called_once_with(headers={"X-API-Key": "test-key"}, timeout=60)
|
||||
assert images == [
|
||||
script.UploadedImage(
|
||||
url="https://images.example/first.jpg",
|
||||
thumb_url="https://images.example/thumb/first.jpg",
|
||||
delete_url="https://hamsterimg.net/first.jpg/delete/token",
|
||||
),
|
||||
script.UploadedImage(
|
||||
url="https://images.example/second.png",
|
||||
thumb_url="https://images.example/thumb/second.png",
|
||||
delete_url="https://hamsterimg.net/second.png/delete/token",
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("status, body, error", [
|
||||
(401, {"error": {"message": "Invalid API key"}}, RuntimeError),
|
||||
(200, {"status_code": 400, "error": {"message": "Invalid image"}}, RuntimeError),
|
||||
(200, {"status_code": 200, "image": {}}, ValueError),
|
||||
(200, {"status_code": 200, "image": {"url": "https://images.example/full.jpg", "thumb": {"url": "https://images.example/thumb.jpg"}}}, ValueError),
|
||||
(200, {"status_code": 200, "image": {"url": "https://images.example/full.jpg", "thumb": {"url": None}, "delete_url": "https://hamsterimg.net/delete/token"}}, ValueError),
|
||||
])
|
||||
def test_hamsterimg_reports_upload_errors(tmp_path: Path, status: int, body: dict, error: type[Exception]) -> None:
|
||||
image_path = tmp_path / "image.jpg"
|
||||
image_path.write_bytes(b"image")
|
||||
|
||||
def respond(request: httpx.Request) -> httpx.Response:
|
||||
return httpx.Response(status_code=status, json=body)
|
||||
|
||||
client = httpx.Client(transport=httpx.MockTransport(handler=respond))
|
||||
with patch.object(script.httpx, "Client", return_value=client):
|
||||
with pytest.raises(error):
|
||||
script.upload_image_to_hamsterimg(api_key="test-key", image_paths=[image_path])
|
||||
|
||||
|
||||
def test_hamsterimg_upload_error_includes_api_reason(tmp_path: Path) -> None:
|
||||
image_path = tmp_path / "image.jpg"
|
||||
image_path.write_bytes(b"image")
|
||||
|
||||
def respond(request: httpx.Request) -> httpx.Response:
|
||||
return httpx.Response(
|
||||
status_code=400,
|
||||
json={"status_code": 400, "error": {"message": "Unable to read image"}},
|
||||
)
|
||||
|
||||
client = httpx.Client(transport=httpx.MockTransport(handler=respond))
|
||||
with patch.object(script.httpx, "Client", return_value=client):
|
||||
with pytest.raises(RuntimeError, match="Unable to read image"):
|
||||
script.upload_image_to_hamsterimg(api_key="test-key", image_paths=[image_path])
|
||||
|
||||
|
||||
def test_hamsterimg_validates_inputs_before_uploading(tmp_path: Path) -> None:
|
||||
image_path = tmp_path / "image.jpg"
|
||||
image_path.write_bytes(b"image")
|
||||
with patch.object(script.httpx, "Client") as client:
|
||||
assert script.upload_image_to_hamsterimg(api_key="test-key", image_paths=[]) == []
|
||||
with pytest.raises(ValueError, match="API key"):
|
||||
script.upload_image_to_hamsterimg(api_key="", image_paths=[image_path])
|
||||
with pytest.raises(FileNotFoundError):
|
||||
script.upload_image_to_hamsterimg(api_key="test-key", image_paths=[image_path, tmp_path / "missing.jpg"])
|
||||
client.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("output_format", [None, "text", "json", "bbcode"])
|
||||
def test_upload_image_command_output(
|
||||
tmp_path: Path, output_format: str | None, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str],
|
||||
) -> None:
|
||||
image_paths = [tmp_path / "first.jpg", tmp_path / "second.png"]
|
||||
for image_path in image_paths:
|
||||
image_path.write_bytes(b"image")
|
||||
images = [
|
||||
script.UploadedImage(
|
||||
url="https://images.example/first.jpg",
|
||||
thumb_url="https://images.example/first.th.jpg",
|
||||
delete_url="https://hamsterimg.net/first/delete/token",
|
||||
),
|
||||
script.UploadedImage(
|
||||
url="https://images.example/second.png",
|
||||
thumb_url="https://images.example/second.th.png",
|
||||
delete_url="https://hamsterimg.net/second/delete/token",
|
||||
),
|
||||
]
|
||||
arguments = ["hamsterimg.py", *map(str, image_paths)]
|
||||
if output_format is not None:
|
||||
arguments.extend(["--format", output_format])
|
||||
monkeypatch.setattr(sys, "argv", arguments)
|
||||
monkeypatch.setenv("HAMSTERIMG_API_KEY", "test-key")
|
||||
with patch.object(script, "upload_image_to_hamsterimg", return_value=images) as upload:
|
||||
assert script.main() == 0
|
||||
result = capsys.readouterr()
|
||||
assert result.err == ""
|
||||
upload.assert_called_once_with(api_key="test-key", image_paths=image_paths)
|
||||
if output_format == "json":
|
||||
assert json.loads(result.out) == [
|
||||
{
|
||||
"url": "https://images.example/first.jpg",
|
||||
"thumb_url": "https://images.example/first.th.jpg",
|
||||
"delete_url": "https://hamsterimg.net/first/delete/token",
|
||||
},
|
||||
{
|
||||
"url": "https://images.example/second.png",
|
||||
"thumb_url": "https://images.example/second.th.png",
|
||||
"delete_url": "https://hamsterimg.net/second/delete/token",
|
||||
},
|
||||
]
|
||||
elif output_format == "bbcode":
|
||||
assert result.out == (
|
||||
"[url=https://images.example/first.jpg][img]https://images.example/first.th.jpg[/img][/url]\n"
|
||||
"[url=https://images.example/second.png][img]https://images.example/second.th.png[/img][/url]\n"
|
||||
)
|
||||
else:
|
||||
assert result.out == "https://images.example/first.jpg\nhttps://images.example/second.png\n"
|
||||
|
||||
|
||||
def test_upload_image_command_failure(
|
||||
tmp_path: Path, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str], caplog: pytest.LogCaptureFixture,
|
||||
) -> None:
|
||||
image_path = tmp_path / "image.jpg"
|
||||
image_path.write_bytes(b"image")
|
||||
monkeypatch.setattr(sys, "argv", ["hamsterimg.py", str(image_path), "--api-key", "explicit-key", "--format", "json"])
|
||||
with patch.object(script, "upload_image_to_hamsterimg", side_effect=RuntimeError("Invalid image")) as upload:
|
||||
assert script.main() == 1
|
||||
assert capsys.readouterr().out == ""
|
||||
assert "Invalid image" in caplog.text
|
||||
upload.assert_called_once_with(api_key="explicit-key", image_paths=[image_path])
|
||||
|
||||
|
||||
def test_upload_image_command_requires_paths(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(sys, "argv", ["hamsterimg.py", "--api-key", "test-key"])
|
||||
with patch.object(script, "upload_image_to_hamsterimg") as upload:
|
||||
with pytest.raises(SystemExit) as error:
|
||||
script.main()
|
||||
assert error.value.code == 2
|
||||
upload.assert_not_called()
|
||||
|
||||
|
||||
def test_upload_image_command_requires_key(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
image_path = tmp_path / "image.jpg"
|
||||
image_path.write_bytes(b"image")
|
||||
monkeypatch.setattr(sys, "argv", ["hamsterimg.py", str(image_path)])
|
||||
monkeypatch.delenv("HAMSTERIMG_API_KEY", raising=False)
|
||||
with patch.object(script, "upload_image_to_hamsterimg") as upload:
|
||||
with pytest.raises(SystemExit) as error:
|
||||
script.main()
|
||||
assert error.value.code == 2
|
||||
upload.assert_not_called()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main(args=[str(Path(__file__).resolve()), *sys.argv[1:]]))
|
||||
@@ -1,80 +0,0 @@
|
||||
import datetime
|
||||
import importlib
|
||||
import json
|
||||
import logging
|
||||
import typing
|
||||
import uuid
|
||||
from typing import Any
|
||||
|
||||
from fastapi import FastAPI, HTTPException
|
||||
from pydantic import BaseModel
|
||||
from starlette import status
|
||||
import sentry_sdk
|
||||
from sentry_sdk.integrations.asgi import SentryAsgiMiddleware
|
||||
import seqlog
|
||||
|
||||
app = FastAPI()
|
||||
|
||||
|
||||
class JsonDateEncoder(json.JSONEncoder):
|
||||
def default(self, o: Any) -> Any:
|
||||
if isinstance(o, datetime.datetime):
|
||||
return o.isoformat()
|
||||
return super().default(o)
|
||||
|
||||
|
||||
seqlog.log_to_seq(
|
||||
server_url="http://log.abdus.dev:5341/",
|
||||
api_key="ePewcJUxe8dST5EiC8Pw",
|
||||
level=logging.INFO,
|
||||
batch_size=10,
|
||||
auto_flush_timeout=10, # seconds
|
||||
override_root_logger=True,
|
||||
json_encoder_class=JsonDateEncoder, # Optional; only specify this if you want to use a custom JSON encoder
|
||||
)
|
||||
sentry_sdk.init(
|
||||
"https://ccb81025166c449d87e0646cede48448@o271257.ingest.sentry.io/6040020",
|
||||
# Set traces_sample_rate to 1.0 to capture 100%
|
||||
# of transactions for performance monitoring.
|
||||
# We recommend adjusting this value in production.
|
||||
traces_sample_rate=1.0,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class HookResponse(BaseModel):
|
||||
id: str
|
||||
result: typing.Optional[dict] = None
|
||||
error: typing.Optional[str] = None
|
||||
|
||||
|
||||
@app.post("/{event}", response_model=HookResponse)
|
||||
def receive_hook(event: str, payload: dict):
|
||||
logger.info(f"Received hook request for {event=} with {payload=}")
|
||||
hook_id = uuid.uuid4().hex
|
||||
|
||||
try:
|
||||
module = importlib.import_module(f"scripts.{event}")
|
||||
func = getattr(module, "main")
|
||||
except ModuleNotFoundError:
|
||||
raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="No such handler")
|
||||
|
||||
if not callable(func):
|
||||
raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail="Invalid script interface")
|
||||
|
||||
try:
|
||||
logger.info(f"Running hook handler {func=}")
|
||||
result = func(payload)
|
||||
return HookResponse(id=hook_id, result=result, error=None)
|
||||
except Exception as e:
|
||||
logger.error(f"Hook handler threw error", exc_info=True)
|
||||
return HookResponse(id=hook_id, result=None, error=str(e))
|
||||
|
||||
|
||||
app = SentryAsgiMiddleware(app)
|
||||
|
||||
if __name__ == "__main__":
|
||||
import uvicorn
|
||||
|
||||
uvicorn.run(app)
|
||||
@@ -1,21 +0,0 @@
|
||||
import logging
|
||||
import typing
|
||||
from pathlib import Path
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
Payload = typing.TypedDict(
|
||||
"Payload",
|
||||
{
|
||||
"gid": str,
|
||||
"num_files": int,
|
||||
"path": str,
|
||||
"files": list[typing.TypedDict("File", {"path": str, "size": int})],
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def main(payload: Payload):
|
||||
logging.info("got {torrent}", torrent=payload)
|
||||
path = Path(payload["path"])
|
||||
@@ -1,26 +0,0 @@
|
||||
import logging
|
||||
from pathlib import Path
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def main(payload: dict = None):
|
||||
if not payload:
|
||||
payload = {}
|
||||
|
||||
links: list[dict] = payload.get("links", [])
|
||||
if not links:
|
||||
logger.info("No links, exiting")
|
||||
return
|
||||
|
||||
logger.info("Got %d links", len(links))
|
||||
|
||||
for it in links:
|
||||
download_path = Path(it["download_path"])
|
||||
url: str = it.get("url", it.get("content_url"))
|
||||
logger.info("Processing downloaded link {url} at {path}", path=download_path, url=url)
|
||||
|
||||
if download_path.suffix.lower() in {".zip", ".rar"}:
|
||||
logger.info("Got an imageset!")
|
||||
|
||||
# raise Exception("oops!")
|
||||
Whitespace-only changes.
Executable
+444
@@ -0,0 +1,444 @@
|
||||
#!/usr/bin/env -S uv run --script
|
||||
# /// script
|
||||
# dependencies = ["Pillow"]
|
||||
# ///
|
||||
|
||||
import argparse
|
||||
import atexit
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
from concurrent.futures import FIRST_COMPLETED, Future, ProcessPoolExecutor, wait
|
||||
from contextlib import ExitStack
|
||||
from dataclasses import dataclass
|
||||
from functools import partial
|
||||
from io import BytesIO
|
||||
from multiprocessing import get_context
|
||||
from pathlib import Path
|
||||
from time import perf_counter
|
||||
from typing import BinaryIO, Callable
|
||||
from zipfile import BadZipFile, ZipFile, ZipInfo, is_zipfile
|
||||
|
||||
from PIL import Image, ImageColor, ImageOps
|
||||
|
||||
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s: %(message)s")
|
||||
worker_archive: ZipFile | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class Source:
|
||||
name: str
|
||||
path: Path | None = None
|
||||
member: ZipInfo | None = None
|
||||
ratio: float = 1.0
|
||||
|
||||
|
||||
@dataclass
|
||||
class Row:
|
||||
sources: list[Source]
|
||||
height: int
|
||||
widths: list[int]
|
||||
|
||||
|
||||
@dataclass
|
||||
class Cell:
|
||||
source: Source
|
||||
position: tuple[int, int]
|
||||
size: tuple[int, int]
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Create a separate justified image contact sheet for each directory or ZIP.",
|
||||
formatter_class=argparse.ArgumentDefaultsHelpFormatter,
|
||||
)
|
||||
parser.add_argument("inputs", nargs="+", type=Path, help="Image directories or ZIP archives; missing paths are skipped")
|
||||
parser.add_argument("--output", type=Path, help="JPEG output for a single existing input; defaults to <input-name>_thumbs.jpg beside each input")
|
||||
parser.add_argument("--width", type=int, default=2500, help="Sheet width in pixels")
|
||||
parser.add_argument("--target-height", type=int, default=300, help="Preferred row height in pixels; rows may be taller or shorter")
|
||||
parser.add_argument("--max-rows-per-page", type=int, help="Maximum rows per output image; additional rows are written to numbered pages")
|
||||
parser.add_argument("--gap", type=int, default=8, help="Gap between images and at sheet edges in pixels")
|
||||
parser.add_argument("--bg-color", default="white", help="Background color as a color name or hex value")
|
||||
parser.add_argument("--workers", type=int, default=os.cpu_count() or 1, help="Number of resize worker processes")
|
||||
arguments = parser.parse_args()
|
||||
if arguments.target_height <= 0 or arguments.width <= 0 or arguments.gap < 0:
|
||||
parser.error("width and target height must be positive; gap must be nonnegative")
|
||||
if arguments.width <= 2 * arguments.gap:
|
||||
parser.error("width must exceed twice the gap")
|
||||
if arguments.workers <= 0:
|
||||
parser.error("workers must be positive")
|
||||
if arguments.max_rows_per_page is not None and arguments.max_rows_per_page <= 0:
|
||||
parser.error("maximum rows per page must be positive")
|
||||
try:
|
||||
ImageColor.getcolor(color=arguments.bg_color, mode="RGB")
|
||||
except ValueError as error:
|
||||
parser.error(f"Invalid background color: {error}")
|
||||
return arguments
|
||||
|
||||
|
||||
def read_image(source: Source, archive: ZipFile | None) -> Image.Image:
|
||||
with ExitStack() as stack:
|
||||
if source.member is not None:
|
||||
if archive is None:
|
||||
raise ValueError("ZIP source requires an open archive")
|
||||
stream = stack.enter_context(archive.open(name=source.member))
|
||||
original = stack.enter_context(Image.open(fp=stream))
|
||||
else:
|
||||
original = stack.enter_context(Image.open(fp=source.path))
|
||||
oriented = ImageOps.exif_transpose(image=original)
|
||||
result = oriented.convert(mode="RGBA")
|
||||
if oriented is not original:
|
||||
oriented.close()
|
||||
return result
|
||||
|
||||
|
||||
def read_dims_jpeg(stream: BinaryIO) -> tuple[int, int]:
|
||||
EXIF_ORIENTATION_TAG = 274
|
||||
ORIENTATIONS_SWAPPING_WIDTH_HEIGHT = {5, 6, 7, 8}
|
||||
with Image.open(fp=stream) as image:
|
||||
width, height = image.size
|
||||
if image.getexif().get(EXIF_ORIENTATION_TAG) in ORIENTATIONS_SWAPPING_WIDTH_HEIGHT:
|
||||
return height, width
|
||||
return width, height
|
||||
|
||||
|
||||
def read_dims(image_path: Path | BinaryIO) -> tuple[int, int]:
|
||||
EXIF_ORIENTATION_TAG = 274
|
||||
ORIENTATIONS_SWAPPING_WIDTH_HEIGHT = {5, 6, 7, 8}
|
||||
with ExitStack() as stack:
|
||||
stream = stack.enter_context(image_path.open(mode="rb")) if isinstance(image_path, Path) else image_path
|
||||
signature = stream.read(12)
|
||||
stream.seek(0)
|
||||
if signature.startswith(b"\xff\xd8"):
|
||||
return read_dims_jpeg(stream=stream)
|
||||
if signature.startswith(b"\x89PNG\r\n\x1a\n"):
|
||||
raise ValueError("PNG images are not supported")
|
||||
if signature[:4] == b"RIFF" and signature[8:12] == b"WEBP":
|
||||
raise ValueError("WebP images are not supported")
|
||||
with Image.open(fp=stream) as image:
|
||||
if image.format == "TIFF":
|
||||
raise ValueError("TIFF images are not supported")
|
||||
width, height = image.size
|
||||
if image.getexif().get(EXIF_ORIENTATION_TAG) in ORIENTATIONS_SWAPPING_WIDTH_HEIGHT:
|
||||
return height, width
|
||||
return width, height
|
||||
|
||||
|
||||
def natural_sort_key(name: str) -> tuple[str | int, ...]:
|
||||
return tuple(int(part) if index % 2 else part.casefold() for index, part in enumerate(re.split(pattern=r"([0-9]+)", string=name)))
|
||||
|
||||
|
||||
def discover_sources(input_path: Path, output_path: Path, archive: ZipFile | None) -> list[Source]:
|
||||
Image.init()
|
||||
extensions = {extension for extension, image_format in Image.registered_extensions().items() if image_format not in {"PNG", "WEBP", "TIFF"}}
|
||||
if archive is not None:
|
||||
candidates = [
|
||||
Source(name=member.filename, member=member)
|
||||
for member in archive.infolist()
|
||||
if not member.is_dir() and Path(member.filename).suffix.lower() in extensions
|
||||
]
|
||||
else:
|
||||
candidates = [
|
||||
Source(name=path.relative_to(input_path).as_posix(), path=path)
|
||||
for path in input_path.rglob("*")
|
||||
if path.is_file()
|
||||
and path.suffix.lower() in extensions
|
||||
and path.resolve() != output_path.resolve()
|
||||
and not (
|
||||
path.parent.resolve() == output_path.parent.resolve()
|
||||
and path.suffix.lower() == output_path.suffix.lower()
|
||||
and re.fullmatch(pattern=rf"{re.escape(output_path.stem)}_page_[0-9]+", string=path.stem)
|
||||
)
|
||||
]
|
||||
sources = []
|
||||
logging.info(f"Reading metadata for {len(candidates)} images")
|
||||
for index, source in enumerate(sorted(candidates, key=lambda source: (natural_sort_key(name=source.name), source.name)), start=1):
|
||||
try:
|
||||
if source.member is not None and archive is not None:
|
||||
with archive.open(name=source.member) as stream:
|
||||
width, height = read_dims(image_path=stream)
|
||||
elif source.path is not None:
|
||||
width, height = read_dims(image_path=source.path)
|
||||
else:
|
||||
raise ValueError("Image source has no file or ZIP member")
|
||||
source.ratio = width / height
|
||||
sources.append(source)
|
||||
except (OSError, ValueError, SyntaxError, RuntimeError, BadZipFile, Image.DecompressionBombError) as error:
|
||||
logging.warning(f"Skipping {source.name}: {error}")
|
||||
if index % 100 == 0 or index == len(candidates):
|
||||
logging.info(f"Metadata {index}/{len(candidates)} ({index / len(candidates):.0%}); {len(sources)} usable images")
|
||||
if not sources:
|
||||
raise ValueError("No usable images found")
|
||||
return sources
|
||||
|
||||
|
||||
def fitted_height(sources: list[Source], width: int, gap: int) -> float:
|
||||
available = width - gap * (len(sources) + 1)
|
||||
return available / sum(source.ratio for source in sources)
|
||||
|
||||
|
||||
def make_row(sources: list[Source], width: int, gap: int, target_height: int, justify: bool) -> Row:
|
||||
available = width - gap * (len(sources) + 1)
|
||||
ideal_height = fitted_height(sources=sources, width=width, gap=gap)
|
||||
height = max(1, round(ideal_height if justify else min(target_height, ideal_height)))
|
||||
ratio_sum = sum(source.ratio for source in sources)
|
||||
total_width = available if justify else min(available, max(len(sources), round(ratio_sum * height)))
|
||||
widths = []
|
||||
cumulative = 0.0
|
||||
previous = 0
|
||||
for index, source in enumerate(sources):
|
||||
cumulative += source.ratio
|
||||
boundary = min(
|
||||
total_width - (len(sources) - index - 1),
|
||||
max(previous + 1, round(total_width * cumulative / ratio_sum)),
|
||||
)
|
||||
widths.append(boundary - previous)
|
||||
previous = boundary
|
||||
return Row(sources=sources, height=height, widths=widths)
|
||||
|
||||
|
||||
def layout_rows(sources: list[Source], width: int, gap: int, target_height: int) -> list[Row]:
|
||||
rows = []
|
||||
start = 0
|
||||
while start < len(sources):
|
||||
end = start + 1
|
||||
while end < len(sources) and fitted_height(sources=sources[start:end], width=width, gap=gap) > target_height:
|
||||
candidate = sources[start : end + 1]
|
||||
if width - gap * (len(candidate) + 1) < len(candidate):
|
||||
break
|
||||
candidate_height = fitted_height(sources=candidate, width=width, gap=gap)
|
||||
current_height = fitted_height(sources=sources[start:end], width=width, gap=gap)
|
||||
if candidate_height <= target_height and abs(current_height - target_height) <= abs(candidate_height - target_height):
|
||||
break
|
||||
end += 1
|
||||
group = sources[start:end]
|
||||
justify = end < len(sources) or fitted_height(sources=group, width=width, gap=gap) <= target_height
|
||||
rows.append(make_row(sources=group, width=width, gap=gap, target_height=target_height, justify=justify))
|
||||
start = end
|
||||
return rows
|
||||
|
||||
|
||||
def initialize_worker(archive_path: Path | None) -> None:
|
||||
global worker_archive
|
||||
if archive_path is not None:
|
||||
worker_archive = ZipFile(file=archive_path)
|
||||
atexit.register(worker_archive.close)
|
||||
|
||||
|
||||
def resize_with_pillow(cell: Cell) -> bytes:
|
||||
with read_image(source=cell.source, archive=worker_archive) as image:
|
||||
with image.resize(size=cell.size, resample=Image.Resampling.LANCZOS) as thumbnail:
|
||||
return thumbnail.tobytes()
|
||||
|
||||
|
||||
def resize_with_vips(cell: Cell, vips_binary: str) -> bytes:
|
||||
input_bytes = None
|
||||
if cell.source.member is not None:
|
||||
if worker_archive is None:
|
||||
raise ValueError("ZIP source requires an open archive")
|
||||
input_bytes = worker_archive.read(name=cell.source.member)
|
||||
operation, input_name = "thumbnail_source", "[descriptor=0]"
|
||||
elif cell.source.path is not None:
|
||||
operation, input_name = "thumbnail", str(cell.source.path)
|
||||
else:
|
||||
raise ValueError("Image source has no file or ZIP member")
|
||||
try:
|
||||
result = subprocess.run(
|
||||
args=[
|
||||
vips_binary,
|
||||
"--vips-concurrency=1",
|
||||
operation,
|
||||
input_name,
|
||||
".png",
|
||||
str(cell.size[0]),
|
||||
"--height",
|
||||
str(cell.size[1]),
|
||||
"--size",
|
||||
"both",
|
||||
],
|
||||
input=input_bytes,
|
||||
check=True,
|
||||
capture_output=True,
|
||||
)
|
||||
except (OSError, subprocess.CalledProcessError) as error:
|
||||
detail = error.stderr.decode(errors="replace").strip() if isinstance(error, subprocess.CalledProcessError) and error.stderr else str(error)
|
||||
logging.warning(f"vips failed for {cell.source.name}; using Pillow: {detail}")
|
||||
return resize_with_pillow(cell=cell)
|
||||
with Image.open(fp=BytesIO(result.stdout)) as image:
|
||||
with image.convert(mode="RGBA") as thumbnail:
|
||||
if thumbnail.size == cell.size:
|
||||
return thumbnail.tobytes()
|
||||
with thumbnail.resize(size=cell.size, resample=Image.Resampling.LANCZOS) as exact_thumbnail:
|
||||
return exact_thumbnail.tobytes()
|
||||
|
||||
|
||||
def resize_thumbnail(cell: Cell, processor: Callable[..., bytes]) -> tuple[bytes, float]:
|
||||
started = perf_counter()
|
||||
logging.info(f"Resizing {cell.source.name} to {cell.size[0]}x{cell.size[1]}")
|
||||
pixels = processor(cell=cell)
|
||||
return pixels, perf_counter() - started
|
||||
|
||||
|
||||
def render_rows(
|
||||
sheet: Image.Image,
|
||||
rows: list[Row],
|
||||
input_path: Path,
|
||||
workers: int,
|
||||
gap: int,
|
||||
processor: Callable[..., bytes],
|
||||
) -> int:
|
||||
cells = []
|
||||
top = gap
|
||||
for row in rows:
|
||||
left = gap
|
||||
for source, cell_width in zip(row.sources, row.widths):
|
||||
cells.append(Cell(source=source, position=(left, top), size=(cell_width, row.height)))
|
||||
left += cell_width + gap
|
||||
top += row.height + gap
|
||||
remaining = iter(cells)
|
||||
pending: dict[Future[tuple[bytes, float]], Cell] = {}
|
||||
processed_count = rendered_count = 0
|
||||
archive_path = None if input_path.is_dir() else input_path
|
||||
with ProcessPoolExecutor(
|
||||
max_workers=workers,
|
||||
mp_context=get_context(method="spawn"),
|
||||
initializer=initialize_worker,
|
||||
initargs=(archive_path,),
|
||||
) as pool:
|
||||
while True:
|
||||
while len(pending) < workers * 2:
|
||||
cell = next(remaining, None)
|
||||
if cell is None:
|
||||
break
|
||||
future = pool.submit(resize_thumbnail, cell=cell, processor=processor)
|
||||
pending[future] = cell
|
||||
if not pending:
|
||||
break
|
||||
completed, _ = wait(fs=pending, return_when=FIRST_COMPLETED)
|
||||
for future in completed:
|
||||
cell = pending.pop(future)
|
||||
processed_count += 1
|
||||
try:
|
||||
pixels, duration = future.result()
|
||||
with Image.frombytes(mode="RGBA", size=cell.size, data=pixels) as thumbnail:
|
||||
sheet.paste(im=thumbnail, box=cell.position, mask=thumbnail)
|
||||
rendered_count += 1
|
||||
logging.info(f"Placed {cell.source.name} (resized in {duration:.1f}s)")
|
||||
except (OSError, ValueError, SyntaxError, RuntimeError, BadZipFile, Image.DecompressionBombError) as error:
|
||||
logging.warning(f"Skipping {cell.source.name} during rendering: {error}")
|
||||
logging.info(f"Progress {processed_count}/{len(cells)} ({processed_count / len(cells):.0%})")
|
||||
return rendered_count
|
||||
|
||||
|
||||
def create_sheet(
|
||||
input_path: Path,
|
||||
output_path: Path,
|
||||
width: int,
|
||||
target_height: int,
|
||||
gap: int,
|
||||
bg_color: str = "white",
|
||||
workers: int | None = None,
|
||||
max_rows_per_page: int | None = None,
|
||||
processor: Callable[..., bytes] = resize_with_pillow,
|
||||
) -> None:
|
||||
if width <= 0 or target_height <= 0 or gap < 0 or width <= 2 * gap:
|
||||
raise ValueError("Invalid sheet dimensions or gap")
|
||||
if output_path.suffix.lower() not in {".jpg", ".jpeg"}:
|
||||
raise ValueError("Output must have a .jpg or .jpeg extension")
|
||||
if not input_path.is_dir() and not (input_path.is_file() and is_zipfile(filename=input_path)):
|
||||
raise ValueError(f"Expected a directory or ZIP archive: {input_path}")
|
||||
if input_path.resolve() == output_path.resolve():
|
||||
raise ValueError("Output must differ from input")
|
||||
if workers is not None and workers <= 0:
|
||||
raise ValueError("Workers must be positive")
|
||||
if max_rows_per_page is not None and max_rows_per_page <= 0:
|
||||
raise ValueError("Maximum rows per page must be positive")
|
||||
background = ImageColor.getcolor(color=bg_color, mode="RGB")
|
||||
with ExitStack() as stack:
|
||||
archive = None if input_path.is_dir() else stack.enter_context(ZipFile(file=input_path))
|
||||
sources = discover_sources(input_path=input_path, output_path=output_path, archive=archive)
|
||||
rows = layout_rows(sources=sources, width=width, gap=gap, target_height=target_height)
|
||||
worker_count = min(workers if workers is not None else os.cpu_count() or 1, len(sources))
|
||||
page_size = max_rows_per_page or len(rows)
|
||||
total_rendered = 0
|
||||
page_count = (len(rows) + page_size - 1) // page_size
|
||||
for page_index, start in enumerate(range(0, len(rows), page_size), start=1):
|
||||
page_rows = rows[start : start + page_size]
|
||||
page_sources = sum(len(row.sources) for row in page_rows)
|
||||
height = sum(row.height for row in page_rows) + gap * (len(page_rows) + 1)
|
||||
page_output = output_path if page_index == 1 else output_path.with_name(
|
||||
f"{output_path.stem}_page_{page_index}{output_path.suffix}"
|
||||
)
|
||||
logging.info(
|
||||
f"Rendering page {page_index}/{page_count}: {page_sources} images in {len(page_rows)} rows "
|
||||
f"on a {width}x{height} sheet using {worker_count} workers"
|
||||
)
|
||||
started = perf_counter()
|
||||
with Image.new(mode="RGB", size=(width, height), color=background) as sheet:
|
||||
rendered_count = render_rows(
|
||||
sheet=sheet,
|
||||
rows=page_rows,
|
||||
input_path=input_path,
|
||||
workers=worker_count,
|
||||
gap=gap,
|
||||
processor=processor,
|
||||
)
|
||||
if not rendered_count:
|
||||
logging.warning(f"Skipping empty page {page_index}: no images could be decoded")
|
||||
continue
|
||||
total_rendered += rendered_count
|
||||
logging.info(f"Rendering finished in {perf_counter() - started:.1f}s; saving {page_output}")
|
||||
page_output.parent.mkdir(parents=True, exist_ok=True)
|
||||
sheet.save(fp=page_output, quality=95)
|
||||
logging.info(f"Wrote {page_output} ({rendered_count} images, {width}x{height})")
|
||||
if not total_rendered:
|
||||
raise ValueError("No usable images could be decoded")
|
||||
logging.info(f"Wrote {total_rendered} images across {page_count} page(s)")
|
||||
|
||||
|
||||
def main() -> None:
|
||||
arguments = parse_args()
|
||||
input_paths = []
|
||||
for path in arguments.inputs:
|
||||
if path.exists():
|
||||
input_paths.append(path.resolve())
|
||||
else:
|
||||
logging.warning(f"Skipping non-existing input: {path}")
|
||||
if not input_paths:
|
||||
logging.error("No existing input paths")
|
||||
sys.exit(1)
|
||||
if arguments.output is not None and len(input_paths) > 1:
|
||||
logging.error("--output requires a single existing input")
|
||||
sys.exit(2)
|
||||
vips_binary = shutil.which(cmd="vips")
|
||||
processor = partial(resize_with_vips, vips_binary=vips_binary) if vips_binary else resize_with_pillow
|
||||
logging.info(f"Resize backend: {vips_binary or 'Pillow'}")
|
||||
failed_count = 0
|
||||
for index, input_path in enumerate(input_paths, start=1):
|
||||
logging.info(f"Processing input {index}/{len(input_paths)}: {input_path}")
|
||||
input_name = input_path.name if input_path.is_dir() else input_path.stem
|
||||
output_path = arguments.output or input_path.with_name(f"{input_name}_thumbs.jpg")
|
||||
try:
|
||||
create_sheet(
|
||||
input_path=input_path,
|
||||
output_path=output_path,
|
||||
width=arguments.width,
|
||||
target_height=arguments.target_height,
|
||||
gap=arguments.gap,
|
||||
bg_color=arguments.bg_color,
|
||||
workers=arguments.workers,
|
||||
max_rows_per_page=arguments.max_rows_per_page,
|
||||
processor=processor,
|
||||
)
|
||||
except (OSError, ValueError, RuntimeError, BadZipFile) as error:
|
||||
failed_count += 1
|
||||
logging.error(f"Could not create contact sheet for {input_path}: {error}")
|
||||
if failed_count:
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Executable
+476
@@ -0,0 +1,476 @@
|
||||
#!/usr/bin/env -S uv run --script
|
||||
# /// script
|
||||
# dependencies = ["pytest", "Pillow"]
|
||||
# ///
|
||||
|
||||
from functools import partial
|
||||
from io import BytesIO
|
||||
from pathlib import Path
|
||||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
from zipfile import ZIP_DEFLATED, ZipFile
|
||||
|
||||
from PIL import Image
|
||||
import pytest
|
||||
|
||||
import imageset_thumbnails as script
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def rendered_sheets(monkeypatch: pytest.MonkeyPatch) -> list[Image.Image]:
|
||||
sheets = []
|
||||
original_save = Image.Image.save
|
||||
|
||||
def capture_sheet(image: Image.Image, **keywords: object) -> None:
|
||||
if Path(keywords["fp"]).suffix == ".jpg":
|
||||
sheets.append(image.copy())
|
||||
original_save(image, **keywords)
|
||||
|
||||
monkeypatch.setattr(Image.Image, "save", capture_sheet)
|
||||
return sheets
|
||||
|
||||
|
||||
def test_directory_and_zip_natural_order(tmp_path: Path) -> None:
|
||||
directory = tmp_path / "images"
|
||||
directory.mkdir()
|
||||
for name, color in [("prefix_10.bmp", "blue"), ("prefix_2.bmp", "green"), ("prefix_1.bmp", "red"), ("nested_10/prefix_1.bmp", "yellow"), ("nested_2/prefix_1.bmp", "black")]:
|
||||
path = directory / name
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
with Image.new(mode="RGB", size=(20, 20), color=color) as image:
|
||||
image.save(fp=path)
|
||||
(directory / "broken.bmp").write_bytes(data=b"invalid")
|
||||
(directory / "notes.txt").write_text(data="not an image")
|
||||
archive_path = tmp_path / "images.zip"
|
||||
with ZipFile(file=archive_path, mode="w") as archive:
|
||||
for path in sorted(directory.rglob("*"), reverse=True):
|
||||
if path.is_file():
|
||||
archive.write(filename=path, arcname=path.relative_to(directory).as_posix())
|
||||
output = tmp_path / "directory.jpg"
|
||||
sources = script.discover_sources(input_path=directory, output_path=output, archive=None)
|
||||
assert [source.name for source in sources] == ["nested_2/prefix_1.bmp", "nested_10/prefix_1.bmp", "prefix_1.bmp", "prefix_2.bmp", "prefix_10.bmp"]
|
||||
script.create_sheet(input_path=directory, output_path=output, width=68, target_height=20, gap=2)
|
||||
zip_output = tmp_path / "zip.jpg"
|
||||
script.create_sheet(input_path=archive_path, output_path=zip_output, width=68, target_height=20, gap=2)
|
||||
with Image.open(fp=output) as first, Image.open(fp=zip_output) as second:
|
||||
assert first.size == second.size
|
||||
assert first.tobytes() == second.tobytes()
|
||||
assert first.getpixel(xy=(10, 10)) == pytest.approx((0, 0, 0), abs=5)
|
||||
assert first.getpixel(xy=(32, 10)) == pytest.approx((255, 255, 0), abs=5)
|
||||
assert first.getpixel(xy=(54, 10)) == pytest.approx((255, 0, 0), abs=5)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("ratios,width,target,gap", [
|
||||
([1.0] * 7, 100, 30, 2),
|
||||
([0.5, 2.0, 1.0, 0.75, 1.5], 100, 25, 3),
|
||||
([100.0], 100, 30, 2),
|
||||
([1.0] * 5, 90, 30, 0),
|
||||
([0.001, 100.0, 0.001], 100, 30, 1),
|
||||
([0.5] * 4, 3, 20, 1),
|
||||
])
|
||||
def test_layout_geometry(ratios: list[float], width: int, target: int, gap: int) -> None:
|
||||
sources = [script.Source(name=str(index), ratio=ratio) for index, ratio in enumerate(ratios)]
|
||||
rows = script.layout_rows(sources=sources, width=width, target_height=target, gap=gap)
|
||||
assert [source for row in rows for source in row.sources] == sources
|
||||
for index, row in enumerate(rows):
|
||||
assert row.height >= 1
|
||||
assert all(cell_width >= 1 for cell_width in row.widths)
|
||||
occupied = sum(row.widths) + gap * (len(row.sources) + 1)
|
||||
assert occupied <= width
|
||||
if index < len(rows) - 1:
|
||||
assert occupied == width
|
||||
|
||||
|
||||
def test_render_gaps_equal_heights_and_last_row(tmp_path: Path, rendered_sheets: list[Image.Image]) -> None:
|
||||
directory = tmp_path / "images"
|
||||
directory.mkdir()
|
||||
for index in range(4):
|
||||
with Image.new(mode="RGB", size=(40, 40), color="red") as image:
|
||||
image.save(fp=directory / f"image_{index}.bmp")
|
||||
output = tmp_path / "sheet.jpg"
|
||||
script.create_sheet(input_path=directory, output_path=output, width=100, target_height=30, gap=2)
|
||||
with rendered_sheets[0] as sheet:
|
||||
assert sheet.size == (100, 67)
|
||||
for left in (2, 35, 67):
|
||||
assert sheet.getpixel(xy=(left, 2)) == (255, 0, 0)
|
||||
assert sheet.getpixel(xy=(left, 32)) == (255, 0, 0)
|
||||
for position in [(0, 0), (33, 10), (65, 10), (99, 10), (10, 33), (32, 40)]:
|
||||
assert sheet.getpixel(xy=position) == (255, 255, 255)
|
||||
assert sheet.getpixel(xy=(2, 35)) == (255, 0, 0)
|
||||
|
||||
|
||||
def test_orientation_and_transparency(tmp_path: Path) -> None:
|
||||
path = tmp_path / "oriented.jpg"
|
||||
with Image.new(mode="RGB", size=(40, 20), color="red") as image:
|
||||
exif = Image.Exif()
|
||||
exif[274] = 6
|
||||
image.save(fp=path, exif=exif)
|
||||
with script.read_image(source=script.Source(name=path.name, path=path), archive=None) as image:
|
||||
assert image.size == (20, 40)
|
||||
path = tmp_path / "transparent.gif"
|
||||
with Image.new(mode="RGBA", size=(10, 10), color=(255, 0, 0, 0)) as image:
|
||||
image.save(fp=path)
|
||||
with script.read_image(source=script.Source(name=path.name, path=path), archive=None) as image:
|
||||
assert image.mode == "RGBA"
|
||||
assert image.getpixel(xy=(0, 0)) == (255, 0, 0, 0)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("bg_color,expected", [("white", (255, 255, 255)), ("black", (0, 0, 0)), ("#123456", (18, 52, 86))])
|
||||
@pytest.mark.parametrize("use_vips", [False, True])
|
||||
def test_sheet_background(
|
||||
tmp_path: Path, rendered_sheets: list[Image.Image], bg_color: str,
|
||||
expected: tuple[int, int, int], use_vips: bool,
|
||||
) -> None:
|
||||
vips_binary = shutil.which(cmd="vips") if use_vips else None
|
||||
if use_vips and vips_binary is None:
|
||||
pytest.skip("vips is not installed")
|
||||
processor = partial(script.resize_with_vips, vips_binary=vips_binary) if vips_binary else script.resize_with_pillow
|
||||
directory = tmp_path / "images"
|
||||
directory.mkdir()
|
||||
with Image.new(mode="RGBA", size=(20, 20), color=(255, 0, 0, 0)) as image:
|
||||
image.save(fp=directory / "transparent.gif")
|
||||
with Image.new(mode="RGB", size=(20, 20), color="red") as image:
|
||||
image.save(fp=directory / "red.bmp")
|
||||
output = tmp_path / "sheet.jpg"
|
||||
script.create_sheet(input_path=directory, output_path=output, width=100, target_height=20, gap=2, bg_color=bg_color, processor=processor)
|
||||
with rendered_sheets[0] as sheet:
|
||||
assert sheet.getpixel(xy=(0, 0)) == expected
|
||||
assert sheet.getpixel(xy=(22, 5)) == expected
|
||||
assert sheet.getpixel(xy=(25, 5)) == expected
|
||||
assert sheet.getpixel(xy=(99, 5)) == expected
|
||||
assert sheet.getpixel(xy=(3, 5)) == (255, 0, 0)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("extension", ["jpg"])
|
||||
@pytest.mark.parametrize("orientation", [1, 6, 8])
|
||||
def test_metadata_dimensions_without_pixel_decoding(
|
||||
tmp_path: Path, monkeypatch: pytest.MonkeyPatch,
|
||||
extension: str, orientation: int,
|
||||
) -> None:
|
||||
directory = tmp_path / "images"
|
||||
directory.mkdir()
|
||||
path = directory / f"image.{extension}"
|
||||
with Image.new(mode="RGB", size=(40, 20), color="red") as image:
|
||||
exif = Image.Exif()
|
||||
exif[274] = orientation
|
||||
image.save(fp=path, exif=exif)
|
||||
archive_path = tmp_path / "images.zip"
|
||||
with ZipFile(file=archive_path, mode="w") as archive:
|
||||
archive.write(filename=path, arcname=path.name)
|
||||
|
||||
def reject_decoding(image: Image.Image, *arguments: object, **keywords: object) -> None:
|
||||
pytest.fail("Metadata discovery must not decode image pixels")
|
||||
|
||||
Image.init()
|
||||
monkeypatch.setattr(Image.Image, "load", reject_decoding)
|
||||
expected = (40, 20) if orientation == 1 else (20, 40)
|
||||
assert script.read_dims(image_path=path) == expected
|
||||
sources = script.discover_sources(input_path=directory, output_path=tmp_path / "sheet.jpg", archive=None)
|
||||
assert sources[0].ratio == expected[0] / expected[1]
|
||||
with ZipFile(file=archive_path) as archive:
|
||||
with archive.open(name=path.name) as stream:
|
||||
assert script.read_dims(image_path=stream) == expected
|
||||
sources = script.discover_sources(input_path=archive_path, output_path=tmp_path / "sheet.jpg", archive=archive)
|
||||
assert sources[0].ratio == expected[0] / expected[1]
|
||||
|
||||
|
||||
def test_rendering_skips_pixel_corruption(tmp_path: Path, caplog: pytest.LogCaptureFixture) -> None:
|
||||
directory = tmp_path / "images"
|
||||
directory.mkdir()
|
||||
valid = directory / "1.bmp"
|
||||
broken = directory / "2.jpg"
|
||||
with Image.new(mode="RGB", size=(20, 20), color="red") as image:
|
||||
image.save(fp=valid)
|
||||
image.save(fp=broken)
|
||||
broken.write_bytes(data=broken.read_bytes()[:-10])
|
||||
assert script.read_dims(image_path=broken) == (20, 20)
|
||||
output = tmp_path / "sheet.jpg"
|
||||
script.create_sheet(input_path=directory, output_path=output, width=100, target_height=20, gap=0)
|
||||
assert "Skipping 2.jpg during rendering" in caplog.text
|
||||
with Image.open(fp=output) as sheet:
|
||||
assert sheet.getpixel(xy=(10, 10)) == pytest.approx((255, 0, 0), abs=5)
|
||||
valid.unlink()
|
||||
with pytest.raises(ValueError, match="No usable images could be decoded"):
|
||||
script.create_sheet(input_path=directory, output_path=tmp_path / "failed.jpg", width=100, target_height=20, gap=0)
|
||||
assert not (tmp_path / "failed.jpg").exists()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("extension,expected_format", [("jpg", "JPEG"), ("jpeg", "JPEG")])
|
||||
def test_output_format_and_exclusion(tmp_path: Path, extension: str, expected_format: str) -> None:
|
||||
with Image.new(mode="RGB", size=(20, 20), color="blue") as image:
|
||||
image.save(fp=tmp_path / "input.bmp")
|
||||
output = tmp_path / f"sheet.{extension}"
|
||||
script.create_sheet(input_path=tmp_path, output_path=output, width=100, target_height=20, gap=0)
|
||||
sources = script.discover_sources(input_path=tmp_path, output_path=output, archive=None)
|
||||
assert [source.name for source in sources] == ["input.bmp"]
|
||||
with Image.open(fp=output) as sheet:
|
||||
assert sheet.format == expected_format
|
||||
assert sheet.size == (100, 20)
|
||||
|
||||
|
||||
def test_first_animation_frame(tmp_path: Path) -> None:
|
||||
path = tmp_path / "animated.gif"
|
||||
with Image.new(mode="RGB", size=(10, 10), color="red") as first, Image.new(mode="RGB", size=(10, 10), color="blue") as second:
|
||||
first.save(fp=path, save_all=True, append_images=[second])
|
||||
with script.read_image(source=script.Source(name=path.name, path=path), archive=None) as image:
|
||||
assert image.getpixel(xy=(0, 0)) == (255, 0, 0, 255)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("input_kind", ["directory", "zip"])
|
||||
def test_cli_process_pool_matches_single_worker(tmp_path: Path, input_kind: str) -> None:
|
||||
directory = tmp_path / "images"
|
||||
directory.mkdir()
|
||||
for index in [10, 2, 1, 3, 4, 5, 6]:
|
||||
with Image.new(mode="RGBA", size=(20 + index * 5, 40), color=(index * 20, 100, 50, 128)) as image:
|
||||
exif = Image.Exif()
|
||||
exif[274] = 6 if index == 2 else 1
|
||||
image.convert(mode="RGB").save(fp=directory / f"image_{index}.jpg", exif=exif)
|
||||
input_path = directory
|
||||
if input_kind == "zip":
|
||||
input_path = tmp_path / "images.zip"
|
||||
with ZipFile(file=input_path, mode="w", compression=ZIP_DEFLATED) as archive:
|
||||
for path in directory.iterdir():
|
||||
archive.write(filename=path, arcname=f"nested/{path.name}")
|
||||
outputs = []
|
||||
for workers in [1, 2]:
|
||||
output = tmp_path / f"workers_{workers}.jpg"
|
||||
result = subprocess.run(
|
||||
args=[sys.executable, str(Path(script.__file__).resolve()), str(input_path),
|
||||
"--output", str(output), "--workers", str(workers), "--width", "180",
|
||||
"--target-height", "40", "--gap", "3", "--bg-color", "#123456"],
|
||||
check=True, capture_output=True, text=True, timeout=30,
|
||||
)
|
||||
assert "Progress 7/7 (100%)" in result.stderr
|
||||
with Image.open(fp=output) as image:
|
||||
outputs.append((image.size, image.tobytes()))
|
||||
assert outputs[0] == outputs[1]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("options", [
|
||||
["--width", "0"], ["--target-height", "0"], ["--gap", "-1"],
|
||||
["--width", "16", "--gap", "8"],
|
||||
["--bg-color", "invalid-color"],
|
||||
["--workers", "0"], ["--workers", "-1"],
|
||||
])
|
||||
def test_invalid_arguments(monkeypatch: pytest.MonkeyPatch, options: list[str]) -> None:
|
||||
monkeypatch.setattr(sys, "argv", ["imageset_thumbnails.py", "images", *options])
|
||||
with pytest.raises(SystemExit) as error:
|
||||
script.parse_args()
|
||||
assert error.value.code == 2
|
||||
|
||||
|
||||
def test_cli_names(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(sys, "argv", ["imageset_thumbnails.py", "images", "--target-height", "180", "--gap", "4", "--bg-color", "#123456", "--workers", "3"])
|
||||
arguments = script.parse_args()
|
||||
assert arguments.target_height == 180
|
||||
assert arguments.gap == 4
|
||||
assert arguments.bg_color == "#123456"
|
||||
assert arguments.workers == 3
|
||||
assert arguments.inputs == [Path("images")]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("bad_input_kind,expected_status", [("missing", 0), ("invalid", 1)])
|
||||
def test_cli_multiple_inputs(tmp_path: Path, bad_input_kind: str, expected_status: int) -> None:
|
||||
directory = tmp_path / "photos"
|
||||
directory.mkdir()
|
||||
with Image.new(mode="RGB", size=(20, 20), color="red") as image:
|
||||
image.save(fp=directory / "image.bmp")
|
||||
archive_path = tmp_path / "archive.zip"
|
||||
with ZipFile(file=archive_path, mode="w", compression=ZIP_DEFLATED) as archive:
|
||||
archive.write(filename=directory / "image.bmp", arcname="image.bmp")
|
||||
bad_input = tmp_path / "bad_input"
|
||||
if bad_input_kind == "invalid":
|
||||
bad_input.write_text(data="not a ZIP")
|
||||
result = subprocess.run(
|
||||
args=[sys.executable, str(Path(script.__file__).resolve()), str(bad_input),
|
||||
str(directory), str(archive_path), "--workers", "2"],
|
||||
check=False, capture_output=True, text=True, timeout=30,
|
||||
)
|
||||
assert result.returncode == expected_status
|
||||
if bad_input_kind == "missing":
|
||||
assert "WARNING: Skipping non-existing input" in result.stderr
|
||||
else:
|
||||
assert "Could not create contact sheet for" in result.stderr
|
||||
for output in [tmp_path / "photos_thumbs.jpg", tmp_path / "archive_thumbs.jpg"]:
|
||||
with Image.open(fp=output) as sheet:
|
||||
assert sheet.size == (2500, 316)
|
||||
|
||||
|
||||
def test_cli_all_inputs_missing(tmp_path: Path) -> None:
|
||||
result = subprocess.run(
|
||||
args=[sys.executable, str(Path(script.__file__).resolve()),
|
||||
str(tmp_path / "missing.zip"), str(tmp_path / "missing_directory")],
|
||||
check=False, capture_output=True, text=True, timeout=30,
|
||||
)
|
||||
assert result.returncode == 1
|
||||
assert result.stderr.count("Skipping non-existing input") == 2
|
||||
assert "No existing input paths" in result.stderr
|
||||
assert list(tmp_path.iterdir()) == []
|
||||
|
||||
|
||||
def test_cli_output_requires_single_existing_input(tmp_path: Path) -> None:
|
||||
first = tmp_path / "first"
|
||||
second = tmp_path / "second"
|
||||
first.mkdir()
|
||||
second.mkdir()
|
||||
output = tmp_path / "sheet.jpg"
|
||||
result = subprocess.run(
|
||||
args=[sys.executable, str(Path(script.__file__).resolve()), str(first),
|
||||
str(second), "--output", str(output)],
|
||||
check=False, capture_output=True, text=True, timeout=30,
|
||||
)
|
||||
assert result.returncode == 2
|
||||
assert "--output requires a single existing input" in result.stderr
|
||||
assert not output.exists()
|
||||
with Image.new(mode="RGB", size=(20, 20), color="blue") as image:
|
||||
image.save(fp=first / "image.bmp")
|
||||
result = subprocess.run(
|
||||
args=[sys.executable, str(Path(script.__file__).resolve()), str(first),
|
||||
str(tmp_path / "missing"), "--output", str(output)],
|
||||
check=True, capture_output=True, text=True, timeout=30,
|
||||
)
|
||||
assert "Skipping non-existing input" in result.stderr
|
||||
assert output.exists()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("use_vips", [False, True])
|
||||
def test_main_selects_processor_once(
|
||||
tmp_path: Path, monkeypatch: pytest.MonkeyPatch, use_vips: bool,
|
||||
) -> None:
|
||||
vips_binary = shutil.which(cmd="vips") if use_vips else None
|
||||
if use_vips and vips_binary is None:
|
||||
pytest.skip("vips is not installed")
|
||||
directories = [tmp_path / "first", tmp_path / "second"]
|
||||
for directory in directories:
|
||||
directory.mkdir()
|
||||
with Image.new(mode="RGB", size=(40, 20), color="red") as image:
|
||||
image.save(fp=directory / "image_1.bmp")
|
||||
image.save(fp=directory / "image_2.bmp")
|
||||
lookups = []
|
||||
|
||||
def find_binary(cmd: str) -> str | None:
|
||||
lookups.append(cmd)
|
||||
return vips_binary
|
||||
|
||||
monkeypatch.setattr(script.shutil, "which", find_binary)
|
||||
monkeypatch.setattr(sys, "argv", ["imageset_thumbnails.py", *map(str, directories),
|
||||
"--width", "100", "--target-height", "20", "--workers", "2"])
|
||||
script.main()
|
||||
assert lookups == ["vips"]
|
||||
for name in ["first", "second"]:
|
||||
with Image.open(fp=tmp_path / f"{name}_thumbs.jpg") as sheet:
|
||||
assert sheet.width == 100
|
||||
|
||||
|
||||
@pytest.mark.parametrize("input_kind", ["directory", "zip"])
|
||||
def test_vips_resizes_with_legacy_options_before_pillow_reads_pixels(
|
||||
tmp_path: Path, monkeypatch: pytest.MonkeyPatch, input_kind: str,
|
||||
) -> None:
|
||||
vips_binary = shutil.which(cmd="vips")
|
||||
if vips_binary is None:
|
||||
pytest.skip("vips is not installed")
|
||||
path = tmp_path / "image.jpg"
|
||||
with Image.new(mode="RGB", size=(400, 200), color="red") as image:
|
||||
image.paste(im="blue", box=(200, 0, 400, 200))
|
||||
exif = Image.Exif()
|
||||
exif[274] = 6
|
||||
image.save(fp=path, exif=exif)
|
||||
cell = script.Cell(source=script.Source(name=path.name, path=path), position=(0, 0), size=(10, 20))
|
||||
archive_path = tmp_path / "images.zip"
|
||||
with ZipFile(file=archive_path, mode="w") as archive:
|
||||
archive.write(filename=path, arcname=path.name)
|
||||
original_open = Image.open
|
||||
original_run = subprocess.run
|
||||
reads = []
|
||||
|
||||
def run_legacy_vips(args: list[str], **keywords: object) -> subprocess.CompletedProcess[str]:
|
||||
assert ".png" in args
|
||||
if input_kind == "zip":
|
||||
assert "thumbnail_source" in args
|
||||
assert keywords["input"] == path.read_bytes()
|
||||
else:
|
||||
assert "thumbnail" in args
|
||||
assert keywords["input"] is None
|
||||
if "--fail-on" in args:
|
||||
raise subprocess.CalledProcessError(returncode=1, cmd=args, stderr=b"Unknown option --fail-on")
|
||||
return original_run(args=args, **keywords)
|
||||
|
||||
def read_thumbnail(fp: BytesIO) -> Image.Image:
|
||||
assert isinstance(fp, BytesIO)
|
||||
image = original_open(fp=fp)
|
||||
reads.append(image.size)
|
||||
return image
|
||||
|
||||
monkeypatch.setattr(script.Image, "open", read_thumbnail)
|
||||
monkeypatch.setattr(script.subprocess, "run", run_legacy_vips)
|
||||
with ZipFile(file=archive_path) as archive:
|
||||
if input_kind == "zip":
|
||||
cell.source = script.Source(name=path.name, member=archive.getinfo(name=path.name))
|
||||
monkeypatch.setattr(script, "worker_archive", archive)
|
||||
pixels = script.resize_with_vips(cell=cell, vips_binary=vips_binary)
|
||||
assert reads == [cell.size]
|
||||
with Image.frombytes(mode="RGBA", size=cell.size, data=pixels) as thumbnail:
|
||||
assert thumbnail.getpixel(xy=(5, 2)) == pytest.approx((255, 0, 0, 255), abs=5)
|
||||
assert thumbnail.getpixel(xy=(5, 17)) == pytest.approx((0, 0, 255, 255), abs=5)
|
||||
|
||||
|
||||
def test_vips_failure_falls_back_to_pillow(
|
||||
tmp_path: Path, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture,
|
||||
) -> None:
|
||||
path = tmp_path / "image.gif"
|
||||
with Image.new(mode="RGBA", size=(40, 20), color=(255, 0, 0, 128)) as image:
|
||||
image.save(fp=path)
|
||||
cell = script.Cell(source=script.Source(name=path.name, path=path), position=(0, 0), size=(10, 5))
|
||||
expected = script.resize_with_pillow(cell=cell)
|
||||
|
||||
def fail_command(args: list[str], **keywords: object) -> subprocess.CompletedProcess[str]:
|
||||
raise subprocess.CalledProcessError(returncode=1, cmd=args, stderr=b"unsupported format")
|
||||
|
||||
monkeypatch.setattr(script.subprocess, "run", fail_command)
|
||||
assert script.resize_with_vips(cell=cell, vips_binary="vips") == expected
|
||||
assert "using Pillow: unsupported format" in caplog.text
|
||||
|
||||
|
||||
@pytest.mark.parametrize("input_kind", ["directory", "zip"])
|
||||
@pytest.mark.parametrize("image_format,format_name", [("PNG", "PNG"), ("WEBP", "WebP"), ("TIFF", "TIFF")])
|
||||
def test_unsupported_inputs_are_skipped(tmp_path: Path, input_kind: str, image_format: str, format_name: str) -> None:
|
||||
directory = tmp_path / "images"
|
||||
directory.mkdir()
|
||||
with Image.new(mode="RGB", size=(20, 20), color="red") as image:
|
||||
image.save(fp=directory / "supported.jpg")
|
||||
image.save(fp=directory / f"ignored.{image_format}", format=image_format)
|
||||
image.save(fp=directory / "disguised.jpg", format=image_format)
|
||||
output = tmp_path / "sheet.jpg"
|
||||
with pytest.raises(ValueError, match=f"{format_name} images are not supported"):
|
||||
script.read_dims(image_path=directory / f"ignored.{image_format}")
|
||||
if input_kind == "zip":
|
||||
archive_path = tmp_path / "images.zip"
|
||||
with ZipFile(file=archive_path, mode="w") as archive:
|
||||
for path in directory.iterdir():
|
||||
archive.write(filename=path, arcname=path.name)
|
||||
with ZipFile(file=archive_path) as archive:
|
||||
sources = script.discover_sources(input_path=archive_path, output_path=output, archive=archive)
|
||||
else:
|
||||
sources = script.discover_sources(input_path=directory, output_path=output, archive=None)
|
||||
assert [source.name for source in sources] == ["supported.jpg"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("extension", ["png", "webp", "tif", "tiff"])
|
||||
def test_unsupported_output_is_rejected(tmp_path: Path, extension: str) -> None:
|
||||
output = tmp_path / f"sheet.{extension}"
|
||||
with pytest.raises(ValueError, match="Output must have a .jpg or .jpeg extension"):
|
||||
script.create_sheet(input_path=tmp_path, output_path=output, width=100, target_height=20, gap=0)
|
||||
assert not output.exists()
|
||||
|
||||
|
||||
def test_empty_and_invalid_input(tmp_path: Path) -> None:
|
||||
output = tmp_path / "sheet.jpg"
|
||||
with pytest.raises(ValueError, match="No usable images"):
|
||||
script.create_sheet(input_path=tmp_path, output_path=output, width=100, target_height=20, gap=0)
|
||||
with pytest.raises(ValueError, match="Expected a directory or ZIP"):
|
||||
script.create_sheet(input_path=tmp_path / "missing", output_path=output, width=100, target_height=20, gap=0)
|
||||
assert not output.exists()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main(args=[str(Path(__file__).resolve()), *sys.argv[1:]]))
|
||||
Executable
+294
@@ -0,0 +1,294 @@
|
||||
#!/usr/bin/env -S uv run --script
|
||||
# /// script
|
||||
# dependencies = ["httpx>=0.27,<1"]
|
||||
# requires-python = ">=3.10"
|
||||
# ///
|
||||
"""Upload images to Immich and derive albums and tags from their filenames."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import hashlib
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
import shutil
|
||||
import struct
|
||||
import sys
|
||||
import threading
|
||||
import uuid
|
||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from typing import Any, Sequence
|
||||
|
||||
import httpx
|
||||
|
||||
|
||||
LOG = logging.getLogger("immichup")
|
||||
TAG_PORTRAIT = "aspect:portrait"
|
||||
TAG_LANDSCAPE = "aspect:landscape"
|
||||
STUDIO_RE = re.compile(r"\b@(\S+)\b")
|
||||
|
||||
|
||||
def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Upload JPEG images to Immich; filename prefixes become albums and tags."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--api-key", default=os.getenv("IMMICH_API_KEY", ""),
|
||||
help="Immich API key (defaults to IMMICH_API_KEY)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--server-url", default=os.getenv("IMMICH_SERVER_URL", ""),
|
||||
help="Immich server URL (defaults to IMMICH_SERVER_URL)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--move-uploaded-to", metavar="DIR",
|
||||
help="move successfully uploaded images into this directory",
|
||||
)
|
||||
parser.add_argument("--debug", action="store_true", help="enable debug logging")
|
||||
parser.add_argument(
|
||||
"--dry-run", action="store_true", help="show planned uploads without uploading",
|
||||
)
|
||||
parser.add_argument("images", nargs="+", type=Path, help="image paths to process")
|
||||
args = parser.parse_args(argv)
|
||||
if not args.dry_run and not args.server_url:
|
||||
parser.error("--server-url or IMMICH_SERVER_URL is required unless --dry-run is used")
|
||||
if not args.dry_run and not args.api_key:
|
||||
parser.error("--api-key or IMMICH_API_KEY is required unless --dry-run is used")
|
||||
return args
|
||||
|
||||
|
||||
def parse_filename(image_path: Path) -> tuple[list[str], str | None] | None:
|
||||
filename = image_path.stem
|
||||
if " -- " not in filename:
|
||||
return None
|
||||
people = [name.strip() for name in filename.split(" -- ", 1)[0].split(",")]
|
||||
if not people:
|
||||
return None
|
||||
match = STUDIO_RE.search(filename)
|
||||
return people, match.group(1) if match else None
|
||||
|
||||
|
||||
def read_jpeg_dimensions(path: Path) -> tuple[int, int]:
|
||||
"""Read width and height from a JPEG frame header without loading the image."""
|
||||
with path.open("rb") as image:
|
||||
if image.read(2) != b"\xff\xd8":
|
||||
raise ValueError("not a valid JPEG file")
|
||||
while True:
|
||||
byte = image.read(1)
|
||||
if not byte:
|
||||
raise ValueError("JPEG has no frame header")
|
||||
if byte != b"\xff":
|
||||
raise ValueError("invalid JPEG format")
|
||||
marker_byte = image.read(1)
|
||||
while marker_byte == b"\xff":
|
||||
marker_byte = image.read(1)
|
||||
if not marker_byte:
|
||||
raise ValueError("truncated JPEG marker")
|
||||
marker = marker_byte[0]
|
||||
if marker in (0xD8, 0xD9) or 0xD0 <= marker <= 0xD7 or marker == 0x01:
|
||||
continue
|
||||
length_data = image.read(2)
|
||||
if len(length_data) != 2:
|
||||
raise ValueError("truncated JPEG segment")
|
||||
length = struct.unpack(">H", length_data)[0]
|
||||
if length < 2:
|
||||
raise ValueError("invalid JPEG segment length")
|
||||
if 0xC0 <= marker <= 0xC3:
|
||||
frame = image.read(length - 2)
|
||||
if len(frame) < 5:
|
||||
raise ValueError("truncated JPEG frame header")
|
||||
height, width = struct.unpack(">HH", frame[1:5])
|
||||
return width, height
|
||||
image.seek(length - 2, os.SEEK_CUR)
|
||||
|
||||
|
||||
class ImmichClient:
|
||||
def __init__(self, server_url: str, api_key: str) -> None:
|
||||
self.http = httpx.Client(
|
||||
base_url=server_url.rstrip("/") + "/api/",
|
||||
headers={"x-api-key": api_key},
|
||||
timeout=60.0,
|
||||
)
|
||||
self.device_id = str(uuid.uuid4())
|
||||
self.album_ids: dict[str, str] = {}
|
||||
self.tag_ids: dict[str, str] = {}
|
||||
self.metadata_lock = threading.RLock()
|
||||
|
||||
def close(self) -> None:
|
||||
self.http.close()
|
||||
|
||||
def _check(self, response: httpx.Response, action: str) -> None:
|
||||
if not response.is_success:
|
||||
raise RuntimeError(f"{action}: HTTP {response.status_code} {response.reason_phrase}")
|
||||
|
||||
def upload_asset(self, path: Path) -> dict[str, Any]:
|
||||
stat = path.stat()
|
||||
with path.open("rb") as image:
|
||||
digest = hashlib.sha1()
|
||||
for chunk in iter(lambda: image.read(1024 * 1024), b""):
|
||||
digest.update(chunk)
|
||||
checksum = digest.hexdigest()
|
||||
modified = datetime.fromtimestamp(stat.st_mtime).astimezone().isoformat()
|
||||
form = {
|
||||
"deviceAssetId": f"{path.name}-{stat.st_size}",
|
||||
"deviceId": self.device_id,
|
||||
"assetType": "image",
|
||||
"fileCreatedAt": modified,
|
||||
"fileModifiedAt": modified,
|
||||
"isFavorite": "false",
|
||||
"fileExtension": path.suffix.lower(),
|
||||
"duration": "0",
|
||||
"isReadOnly": "false",
|
||||
"isArchived": "false",
|
||||
}
|
||||
with path.open("rb") as image:
|
||||
response = self.http.post(
|
||||
"assets",
|
||||
headers={"x-immich-checksum": checksum},
|
||||
data=form,
|
||||
files={"assetData": (path.name, image)},
|
||||
)
|
||||
self._check(response, "failed to upload asset")
|
||||
return response.json()
|
||||
|
||||
def get_or_create_album(self, name: str) -> str:
|
||||
with self.metadata_lock:
|
||||
if name in self.album_ids:
|
||||
return self.album_ids[name]
|
||||
response = self.http.get("albums")
|
||||
self._check(response, "failed to get albums")
|
||||
for album in response.json():
|
||||
self.album_ids[album["albumName"]] = album["id"]
|
||||
if name not in self.album_ids:
|
||||
response = self.http.post("albums", json={"albumName": name})
|
||||
self._check(response, "failed to create album")
|
||||
album = response.json()
|
||||
self.album_ids[name] = album["id"]
|
||||
return self.album_ids[name]
|
||||
|
||||
def create_tags(self, names: list[str]) -> dict[str, str]:
|
||||
with self.metadata_lock:
|
||||
missing = list(dict.fromkeys(name for name in names if name not in self.tag_ids))
|
||||
if missing:
|
||||
response = self.http.put("tags", json={"tags": missing})
|
||||
self._check(response, "failed to create tags")
|
||||
self.tag_ids.update({tag["name"]: tag["id"] for tag in response.json()})
|
||||
return {name: self.tag_ids[name] for name in names if name in self.tag_ids}
|
||||
|
||||
def add_asset_to_album(self, album_id: str, asset_id: str) -> None:
|
||||
response = self.http.put(f"albums/{album_id}/assets", json={"ids": [asset_id]})
|
||||
self._check(response, "failed to add asset to album")
|
||||
|
||||
def link_tag_to_asset(self, tag_id: str, asset_id: str) -> None:
|
||||
response = self.http.put(f"tags/{tag_id}/assets", json={"ids": [asset_id]})
|
||||
self._check(response, "failed to link tag to asset")
|
||||
|
||||
def upload_image(self, path: Path, albums: list[str], tags: list[str]) -> str:
|
||||
asset = self.upload_asset(path)
|
||||
asset_id = asset["id"]
|
||||
for album in albums:
|
||||
try:
|
||||
self.add_asset_to_album(self.get_or_create_album(album), asset_id)
|
||||
except (httpx.HTTPError, RuntimeError, KeyError, ValueError) as error:
|
||||
LOG.error("Could not add %s to album %r: %s", path, album, error)
|
||||
try:
|
||||
tag_ids = self.create_tags(tags)
|
||||
for tag, tag_id in tag_ids.items():
|
||||
try:
|
||||
self.link_tag_to_asset(tag_id, asset_id)
|
||||
except (httpx.HTTPError, RuntimeError) as error:
|
||||
LOG.error("Could not apply tag %r to %s: %s", tag, path, error)
|
||||
except (httpx.HTTPError, RuntimeError, KeyError, ValueError) as error:
|
||||
LOG.error("Could not create tags for %s: %s", path, error)
|
||||
return asset.get("status") or "uploaded"
|
||||
|
||||
|
||||
def process_image(path: Path, client: ImmichClient | None, dry_run: bool) -> bool:
|
||||
try:
|
||||
parsed = parse_filename(path)
|
||||
if parsed is None:
|
||||
raise ValueError(f"failed to parse filename: {path}")
|
||||
people, studio = parsed
|
||||
width, height = read_jpeg_dimensions(path)
|
||||
albums = [person.lower() for person in people]
|
||||
tags = list(albums)
|
||||
if studio:
|
||||
tags.append(studio.lower())
|
||||
tags.append(TAG_PORTRAIT if width <= height else TAG_LANDSCAPE)
|
||||
if dry_run:
|
||||
LOG.info("[DRY RUN] Would upload %s (albums=%s, tags=%s)", path, albums, tags)
|
||||
return True
|
||||
assert client is not None
|
||||
status = client.upload_image(path, albums, tags)
|
||||
LOG.info("Processed %s (status=%s)", path, status)
|
||||
return True
|
||||
except (OSError, ValueError, RuntimeError, httpx.HTTPError, KeyError) as error:
|
||||
LOG.error("Failed to process %s: %s", path, error)
|
||||
return False
|
||||
|
||||
|
||||
def run(args: argparse.Namespace) -> int:
|
||||
logging.basicConfig(
|
||||
level=logging.DEBUG if args.debug else logging.INFO,
|
||||
format="%(levelname)s: %(message)s",
|
||||
)
|
||||
paths: list[Path] = []
|
||||
failures = 0
|
||||
for path in args.images:
|
||||
if not path.exists():
|
||||
LOG.error("Image does not exist, skipping: %s", path)
|
||||
failures += 1
|
||||
elif path.is_dir():
|
||||
LOG.warning("Skipping directory: %s", path)
|
||||
else:
|
||||
paths.append(path)
|
||||
LOG.info("Found %d images to process", len(paths))
|
||||
|
||||
if args.move_uploaded_to:
|
||||
try:
|
||||
Path(args.move_uploaded_to).mkdir(parents=True, exist_ok=True)
|
||||
except OSError as error:
|
||||
LOG.error("Failed to create destination directory: %s", error)
|
||||
return 1
|
||||
|
||||
client = None if args.dry_run else ImmichClient(args.server_url, args.api_key)
|
||||
try:
|
||||
with ThreadPoolExecutor(max_workers=5) as pool:
|
||||
futures = {pool.submit(process_image, path, client, args.dry_run): path for path in paths}
|
||||
for future in as_completed(futures):
|
||||
path = futures[future]
|
||||
try:
|
||||
succeeded = future.result()
|
||||
except Exception:
|
||||
LOG.exception("Unexpected error processing %s", path)
|
||||
succeeded = False
|
||||
if not succeeded:
|
||||
failures += 1
|
||||
continue
|
||||
if args.move_uploaded_to:
|
||||
destination = Path(args.move_uploaded_to) / path.name
|
||||
if args.dry_run:
|
||||
LOG.info("[DRY RUN] Would move %s to %s", path, destination.parent)
|
||||
else:
|
||||
try:
|
||||
shutil.move(str(path), str(destination))
|
||||
LOG.info("Moved %s to %s", path, destination)
|
||||
except OSError as error:
|
||||
LOG.error("Failed to move %s: %s", path, error)
|
||||
failures += 1
|
||||
finally:
|
||||
if client is not None:
|
||||
client.close()
|
||||
LOG.info("Completed processing all images")
|
||||
return 1 if failures else 0
|
||||
|
||||
|
||||
def main() -> int:
|
||||
return run(parse_args())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
+84
-28
@@ -1,6 +1,6 @@
|
||||
#!/usr/bin/env -S uv run
|
||||
#!/usr/bin/env -S uv run --script
|
||||
# /// script
|
||||
# requires-python = ">=3.8"
|
||||
# requires-python = ">=3.11"
|
||||
# dependencies = [
|
||||
# "watchdog"
|
||||
# ]
|
||||
@@ -12,20 +12,20 @@ import logging
|
||||
import os
|
||||
from pathlib import Path
|
||||
import re
|
||||
import stat
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
from typing import Self
|
||||
from watchdog import observers
|
||||
from watchdog.observers.polling import PollingObserver
|
||||
from watchdog import events
|
||||
|
||||
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s: %(message)s", datefmt="%Y-%m-%d %H:%M:%S")
|
||||
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s: %(message)s")
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class EventHandler:
|
||||
pattern: re.Pattern
|
||||
pattern: re.Pattern[str]
|
||||
script_path: Path
|
||||
|
||||
@classmethod
|
||||
@@ -51,24 +51,24 @@ class EventHandler:
|
||||
|
||||
|
||||
class ChangeDispatchHandler(events.FileSystemEventHandler):
|
||||
def __init__(self, base_dir: Path, handlers_config: dict[str, EventHandler], ignore_patterns: list[re.Pattern]):
|
||||
def __init__(self, base_dir: Path, handlers_config: dict[str, list[EventHandler]], ignore_patterns: list[re.Pattern[str]]) -> None:
|
||||
super().__init__()
|
||||
self.base_dir = base_dir
|
||||
self.handlers_config = handlers_config
|
||||
self.ignore_patterns = ignore_patterns if ignore_patterns else []
|
||||
|
||||
def _execute_script(self, script_path: str, event_name: str, src_path: str, dest_path: str = None):
|
||||
def _execute_script(self, script_path: Path, event_name: str, src_path: Path, dest_path: Path | None = None) -> None:
|
||||
env_vars = os.environ.copy()
|
||||
env_vars["MONITORED_DIR"] = str(self.base_dir)
|
||||
env_vars["EVENT_TYPE"] = event_name.upper()
|
||||
env_vars["SRC_PATH"] = str(Path(src_path).resolve())
|
||||
env_vars["SRC_PATH"] = str(src_path.absolute())
|
||||
args = [
|
||||
event_name,
|
||||
str(Path(src_path).resolve()),
|
||||
*([str(Path(dest_path).resolve())] if dest_path else []),
|
||||
str(src_path.absolute()),
|
||||
*([str(dest_path.absolute())] if dest_path else []),
|
||||
]
|
||||
if dest_path:
|
||||
env_vars["DEST_PATH"] = str(Path(dest_path).resolve())
|
||||
env_vars["DEST_PATH"] = str(dest_path.absolute())
|
||||
elif "DEST_PATH" in env_vars:
|
||||
del env_vars["DEST_PATH"]
|
||||
|
||||
@@ -93,7 +93,7 @@ class ChangeDispatchHandler(events.FileSystemEventHandler):
|
||||
logging.error(f"An unexpected error occurred while running script '{script_path}': {e}")
|
||||
|
||||
def _is_ignored(self, path_to_check: Path) -> bool:
|
||||
abs_path_to_check = path_to_check.resolve()
|
||||
abs_path_to_check = path_to_check.absolute()
|
||||
for pattern in self.ignore_patterns:
|
||||
if pattern.search(str(abs_path_to_check)):
|
||||
return True
|
||||
@@ -107,7 +107,7 @@ class ChangeDispatchHandler(events.FileSystemEventHandler):
|
||||
pass
|
||||
return False
|
||||
|
||||
def _process_event_type(self, event_name: str, path_to_filter_on: Path, src_path: Path, dest_path: Path | None = None):
|
||||
def _process_event_type(self, event_name: str, path_to_filter_on: Path, src_path: Path, dest_path: Path | None = None) -> None:
|
||||
if self._is_ignored(path_to_filter_on):
|
||||
logging.debug(f"Path '{path_to_filter_on}' is ignored by configured patterns.")
|
||||
return
|
||||
@@ -127,27 +127,27 @@ class ChangeDispatchHandler(events.FileSystemEventHandler):
|
||||
dest_path=dest_path,
|
||||
)
|
||||
|
||||
def on_created(self, event: events.FileSystemEvent):
|
||||
def on_created(self, event: events.FileSystemEvent) -> None:
|
||||
what = "directory" if event.is_directory else "file"
|
||||
logging.debug(f"Watchdog event: CREATED {what} at {event.src_path}")
|
||||
|
||||
self._process_event_type(
|
||||
event.event_type,
|
||||
event_name=event.event_type,
|
||||
path_to_filter_on=Path(event.src_path),
|
||||
src_path=Path(event.src_path),
|
||||
)
|
||||
|
||||
def on_deleted(self, event: events.FileSystemEvent):
|
||||
def on_deleted(self, event: events.FileSystemEvent) -> None:
|
||||
what = "directory" if event.is_directory else "file"
|
||||
logging.debug(f"Watchdog event: DELETED {what} at {event.src_path}")
|
||||
|
||||
self._process_event_type(
|
||||
event.event_type,
|
||||
event_name=event.event_type,
|
||||
path_to_filter_on=Path(event.src_path),
|
||||
src_path=Path(event.src_path),
|
||||
)
|
||||
|
||||
def on_moved(self, event: events.FileSystemEvent):
|
||||
def on_moved(self, event: events.FileSystemMovedEvent) -> None:
|
||||
what = "directory" if event.is_directory else "file"
|
||||
logging.debug(f"Watchdog event: MOVED {what} from {event.src_path} to {event.dest_path}")
|
||||
|
||||
@@ -155,19 +155,62 @@ class ChangeDispatchHandler(events.FileSystemEventHandler):
|
||||
dest_path = Path(event.dest_path)
|
||||
|
||||
self._process_event_type(
|
||||
event.event_type,
|
||||
event_name=event.event_type,
|
||||
path_to_filter_on=dest_path,
|
||||
src_path=src_path,
|
||||
dest_path=dest_path,
|
||||
)
|
||||
|
||||
|
||||
def parse_args():
|
||||
def take_snapshot(base_dir: Path, recursive: bool) -> dict[Path, bool]:
|
||||
if not stat.S_ISDIR(base_dir.stat().st_mode):
|
||||
raise NotADirectoryError(f"Monitored path '{base_dir}' is not a directory")
|
||||
snapshot: dict[Path, bool] = {}
|
||||
directories = [base_dir]
|
||||
while directories:
|
||||
directory = directories.pop()
|
||||
for path in directory.iterdir():
|
||||
is_directory = stat.S_ISDIR(path.lstat().st_mode)
|
||||
snapshot[path] = is_directory
|
||||
if recursive and is_directory:
|
||||
directories.append(path)
|
||||
return snapshot
|
||||
|
||||
|
||||
def dispatch_changes(previous: dict[Path, bool], current: dict[Path, bool], event_handler: ChangeDispatchHandler) -> None:
|
||||
deleted = previous.keys() - current.keys()
|
||||
created = current.keys() - previous.keys()
|
||||
replaced = {path for path in previous.keys() & current.keys() if previous[path] != current[path]}
|
||||
for path in sorted(deleted | replaced):
|
||||
event_handler._process_event_type(event_name="deleted", path_to_filter_on=path, src_path=path)
|
||||
for path in sorted(created | replaced):
|
||||
event_handler._process_event_type(event_name="created", path_to_filter_on=path, src_path=path)
|
||||
|
||||
|
||||
def monitor_paths(base_dir: Path, recursive: bool, interval: float, event_handler: ChangeDispatchHandler) -> None:
|
||||
previous: dict[Path, bool] | None = None
|
||||
while True:
|
||||
try:
|
||||
current = take_snapshot(base_dir=base_dir, recursive=recursive)
|
||||
except OSError as error:
|
||||
logging.warning(f"Directory scan failed; retaining previous snapshot and retrying in {interval}s: {error}")
|
||||
else:
|
||||
if previous is not None:
|
||||
dispatch_changes(previous=previous, current=current, event_handler=event_handler)
|
||||
previous = current
|
||||
time.sleep(interval)
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Monitors a folder for changes and executes handlers based on event type and regex filters.", formatter_class=argparse.RawTextHelpFormatter
|
||||
description="Monitors a folder for changes and executes handlers based on event type and regex filters.", formatter_class=argparse.ArgumentDefaultsHelpFormatter
|
||||
)
|
||||
parser.add_argument("dir", type=Path, help="The directory to monitor.")
|
||||
parser.add_argument("--poll", type=int, default=10, help="Polling interval in seconds.")
|
||||
parser.add_argument("--poll", type=float, default=10, help="Polling interval in seconds; must be positive.")
|
||||
parser.add_argument(
|
||||
"--polling-mode", choices=["path", "inode"], default="path",
|
||||
help="Path polling tolerates changing SSHFS inode numbers and retries failed scans. Renames emit deleted/created events. Inode polling supports moved events and requires stable inode numbers.",
|
||||
)
|
||||
parser.add_argument("--recursive", action="store_true", help="Monitor the directory recursively.")
|
||||
parser.add_argument(
|
||||
"--created",
|
||||
@@ -206,10 +249,15 @@ def parse_args():
|
||||
default=[],
|
||||
)
|
||||
parser.add_argument("--debug", action="store_true", help="Enable debug logging.")
|
||||
return parser.parse_args()
|
||||
args = parser.parse_args()
|
||||
if not 0 < args.poll < float("inf"):
|
||||
parser.error("--poll must be a finite positive number")
|
||||
if args.moved_handlers and args.polling_mode != "inode":
|
||||
parser.error("--moved requires --polling-mode inode; for SSHFS, use --created and --deleted instead")
|
||||
return args
|
||||
|
||||
|
||||
def main():
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
|
||||
if args.debug:
|
||||
@@ -220,7 +268,7 @@ def main():
|
||||
logging.error(f"Monitored directory '{args.dir}' (resolved to '{args.dir.resolve()}') not found or not a directory.")
|
||||
sys.exit(1)
|
||||
|
||||
handlers: dict[str, EventHandler] = {
|
||||
handlers: dict[str, list[EventHandler]] = {
|
||||
"created": args.created_handlers,
|
||||
"deleted": args.deleted_handlers,
|
||||
"moved": args.moved_handlers,
|
||||
@@ -235,7 +283,6 @@ def main():
|
||||
if handlers_list:
|
||||
logging.info(f"Registered {len(handlers_list)} handler(s) for '{event_type}' events.")
|
||||
for handler in handlers_list:
|
||||
handler: EventHandler
|
||||
logging.info(f" - Filter: '{handler.pattern.pattern}', Script: '{handler.script_path}'")
|
||||
|
||||
if args.ignore_patterns:
|
||||
@@ -243,13 +290,22 @@ def main():
|
||||
for pattern in args.ignore_patterns:
|
||||
logging.info(f" - {pattern.pattern}")
|
||||
|
||||
base_dir = args.dir.resolve()
|
||||
event_handler = ChangeDispatchHandler(
|
||||
base_dir=args.dir.resolve(),
|
||||
base_dir=base_dir,
|
||||
handlers_config=handlers,
|
||||
ignore_patterns=args.ignore_patterns,
|
||||
)
|
||||
if args.polling_mode == "path":
|
||||
logging.info("Monitor is now running with path polling. Press Ctrl+C to stop.")
|
||||
try:
|
||||
monitor_paths(base_dir=base_dir, recursive=args.recursive, interval=args.poll, event_handler=event_handler)
|
||||
except KeyboardInterrupt:
|
||||
logging.info("Monitor has been shut down.")
|
||||
return
|
||||
|
||||
observer = PollingObserver(timeout=args.poll)
|
||||
observer.schedule(event_handler, args.dir, recursive=args.recursive)
|
||||
observer.schedule(event_handler=event_handler, path=str(base_dir), recursive=args.recursive)
|
||||
observer.start()
|
||||
logging.info("Monitor is now running. Press Ctrl+C to stop.")
|
||||
|
||||
|
||||
@@ -0,0 +1,91 @@
|
||||
#!/usr/bin/env -S uv run --script
|
||||
# /// script
|
||||
# dependencies = ["pytest", "watchdog"]
|
||||
# ///
|
||||
|
||||
from pathlib import Path
|
||||
import stat
|
||||
import sys
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import Mock, call, patch
|
||||
|
||||
import pytest
|
||||
|
||||
import monitor_dir as script
|
||||
|
||||
|
||||
@pytest.mark.parametrize("recursive", [False, True])
|
||||
def test_recursive_scan_does_not_follow_symlinks(tmp_path: Path, recursive: bool) -> None:
|
||||
child = tmp_path / "child"
|
||||
child.mkdir()
|
||||
file_path = child / "file.txt"
|
||||
file_path.touch()
|
||||
link = tmp_path / "loop"
|
||||
link.symlink_to(target=tmp_path, target_is_directory=True)
|
||||
|
||||
expected = {child: True, link: False}
|
||||
if recursive:
|
||||
expected[file_path] = False
|
||||
assert script.take_snapshot(base_dir=tmp_path, recursive=recursive) == expected
|
||||
|
||||
|
||||
def test_changing_inode_numbers_do_not_emit_events(tmp_path: Path) -> None:
|
||||
file_path = tmp_path / "file.txt"
|
||||
file_path.touch()
|
||||
handler = Mock()
|
||||
with patch.object(
|
||||
target=Path, attribute="lstat",
|
||||
side_effect=[
|
||||
SimpleNamespace(st_mode=stat.S_IFREG, st_ino=1),
|
||||
SimpleNamespace(st_mode=stat.S_IFREG, st_ino=999),
|
||||
],
|
||||
):
|
||||
previous = script.take_snapshot(base_dir=tmp_path, recursive=False)
|
||||
current = script.take_snapshot(base_dir=tmp_path, recursive=False)
|
||||
|
||||
script.dispatch_changes(previous=previous, current=current, event_handler=handler)
|
||||
handler._process_event_type.assert_not_called()
|
||||
|
||||
|
||||
def test_rename_emits_deleted_and_created() -> None:
|
||||
source = Path("/watched/before.txt")
|
||||
destination = Path("/watched/after.txt")
|
||||
handler = Mock()
|
||||
script.dispatch_changes(previous={source: False}, current={destination: False}, event_handler=handler)
|
||||
assert handler._process_event_type.call_args_list == [
|
||||
call(event_name="deleted", path_to_filter_on=source, src_path=source),
|
||||
call(event_name="created", path_to_filter_on=destination, src_path=destination),
|
||||
]
|
||||
|
||||
|
||||
def test_failed_scan_preserves_last_snapshot() -> None:
|
||||
base_dir = Path("/watched")
|
||||
existing = base_dir / "existing.txt"
|
||||
created = base_dir / "created.txt"
|
||||
handler = Mock()
|
||||
with patch.object(
|
||||
target=script, attribute="take_snapshot",
|
||||
side_effect=[{existing: False}, OSError("SSHFS disconnected"), {existing: False, created: False}],
|
||||
), patch.object(
|
||||
target=script.time, attribute="sleep", side_effect=[None, None, KeyboardInterrupt],
|
||||
):
|
||||
with pytest.raises(KeyboardInterrupt):
|
||||
script.monitor_paths(base_dir=base_dir, recursive=True, interval=1, event_handler=handler)
|
||||
handler._process_event_type.assert_called_once_with(
|
||||
event_name="created", path_to_filter_on=created, src_path=created,
|
||||
)
|
||||
|
||||
|
||||
def test_partial_scan_is_discarded(tmp_path: Path) -> None:
|
||||
(tmp_path / "first.txt").touch()
|
||||
(tmp_path / "second.txt").touch()
|
||||
with patch.object(
|
||||
target=Path, attribute="lstat",
|
||||
side_effect=[SimpleNamespace(st_mode=stat.S_IFREG), PermissionError("scan interrupted")],
|
||||
):
|
||||
with pytest.raises(PermissionError):
|
||||
script.take_snapshot(base_dir=tmp_path, recursive=True)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main(args=[str(Path(__file__).resolve()), *sys.argv[1:]]))
|
||||
+200
-185
@@ -1,9 +1,8 @@
|
||||
#!/usr/bin/env -S uv run --script
|
||||
# /// script
|
||||
# dependencies = ["psycopg[binary]", "click", "httpx"]
|
||||
# dependencies = ["click", "httpx"]
|
||||
# ///
|
||||
|
||||
import argparse
|
||||
import contextlib
|
||||
import dataclasses
|
||||
import datetime
|
||||
@@ -19,22 +18,22 @@ import socket
|
||||
import subprocess
|
||||
import threading
|
||||
import time
|
||||
import typing
|
||||
from collections.abc import Callable
|
||||
from decimal import ROUND_CEILING, ROUND_FLOOR, Decimal
|
||||
from functools import partial, wraps
|
||||
from os import getenv
|
||||
from pathlib import Path
|
||||
from typing import Any, Callable
|
||||
from typing import Any
|
||||
from urllib.parse import parse_qs
|
||||
from wsgiref.simple_server import WSGIRequestHandler, make_server
|
||||
|
||||
import click
|
||||
import httpx
|
||||
import psycopg
|
||||
from psycopg import sql as pgsql
|
||||
|
||||
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s: %(message)s")
|
||||
|
||||
HandlerFunc = Callable[["Request"], "Response"]
|
||||
client = httpx.Client(base_url="http://klein:18543", timeout=30)
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
@@ -70,8 +69,7 @@ def ffprobe(video_path: Path) -> FfprobeResult:
|
||||
str(video_path),
|
||||
],
|
||||
text=True,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.PIPE,
|
||||
capture_output=True,
|
||||
)
|
||||
try:
|
||||
proc.check_returncode()
|
||||
@@ -79,9 +77,7 @@ def ffprobe(video_path: Path) -> FfprobeResult:
|
||||
logging.error(f"failed to run ffprobe. stderr={e.stderr}")
|
||||
raise
|
||||
output: dict = json.loads(proc.stdout)
|
||||
video_stream: dict = [
|
||||
s for s in output.get("streams", []) if s.get("codec_type") == "video"
|
||||
][0]
|
||||
video_stream: dict = [s for s in output.get("streams", []) if s.get("codec_type") == "video"][0]
|
||||
video_format: dict = output.get("format", {})
|
||||
width, height = video_stream.get("width"), video_stream.get("height")
|
||||
codec = video_stream.get("codec_name")
|
||||
@@ -94,9 +90,7 @@ def ffprobe(video_path: Path) -> FfprobeResult:
|
||||
tags = video_format.get("tags", {})
|
||||
|
||||
try:
|
||||
sample_w, sample_h = list(
|
||||
map(int, video_stream["sample_aspect_ratio"].split(":"))
|
||||
)
|
||||
sample_w, sample_h = list(map(int, video_stream["sample_aspect_ratio"].split(":")))
|
||||
sample_ar = (sample_w / sample_h) or 1.0
|
||||
except:
|
||||
sample_ar = 1.0
|
||||
@@ -151,9 +145,7 @@ def retry(max_attempts: int = 3, delay: float = 1.0):
|
||||
except Exception as e:
|
||||
last_exception = e
|
||||
if attempt < max_attempts:
|
||||
logging.warning(
|
||||
f"Attempt {attempt} failed: {e}. Retrying in {delay}s..."
|
||||
)
|
||||
logging.warning(f"Attempt {attempt} failed: {e}. Retrying in {delay}s...")
|
||||
time.sleep(delay)
|
||||
else:
|
||||
logging.error(f"All {max_attempts} attempts failed.")
|
||||
@@ -164,49 +156,55 @@ def retry(max_attempts: int = 3, delay: float = 1.0):
|
||||
return decorator
|
||||
|
||||
|
||||
def video_exists(con: psycopg.Connection, file_path: str) -> bool:
|
||||
stmt = con.execute(
|
||||
"UPDATE videos SET last_seen_at = NOW() WHERE file_path = %s RETURNING true AS exists",
|
||||
[file_path],
|
||||
def url_from_path(input_path: str) -> str:
|
||||
return input_path.replace("/mnt/box/files/", "https://u201686.your-storagebox.de/files/")
|
||||
|
||||
|
||||
def video_exists(file_path: str) -> bool:
|
||||
response = client.patch(
|
||||
"/videos",
|
||||
params={"file_path": f"eq.{file_path}", "select": "file_path"},
|
||||
json={"last_seen_at": "now"},
|
||||
headers={"Prefer": "return=representation"},
|
||||
)
|
||||
row = stmt.fetchone()
|
||||
if not row:
|
||||
return False
|
||||
return row["exists"]
|
||||
response.raise_for_status()
|
||||
return bool(response.json())
|
||||
|
||||
|
||||
def save_video(
|
||||
con: psycopg.Connection,
|
||||
category: str,
|
||||
file_path: str,
|
||||
ffprobe_data: dict,
|
||||
video_hash: str,
|
||||
):
|
||||
sql = """
|
||||
INSERT INTO videos(category, file_path, ffprobe, hash_partial)
|
||||
VALUES (%s, %s, %s, %s)
|
||||
ON CONFLICT(file_path) DO UPDATE SET
|
||||
last_seen_at = NOW();"""
|
||||
con.execute(
|
||||
sql,
|
||||
[
|
||||
category,
|
||||
file_path,
|
||||
json.dumps(ffprobe_data),
|
||||
video_hash,
|
||||
],
|
||||
response = client.post(
|
||||
"/videos",
|
||||
params={"on_conflict": "file_path", "select": "file_path"},
|
||||
json={
|
||||
"category": category,
|
||||
"file_path": file_path,
|
||||
"ffprobe": ffprobe_data,
|
||||
"hash_partial": video_hash,
|
||||
},
|
||||
headers={"Prefer": "resolution=ignore-duplicates,return=representation"},
|
||||
)
|
||||
response.raise_for_status()
|
||||
if not response.json():
|
||||
response = client.patch(
|
||||
"/videos",
|
||||
params={"file_path": f"eq.{file_path}"},
|
||||
json={"last_seen_at": "now"},
|
||||
)
|
||||
response.raise_for_status()
|
||||
|
||||
|
||||
def save_error(con: psycopg.Connection, file_path: str, stderr: str):
|
||||
sql = """INSERT INTO scan_state(file_path, stderr) VALUES (%s, %s) ON conflict do nothing"""
|
||||
con.execute(
|
||||
sql,
|
||||
[
|
||||
file_path,
|
||||
stderr,
|
||||
],
|
||||
def save_error(file_path: str, stderr: str) -> None:
|
||||
response = client.post(
|
||||
"/scan_state",
|
||||
json={"file_path": file_path, "stderr": stderr},
|
||||
headers={"Prefer": "resolution=ignore-duplicates"},
|
||||
)
|
||||
response.raise_for_status()
|
||||
|
||||
|
||||
def category_for_file(file_path: Path) -> str:
|
||||
@@ -218,7 +216,7 @@ def category_for_file(file_path: Path) -> str:
|
||||
return "local"
|
||||
|
||||
|
||||
def scan_videos(con: psycopg.Connection, video_paths: list[Path]) -> None:
|
||||
def scan_videos(video_paths: list[Path]) -> None:
|
||||
new_videos = []
|
||||
for i, f in enumerate(video_paths):
|
||||
if "/_picks/" in str(f):
|
||||
@@ -226,7 +224,7 @@ def scan_videos(con: psycopg.Connection, video_paths: list[Path]) -> None:
|
||||
progress = f"[{i + 1}/{len(video_paths)}]"
|
||||
if f.is_dir():
|
||||
# Recursively find files in directory
|
||||
scan_videos(con, list(f.glob("**/*")))
|
||||
scan_videos(list(f.glob("**/*")))
|
||||
continue
|
||||
|
||||
if f.suffix.lower() not in [".mp4", ".mkv", ".avi", ".mov", ".wmv", ".webm"]:
|
||||
@@ -234,7 +232,7 @@ def scan_videos(con: psycopg.Connection, video_paths: list[Path]) -> None:
|
||||
|
||||
file_path = str(f.absolute())
|
||||
|
||||
if video_exists(con, file_path=file_path):
|
||||
if video_exists(file_path=file_path):
|
||||
logging.debug(f"{progress} video already exists {f=}")
|
||||
continue
|
||||
|
||||
@@ -243,27 +241,25 @@ def scan_videos(con: psycopg.Connection, video_paths: list[Path]) -> None:
|
||||
probed = ffprobe(f)
|
||||
except Exception as e:
|
||||
logging.error(f"Failed to probe {f}: {e}")
|
||||
save_error(con, file_path=file_path, stderr=str(e))
|
||||
save_error(file_path=file_path, stderr=str(e))
|
||||
continue
|
||||
|
||||
file_hash = hash_partial(f)
|
||||
|
||||
save_video(
|
||||
category=category_for_file(f),
|
||||
con=con,
|
||||
file_path=file_path,
|
||||
ffprobe_data=probed.__dict__,
|
||||
video_hash=file_hash,
|
||||
)
|
||||
logging.info(f"{progress} saved: {f.name}")
|
||||
new_videos.append(file_path)
|
||||
con.commit()
|
||||
|
||||
if new_videos:
|
||||
parse_releases(con, new_videos)
|
||||
parse_releases(new_videos)
|
||||
|
||||
|
||||
def parse_releases(con: psycopg.Connection, file_paths: list[str]) -> None:
|
||||
def parse_releases(file_paths: list[str]) -> None:
|
||||
api_key = os.getenv("OPENROUTER_API_KEY")
|
||||
if not api_key:
|
||||
logging.warning("OPENROUTER_API_KEY not set, skipping filename parsing")
|
||||
@@ -272,9 +268,7 @@ def parse_releases(con: psycopg.Connection, file_paths: list[str]) -> None:
|
||||
# Chunk file_paths by 20 items to avoid prompt token limits
|
||||
for i, chunk in enumerate(chunked(file_paths, 20), 1):
|
||||
filenames = [Path(p).name for p in chunk]
|
||||
logging.info(
|
||||
f"Parsing filenames for {len(filenames)} videos (chunk {i}) via OpenRouter..."
|
||||
)
|
||||
logging.info(f"Parsing filenames for {len(filenames)} videos (chunk {i}) via OpenRouter...")
|
||||
|
||||
try:
|
||||
parsed_data = parse_filenames_with_ai(filenames, api_key)
|
||||
@@ -294,16 +288,15 @@ def parse_releases(con: psycopg.Connection, file_paths: list[str]) -> None:
|
||||
|
||||
# Update the release column
|
||||
release_info = {k: v for k, v in item.items() if k != "filename"}
|
||||
con.execute(
|
||||
"UPDATE videos SET release = %s WHERE file_path = %s",
|
||||
[json.dumps(release_info), full_path],
|
||||
response = client.patch(
|
||||
"/videos",
|
||||
params={"file_path": f"eq.{full_path}"},
|
||||
json={"release": release_info},
|
||||
)
|
||||
con.commit()
|
||||
response.raise_for_status()
|
||||
logging.info(f"Successfully updated release info for chunk {i}")
|
||||
except Exception as e:
|
||||
logging.error(
|
||||
f"Failed to parse filenames or update database for chunk {i}: {e}"
|
||||
)
|
||||
logging.error(f"Failed to parse filenames or update database for chunk {i}: {e}")
|
||||
|
||||
|
||||
@retry(max_attempts=3, delay=0.1)
|
||||
@@ -357,11 +350,7 @@ HookupHotshot.26.02.06.Episode.453.Shrooms.Q.XXX.720p.HEVC.x265.PRT.mp4 -> studi
|
||||
# Clean output: omit empty strings and empty arrays
|
||||
cleaned_data = []
|
||||
for entry in parsed_data:
|
||||
cleaned_entry = {
|
||||
k: v
|
||||
for k, v in entry.items()
|
||||
if v != "" and not (isinstance(v, list) and not v)
|
||||
}
|
||||
cleaned_entry = {k: v for k, v in entry.items() if v != "" and not (isinstance(v, list) and not v)}
|
||||
actors = cleaned_entry.get("actors", [])
|
||||
if isinstance(actors, list) and len(actors) > 0:
|
||||
actors = [a for a in actors if a and "@" not in a]
|
||||
@@ -523,51 +512,58 @@ class TinyAPI:
|
||||
thread.join(timeout=5)
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def connect_db() -> typing.Generator[psycopg.Connection, typing.Any, typing.Any]:
|
||||
with psycopg.connect(
|
||||
"postgres://abdus:abdus@db.abdus.dev:5444/snot?sslmode=disable"
|
||||
) as conn:
|
||||
conn.row_factory = psycopg.rows.dict_row
|
||||
yield conn
|
||||
|
||||
|
||||
def handle_search(req: Request, conn: psycopg.Connection) -> Response:
|
||||
def handle_search(req: Request) -> Response:
|
||||
query: str = req.payload.get("query", "")
|
||||
if not query:
|
||||
return Response.error("Query is empty")
|
||||
|
||||
where_sql, where_params = filter_to_where(query)
|
||||
|
||||
query_stmt = pgsql.SQL(
|
||||
"""
|
||||
with q as (select
|
||||
id,
|
||||
filename_from_path(file_path) as file_name,
|
||||
file_path,
|
||||
actors,
|
||||
size_bytes / 1048576 as size_mb,
|
||||
(ffprobe->>'width') || 'x' || (ffprobe->>'height') as resolution,
|
||||
duration_human,
|
||||
created_at
|
||||
from videos)
|
||||
select * from q
|
||||
where {filter}
|
||||
order by created_at desc
|
||||
"""
|
||||
).format(filter=pgsql.SQL(where_sql))
|
||||
with conn.cursor() as cursor:
|
||||
rows = cursor.execute(query_stmt, where_params).fetchall()
|
||||
return Response.from_json(
|
||||
[
|
||||
{
|
||||
**row,
|
||||
"download_url": "https://u201686:T6672ICVoWAedECH@u201686.your-storagebox.de/files"
|
||||
+ row["file_path"].replace("/mnt/box/files", ""),
|
||||
}
|
||||
for row in rows
|
||||
]
|
||||
params = filter_to_postgrest(query)
|
||||
params.update(
|
||||
{
|
||||
"select": "id,file_path,file_name,actors,size_bytes,duration_human,created_at,ffprobe",
|
||||
"order": "created_at.desc,id.desc",
|
||||
}
|
||||
)
|
||||
rows = fetch_videos(params=params)
|
||||
results = []
|
||||
for row in rows:
|
||||
probe = row.get("ffprobe") or {}
|
||||
width = probe.get("width")
|
||||
height = probe.get("height")
|
||||
resolution = f"{width}x{height}" if width is not None and height is not None else None
|
||||
results.append(
|
||||
{
|
||||
"id": row["id"],
|
||||
"file_name": row["file_name"],
|
||||
"file_path": row["file_path"],
|
||||
"actors": row.get("actors"),
|
||||
"size_mb": row["size_bytes"] // 1048576,
|
||||
"resolution": resolution,
|
||||
"duration_human": row.get("duration_human"),
|
||||
"created_at": row.get("created_at"),
|
||||
"download_url": url_from_path(input_path=row["file_path"]),
|
||||
}
|
||||
)
|
||||
return Response.from_json(results)
|
||||
|
||||
|
||||
def fetch_videos(params: dict[str, str]) -> list[dict[str, Any]]:
|
||||
rows: list[dict[str, Any]] = []
|
||||
total: int | None = None
|
||||
while total is None or len(rows) < total:
|
||||
response = client.get(
|
||||
"/videos",
|
||||
params={**params, "limit": "1000", "offset": str(len(rows))},
|
||||
headers={"Prefer": "count=exact"} if total is None else None,
|
||||
)
|
||||
response.raise_for_status()
|
||||
page = response.json()
|
||||
if total is None:
|
||||
total = int(response.headers["Content-Range"].rsplit("/", 1)[1])
|
||||
if not page:
|
||||
break
|
||||
rows.extend(page)
|
||||
return rows
|
||||
|
||||
|
||||
def tokenize_filter_groups(filter_query: str) -> list[list[str]]:
|
||||
@@ -587,58 +583,72 @@ def tokenize_filter_groups(filter_query: str) -> list[list[str]]:
|
||||
return [group for group in groups if group]
|
||||
|
||||
|
||||
def token_to_condition(token: str) -> tuple[str | None, list[typing.Any]]:
|
||||
def token_to_conditions(token: str) -> list[str]:
|
||||
if ":" in token:
|
||||
key, value = token.split(":", 1)
|
||||
key = key.strip().lower()
|
||||
value = value.strip()
|
||||
if not value:
|
||||
return None, []
|
||||
return []
|
||||
|
||||
if key == "file_name":
|
||||
return "file_name ILIKE %s", [f"%{value}%"]
|
||||
return [f"file_name.ilike.{postgrest_quote(f'%{value}%')}"]
|
||||
|
||||
if key == "size_mb":
|
||||
match = re.match(
|
||||
r"^(>=|<=|>|<|=)?\s*(\d+(?:\.\d+)?)$", value, re.IGNORECASE
|
||||
)
|
||||
match = re.match(r"^(>=|<=|>|<|=)?\s*(\d+(?:\.\d+)?)$", value, re.IGNORECASE)
|
||||
if not match:
|
||||
return None, []
|
||||
return []
|
||||
op = match.group(1) or "="
|
||||
num = float(match.group(2))
|
||||
return f"size_mb {op} %s", [num]
|
||||
size_mb = Decimal(match.group(2))
|
||||
floor_mb = int(size_mb.to_integral_value(rounding=ROUND_FLOOR))
|
||||
ceil_mb = int(size_mb.to_integral_value(rounding=ROUND_CEILING))
|
||||
if op == ">":
|
||||
return [f"size_bytes.gte.{(floor_mb + 1) * 1048576}"]
|
||||
if op == ">=":
|
||||
return [f"size_bytes.gte.{ceil_mb * 1048576}"]
|
||||
if op == "<":
|
||||
return [f"size_bytes.lt.{ceil_mb * 1048576}"]
|
||||
if op == "<=":
|
||||
return [f"size_bytes.lt.{(floor_mb + 1) * 1048576}"]
|
||||
if size_mb != floor_mb:
|
||||
return ["size_bytes.gt.0", "size_bytes.lt.0"]
|
||||
return [
|
||||
f"size_bytes.gte.{floor_mb * 1048576}",
|
||||
f"size_bytes.lt.{(floor_mb + 1) * 1048576}",
|
||||
]
|
||||
|
||||
if key == "actor":
|
||||
return "actors @> %s::text[]", [[value]]
|
||||
array_value = value.replace("\\", "\\\\").replace('"', '\\"')
|
||||
return [f'actors.cs.{{"{array_value}"}}']
|
||||
|
||||
return None, []
|
||||
return []
|
||||
|
||||
return "file_name ILIKE %s", [f"%{token}%"]
|
||||
return [f"file_name.ilike.{postgrest_quote(f'%{token}%')}"]
|
||||
|
||||
|
||||
def filter_to_where(filter_query: str) -> tuple[str, list[typing.Any]]:
|
||||
def postgrest_quote(value: str) -> str:
|
||||
if re.search(r"[,(){}\\\"]", value):
|
||||
return '"' + value.replace("\\", "\\\\").replace('"', '\\"') + '"'
|
||||
return value
|
||||
|
||||
|
||||
def filter_to_postgrest(filter_query: str) -> dict[str, str]:
|
||||
groups = tokenize_filter_groups(filter_query=filter_query)
|
||||
if not groups:
|
||||
return "TRUE", []
|
||||
return {}
|
||||
|
||||
or_clauses: list[str] = []
|
||||
params: list[typing.Any] = []
|
||||
|
||||
for group in groups:
|
||||
and_clauses: list[str] = []
|
||||
for token in group:
|
||||
clause, clause_params = token_to_condition(token=token)
|
||||
if not clause:
|
||||
continue
|
||||
and_clauses.append(clause)
|
||||
params.extend(clause_params)
|
||||
and_clauses.extend(token_to_conditions(token=token))
|
||||
if and_clauses:
|
||||
or_clauses.append("(" + " AND ".join(and_clauses) + ")")
|
||||
or_clauses.append("and(" + ",".join(and_clauses) + ")")
|
||||
|
||||
if not or_clauses:
|
||||
return "TRUE", []
|
||||
return {}
|
||||
|
||||
return " OR ".join(or_clauses), params
|
||||
return {"or": "(" + ",".join(or_clauses) + ")"}
|
||||
|
||||
|
||||
def handle_home(req: Request, query: str = "") -> Response:
|
||||
@@ -680,12 +690,8 @@ def handle_assets(req: Request, base_path: Path) -> Response:
|
||||
def run_htmlpopup(address: str):
|
||||
exe_path = shutil.which("htmlpopup")
|
||||
if not exe_path:
|
||||
raise RuntimeError(
|
||||
"htmlpopup executable not found in PATH. Please install it to use the HTML popup feature."
|
||||
)
|
||||
proc = subprocess.Popen(
|
||||
[exe_path, "--title", "s·m·u·t", address], text=True, stdout=subprocess.PIPE
|
||||
)
|
||||
raise RuntimeError("htmlpopup executable not found in PATH. Please install it to use the HTML popup feature.")
|
||||
proc = subprocess.Popen([exe_path, "--title", "s·m·u·t", address], text=True, stdout=subprocess.PIPE)
|
||||
try:
|
||||
yield
|
||||
proc.wait()
|
||||
@@ -711,21 +717,20 @@ def cli(ctx: click.Context):
|
||||
@click.option("--query", type=str, default="", help="Initial query string")
|
||||
def serve(port: int, query: str):
|
||||
"""Start the web UI"""
|
||||
with connect_db() as conn:
|
||||
handlers = {
|
||||
"POST /search": partial(handle_search, conn=conn),
|
||||
"GET /": partial(handle_home, query=query),
|
||||
"GET /.*": partial(handle_assets, base_path=Path(__file__).parent),
|
||||
}
|
||||
api = TinyAPI(handlers=handlers)
|
||||
handlers = {
|
||||
"POST /search": handle_search,
|
||||
"GET /": partial(handle_home, query=query),
|
||||
"GET /.*": partial(handle_assets, base_path=Path(__file__).parent),
|
||||
}
|
||||
api = TinyAPI(handlers=handlers)
|
||||
|
||||
with api.serve(port=port) as server_url:
|
||||
logging.info(f"Server running at {server_url}")
|
||||
logging.info("Available endpoints:")
|
||||
for path in handlers.keys():
|
||||
logging.info(f" {path}")
|
||||
with run_htmlpopup(server_url):
|
||||
logging.info("HTML popup started. Press Ctrl+C to stop.")
|
||||
with api.serve(port=port) as server_url:
|
||||
logging.info(f"Server running at {server_url}")
|
||||
logging.info("Available endpoints:")
|
||||
for path in handlers:
|
||||
logging.info(f" {path}")
|
||||
with run_htmlpopup(server_url):
|
||||
logging.info("HTML popup started. Press Ctrl+C to stop.")
|
||||
|
||||
|
||||
@cli.command()
|
||||
@@ -736,8 +741,7 @@ def scan(paths: list[Path]):
|
||||
click.echo("No paths provided to scan.")
|
||||
return
|
||||
|
||||
with connect_db() as con:
|
||||
scan_videos(con=con, video_paths=list(paths))
|
||||
scan_videos(video_paths=list(paths))
|
||||
|
||||
|
||||
@cli.command("mark-deletion")
|
||||
@@ -748,37 +752,48 @@ def mark_deletion(paths: list[Path]):
|
||||
click.echo("No paths provided.")
|
||||
return
|
||||
|
||||
with connect_db() as con:
|
||||
for path in paths:
|
||||
file_name = path.name
|
||||
rows = con.execute(
|
||||
"UPDATE videos SET marked_for_deletion_at = NOW() WHERE file_name = %s RETURNING file_path",
|
||||
[file_name],
|
||||
).fetchall()
|
||||
if rows:
|
||||
for row in rows:
|
||||
click.echo(f"Marked for deletion: {row['file_path']}")
|
||||
else:
|
||||
click.echo(f"Not found in database: {file_name}", err=True)
|
||||
continue
|
||||
for path in paths:
|
||||
file_name = path.name
|
||||
response = client.patch(
|
||||
"/videos",
|
||||
params={
|
||||
"file_name": f"eq.{postgrest_quote(file_name)}",
|
||||
"select": "file_path",
|
||||
},
|
||||
json={"marked_for_deletion_at": "now"},
|
||||
headers={"Prefer": "return=representation"},
|
||||
)
|
||||
response.raise_for_status()
|
||||
rows = response.json()
|
||||
if rows:
|
||||
for row in rows:
|
||||
click.echo(f"Marked for deletion: {row['file_path']}")
|
||||
else:
|
||||
click.echo(f"Not found in database: {file_name}", err=True)
|
||||
continue
|
||||
|
||||
if not click.get_text_stream("stdin").isatty():
|
||||
continue
|
||||
if not click.get_text_stream("stdin").isatty():
|
||||
continue
|
||||
|
||||
stem = path.stem
|
||||
siblings = con.execute(
|
||||
"SELECT file_path, file_name FROM videos WHERE file_name != %s AND file_name LIKE %s AND marked_for_deletion_at IS NULL",
|
||||
[file_name, f"{stem}.%"],
|
||||
).fetchall()
|
||||
for sibling in siblings:
|
||||
if click.confirm(f" Also mark {sibling['file_name']}?", default=False):
|
||||
con.execute(
|
||||
"UPDATE videos SET marked_for_deletion_at = NOW() WHERE file_path = %s",
|
||||
[sibling["file_path"]],
|
||||
)
|
||||
click.echo(f" Marked for deletion: {sibling['file_path']}")
|
||||
|
||||
con.commit()
|
||||
stem = path.stem
|
||||
siblings = fetch_videos(
|
||||
params={
|
||||
"select": "file_path,file_name",
|
||||
"file_name": f"like.{postgrest_quote(f'{stem}.%')}",
|
||||
"marked_for_deletion_at": "is.null",
|
||||
"order": "file_path.asc",
|
||||
}
|
||||
)
|
||||
siblings = [row for row in siblings if row["file_name"] != file_name]
|
||||
for sibling in siblings:
|
||||
if click.confirm(f" Also mark {sibling['file_name']}?", default=False):
|
||||
response = client.patch(
|
||||
"/videos",
|
||||
params={"file_path": f"eq.{sibling['file_path']}"},
|
||||
json={"marked_for_deletion_at": "now"},
|
||||
)
|
||||
response.raise_for_status()
|
||||
click.echo(f" Marked for deletion: {sibling['file_path']}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
+261
-469
@@ -1,501 +1,293 @@
|
||||
#!/usr/bin/env python3.9
|
||||
import argparse
|
||||
#!/usr/bin/env -S uv run --script
|
||||
# /// script
|
||||
# dependencies = ["click", "httpx", "torf"]
|
||||
# ///
|
||||
|
||||
import base64
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
import select
|
||||
import shutil
|
||||
import sys
|
||||
import time
|
||||
from turtle import back
|
||||
import typing
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
import ffmpeg
|
||||
from ffmpeg_strip import clean_video
|
||||
import click
|
||||
import httpx
|
||||
import torf
|
||||
|
||||
logging.basicConfig(level=logging.INFO, format=f"%(asctime)s {logging.BASIC_FORMAT}")
|
||||
ANNOUNCE_URLS = ("http://tracker.empornium.sx:2710/tegqucis10uanp672qh6nhn393xlkncs/announce",)
|
||||
AUTHOR = "zzzp"
|
||||
|
||||
|
||||
def force_import(module: str):
|
||||
import importlib
|
||||
import subprocess
|
||||
import sys
|
||||
|
||||
try:
|
||||
return importlib.import_module(module)
|
||||
except ModuleNotFoundError:
|
||||
subprocess.run([sys.executable, "-m", "pip", "install", module])
|
||||
importlib.invalidate_caches()
|
||||
return importlib.import_module(module)
|
||||
|
||||
|
||||
try:
|
||||
import torf
|
||||
except ImportError:
|
||||
torf = force_import("torf")
|
||||
|
||||
try:
|
||||
import httpx
|
||||
except ImportError:
|
||||
httpx = force_import("httpx")
|
||||
|
||||
try:
|
||||
import inquirer
|
||||
except ImportError:
|
||||
inquirer = force_import("inquirer")
|
||||
|
||||
TORRENT_TRACKER_URL = os.getenv(
|
||||
"TORRENT_TRACKER_URL", "http://tracker.empornium.sx:2710/tegqucis10uanp672qh6nhn393xlkncs/announce"
|
||||
)
|
||||
TORRENT_CREATED_BY = os.getenv("TORRENT_CREATED_BY", "zzzp")
|
||||
TORRENT_DIR = Path(os.getenv("TORRENT_DIR", "/mnt/box/files/_torrents/_new2/"))
|
||||
|
||||
VIDEO_EXTENSIONS = {".mp4", ".mpg", ".mkv"}
|
||||
def format_size(size: int) -> str:
|
||||
if size < 1000:
|
||||
return f"{size:,} bytes"
|
||||
scaled_size = float(size)
|
||||
for unit in ("KB", "MB", "GB", "TB", "PB", "EB"):
|
||||
scaled_size /= 1000
|
||||
if scaled_size < 1000 or unit == "EB":
|
||||
return f"{scaled_size:.2f} {unit} ({size:,} bytes)"
|
||||
raise ValueError(f"Cannot format size: {size}")
|
||||
|
||||
|
||||
def make_torrent(source_dir: Path, tracker_url: str = TORRENT_TRACKER_URL) -> Path:
|
||||
torrent_path = source_dir / f"{source_dir.name}.torrent"
|
||||
if torrent_path.is_file():
|
||||
logging.info("Torrent file is already created")
|
||||
return torrent_path
|
||||
|
||||
t = torf.Torrent(
|
||||
path=str(source_dir.resolve()),
|
||||
name=source_dir.resolve().name,
|
||||
trackers=[tracker_url],
|
||||
private=True,
|
||||
created_by=TORRENT_CREATED_BY,
|
||||
exclude_globs=["*.txt", "post.txt", "*.torrent", "*.gif"],
|
||||
def create_torrent(
|
||||
paths: tuple[Path, ...],
|
||||
output: Path | None = None,
|
||||
announce_urls: tuple[str, ...] = ANNOUNCE_URLS,
|
||||
author: str = AUTHOR,
|
||||
private: bool = True,
|
||||
name: str | None = None,
|
||||
overwrite: bool = False,
|
||||
) -> Path:
|
||||
sources = list(dict.fromkeys(path.expanduser().absolute() for path in paths))
|
||||
if not sources:
|
||||
raise ValueError("Provide at least one file or directory")
|
||||
for source in sources:
|
||||
if not source.is_file() and not source.is_dir():
|
||||
raise ValueError(f"Not a file or directory: {source}")
|
||||
if len({source.parent for source in sources}) > 1:
|
||||
raise ValueError("All inputs must be in the same directory")
|
||||
torrent = torf.Torrent(
|
||||
trackers=announce_urls,
|
||||
created_by=author,
|
||||
private=private,
|
||||
creation_date=int(time.time()),
|
||||
)
|
||||
t.generate()
|
||||
t.write(torrent_path, overwrite=True)
|
||||
if len(sources) == 1:
|
||||
torrent.path = sources[0]
|
||||
else:
|
||||
torrent.filepaths = sources
|
||||
if name is not None:
|
||||
torrent.name = name
|
||||
if torrent.size <= 0:
|
||||
raise ValueError("The selected paths contain no nonempty files")
|
||||
torrent_path = output or Path.cwd() / f"{torrent.name}.torrent"
|
||||
if torrent_path.exists() and not overwrite:
|
||||
raise FileExistsError(f"Output already exists: {torrent_path}")
|
||||
if torrent_path.absolute() in {Path(path).absolute() for path in torrent.filepaths}:
|
||||
raise ValueError("Output must not overwrite a source file")
|
||||
with click.progressbar(length=torrent.pieces, label="Hashing", show_percent=True, file=sys.stderr) as progress:
|
||||
|
||||
def report_progress(
|
||||
torrent: torf.Torrent,
|
||||
filepath: str,
|
||||
pieces_done: int,
|
||||
pieces_total: int,
|
||||
) -> None:
|
||||
progress.update(n_steps=pieces_done - progress.pos)
|
||||
|
||||
if not torrent.generate(callback=report_progress, interval=0.2):
|
||||
raise ValueError("Torrent hashing did not complete")
|
||||
torrent.write(filepath=torrent_path, overwrite=overwrite)
|
||||
return torrent_path
|
||||
|
||||
|
||||
def human_size(size: int) -> str:
|
||||
suffix = "B"
|
||||
for unit in ["", "K", "M", "G", "T", "P", "E", "Z"]:
|
||||
if abs(size) < 1024.0:
|
||||
return "%3.2f%s%s" % (size, unit, suffix)
|
||||
size /= 1024.0
|
||||
return "%.2f%s%s" % (size, "Yi", suffix)
|
||||
def modify_torrent(
|
||||
torrent_path: Path,
|
||||
output: Path | None = None,
|
||||
announce_urls: tuple[str, ...] = (),
|
||||
author: str | None = None,
|
||||
private: bool | None = None,
|
||||
name: str | None = None,
|
||||
) -> Path:
|
||||
torrent = torf.Torrent.read(filepath=torrent_path)
|
||||
if announce_urls:
|
||||
torrent.trackers = announce_urls
|
||||
if author is not None:
|
||||
torrent.created_by = author
|
||||
if private is not None:
|
||||
torrent.private = private
|
||||
if name is not None:
|
||||
torrent.name = name
|
||||
destination = output or torrent_path
|
||||
torrent.write(filepath=destination, overwrite=destination.absolute() == torrent_path.absolute())
|
||||
return destination
|
||||
|
||||
|
||||
def find_video(root_dir: Path) -> Optional[Path]:
|
||||
for ext in VIDEO_EXTENSIONS:
|
||||
for f in root_dir.rglob(f"*{ext}"):
|
||||
return f
|
||||
return None
|
||||
|
||||
|
||||
def pick_name(video_path: Path) -> str:
|
||||
try:
|
||||
if parse_filename(video_path.stem):
|
||||
return video_path.stem
|
||||
except ValueError:
|
||||
pass
|
||||
|
||||
candidates = set()
|
||||
|
||||
parent = video_path.parent
|
||||
while parent.exists():
|
||||
try:
|
||||
if parse_filename(parent.name):
|
||||
candidates.add(parent.name)
|
||||
except ValueError:
|
||||
pass
|
||||
if parent.parent == parent:
|
||||
break
|
||||
parent = parent.parent
|
||||
|
||||
if len(candidates) == 1:
|
||||
return candidates.pop()
|
||||
elif len(candidates) > 1:
|
||||
questions = [
|
||||
inquirer.List(
|
||||
"best_name",
|
||||
message="Pick the best filename",
|
||||
choices=[*candidates, "<custom>"],
|
||||
),
|
||||
]
|
||||
|
||||
answers = inquirer.prompt(questions) or {}
|
||||
best_name = answers.get("best_name")
|
||||
if best_name != "<custom>":
|
||||
return best_name
|
||||
|
||||
while True:
|
||||
try:
|
||||
answers = inquirer.prompt([inquirer.Editor("best_name", message="New name", default=video_path.stem)])
|
||||
if not answers:
|
||||
raise Exception("Cancelled")
|
||||
best_name = answers["best_name"].splitlines(keepends=False)[0].strip()
|
||||
parse_filename(best_name)
|
||||
return best_name
|
||||
except ValueError:
|
||||
continue
|
||||
|
||||
|
||||
def move_video(video_path: Path, new_name: typing.Optional[str], link: bool = False) -> Path:
|
||||
"""
|
||||
Move video to the torrent directory and returns the video path
|
||||
"""
|
||||
if new_name:
|
||||
best_name = Path(new_name).name
|
||||
else:
|
||||
best_name = pick_name(video_path)
|
||||
assert not best_name.startswith('.')
|
||||
logging.info("using name: %s", best_name)
|
||||
|
||||
best_name = best_name.removesuffix('.nometadata')
|
||||
if Path(best_name).suffix in VIDEO_EXTENSIONS:
|
||||
best_name = Path(best_name).stem
|
||||
|
||||
renamed_path = video_path.with_stem(best_name)
|
||||
if link:
|
||||
video_path = video_path.link_to(renamed_path)
|
||||
else:
|
||||
video_path = video_path.rename(renamed_path)
|
||||
video_path = renamed_path
|
||||
|
||||
new_loc = TORRENT_DIR / video_path.stem
|
||||
new_loc.mkdir(parents=True, exist_ok=True)
|
||||
target_path = new_loc / video_path.name
|
||||
video_path.rename(target_path)
|
||||
|
||||
for f in video_path.parent.glob("*.jpg"):
|
||||
if f.stem.startswith(video_path.stem):
|
||||
f.rename(new_loc / f.name)
|
||||
break
|
||||
|
||||
return target_path
|
||||
|
||||
|
||||
def upload_image(image_path: Path) -> str:
|
||||
session = httpx.Client()
|
||||
res = session.get("https://jerking.empornium.ph/?agree-consent", follow_redirects=True)
|
||||
try:
|
||||
auth_token = re.search(r'name="auth_token" value="([^"]+)"', res.text).group(1)
|
||||
except AttributeError:
|
||||
auth_token = re.search(r'auth_token = "([^"]+)"', res.text).group(1)
|
||||
|
||||
with image_path.open("rb") as f:
|
||||
res = session.post(
|
||||
"https://jerking.empornium.ph/json",
|
||||
headers={"accept": "application/json"},
|
||||
data={
|
||||
"thumb_width": "160",
|
||||
"thumb_height": "160",
|
||||
"thumb_crop": "false",
|
||||
"medium_width": "500",
|
||||
"medium_crop": "false",
|
||||
"type": "file",
|
||||
"action": "upload",
|
||||
"timestamp": str(round(time.time() * 1000)),
|
||||
"auth_token": auth_token,
|
||||
"nsfw": "0",
|
||||
},
|
||||
files={
|
||||
"source": ("thumbs.jpg", f),
|
||||
},
|
||||
)
|
||||
|
||||
data = res.json()
|
||||
image_url = data["image"]["url"]
|
||||
thumbnail_url = data["image"]["display_url"]
|
||||
|
||||
return image_url
|
||||
|
||||
|
||||
ParsedFilename = typing.TypedDict(
|
||||
'ParsedFilename', {'actors': list[str], 'studio': str, 'date': str, 'title': str, 'tags': list[str]}
|
||||
)
|
||||
|
||||
|
||||
def parse_filename(filename: str) -> ParsedFilename:
|
||||
parsed = {
|
||||
"tags": [],
|
||||
}
|
||||
filename = re.sub(r"\s*\[\d+[^]]+]", "", filename)
|
||||
if m := re.search(r"(?P<actors>.+)\s+-+\s+@(?P<studio>\S+)\s+-+\s+(?P<title>.+)\s+-+\s+(?P<date>[\d-]+)", filename):
|
||||
parsed.update(m.groupdict())
|
||||
elif m := re.search(r"(?P<actors>.+)\s+-+\s+@(?P<studio>\S+)\s+-+\s+(?P<date>[\d-]+)", filename):
|
||||
parsed.update(m.groupdict())
|
||||
elif m := re.search(r"(?P<actors>.+)\s+-+\s+(?P<title>.+)\s+-+\s+(?P<date>[\d-]+)", filename):
|
||||
parsed.update(m.groupdict())
|
||||
elif m := re.search(r"(?P<actors>.+)\s+-+\s+(?P<title>\D.+)", filename):
|
||||
parsed.update(m.groupdict())
|
||||
else:
|
||||
raise ValueError("Unknown filename format")
|
||||
|
||||
if "actors" in parsed:
|
||||
names = re.split(r"\s*,\s*", parsed["actors"])
|
||||
with_aliases = []
|
||||
for it in names:
|
||||
with_aliases.extend(re.split(r"\s+aka\s+", it))
|
||||
parsed["actors"] = names
|
||||
parsed["tags"].extend([it.lower().replace(" ", ".") for it in with_aliases])
|
||||
if "studio" in parsed:
|
||||
parsed["tags"].append(parsed["studio"].lower() + ".com")
|
||||
if "date" in parsed:
|
||||
year, ym = parsed["date"][:4], parsed["date"][:7].replace("-", ".")
|
||||
parsed["tags"].append(year)
|
||||
parsed["tags"].append(ym)
|
||||
|
||||
return parsed
|
||||
|
||||
|
||||
def generate_title(parsed: ParsedFilename) -> str:
|
||||
parts = []
|
||||
if actors := parsed.get('actors'):
|
||||
parts.append(', '.join(actors))
|
||||
if studio := parsed.get('studio'):
|
||||
parts.append(f'@{studio}')
|
||||
if title := parsed.get('title'):
|
||||
parts.append(title)
|
||||
if date := parsed.get('date'):
|
||||
parts.append(date)
|
||||
return ' -- '.join(parts)
|
||||
|
||||
|
||||
def generate_post_bbcode(video_path: Path, thumbnails_path: Path) -> str:
|
||||
try:
|
||||
image_url = upload_image(thumbnails_path)
|
||||
image_bbcode = f"[img]{image_url}[/img]"
|
||||
except:
|
||||
logging.exception("Failed to upload image")
|
||||
image_bbcode = ""
|
||||
|
||||
parsed = parse_filename(video_path.stem)
|
||||
probe = ffmpeg.ffprobe(video_path)
|
||||
|
||||
if 1900 <= probe.width <= 2200:
|
||||
hd = "1080p"
|
||||
elif 1200 <= probe.width <= 1400:
|
||||
hd = "720p"
|
||||
elif 3000 <= probe.width:
|
||||
hd = "4K"
|
||||
else:
|
||||
hd = None
|
||||
|
||||
if hd:
|
||||
parsed["tags"].append(hd.lower())
|
||||
if probe.codec == "hevc":
|
||||
parsed["tags"].extend(["x265", "x265.reencode", "hevc.x265"])
|
||||
|
||||
table = {
|
||||
"Duration": probe.duration_human,
|
||||
"Format": video_path.suffix.strip("."),
|
||||
"Filesize": human_size(video_path.stat().st_size),
|
||||
"Resolution": f"{probe.width}x{probe.height}",
|
||||
"Codec": probe.codec,
|
||||
"Bit rate": f"{probe.bitrate // 1000} kbit/s",
|
||||
"FPS": probe.fps,
|
||||
}
|
||||
release_date = parsed.get("date", "")
|
||||
|
||||
if release_date:
|
||||
table = {
|
||||
"Release Date": release_date,
|
||||
**table,
|
||||
}
|
||||
|
||||
lines = [
|
||||
"[table=nball]",
|
||||
*(f"[tr][th=20]{k}[/th][td]{v}[/td][/tr]" for k, v in table.items()),
|
||||
"[/table]",
|
||||
]
|
||||
table_bbcode = "\n".join(lines)
|
||||
|
||||
m_performers = re.search(r"(.+?) -+\s", video_path.stem)
|
||||
performers = m_performers.group(1) if m_performers else ""
|
||||
|
||||
title = generate_title(parsed)
|
||||
|
||||
metadata = {
|
||||
**parsed,
|
||||
"is_hevc": probe.codec == "hevc",
|
||||
"hd": hd,
|
||||
}
|
||||
metadata_json = json.dumps(metadata)
|
||||
|
||||
return f"""
|
||||
{metadata_json}
|
||||
---
|
||||
[b]{title}[/b]
|
||||
|
||||
[cast]
|
||||
{performers}
|
||||
|
||||
[details]
|
||||
|
||||
[info]
|
||||
{table_bbcode}
|
||||
|
||||
[screens]
|
||||
{image_bbcode}
|
||||
""".strip()
|
||||
|
||||
|
||||
def add_torrent(torrent_path: Path):
|
||||
import time
|
||||
from functools import partial
|
||||
|
||||
make_id = partial(time.time_ns)
|
||||
|
||||
# log in
|
||||
logging.debug("Connecting to Deluge")
|
||||
session = httpx.Client(base_url="https://t.zzzp.win/", timeout=30)
|
||||
res = session.post("/json", json={"method": "auth.login", "params": ["xAsametk50"], "id": make_id()})
|
||||
res.raise_for_status()
|
||||
|
||||
# find first available host
|
||||
res = session.post("/json", json={"method": "web.get_hosts", "params": [], "id": make_id()})
|
||||
res.raise_for_status()
|
||||
host_id = res.json()["result"][0][0]
|
||||
|
||||
# connect to a host
|
||||
res = session.post("/json", json={"method": "web.connect", "params": [host_id], "id": make_id()})
|
||||
res.raise_for_status()
|
||||
|
||||
# upload torrent
|
||||
logging.debug("Uploading torrent file")
|
||||
with torrent_path.open("rb") as f:
|
||||
res = session.post("/upload", files={"file": f})
|
||||
res.raise_for_status()
|
||||
remote_path: str = res.json()["files"][0]
|
||||
|
||||
# add torrent
|
||||
logging.debug("Adding torrent file")
|
||||
res = session.post(
|
||||
"/json",
|
||||
json={
|
||||
"method": "web.add_torrents",
|
||||
"params": [
|
||||
[
|
||||
{
|
||||
"path": remote_path,
|
||||
"options": {
|
||||
"file_priorities": [1],
|
||||
"add_paused": True,
|
||||
"sequential_download": False,
|
||||
"pre_allocate_storage": False,
|
||||
"download_location": "/dl/_new2",
|
||||
"move_completed": False,
|
||||
"move_completed_path": "/root/Downloads",
|
||||
"prioritize_first_last_pieces": True,
|
||||
"seed_mode": True,
|
||||
"super_seeding": False,
|
||||
},
|
||||
}
|
||||
]
|
||||
],
|
||||
"id": make_id(),
|
||||
def enqueue_torrent(
|
||||
transmission_url: str,
|
||||
credentials: str | tuple[str, str] | None,
|
||||
torrent_path: Path,
|
||||
download_path: str | Path,
|
||||
) -> str:
|
||||
if isinstance(credentials, str):
|
||||
username, separator, password = credentials.partition(":")
|
||||
if not separator:
|
||||
raise ValueError("Credentials must be username:password")
|
||||
credentials = (username, password)
|
||||
auth = httpx.BasicAuth(*credentials) if credentials is not None else None
|
||||
rpc_url = transmission_url.rstrip("/")
|
||||
if not httpx.URL(url=rpc_url).path.endswith("/rpc"):
|
||||
rpc_url = f"{rpc_url}/transmission/rpc"
|
||||
payload = {
|
||||
"method": "torrent-add",
|
||||
"arguments": {
|
||||
"download-dir": str(download_path),
|
||||
"metainfo": base64.b64encode(torrent_path.read_bytes()).decode("ascii"),
|
||||
"paused": False,
|
||||
},
|
||||
)
|
||||
|
||||
res.raise_for_status()
|
||||
torrent_id = res.json()["result"][0][1]
|
||||
|
||||
# force recheck
|
||||
# logging.info("Triggering a forced recheck")
|
||||
# time.sleep(5)
|
||||
# try:
|
||||
# _ = session.post(
|
||||
# "/json", json={"method": "core.force_recheck", "params": [[torrent_id]], "id": make_id()}, timeout=0.1
|
||||
# )
|
||||
# except httpx.HTTPError:
|
||||
# pass
|
||||
|
||||
logging.debug("Torrent file has been added successfully.")
|
||||
}
|
||||
with httpx.Client(auth=auth, timeout=20) as client:
|
||||
for _ in range(3):
|
||||
response = client.post(url=rpc_url, json=payload)
|
||||
if response.status_code == 409:
|
||||
client.headers["X-Transmission-Session-Id"] = response.headers["X-Transmission-Session-Id"]
|
||||
continue
|
||||
response.raise_for_status()
|
||||
result = response.json()
|
||||
if result["result"] not in ("success", "duplicate-torrent"):
|
||||
raise ValueError(f"Transmission rejected torrent: {result['result']}")
|
||||
added = result["arguments"].get("torrent-added") or result["arguments"].get("torrent-duplicate")
|
||||
if not added:
|
||||
raise ValueError("Transmission returned no torrent details")
|
||||
return added["hashString"]
|
||||
raise ValueError("Transmission session negotiation failed")
|
||||
|
||||
|
||||
def parse_args(argv: list[str]):
|
||||
arger = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
|
||||
arger.add_argument("dir_or_video", type=Path, help="Directory or video file to add")
|
||||
arger.add_argument("--name", "-n", help="Torrent name")
|
||||
arger.add_argument("--no-metadata", action='store_true', dest='remove_metadata', help="Remove all metadata")
|
||||
arger.add_argument("--post-only", action="store_true", help="Generate only the post bbcode")
|
||||
arger.add_argument("--link", action="store_true", help="Link video file instead of moving")
|
||||
arger.add_argument("--thumb-path", dest='thumbnail_path', type=Path, help="Path to thumbnail sheet")
|
||||
if not argv:
|
||||
arger.print_help()
|
||||
sys.exit(0)
|
||||
return arger.parse_args(argv)
|
||||
@click.group()
|
||||
def main() -> None:
|
||||
"""Create, modify, inspect, and enqueue torrent files."""
|
||||
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s: %(message)s")
|
||||
|
||||
|
||||
def validate_name(base_name: str):
|
||||
if len(base_name) > (128 - len(".mp4")):
|
||||
raise ValueError(f"Filename is too long: {base_name} ({len(base_name)} chars)")
|
||||
if "?" in base_name:
|
||||
raise ValueError(f"Filename contains invalid characters: {base_name}")
|
||||
@main.command(name="create")
|
||||
@click.argument("paths", nargs=-1, required=True, type=click.Path(exists=True, path_type=Path))
|
||||
@click.option("--output", type=click.Path(dir_okay=False, path_type=Path), help="Destination; defaults to NAME.torrent in the current directory.")
|
||||
@click.option("--announce-url", "announce_urls", multiple=True, help="Tracker URL; repeat for multiple trackers. Defaults to the hardcoded URLs.")
|
||||
@click.option("--author", default=AUTHOR, show_default=True)
|
||||
@click.option("--private/--no-private", default=True, show_default=True)
|
||||
@click.option("--name", help="Torrent name; defaults to the file or common directory name.")
|
||||
@click.option("--overwrite", is_flag=True, help="Replace an existing output torrent.")
|
||||
def create(
|
||||
paths: tuple[Path, ...],
|
||||
output: Path | None,
|
||||
announce_urls: tuple[str, ...],
|
||||
author: str,
|
||||
private: bool,
|
||||
name: str | None,
|
||||
overwrite: bool,
|
||||
) -> None:
|
||||
"""Create one torrent from FILES or DIRECTORIES, retaining their relative paths."""
|
||||
try:
|
||||
destination = create_torrent(
|
||||
paths=paths,
|
||||
output=output,
|
||||
announce_urls=announce_urls or ANNOUNCE_URLS,
|
||||
author=author,
|
||||
private=private,
|
||||
name=name,
|
||||
overwrite=overwrite,
|
||||
)
|
||||
except (OSError, ValueError, torf.TorfError) as error:
|
||||
raise click.ClickException(message=str(error)) from error
|
||||
logging.info(f"Saved torrent at {destination}")
|
||||
|
||||
|
||||
def main():
|
||||
args = parse_args(sys.argv[1:])
|
||||
dir_or_video: Path = args.dir_or_video.expanduser().resolve()
|
||||
if dir_or_video.is_dir():
|
||||
dir_path = dir_or_video
|
||||
if not dir_path.is_relative_to(TORRENT_DIR):
|
||||
target_path = TORRENT_DIR / dir_path.name
|
||||
logging.info(f"Linking to {target_path}")
|
||||
shutil.copytree(src=dir_path, dst=target_path, copy_function=os.link, symlinks=False)
|
||||
dir_path = target_path
|
||||
torrent_path = make_torrent(dir_path)
|
||||
logging.info(f"Saved torrent at {torrent_path}")
|
||||
@main.command(name="modify")
|
||||
@click.argument("torrent_path", type=click.Path(exists=True, dir_okay=False, path_type=Path))
|
||||
@click.option("--output", type=click.Path(dir_okay=False, path_type=Path), help="Destination; defaults to editing the input torrent in place.")
|
||||
@click.option("--announce-url", "announce_urls", multiple=True, help="Replace trackers; repeat for multiple URLs.")
|
||||
@click.option("--author", help="Replace the creator.")
|
||||
@click.option("--private/--no-private", default=None, help="Set or clear the private flag.")
|
||||
@click.option("--name", help="Replace the torrent name.")
|
||||
def modify(
|
||||
torrent_path: Path,
|
||||
output: Path | None,
|
||||
announce_urls: tuple[str, ...],
|
||||
author: str | None,
|
||||
private: bool | None,
|
||||
name: str | None,
|
||||
) -> None:
|
||||
"""Change metadata without rehashing the content; omitted fields stay unchanged."""
|
||||
try:
|
||||
destination = modify_torrent(
|
||||
torrent_path=torrent_path,
|
||||
output=output,
|
||||
announce_urls=announce_urls,
|
||||
author=author,
|
||||
private=private,
|
||||
name=name,
|
||||
)
|
||||
except (OSError, ValueError, torf.TorfError) as error:
|
||||
raise click.ClickException(message=str(error)) from error
|
||||
logging.info(f"Saved torrent at {destination}")
|
||||
|
||||
logging.info("Adding torrent file to Deluge")
|
||||
add_torrent(torrent_path)
|
||||
|
||||
@main.command(name="enqueue")
|
||||
@click.argument("torrent_path", type=click.Path(exists=True, dir_okay=False, path_type=Path))
|
||||
@click.option("--transmission-url", "--host", required=True, envvar="TRANSMISSION_HOST", help="Transmission host or RPC URL; also reads TRANSMISSION_HOST.")
|
||||
@click.option("--credentials", "--auth", envvar="TRANSMISSION_AUTH", help="Username:password; also reads TRANSMISSION_AUTH.")
|
||||
@click.option("--download-path", required=True, type=click.Path(path_type=Path), help="Download directory on the remote Transmission host.")
|
||||
def enqueue(
|
||||
torrent_path: Path,
|
||||
transmission_url: str,
|
||||
credentials: str | None,
|
||||
download_path: Path,
|
||||
) -> None:
|
||||
"""Add a torrent to Transmission at the given remote download directory."""
|
||||
try:
|
||||
torrent_hash = enqueue_torrent(
|
||||
transmission_url=transmission_url,
|
||||
credentials=credentials,
|
||||
torrent_path=torrent_path,
|
||||
download_path=download_path,
|
||||
)
|
||||
except (OSError, ValueError, httpx.HTTPError) as error:
|
||||
raise click.ClickException(message=str(error)) from error
|
||||
logging.info(f"Enqueued torrent {torrent_hash} at {download_path}")
|
||||
|
||||
|
||||
|
||||
|
||||
@main.command(name="inspect")
|
||||
@click.argument("torrent_path", type=click.Path(exists=True, dir_okay=False, path_type=Path))
|
||||
@click.option("--json", "json_output", is_flag=True, help="Output JSON instead of plain text.")
|
||||
def inspect(torrent_path: Path, json_output: bool) -> None:
|
||||
"""Show torrent properties and content paths with readable sizes."""
|
||||
try:
|
||||
torrent = torf.Torrent.read(filepath=torrent_path)
|
||||
except (OSError, ValueError, torf.TorfError) as error:
|
||||
raise click.ClickException(message=str(error)) from error
|
||||
properties = {
|
||||
"name": torrent.name,
|
||||
"author": torrent.created_by,
|
||||
"private": torrent.private,
|
||||
"infohash": torrent.infohash,
|
||||
"size": torrent.size,
|
||||
"piece_size": torrent.piece_size,
|
||||
"pieces": torrent.pieces,
|
||||
"creation_date": torrent.creation_date.isoformat() if torrent.creation_date else None,
|
||||
"comment": torrent.comment,
|
||||
"source": torrent.source,
|
||||
"announce_urls": [list(tier) for tier in torrent.trackers],
|
||||
"webseeds": list(torrent.webseeds),
|
||||
"httpseeds": list(torrent.httpseeds),
|
||||
}
|
||||
files = [{"path": str(file), "size": file.size} for file in torrent.files]
|
||||
if json_output:
|
||||
click.echo(message=json.dumps({"properties": properties, "files": files}, indent=2, ensure_ascii=False))
|
||||
return
|
||||
|
||||
if not dir_or_video.is_file():
|
||||
logging.error("Invalid file or directory")
|
||||
return
|
||||
|
||||
video_path = dir_or_video
|
||||
if video_path.suffix not in VIDEO_EXTENSIONS:
|
||||
logging.error('not a video')
|
||||
return
|
||||
|
||||
if args.remove_metadata:
|
||||
logging.info('Removing video metadata')
|
||||
save_path = video_path.with_stem(f'{video_path.stem}.nometadata')
|
||||
video_path = clean_video(video_path=video_path, save_path=save_path)
|
||||
|
||||
video_path = move_video(video_path, new_name=args.name, link=args.link)
|
||||
dir_path = video_path.parent
|
||||
|
||||
if p := args.thumbnail_path:
|
||||
thumbnail_path = video_path.with_suffix('.jpg')
|
||||
p.link_to(thumbnail_path)
|
||||
else:
|
||||
logging.info("Creating thumbnail tile")
|
||||
old_thumb_path = video_path.with_suffix(".thumbnail.jpg")
|
||||
thumbnail_path = video_path.with_suffix(".jpg")
|
||||
if old_thumb_path.is_file():
|
||||
old_thumb_path.rename(thumbnail_path)
|
||||
image_jpeg = ffmpeg.make_thumbnail_tile(video_path)
|
||||
thumbnail_path.write_bytes(image_jpeg)
|
||||
logging.info(f"Saved thumbnails at {thumbnail_path}")
|
||||
|
||||
post_bbcode = generate_post_bbcode(video_path, thumbnail_path)
|
||||
video_path.with_name("post.txt").write_text(post_bbcode)
|
||||
print(post_bbcode)
|
||||
|
||||
if args.post_only:
|
||||
return
|
||||
|
||||
total_files = len(list(dir_path.rglob("*")))
|
||||
logging.info(f"Creating torrent file from {dir_path}. Total files: {total_files}")
|
||||
torrent_path = make_torrent(dir_path)
|
||||
logging.info(f"Saved torrent at {torrent_path}")
|
||||
|
||||
logging.info("Adding torrent file to Deluge")
|
||||
add_torrent(torrent_path)
|
||||
for key, value in properties.items():
|
||||
label = key.replace("_", " ").capitalize()
|
||||
if isinstance(value, list):
|
||||
click.echo(message=f"{label}:" if value else f"{label}: (none)")
|
||||
for index, item in enumerate(value, start=1):
|
||||
if isinstance(item, list):
|
||||
click.echo(message=f" Tier {index}: {', '.join(item)}")
|
||||
else:
|
||||
click.echo(message=f" {item}")
|
||||
elif value is None:
|
||||
click.echo(message=f"{label}: (unset)")
|
||||
elif isinstance(value, bool):
|
||||
click.echo(message=f"{label}: {'yes' if value else 'no'}")
|
||||
elif key in ("size", "piece_size"):
|
||||
click.echo(message=f"{label}: {format_size(size=value)}")
|
||||
else:
|
||||
click.echo(message=f"{label}: {value}")
|
||||
click.echo(message=f"\nFiles ({len(files)}):")
|
||||
for file in files:
|
||||
click.echo(message=f" {file['path']}: {format_size(size=file['size'])}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
@@ -0,0 +1,247 @@
|
||||
#!/usr/bin/env -S uv run --script
|
||||
# /// script
|
||||
# dependencies = ["click", "httpx", "torf", "pytest"]
|
||||
# ///
|
||||
|
||||
import base64
|
||||
import hashlib
|
||||
import json
|
||||
from pathlib import Path
|
||||
import sys
|
||||
from unittest.mock import patch
|
||||
|
||||
from click.testing import CliRunner
|
||||
import httpx
|
||||
import pytest
|
||||
import torf
|
||||
|
||||
import torrentify as script
|
||||
|
||||
|
||||
@pytest.mark.parametrize("kind", ["file", "directory", "mixed"])
|
||||
def test_create_selected_content(tmp_path: Path, kind: str) -> None:
|
||||
source = tmp_path / "first.txt"
|
||||
source.write_bytes(b"first")
|
||||
directory = tmp_path / "folder"
|
||||
directory.mkdir()
|
||||
(directory / "nested.txt").write_bytes(b"nested")
|
||||
(tmp_path / "unselected.txt").write_bytes(b"unselected")
|
||||
paths = {"file": (source,), "directory": (directory,), "mixed": (source, directory)}[kind]
|
||||
output = tmp_path / "result.torrent"
|
||||
result = CliRunner().invoke(cli=script.main, args=["create", *map(str, paths), "--output", str(output)])
|
||||
assert result.exit_code == 0, result.output
|
||||
torrent = torf.Torrent.read(filepath=output)
|
||||
assert torrent.private is True
|
||||
assert torrent.created_by == script.AUTHOR
|
||||
assert list(torrent.trackers.flat) == list(script.ANNOUNCE_URLS)
|
||||
expected = {"file": b"first", "directory": b"nested", "mixed": b"firstnested"}[kind]
|
||||
assert torrent.metainfo["info"]["pieces"] == hashlib.sha1(expected).digest()
|
||||
if kind == "mixed":
|
||||
assert [item["path"] for item in torrent.metainfo["info"]["files"]] == [
|
||||
["first.txt"], ["folder", "nested.txt"],
|
||||
]
|
||||
if kind == "directory":
|
||||
assert torrent.name == "folder"
|
||||
assert torrent.metainfo["info"]["files"][0]["path"] == ["nested.txt"]
|
||||
|
||||
|
||||
def test_reject_inputs_in_different_directories(tmp_path: Path) -> None:
|
||||
source = tmp_path / "first.txt"
|
||||
source.write_bytes(b"first")
|
||||
directory = tmp_path / "folder"
|
||||
directory.mkdir()
|
||||
nested = directory / "nested.txt"
|
||||
nested.write_bytes(b"nested")
|
||||
result = CliRunner().invoke(cli=script.main, args=["create", str(source), str(nested)])
|
||||
assert result.exit_code == 1
|
||||
assert "All inputs must be in the same directory" in result.output
|
||||
assert not list(tmp_path.glob("*.torrent"))
|
||||
|
||||
|
||||
def test_create_rejects_empty_and_existing_output(tmp_path: Path) -> None:
|
||||
source = tmp_path / "file.txt"
|
||||
source.touch()
|
||||
with pytest.raises(ValueError, match="no nonempty files"):
|
||||
script.create_torrent(paths=(source,), output=tmp_path / "result.torrent")
|
||||
source.write_bytes(b"content")
|
||||
output = script.create_torrent(paths=(source,), output=tmp_path / "result.torrent")
|
||||
original = output.read_bytes()
|
||||
with pytest.raises(FileExistsError):
|
||||
script.create_torrent(paths=(source,), output=output)
|
||||
assert output.read_bytes() == original
|
||||
with pytest.raises(ValueError, match="overwrite a source"):
|
||||
script.create_torrent(paths=(source,), output=source, overwrite=True)
|
||||
assert source.read_bytes() == b"content"
|
||||
|
||||
|
||||
def test_modify_metadata_without_content(tmp_path: Path) -> None:
|
||||
source = tmp_path / "content.txt"
|
||||
source.write_bytes(b"content")
|
||||
output = script.create_torrent(paths=(source,), output=tmp_path / "original.torrent")
|
||||
before = torf.Torrent.read(filepath=output)
|
||||
source.unlink()
|
||||
result = CliRunner().invoke(cli=script.main, args=[
|
||||
"modify", str(output), "--announce-url", "https://tracker.example/announce",
|
||||
"--announce-url", "udp://tracker.example:80/announce", "--author", "new author",
|
||||
"--name", "new name", "--no-private",
|
||||
])
|
||||
assert result.exit_code == 0, result.output
|
||||
after = torf.Torrent.read(filepath=output)
|
||||
assert after.name == "new name"
|
||||
assert after.created_by == "new author"
|
||||
assert after.private is False
|
||||
assert list(after.trackers.flat) == ["https://tracker.example/announce", "udp://tracker.example:80/announce"]
|
||||
assert after.metainfo["info"]["pieces"] == before.metainfo["info"]["pieces"]
|
||||
|
||||
|
||||
def test_modify_preserves_omitted_fields_and_original(tmp_path: Path) -> None:
|
||||
source = tmp_path / "content.txt"
|
||||
source.write_bytes(b"content")
|
||||
original = script.create_torrent(paths=(source,), output=tmp_path / "original.torrent")
|
||||
original_bytes = original.read_bytes()
|
||||
output = tmp_path / "modified.torrent"
|
||||
script.modify_torrent(torrent_path=original, output=output, author="replacement")
|
||||
before = torf.Torrent.read(filepath=original)
|
||||
after = torf.Torrent.read(filepath=output)
|
||||
assert original.read_bytes() == original_bytes
|
||||
assert after.infohash == before.infohash
|
||||
assert after.private == before.private
|
||||
assert after.trackers == before.trackers
|
||||
assert after.name == before.name
|
||||
with pytest.raises(torf.WriteError):
|
||||
script.modify_torrent(torrent_path=original, output=output)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("added_key", ["torrent-added", "torrent-duplicate"])
|
||||
@pytest.mark.parametrize("url", ["http://transmission:9091", "http://transmission:9091/transmission/rpc"])
|
||||
def test_enqueue_negotiates_session_and_sets_remote_directory(tmp_path: Path, added_key: str, url: str) -> None:
|
||||
torrent_path = tmp_path / "content.torrent"
|
||||
torrent_path.write_bytes(b"torrent data")
|
||||
requests = []
|
||||
|
||||
def respond(request: httpx.Request) -> httpx.Response:
|
||||
requests.append(request)
|
||||
if len(requests) == 1:
|
||||
return httpx.Response(status_code=409, headers={"X-Transmission-Session-Id": "session-token"})
|
||||
return httpx.Response(status_code=200, json={
|
||||
"result": "success", "arguments": {added_key: {"hashString": "torrent-hash"}},
|
||||
})
|
||||
|
||||
client = httpx.Client(transport=httpx.MockTransport(handler=respond), auth=httpx.BasicAuth("user", "pass:word"))
|
||||
with patch.object(script.httpx, "Client", return_value=client) as factory:
|
||||
result = script.enqueue_torrent(
|
||||
transmission_url=url, credentials="user:pass:word",
|
||||
torrent_path=torrent_path, download_path="/remote/downloads",
|
||||
)
|
||||
assert result == "torrent-hash"
|
||||
assert factory.call_args.kwargs["auth"]._auth_header == client.auth._auth_header
|
||||
assert len(requests) == 2
|
||||
assert requests[1].headers["X-Transmission-Session-Id"] == "session-token"
|
||||
assert str(requests[1].url) == "http://transmission:9091/transmission/rpc"
|
||||
payload = json.loads(requests[1].content)
|
||||
assert payload["method"] == "torrent-add"
|
||||
assert payload["arguments"]["download-dir"] == "/remote/downloads"
|
||||
assert payload["arguments"]["paused"] is False
|
||||
assert base64.b64decode(payload["arguments"]["metainfo"]) == torrent_path.read_bytes()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("status, result", [(401, "success"), (200, "invalid torrent"), (409, "success")])
|
||||
def test_enqueue_reports_failures(tmp_path: Path, status: int, result: str) -> None:
|
||||
torrent_path = tmp_path / "content.torrent"
|
||||
torrent_path.write_bytes(b"torrent data")
|
||||
|
||||
def respond(request: httpx.Request) -> httpx.Response:
|
||||
return httpx.Response(status_code=status, headers={"X-Transmission-Session-Id": "token"}, json={"result": result})
|
||||
|
||||
client = httpx.Client(transport=httpx.MockTransport(handler=respond))
|
||||
with patch.object(script.httpx, "Client", return_value=client):
|
||||
with pytest.raises((httpx.HTTPStatusError, ValueError)):
|
||||
script.enqueue_torrent(
|
||||
transmission_url="http://transmission:9091", credentials=None,
|
||||
torrent_path=torrent_path, download_path="/remote/downloads",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("use_environment", [False, True])
|
||||
def test_enqueue_command(tmp_path: Path, use_environment: bool) -> None:
|
||||
torrent_path = tmp_path / "content.torrent"
|
||||
torrent_path.write_bytes(b"torrent data")
|
||||
arguments = ["enqueue", str(torrent_path), "--download-path", "/remote/downloads"]
|
||||
environment = {"TRANSMISSION_HOST": "http://transmission:9091", "TRANSMISSION_AUTH": "user:password"}
|
||||
if not use_environment:
|
||||
arguments.extend(["--transmission-url", environment["TRANSMISSION_HOST"], "--credentials", environment["TRANSMISSION_AUTH"]])
|
||||
with patch.object(script, "enqueue_torrent", return_value="torrent-hash") as enqueue:
|
||||
result = CliRunner().invoke(cli=script.main, args=arguments, env=environment if use_environment else {})
|
||||
assert result.exit_code == 0, result.output
|
||||
enqueue.assert_called_once_with(
|
||||
transmission_url="http://transmission:9091", credentials="user:password",
|
||||
torrent_path=torrent_path, download_path=Path("/remote/downloads"),
|
||||
)
|
||||
|
||||
|
||||
def test_enqueue_command_reports_error(tmp_path: Path) -> None:
|
||||
torrent_path = tmp_path / "content.torrent"
|
||||
torrent_path.write_bytes(b"torrent data")
|
||||
with patch.object(script, "enqueue_torrent", side_effect=ValueError("Transmission rejected torrent")):
|
||||
result = CliRunner().invoke(cli=script.main, args=[
|
||||
"enqueue", str(torrent_path), "--host", "http://transmission:9091",
|
||||
"--download-path", "/remote/downloads",
|
||||
])
|
||||
assert result.exit_code == 1
|
||||
assert "Error: Transmission rejected torrent" in result.output
|
||||
|
||||
|
||||
@pytest.mark.parametrize("directory", [False, True])
|
||||
@pytest.mark.parametrize("json_output", [False, True])
|
||||
def test_inspect_reports_properties_and_files_without_content(
|
||||
tmp_path: Path, directory: bool, json_output: bool,
|
||||
) -> None:
|
||||
folder = tmp_path / "content"
|
||||
folder.mkdir()
|
||||
source = folder / "first.txt"
|
||||
source.write_bytes(b"first")
|
||||
if directory:
|
||||
(folder / "second.txt").write_bytes(b"second")
|
||||
output = script.create_torrent(paths=(folder if directory else source,), output=tmp_path / "content.torrent")
|
||||
original = output.read_bytes()
|
||||
source.unlink()
|
||||
arguments = ["inspect", str(output)]
|
||||
if json_output:
|
||||
arguments.append("--json")
|
||||
result = CliRunner().invoke(cli=script.main, args=arguments)
|
||||
assert result.exit_code == 0, result.output
|
||||
assert result.stderr == ""
|
||||
if json_output:
|
||||
report = json.loads(result.stdout)
|
||||
assert report["properties"]["author"] == script.AUTHOR
|
||||
assert report["properties"]["private"] is True
|
||||
assert report["properties"]["announce_urls"] == [list(script.ANNOUNCE_URLS)]
|
||||
assert report["properties"]["size"] == (11 if directory else 5)
|
||||
assert report["files"] == (
|
||||
[{"path": "content/first.txt", "size": 5}, {"path": "content/second.txt", "size": 6}]
|
||||
if directory else [{"path": "first.txt", "size": 5}]
|
||||
)
|
||||
else:
|
||||
assert f"Author: {script.AUTHOR}\n" in result.stdout
|
||||
assert "Private: yes\n" in result.stdout
|
||||
assert f"Size: {11 if directory else 5} bytes\n" in result.stdout
|
||||
assert f"Tier 1: {script.ANNOUNCE_URLS[0]}\n" in result.stdout
|
||||
assert f"Files ({2 if directory else 1}):\n" in result.stdout
|
||||
assert f" {'content/' if directory else ''}first.txt: 5 bytes\n" in result.stdout
|
||||
if directory:
|
||||
assert " content/second.txt: 6 bytes\n" in result.stdout
|
||||
assert output.read_bytes() == original
|
||||
|
||||
|
||||
def test_inspect_rejects_invalid_torrent(tmp_path: Path) -> None:
|
||||
torrent_path = tmp_path / "invalid.torrent"
|
||||
torrent_path.write_bytes(b"invalid")
|
||||
result = CliRunner().invoke(cli=script.main, args=["inspect", str(torrent_path)])
|
||||
assert result.exit_code == 1
|
||||
assert "Error:" in result.output
|
||||
|
||||
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main(args=[str(Path(__file__).resolve()), *sys.argv[1:]]))
|
||||
Reference in new issue
Block a user