#!/usr/bin/env python3 """prime_client.py — a template client for the PRIME web server. Submits one structure, waits for the job, and writes the results. Batch submission (an archive, or several ids at once) is disabled server-side. Standard library only: no installation, no account, no API key. # one structure, by PDB id python3 prime_client.py --id 12CA --metals ZN # one structure, from a local file, every supported ion python3 prime_client.py --file 12ca.cif --metals ALL -o out/ Exit status is 0 when the job succeeded, 1 otherwise. """ import argparse import json import mimetypes import os import sys import time import urllib.error import urllib.request import uuid BASE = os.environ.get("PRIME_URL", "https://onodalab.ees.hokudai.ac.jp/prime") # --------------------------------------------------------------- transport def _multipart(fields, files): """Encode a multipart/form-data body without requiring `requests`.""" boundary = "----prime-" + uuid.uuid4().hex out = [] for k, v in fields.items(): if v is None or v == "": continue out.append(f"--{boundary}\r\nContent-Disposition: form-data; name=\"{k}\"\r\n\r\n{v}\r\n".encode()) for k, path in files.items(): name = os.path.basename(path) ctype = mimetypes.guess_type(name)[0] or "application/octet-stream" with open(path, "rb") as fh: body = fh.read() out.append( f"--{boundary}\r\nContent-Disposition: form-data; name=\"{k}\"; " f"filename=\"{name}\"\r\nContent-Type: {ctype}\r\n\r\n".encode() + body + b"\r\n") out.append(f"--{boundary}--\r\n".encode()) return b"".join(out), f"multipart/form-data; boundary={boundary}" def _request(url, data=None, ctype=None, timeout=120): req = urllib.request.Request(url, data=data) if ctype: req.add_header("Content-Type", ctype) try: with urllib.request.urlopen(req, timeout=timeout) as r: return r.status, r.read() except urllib.error.HTTPError as e: return e.code, e.read() def _json(url, **kw): status, body = _request(url, **kw) try: return status, json.loads(body.decode()) except ValueError: return status, {"message": body[:400].decode(errors="replace")} # ------------------------------------------------------------------- calls def submit(metals, path=None, ids=None, chain=None, top_k=None, max_rot=None, cluster_dist=None, rot_seed=0): """Start a job. Returns its id.""" fields = {"metal_type": metals, "chain": chain, "top_k": top_k, "max_rot": max_rot, "rot_seed": rot_seed, "cluster_dist": cluster_dist} if ids: fields["protein_id"] = ids files = {"structure_file": path} if path else {} body, ctype = _multipart(fields, files) status, j = _json(f"{BASE}/api/jobs", data=body, ctype=ctype, timeout=300) if status >= 300 or "job_id" not in j: sys.exit(f"submission failed ({status}): {j.get('message', j)}") return j["job_id"] def wait(job_id, poll=2.0, quiet=False): """Block until the job leaves the queue, printing progress as it goes.""" while True: status, j = _json(f"{BASE}/api/jobs/{job_id}") st = j.get("status") if st in ("succeeded", "failed", "canceled", "expired"): return st if not quiet: p = j.get("progress") or {} pos = f" (queue position {j['position']})" if j.get("position") else "" msg = p.get("message") or st or "working" sys.stderr.write(f"\r{msg}{pos} pct={p.get('pct', 0):.0f}% ") sys.stderr.flush() time.sleep(poll) def fetch(job_id): status, j = _json(f"{BASE}/api/jobs/{job_id}/result", timeout=300) if status >= 300: sys.exit(f"could not fetch the result ({status}): {j.get('message', j)}") return j def download(job_id, name, outdir): """Save one of the job's artifacts (predicted.pdb, results.csv, ...).""" status, body = _request(f"{BASE}/artifact/{job_id}/{name}", timeout=300) if status >= 300: return None dest = os.path.join(outdir, name) os.makedirs(os.path.dirname(dest) or ".", exist_ok=True) with open(dest, "wb") as fh: fh.write(body) return dest # -------------------------------------------------------------------- main def main(): ap = argparse.ArgumentParser(description="PRIME web-server client") src = ap.add_mutually_exclusive_group(required=True) src.add_argument("--file", help="one structure (.pdb/.cif)") src.add_argument("--id", dest="ids", help="one PDB or UniProt id") ap.add_argument("--metals", default="ZN", help="ion, comma-separated list, or ALL (default: ZN)") ap.add_argument("--chain", default=None, help="restrict to one chain id") ap.add_argument("--top-k", dest="top_k", default=None, help="probes kept per chain (default 0, all probes)") ap.add_argument("--max-rot", dest="max_rot", default=None, help="rotations per probe, 1-5 (default 5)") ap.add_argument("--rot-seed", dest="rot_seed", type=int, default=0, help="seed for the fixed rotation set (default 0)") ap.add_argument("--cluster-dist", dest="cluster_dist", default=None, help="distance in A within which predicted positions merge " "into one site, 0.5-10 (default 3.0; try 1.0 for ion " "channels)") ap.add_argument("-o", "--outdir", default=".", help="where to write outputs") ap.add_argument("-q", "--quiet", action="store_true") a = ap.parse_args() job = submit(a.metals, path=a.file, ids=a.ids, chain=a.chain, top_k=a.top_k, max_rot=a.max_rot, cluster_dist=a.cluster_dist, rot_seed=a.rot_seed) print(f"job {job}\nresult page: {BASE}/results/{job}", file=sys.stderr) st = wait(job, quiet=a.quiet) if not a.quiet: sys.stderr.write("\n") if st != "succeeded": sys.exit(f"job {st}") res = fetch(job) os.makedirs(a.outdir, exist_ok=True) with open(os.path.join(a.outdir, f"{job}.json"), "w") as fh: json.dump(res, fh, indent=1) if res.get("batch"): stats = res.get("stats") or {} print(f"{stats.get('n_structures', '?')} structures, " f"{res.get('count', 0)} sites, {stats.get('seconds', '?')} s") for name in ("results.csv", "results.json", "predictions.zip"): got = download(job, name, a.outdir) if got: print(f"wrote {got}") return sites = res.get("binding_sites") or [] print(f"{len(sites)} site(s)") for s in sites: x, y, z = s["position"] print(f" {s['element']:>4s} chain {s.get('chain', '?'):<2s} " f"{x:8.3f} {y:8.3f} {z:8.3f} p={s['confidence']:.3f}") for name in ("predicted.pdb", "complex.pdb", "sites.csv"): got = download(job, name, a.outdir) if got: print(f"wrote {got}") if __name__ == "__main__": main()