FrameBridge 0.3.1 community source release
This commit is contained in:
commit
b23819fff1
1354 files changed
+343716
No files matched your search
+425
@@ -0,0 +1,425 @@
|
||||
#!/usr/bin/env python3
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import socket
|
||||
import stat
|
||||
import struct
|
||||
import subprocess
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
VERSION = 1
|
||||
MAX_SCRIPT_BYTES = 262144
|
||||
MAX_REQUEST_BYTES = 393216
|
||||
JOB_RE = re.compile(r"^[A-Za-z0-9._-]{1,80}$")
|
||||
|
||||
HOME = Path.home()
|
||||
FBROOT = HOME / ".local/share/framebridge/fbterm"
|
||||
ENGINE = FBROOT / "bin/fbterm_engine.py"
|
||||
|
||||
ROOT = Path(
|
||||
os.environ.get(
|
||||
"FRAMEBRIDGE_TOOL_ROOT",
|
||||
str(HOME / ".local/share/framebridge/tool-adapter")
|
||||
)
|
||||
)
|
||||
SOCKET_PATH = ROOT / "adapter.sock"
|
||||
INBOX = ROOT / "inbox"
|
||||
PID_FILE = ROOT / "adapter.pid"
|
||||
NATIVE_HOST = HOME / "Descargas/framebridge_v0_1/host/framebridge_host.py"
|
||||
|
||||
|
||||
def prepare_root():
|
||||
ROOT.mkdir(parents=True, exist_ok=True)
|
||||
INBOX.mkdir(parents=True, exist_ok=True)
|
||||
os.chmod(ROOT, 0o700)
|
||||
os.chmod(INBOX, 0o700)
|
||||
|
||||
|
||||
def recv_request(conn):
|
||||
data = bytearray()
|
||||
|
||||
while True:
|
||||
chunk = conn.recv(65536)
|
||||
if not chunk:
|
||||
break
|
||||
|
||||
data.extend(chunk)
|
||||
|
||||
if len(data) > MAX_REQUEST_BYTES:
|
||||
raise ValueError("REQUEST_TOO_LARGE")
|
||||
|
||||
if b"\n" in data:
|
||||
break
|
||||
|
||||
raw = bytes(data).split(b"\n", 1)[0]
|
||||
|
||||
if not raw:
|
||||
raise ValueError("EMPTY_REQUEST")
|
||||
|
||||
return json.loads(raw.decode("utf-8"))
|
||||
|
||||
|
||||
def send_response(conn, obj):
|
||||
raw = (
|
||||
json.dumps(
|
||||
obj,
|
||||
ensure_ascii=True,
|
||||
sort_keys=True,
|
||||
separators=(",", ":"),
|
||||
)
|
||||
+ "\n"
|
||||
).encode("utf-8")
|
||||
|
||||
try:
|
||||
conn.sendall(raw)
|
||||
except (BrokenPipeError, ConnectionResetError):
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def peer_uid(conn):
|
||||
fmt = "3i"
|
||||
raw = conn.getsockopt(
|
||||
socket.SOL_SOCKET,
|
||||
socket.SO_PEERCRED,
|
||||
struct.calcsize(fmt),
|
||||
)
|
||||
_pid, uid, _gid = struct.unpack(fmt, raw)
|
||||
return uid
|
||||
|
||||
|
||||
def last_result_final(stdout_path):
|
||||
try:
|
||||
lines = Path(stdout_path).read_text(
|
||||
encoding="utf-8",
|
||||
errors="replace",
|
||||
).splitlines()
|
||||
except OSError:
|
||||
return None
|
||||
|
||||
for line in reversed(lines):
|
||||
if line.startswith("RESULTADO_FINAL="):
|
||||
return line.split("=", 1)[1].strip()[:2000]
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def engine_result(stdout):
|
||||
for line in reversed(stdout.splitlines()):
|
||||
line = line.strip()
|
||||
|
||||
if not line.startswith("{"):
|
||||
continue
|
||||
|
||||
try:
|
||||
obj = json.loads(line)
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
|
||||
if isinstance(obj, dict):
|
||||
return obj
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def native_cua_roundtrip(payload, timeout_s=60):
|
||||
if not NATIVE_HOST.is_file():
|
||||
raise ValueError("FRAMEBRIDGE_NATIVE_HOST_NOT_FOUND")
|
||||
|
||||
raw = json.dumps(
|
||||
payload,
|
||||
ensure_ascii=True,
|
||||
separators=(",", ":"),
|
||||
).encode("utf-8")
|
||||
|
||||
frame = struct.pack("<I", len(raw)) + raw
|
||||
|
||||
try:
|
||||
proc = subprocess.run(
|
||||
[str(NATIVE_HOST)],
|
||||
input=frame,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.PIPE,
|
||||
timeout=timeout_s,
|
||||
check=False,
|
||||
)
|
||||
except subprocess.TimeoutExpired:
|
||||
raise ValueError("CUA_NATIVE_HOST_TIMEOUT")
|
||||
|
||||
out = proc.stdout or b""
|
||||
|
||||
if len(out) < 4:
|
||||
raise ValueError("CUA_NATIVE_RESPONSE_HEADER_MISSING")
|
||||
|
||||
size = struct.unpack("<I", out[:4])[0]
|
||||
|
||||
if size <= 0 or size > 2_000_000:
|
||||
raise ValueError("CUA_NATIVE_RESPONSE_SIZE_INVALID")
|
||||
|
||||
body = out[4:4 + size]
|
||||
|
||||
if len(body) != size:
|
||||
raise ValueError("CUA_NATIVE_RESPONSE_TRUNCATED")
|
||||
|
||||
try:
|
||||
response = json.loads(body.decode("utf-8"))
|
||||
except Exception:
|
||||
raise ValueError("CUA_NATIVE_RESPONSE_INVALID_JSON")
|
||||
|
||||
if not isinstance(response, dict):
|
||||
raise ValueError("CUA_NATIVE_RESPONSE_NOT_OBJECT")
|
||||
|
||||
return response
|
||||
|
||||
|
||||
def handle_cua_request(req, job_id):
|
||||
timeout_s = int(req.get("timeout_s", 60))
|
||||
|
||||
if not 1 <= timeout_s <= 180:
|
||||
raise ValueError("INVALID_CUA_TIMEOUT")
|
||||
|
||||
cua = req.get("cua")
|
||||
|
||||
if not isinstance(cua, dict):
|
||||
raise ValueError("CUA_PAYLOAD_REQUIRED")
|
||||
|
||||
payload = dict(cua)
|
||||
payload["op"] = "cua"
|
||||
payload["id"] = job_id
|
||||
payload["target"] = "ASUS"
|
||||
payload["authorized"] = "YES"
|
||||
|
||||
response = native_cua_roundtrip(
|
||||
payload,
|
||||
timeout_s=timeout_s,
|
||||
)
|
||||
|
||||
return {
|
||||
"version": VERSION,
|
||||
"ok": bool(response.get("executed")),
|
||||
"job_id": job_id,
|
||||
"mode": "CUA",
|
||||
"action": str(payload.get("action", ""))[:32],
|
||||
"channel": str(response.get("channel", "CUA-DIRECT")),
|
||||
"exit_code": response.get("exit_code"),
|
||||
"frame_output": str(response.get("output", "")),
|
||||
"error": str(response.get("error", "")),
|
||||
}
|
||||
|
||||
|
||||
def handle(req):
|
||||
if not isinstance(req, dict):
|
||||
raise ValueError("REQUEST_NOT_OBJECT")
|
||||
|
||||
if req.get("version") != VERSION:
|
||||
raise ValueError("UNSUPPORTED_VERSION")
|
||||
|
||||
job_id = str(req.get("job_id", ""))
|
||||
|
||||
if not JOB_RE.fullmatch(job_id):
|
||||
raise ValueError("INVALID_JOB_ID")
|
||||
|
||||
mode = str(req.get("mode", "EXEC")).upper()
|
||||
|
||||
if mode != "EXEC":
|
||||
raise ValueError("UNSUPPORTED_MODE")
|
||||
|
||||
if req.get("authorized") is not True:
|
||||
raise ValueError("AUTHORIZED_TRUE_REQUIRED")
|
||||
|
||||
request_type = str(
|
||||
req.get("request_type", "exec")
|
||||
).lower().strip()
|
||||
|
||||
if request_type == "cua":
|
||||
return handle_cua_request(req, job_id)
|
||||
|
||||
if request_type != "exec":
|
||||
raise ValueError("UNSUPPORTED_REQUEST_TYPE")
|
||||
|
||||
script = req.get("script")
|
||||
|
||||
if not isinstance(script, str) or not script:
|
||||
raise ValueError("SCRIPT_REQUIRED")
|
||||
|
||||
script_bytes = script.encode("utf-8")
|
||||
|
||||
if len(script_bytes) > MAX_SCRIPT_BYTES:
|
||||
raise ValueError("SCRIPT_TOO_LARGE")
|
||||
|
||||
timeout_s = int(req.get("timeout_s", 120))
|
||||
|
||||
if not 1 <= timeout_s <= 600:
|
||||
raise ValueError("INVALID_TIMEOUT")
|
||||
|
||||
digest = hashlib.sha256(script_bytes).hexdigest()
|
||||
|
||||
expected = req.get("sha256")
|
||||
|
||||
if expected is not None:
|
||||
expected = str(expected).lower()
|
||||
|
||||
if digest != expected:
|
||||
raise ValueError("SCRIPT_SHA256_MISMATCH")
|
||||
|
||||
command_path = INBOX / f"{job_id}.bash"
|
||||
temp_path = INBOX / f".{job_id}.tmp"
|
||||
|
||||
temp_path.write_bytes(script_bytes)
|
||||
os.chmod(temp_path, 0o600)
|
||||
temp_path.replace(command_path)
|
||||
|
||||
proc = subprocess.run(
|
||||
[
|
||||
sys.executable,
|
||||
str(ENGINE),
|
||||
"exec",
|
||||
job_id,
|
||||
str(timeout_s),
|
||||
str(command_path),
|
||||
],
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.PIPE,
|
||||
text=True,
|
||||
timeout=timeout_s + 45,
|
||||
check=False,
|
||||
env={
|
||||
**os.environ,
|
||||
"FBTERM_ROOT": str(FBROOT),
|
||||
"FBTERM_SESSION": os.environ.get(
|
||||
"FBTERM_SESSION",
|
||||
"framebridge-fbterm-r1",
|
||||
),
|
||||
"FBTERM_TITLE": os.environ.get(
|
||||
"FBTERM_TITLE",
|
||||
"FrameBridge Terminal R1",
|
||||
),
|
||||
},
|
||||
)
|
||||
|
||||
result = engine_result(proc.stdout)
|
||||
|
||||
if result is None:
|
||||
return {
|
||||
"version": VERSION,
|
||||
"ok": False,
|
||||
"job_id": job_id,
|
||||
"adapter_error": "ENGINE_RESULT_NOT_FOUND",
|
||||
"engine_rc": proc.returncode,
|
||||
"engine_stderr": proc.stderr[-2000:],
|
||||
"command_sha256": digest,
|
||||
}
|
||||
|
||||
stdout_path = result.get("STDOUT_PATH")
|
||||
stderr_path = result.get("STDERR_PATH")
|
||||
|
||||
return {
|
||||
"version": VERSION,
|
||||
"ok": True,
|
||||
"job_id": job_id,
|
||||
"mode": mode,
|
||||
"command_sha256": digest,
|
||||
"engine_rc": proc.returncode,
|
||||
"exit_code": result.get("EXIT"),
|
||||
"result_final": (
|
||||
last_result_final(stdout_path)
|
||||
if stdout_path
|
||||
else None
|
||||
),
|
||||
"stdout_path": stdout_path,
|
||||
"stderr_path": stderr_path,
|
||||
"stdout_sha256": result.get("STDOUT_SHA256"),
|
||||
"stderr_sha256": result.get("STDERR_SHA256"),
|
||||
"command_hash_gate": result.get(
|
||||
"COMMAND_HASH_GATE"
|
||||
),
|
||||
"pty": result.get("PTY"),
|
||||
"pane_pid": result.get("PANE_PID"),
|
||||
"pane_tty": result.get("PANE_TTY"),
|
||||
"shell_ready": result.get("SHELL_READY"),
|
||||
}
|
||||
|
||||
|
||||
def serve():
|
||||
prepare_root()
|
||||
|
||||
if not ENGINE.is_file():
|
||||
raise RuntimeError("FBTERM_ENGINE_NOT_FOUND")
|
||||
|
||||
if SOCKET_PATH.exists():
|
||||
st = SOCKET_PATH.lstat()
|
||||
|
||||
if not stat.S_ISSOCK(st.st_mode):
|
||||
raise RuntimeError("SOCKET_PATH_OCCUPIED")
|
||||
|
||||
SOCKET_PATH.unlink()
|
||||
|
||||
server = socket.socket(
|
||||
socket.AF_UNIX,
|
||||
socket.SOCK_STREAM,
|
||||
)
|
||||
|
||||
server.bind(str(SOCKET_PATH))
|
||||
os.chmod(SOCKET_PATH, 0o600)
|
||||
server.listen(4)
|
||||
|
||||
PID_FILE.write_text(
|
||||
str(os.getpid()) + "\n",
|
||||
encoding="ascii",
|
||||
)
|
||||
os.chmod(PID_FILE, 0o600)
|
||||
|
||||
try:
|
||||
while True:
|
||||
conn, _ = server.accept()
|
||||
|
||||
with conn:
|
||||
try:
|
||||
if peer_uid(conn) != os.getuid():
|
||||
raise PermissionError(
|
||||
"PEER_UID_MISMATCH"
|
||||
)
|
||||
|
||||
req = recv_request(conn)
|
||||
resp = handle(req)
|
||||
|
||||
except Exception as exc:
|
||||
resp = {
|
||||
"version": VERSION,
|
||||
"ok": False,
|
||||
"adapter_error": (
|
||||
f"{type(exc).__name__}:"
|
||||
f"{exc}"
|
||||
)[:2000],
|
||||
}
|
||||
|
||||
send_response(conn, resp)
|
||||
|
||||
finally:
|
||||
server.close()
|
||||
|
||||
try:
|
||||
SOCKET_PATH.unlink()
|
||||
except FileNotFoundError:
|
||||
pass
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
action = sys.argv[1] if len(sys.argv) > 1 else ""
|
||||
|
||||
if action == "serve":
|
||||
serve()
|
||||
else:
|
||||
print(
|
||||
json.dumps(
|
||||
{
|
||||
"ok": False,
|
||||
"error": "USAGE: serve",
|
||||
},
|
||||
sort_keys=True,
|
||||
)
|
||||
)
|
||||
raise SystemExit(64)
|
||||
File diff suppressed because it is too large.
Load diff
+1
@@ -0,0 +1 @@
|
||||
__version__ = "0.1.0-alpha.1"
|
||||
+276
@@ -0,0 +1,276 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import hashlib
|
||||
import json
|
||||
import secrets
|
||||
import socket
|
||||
import threading
|
||||
import time
|
||||
import urllib.parse
|
||||
from http.server import BaseHTTPRequestHandler, HTTPServer
|
||||
from typing import Any
|
||||
|
||||
AUTHORIZATION_ENDPOINT="https://auth.openai.com/api/accounts/authorize"
|
||||
TOKEN_ENDPOINT="https://auth.openai.com/api/accounts/oauth/token"
|
||||
DISCOVERY_ENDPOINT="https://auth.openai.com/.well-known/openid-configuration"
|
||||
RESOURCE="https://api.openai.com/v1"
|
||||
DYNAMIC_CLIENT_ID="dynamic_agent_client"
|
||||
DEFAULT_SCOPES=("openid","profile","email","offline_access","resource.invoke","chatgpt.tokens.use.direct")
|
||||
AGENT_NAME="FrameBridge"
|
||||
|
||||
|
||||
class ChatGPTAuthError(RuntimeError):
|
||||
pass
|
||||
|
||||
|
||||
def discover_openai(*, requests_module=None) -> dict[str,str]:
|
||||
if requests_module is None:
|
||||
import requests as requests_module
|
||||
response=requests_module.get(DISCOVERY_ENDPOINT,timeout=15)
|
||||
if response.status_code!=200:
|
||||
raise ChatGPTAuthError("CHATGPT_DISCOVERY_FAILED")
|
||||
value=response.json()
|
||||
if not isinstance(value,dict):
|
||||
raise ChatGPTAuthError("CHATGPT_DISCOVERY_INVALID")
|
||||
issuer=str(value.get("issuer") or "")
|
||||
auth=str(value.get("authorization_endpoint") or "")
|
||||
token=str(value.get("token_endpoint") or "")
|
||||
jwks=str(value.get("jwks_uri") or "")
|
||||
if issuer!="https://auth.openai.com":
|
||||
raise ChatGPTAuthError("CHATGPT_DISCOVERY_INVALID")
|
||||
for url in (auth,token,jwks):
|
||||
if not url.startswith("https://auth.openai.com/"):
|
||||
raise ChatGPTAuthError("CHATGPT_DISCOVERY_INVALID")
|
||||
return {
|
||||
"issuer":issuer,
|
||||
"authorization_endpoint":auth,
|
||||
"token_endpoint":token,
|
||||
"jwks_uri":jwks,
|
||||
}
|
||||
|
||||
|
||||
def _b64url(data: bytes) -> str:
|
||||
return base64.urlsafe_b64encode(data).decode("ascii").rstrip("=")
|
||||
|
||||
|
||||
def _random(n: int=48) -> str:
|
||||
return secrets.token_urlsafe(n)
|
||||
|
||||
|
||||
def account_binding_id(issuer: str, subject: str) -> str:
|
||||
raw=(str(issuer)+"\0"+str(subject)).encode("utf-8")
|
||||
return "fbacct_"+hashlib.sha256(raw).hexdigest()[:40]
|
||||
|
||||
|
||||
def build_authorization(
|
||||
host_id: str,
|
||||
redirect_uri: str,
|
||||
*,
|
||||
scopes: tuple[str,...]=DEFAULT_SCOPES,
|
||||
authorization_endpoint: str=AUTHORIZATION_ENDPOINT,
|
||||
client_id: str=DYNAMIC_CLIENT_ID,
|
||||
) -> tuple[str,dict[str,str]]:
|
||||
host_id=str(host_id).strip()
|
||||
if not host_id:
|
||||
raise ChatGPTAuthError("CHATGPT_HOST_ID_REQUIRED")
|
||||
state=_random(32)
|
||||
nonce=_random(32)
|
||||
verifier=_random(64)
|
||||
challenge=_b64url(hashlib.sha256(verifier.encode("ascii")).digest())
|
||||
client_id=str(client_id or DYNAMIC_CLIENT_ID)
|
||||
tx={
|
||||
"state":state,
|
||||
"nonce":nonce,
|
||||
"code_verifier":verifier,
|
||||
"redirect_uri":str(redirect_uri),
|
||||
"created_at":str(time.time()),
|
||||
"client_id":client_id,
|
||||
}
|
||||
query={
|
||||
"client_id":client_id,
|
||||
"ext_agent_host_id":host_id,
|
||||
"response_type":"code",
|
||||
"redirect_uri":str(redirect_uri),
|
||||
"scope":" ".join(scopes),
|
||||
"resource":RESOURCE,
|
||||
"state":state,
|
||||
"nonce":nonce,
|
||||
"code_challenge_method":"S256",
|
||||
"code_challenge":challenge,
|
||||
}
|
||||
if client_id==DYNAMIC_CLIENT_ID:
|
||||
query["agent_name_hint"]=AGENT_NAME
|
||||
return str(authorization_endpoint)+"?"+urllib.parse.urlencode(query),tx
|
||||
|
||||
|
||||
class _Callback:
|
||||
def __init__(self):
|
||||
self.event=threading.Event()
|
||||
self.query: dict[str,str]={}
|
||||
|
||||
|
||||
def _handler(callback: _Callback):
|
||||
class Handler(BaseHTTPRequestHandler):
|
||||
def log_message(self,fmt,*args):
|
||||
return
|
||||
def do_GET(self):
|
||||
p=urllib.parse.urlparse(self.path)
|
||||
if p.path!="/auth/callback":
|
||||
self.send_response(404);self.end_headers();return
|
||||
q=urllib.parse.parse_qs(p.query,keep_blank_values=True)
|
||||
callback.query={k:(v[0] if v else "") for k,v in q.items()}
|
||||
callback.event.set()
|
||||
body=(
|
||||
b"<!doctype html><meta charset=utf-8>"
|
||||
b"<title>FrameBridge</title>"
|
||||
b"<h2>FrameBridge</h2>"
|
||||
b"<p>ChatGPT authentication returned to FrameBridge.</p>"
|
||||
b"<p>You may close this browser tab.</p>"
|
||||
)
|
||||
self.send_response(200)
|
||||
self.send_header("Content-Type","text/html; charset=utf-8")
|
||||
self.send_header("Cache-Control","no-store")
|
||||
self.send_header("Content-Length",str(len(body)))
|
||||
self.end_headers()
|
||||
self.wfile.write(body)
|
||||
return Handler
|
||||
|
||||
|
||||
def start_loopback_listener() -> tuple[HTTPServer,_Callback,str]:
|
||||
cb=_Callback()
|
||||
server=HTTPServer(("127.0.0.1",0),_handler(cb))
|
||||
port=int(server.server_address[1])
|
||||
uri=f"http://127.0.0.1:{port}/auth/callback"
|
||||
thread=threading.Thread(target=server.serve_forever,daemon=True)
|
||||
thread.start()
|
||||
return server,cb,uri
|
||||
|
||||
|
||||
def exchange_and_validate(
|
||||
callback_query: dict[str,str],
|
||||
tx: dict[str,str],
|
||||
*,
|
||||
requests_module=None,
|
||||
jwt_module=None,
|
||||
provider_metadata: dict[str,str] | None=None,
|
||||
) -> dict[str,Any]:
|
||||
if callback_query.get("state") != tx.get("state"):
|
||||
raise ChatGPTAuthError("CHATGPT_STATE_MISMATCH")
|
||||
if callback_query.get("error"):
|
||||
raise ChatGPTAuthError("CHATGPT_AUTH_DENIED")
|
||||
code=str(callback_query.get("code") or "")
|
||||
requested_client_id=str(tx.get("client_id") or DYNAMIC_CLIENT_ID)
|
||||
callback_client_id=str(callback_query.get("client_id") or "")
|
||||
if requested_client_id==DYNAMIC_CLIENT_ID:
|
||||
issued_client_id=callback_client_id
|
||||
if not code or not issued_client_id or issued_client_id==DYNAMIC_CLIENT_ID:
|
||||
raise ChatGPTAuthError("CHATGPT_CALLBACK_INCOMPLETE")
|
||||
else:
|
||||
issued_client_id=requested_client_id
|
||||
if callback_client_id and callback_client_id!=requested_client_id:
|
||||
raise ChatGPTAuthError("CHATGPT_AUTH_CLIENT_MISMATCH")
|
||||
if not code:
|
||||
raise ChatGPTAuthError("CHATGPT_CALLBACK_INCOMPLETE")
|
||||
|
||||
if requests_module is None:
|
||||
import requests as requests_module
|
||||
if jwt_module is None:
|
||||
import jwt as jwt_module
|
||||
|
||||
provider_metadata=provider_metadata or discover_openai(requests_module=requests_module)
|
||||
token_endpoint=provider_metadata["token_endpoint"]
|
||||
issuer=provider_metadata["issuer"]
|
||||
jwks_uri=provider_metadata["jwks_uri"]
|
||||
|
||||
token_resp=requests_module.post(
|
||||
token_endpoint,
|
||||
data={
|
||||
"grant_type":"authorization_code",
|
||||
"client_id":issued_client_id,
|
||||
"code":code,
|
||||
"code_verifier":tx["code_verifier"],
|
||||
"redirect_uri":tx["redirect_uri"],
|
||||
"resource":RESOURCE,
|
||||
},
|
||||
timeout=20,
|
||||
)
|
||||
if token_resp.status_code != 200:
|
||||
raise ChatGPTAuthError("CHATGPT_TOKEN_EXCHANGE_FAILED")
|
||||
tokens=token_resp.json()
|
||||
id_token=str(tokens.get("id_token") or "")
|
||||
if not id_token:
|
||||
raise ChatGPTAuthError("CHATGPT_ID_TOKEN_MISSING")
|
||||
|
||||
jwk_client=jwt_module.PyJWKClient(jwks_uri)
|
||||
signing_key=jwk_client.get_signing_key_from_jwt(id_token)
|
||||
alg=str(getattr(signing_key,"algorithm_name","") or "")
|
||||
if not alg:
|
||||
try:
|
||||
alg=str(jwt_module.get_unverified_header(id_token).get("alg") or "")
|
||||
except Exception:
|
||||
alg=""
|
||||
if alg not in {"RS256","ES256","EdDSA"}:
|
||||
raise ChatGPTAuthError("CHATGPT_ID_TOKEN_ALG_INVALID")
|
||||
|
||||
claims=jwt_module.decode(
|
||||
id_token,
|
||||
signing_key.key,
|
||||
algorithms=[alg],
|
||||
audience=issued_client_id,
|
||||
issuer=issuer,
|
||||
options={"require":["sub","exp","iat"]},
|
||||
leeway=5,
|
||||
)
|
||||
if claims.get("nonce") != tx.get("nonce"):
|
||||
raise ChatGPTAuthError("CHATGPT_NONCE_MISMATCH")
|
||||
subject=str(claims.get("sub") or "")
|
||||
if not subject:
|
||||
raise ChatGPTAuthError("CHATGPT_SUBJECT_MISSING")
|
||||
|
||||
granted=str(tokens.get("scope") or callback_query.get("scope") or "")
|
||||
scopes=sorted({x for x in granted.split() if x})
|
||||
email=str(claims.get("email") or "")
|
||||
|
||||
# No OAuth access/refresh/ID token is persisted or returned.
|
||||
return {
|
||||
"ok":True,
|
||||
"issuer":issuer,
|
||||
"client_id":issued_client_id,
|
||||
"account_binding_id":account_binding_id(issuer,subject),
|
||||
"subject_sha256":hashlib.sha256(subject.encode("utf-8")).hexdigest(),
|
||||
"email_sha256":hashlib.sha256(email.encode("utf-8")).hexdigest() if email else None,
|
||||
"scopes":scopes,
|
||||
"identity_only": "chatgpt.tokens.use.direct" not in scopes,
|
||||
}
|
||||
|
||||
|
||||
def authenticate_interactive(
|
||||
host_id: str,
|
||||
*,
|
||||
browser_opener=None,
|
||||
timeout_seconds: int=300,
|
||||
client_id: str=DYNAMIC_CLIENT_ID,
|
||||
) -> dict[str,Any]:
|
||||
server,cb,redirect_uri=start_loopback_listener()
|
||||
try:
|
||||
provider=discover_openai()
|
||||
url,tx=build_authorization(
|
||||
host_id,
|
||||
redirect_uri,
|
||||
authorization_endpoint=provider["authorization_endpoint"],
|
||||
client_id=client_id,
|
||||
)
|
||||
if browser_opener is None:
|
||||
import webbrowser
|
||||
opened=webbrowser.open(url,new=1,autoraise=True)
|
||||
else:
|
||||
opened=bool(browser_opener(url))
|
||||
if not opened:
|
||||
raise ChatGPTAuthError("CHATGPT_BROWSER_OPEN_FAILED")
|
||||
if not cb.event.wait(timeout_seconds):
|
||||
raise ChatGPTAuthError("CHATGPT_AUTH_TIMEOUT")
|
||||
return exchange_and_validate(cb.query,tx,provider_metadata=provider)
|
||||
finally:
|
||||
server.shutdown()
|
||||
server.server_close()
|
||||
+224
@@ -0,0 +1,224 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
import re
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
from .product_core import ProductError
|
||||
from .seed_capabilities import execute_capability, list_capabilities
|
||||
from .seed_core import SeedContext
|
||||
|
||||
FRAME_ID_RE = re.compile(r"^[A-Za-z0-9._-]{1,80}$")
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ClassicRequest:
|
||||
frame_id: str
|
||||
target: str
|
||||
mode: str
|
||||
lang: str
|
||||
body: dict[str, Any]
|
||||
body_sha256: str
|
||||
|
||||
|
||||
def parse_frame(text: str) -> ClassicRequest:
|
||||
raw = str(text).replace("\r\n", "\n").replace("\r", "\n")
|
||||
lines = raw.splitlines()
|
||||
|
||||
try:
|
||||
start = lines.index("FRAMEEXEC/1")
|
||||
except ValueError:
|
||||
raise ProductError("CLASSIC_FRAME_START_MISSING")
|
||||
|
||||
try:
|
||||
end = lines.index("FRAMEEND/1", start + 1)
|
||||
except ValueError:
|
||||
raise ProductError("CLASSIC_FRAME_END_MISSING")
|
||||
|
||||
frame_lines = lines[start + 1:end]
|
||||
headers: dict[str, str] = {}
|
||||
body_start = None
|
||||
|
||||
for i, line in enumerate(frame_lines):
|
||||
if not line.strip():
|
||||
body_start = i + 1
|
||||
break
|
||||
if "=" not in line:
|
||||
raise ProductError("CLASSIC_HEADER_INVALID")
|
||||
key, value = line.split("=", 1)
|
||||
headers[key.strip().upper()] = value.strip()
|
||||
|
||||
if body_start is None:
|
||||
raise ProductError("CLASSIC_BODY_SEPARATOR_MISSING")
|
||||
|
||||
frame_id = headers.get("ID", "")
|
||||
target = headers.get("TARGET", "")
|
||||
mode = headers.get("MODE", "")
|
||||
lang = headers.get("LANG", "")
|
||||
|
||||
if not FRAME_ID_RE.fullmatch(frame_id):
|
||||
raise ProductError("CLASSIC_ID_INVALID")
|
||||
if not target:
|
||||
raise ProductError("CLASSIC_TARGET_REQUIRED")
|
||||
if mode != "EXEC":
|
||||
raise ProductError("CLASSIC_MODE_UNSUPPORTED")
|
||||
if lang != "framebridge-cap-v1":
|
||||
raise ProductError("CLASSIC_LANG_UNSUPPORTED")
|
||||
|
||||
body_text = "\n".join(frame_lines[body_start:]).strip()
|
||||
if not body_text:
|
||||
raise ProductError("CLASSIC_BODY_MISSING")
|
||||
if len(body_text.encode("utf-8")) > 32768:
|
||||
raise ProductError("CLASSIC_BODY_TOO_LARGE")
|
||||
|
||||
try:
|
||||
body = json.loads(body_text)
|
||||
except Exception:
|
||||
raise ProductError("CLASSIC_BODY_JSON_INVALID")
|
||||
|
||||
if not isinstance(body, dict):
|
||||
raise ProductError("CLASSIC_BODY_NOT_OBJECT")
|
||||
|
||||
canonical = json.dumps(
|
||||
body,
|
||||
sort_keys=True,
|
||||
separators=(",", ":"),
|
||||
ensure_ascii=True,
|
||||
).encode("utf-8")
|
||||
|
||||
return ClassicRequest(
|
||||
frame_id=frame_id,
|
||||
target=target,
|
||||
mode=mode,
|
||||
lang=lang,
|
||||
body=body,
|
||||
body_sha256=hashlib.sha256(canonical).hexdigest(),
|
||||
)
|
||||
|
||||
|
||||
def _target_matches(seed: SeedContext, target: str) -> bool:
|
||||
identity = seed.load_identity()
|
||||
if identity is None:
|
||||
raise ProductError("SEED_NOT_INITIALIZED")
|
||||
|
||||
accepted = {
|
||||
str(identity.get("hostname") or ""),
|
||||
str(identity.get("node_id") or ""),
|
||||
str(identity.get("node_id_short") or ""),
|
||||
}
|
||||
return target in accepted
|
||||
|
||||
|
||||
def execute_frame(seed: SeedContext, text: str) -> dict[str, Any]:
|
||||
request = parse_frame(text)
|
||||
|
||||
if not _target_matches(seed, request.target):
|
||||
raise ProductError("CLASSIC_TARGET_MISMATCH")
|
||||
|
||||
capability = str(request.body.get("capability") or "").strip()
|
||||
args = request.body.get("args") or {}
|
||||
|
||||
if "script" in request.body or "shell" in request.body or "command" in request.body:
|
||||
raise ProductError("CLASSIC_RAW_SHELL_FORBIDDEN")
|
||||
|
||||
started = time.monotonic()
|
||||
try:
|
||||
result = execute_capability(capability, args)
|
||||
status = "PASS"
|
||||
exit_code = 0
|
||||
except Exception as exc:
|
||||
result = {
|
||||
"ok": False,
|
||||
"error": str(exc),
|
||||
"error_type": type(exc).__name__,
|
||||
}
|
||||
status = "ERROR"
|
||||
exit_code = 1
|
||||
|
||||
duration_ms = int((time.monotonic() - started) * 1000)
|
||||
|
||||
seed.audit.append(
|
||||
"classic." + capability,
|
||||
status,
|
||||
request={
|
||||
"id": request.frame_id,
|
||||
"target": request.target,
|
||||
"lang": request.lang,
|
||||
"body_sha256": request.body_sha256,
|
||||
"capability": capability,
|
||||
"args": args,
|
||||
},
|
||||
result={
|
||||
"exit": exit_code,
|
||||
"duration_ms": duration_ms,
|
||||
},
|
||||
principal="chat-manual",
|
||||
)
|
||||
|
||||
return {
|
||||
"id": request.frame_id,
|
||||
"exit": exit_code,
|
||||
"channel": "SEED-CLASSIC-TTY",
|
||||
"capability": capability,
|
||||
"body_sha256": request.body_sha256,
|
||||
"duration_ms": duration_ms,
|
||||
"result": result,
|
||||
}
|
||||
|
||||
|
||||
def render_result(value: dict[str, Any]) -> str:
|
||||
payload = json.dumps(
|
||||
value.get("result"),
|
||||
indent=2,
|
||||
sort_keys=True,
|
||||
ensure_ascii=True,
|
||||
)
|
||||
return "\n".join([
|
||||
"FRAMERESULT/1",
|
||||
"ID=" + str(value.get("id")),
|
||||
"EXIT=" + str(value.get("exit")),
|
||||
"CHANNEL=" + str(value.get("channel")),
|
||||
"CAPABILITY=" + str(value.get("capability")),
|
||||
"BODY_SHA256=" + str(value.get("body_sha256")),
|
||||
"DURATION_MS=" + str(value.get("duration_ms")),
|
||||
"",
|
||||
payload,
|
||||
"FRAMEEND/1",
|
||||
])
|
||||
|
||||
|
||||
def render_snapshot(seed: SeedContext) -> str:
|
||||
identity = seed.load_identity()
|
||||
config = seed.load_config()
|
||||
if identity is None or config is None:
|
||||
raise ProductError("SEED_NOT_INITIALIZED")
|
||||
|
||||
snapshot = {
|
||||
"protocol": "framebridge-cap-v1",
|
||||
"hostname": identity.get("hostname"),
|
||||
"node_id": identity.get("node_id"),
|
||||
"node_id_short": identity.get("node_id_short"),
|
||||
"provider": config.get("provider"),
|
||||
"connection_mode": config.get("connection_mode", "MANUAL"),
|
||||
"capability_profile": config.get("capability_profile"),
|
||||
"capabilities": list_capabilities(),
|
||||
"instruction": (
|
||||
"Return only one FRAMEEXEC/1 frame using LANG=framebridge-cap-v1. "
|
||||
"Use only a listed READONLY capability. Never include raw shell."
|
||||
),
|
||||
}
|
||||
|
||||
payload = json.dumps(snapshot, indent=2, sort_keys=True, ensure_ascii=True)
|
||||
|
||||
return "\n".join([
|
||||
"FRAMEBRIDGE-SNAPSHOT/1",
|
||||
"TARGET=" + str(identity.get("node_id_short")),
|
||||
"PROVIDER=" + str(config.get("provider")),
|
||||
"CONNECTION=MANUAL",
|
||||
"",
|
||||
payload,
|
||||
"FRAMEEND/1",
|
||||
])
|
||||
+133
@@ -0,0 +1,133 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from .memory_store import MemoryStore
|
||||
from .product_core import ProductError
|
||||
|
||||
|
||||
BANNER = """
|
||||
==============================================================
|
||||
FrameBridge Console
|
||||
==============================================================
|
||||
""".strip("\n")
|
||||
|
||||
|
||||
HELP = """
|
||||
Commands:
|
||||
help Show this help
|
||||
status Show node/Memory summary
|
||||
manifest Show the AI bootstrap manifest
|
||||
memory <domain> Query normalized Memory
|
||||
phone Show official mobile-app guidance
|
||||
provider Show configured AI provider
|
||||
objectives Show future objective-prompt roadmap
|
||||
quit / exit Leave the console
|
||||
|
||||
Anything else is treated as a chat message.
|
||||
""".strip("\n")
|
||||
|
||||
|
||||
PHONE_HINT = "📱 You can also continue from the official ChatGPT app on your phone."
|
||||
|
||||
|
||||
def _setup(ctx) -> dict[str, Any]:
|
||||
cfg = ctx.load_config() or {}
|
||||
return cfg.get("setup") or {}
|
||||
|
||||
|
||||
def session_header(ctx) -> str:
|
||||
setup = _setup(ctx)
|
||||
provider = setup.get("provider") or (ctx.load_config() or {}).get("provider") or "unknown"
|
||||
memory = MemoryStore(ctx).verify()
|
||||
identity = ctx.load_identity() or {}
|
||||
connected = os.environ.get("FRAMEBRIDGE_AGENT_CONNECTED") == "1"
|
||||
icon = "🤖" if connected else "○"
|
||||
lines = [
|
||||
BANNER,
|
||||
f"Node: {identity.get('hostname') or 'unknown'} ({identity.get('node_id_short') or '--------'})",
|
||||
f"AI provider: {provider}",
|
||||
f"Agent: {icon} {'connected' if connected else 'not connected'}",
|
||||
f"Memory: {memory.get('protocol')} R/W | events={memory.get('events')}",
|
||||
PHONE_HINT,
|
||||
"",
|
||||
"Type help for commands.",
|
||||
]
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
def _chat_response(ctx, text: str) -> str:
|
||||
setup = _setup(ctx)
|
||||
provider = setup.get("provider") or "chatgpt"
|
||||
# Chat transcripts are intentionally NOT written to technical Memory.
|
||||
# Only normalized technical facts/decisions/incidents/artifacts belong in FBM.
|
||||
if provider == "chatgpt":
|
||||
if os.environ.get("FRAMEBRIDGE_AGENT_CONNECTED") == "1":
|
||||
return (
|
||||
"A remote ChatGPT session is connected to this node. "
|
||||
"Console-to-agent conversational transport is not enabled in this candidate yet. "
|
||||
"Continue in the official ChatGPT app while this console remains available locally."
|
||||
)
|
||||
return (
|
||||
"ChatGPT companion mode is configured, but no remote agent session is currently attached.\n"
|
||||
+ PHONE_HINT + "\n"
|
||||
"Open ChatGPT and connect to this FrameBridge node. The local console will remain active."
|
||||
)
|
||||
return (
|
||||
f"Provider '{provider}' is configured, but direct conversational AgentAdapter transport "
|
||||
"is not implemented in this candidate."
|
||||
)
|
||||
|
||||
|
||||
def execute_console_command(ctx, line: str) -> tuple[bool, str]:
|
||||
raw = str(line).strip()
|
||||
if not raw:
|
||||
return True, ""
|
||||
if raw in {"quit", "exit"}:
|
||||
return False, "Goodbye."
|
||||
if raw == "help":
|
||||
return True, HELP
|
||||
if raw == "phone":
|
||||
return True, PHONE_HINT
|
||||
if raw == "provider":
|
||||
setup = _setup(ctx)
|
||||
return True, json.dumps({"provider": setup.get("provider"), "mode": "companion" if setup.get("provider") == "chatgpt" else "adapter"}, indent=2)
|
||||
if raw == "status":
|
||||
return True, json.dumps({"status": ctx.status(), "memory": MemoryStore(ctx).stats()}, indent=2, sort_keys=True)
|
||||
if raw == "manifest":
|
||||
return True, json.dumps(MemoryStore(ctx).bootstrap(), indent=2, sort_keys=True)
|
||||
if raw.startswith("memory "):
|
||||
domain = raw.split(None, 1)[1].strip()
|
||||
return True, json.dumps(MemoryStore(ctx).query(domain, 20), indent=2, sort_keys=True)
|
||||
if raw == "objectives":
|
||||
return True, (
|
||||
"Objective prompts are planned for a later update. Examples: install a desktop, "
|
||||
"create a server, configure a workstation, prepare a NAS. They are NOT active in this candidate."
|
||||
)
|
||||
return True, _chat_response(ctx, raw)
|
||||
|
||||
|
||||
def run_console_chat(ctx) -> int:
|
||||
try:
|
||||
MemoryStore(ctx).require_ready()
|
||||
except Exception as exc:
|
||||
print("FrameBridge Console blocked: mandatory Memory is not healthy.")
|
||||
print(str(exc))
|
||||
return 2
|
||||
print(session_header(ctx))
|
||||
while True:
|
||||
try:
|
||||
line = input("\nYou > ")
|
||||
except (EOFError, KeyboardInterrupt):
|
||||
print("\nGoodbye.")
|
||||
return 0
|
||||
keep, out = execute_console_command(ctx, line)
|
||||
if out:
|
||||
print("\n" + ("AI > " if line.strip() not in {"help","status","manifest","phone","provider","objectives"} and not line.strip().startswith("memory ") else "") + out)
|
||||
if not keep:
|
||||
return 0
|
||||
+88
@@ -0,0 +1,88 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
|
||||
from .classic_tty import execute_frame, render_result, render_snapshot
|
||||
from .product_core import ProductError
|
||||
from .seed_core import SeedContext
|
||||
|
||||
|
||||
BANNER = """FrameBridge Seed // Classic TTY
|
||||
Modo: MANUAL / CHAT_ONLY
|
||||
Seguridad: READONLY typed capabilities
|
||||
Comandos locales: /snapshot /status /audit /help /quit
|
||||
Pegue un FRAMEEXEC/1 completo para ejecutarlo.
|
||||
"""
|
||||
|
||||
|
||||
def _read_frame(first_line: str) -> str:
|
||||
lines = [first_line.rstrip("\n")]
|
||||
while True:
|
||||
raw = sys.stdin.readline()
|
||||
if raw == "":
|
||||
raise ProductError("CLASSIC_FRAME_EOF_BEFORE_END")
|
||||
line = raw.rstrip("\n")
|
||||
lines.append(line)
|
||||
if line == "FRAMEEND/1":
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
def run_console(seed: SeedContext) -> int:
|
||||
if seed.load_identity() is None:
|
||||
print("ERROR: Seed no inicializado. Ejecute: framebridge init --provider chatgpt")
|
||||
return 2
|
||||
|
||||
print(BANNER)
|
||||
print("--- SNAPSHOT PARA SU IA ---")
|
||||
print(render_snapshot(seed))
|
||||
print("--- FIN SNAPSHOT ---")
|
||||
|
||||
while True:
|
||||
if sys.stdin.isatty():
|
||||
print("fb> ", end="", flush=True)
|
||||
|
||||
raw = sys.stdin.readline()
|
||||
if raw == "":
|
||||
return 0
|
||||
|
||||
line = raw.rstrip("\n")
|
||||
stripped = line.strip()
|
||||
|
||||
if not stripped:
|
||||
continue
|
||||
|
||||
if stripped == "/quit":
|
||||
return 0
|
||||
if stripped == "/help":
|
||||
print(BANNER)
|
||||
continue
|
||||
if stripped == "/snapshot":
|
||||
print(render_snapshot(seed))
|
||||
continue
|
||||
if stripped == "/status":
|
||||
import json
|
||||
print(json.dumps(seed.status(), indent=2, sort_keys=True, ensure_ascii=True))
|
||||
continue
|
||||
if stripped == "/audit":
|
||||
import json
|
||||
print(json.dumps(seed.audit.verify(), indent=2, sort_keys=True, ensure_ascii=True))
|
||||
continue
|
||||
|
||||
if stripped != "FRAMEEXEC/1":
|
||||
print("ERROR: comando desconocido. Use /help o pegue FRAMEEXEC/1.")
|
||||
continue
|
||||
|
||||
try:
|
||||
frame = _read_frame(line)
|
||||
result = execute_frame(seed, frame)
|
||||
print(render_result(result))
|
||||
except ProductError as exc:
|
||||
print("FRAMERESULT/1")
|
||||
print("ID=unknown")
|
||||
print("EXIT=2")
|
||||
print("CHANNEL=SEED-CLASSIC-TTY")
|
||||
print("")
|
||||
print('{"ok":false,"error":' + __import__("json").dumps(str(exc)) + "}")
|
||||
print("FRAMEEND/1")
|
||||
|
||||
return 0
|
||||
+121
@@ -0,0 +1,121 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
import socket
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from .memory_store import MemoryStore
|
||||
from .product_core import ProductError
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class SetupChoice:
|
||||
privilege_mode: str
|
||||
provider: str
|
||||
ui_language: str
|
||||
|
||||
|
||||
VALID_PRIVILEGE_MODES = {"totp", "sudo_tty"}
|
||||
VALID_PROVIDERS = {"chatgpt", "ollama", "lmstudio", "other"}
|
||||
VALID_LANGUAGES = {"en", "es", "zh", "fr", "de", "pt", "ja", "ko"}
|
||||
|
||||
|
||||
def detect_system() -> dict[str, Any]:
|
||||
from .seed_capabilities import hardware_summary, system_status, storage_inspect, network_inspect
|
||||
system = system_status({})
|
||||
hardware = hardware_summary({})
|
||||
storage = storage_inspect({})
|
||||
network = network_inspect({})
|
||||
return {
|
||||
"hostname": socket.gethostname(),
|
||||
"system": system,
|
||||
"hardware": hardware,
|
||||
"storage": storage,
|
||||
"network": network,
|
||||
}
|
||||
|
||||
|
||||
def validate_choice(choice: SetupChoice) -> SetupChoice:
|
||||
pm = str(choice.privilege_mode).strip().lower()
|
||||
provider = str(choice.provider).strip().lower()
|
||||
lang = str(choice.ui_language).strip().lower()
|
||||
if pm not in VALID_PRIVILEGE_MODES:
|
||||
raise ProductError("SETUP_PRIVILEGE_MODE_INVALID")
|
||||
if provider not in VALID_PROVIDERS:
|
||||
raise ProductError("SETUP_PROVIDER_INVALID")
|
||||
if lang not in VALID_LANGUAGES:
|
||||
raise ProductError("SETUP_LANGUAGE_INVALID")
|
||||
return SetupChoice(pm, provider, lang)
|
||||
|
||||
|
||||
def save_setup_config(ctx, choice: SetupChoice) -> dict[str, Any]:
|
||||
choice = validate_choice(choice)
|
||||
config = ctx.load_config() or {}
|
||||
config["provider"] = choice.provider
|
||||
config["setup"] = {
|
||||
"completed": True,
|
||||
"privilege_mode": choice.privilege_mode,
|
||||
"provider": choice.provider,
|
||||
"ui_language": choice.ui_language,
|
||||
"chat_mode": "companion" if choice.provider == "chatgpt" else "agent_adapter",
|
||||
"pairing_state": (
|
||||
(config.get("pairing") or {}).get("state", "not_paired")
|
||||
if choice.provider == "chatgpt"
|
||||
else "adapter_pending"
|
||||
),
|
||||
}
|
||||
path = Path(ctx.config_path)
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
tmp = path.with_name(path.name + ".setup.tmp")
|
||||
tmp.write_text(json.dumps(config, indent=2, sort_keys=True) + "\n", encoding="utf-8")
|
||||
os.chmod(tmp, 0o600)
|
||||
os.replace(tmp, path)
|
||||
return config
|
||||
|
||||
|
||||
def setup_summary(ctx) -> dict[str, Any]:
|
||||
config = ctx.load_config() or {}
|
||||
setup = config.get("setup") or {}
|
||||
memory = MemoryStore(ctx).verify() if (config.get("memory") or {}).get("required") else {"ok": False, "required": False}
|
||||
return {
|
||||
"ok": True,
|
||||
"setup": setup,
|
||||
"memory": memory,
|
||||
"provider": config.get("provider"),
|
||||
"seed_version": config.get("seed_version"),
|
||||
}
|
||||
|
||||
|
||||
def render_privilege_explanation(mode: str) -> list[str]:
|
||||
if mode == "totp":
|
||||
return [
|
||||
"Authenticator / QR TOTP selected.",
|
||||
"FrameBridge will show enrollment QR locally.",
|
||||
"The TOTP secret and OTP codes never leave this machine.",
|
||||
"Privileged CHANGE operations still use typed plans and one-shot approval.",
|
||||
]
|
||||
return [
|
||||
"Local sudo password selected.",
|
||||
"FrameBridge never stores or forwards your sudo password.",
|
||||
"sudo prompts are handled directly by the local terminal.",
|
||||
"The AI does not receive general sudo or raw root shell access.",
|
||||
]
|
||||
|
||||
|
||||
def provider_instructions(provider: str) -> list[str]:
|
||||
if provider == "chatgpt":
|
||||
return [
|
||||
"ChatGPT companion mode selected.",
|
||||
"Use the official ChatGPT app or web interface for your OpenAI account.",
|
||||
"FrameBridge does not ask for or store your ChatGPT password.",
|
||||
"After pairing, this local console remains the primary on-machine FrameBridge interface.",
|
||||
]
|
||||
if provider == "ollama":
|
||||
return ["Ollama selected. Local AgentAdapter integration will be configured when available."]
|
||||
if provider == "lmstudio":
|
||||
return ["LM Studio selected. Local AgentAdapter integration will be configured when available."]
|
||||
return ["Custom provider selected. Configure an AgentAdapter after setup."]
|
||||
+548
@@ -0,0 +1,548 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import struct
|
||||
import tempfile
|
||||
import time
|
||||
import uuid
|
||||
import zlib
|
||||
from pathlib import Path
|
||||
from typing import Any, Iterator
|
||||
|
||||
from .product_core import ProductError, redact
|
||||
from .seed_capabilities import hardware_summary, list_capabilities, network_inspect, storage_inspect, system_status
|
||||
|
||||
FBAM_PROTOCOL = "FBAM/1"
|
||||
FBM_PROTOCOL = "FBM/1"
|
||||
FBM_SCHEMA = 1
|
||||
CODEC_ZLIB = 1
|
||||
HEADER = struct.Struct(">4sBBII32s")
|
||||
MAGIC_MANIFEST = b"FBAM"
|
||||
MAGIC_STATE = b"FBM1"
|
||||
MAGIC_RECORD = b"FBMR"
|
||||
MAGIC_INDEX = b"FBI1"
|
||||
MAX_EVENT_RAW = 65536
|
||||
MAX_QUERY_EVENTS = 100
|
||||
|
||||
READ_CAPABILITIES = tuple(x["name"] for x in list_capabilities())
|
||||
CHANGE_CAPABILITIES = (
|
||||
"workspace.write_text",
|
||||
"package.install",
|
||||
"package.remove",
|
||||
"service.manage",
|
||||
"system.hostname",
|
||||
"system.timezone",
|
||||
"package.refresh",
|
||||
"package.upgrade",
|
||||
"system.file.write_text",
|
||||
"system.directory.ensure",
|
||||
"system.file.mode",
|
||||
"system.daemon_reload",
|
||||
)
|
||||
MEMORY_CAPABILITIES = ("memory.bootstrap", "memory.query", "memory.append", "memory.checkpoint", "memory.stats")
|
||||
|
||||
RECORD_TYPE_IDS = {
|
||||
"checkpoint": 1,
|
||||
"change": 2,
|
||||
"incident": 3,
|
||||
"decision": 4,
|
||||
"artifact": 5,
|
||||
"recovery": 6,
|
||||
"certification": 7,
|
||||
"note": 8,
|
||||
}
|
||||
RECORD_TYPE_NAMES = {v: k for k, v in RECORD_TYPE_IDS.items()}
|
||||
SOURCE_IDS = {"system": 1, "agent": 2, "operator": 3, "import": 4}
|
||||
SOURCE_NAMES = {v: k for k, v in SOURCE_IDS.items()}
|
||||
FIELD_REGISTRY = {
|
||||
1: "schema", 2: "record_type", 3: "timestamp_ns", 4: "node_id",
|
||||
5: "identity", 6: "platform", 7: "hardware", 8: "storage", 9: "network",
|
||||
10: "software", 11: "capabilities", 12: "policy", 13: "connections",
|
||||
14: "plans", 15: "audit", 16: "changes", 17: "incidents", 18: "decisions",
|
||||
19: "artifacts", 20: "recovery", 21: "relationships", 22: "pending",
|
||||
23: "freshness", 24: "provenance",
|
||||
}
|
||||
DOMAIN_TO_FIELD = {name: fid for fid, name in FIELD_REGISTRY.items() if fid >= 5}
|
||||
FORBIDDEN_KEY_RE = re.compile(r"(?i)(password|passwd|secret|token|private[_-]?key|otp|credential|client[_-]?secret|api[_-]?key)")
|
||||
|
||||
|
||||
def _uint_head(major: int, value: int) -> bytes:
|
||||
if value < 0:
|
||||
raise ValueError("CBOR_UINT_NEGATIVE")
|
||||
if value < 24:
|
||||
return bytes([(major << 5) | value])
|
||||
if value <= 0xFF:
|
||||
return bytes([(major << 5) | 24, value])
|
||||
if value <= 0xFFFF:
|
||||
return bytes([(major << 5) | 25]) + struct.pack(">H", value)
|
||||
if value <= 0xFFFFFFFF:
|
||||
return bytes([(major << 5) | 26]) + struct.pack(">I", value)
|
||||
if value <= 0xFFFFFFFFFFFFFFFF:
|
||||
return bytes([(major << 5) | 27]) + struct.pack(">Q", value)
|
||||
raise ValueError("CBOR_INTEGER_TOO_LARGE")
|
||||
|
||||
|
||||
def cbor_encode(value: Any) -> bytes:
|
||||
if value is None:
|
||||
return b"\xf6"
|
||||
if value is False:
|
||||
return b"\xf4"
|
||||
if value is True:
|
||||
return b"\xf5"
|
||||
if isinstance(value, int) and not isinstance(value, bool):
|
||||
return _uint_head(0, value) if value >= 0 else _uint_head(1, -1 - value)
|
||||
if isinstance(value, float):
|
||||
return b"\xfb" + struct.pack(">d", value)
|
||||
if isinstance(value, bytes):
|
||||
return _uint_head(2, len(value)) + value
|
||||
if isinstance(value, str):
|
||||
raw = value.encode("utf-8")
|
||||
return _uint_head(3, len(raw)) + raw
|
||||
if isinstance(value, (list, tuple)):
|
||||
parts = [cbor_encode(x) for x in value]
|
||||
return _uint_head(4, len(parts)) + b"".join(parts)
|
||||
if isinstance(value, dict):
|
||||
items = []
|
||||
for key, item in value.items():
|
||||
ek = cbor_encode(key)
|
||||
items.append((len(ek), ek, cbor_encode(item)))
|
||||
items.sort(key=lambda x: (x[0], x[1]))
|
||||
return _uint_head(5, len(items)) + b"".join(ek + ev for _, ek, ev in items)
|
||||
raise TypeError("CBOR_TYPE_UNSUPPORTED:" + type(value).__name__)
|
||||
|
||||
|
||||
def _read_ai(data: bytes, offset: int, ai: int) -> tuple[int, int]:
|
||||
if ai < 24:
|
||||
return ai, offset
|
||||
if ai == 24:
|
||||
if offset + 1 > len(data): raise ValueError("CBOR_TRUNCATED")
|
||||
return data[offset], offset + 1
|
||||
if ai == 25:
|
||||
if offset + 2 > len(data): raise ValueError("CBOR_TRUNCATED")
|
||||
return struct.unpack_from(">H", data, offset)[0], offset + 2
|
||||
if ai == 26:
|
||||
if offset + 4 > len(data): raise ValueError("CBOR_TRUNCATED")
|
||||
return struct.unpack_from(">I", data, offset)[0], offset + 4
|
||||
if ai == 27:
|
||||
if offset + 8 > len(data): raise ValueError("CBOR_TRUNCATED")
|
||||
return struct.unpack_from(">Q", data, offset)[0], offset + 8
|
||||
raise ValueError("CBOR_INDEFINITE_NOT_SUPPORTED")
|
||||
|
||||
|
||||
def _decode_one(data: bytes, offset: int = 0) -> tuple[Any, int]:
|
||||
if offset >= len(data):
|
||||
raise ValueError("CBOR_TRUNCATED")
|
||||
first = data[offset]; offset += 1
|
||||
major, ai = first >> 5, first & 31
|
||||
if major in (0, 1):
|
||||
n, offset = _read_ai(data, offset, ai)
|
||||
return (n if major == 0 else -1 - n), offset
|
||||
if major in (2, 3):
|
||||
n, offset = _read_ai(data, offset, ai)
|
||||
end = offset + n
|
||||
if end > len(data): raise ValueError("CBOR_TRUNCATED")
|
||||
raw = data[offset:end]
|
||||
return (raw if major == 2 else raw.decode("utf-8")), end
|
||||
if major == 4:
|
||||
n, offset = _read_ai(data, offset, ai)
|
||||
out = []
|
||||
for _ in range(n):
|
||||
item, offset = _decode_one(data, offset); out.append(item)
|
||||
return out, offset
|
||||
if major == 5:
|
||||
n, offset = _read_ai(data, offset, ai)
|
||||
out = {}
|
||||
for _ in range(n):
|
||||
key, offset = _decode_one(data, offset)
|
||||
val, offset = _decode_one(data, offset)
|
||||
out[key] = val
|
||||
return out, offset
|
||||
if major == 7:
|
||||
if ai == 20: return False, offset
|
||||
if ai == 21: return True, offset
|
||||
if ai == 22: return None, offset
|
||||
if ai == 27:
|
||||
if offset + 8 > len(data): raise ValueError("CBOR_TRUNCATED")
|
||||
return struct.unpack_from(">d", data, offset)[0], offset + 8
|
||||
raise ValueError("CBOR_VALUE_UNSUPPORTED")
|
||||
|
||||
|
||||
def cbor_decode(data: bytes) -> Any:
|
||||
value, offset = _decode_one(bytes(data), 0)
|
||||
if offset != len(data):
|
||||
raise ValueError("CBOR_TRAILING_DATA")
|
||||
return value
|
||||
|
||||
|
||||
def _pack_container(magic: bytes, value: Any) -> bytes:
|
||||
raw = cbor_encode(value)
|
||||
comp = zlib.compress(raw, 9)
|
||||
digest = hashlib.sha256(raw).digest()
|
||||
return HEADER.pack(magic, FBM_SCHEMA, CODEC_ZLIB, len(raw), len(comp), digest) + comp
|
||||
|
||||
|
||||
def _unpack_container(data: bytes, expected_magic: bytes) -> tuple[Any, dict[str, Any]]:
|
||||
if len(data) < HEADER.size:
|
||||
raise ProductError("MEMORY_CONTAINER_TRUNCATED")
|
||||
magic, schema, codec, raw_len, comp_len, digest = HEADER.unpack_from(data, 0)
|
||||
if magic != expected_magic: raise ProductError("MEMORY_MAGIC_INVALID")
|
||||
if schema != FBM_SCHEMA: raise ProductError("MEMORY_SCHEMA_UNSUPPORTED")
|
||||
if codec != CODEC_ZLIB: raise ProductError("MEMORY_CODEC_UNSUPPORTED")
|
||||
if len(data) != HEADER.size + comp_len: raise ProductError("MEMORY_CONTAINER_SIZE_INVALID")
|
||||
try: raw = zlib.decompress(data[HEADER.size:])
|
||||
except Exception as exc: raise ProductError("MEMORY_DECOMPRESS_FAILED") from exc
|
||||
if len(raw) != raw_len: raise ProductError("MEMORY_RAW_SIZE_INVALID")
|
||||
if hashlib.sha256(raw).digest() != digest: raise ProductError("MEMORY_DIGEST_MISMATCH")
|
||||
try: value = cbor_decode(raw)
|
||||
except Exception as exc: raise ProductError("MEMORY_CBOR_INVALID") from exc
|
||||
return value, {"raw_bytes": raw_len, "compressed_bytes": comp_len, "raw_sha256": digest.hex()}
|
||||
|
||||
|
||||
def _atomic_write(path: Path, data: bytes) -> None:
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
os.chmod(path.parent, 0o700)
|
||||
fd, tmp = tempfile.mkstemp(prefix=".fbm-", dir=str(path.parent))
|
||||
try:
|
||||
with os.fdopen(fd, "wb") as handle:
|
||||
handle.write(data); handle.flush(); os.fsync(handle.fileno())
|
||||
os.chmod(tmp, 0o600)
|
||||
os.replace(tmp, path)
|
||||
finally:
|
||||
if os.path.exists(tmp): os.unlink(tmp)
|
||||
|
||||
|
||||
def _atomic_write_json(path: Path, value: dict[str, Any]) -> None:
|
||||
data = (json.dumps(value, sort_keys=True, indent=2, ensure_ascii=True) + "\n").encode("utf-8")
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
fd, tmp = tempfile.mkstemp(prefix=".fbm-config-", dir=str(path.parent))
|
||||
try:
|
||||
with os.fdopen(fd, "wb") as handle:
|
||||
handle.write(data); handle.flush(); os.fsync(handle.fileno())
|
||||
os.chmod(tmp, 0o600)
|
||||
os.replace(tmp, path)
|
||||
finally:
|
||||
if os.path.exists(tmp): os.unlink(tmp)
|
||||
|
||||
|
||||
def _uuid_bytes(value: str) -> bytes:
|
||||
try: return uuid.UUID(str(value)).bytes
|
||||
except Exception: return hashlib.sha256(str(value).encode("utf-8")).digest()[:16]
|
||||
|
||||
|
||||
def _uuid_text(value: bytes) -> str:
|
||||
try: return str(uuid.UUID(bytes=bytes(value)))
|
||||
except Exception: return bytes(value).hex()
|
||||
|
||||
|
||||
def _human(value: Any) -> Any:
|
||||
if isinstance(value, bytes):
|
||||
return value.hex()
|
||||
if isinstance(value, dict):
|
||||
return {str(k): _human(v) for k, v in value.items()}
|
||||
if isinstance(value, list):
|
||||
return [_human(v) for v in value]
|
||||
return value
|
||||
|
||||
|
||||
def _safe_payload(value: Any, depth: int = 0) -> Any:
|
||||
if depth > 8: raise ProductError("MEMORY_PAYLOAD_DEPTH_LIMIT")
|
||||
if isinstance(value, dict):
|
||||
out = {}
|
||||
for key, item in value.items():
|
||||
key = str(key)
|
||||
if len(key) > 80: raise ProductError("MEMORY_KEY_TOO_LONG")
|
||||
out[key] = "[REDACTED]" if FORBIDDEN_KEY_RE.search(key) else _safe_payload(item, depth + 1)
|
||||
return out
|
||||
if isinstance(value, list):
|
||||
if len(value) > 256: raise ProductError("MEMORY_LIST_TOO_LONG")
|
||||
return [_safe_payload(x, depth + 1) for x in value]
|
||||
if isinstance(value, str):
|
||||
return str(redact(value))[:4096]
|
||||
if isinstance(value, (int, float, bool)) or value is None:
|
||||
return value
|
||||
return str(value)[:4096]
|
||||
|
||||
|
||||
class MemoryStore:
|
||||
def __init__(self, seed):
|
||||
self.seed = seed
|
||||
self.root = Path(seed.state_dir) / "memory"
|
||||
self.state_path = self.root / "node.fbm"
|
||||
self.events_path = self.root / "events.fbm"
|
||||
self.index_path = self.root / "index.fbi"
|
||||
self.manifest_path = self.root / "manifest.fbam"
|
||||
|
||||
def _identity_config(self) -> tuple[dict[str, Any], dict[str, Any]]:
|
||||
identity = self.seed.load_identity()
|
||||
config = self.seed.load_config()
|
||||
if not isinstance(identity, dict) or not isinstance(config, dict):
|
||||
raise ProductError("MEMORY_SEED_NOT_INITIALIZED")
|
||||
return identity, config
|
||||
|
||||
def _plans(self) -> list[dict[str, Any]]:
|
||||
out = []
|
||||
plans_dir = Path(self.seed.state_dir) / "plans"
|
||||
if not plans_dir.is_dir(): return out
|
||||
import json
|
||||
for path in sorted(plans_dir.glob("*.json")):
|
||||
try:
|
||||
value = json.loads(path.read_text(encoding="utf-8"))
|
||||
plan = value.get("plan") or {}
|
||||
out.append({1: path.stem, 2: value.get("state"), 3: value.get("plan_sha256"), 4: plan.get("title")})
|
||||
except Exception:
|
||||
out.append({1: path.stem, 2: "INVALID"})
|
||||
return out
|
||||
|
||||
@staticmethod
|
||||
def _normalize_storage(value: dict[str, Any]) -> dict[int, Any]:
|
||||
devices = []
|
||||
for item in (value.get("blockdevices") or [])[:128]:
|
||||
if not isinstance(item, dict): continue
|
||||
devices.append({
|
||||
1: item.get("name"), 2: item.get("type"), 3: item.get("size"),
|
||||
4: item.get("fstype"), 5: item.get("mountpoints"), 6: item.get("rota"), 7: item.get("model"),
|
||||
})
|
||||
return {1: value.get("ok") is True, 2: devices}
|
||||
|
||||
@staticmethod
|
||||
def _normalize_network(value: dict[str, Any]) -> dict[int, Any]:
|
||||
addresses = []
|
||||
for item in (value.get("addresses") or [])[:64]:
|
||||
if not isinstance(item, dict): continue
|
||||
infos = []
|
||||
for a in (item.get("addr_info") or [])[:32]:
|
||||
if not isinstance(a, dict): continue
|
||||
infos.append({1:a.get("family"),2:a.get("local"),3:a.get("prefixlen"),4:a.get("scope")})
|
||||
addresses.append({1:item.get("ifname"),2:item.get("operstate"),3:infos})
|
||||
routes = []
|
||||
for r in (value.get("routes") or [])[:128]:
|
||||
if not isinstance(r, dict): continue
|
||||
routes.append({1:r.get("dst"),2:r.get("gateway"),3:r.get("dev"),4:r.get("protocol"),5:r.get("metric")})
|
||||
return {1:value.get("ok") is True,2:addresses,3:routes}
|
||||
|
||||
def _collect_state(self) -> dict[int, Any]:
|
||||
identity, config = self._identity_config()
|
||||
doctor = self.seed.doctor()
|
||||
sys_state = system_status({})
|
||||
hw = hardware_summary({})
|
||||
storage = storage_inspect({})
|
||||
network = network_inspect({})
|
||||
audit = self.seed.audit.verify()
|
||||
plans = self._plans()
|
||||
pending = [p for p in plans if p.get(2) not in ("ROLLED_BACK",)]
|
||||
memory_cfg = config.get("memory") or {}
|
||||
return {
|
||||
1: FBM_SCHEMA,
|
||||
2: 1,
|
||||
3: time.time_ns(),
|
||||
4: _uuid_bytes(identity.get("node_id", "")),
|
||||
5: {
|
||||
1: identity.get("hostname"),
|
||||
2: identity.get("node_id_short"),
|
||||
3: config.get("provider"),
|
||||
4: identity.get("node_host_id"),
|
||||
5: identity.get("node_public_key_sha256"),
|
||||
},
|
||||
6: {1: sys_state.get("os_id"), 2: (doctor.get("checks") or {}).get("os_version"), 3: sys_state.get("os_name"), 4: sys_state.get("kernel"), 5: sys_state.get("architecture"), 6: (doctor.get("checks") or {}).get("python")},
|
||||
7: {1: hw.get("cpu_model"), 2: hw.get("logical_cpus"), 3: hw.get("memory_total_kib"), 4: hw.get("architecture")},
|
||||
8: self._normalize_storage(storage),
|
||||
9: self._normalize_network(network),
|
||||
10: {1: config.get("seed_version"), 2: (doctor.get("checks") or {}).get("python")},
|
||||
11: {1: list(READ_CAPABILITIES), 2: list(CHANGE_CAPABILITIES), 3: list(MEMORY_CAPABILITIES)},
|
||||
12: {1: config.get("capability_profile"), 2: True, 3: "plan-required"},
|
||||
13: {1: config.get("connection_name"), 2: config.get("connection_mode"), 3: (config.get("pairing") or {}).get("state"), 4: (config.get("pairing") or {}).get("transport")},
|
||||
14: plans,
|
||||
15: {1: audit.get("ok") is True, 2: audit.get("entries"), 3: audit.get("last_hash")},
|
||||
21: [],
|
||||
22: pending,
|
||||
23: {1: time.time_ns(), 2: "live-collection"},
|
||||
24: {1: "framebridge-node", 2: memory_cfg.get("authority", "deployment-defined"), 3: "node-operational"},
|
||||
}
|
||||
|
||||
def _manifest_object(self, state: dict[int, Any]) -> dict[int, Any]:
|
||||
identity, config = self._identity_config()
|
||||
return {
|
||||
1: 1,
|
||||
2: "FrameBridge",
|
||||
3: FBAM_PROTOCOL,
|
||||
4: _uuid_bytes(identity.get("node_id", "")),
|
||||
5: identity.get("node_id_short"),
|
||||
6: identity.get("hostname"),
|
||||
7: config.get("seed_version"),
|
||||
8: state.get(6),
|
||||
9: state.get(7),
|
||||
10: state.get(11),
|
||||
11: {1: config.get("capability_profile"), 2: True},
|
||||
12: {1: True, 2: "read_write", 3: FBM_PROTOCOL, 4: list(MEMORY_CAPABILITIES)},
|
||||
13: {1: (config.get("memory") or {}).get("authority", "deployment-defined"), 2: "node-operational"},
|
||||
14: time.time_ns(),
|
||||
}
|
||||
|
||||
def _read_container(self, path: Path, magic: bytes) -> tuple[Any, dict[str, Any]]:
|
||||
try: data = path.read_bytes()
|
||||
except FileNotFoundError as exc: raise ProductError("MEMORY_FILE_MISSING") from exc
|
||||
return _unpack_container(data, magic)
|
||||
|
||||
def _load_index(self) -> dict[int, Any]:
|
||||
if not self.index_path.is_file(): return {1:1,2:0,3:b"",4:[]}
|
||||
value, _ = self._read_container(self.index_path, MAGIC_INDEX)
|
||||
if not isinstance(value, dict): raise ProductError("MEMORY_INDEX_INVALID")
|
||||
return value
|
||||
|
||||
def initialize(self) -> dict[str, Any]:
|
||||
_, config = self._identity_config()
|
||||
desired = {
|
||||
"required": True, "protocol": FBM_PROTOCOL, "manifest_protocol": FBAM_PROTOCOL,
|
||||
"mode": "read_write", "authority": (config.get("memory") or {}).get("authority", "deployment-defined"),
|
||||
}
|
||||
if config.get("memory") != desired:
|
||||
config["memory"] = desired
|
||||
_atomic_write_json(Path(self.seed.config_path), config)
|
||||
self.root.mkdir(parents=True, exist_ok=True); os.chmod(self.root, 0o700)
|
||||
if not self.events_path.exists():
|
||||
self.events_path.touch(mode=0o600); os.chmod(self.events_path, 0o600)
|
||||
if not self.index_path.exists():
|
||||
_atomic_write(self.index_path, _pack_container(MAGIC_INDEX, {1:1,2:0,3:b"",4:[]}))
|
||||
checkpoint = self.checkpoint()
|
||||
if self._load_index().get(2, 0) == 0:
|
||||
self._append_event_raw("checkpoint", {"action":"memory.initialize","protocol":FBM_PROTOCOL}, "system")
|
||||
return {"ok":True,"protocol":FBM_PROTOCOL,"manifest_protocol":FBAM_PROTOCOL,"required":True,"mode":"read_write","checkpoint":checkpoint,"verify":self.verify()}
|
||||
|
||||
def checkpoint(self) -> dict[str, Any]:
|
||||
state = self._collect_state()
|
||||
manifest = self._manifest_object(state)
|
||||
state_blob = _pack_container(MAGIC_STATE, state)
|
||||
manifest_blob = _pack_container(MAGIC_MANIFEST, manifest)
|
||||
_atomic_write(self.state_path, state_blob)
|
||||
_atomic_write(self.manifest_path, manifest_blob)
|
||||
return {
|
||||
"ok": True,
|
||||
"state_sha256": hashlib.sha256(state_blob).hexdigest(),
|
||||
"manifest_sha256": hashlib.sha256(manifest_blob).hexdigest(),
|
||||
"state_bytes": len(state_blob),
|
||||
"manifest_bytes": len(manifest_blob),
|
||||
}
|
||||
|
||||
def ready(self) -> bool:
|
||||
try: return self.verify().get("ok") is True
|
||||
except Exception: return False
|
||||
|
||||
def require_ready(self) -> None:
|
||||
_, config = self._identity_config()
|
||||
required = bool((config.get("memory") or {}).get("required"))
|
||||
if required and not self.ready(): raise ProductError("MEMORY_REQUIRED_NOT_READY")
|
||||
|
||||
def _append_event_raw(self, record_type: str, data: dict[str, Any], source: str) -> dict[str, Any]:
|
||||
if record_type not in RECORD_TYPE_IDS: raise ProductError("MEMORY_RECORD_TYPE_INVALID")
|
||||
if source not in SOURCE_IDS: raise ProductError("MEMORY_SOURCE_INVALID")
|
||||
identity, _ = self._identity_config()
|
||||
safe = _safe_payload(data)
|
||||
if not isinstance(safe, dict): raise ProductError("MEMORY_PAYLOAD_NOT_OBJECT")
|
||||
index = self._load_index()
|
||||
prev = index.get(3) or b""
|
||||
record = {1:FBM_SCHEMA,2:RECORD_TYPE_IDS[record_type],3:time.time_ns(),4:_uuid_bytes(identity.get("node_id","")),5:safe,6:SOURCE_IDS[source],7:prev}
|
||||
raw = cbor_encode(record)
|
||||
if len(raw) > MAX_EVENT_RAW: raise ProductError("MEMORY_EVENT_TOO_LARGE")
|
||||
block = _pack_container(MAGIC_RECORD, record)
|
||||
self.events_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
offset = self.events_path.stat().st_size if self.events_path.exists() else 0
|
||||
with self.events_path.open("ab") as handle:
|
||||
handle.write(block); handle.flush(); os.fsync(handle.fileno())
|
||||
os.chmod(self.events_path, 0o600)
|
||||
record_hash = hashlib.sha256(raw).digest()
|
||||
offsets = list(index.get(4) or [])
|
||||
offsets.append(offset)
|
||||
offsets = offsets[-1024:]
|
||||
new_index = {1:1,2:int(index.get(2) or 0)+1,3:record_hash,4:offsets}
|
||||
_atomic_write(self.index_path, _pack_container(MAGIC_INDEX, new_index))
|
||||
return {"ok":True,"record_type":record_type,"event_number":new_index[2],"event_sha256":record_hash.hex(),"source":source}
|
||||
|
||||
def append(self, record_type: str, data: dict[str, Any], source: str = "agent") -> dict[str, Any]:
|
||||
self.require_ready()
|
||||
return self._append_event_raw(record_type, data, source)
|
||||
|
||||
def _iter_events(self) -> Iterator[tuple[dict[int, Any], bytes, int]]:
|
||||
if not self.events_path.exists(): return
|
||||
data = self.events_path.read_bytes(); off = 0; number = 0
|
||||
while off < len(data):
|
||||
if off + HEADER.size > len(data): raise ProductError("MEMORY_EVENT_STREAM_TRUNCATED")
|
||||
_, _, _, _, comp_len, _ = HEADER.unpack_from(data, off)
|
||||
end = off + HEADER.size + comp_len
|
||||
if end > len(data): raise ProductError("MEMORY_EVENT_STREAM_TRUNCATED")
|
||||
value, _ = _unpack_container(data[off:end], MAGIC_RECORD)
|
||||
if not isinstance(value, dict): raise ProductError("MEMORY_EVENT_INVALID")
|
||||
number += 1
|
||||
yield value, hashlib.sha256(cbor_encode(value)).digest(), off
|
||||
off = end
|
||||
|
||||
def verify(self) -> dict[str, Any]:
|
||||
state, sm = self._read_container(self.state_path, MAGIC_STATE)
|
||||
manifest, mm = self._read_container(self.manifest_path, MAGIC_MANIFEST)
|
||||
index = self._load_index()
|
||||
prev = b""; count = 0; offsets = []
|
||||
for record, record_hash, offset in self._iter_events():
|
||||
if record.get(7, b"") != prev: raise ProductError("MEMORY_EVENT_CHAIN_INVALID")
|
||||
prev = record_hash; count += 1; offsets.append(offset)
|
||||
if int(index.get(2) or 0) != count: raise ProductError("MEMORY_INDEX_COUNT_MISMATCH")
|
||||
if (index.get(3) or b"") != prev: raise ProductError("MEMORY_INDEX_HASH_MISMATCH")
|
||||
identity, _ = self._identity_config()
|
||||
if state.get(4) != _uuid_bytes(identity.get("node_id","")): raise ProductError("MEMORY_NODE_ID_MISMATCH")
|
||||
if manifest.get(4) != state.get(4): raise ProductError("MEMORY_MANIFEST_NODE_MISMATCH")
|
||||
return {
|
||||
"ok":True,"protocol":FBM_PROTOCOL,"manifest_protocol":FBAM_PROTOCOL,"events":count,
|
||||
"event_head_sha256":prev.hex() if prev else "GENESIS",
|
||||
"state_raw_bytes":sm["raw_bytes"],"state_compressed_bytes":sm["compressed_bytes"],
|
||||
"manifest_raw_bytes":mm["raw_bytes"],"manifest_compressed_bytes":mm["compressed_bytes"],
|
||||
}
|
||||
|
||||
def bootstrap(self) -> dict[str, Any]:
|
||||
self.require_ready()
|
||||
manifest, _ = self._read_container(self.manifest_path, MAGIC_MANIFEST)
|
||||
platform = manifest.get(8) or {}; hardware = manifest.get(9) or {}; caps = manifest.get(10) or {}; mem = manifest.get(12) or {}; authority = manifest.get(13) or {}
|
||||
return {
|
||||
"ok": True,
|
||||
"protocol": FBAM_PROTOCOL,
|
||||
"framebridge": {"product":"FrameBridge","seed_version":manifest.get(7)},
|
||||
"node": {"node_id":_uuid_text(manifest.get(4,b"")),"node_id_short":manifest.get(5),"hostname":manifest.get(6)},
|
||||
"os": {"id":platform.get(1),"version":platform.get(2),"name":platform.get(3),"kernel":platform.get(4),"architecture":platform.get(5)},
|
||||
"hardware": {"cpu_model":hardware.get(1),"logical_cpus":hardware.get(2),"memory_total_kib":hardware.get(3),"architecture":hardware.get(4)},
|
||||
"capabilities": {"read":caps.get(1) or [],"change":caps.get(2) or [],"memory":caps.get(3) or []},
|
||||
"policy": {"profile":(manifest.get(11) or {}).get(1),"change_requires_plan":True},
|
||||
"memory": {"required":mem.get(1) is True,"mode":mem.get(2),"protocol":mem.get(3),"capabilities":mem.get(4) or [],"rule":"READ_BEFORE_DEPENDENT_ACTION__WRITE_AFTER_MATERIAL_CHANGE"},
|
||||
"authority": {"mode":authority.get(1),"node_role":authority.get(2)},
|
||||
"generated_at_ns": manifest.get(14),
|
||||
"manifest_sha256": hashlib.sha256(self.manifest_path.read_bytes()).hexdigest(),
|
||||
}
|
||||
|
||||
def query(self, domain: str = "summary", limit: int = 20) -> dict[str, Any]:
|
||||
self.require_ready()
|
||||
domain = str(domain).strip().lower()
|
||||
limit = max(1, min(int(limit), MAX_QUERY_EVENTS))
|
||||
state, _ = self._read_container(self.state_path, MAGIC_STATE)
|
||||
if domain == "summary":
|
||||
verify = self.verify()
|
||||
return {"ok":True,"domain":"summary","node_id":_uuid_text(state.get(4,b"")),"hostname":(state.get(5) or {}).get(1),"seed_version":(state.get(10) or {}).get(1),"audit":_human(state.get(15)),"pending":_human(state.get(22)),"memory":verify}
|
||||
if domain == "events":
|
||||
items = []
|
||||
for record, record_hash, _ in self._iter_events():
|
||||
items.append({"record_type":RECORD_TYPE_NAMES.get(record.get(2),"unknown"),"timestamp_ns":record.get(3),"source":SOURCE_NAMES.get(record.get(6),"unknown"),"data":_human(record.get(5)),"sha256":record_hash.hex()})
|
||||
return {"ok":True,"domain":"events","events":items[-limit:]}
|
||||
if domain not in DOMAIN_TO_FIELD: raise ProductError("MEMORY_QUERY_DOMAIN_INVALID")
|
||||
return {"ok":True,"domain":domain,"value":_human(state.get(DOMAIN_TO_FIELD[domain]))}
|
||||
|
||||
def stats(self) -> dict[str, Any]:
|
||||
self.require_ready(); verify = self.verify()
|
||||
return {
|
||||
"ok":True,"protocol":FBM_PROTOCOL,
|
||||
"files": {
|
||||
"node.fbm":self.state_path.stat().st_size,
|
||||
"events.fbm":self.events_path.stat().st_size,
|
||||
"index.fbi":self.index_path.stat().st_size,
|
||||
"manifest.fbam":self.manifest_path.stat().st_size,
|
||||
},
|
||||
"events":verify.get("events"),"event_head_sha256":verify.get("event_head_sha256"),
|
||||
"state_compression_ratio":round(verify["state_compressed_bytes"] / max(1, verify["state_raw_bytes"]), 4),
|
||||
}
|
||||
+353
@@ -0,0 +1,353 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import sys
|
||||
import time
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any, Callable
|
||||
|
||||
from .memory_store import MemoryStore
|
||||
from .product_core import ProductError
|
||||
from .network_security import inspect_network
|
||||
from .pairing import PairingError, PairingStore, public_pairing_view
|
||||
from .pairing_client import PairingClientError, cancel_pairing_remote
|
||||
from .pairing_window import close_pairing_window, pairing_window_status
|
||||
from .seed_core import SeedContext, atomic_write_json
|
||||
|
||||
|
||||
MONITOR_SCHEMA = 1
|
||||
DEFAULT_INTERVAL_SECONDS = 1.0
|
||||
|
||||
|
||||
def _candidate_binding(candidate: dict[str, Any] | None) -> dict[str, Any]:
|
||||
c=candidate or {}
|
||||
return {
|
||||
"interface":str(c.get("interface") or ""),
|
||||
"source_subnet":str(c.get("subnet") or ""),
|
||||
"destination_address":str(c.get("address") or ""),
|
||||
"network_fingerprint":str(c.get("network_fingerprint") or ""),
|
||||
}
|
||||
|
||||
|
||||
def _window_binding(window: dict[str, Any]) -> dict[str, Any]:
|
||||
return {
|
||||
"interface":str(window.get("interface") or ""),
|
||||
"source_subnet":str(window.get("source_subnet") or ""),
|
||||
"destination_address":str(window.get("destination_address") or ""),
|
||||
"network_fingerprint":str(window.get("network_fingerprint") or ""),
|
||||
}
|
||||
|
||||
|
||||
def compare_binding(window: dict[str, Any], candidate: dict[str, Any] | None) -> dict[str, Any]:
|
||||
old=_window_binding(window)
|
||||
new=_candidate_binding(candidate)
|
||||
reasons=[]
|
||||
if not candidate:
|
||||
reasons.append("NO_ELIGIBLE_PAIRING_CANDIDATE")
|
||||
else:
|
||||
for key,label in (
|
||||
("interface","INTERFACE_CHANGED"),
|
||||
("source_subnet","SUBNET_CHANGED"),
|
||||
("destination_address","ADDRESS_CHANGED"),
|
||||
("network_fingerprint","NETWORK_FINGERPRINT_CHANGED"),
|
||||
):
|
||||
if old.get(key)!=new.get(key):
|
||||
reasons.append(label)
|
||||
return {
|
||||
"changed":bool(reasons),
|
||||
"reasons":reasons,
|
||||
"old":old,
|
||||
"new":new,
|
||||
}
|
||||
|
||||
|
||||
def _sync_pairing_config(ctx: SeedContext, view: dict[str, Any]) -> None:
|
||||
config=ctx.load_config() or {}
|
||||
pairing=config.get("pairing") or {}
|
||||
pairing["state"]=view.get("state","expired")
|
||||
if view.get("expiration_reason"):
|
||||
pairing["expiration_reason"]=view.get("expiration_reason")
|
||||
pairing["requires_new_local_action"]=True
|
||||
config["pairing"]=pairing
|
||||
config["connection_mode"]="PAIRING"
|
||||
setup=config.get("setup")
|
||||
if isinstance(setup,dict):
|
||||
setup["pairing_state"]=view.get("state","expired")
|
||||
atomic_write_json(ctx.config_path,config)
|
||||
|
||||
|
||||
def _invalidate_pending_pair(ctx: SeedContext, reason: str) -> dict[str, Any] | None:
|
||||
store=PairingStore(ctx.state_dir)
|
||||
current=store.load()
|
||||
if current.get("state")!="pair_requested":
|
||||
return None
|
||||
view=store.invalidate_pending(reason)
|
||||
_sync_pairing_config(ctx,view)
|
||||
return view
|
||||
|
||||
|
||||
def reconcile_broker_cancel(
|
||||
ctx: SeedContext,
|
||||
*,
|
||||
cancel_func=cancel_pairing_remote,
|
||||
retry_seconds: float=5.0,
|
||||
) -> dict[str, Any]:
|
||||
store=PairingStore(ctx.state_dir)
|
||||
current=store.load()
|
||||
if current.get("broker_cancel_required") is not True:
|
||||
return {"ok":True,"state":"NOT_REQUIRED","cancelled":False}
|
||||
pair_id=str(current.get("pair_id") or "")
|
||||
reason=str(current.get("broker_cancel_reason") or current.get("expiration_reason") or "reconcile")
|
||||
last=str(current.get("broker_cancel_last_attempt_at") or "")
|
||||
if last:
|
||||
try:
|
||||
last_dt=datetime.fromisoformat(last.replace("Z","+00:00")).astimezone(timezone.utc)
|
||||
age=(datetime.now(timezone.utc)-last_dt).total_seconds()
|
||||
if age<max(1.0,float(retry_seconds)):
|
||||
return {
|
||||
"ok":True,"state":"RETRY_DEFERRED","cancelled":False,
|
||||
"pair_id":pair_id,"retry_after":round(max(0.0,float(retry_seconds)-age),3),
|
||||
}
|
||||
except Exception:
|
||||
pass
|
||||
try:
|
||||
result=cancel_func(state_dir=ctx.state_dir,reason=reason,record=current)
|
||||
except PairingClientError as exc:
|
||||
store.note_broker_cancel_error(pair_id,str(exc))
|
||||
return {
|
||||
"ok":False,"state":"BROKER_CANCEL_PENDING","cancelled":False,
|
||||
"pair_id":pair_id,"error":str(exc),
|
||||
}
|
||||
updated=store.mark_broker_cancelled(pair_id,result)
|
||||
ctx.audit.append(
|
||||
"pair.broker_cancelled","PASS",
|
||||
request={"pair_id":pair_id,"reason":reason},
|
||||
result={"broker_state":result.get("state"),"idempotent":result.get("idempotent",False)},
|
||||
principal="network-monitor",
|
||||
)
|
||||
memory=MemoryStore(ctx)
|
||||
if memory.ready():
|
||||
memory.append(
|
||||
"recovery",
|
||||
{"operation":"pair.broker_cancelled","pair_id":pair_id,"reason":reason,"broker_state":result.get("state")},
|
||||
source="system",
|
||||
)
|
||||
return {
|
||||
"ok":True,"state":"BROKER_CANCELLED","cancelled":True,
|
||||
"pair_id":pair_id,"broker_state":result.get("state"),
|
||||
"local_state":updated.get("state"),
|
||||
}
|
||||
|
||||
|
||||
def _lifecycle_state_path(ctx: SeedContext):
|
||||
return ctx.state_dir/"network-lifecycle.json"
|
||||
|
||||
|
||||
def _sync_expired_window(ctx: SeedContext, window: dict[str, Any]) -> dict[str, Any]:
|
||||
window_id=str(window.get("window_id") or "")
|
||||
if window.get("state")=="expired" and window.get("rule_present") is not False:
|
||||
raise ProductError("NETWORK_WINDOW_EXPIRY_UNVERIFIED")
|
||||
if not window_id or window.get("state")!="expired" or window.get("close_reason")!="expired":
|
||||
return {"ok":True,"synced":False}
|
||||
path=_lifecycle_state_path(ctx)
|
||||
prior={}
|
||||
if path.is_file():
|
||||
try:
|
||||
prior=json.loads(path.read_text())
|
||||
except Exception:
|
||||
prior={}
|
||||
if prior.get("last_expiry_window_id")==window_id:
|
||||
return {"ok":True,"synced":False,"already_synced":True,"window_id":window_id}
|
||||
pair_view=_invalidate_pending_pair(ctx,"window_expired")
|
||||
broker_cancel=reconcile_broker_cancel(ctx)
|
||||
ctx.audit.append(
|
||||
"network.pairing_window.expired","PASS",
|
||||
request={"window_id":window_id,"expires_at":window.get("expires_at")},
|
||||
result={
|
||||
"rule_present":window.get("rule_present"),
|
||||
"pair_state":(pair_view or {}).get("state"),
|
||||
"requires_new_local_action":bool(pair_view),
|
||||
"broker_cancel_state":broker_cancel.get("state"),
|
||||
},
|
||||
principal="network-monitor",
|
||||
)
|
||||
memory=MemoryStore(ctx)
|
||||
if memory.ready():
|
||||
memory.append(
|
||||
"recovery",
|
||||
{
|
||||
"operation":"network.pairing_window.expired",
|
||||
"window_id":window_id,
|
||||
"rule_present":window.get("rule_present"),
|
||||
"pair_state":(pair_view or {}).get("state"),
|
||||
"requires_new_local_action":bool(pair_view),
|
||||
"broker_cancel_state":broker_cancel.get("state"),
|
||||
},source="system",
|
||||
)
|
||||
atomic_write_json(path,{"schema":MONITOR_SCHEMA,"last_expiry_window_id":window_id,"last_expiry_synced_at":time.time()})
|
||||
return {
|
||||
"ok":True,"synced":True,"window_id":window_id,
|
||||
"pair_state":(pair_view or {}).get("state"),
|
||||
"requires_new_local_action":bool(pair_view),
|
||||
"broker_cancel":broker_cancel,
|
||||
}
|
||||
|
||||
|
||||
|
||||
def _record_invalidation(
|
||||
ctx: SeedContext,
|
||||
*,
|
||||
reason: str,
|
||||
comparison: dict[str, Any],
|
||||
window: dict[str, Any],
|
||||
closed: dict[str, Any],
|
||||
pair_view: dict[str, Any] | None,
|
||||
) -> None:
|
||||
result={
|
||||
"window_id":window.get("window_id"),
|
||||
"close_reason":closed.get("close_reason"),
|
||||
"rule_present":closed.get("rule_present"),
|
||||
"pair_state":(pair_view or {}).get("state"),
|
||||
"requires_new_local_action":True,
|
||||
"reasons":comparison.get("reasons") or [reason],
|
||||
}
|
||||
ctx.audit.append(
|
||||
"network.pairing_window.invalidated",
|
||||
"PASS",
|
||||
request={
|
||||
"reason":reason,
|
||||
"old":comparison.get("old"),
|
||||
"new":comparison.get("new"),
|
||||
},
|
||||
result=result,
|
||||
principal="network-monitor",
|
||||
)
|
||||
store=MemoryStore(ctx)
|
||||
if store.ready():
|
||||
store.append(
|
||||
"incident",
|
||||
{
|
||||
"operation":"network.pairing_window.invalidated",
|
||||
"reason":reason,
|
||||
"reasons":comparison.get("reasons") or [reason],
|
||||
"window_id":window.get("window_id"),
|
||||
"old":comparison.get("old"),
|
||||
"new":comparison.get("new"),
|
||||
"pair_state":(pair_view or {}).get("state"),
|
||||
"requires_new_local_action":True,
|
||||
},
|
||||
source="system",
|
||||
)
|
||||
|
||||
|
||||
def invalidate_active_window(
|
||||
*,
|
||||
ctx: SeedContext | None=None,
|
||||
reason: str,
|
||||
comparison: dict[str, Any] | None=None,
|
||||
status_func: Callable[..., dict[str, Any]]=pairing_window_status,
|
||||
close_func: Callable[..., dict[str, Any]]=close_pairing_window,
|
||||
) -> dict[str, Any]:
|
||||
ctx=ctx or SeedContext()
|
||||
window=status_func(non_interactive=True)
|
||||
if window.get("active") is not True:
|
||||
return {"ok":True,"schema":MONITOR_SCHEMA,"state":"NO_ACTIVE_PAIRING_WINDOW","invalidated":False}
|
||||
if comparison is None:
|
||||
comparison={"changed":True,"reasons":[reason.upper()],"old":_window_binding(window),"new":{}}
|
||||
closed=close_func(reason,non_interactive=True)
|
||||
if closed.get("ok") is not True or closed.get("rule_present") is not False or closed.get("active") is True:
|
||||
raise ProductError("NETWORK_WINDOW_CLOSE_UNVERIFIED")
|
||||
pair_view=_invalidate_pending_pair(ctx,reason)
|
||||
broker_cancel=reconcile_broker_cancel(ctx)
|
||||
_record_invalidation(
|
||||
ctx,reason=reason,comparison=comparison,window=window,closed=closed,pair_view=pair_view,
|
||||
)
|
||||
return {
|
||||
"ok":True,"schema":MONITOR_SCHEMA,"state":"INVALIDATED","invalidated":True,
|
||||
"reason":reason,"reasons":comparison.get("reasons") or [],
|
||||
"window_id":window.get("window_id"),"rule_present":closed.get("rule_present"),
|
||||
"pair_state":(pair_view or {}).get("state"),"requires_new_local_action":True,
|
||||
"broker_cancel":broker_cancel,"old":comparison.get("old"),"new":comparison.get("new"),
|
||||
}
|
||||
|
||||
|
||||
|
||||
def reconcile_once(
|
||||
*,
|
||||
ctx: SeedContext | None=None,
|
||||
snapshot: dict[str, Any] | None=None,
|
||||
status_func: Callable[..., dict[str, Any]]=pairing_window_status,
|
||||
close_func: Callable[..., dict[str, Any]]=close_pairing_window,
|
||||
) -> dict[str, Any]:
|
||||
ctx=ctx or SeedContext()
|
||||
broker_cancel=reconcile_broker_cancel(ctx)
|
||||
window=status_func(non_interactive=True)
|
||||
if window.get("active") is not True:
|
||||
expiry_sync=_sync_expired_window(ctx,window)
|
||||
return {
|
||||
"ok":True,"schema":MONITOR_SCHEMA,
|
||||
"state":"PAIRING_WINDOW_EXPIRED_SYNCED" if expiry_sync.get("synced") else "NO_ACTIVE_PAIRING_WINDOW",
|
||||
"invalidated":False,"expiry_sync":expiry_sync,"broker_cancel":broker_cancel,
|
||||
}
|
||||
try:
|
||||
snap=snapshot if snapshot is not None else inspect_network()
|
||||
if not isinstance(snap,dict): raise ValueError("NETWORK_SNAPSHOT_NOT_OBJECT")
|
||||
except Exception:
|
||||
return invalidate_active_window(
|
||||
ctx=ctx,reason="reconcile",
|
||||
comparison={"changed":True,"reasons":["NETWORK_INSPECTION_FAILED"],"old":_window_binding(window),"new":{}},
|
||||
status_func=lambda **_:window,close_func=close_func,
|
||||
)
|
||||
candidate=snap.get("pairing_candidate")
|
||||
comparison=compare_binding(window,candidate)
|
||||
if not comparison["changed"]:
|
||||
return {
|
||||
"ok":True,"schema":MONITOR_SCHEMA,"state":"STABLE","invalidated":False,
|
||||
"window_id":window.get("window_id"),"binding":comparison["old"],"broker_cancel":broker_cancel,
|
||||
}
|
||||
return invalidate_active_window(
|
||||
ctx=ctx,reason="network_change",comparison=comparison,
|
||||
status_func=lambda **_:window,close_func=close_func,
|
||||
)
|
||||
|
||||
|
||||
|
||||
def run_forever(interval: float=DEFAULT_INTERVAL_SECONDS) -> None:
|
||||
interval=max(0.25,float(interval))
|
||||
ctx=SeedContext()
|
||||
while True:
|
||||
try:
|
||||
result=reconcile_once(ctx=ctx)
|
||||
if result.get("invalidated") is True:
|
||||
print(json.dumps(result,sort_keys=True),flush=True)
|
||||
except Exception as exc:
|
||||
print(json.dumps({
|
||||
"ok":False,
|
||||
"schema":MONITOR_SCHEMA,
|
||||
"state":"MONITOR_ERROR",
|
||||
"error":type(exc).__name__,
|
||||
"message":str(exc),
|
||||
},sort_keys=True),file=sys.stderr,flush=True)
|
||||
time.sleep(interval)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser=argparse.ArgumentParser(prog="framebridge-network-monitor")
|
||||
sub=parser.add_subparsers(dest="command",required=True)
|
||||
run=sub.add_parser("run")
|
||||
run.add_argument("--interval",type=float,default=DEFAULT_INTERVAL_SECONDS)
|
||||
sub.add_parser("once")
|
||||
close=sub.add_parser("close")
|
||||
close.add_argument("--reason",choices=("service_stop","reconcile"),default="service_stop")
|
||||
args=parser.parse_args()
|
||||
if args.command=="run":
|
||||
run_forever(args.interval)
|
||||
elif args.command=="once":
|
||||
print(json.dumps(reconcile_once(),sort_keys=True,indent=2))
|
||||
else:
|
||||
print(json.dumps(invalidate_active_window(reason=args.reason),sort_keys=True,indent=2))
|
||||
|
||||
|
||||
if __name__=="__main__":
|
||||
main()
|
||||
+513
@@ -0,0 +1,513 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import ipaddress
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
import subprocess
|
||||
from typing import Any
|
||||
|
||||
SCHEMA = 1
|
||||
PAIRING_WINDOW_TTL_SECONDS = 600
|
||||
|
||||
RFC1918 = (
|
||||
ipaddress.ip_network("10.0.0.0/8"),
|
||||
ipaddress.ip_network("172.16.0.0/12"),
|
||||
ipaddress.ip_network("192.168.0.0/16"),
|
||||
)
|
||||
ULA_V6 = ipaddress.ip_network("fc00::/7")
|
||||
|
||||
OVERLAY_NAMES = {"tailscale0"}
|
||||
VIRTUAL_PREFIXES = (
|
||||
"docker", "br-", "virbr", "veth", "lxc", "waydroid",
|
||||
"tap", "tun", "zt", "wg",
|
||||
)
|
||||
|
||||
|
||||
def _run(argv: list[str], timeout: int = 5) -> subprocess.CompletedProcess:
|
||||
return subprocess.run(
|
||||
argv,
|
||||
text=True,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.PIPE,
|
||||
timeout=timeout,
|
||||
check=False,
|
||||
)
|
||||
|
||||
|
||||
def _run_json(argv: list[str]) -> list[dict[str, Any]]:
|
||||
try:
|
||||
cp = _run(argv)
|
||||
except Exception:
|
||||
return []
|
||||
if cp.returncode != 0:
|
||||
return []
|
||||
try:
|
||||
value = json.loads(cp.stdout or "[]")
|
||||
except Exception:
|
||||
return []
|
||||
return value if isinstance(value, list) else []
|
||||
|
||||
|
||||
def address_class(address: str) -> str:
|
||||
try:
|
||||
ip = ipaddress.ip_address(str(address))
|
||||
except ValueError:
|
||||
return "INVALID"
|
||||
if ip.is_loopback:
|
||||
return "LOOPBACK"
|
||||
if ip.version == 4:
|
||||
if any(ip in net for net in RFC1918):
|
||||
return "PRIVATE_LAN"
|
||||
if ip.is_link_local:
|
||||
return "LINK_LOCAL"
|
||||
if ip.is_global:
|
||||
return "PUBLIC"
|
||||
return "SPECIAL"
|
||||
if ip in ULA_V6:
|
||||
return "ULA"
|
||||
if ip.is_link_local:
|
||||
return "LINK_LOCAL"
|
||||
if ip.is_global:
|
||||
return "PUBLIC"
|
||||
return "SPECIAL"
|
||||
|
||||
|
||||
def classify_interface(
|
||||
name: str,
|
||||
*,
|
||||
has_device: bool = False,
|
||||
wireless: bool = False,
|
||||
) -> str:
|
||||
n = str(name or "")
|
||||
if n == "lo":
|
||||
return "LOOPBACK"
|
||||
if n in OVERLAY_NAMES:
|
||||
return "PRIVATE_OVERLAY"
|
||||
if n.startswith(VIRTUAL_PREFIXES):
|
||||
return "VIRTUAL_LOCAL"
|
||||
if wireless or n.startswith(("wl", "wlan")):
|
||||
return "PHYSICAL_LAN"
|
||||
if has_device or n.startswith(("en", "eth")):
|
||||
return "PHYSICAL_LAN"
|
||||
return "UNKNOWN"
|
||||
|
||||
|
||||
def _sysfs_interface_kind(name: str) -> str:
|
||||
base = Path("/sys/class/net") / name
|
||||
return classify_interface(
|
||||
name,
|
||||
has_device=(base / "device").exists(),
|
||||
wireless=(base / "wireless").exists(),
|
||||
)
|
||||
|
||||
|
||||
def _prefix_subnet(address: str, prefixlen: int) -> str:
|
||||
try:
|
||||
return str(ipaddress.ip_interface(f"{address}/{int(prefixlen)}").network)
|
||||
except Exception:
|
||||
return ""
|
||||
|
||||
|
||||
def _default_routes() -> list[dict[str, Any]]:
|
||||
routes = _run_json(["ip", "-j", "-4", "route", "show", "default"])
|
||||
clean: list[dict[str, Any]] = []
|
||||
for r in routes:
|
||||
if not isinstance(r, dict):
|
||||
continue
|
||||
clean.append(
|
||||
{
|
||||
"dev": str(r.get("dev") or ""),
|
||||
"gateway": str(r.get("gateway") or ""),
|
||||
"metric": int(r.get("metric") or 0),
|
||||
"prefsrc": str(r.get("prefsrc") or ""),
|
||||
}
|
||||
)
|
||||
return clean
|
||||
|
||||
|
||||
def _network_manager_identity(ifname: str) -> dict[str, str]:
|
||||
result = {"connection": "", "connection_uuid": ""}
|
||||
try:
|
||||
cp = _run(["nmcli", "-g", "GENERAL.CONNECTION", "device", "show", ifname])
|
||||
except Exception:
|
||||
return result
|
||||
if cp.returncode != 0:
|
||||
return result
|
||||
connection = (cp.stdout or "").strip().splitlines()
|
||||
if not connection:
|
||||
return result
|
||||
name = connection[0].strip()
|
||||
if not name or name == "--":
|
||||
return result
|
||||
result["connection"] = name
|
||||
try:
|
||||
cp2 = _run(["nmcli", "-g", "connection.uuid", "connection", "show", name])
|
||||
except Exception:
|
||||
return result
|
||||
if cp2.returncode == 0:
|
||||
result["connection_uuid"] = (cp2.stdout or "").strip().splitlines()[0] if (cp2.stdout or "").strip() else ""
|
||||
return result
|
||||
|
||||
|
||||
def collect_interfaces() -> list[dict[str, Any]]:
|
||||
links = _run_json(["ip", "-j", "link", "show"])
|
||||
addrs = _run_json(["ip", "-j", "addr", "show"])
|
||||
addr_by_name = {str(x.get("ifname") or ""): x for x in addrs if isinstance(x, dict)}
|
||||
out: list[dict[str, Any]] = []
|
||||
for link in links:
|
||||
if not isinstance(link, dict):
|
||||
continue
|
||||
name = str(link.get("ifname") or "")
|
||||
if not name:
|
||||
continue
|
||||
a = addr_by_name.get(name, {})
|
||||
ipv4 = []
|
||||
ipv6 = []
|
||||
for ai in a.get("addr_info") or []:
|
||||
if not isinstance(ai, dict):
|
||||
continue
|
||||
local = str(ai.get("local") or "")
|
||||
if not local:
|
||||
continue
|
||||
item = {
|
||||
"address": local,
|
||||
"prefixlen": int(ai.get("prefixlen") or 0),
|
||||
"class": address_class(local),
|
||||
"scope": str(ai.get("scope") or ""),
|
||||
}
|
||||
if ai.get("family") == "inet":
|
||||
ipv4.append(item)
|
||||
elif ai.get("family") == "inet6":
|
||||
ipv6.append(item)
|
||||
kind = _sysfs_interface_kind(name)
|
||||
out.append(
|
||||
{
|
||||
"name": name,
|
||||
"kind": kind,
|
||||
"operstate": str(link.get("operstate") or ""),
|
||||
"flags": list(link.get("flags") or []),
|
||||
"ipv4": ipv4,
|
||||
"ipv6": ipv6,
|
||||
}
|
||||
)
|
||||
return out
|
||||
|
||||
|
||||
def choose_pairing_candidate(
|
||||
interfaces: list[dict[str, Any]],
|
||||
routes: list[dict[str, Any]],
|
||||
) -> dict[str, Any] | None:
|
||||
"""Select only an unambiguous physical private-LAN binding.
|
||||
|
||||
Metric zero is valid. Route preferred source is honored. Equal-priority
|
||||
candidates that differ in interface/address/prefix/gateway are blocked,
|
||||
not resolved by comparing payload dictionaries or snapshot order.
|
||||
"""
|
||||
route_by_dev: dict[str, list[dict[str, Any]]] = {}
|
||||
for route in routes:
|
||||
route_by_dev.setdefault(str(route.get("dev") or ""), []).append(route)
|
||||
|
||||
choices: dict[tuple, tuple[tuple, dict[str, Any]]] = {}
|
||||
for iface in interfaces:
|
||||
if iface.get("kind") != "PHYSICAL_LAN":
|
||||
continue
|
||||
if str(iface.get("operstate") or "").upper() not in {"UP", "UNKNOWN"}:
|
||||
continue
|
||||
name = str(iface.get("name") or "")
|
||||
if not name:
|
||||
continue
|
||||
dev_routes = route_by_dev.get(name, [])
|
||||
if dev_routes:
|
||||
minimum = min(int(r.get("metric") or 0) for r in dev_routes)
|
||||
best_routes = [r for r in dev_routes if int(r.get("metric") or 0) == minimum]
|
||||
else:
|
||||
best_routes = [{}]
|
||||
for selected_route in best_routes:
|
||||
preferred = str(selected_route.get("prefsrc") or "")
|
||||
metric = int(selected_route.get("metric") or 0) if dev_routes else None
|
||||
for addr in iface.get("ipv4") or []:
|
||||
address = str(addr.get("address") or "")
|
||||
if address_class(address) != "PRIVATE_LAN":
|
||||
continue
|
||||
if preferred and address != preferred:
|
||||
continue
|
||||
prefix = addr.get("prefixlen")
|
||||
if isinstance(prefix, bool) or not isinstance(prefix, int) or not 1 <= prefix <= 32:
|
||||
continue
|
||||
gateway = str(selected_route.get("gateway") or "")
|
||||
item = {
|
||||
"interface": name,
|
||||
"interface_kind": "PHYSICAL_LAN",
|
||||
"address": address,
|
||||
"prefixlen": prefix,
|
||||
"subnet": _prefix_subnet(address, prefix),
|
||||
"address_class": "PRIVATE_LAN",
|
||||
"gateway": gateway,
|
||||
"route_metric": metric,
|
||||
}
|
||||
rank = (0 if dev_routes else 1, metric if metric is not None else 0,
|
||||
0 if preferred else 1)
|
||||
identity = (name, address, prefix, gateway, metric)
|
||||
choices[identity] = (rank, item)
|
||||
if not choices:
|
||||
return None
|
||||
best_rank = min(rank for rank, item in choices.values())
|
||||
top = [item for rank, item in choices.values() if rank == best_rank]
|
||||
if len(top) != 1:
|
||||
return None
|
||||
selected = dict(top[0])
|
||||
selected.update(_network_manager_identity(selected["interface"]))
|
||||
binding = {key: selected.get(key) for key in (
|
||||
"interface", "address", "prefixlen", "gateway", "route_metric", "connection_uuid"
|
||||
)}
|
||||
selected["network_fingerprint_schema"] = 2
|
||||
raw = json.dumps(binding, sort_keys=True, separators=(",", ":")).encode("utf-8")
|
||||
selected["network_fingerprint"] = hashlib.sha256(raw).hexdigest()
|
||||
return selected
|
||||
|
||||
|
||||
def _service_active(name: str) -> bool:
|
||||
try:
|
||||
cp = _run(["systemctl", "is-active", "--quiet", name])
|
||||
except Exception:
|
||||
return False
|
||||
return cp.returncode == 0
|
||||
|
||||
|
||||
def _ufw_config_enabled(path: Path=Path("/etc/ufw/ufw.conf")) -> bool:
|
||||
try:
|
||||
text=path.read_text(encoding="utf-8",errors="replace")
|
||||
except OSError:
|
||||
return False
|
||||
for raw in text.splitlines():
|
||||
line=raw.strip()
|
||||
if not line or line.startswith("#"):
|
||||
continue
|
||||
if line.upper().startswith("ENABLED="):
|
||||
return line.split("=",1)[1].strip().lower() in {"yes","true","1"}
|
||||
return False
|
||||
|
||||
|
||||
def firewall_backend() -> dict[str, Any]:
|
||||
ufw_systemd_active=_service_active("ufw")
|
||||
ufw_config_enabled=_ufw_config_enabled()
|
||||
if ufw_systemd_active or ufw_config_enabled:
|
||||
return {
|
||||
"backend": "ufw",
|
||||
"active": True,
|
||||
"systemd_active": ufw_systemd_active,
|
||||
"config_enabled": ufw_config_enabled,
|
||||
"managed_pairing_rules": True,
|
||||
}
|
||||
if _service_active("firewalld"):
|
||||
return {"backend": "firewalld", "active": True, "managed_pairing_rules": False}
|
||||
if _service_active("nftables"):
|
||||
return {"backend": "nftables", "active": True, "managed_pairing_rules": False}
|
||||
return {"backend": "none", "active": False, "managed_pairing_rules": False}
|
||||
|
||||
|
||||
def ufw_pairing_plan(candidate: dict[str, Any], port: int) -> dict[str, Any]:
|
||||
iface = str(candidate.get("interface") or "")
|
||||
subnet = str(candidate.get("subnet") or "")
|
||||
address = str(candidate.get("address") or "")
|
||||
fp = str(candidate.get("network_fingerprint") or "")
|
||||
if not iface or not subnet or not address or not fp:
|
||||
return {"ok": False, "error": "PAIRING_CANDIDATE_INCOMPLETE"}
|
||||
return {
|
||||
"ok": True,
|
||||
"backend": "ufw",
|
||||
"execution_backend": "iptables-runtime-inside-ufw",
|
||||
"chain": "ufw-user-input",
|
||||
"execute": False,
|
||||
"operation": "runtime_scoped_accept",
|
||||
"interface": iface,
|
||||
"source_subnet": subnet,
|
||||
"destination_address": address,
|
||||
"port": int(port),
|
||||
"protocol": "tcp",
|
||||
"network_fingerprint": fp,
|
||||
"ttl_seconds": PAIRING_WINDOW_TTL_SECONDS,
|
||||
"requires_local_tty": True,
|
||||
"runtime_only": True,
|
||||
"reboot_fail_closed": True,
|
||||
"automatic_expiry": True,
|
||||
"persistent_ufw_rule": False,
|
||||
"rule_comment_template": "FrameBridge:<window_id>:pairing",
|
||||
}
|
||||
|
||||
|
||||
def inspect_network() -> dict[str, Any]:
|
||||
interfaces = collect_interfaces()
|
||||
routes = _default_routes()
|
||||
candidate = choose_pairing_candidate(interfaces, routes)
|
||||
fw = firewall_backend()
|
||||
return {
|
||||
"ok": True,
|
||||
"schema": SCHEMA,
|
||||
"interfaces": interfaces,
|
||||
"default_routes": routes,
|
||||
"pairing_candidate": candidate,
|
||||
"firewall": fw,
|
||||
"policy": {
|
||||
"local_pairing_transport": "LAN_HTTP",
|
||||
"tailscale_required": False,
|
||||
"public_http_allowed": False,
|
||||
"wildcard_pairing_bind_allowed": False,
|
||||
"local_tty_required": True,
|
||||
"pairing_window_ttl_seconds": PAIRING_WINDOW_TTL_SECONDS,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def pairing_policy(port: int = 18765) -> dict[str, Any]:
|
||||
snap = inspect_network()
|
||||
candidate = snap.get("pairing_candidate")
|
||||
fw = snap.get("firewall") or {}
|
||||
if not candidate:
|
||||
return {
|
||||
"ok": False,
|
||||
"schema": SCHEMA,
|
||||
"state": "BLOCKED",
|
||||
"reason": "NO_ELIGIBLE_PHYSICAL_PRIVATE_LAN",
|
||||
"execute": False,
|
||||
"snapshot": snap,
|
||||
}
|
||||
if not fw.get("active"):
|
||||
return {
|
||||
"ok": False,
|
||||
"schema": SCHEMA,
|
||||
"state": "BLOCKED",
|
||||
"reason": "NO_ACTIVE_FIREWALL_BACKEND",
|
||||
"execute": False,
|
||||
"snapshot": snap,
|
||||
}
|
||||
if fw.get("backend") != "ufw":
|
||||
return {
|
||||
"ok": False,
|
||||
"schema": SCHEMA,
|
||||
"state": "BLOCKED",
|
||||
"reason": "FIREWALL_BACKEND_NOT_YET_MANAGED",
|
||||
"execute": False,
|
||||
"snapshot": snap,
|
||||
}
|
||||
return {
|
||||
"ok": True,
|
||||
"schema": SCHEMA,
|
||||
"state": "ELIGIBLE_REQUIRES_LOCAL_TTY",
|
||||
"execute": False,
|
||||
"candidate": candidate,
|
||||
"firewall_plan": ufw_pairing_plan(candidate, int(port)),
|
||||
"network_change_policy": "INVALIDATE_PENDING_PAIR_AND_CLOSE_OLD_WINDOW",
|
||||
}
|
||||
|
||||
CGNAT_V4 = ipaddress.ip_network("100.64.0.0/10")
|
||||
FORWARD_PROTOCOLS = {"any", "tcp", "udp"}
|
||||
|
||||
|
||||
def _valid_interface_name(value: str) -> bool:
|
||||
import re
|
||||
return re.fullmatch(r"[A-Za-z0-9_.:-]{1,32}", str(value or "")) is not None
|
||||
|
||||
|
||||
def _normalize_ipv4_network(value: str) -> str:
|
||||
try:
|
||||
net = ipaddress.ip_network(str(value), strict=False)
|
||||
except ValueError as exc:
|
||||
raise ValueError("NETWORK_CIDR_INVALID") from exc
|
||||
if net.version != 4:
|
||||
raise ValueError("NETWORK_IPV4_REQUIRED")
|
||||
if net.is_multicast or net.is_loopback or net.is_link_local:
|
||||
raise ValueError("NETWORK_CIDR_UNSAFE")
|
||||
return str(net)
|
||||
|
||||
|
||||
def _source_network_allowed(value: str) -> bool:
|
||||
net = ipaddress.ip_network(value, strict=False)
|
||||
if net.version != 4:
|
||||
return False
|
||||
if net.subnet_of(CGNAT_V4):
|
||||
return True
|
||||
return any(net.subnet_of(x) for x in RFC1918)
|
||||
|
||||
|
||||
def forwarding_policy(
|
||||
*,
|
||||
ingress_interface: str,
|
||||
egress_interface: str,
|
||||
source_subnet: str,
|
||||
destination_subnet: str,
|
||||
protocol: str = "any",
|
||||
destination_port: int = 0,
|
||||
nat: bool = False,
|
||||
lifetime: str = "runtime",
|
||||
) -> dict[str, Any]:
|
||||
ingress_interface = str(ingress_interface or "").strip()
|
||||
egress_interface = str(egress_interface or "").strip()
|
||||
protocol = str(protocol or "any").strip().lower()
|
||||
lifetime = str(lifetime or "runtime").strip().lower()
|
||||
try:
|
||||
destination_port = int(destination_port)
|
||||
except Exception:
|
||||
return {"ok": False, "state": "BLOCKED", "reason": "NETWORK_FORWARD_PORT_INVALID", "execute": False}
|
||||
|
||||
if not _valid_interface_name(ingress_interface) or not _valid_interface_name(egress_interface):
|
||||
return {"ok": False, "state": "BLOCKED", "reason": "NETWORK_FORWARD_INTERFACE_INVALID", "execute": False}
|
||||
if ingress_interface == egress_interface:
|
||||
return {"ok": False, "state": "BLOCKED", "reason": "NETWORK_FORWARD_INTERFACES_MUST_DIFFER", "execute": False}
|
||||
|
||||
try:
|
||||
source_subnet = _normalize_ipv4_network(source_subnet)
|
||||
destination_subnet = _normalize_ipv4_network(destination_subnet)
|
||||
except ValueError as exc:
|
||||
return {"ok": False, "state": "BLOCKED", "reason": str(exc), "execute": False}
|
||||
|
||||
if not _source_network_allowed(source_subnet):
|
||||
return {"ok": False, "state": "BLOCKED", "reason": "NETWORK_FORWARD_SOURCE_NOT_PRIVATE", "execute": False}
|
||||
|
||||
if protocol not in FORWARD_PROTOCOLS:
|
||||
return {"ok": False, "state": "BLOCKED", "reason": "NETWORK_FORWARD_PROTOCOL_INVALID", "execute": False}
|
||||
if destination_port < 0 or destination_port > 65535:
|
||||
return {"ok": False, "state": "BLOCKED", "reason": "NETWORK_FORWARD_PORT_INVALID", "execute": False}
|
||||
if protocol == "any" and destination_port != 0:
|
||||
return {"ok": False, "state": "BLOCKED", "reason": "NETWORK_FORWARD_PORT_REQUIRES_PROTOCOL", "execute": False}
|
||||
if lifetime != "runtime":
|
||||
return {"ok": False, "state": "BLOCKED", "reason": "NETWORK_FORWARD_PERSISTENCE_REQUIRES_SEPARATE_CAPABILITY", "execute": False}
|
||||
|
||||
fw = firewall_backend()
|
||||
if not fw.get("active"):
|
||||
return {"ok": False, "state": "BLOCKED", "reason": "NO_ACTIVE_FIREWALL_BACKEND", "execute": False}
|
||||
if fw.get("backend") != "ufw":
|
||||
return {"ok": False, "state": "BLOCKED", "reason": "FIREWALL_BACKEND_NOT_YET_MANAGED", "execute": False}
|
||||
|
||||
names = {x.get("name") for x in collect_interfaces()}
|
||||
if ingress_interface not in names or egress_interface not in names:
|
||||
return {"ok": False, "state": "BLOCKED", "reason": "NETWORK_FORWARD_INTERFACE_NOT_PRESENT", "execute": False}
|
||||
|
||||
args = {
|
||||
"ingress_interface": ingress_interface,
|
||||
"egress_interface": egress_interface,
|
||||
"source_subnet": source_subnet,
|
||||
"destination_subnet": destination_subnet,
|
||||
"protocol": protocol,
|
||||
"destination_port": destination_port,
|
||||
"nat": bool(nat),
|
||||
"lifetime": "runtime",
|
||||
}
|
||||
return {
|
||||
"ok": True,
|
||||
"schema": SCHEMA,
|
||||
"state": "ELIGIBLE_REQUIRES_TYPED_PRIVILEGED_PLAN",
|
||||
"execute": False,
|
||||
"capability": "network.forward",
|
||||
"args": args,
|
||||
"security": {
|
||||
"local_tty_privileged_required": True,
|
||||
"independent_rollback_approval_required": True,
|
||||
"runtime_only": True,
|
||||
"reboot_fail_closed": True,
|
||||
"nat_explicit": bool(nat),
|
||||
"wildcard_source_forbidden": True,
|
||||
},
|
||||
}
|
||||
+239
@@ -0,0 +1,239 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
import stat
|
||||
import struct
|
||||
import subprocess
|
||||
from pathlib import Path
|
||||
from typing import Any, Callable
|
||||
|
||||
NODE_KEY_NAMESPACE = "framebridge-pairing"
|
||||
NODE_KEY_ALGORITHM = "ssh-ed25519"
|
||||
|
||||
|
||||
class NodeIdentityError(RuntimeError):
|
||||
pass
|
||||
|
||||
|
||||
def _b64url(data: bytes) -> str:
|
||||
return base64.urlsafe_b64encode(data).decode("ascii").rstrip("=")
|
||||
|
||||
|
||||
def _ssh_string(blob: bytes, offset: int) -> tuple[bytes, int]:
|
||||
if offset + 4 > len(blob):
|
||||
raise NodeIdentityError("NODE_PUBLIC_KEY_INVALID")
|
||||
size = struct.unpack(">I", blob[offset:offset + 4])[0]
|
||||
offset += 4
|
||||
end = offset + size
|
||||
if end > len(blob):
|
||||
raise NodeIdentityError("NODE_PUBLIC_KEY_INVALID")
|
||||
return blob[offset:end], end
|
||||
|
||||
|
||||
def read_ed25519_public_key(public_path: Path) -> dict[str, Any]:
|
||||
path = Path(public_path)
|
||||
if not path.is_file() or path.is_symlink():
|
||||
raise NodeIdentityError("NODE_PUBLIC_KEY_MISSING")
|
||||
parts = path.read_text(encoding="utf-8").strip().split()
|
||||
if len(parts) < 2 or parts[0] != NODE_KEY_ALGORITHM:
|
||||
raise NodeIdentityError("NODE_PUBLIC_KEY_INVALID")
|
||||
try:
|
||||
blob = base64.b64decode(parts[1], validate=True)
|
||||
except Exception as exc:
|
||||
raise NodeIdentityError("NODE_PUBLIC_KEY_INVALID") from exc
|
||||
alg, off = _ssh_string(blob, 0)
|
||||
raw, off = _ssh_string(blob, off)
|
||||
if off != len(blob) or alg != b"ssh-ed25519" or len(raw) != 32:
|
||||
raise NodeIdentityError("NODE_PUBLIC_KEY_INVALID")
|
||||
x = _b64url(raw)
|
||||
jwk = {"crv": "Ed25519", "kty": "OKP", "x": x}
|
||||
canonical = json.dumps(
|
||||
jwk, sort_keys=True, separators=(",", ":"), ensure_ascii=True
|
||||
).encode("ascii")
|
||||
thumbprint = _b64url(hashlib.sha256(canonical).digest())
|
||||
host_id = (
|
||||
"urn:ietf:params:oauth:jwk-thumbprint:sha-256:" + thumbprint
|
||||
)
|
||||
return {
|
||||
"algorithm": "Ed25519",
|
||||
"public_jwk": jwk,
|
||||
"jwk_thumbprint_sha256": thumbprint,
|
||||
"host_id": host_id,
|
||||
"public_key_sha256": hashlib.sha256(blob).hexdigest(),
|
||||
}
|
||||
|
||||
|
||||
def ensure_node_keypair(
|
||||
state_dir: Path,
|
||||
*,
|
||||
comment: str = "framebridge-node",
|
||||
runner: Callable[..., subprocess.CompletedProcess[str]] = subprocess.run,
|
||||
) -> dict[str, Any]:
|
||||
root = Path(state_dir) / "node-identity"
|
||||
private_path = root / "node_ed25519"
|
||||
public_path = root / "node_ed25519.pub"
|
||||
|
||||
if root.exists() and root.is_symlink():
|
||||
raise NodeIdentityError("NODE_KEY_DIR_SYMLINK_FORBIDDEN")
|
||||
root.mkdir(parents=True, exist_ok=True, mode=0o700)
|
||||
os.chmod(root, 0o700)
|
||||
|
||||
private_exists = private_path.exists()
|
||||
public_exists = public_path.exists()
|
||||
if private_exists != public_exists:
|
||||
raise NodeIdentityError("NODE_KEYPAIR_INCOMPLETE")
|
||||
|
||||
if not private_exists:
|
||||
proc = runner(
|
||||
[
|
||||
"/usr/bin/ssh-keygen",
|
||||
"-q",
|
||||
"-t",
|
||||
"ed25519",
|
||||
"-N",
|
||||
"",
|
||||
"-C",
|
||||
str(comment),
|
||||
"-f",
|
||||
str(private_path),
|
||||
],
|
||||
text=True,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.PIPE,
|
||||
check=False,
|
||||
)
|
||||
if proc.returncode != 0:
|
||||
raise NodeIdentityError("NODE_KEYGEN_FAILED")
|
||||
|
||||
if (
|
||||
not private_path.is_file()
|
||||
or private_path.is_symlink()
|
||||
or not public_path.is_file()
|
||||
or public_path.is_symlink()
|
||||
):
|
||||
raise NodeIdentityError("NODE_KEYPAIR_INVALID")
|
||||
|
||||
os.chmod(private_path, 0o600)
|
||||
os.chmod(public_path, 0o644)
|
||||
mode = stat.S_IMODE(private_path.stat().st_mode)
|
||||
if mode != 0o600:
|
||||
raise NodeIdentityError("NODE_PRIVATE_KEY_MODE_INVALID")
|
||||
|
||||
public = read_ed25519_public_key(public_path)
|
||||
return {
|
||||
"ok": True,
|
||||
"private_key_path": str(private_path),
|
||||
"public_key_path": str(public_path),
|
||||
"public_key_ssh": _safe_public_line(public_path),
|
||||
**public,
|
||||
}
|
||||
|
||||
|
||||
def _safe_public_line(public_path: Path) -> str:
|
||||
parts=Path(public_path).read_text(encoding="utf-8").strip().split()
|
||||
if len(parts)<2 or parts[0] != NODE_KEY_ALGORITHM:
|
||||
raise NodeIdentityError("NODE_PUBLIC_KEY_INVALID")
|
||||
return parts[0] + " " + parts[1]
|
||||
|
||||
|
||||
def sign_pairing_payload(
|
||||
private_path: Path,
|
||||
payload: bytes,
|
||||
*,
|
||||
runner: Callable[..., subprocess.CompletedProcess[str]] = subprocess.run,
|
||||
) -> str:
|
||||
private_path=Path(private_path)
|
||||
if not private_path.is_file() or private_path.is_symlink():
|
||||
raise NodeIdentityError("NODE_PRIVATE_KEY_MISSING")
|
||||
import tempfile
|
||||
with tempfile.TemporaryDirectory(prefix="framebridge-pair-sign-") as td:
|
||||
msg=Path(td)/"message"
|
||||
msg.write_bytes(bytes(payload))
|
||||
proc=runner(
|
||||
[
|
||||
"/usr/bin/ssh-keygen","-Y","sign",
|
||||
"-f",str(private_path),
|
||||
"-n",NODE_KEY_NAMESPACE,
|
||||
str(msg),
|
||||
],
|
||||
text=True,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.PIPE,
|
||||
check=False,
|
||||
)
|
||||
sig=Path(str(msg)+".sig")
|
||||
if proc.returncode!=0 or not sig.is_file():
|
||||
raise NodeIdentityError("NODE_PAIR_SIGNATURE_FAILED")
|
||||
return base64.b64encode(sig.read_bytes()).decode("ascii")
|
||||
|
||||
|
||||
def verify_pairing_payload(
|
||||
public_key_ssh: str,
|
||||
payload: bytes,
|
||||
signature_b64: str,
|
||||
*,
|
||||
runner: Callable[..., subprocess.CompletedProcess[str]] = subprocess.run,
|
||||
) -> bool:
|
||||
import tempfile
|
||||
try:
|
||||
signature=base64.b64decode(str(signature_b64),validate=True)
|
||||
except Exception:
|
||||
return False
|
||||
public_key_ssh=str(public_key_ssh).strip()
|
||||
if not public_key_ssh.startswith("ssh-ed25519 "):
|
||||
return False
|
||||
with tempfile.TemporaryDirectory(prefix="framebridge-pair-verify-") as td:
|
||||
root=Path(td)
|
||||
msg=root/"message"
|
||||
sig=root/"message.sig"
|
||||
allowed=root/"allowed_signers"
|
||||
msg.write_bytes(bytes(payload))
|
||||
sig.write_bytes(signature)
|
||||
allowed.write_text(
|
||||
"framebridge-node " + public_key_ssh + "\n",
|
||||
encoding="utf-8",
|
||||
)
|
||||
proc=runner(
|
||||
[
|
||||
"/usr/bin/ssh-keygen","-Y","verify",
|
||||
"-f",str(allowed),
|
||||
"-I","framebridge-node",
|
||||
"-n",NODE_KEY_NAMESPACE,
|
||||
"-s",str(sig),
|
||||
],
|
||||
input=bytes(payload),
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.PIPE,
|
||||
check=False,
|
||||
)
|
||||
return proc.returncode==0
|
||||
|
||||
|
||||
def parse_ed25519_public_key_line(line: str) -> dict[str, Any]:
|
||||
parts=str(line).strip().split()
|
||||
if len(parts)<2 or parts[0] != NODE_KEY_ALGORITHM:
|
||||
raise NodeIdentityError("NODE_PUBLIC_KEY_INVALID")
|
||||
try:
|
||||
blob=base64.b64decode(parts[1],validate=True)
|
||||
except Exception as exc:
|
||||
raise NodeIdentityError("NODE_PUBLIC_KEY_INVALID") from exc
|
||||
alg,off=_ssh_string(blob,0)
|
||||
raw,off=_ssh_string(blob,off)
|
||||
if off!=len(blob) or alg!=b"ssh-ed25519" or len(raw)!=32:
|
||||
raise NodeIdentityError("NODE_PUBLIC_KEY_INVALID")
|
||||
x=_b64url(raw)
|
||||
jwk={"crv":"Ed25519","kty":"OKP","x":x}
|
||||
canonical=json.dumps(
|
||||
jwk,sort_keys=True,separators=(",",":"),ensure_ascii=True
|
||||
).encode("ascii")
|
||||
thumbprint=_b64url(hashlib.sha256(canonical).digest())
|
||||
return {
|
||||
"algorithm":"Ed25519",
|
||||
"public_jwk":jwk,
|
||||
"jwk_thumbprint_sha256":thumbprint,
|
||||
"host_id":"urn:ietf:params:oauth:jwk-thumbprint:sha-256:"+thumbprint,
|
||||
"public_key_sha256":hashlib.sha256(blob).hexdigest(),
|
||||
}
|
||||
+287
@@ -0,0 +1,287 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import secrets
|
||||
import time
|
||||
import urllib.error
|
||||
import urllib.request
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from .node_identity import ensure_node_keypair, sign_pairing_payload
|
||||
from .pairing import PairingStore, canonical
|
||||
from .pairing_client import PairingClientError, confirm_pairing_remote
|
||||
from .product_core import ProductError
|
||||
from .seed_capabilities import execute_capability, list_capabilities
|
||||
from .seed_core import SeedContext, atomic_write_json
|
||||
from .plan_engine import PlanEngine
|
||||
from .memory_store import MemoryStore
|
||||
|
||||
READ_CAPS={x["name"] for x in list_capabilities() if x.get("risk")=="READONLY"}
|
||||
|
||||
|
||||
class OutboundAgentError(RuntimeError):
|
||||
pass
|
||||
|
||||
|
||||
def _post(url: str, value: dict[str,Any], timeout: float=10.0) -> dict[str,Any]:
|
||||
req=urllib.request.Request(
|
||||
url,
|
||||
data=json.dumps(value,sort_keys=True,separators=(",",":")).encode("utf-8"),
|
||||
headers={"Content-Type":"application/json","Accept":"application/json"},
|
||||
method="POST",
|
||||
)
|
||||
try:
|
||||
with urllib.request.urlopen(req,timeout=timeout) as resp:
|
||||
out=json.loads(resp.read(1024*1024).decode("utf-8"))
|
||||
except urllib.error.HTTPError as exc:
|
||||
detail=""
|
||||
try:
|
||||
body=json.loads(exc.read(1024*1024).decode("utf-8"))
|
||||
if isinstance(body,dict):
|
||||
detail=str(body.get("error") or "")
|
||||
except Exception:
|
||||
detail=""
|
||||
if detail=="node_revoked":
|
||||
raise OutboundAgentError("OUTBOUND_NODE_REVOKED") from exc
|
||||
raise OutboundAgentError("OUTBOUND_BROKER_REJECTED") from exc
|
||||
except Exception as exc:
|
||||
raise OutboundAgentError("OUTBOUND_BROKER_UNREACHABLE") from exc
|
||||
if not isinstance(out,dict):
|
||||
raise OutboundAgentError("OUTBOUND_BROKER_RESPONSE_INVALID")
|
||||
return out
|
||||
|
||||
|
||||
def _signed_envelope(ctx: SeedContext, kind: str, payload: dict[str,Any]) -> dict[str,Any]:
|
||||
identity=ctx.load_identity() or {}
|
||||
node_id=str(identity.get("node_id") or "")
|
||||
host_id=str(identity.get("node_host_id") or "")
|
||||
if not node_id or not host_id:
|
||||
raise OutboundAgentError("OUTBOUND_NODE_IDENTITY_MISSING")
|
||||
crypto=ensure_node_keypair(ctx.state_dir,comment="framebridge-"+node_id.replace("-","")[:8])
|
||||
base={
|
||||
"schema":1,
|
||||
"kind":str(kind),
|
||||
"node_id":node_id,
|
||||
"node_host_id":host_id,
|
||||
"timestamp":time.time(),
|
||||
"nonce":secrets.token_hex(16),
|
||||
"payload":payload,
|
||||
}
|
||||
sig=sign_pairing_payload(crypto["private_key_path"],canonical(base))
|
||||
return {
|
||||
**base,
|
||||
"node_proof":{
|
||||
"format":"sshsig",
|
||||
"namespace":"framebridge-pairing",
|
||||
"signature_b64":sig,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def broker_base(ctx: SeedContext) -> str:
|
||||
config=ctx.load_config() or {}
|
||||
pairing=config.get("pairing") or {}
|
||||
base=str(pairing.get("broker_base_url") or "").rstrip("/")
|
||||
if not base:
|
||||
raise OutboundAgentError("OUTBOUND_BROKER_NOT_CONFIGURED")
|
||||
return base
|
||||
|
||||
|
||||
def sync_local_revocation(ctx: SeedContext) -> dict[str,Any]:
|
||||
store=PairingStore(ctx.state_dir)
|
||||
current=store.load()
|
||||
if current.get("state")=="paired":
|
||||
view=store.revoke()
|
||||
elif current.get("state")=="revoked":
|
||||
view=current
|
||||
else:
|
||||
raise OutboundAgentError("OUTBOUND_LOCAL_PAIRING_STATE_INVALID")
|
||||
|
||||
config=ctx.load_config() or {}
|
||||
pairing=config.get("pairing") or {}
|
||||
pairing["state"]="revoked"
|
||||
pairing["transport"]="broker_outbound"
|
||||
if view.get("revoked_at"):
|
||||
pairing["revoked_at"]=view.get("revoked_at")
|
||||
config["pairing"]=pairing
|
||||
config["connection_mode"]="PAIRING"
|
||||
setup=config.get("setup")
|
||||
if isinstance(setup,dict):
|
||||
setup["pairing_state"]="revoked"
|
||||
atomic_write_json(ctx.config_path,config)
|
||||
|
||||
ctx.audit.append(
|
||||
"pair.revoked",
|
||||
"PASS",
|
||||
request={"source":"broker"},
|
||||
result={"state":"revoked"},
|
||||
principal="pairing-broker",
|
||||
)
|
||||
MemoryStore(ctx).append(
|
||||
"recovery",
|
||||
{"event":"pair.revoked","state":"revoked","source":"broker"},
|
||||
source="system",
|
||||
)
|
||||
return {"ok":True,"state":"revoked"}
|
||||
|
||||
|
||||
def execute_plan_operation(ctx: SeedContext, operation: str, payload: dict[str,Any]) -> dict[str,Any]:
|
||||
if operation not in {"plan.import","plan.status","plan.apply","plan.rollback"}:
|
||||
raise OutboundAgentError("OUTBOUND_PLAN_OPERATION_NOT_ALLOWED")
|
||||
if not isinstance(payload,dict):
|
||||
raise OutboundAgentError("OUTBOUND_PLAN_PAYLOAD_INVALID")
|
||||
engine=PlanEngine(ctx)
|
||||
|
||||
if operation=="plan.import":
|
||||
plan=payload.get("plan")
|
||||
if not isinstance(plan,dict):
|
||||
raise OutboundAgentError("OUTBOUND_PLAN_DOCUMENT_INVALID")
|
||||
return engine.import_document(plan)
|
||||
|
||||
plan_id=str(payload.get("plan_id") or "")
|
||||
if not plan_id:
|
||||
raise OutboundAgentError("OUTBOUND_PLAN_ID_MISSING")
|
||||
|
||||
if operation=="plan.status":
|
||||
return engine.show(plan_id)
|
||||
|
||||
expected=str(payload.get("plan_sha256") or "")
|
||||
shown=engine.show(plan_id)
|
||||
if shown.get("plan_sha256")!=expected:
|
||||
raise OutboundAgentError("OUTBOUND_PLAN_SHA_MISMATCH")
|
||||
|
||||
if operation=="plan.apply":
|
||||
return engine.apply(plan_id)
|
||||
return engine.rollback(plan_id)
|
||||
|
||||
|
||||
def run_once(ctx: SeedContext | None=None) -> dict[str,Any]:
|
||||
ctx=ctx or SeedContext()
|
||||
store=PairingStore(ctx.state_dir)
|
||||
local_pair=store.load()
|
||||
if local_pair.get("state")=="paired" and local_pair.get("broker_confirm_required") is True:
|
||||
try:
|
||||
confirm_pairing_remote(state_dir=ctx.state_dir,record=local_pair)
|
||||
except PairingClientError as exc:
|
||||
raise OutboundAgentError("OUTBOUND_PAIR_CONFIRM_PENDING") from exc
|
||||
local_pair=store.load()
|
||||
config=ctx.load_config() or {}
|
||||
pairing_cfg=config.get("pairing") or {}
|
||||
pairing_cfg["broker_confirm_required"]=False
|
||||
pairing_cfg["broker_confirmed_at"]=local_pair.get("broker_confirmed_at")
|
||||
config["pairing"]=pairing_cfg
|
||||
atomic_write_json(ctx.config_path,config)
|
||||
config=ctx.load_config() or {}
|
||||
pairing=config.get("pairing") or {}
|
||||
if pairing.get("state")!="paired":
|
||||
raise OutboundAgentError("OUTBOUND_NODE_NOT_PAIRED")
|
||||
if store.load().get("broker_confirm_required") is True:
|
||||
raise OutboundAgentError("OUTBOUND_PAIR_CONFIRM_PENDING")
|
||||
base=broker_base(ctx)
|
||||
poll=_post(
|
||||
base+"/fbpair/v1/node/poll",
|
||||
_signed_envelope(ctx,"node.poll",{}),
|
||||
)
|
||||
if poll.get("ok") is not True:
|
||||
raise OutboundAgentError("OUTBOUND_POLL_REJECTED")
|
||||
req=poll.get("request")
|
||||
if not isinstance(req,dict):
|
||||
return {"ok":True,"idle":True}
|
||||
rid=str(req.get("request_id") or "")
|
||||
operation=str(req.get("operation") or "")
|
||||
if operation:
|
||||
payload=req.get("payload") or {}
|
||||
try:
|
||||
value=execute_plan_operation(ctx,operation,payload)
|
||||
result={"ok":True,"operation":operation,"result":value}
|
||||
ctx.audit.append(
|
||||
"outbound."+operation,
|
||||
"PASS",
|
||||
request={"request_id":rid,"operation":operation},
|
||||
result={"ok":True},
|
||||
principal="paired-ai",
|
||||
)
|
||||
except Exception as exc:
|
||||
result={"ok":False,"operation":operation,"error":str(exc)}
|
||||
ctx.audit.append(
|
||||
"outbound."+operation,
|
||||
"FAIL",
|
||||
request={"request_id":rid,"operation":operation},
|
||||
result={"ok":False,"error":type(exc).__name__},
|
||||
principal="paired-ai",
|
||||
)
|
||||
cap=operation
|
||||
else:
|
||||
cap=str(req.get("capability") or "")
|
||||
args=req.get("args") or {}
|
||||
if cap not in READ_CAPS:
|
||||
result={"ok":False,"error":"CAPABILITY_NOT_ALLOWED","capability":cap}
|
||||
elif not isinstance(args,dict):
|
||||
result={"ok":False,"error":"CAPABILITY_ARGS_NOT_OBJECT","capability":cap}
|
||||
else:
|
||||
try:
|
||||
value=execute_capability(cap,args)
|
||||
result={"ok":True,"capability":cap,"result":value}
|
||||
ctx.audit.append(
|
||||
"outbound."+cap,
|
||||
"PASS",
|
||||
request={"request_id":rid,"capability":cap},
|
||||
result={"ok":True},
|
||||
principal="paired-ai",
|
||||
)
|
||||
except Exception as exc:
|
||||
result={"ok":False,"capability":cap,"error":str(exc)}
|
||||
ctx.audit.append(
|
||||
"outbound."+cap,
|
||||
"FAIL",
|
||||
request={"request_id":rid,"capability":cap},
|
||||
result={"ok":False,"error":type(exc).__name__},
|
||||
principal="paired-ai",
|
||||
)
|
||||
response=_post(
|
||||
base+"/fbpair/v1/node/result",
|
||||
_signed_envelope(
|
||||
ctx,
|
||||
"node.result",
|
||||
{"request_id":rid,"result":result},
|
||||
),
|
||||
)
|
||||
if response.get("ok") is not True:
|
||||
raise OutboundAgentError("OUTBOUND_RESULT_REJECTED")
|
||||
return {"ok":True,"idle":False,"request_id":rid,"capability":cap}
|
||||
|
||||
|
||||
def run_forever(interval: float=1.0) -> None:
|
||||
interval=max(0.5,min(float(interval),30.0))
|
||||
ctx=SeedContext()
|
||||
while True:
|
||||
try:
|
||||
run_once(ctx)
|
||||
except OutboundAgentError as exc:
|
||||
if str(exc)=="OUTBOUND_NODE_REVOKED":
|
||||
try:
|
||||
sync_local_revocation(ctx)
|
||||
except Exception:
|
||||
pass
|
||||
except Exception:
|
||||
pass
|
||||
time.sleep(interval)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
import argparse
|
||||
parser=argparse.ArgumentParser(prog="framebridge-outbound-agent")
|
||||
sub=parser.add_subparsers(dest="command",required=True)
|
||||
sub.add_parser("once")
|
||||
run=sub.add_parser("run")
|
||||
run.add_argument("--interval",type=float,default=1.0)
|
||||
args=parser.parse_args()
|
||||
if args.command=="once":
|
||||
print(json.dumps(run_once(),sort_keys=True))
|
||||
else:
|
||||
run_forever(args.interval)
|
||||
|
||||
|
||||
if __name__=="__main__":
|
||||
main()
|
||||
+484
@@ -0,0 +1,484 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import ipaddress
|
||||
import json
|
||||
import fcntl
|
||||
import os
|
||||
import secrets
|
||||
import tempfile
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from contextlib import contextmanager
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
from urllib.parse import urlparse
|
||||
|
||||
FBPAIR_PROTOCOL = "FBPAIR/1"
|
||||
PAIR_STATES = {
|
||||
"not_paired",
|
||||
"pair_requested",
|
||||
"paired",
|
||||
"denied",
|
||||
"expired",
|
||||
"revoked",
|
||||
}
|
||||
PROVIDERS = {"chatgpt", "ollama", "lmstudio", "other"}
|
||||
|
||||
|
||||
class PairingError(RuntimeError):
|
||||
pass
|
||||
|
||||
|
||||
def utc_now() -> str:
|
||||
return datetime.now(timezone.utc).isoformat()
|
||||
|
||||
|
||||
def canonical(value: Any) -> bytes:
|
||||
return json.dumps(
|
||||
value, sort_keys=True, separators=(",", ":"), ensure_ascii=True
|
||||
).encode("ascii")
|
||||
|
||||
|
||||
def sha256_hex(value: Any) -> str:
|
||||
return hashlib.sha256(canonical(value)).hexdigest()
|
||||
|
||||
|
||||
def _is_local_ip(host: str) -> bool:
|
||||
try:
|
||||
addr=ipaddress.ip_address(str(host))
|
||||
except ValueError:
|
||||
return False
|
||||
if addr.is_loopback or addr.is_link_local:
|
||||
return True
|
||||
if addr.version==4:
|
||||
return any(addr in net for net in (
|
||||
ipaddress.ip_network("10.0.0.0/8"),
|
||||
ipaddress.ip_network("172.16.0.0/12"),
|
||||
ipaddress.ip_network("192.168.0.0/16"),
|
||||
))
|
||||
return addr in ipaddress.ip_network("fc00::/7")
|
||||
|
||||
|
||||
def _https_or_loopback_http(uri: str, allow_insecure_dev: bool=False) -> bool:
|
||||
try:
|
||||
p = urlparse(str(uri))
|
||||
except Exception:
|
||||
return False
|
||||
if p.scheme == "https" and bool(p.netloc):
|
||||
return True
|
||||
if p.scheme == "http" and bool(p.netloc):
|
||||
host=str(p.hostname or "")
|
||||
if _is_local_ip(host):
|
||||
return True
|
||||
if allow_insecure_dev:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def create_pair_request(
|
||||
*,
|
||||
node_id: str,
|
||||
node_label: str,
|
||||
node_host_id: str,
|
||||
public_jwk: dict[str, Any],
|
||||
provider: str,
|
||||
capability_manifest_sha256: str,
|
||||
) -> dict[str, Any]:
|
||||
provider = str(provider).strip().lower()
|
||||
if provider not in PROVIDERS:
|
||||
raise PairingError("PAIR_PROVIDER_UNSUPPORTED")
|
||||
node_label=str(node_label).strip()
|
||||
if not node_id or not node_host_id or not node_label or len(node_label)>253:
|
||||
raise PairingError("PAIR_NODE_IDENTITY_REQUIRED")
|
||||
manifest_sha = str(capability_manifest_sha256).strip().lower()
|
||||
if len(manifest_sha) != 64 or any(c not in "0123456789abcdef" for c in manifest_sha):
|
||||
raise PairingError("PAIR_CAPABILITY_MANIFEST_SHA_INVALID")
|
||||
nonce = secrets.token_urlsafe(32)
|
||||
request = {
|
||||
"protocol": FBPAIR_PROTOCOL,
|
||||
"schema": 1,
|
||||
"node_id": str(node_id),
|
||||
"node_label": node_label,
|
||||
"node_host_id": str(node_host_id),
|
||||
"public_jwk": public_jwk,
|
||||
"provider": provider,
|
||||
"capability_manifest_sha256": manifest_sha,
|
||||
"node_nonce": nonce,
|
||||
"requested_at": utc_now(),
|
||||
}
|
||||
request["request_sha256"] = sha256_hex(request)
|
||||
return request
|
||||
|
||||
|
||||
def attach_node_proof(
|
||||
request: dict[str, Any],
|
||||
*,
|
||||
public_key_ssh: str,
|
||||
signature_b64: str,
|
||||
) -> dict[str, Any]:
|
||||
if request.get("protocol") != FBPAIR_PROTOCOL:
|
||||
raise PairingError("PAIR_REQUEST_PROTOCOL_INVALID")
|
||||
if not str(public_key_ssh).startswith("ssh-ed25519 "):
|
||||
raise PairingError("PAIR_PUBLIC_KEY_INVALID")
|
||||
if not signature_b64:
|
||||
raise PairingError("PAIR_NODE_PROOF_REQUIRED")
|
||||
value=dict(request)
|
||||
value["public_key_ssh"]=str(public_key_ssh)
|
||||
value["node_proof"]={
|
||||
"format":"sshsig",
|
||||
"namespace":"framebridge-pairing",
|
||||
"signature_b64":str(signature_b64),
|
||||
}
|
||||
return value
|
||||
|
||||
|
||||
def pairing_signing_payload(request: dict[str, Any]) -> bytes:
|
||||
value=dict(request)
|
||||
value.pop("node_proof",None)
|
||||
value.pop("public_key_ssh",None)
|
||||
return canonical(value)
|
||||
|
||||
|
||||
def validate_device_authorization(
|
||||
request: dict[str, Any],
|
||||
response: dict[str, Any],
|
||||
*,
|
||||
allow_insecure_dev: bool=False,
|
||||
) -> dict[str, Any]:
|
||||
if not isinstance(response, dict):
|
||||
raise PairingError("PAIR_DEVICE_RESPONSE_INVALID")
|
||||
required = {
|
||||
"pair_id",
|
||||
"device_code",
|
||||
"user_code",
|
||||
"verification_uri",
|
||||
"expires_in",
|
||||
}
|
||||
if not required.issubset(response):
|
||||
raise PairingError("PAIR_DEVICE_RESPONSE_INCOMPLETE")
|
||||
pair_id = str(response["pair_id"])
|
||||
device_code = str(response["device_code"])
|
||||
user_code = str(response["user_code"])
|
||||
verification_uri = str(response["verification_uri"])
|
||||
complete = str(response.get("verification_uri_complete") or "")
|
||||
try:
|
||||
expires_in = int(response["expires_in"])
|
||||
interval = int(response.get("interval", 5))
|
||||
except Exception as exc:
|
||||
raise PairingError("PAIR_DEVICE_RESPONSE_INVALID") from exc
|
||||
if not pair_id or len(device_code) < 24 or not user_code:
|
||||
raise PairingError("PAIR_DEVICE_RESPONSE_INVALID")
|
||||
if not _https_or_loopback_http(verification_uri, allow_insecure_dev):
|
||||
raise PairingError("PAIR_VERIFICATION_URI_INVALID")
|
||||
if complete and not _https_or_loopback_http(complete, allow_insecure_dev):
|
||||
raise PairingError("PAIR_VERIFICATION_URI_COMPLETE_INVALID")
|
||||
if not (30 <= expires_in <= 1800):
|
||||
raise PairingError("PAIR_EXPIRES_INVALID")
|
||||
if not (1 <= interval <= 60):
|
||||
raise PairingError("PAIR_INTERVAL_INVALID")
|
||||
try:
|
||||
requested_dt=datetime.fromisoformat(
|
||||
str(request["requested_at"]).replace("Z","+00:00")
|
||||
).astimezone(timezone.utc)
|
||||
except Exception as exc:
|
||||
raise PairingError("PAIR_REQUEST_TIME_INVALID") from exc
|
||||
expires_at=(requested_dt+timedelta(seconds=expires_in)).isoformat()
|
||||
|
||||
return {
|
||||
"protocol": FBPAIR_PROTOCOL,
|
||||
"schema": 1,
|
||||
"state": "pair_requested",
|
||||
"pair_id": pair_id,
|
||||
"request_sha256": request["request_sha256"],
|
||||
"node_id": request["node_id"],
|
||||
"node_label": request["node_label"],
|
||||
"node_host_id": request["node_host_id"],
|
||||
"provider": request["provider"],
|
||||
"device_code": device_code,
|
||||
"device_code_sha256": hashlib.sha256(device_code.encode("utf-8")).hexdigest(),
|
||||
"user_code": user_code,
|
||||
"verification_uri": verification_uri,
|
||||
"verification_uri_complete": complete or None,
|
||||
"expires_in": expires_in,
|
||||
"expires_at": expires_at,
|
||||
"interval": interval,
|
||||
"requested_at": request["requested_at"],
|
||||
"last_poll_at": None,
|
||||
}
|
||||
|
||||
|
||||
def public_pairing_view(record: dict[str, Any]) -> dict[str, Any]:
|
||||
value = dict(record)
|
||||
value.pop("device_code", None)
|
||||
value.pop("account_subject", None)
|
||||
value.pop("provider_access_token", None)
|
||||
value.pop("provider_refresh_token", None)
|
||||
return value
|
||||
|
||||
|
||||
def pairing_qr_payload(record: dict[str, Any]) -> str:
|
||||
if record.get("state") != "pair_requested":
|
||||
raise PairingError("PAIR_NOT_PENDING")
|
||||
return str(
|
||||
record.get("verification_uri_complete")
|
||||
or record.get("verification_uri")
|
||||
or ""
|
||||
)
|
||||
|
||||
|
||||
def validate_pair_receipt(
|
||||
pending: dict[str, Any],
|
||||
receipt: dict[str, Any],
|
||||
) -> dict[str, Any]:
|
||||
if pending.get("state") != "pair_requested":
|
||||
raise PairingError("PAIR_NOT_PENDING")
|
||||
if not isinstance(receipt, dict):
|
||||
raise PairingError("PAIR_RECEIPT_INVALID")
|
||||
checks = {
|
||||
"protocol": receipt.get("protocol") == FBPAIR_PROTOCOL,
|
||||
"schema": receipt.get("schema") == 1,
|
||||
"state": receipt.get("state") == "paired",
|
||||
"pair_id": receipt.get("pair_id") == pending.get("pair_id"),
|
||||
"request_sha256": receipt.get("request_sha256") == pending.get("request_sha256"),
|
||||
"node_id": receipt.get("node_id") == pending.get("node_id"),
|
||||
"node_label": receipt.get("node_label") == pending.get("node_label"),
|
||||
"node_host_id": receipt.get("node_host_id") == pending.get("node_host_id"),
|
||||
"provider": receipt.get("provider") == pending.get("provider"),
|
||||
}
|
||||
if not all(checks.values()):
|
||||
raise PairingError("PAIR_RECEIPT_BINDING_MISMATCH")
|
||||
broker_issuer = str(receipt.get("broker_issuer") or "")
|
||||
account_binding_id = str(receipt.get("account_binding_id") or "")
|
||||
grant_id = str(receipt.get("grant_id") or "")
|
||||
if (
|
||||
not broker_issuer
|
||||
or not account_binding_id
|
||||
or len(account_binding_id)>160
|
||||
or not grant_id
|
||||
or len(grant_id)>160
|
||||
):
|
||||
raise PairingError("PAIR_ACCOUNT_BINDING_MISSING")
|
||||
return {
|
||||
"protocol": FBPAIR_PROTOCOL,
|
||||
"schema": 1,
|
||||
"state": "paired",
|
||||
"pair_id": pending["pair_id"],
|
||||
"request_sha256": pending["request_sha256"],
|
||||
"node_id": pending["node_id"],
|
||||
"node_label": pending["node_label"],
|
||||
"node_host_id": pending["node_host_id"],
|
||||
"provider": pending["provider"],
|
||||
"broker_issuer": broker_issuer,
|
||||
"account_binding_id": account_binding_id,
|
||||
"grant_id": grant_id,
|
||||
"paired_at": str(receipt.get("paired_at") or utc_now()),
|
||||
"scopes": sorted({str(x) for x in (receipt.get("scopes") or [])}),
|
||||
"revocation_id": str(receipt.get("revocation_id") or ""),
|
||||
"node_grant_binding": str(receipt.get("node_grant_binding") or "proof-of-possession-required"),
|
||||
}
|
||||
|
||||
|
||||
class PairingStore:
|
||||
def __init__(self, state_dir: Path):
|
||||
self.path = Path(state_dir) / "pairing.json"
|
||||
self.lock_path = Path(state_dir) / ".pairing.lock"
|
||||
|
||||
@contextmanager
|
||||
def _exclusive(self):
|
||||
self.lock_path.parent.mkdir(parents=True, exist_ok=True, mode=0o700)
|
||||
flags=os.O_RDWR|os.O_CREAT|getattr(os,"O_NOFOLLOW",0)
|
||||
fd=os.open(self.lock_path,flags,0o600)
|
||||
try:
|
||||
fcntl.flock(fd,fcntl.LOCK_EX)
|
||||
yield
|
||||
finally:
|
||||
try:
|
||||
fcntl.flock(fd,fcntl.LOCK_UN)
|
||||
finally:
|
||||
os.close(fd)
|
||||
|
||||
def _load_unlocked(self) -> dict[str, Any]:
|
||||
if not self.path.is_file():
|
||||
return {"protocol": FBPAIR_PROTOCOL, "schema": 1, "state": "not_paired"}
|
||||
value = json.loads(self.path.read_text(encoding="utf-8"))
|
||||
if not isinstance(value, dict) or value.get("state") not in PAIR_STATES:
|
||||
raise PairingError("PAIR_STATE_INVALID")
|
||||
return value
|
||||
|
||||
def load(self) -> dict[str, Any]:
|
||||
# Atomic replacement makes a lock-free reader see old-or-new complete JSON.
|
||||
return self._load_unlocked()
|
||||
|
||||
def _save_unlocked(self, value: dict[str, Any]) -> None:
|
||||
if value.get("state") not in PAIR_STATES:
|
||||
raise PairingError("PAIR_STATE_INVALID")
|
||||
self.path.parent.mkdir(parents=True, exist_ok=True, mode=0o700)
|
||||
fd, tmp = tempfile.mkstemp(prefix=".pairing.", dir=str(self.path.parent))
|
||||
try:
|
||||
with os.fdopen(fd, "w", encoding="utf-8") as handle:
|
||||
json.dump(value, handle, indent=2, sort_keys=True)
|
||||
handle.write("\n")
|
||||
handle.flush()
|
||||
os.fsync(handle.fileno())
|
||||
os.chmod(tmp, 0o600)
|
||||
os.replace(tmp, self.path)
|
||||
finally:
|
||||
if os.path.exists(tmp):
|
||||
os.unlink(tmp)
|
||||
|
||||
def save(self, value: dict[str, Any]) -> None:
|
||||
with self._exclusive():
|
||||
self._save_unlocked(value)
|
||||
|
||||
def begin(self, pending: dict[str, Any]) -> dict[str, Any]:
|
||||
with self._exclusive():
|
||||
current = self._load_unlocked()
|
||||
if current.get("state") == "paired":
|
||||
raise PairingError("PAIR_ALREADY_PAIRED")
|
||||
if current.get("state") == "pair_requested":
|
||||
raise PairingError("PAIR_REQUEST_ALREADY_PENDING")
|
||||
self._save_unlocked(pending)
|
||||
return public_pairing_view(pending)
|
||||
|
||||
def update_pending_fields(self, pair_id: str, fields: dict[str, Any]) -> dict[str, Any]:
|
||||
pair_id=str(pair_id or "")
|
||||
if not pair_id:
|
||||
raise PairingError("PAIR_ID_REQUIRED")
|
||||
with self._exclusive():
|
||||
current=self._load_unlocked()
|
||||
if current.get("state")!="pair_requested" or str(current.get("pair_id") or "")!=pair_id:
|
||||
return current
|
||||
current=dict(current)
|
||||
current.update(dict(fields or {}))
|
||||
self._save_unlocked(current)
|
||||
return current
|
||||
|
||||
def complete(self, receipt: dict[str, Any]) -> dict[str, Any]:
|
||||
with self._exclusive():
|
||||
current = self._load_unlocked()
|
||||
paired = validate_pair_receipt(current, receipt)
|
||||
paired["broker_base_url"]=str(current.get("broker_base_url") or "")
|
||||
paired["broker_confirm_required"]=True
|
||||
self._save_unlocked(paired)
|
||||
return public_pairing_view(paired)
|
||||
|
||||
def invalidate_unconfirmed_pair(self, pair_id: str, broker_state: str) -> dict[str, Any]:
|
||||
pair_id=str(pair_id or "")
|
||||
broker_state=str(broker_state or "unknown").lower()
|
||||
with self._exclusive():
|
||||
current=self._load_unlocked()
|
||||
if current.get("state")!="paired" or current.get("broker_confirm_required") is not True or str(current.get("pair_id") or "")!=pair_id:
|
||||
return current
|
||||
current=dict(current)
|
||||
current["state"]="expired"
|
||||
current["expired_at"]=utc_now()
|
||||
current["expiration_reason"]="broker_confirm_"+broker_state
|
||||
current["requires_new_local_action"]=True
|
||||
current["broker_confirm_required"]=False
|
||||
current["broker_confirm_state"]=broker_state
|
||||
self._save_unlocked(current)
|
||||
return current
|
||||
|
||||
def mark_broker_confirmed(self, pair_id: str, result: dict[str, Any]) -> dict[str, Any]:
|
||||
pair_id=str(pair_id or "")
|
||||
with self._exclusive():
|
||||
current=self._load_unlocked()
|
||||
if current.get("state")!="paired" or str(current.get("pair_id") or "")!=pair_id:
|
||||
return current
|
||||
current=dict(current)
|
||||
current["broker_confirm_required"]=False
|
||||
current["broker_confirmed_at"]=utc_now()
|
||||
current["broker_confirm_state"]=str((result or {}).get("state") or "paired")
|
||||
self._save_unlocked(current)
|
||||
return current
|
||||
|
||||
def transition_pending(self, pair_id: str, new_state: str, fields: dict[str, Any] | None=None) -> dict[str, Any]:
|
||||
pair_id=str(pair_id or "")
|
||||
if new_state not in {"denied","expired"}:
|
||||
raise PairingError("PAIR_TRANSITION_STATE_INVALID")
|
||||
with self._exclusive():
|
||||
current=self._load_unlocked()
|
||||
if current.get("state")!="pair_requested" or str(current.get("pair_id") or "")!=pair_id:
|
||||
return current
|
||||
current=dict(current)
|
||||
current.pop("device_code",None)
|
||||
current["state"]=new_state
|
||||
current.update(dict(fields or {}))
|
||||
self._save_unlocked(current)
|
||||
return current
|
||||
|
||||
def cancel_pending(self) -> dict[str, Any]:
|
||||
with self._exclusive():
|
||||
current = self._load_unlocked()
|
||||
if current.get("state") != "pair_requested":
|
||||
raise PairingError("PAIR_NOT_PENDING")
|
||||
current=dict(current)
|
||||
current.pop("device_code", None)
|
||||
current["state"] = "revoked"
|
||||
current["revoked_at"] = utc_now()
|
||||
current["revocation_reason"] = "local_operator_cancelled"
|
||||
if current.get("pair_id") and current.get("broker_base_url"):
|
||||
current["broker_cancel_required"]=True
|
||||
current["broker_cancel_reason"]="local_operator_cancelled"
|
||||
self._save_unlocked(current)
|
||||
return public_pairing_view(current)
|
||||
|
||||
def invalidate_pending(self, reason: str) -> dict[str, Any]:
|
||||
reason=str(reason or "").strip().lower()
|
||||
if reason not in {"network_change","service_stop","reconcile","window_expired","window_closed"}:
|
||||
raise PairingError("PAIR_INVALIDATION_REASON_INVALID")
|
||||
with self._exclusive():
|
||||
current=self._load_unlocked()
|
||||
if current.get("state") != "pair_requested":
|
||||
raise PairingError("PAIR_NOT_PENDING")
|
||||
current=dict(current)
|
||||
current.pop("device_code",None)
|
||||
current["state"]="expired"
|
||||
current["expired_at"]=utc_now()
|
||||
current["expiration_reason"]=reason
|
||||
current["requires_new_local_action"]=True
|
||||
if current.get("pair_id") and current.get("broker_base_url"):
|
||||
current["broker_cancel_required"]=True
|
||||
current["broker_cancel_reason"]=reason
|
||||
self._save_unlocked(current)
|
||||
return public_pairing_view(current)
|
||||
|
||||
def mark_broker_cancelled(self, pair_id: str, result: dict[str, Any]) -> dict[str, Any]:
|
||||
pair_id=str(pair_id or "")
|
||||
with self._exclusive():
|
||||
current=self._load_unlocked()
|
||||
if str(current.get("pair_id") or "")!=pair_id:
|
||||
return current
|
||||
current=dict(current)
|
||||
current["broker_cancel_required"]=False
|
||||
current["broker_cancel_state"]=str((result or {}).get("state") or "cancelled")
|
||||
current["broker_cancelled_at"]=utc_now()
|
||||
current.pop("broker_cancel_last_error",None)
|
||||
self._save_unlocked(current)
|
||||
return current
|
||||
|
||||
def note_broker_cancel_error(self, pair_id: str, error: str) -> dict[str, Any]:
|
||||
pair_id=str(pair_id or "")
|
||||
with self._exclusive():
|
||||
current=self._load_unlocked()
|
||||
if str(current.get("pair_id") or "")!=pair_id:
|
||||
return current
|
||||
current=dict(current)
|
||||
current["broker_cancel_required"]=True
|
||||
current["broker_cancel_last_error"]=str(error or "PAIR_BROKER_CANCEL_FAILED")[:128]
|
||||
current["broker_cancel_attempts"]=int(current.get("broker_cancel_attempts") or 0)+1
|
||||
current["broker_cancel_last_attempt_at"]=utc_now()
|
||||
self._save_unlocked(current)
|
||||
return current
|
||||
|
||||
def revoke(self) -> dict[str, Any]:
|
||||
with self._exclusive():
|
||||
current = self._load_unlocked()
|
||||
if current.get("state") != "paired":
|
||||
raise PairingError("PAIR_NOT_PAIRED")
|
||||
current=dict(current)
|
||||
current["state"] = "revoked"
|
||||
current["revoked_at"] = utc_now()
|
||||
self._save_unlocked(current)
|
||||
return public_pairing_view(current)
|
||||
|
||||
+380
@@ -0,0 +1,380 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import ipaddress
|
||||
import json
|
||||
import secrets
|
||||
import ssl
|
||||
import time
|
||||
from datetime import datetime, timezone
|
||||
import urllib.error
|
||||
import urllib.parse
|
||||
import urllib.request
|
||||
from pathlib import Path
|
||||
from typing import Any, Callable
|
||||
|
||||
from .node_identity import ensure_node_keypair, sign_pairing_payload
|
||||
from .pairing import (
|
||||
PairingError,
|
||||
PairingStore,
|
||||
attach_node_proof,
|
||||
create_pair_request,
|
||||
pairing_signing_payload,
|
||||
public_pairing_view,
|
||||
validate_device_authorization,
|
||||
)
|
||||
|
||||
|
||||
class PairingClientError(RuntimeError):
|
||||
pass
|
||||
|
||||
|
||||
def _is_local_ip(host: str) -> bool:
|
||||
try:
|
||||
addr=ipaddress.ip_address(str(host))
|
||||
except ValueError:
|
||||
return False
|
||||
if addr.is_loopback or addr.is_link_local:
|
||||
return True
|
||||
if addr.version==4:
|
||||
return any(addr in net for net in (
|
||||
ipaddress.ip_network("10.0.0.0/8"),
|
||||
ipaddress.ip_network("172.16.0.0/12"),
|
||||
ipaddress.ip_network("192.168.0.0/16"),
|
||||
))
|
||||
return addr in ipaddress.ip_network("fc00::/7")
|
||||
|
||||
|
||||
def _broker_base_ok(url: str, allow_insecure_dev: bool) -> bool:
|
||||
try:
|
||||
p=urllib.parse.urlparse(str(url))
|
||||
except Exception:
|
||||
return False
|
||||
if p.scheme=="https" and p.netloc:
|
||||
return True
|
||||
if p.scheme=="http" and p.netloc:
|
||||
host=str(p.hostname or "")
|
||||
if _is_local_ip(host):
|
||||
return True
|
||||
if allow_insecure_dev:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def _endpoint(base: str, path: str) -> str:
|
||||
return str(base).rstrip("/") + "/" + path.lstrip("/")
|
||||
|
||||
|
||||
def post_json(
|
||||
url: str,
|
||||
value: dict[str, Any],
|
||||
*,
|
||||
timeout: float=10.0,
|
||||
opener: Callable[..., Any] | None=None,
|
||||
) -> dict[str, Any]:
|
||||
data=json.dumps(
|
||||
value, sort_keys=True, separators=(",",":"), ensure_ascii=True
|
||||
).encode("utf-8")
|
||||
req=urllib.request.Request(
|
||||
url,
|
||||
data=data,
|
||||
method="POST",
|
||||
headers={
|
||||
"Content-Type":"application/json",
|
||||
"Accept":"application/json",
|
||||
"User-Agent":"FrameBridge-Linux-Core/0.5",
|
||||
},
|
||||
)
|
||||
try:
|
||||
if opener is None:
|
||||
with urllib.request.urlopen(req,timeout=timeout) as resp:
|
||||
raw=resp.read(1024*1024)
|
||||
else:
|
||||
raw=opener(req,timeout=timeout)
|
||||
if hasattr(raw,"read"):
|
||||
raw=raw.read(1024*1024)
|
||||
result=json.loads(bytes(raw).decode("utf-8"))
|
||||
except urllib.error.HTTPError as exc:
|
||||
try:
|
||||
raw=exc.read(1024*1024)
|
||||
result=json.loads(bytes(raw).decode("utf-8"))
|
||||
except Exception as parse_exc:
|
||||
raise PairingClientError("PAIR_BROKER_RESPONSE_INVALID") from parse_exc
|
||||
if not isinstance(result,dict):
|
||||
raise PairingClientError("PAIR_BROKER_RESPONSE_INVALID")
|
||||
return result
|
||||
except (urllib.error.URLError,TimeoutError,OSError) as exc:
|
||||
raise PairingClientError("PAIR_BROKER_UNREACHABLE") from exc
|
||||
except Exception as exc:
|
||||
raise PairingClientError("PAIR_BROKER_RESPONSE_INVALID") from exc
|
||||
if not isinstance(result,dict):
|
||||
raise PairingClientError("PAIR_BROKER_RESPONSE_INVALID")
|
||||
return result
|
||||
|
||||
|
||||
def start_pairing(
|
||||
*,
|
||||
state_dir: Path,
|
||||
node_id: str,
|
||||
node_label: str,
|
||||
node_host_id: str,
|
||||
public_jwk: dict[str, Any],
|
||||
provider: str,
|
||||
capability_manifest_sha256: str,
|
||||
broker_base_url: str,
|
||||
allow_insecure_dev: bool=False,
|
||||
poster: Callable[..., dict[str,Any]]=post_json,
|
||||
) -> dict[str,Any]:
|
||||
if not _broker_base_ok(broker_base_url,allow_insecure_dev):
|
||||
raise PairingClientError("PAIR_BROKER_URL_INVALID")
|
||||
crypto=ensure_node_keypair(Path(state_dir),comment="framebridge-"+str(node_id)[:8])
|
||||
request=create_pair_request(
|
||||
node_id=node_id,
|
||||
node_label=node_label,
|
||||
node_host_id=node_host_id,
|
||||
public_jwk=public_jwk,
|
||||
provider=provider,
|
||||
capability_manifest_sha256=capability_manifest_sha256,
|
||||
)
|
||||
signature=sign_pairing_payload(
|
||||
crypto["private_key_path"],
|
||||
pairing_signing_payload(request),
|
||||
)
|
||||
signed=attach_node_proof(
|
||||
request,
|
||||
public_key_ssh=crypto["public_key_ssh"],
|
||||
signature_b64=signature,
|
||||
)
|
||||
response=poster(
|
||||
_endpoint(broker_base_url,"fbpair/v1/device_authorization"),
|
||||
signed,
|
||||
)
|
||||
pending=validate_device_authorization(
|
||||
request,response,allow_insecure_dev=allow_insecure_dev
|
||||
)
|
||||
pending["broker_base_url"]=str(broker_base_url).rstrip("/")
|
||||
PairingStore(Path(state_dir)).begin(pending)
|
||||
return public_pairing_view(pending)
|
||||
|
||||
|
||||
def confirm_pairing_remote(
|
||||
*,
|
||||
state_dir: Path,
|
||||
allow_insecure_dev: bool=False,
|
||||
poster: Callable[..., dict[str,Any]]=post_json,
|
||||
record: dict[str,Any] | None=None,
|
||||
) -> dict[str,Any]:
|
||||
store=PairingStore(Path(state_dir))
|
||||
current=dict(record) if isinstance(record,dict) else store.load()
|
||||
if current.get("state")!="paired":
|
||||
raise PairingClientError("PAIR_CONFIRM_STATE_INVALID")
|
||||
pair_id=str(current.get("pair_id") or "")
|
||||
node_id=str(current.get("node_id") or "")
|
||||
node_host_id=str(current.get("node_host_id") or "")
|
||||
request_sha256=str(current.get("request_sha256") or "")
|
||||
broker=str(current.get("broker_base_url") or "")
|
||||
if not pair_id or not node_id or not node_host_id or len(request_sha256)!=64:
|
||||
raise PairingClientError("PAIR_CONFIRM_STATE_INCOMPLETE")
|
||||
if not _broker_base_ok(broker,allow_insecure_dev):
|
||||
raise PairingClientError("PAIR_BROKER_URL_INVALID")
|
||||
crypto=ensure_node_keypair(Path(state_dir),comment="framebridge-"+node_id.replace("-","")[:8])
|
||||
base={
|
||||
"protocol":"FBPAIR/1",
|
||||
"schema":1,
|
||||
"operation":"pair.confirm",
|
||||
"pair_id":pair_id,
|
||||
"request_sha256":request_sha256,
|
||||
"node_id":node_id,
|
||||
"node_host_id":node_host_id,
|
||||
"confirm_nonce":secrets.token_urlsafe(24),
|
||||
"confirmed_at":datetime.now(timezone.utc).isoformat(),
|
||||
}
|
||||
signature=sign_pairing_payload(crypto["private_key_path"],pairing_signing_payload(base))
|
||||
signed=attach_node_proof(
|
||||
base,public_key_ssh=crypto["public_key_ssh"],signature_b64=signature
|
||||
)
|
||||
response=poster(_endpoint(broker,"fbpair/v1/confirm"),signed)
|
||||
if isinstance(response,dict) and response.get("ok") is not True and str(response.get("state") or "") in {"cancelled","denied","expired"}:
|
||||
store.invalidate_unconfirmed_pair(pair_id,str(response.get("state")))
|
||||
return response
|
||||
if not isinstance(response,dict) or response.get("ok") is not True or response.get("state")!="paired":
|
||||
error=str((response or {}).get("error") or "pair_confirm_failed").upper() if isinstance(response,dict) else "PAIR_CONFIRM_RESPONSE_INVALID"
|
||||
raise PairingClientError("PAIR_BROKER_"+error)
|
||||
store.mark_broker_confirmed(pair_id,response)
|
||||
return response
|
||||
|
||||
|
||||
def cancel_pairing_remote(
|
||||
*,
|
||||
state_dir: Path,
|
||||
reason: str,
|
||||
allow_insecure_dev: bool=False,
|
||||
poster: Callable[..., dict[str,Any]]=post_json,
|
||||
record: dict[str,Any] | None=None,
|
||||
) -> dict[str,Any]:
|
||||
reason=str(reason or "").strip().lower()
|
||||
if reason not in {
|
||||
"network_change","service_stop","reconcile","window_expired",
|
||||
"window_closed","local_operator_cancelled",
|
||||
}:
|
||||
raise PairingClientError("PAIR_CANCEL_REASON_INVALID")
|
||||
store=PairingStore(Path(state_dir))
|
||||
current=dict(record) if isinstance(record,dict) else store.load()
|
||||
pair_id=str(current.get("pair_id") or "")
|
||||
node_id=str(current.get("node_id") or "")
|
||||
node_host_id=str(current.get("node_host_id") or "")
|
||||
broker=str(current.get("broker_base_url") or "")
|
||||
if not pair_id or not node_id or not node_host_id:
|
||||
raise PairingClientError("PAIR_CANCEL_STATE_INCOMPLETE")
|
||||
if not _broker_base_ok(broker,allow_insecure_dev):
|
||||
raise PairingClientError("PAIR_BROKER_URL_INVALID")
|
||||
crypto=ensure_node_keypair(Path(state_dir),comment="framebridge-"+node_id.replace("-","")[:8])
|
||||
base={
|
||||
"protocol":"FBPAIR/1",
|
||||
"schema":1,
|
||||
"operation":"pair.cancel",
|
||||
"pair_id":pair_id,
|
||||
"node_id":node_id,
|
||||
"node_host_id":node_host_id,
|
||||
"reason":reason,
|
||||
"cancel_nonce":secrets.token_urlsafe(24),
|
||||
"cancelled_at":datetime.now(timezone.utc).isoformat(),
|
||||
}
|
||||
signature=sign_pairing_payload(crypto["private_key_path"],pairing_signing_payload(base))
|
||||
signed=attach_node_proof(
|
||||
base,
|
||||
public_key_ssh=crypto["public_key_ssh"],
|
||||
signature_b64=signature,
|
||||
)
|
||||
response=poster(_endpoint(broker,"fbpair/v1/cancel"),signed)
|
||||
if not isinstance(response,dict):
|
||||
raise PairingClientError("PAIR_CANCEL_RESPONSE_INVALID")
|
||||
if response.get("ok") is True and str(response.get("state") or "") in {
|
||||
"cancelled","expired","denied",
|
||||
}:
|
||||
return response
|
||||
error=str(response.get("error") or "pair_cancel_failed").upper()
|
||||
raise PairingClientError("PAIR_BROKER_"+error)
|
||||
|
||||
|
||||
def poll_pairing(
|
||||
*,
|
||||
state_dir: Path,
|
||||
allow_insecure_dev: bool=False,
|
||||
poster: Callable[..., dict[str,Any]]=post_json,
|
||||
) -> dict[str,Any]:
|
||||
store=PairingStore(Path(state_dir))
|
||||
pending=store.load()
|
||||
if pending.get("state")=="paired" and pending.get("broker_confirm_required") is True:
|
||||
try:
|
||||
confirm_pairing_remote(
|
||||
state_dir=Path(state_dir),allow_insecure_dev=allow_insecure_dev,
|
||||
poster=poster,record=pending,
|
||||
)
|
||||
except PairingClientError:
|
||||
result=public_pairing_view(store.load())
|
||||
result["broker_status"]="confirm_pending"
|
||||
return result
|
||||
return public_pairing_view(store.load())
|
||||
if pending.get("state")!="pair_requested":
|
||||
raise PairingClientError("PAIR_NOT_PENDING")
|
||||
pair_id=str(pending.get("pair_id") or "")
|
||||
broker=str(pending.get("broker_base_url") or "")
|
||||
if not _broker_base_ok(broker,allow_insecure_dev):
|
||||
raise PairingClientError("PAIR_BROKER_URL_INVALID")
|
||||
|
||||
now=datetime.now(timezone.utc)
|
||||
try:
|
||||
expires=datetime.fromisoformat(
|
||||
str(pending.get("expires_at") or "").replace("Z","+00:00")
|
||||
).astimezone(timezone.utc)
|
||||
except Exception as exc:
|
||||
raise PairingClientError("PAIR_EXPIRY_STATE_INVALID") from exc
|
||||
if now>=expires:
|
||||
current=store.transition_pending(
|
||||
pair_id,"expired",
|
||||
{"expired_at":now.isoformat(),"expiration_reason":"pairing_timeout"},
|
||||
)
|
||||
return public_pairing_view(current)
|
||||
|
||||
last=pending.get("last_poll_at")
|
||||
interval=max(1,int(pending.get("interval") or 5))
|
||||
if last:
|
||||
try:
|
||||
last_dt=datetime.fromisoformat(
|
||||
str(last).replace("Z","+00:00")
|
||||
).astimezone(timezone.utc)
|
||||
except Exception as exc:
|
||||
raise PairingClientError("PAIR_POLL_STATE_INVALID") from exc
|
||||
remaining=interval-(now-last_dt).total_seconds()
|
||||
if remaining>0:
|
||||
result=public_pairing_view(pending)
|
||||
result["broker_status"]="local_poll_interval"
|
||||
result["retry_after"]=max(1,int(remaining+0.999))
|
||||
return result
|
||||
|
||||
pending=store.update_pending_fields(pair_id,{"last_poll_at":now.isoformat()})
|
||||
if pending.get("state")!="pair_requested":
|
||||
return public_pairing_view(pending)
|
||||
device_code=str(pending.get("device_code") or "")
|
||||
if not device_code:
|
||||
return public_pairing_view(pending)
|
||||
|
||||
response=poster(
|
||||
_endpoint(broker,"fbpair/v1/token"),
|
||||
{
|
||||
"protocol":"FBPAIR/1",
|
||||
"pair_id":pair_id,
|
||||
"device_code":device_code,
|
||||
},
|
||||
)
|
||||
|
||||
# A network monitor may have invalidated this pair while the POST was in flight.
|
||||
current=store.load()
|
||||
if current.get("state")!="pair_requested" or str(current.get("pair_id") or "")!=pair_id:
|
||||
return public_pairing_view(current)
|
||||
|
||||
error=str(response.get("error") or "")
|
||||
if error in {"authorization_pending","slow_down"}:
|
||||
if error=="slow_down":
|
||||
current=store.update_pending_fields(
|
||||
pair_id,{"interval":min(60,int(current.get("interval") or 5)+5)}
|
||||
)
|
||||
result=public_pairing_view(current)
|
||||
result["broker_status"]=error
|
||||
return result
|
||||
if error=="access_denied":
|
||||
current=store.transition_pending(
|
||||
pair_id,"denied",{"denied_at":datetime.now(timezone.utc).isoformat()}
|
||||
)
|
||||
return public_pairing_view(current)
|
||||
if error in {"expired_token","expired"}:
|
||||
current=store.transition_pending(
|
||||
pair_id,"expired",
|
||||
{"expired_at":datetime.now(timezone.utc).isoformat(),"expiration_reason":"broker_expired"},
|
||||
)
|
||||
return public_pairing_view(current)
|
||||
if error:
|
||||
raise PairingClientError("PAIR_BROKER_"+error.upper())
|
||||
try:
|
||||
store.complete(response)
|
||||
except PairingError as exc:
|
||||
# If local invalidation won the race, do not revive the pairing.
|
||||
current=store.load()
|
||||
if current.get("state")!="pair_requested":
|
||||
return public_pairing_view(current)
|
||||
raise PairingClientError(str(exc)) from exc
|
||||
paired=store.load()
|
||||
try:
|
||||
confirm_pairing_remote(
|
||||
state_dir=Path(state_dir),
|
||||
allow_insecure_dev=allow_insecure_dev,
|
||||
poster=poster,
|
||||
record=paired,
|
||||
)
|
||||
except PairingClientError:
|
||||
current=public_pairing_view(store.load())
|
||||
current["broker_status"]="confirm_pending"
|
||||
return current
|
||||
return public_pairing_view(store.load())
|
||||
|
||||
+77
@@ -0,0 +1,77 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import subprocess
|
||||
import sys
|
||||
from typing import Any
|
||||
|
||||
from .network_security import PAIRING_WINDOW_TTL_SECONDS, pairing_policy
|
||||
from .product_core import ProductError
|
||||
|
||||
SUDO = "/usr/bin/sudo"
|
||||
PRIV_HELPER = "/usr/local/libexec/framebridge-seed-priv"
|
||||
MIN_TTL = 30
|
||||
MAX_TTL = PAIRING_WINDOW_TTL_SECONDS
|
||||
|
||||
|
||||
def _decode_helper(proc: subprocess.CompletedProcess[str]) -> dict[str, Any]:
|
||||
try:
|
||||
value=json.loads(proc.stdout or "{}")
|
||||
except Exception as exc:
|
||||
raise ProductError("PAIRING_WINDOW_HELPER_INVALID_JSON") from exc
|
||||
if not isinstance(value,dict):
|
||||
raise ProductError("PAIRING_WINDOW_HELPER_INVALID_RESPONSE")
|
||||
if proc.returncode!=0 or value.get("ok") is not True:
|
||||
raise ProductError(str(value.get("error") or "PAIRING_WINDOW_HELPER_FAILED"))
|
||||
return value
|
||||
|
||||
|
||||
def _run_helper(args: list[str], *, interactive: bool=True) -> dict[str, Any]:
|
||||
argv=[SUDO]
|
||||
if not interactive:
|
||||
argv.append("-n")
|
||||
argv.extend([PRIV_HELPER,*args])
|
||||
proc=subprocess.run(
|
||||
argv,
|
||||
text=True,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=None,
|
||||
check=False,
|
||||
shell=False,
|
||||
)
|
||||
return _decode_helper(proc)
|
||||
|
||||
|
||||
def open_pairing_window(port: int=18765, ttl_seconds: int=PAIRING_WINDOW_TTL_SECONDS) -> dict[str, Any]:
|
||||
if not sys.stdin.isatty():
|
||||
raise ProductError("PAIRING_WINDOW_LOCAL_TTY_REQUIRED")
|
||||
if isinstance(ttl_seconds,bool) or not isinstance(ttl_seconds,int) or not (MIN_TTL <= ttl_seconds <= MAX_TTL):
|
||||
raise ProductError("PAIRING_WINDOW_TTL_INVALID")
|
||||
policy=pairing_policy(port)
|
||||
if policy.get("ok") is not True:
|
||||
raise ProductError(str(policy.get("reason") or "PAIRING_WINDOW_POLICY_BLOCKED"))
|
||||
candidate=policy.get("candidate") or {}
|
||||
required=("interface","subnet","address","network_fingerprint")
|
||||
if not all(candidate.get(k) for k in required):
|
||||
raise ProductError("PAIRING_WINDOW_CANDIDATE_INCOMPLETE")
|
||||
result=_run_helper([
|
||||
"pairing-window-open",
|
||||
"--interface",str(candidate["interface"]),
|
||||
"--source-subnet",str(candidate["subnet"]),
|
||||
"--destination-address",str(candidate["address"]),
|
||||
"--port",str(int(port)),
|
||||
"--network-fingerprint",str(candidate["network_fingerprint"]),
|
||||
"--ttl",str(int(ttl_seconds)),
|
||||
])
|
||||
result["policy_state"]=policy.get("state")
|
||||
result["transport"]="LAN_HTTP"
|
||||
result["listen_url_base"]=f"http://{candidate['address']}:{int(port)}"
|
||||
return result
|
||||
|
||||
|
||||
def pairing_window_status(*, non_interactive: bool=False) -> dict[str, Any]:
|
||||
return _run_helper(["pairing-window-status"],interactive=not non_interactive)
|
||||
|
||||
|
||||
def close_pairing_window(reason: str="user", *, non_interactive: bool=False) -> dict[str, Any]:
|
||||
return _run_helper(["pairing-window-close","--reason",str(reason)],interactive=not non_interactive)
|
||||
+1388
File diff suppressed because it is too large.
Load diff
+1540
File diff suppressed because it is too large.
Load diff
+349
@@ -0,0 +1,349 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import sqlite3
|
||||
import tempfile
|
||||
import threading
|
||||
import uuid
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timezone
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from . import __version__
|
||||
|
||||
SECRET_KEY_RE = re.compile(r"(secret|token|password|passwd|api[_-]?key|private[_-]?key|credential|authorization|cookie)", re.I)
|
||||
TOKEN_RE = re.compile(r"(?i)\\b(?:sk-[A-Za-z0-9_-]{8,}|gh[pousr]_[A-Za-z0-9_]{10,}|github_pat_[A-Za-z0-9_]{10,}|xox[baprs]-[A-Za-z0-9-]{10,})\\b")
|
||||
|
||||
class ProductError(RuntimeError):
|
||||
pass
|
||||
|
||||
class PolicyDenied(ProductError):
|
||||
pass
|
||||
|
||||
def utc_now() -> str:
|
||||
return datetime.now(timezone.utc).isoformat()
|
||||
|
||||
def sha256_bytes(data: bytes) -> str:
|
||||
return hashlib.sha256(data).hexdigest()
|
||||
|
||||
def sha256_text(text: str) -> str:
|
||||
return sha256_bytes(text.encode("utf-8"))
|
||||
|
||||
def canonical_json(value: Any) -> str:
|
||||
return json.dumps(value, sort_keys=True, separators=(",", ":"), ensure_ascii=True)
|
||||
|
||||
def redact(value: Any, depth: int = 0) -> Any:
|
||||
if depth > 8:
|
||||
return "[DEPTH_LIMIT]"
|
||||
if isinstance(value, dict):
|
||||
out = {}
|
||||
for key, item in value.items():
|
||||
out[str(key)] = "[REDACTED]" if SECRET_KEY_RE.search(str(key)) else redact(item, depth + 1)
|
||||
return out
|
||||
if isinstance(value, list):
|
||||
return [redact(x, depth + 1) for x in value[:200]]
|
||||
if isinstance(value, tuple):
|
||||
return [redact(x, depth + 1) for x in value[:200]]
|
||||
if isinstance(value, str):
|
||||
clean = TOKEN_RE.sub("[REDACTED TOKEN]", value)
|
||||
return clean[:8192] + ("[TRUNCATED]" if len(clean) > 8192 else "")
|
||||
if isinstance(value, (int, float, bool)) or value is None:
|
||||
return value
|
||||
return redact(str(value), depth + 1)
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class PolicyDecision:
|
||||
allowed: bool
|
||||
mode: str
|
||||
reason: str
|
||||
|
||||
class PolicyEngine:
|
||||
VALID = {"allow", "deny", "require_authorized"}
|
||||
def __init__(self, policy: dict[str, Any] | None = None):
|
||||
policy = policy or {}
|
||||
self.default = str(policy.get("default", "deny")).strip().lower()
|
||||
if self.default not in self.VALID:
|
||||
self.default = "deny"
|
||||
self.tools = {}
|
||||
for key, value in (policy.get("tools") or {}).items():
|
||||
mode = str(value).strip().lower()
|
||||
self.tools[str(key)] = mode if mode in self.VALID else "deny"
|
||||
def decide(self, tool: str, authorized: bool = False) -> PolicyDecision:
|
||||
mode = self.tools.get(tool, self.default)
|
||||
if mode == "allow":
|
||||
return PolicyDecision(True, mode, "POLICY_ALLOW")
|
||||
if mode == "require_authorized" and authorized is True:
|
||||
return PolicyDecision(True, mode, "AUTHORIZED_TRUE")
|
||||
if mode == "require_authorized":
|
||||
return PolicyDecision(False, mode, "AUTHORIZED_TRUE_REQUIRED")
|
||||
return PolicyDecision(False, mode, "POLICY_DENY")
|
||||
|
||||
class AuditLog:
|
||||
def __init__(self, path: Path):
|
||||
self.path = Path(path)
|
||||
self.path.parent.mkdir(parents=True, exist_ok=True)
|
||||
self._lock = threading.Lock()
|
||||
def verify(self) -> dict[str, Any]:
|
||||
previous = "GENESIS"
|
||||
count = 0
|
||||
if not self.path.exists():
|
||||
return {"ok": True, "entries": 0, "last_hash": previous}
|
||||
try:
|
||||
with self.path.open("r", encoding="utf-8") as handle:
|
||||
for number, raw in enumerate(handle, 1):
|
||||
raw = raw.strip()
|
||||
if not raw:
|
||||
continue
|
||||
item = json.loads(raw)
|
||||
event_hash = str(item.pop("event_hash"))
|
||||
if item.get("prev_hash") != previous:
|
||||
return {"ok": False, "entries": count, "line": number, "error": "PREV_HASH_MISMATCH"}
|
||||
expected = sha256_text(previous + canonical_json(item))
|
||||
if event_hash != expected:
|
||||
return {"ok": False, "entries": count, "line": number, "error": "EVENT_HASH_MISMATCH"}
|
||||
previous = event_hash
|
||||
count += 1
|
||||
except Exception as exc:
|
||||
return {"ok": False, "entries": count, "error": type(exc).__name__}
|
||||
return {"ok": True, "entries": count, "last_hash": previous}
|
||||
def append(self, tool: str, status: str, request: Any = None, result: Any = None, principal: str = "chatgpt") -> dict[str, Any]:
|
||||
with self._lock:
|
||||
current = self.verify()
|
||||
if not current.get("ok"):
|
||||
raise ProductError("AUDIT_CHAIN_INVALID")
|
||||
previous = str(current.get("last_hash", "GENESIS"))
|
||||
event = {"event_id": uuid.uuid4().hex, "timestamp": utc_now(), "principal": str(principal), "tool": str(tool), "status": str(status), "request": redact(request), "result": redact(result), "prev_hash": previous}
|
||||
event_hash = sha256_text(previous + canonical_json(event))
|
||||
stored = dict(event)
|
||||
stored["event_hash"] = event_hash
|
||||
with self.path.open("a", encoding="utf-8") as handle:
|
||||
handle.write(canonical_json(stored) + "\n")
|
||||
handle.flush()
|
||||
os.fsync(handle.fileno())
|
||||
return {"event_id": event["event_id"], "event_hash": event_hash}
|
||||
def status(self) -> dict[str, Any]:
|
||||
value = self.verify()
|
||||
value["path"] = str(self.path)
|
||||
return value
|
||||
|
||||
def tail(self, limit: int = 20) -> dict[str, Any]:
|
||||
limit = max(1, min(int(limit), 100))
|
||||
if not self.path.exists():
|
||||
return {"ok": True, "events": []}
|
||||
events = []
|
||||
with self.path.open("r", encoding="utf-8") as handle:
|
||||
for raw in handle:
|
||||
raw = raw.strip()
|
||||
if not raw:
|
||||
continue
|
||||
try:
|
||||
item = json.loads(raw)
|
||||
except Exception:
|
||||
continue
|
||||
events.append({
|
||||
"timestamp": item.get("timestamp"),
|
||||
"tool": item.get("tool"),
|
||||
"status": item.get("status"),
|
||||
"event_hash": item.get("event_hash"),
|
||||
"request": item.get("request"),
|
||||
"result": item.get("result"),
|
||||
})
|
||||
return {"ok": True, "events": events[-limit:]}
|
||||
|
||||
|
||||
class SafeFiles:
|
||||
def __init__(self, roots: list[str]):
|
||||
self.roots = [Path(x).expanduser().resolve(strict=False) for x in roots]
|
||||
def _resolve(self, value: str) -> Path:
|
||||
if not self.roots:
|
||||
raise ProductError("NO_ALLOWED_ROOTS")
|
||||
candidate = Path(value).expanduser()
|
||||
if not candidate.is_absolute():
|
||||
candidate = self.roots[0] / candidate
|
||||
resolved = candidate.resolve(strict=False)
|
||||
for root in self.roots:
|
||||
try:
|
||||
if os.path.commonpath([str(root), str(resolved)]) == str(root):
|
||||
return resolved
|
||||
except ValueError:
|
||||
pass
|
||||
raise ProductError("PATH_OUTSIDE_ALLOWED_ROOTS")
|
||||
def list_dir(self, path: str = ".", limit: int = 200) -> dict[str, Any]:
|
||||
target = self._resolve(path)
|
||||
if not target.is_dir():
|
||||
raise ProductError("DIRECTORY_NOT_FOUND")
|
||||
limit = max(1, min(int(limit), 500))
|
||||
items = []
|
||||
for entry in sorted(target.iterdir(), key=lambda p: p.name.lower())[:limit]:
|
||||
info = entry.lstat()
|
||||
items.append({"name": entry.name, "type": "symlink" if entry.is_symlink() else ("dir" if entry.is_dir() else "file"), "size": info.st_size})
|
||||
return {"path": str(target), "items": items, "count": len(items)}
|
||||
def read_text(self, path: str, max_bytes: int = 65536) -> dict[str, Any]:
|
||||
target = self._resolve(path)
|
||||
if not target.is_file():
|
||||
raise ProductError("FILE_NOT_FOUND")
|
||||
max_bytes = max(1, min(int(max_bytes), 262144))
|
||||
size = target.stat().st_size
|
||||
if size > max_bytes:
|
||||
raise ProductError("FILE_TOO_LARGE")
|
||||
data = target.read_bytes()
|
||||
if b"\\x00" in data:
|
||||
raise ProductError("BINARY_FILE_NOT_ALLOWED")
|
||||
text = data.decode("utf-8", errors="strict")
|
||||
return {"path": str(target), "size": len(data), "sha256": sha256_bytes(data), "text": text}
|
||||
def write_text(self, path: str, text: str, expected_sha256: str = "", max_bytes: int = 262144) -> dict[str, Any]:
|
||||
target = self._resolve(path)
|
||||
parent = target.parent
|
||||
if not parent.is_dir():
|
||||
raise ProductError("PARENT_DIRECTORY_NOT_FOUND")
|
||||
if target.exists() and target.is_symlink():
|
||||
raise ProductError("SYMLINK_WRITE_BLOCKED")
|
||||
data = str(text).encode("utf-8")
|
||||
if len(data) > int(max_bytes):
|
||||
raise ProductError("WRITE_TOO_LARGE")
|
||||
before_sha = ""
|
||||
if target.exists():
|
||||
before_sha = sha256_bytes(target.read_bytes())
|
||||
if expected_sha256 and before_sha != expected_sha256:
|
||||
raise ProductError("EXPECTED_SHA256_MISMATCH")
|
||||
elif expected_sha256:
|
||||
raise ProductError("EXPECTED_SHA256_WITHOUT_EXISTING_FILE")
|
||||
fd, tmp_name = tempfile.mkstemp(prefix=".framebridge-", dir=str(parent))
|
||||
try:
|
||||
with os.fdopen(fd, "wb") as handle:
|
||||
handle.write(data)
|
||||
handle.flush()
|
||||
os.fsync(handle.fileno())
|
||||
os.replace(tmp_name, target)
|
||||
finally:
|
||||
if os.path.exists(tmp_name):
|
||||
os.unlink(tmp_name)
|
||||
return {"path": str(target), "size": len(data), "sha256": sha256_bytes(data), "previous_sha256": before_sha}
|
||||
|
||||
class MemoryIndex:
|
||||
ALLOWED_SUFFIXES = {".md", ".txt", ".json", ".yaml", ".yml"}
|
||||
BLOCKED_NAMES = {"rclone.conf", "runtime.env", ".env"}
|
||||
BLOCKED_SUFFIXES = {".pem", ".key", ".p12", ".pfx"}
|
||||
def __init__(self, db_path: Path, sources: list[str]):
|
||||
self.db_path = Path(db_path)
|
||||
self.db_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
self.sources = [Path(x).expanduser().resolve(strict=False) for x in sources]
|
||||
self._init_db()
|
||||
def _connect(self):
|
||||
con = sqlite3.connect(self.db_path)
|
||||
con.row_factory = sqlite3.Row
|
||||
return con
|
||||
def _init_db(self):
|
||||
with self._connect() as con:
|
||||
con.execute("CREATE TABLE IF NOT EXISTS documents (source_path TEXT PRIMARY KEY, title TEXT NOT NULL, sha256 TEXT NOT NULL, mtime REAL NOT NULL, body TEXT NOT NULL, indexed_at TEXT NOT NULL)")
|
||||
con.execute("CREATE VIRTUAL TABLE IF NOT EXISTS docs_fts USING fts5(source_path UNINDEXED, title, body)")
|
||||
def _safe_file(self, path: Path, root: Path) -> bool:
|
||||
try:
|
||||
resolved = path.resolve(strict=True)
|
||||
if os.path.commonpath([str(root), str(resolved)]) != str(root):
|
||||
return False
|
||||
except Exception:
|
||||
return False
|
||||
if path.is_symlink() or not path.is_file():
|
||||
return False
|
||||
if path.name in self.BLOCKED_NAMES or path.suffix.lower() in self.BLOCKED_SUFFIXES:
|
||||
return False
|
||||
if path.suffix.lower() not in self.ALLOWED_SUFFIXES:
|
||||
return False
|
||||
try:
|
||||
return path.stat().st_size <= 2_000_000
|
||||
except OSError:
|
||||
return False
|
||||
def refresh(self) -> dict[str, Any]:
|
||||
docs = []
|
||||
for root in self.sources:
|
||||
if not root.is_dir():
|
||||
continue
|
||||
for path in sorted(root.rglob("*")):
|
||||
if not self._safe_file(path, root):
|
||||
continue
|
||||
try:
|
||||
body = path.read_text(encoding="utf-8")
|
||||
except Exception:
|
||||
continue
|
||||
data = body.encode("utf-8")
|
||||
lines = body.splitlines()
|
||||
title = lines[0].lstrip("# ").strip() if lines else path.name
|
||||
docs.append({"source_path": str(path.resolve()), "title": title[:500] or path.name, "sha256": sha256_bytes(data), "mtime": path.stat().st_mtime, "body": body, "indexed_at": utc_now()})
|
||||
with self._connect() as con:
|
||||
con.execute("DELETE FROM documents")
|
||||
con.execute("DELETE FROM docs_fts")
|
||||
for item in docs:
|
||||
con.execute("INSERT INTO documents(source_path,title,sha256,mtime,body,indexed_at) VALUES (?,?,?,?,?,?)", (item["source_path"], item["title"], item["sha256"], item["mtime"], item["body"], item["indexed_at"]))
|
||||
con.execute("INSERT INTO docs_fts(source_path,title,body) VALUES (?,?,?)", (item["source_path"], item["title"], item["body"]))
|
||||
return self.status()
|
||||
def status(self) -> dict[str, Any]:
|
||||
with self._connect() as con:
|
||||
row = con.execute("SELECT COUNT(*) AS n, MAX(indexed_at) AS updated FROM documents").fetchone()
|
||||
return {"db_path": str(self.db_path), "documents": int(row["n"]), "last_indexed_at": row["updated"], "sources": [str(x) for x in self.sources]}
|
||||
def search(self, query: str, limit: int = 8) -> dict[str, Any]:
|
||||
tokens = re.findall(r"[A-Za-z0-9_]{2,}", str(query))
|
||||
if not tokens:
|
||||
return {"query": str(query), "results": []}
|
||||
expression = " AND ".join('"' + token.replace('"', '') + '"' for token in tokens[:12])
|
||||
limit = max(1, min(int(limit), 20))
|
||||
with self._connect() as con:
|
||||
rows = con.execute("SELECT source_path,title,snippet(docs_fts,2,'[',']','...',24) AS snippet,bm25(docs_fts) AS rank FROM docs_fts WHERE docs_fts MATCH ? ORDER BY rank LIMIT ?", (expression, limit)).fetchall()
|
||||
return {"query": str(query), "results": [{"source_path": row["source_path"], "title": row["title"], "snippet": row["snippet"], "rank": row["rank"]} for row in rows]}
|
||||
def read(self, source_path: str, start_line: int = 1, max_lines: int = 120) -> dict[str, Any]:
|
||||
start_line = max(1, int(start_line))
|
||||
max_lines = max(1, min(int(max_lines), 400))
|
||||
with self._connect() as con:
|
||||
row = con.execute("SELECT source_path,title,sha256,body,indexed_at FROM documents WHERE source_path=?", (str(source_path),)).fetchone()
|
||||
if row is None:
|
||||
raise ProductError("MEMORY_DOCUMENT_NOT_FOUND")
|
||||
lines = row["body"].splitlines()
|
||||
selected = lines[start_line - 1:start_line - 1 + max_lines]
|
||||
return {"source_path": row["source_path"], "title": row["title"], "sha256": row["sha256"], "indexed_at": row["indexed_at"], "start_line": start_line, "end_line": start_line + len(selected) - 1 if selected else start_line - 1, "text": "\n".join(selected)}
|
||||
|
||||
DEFAULT_POLICY = {"default": "deny", "tools": {"framebridge_product_status": "allow", "framebridge_files_list": "allow", "framebridge_files_read": "allow", "framebridge_files_write": "require_authorized", "framebridge_memory_status": "allow", "framebridge_memory_refresh": "require_authorized", "framebridge_memory_search": "allow", "framebridge_memory_read": "allow", "framebridge_audit_verify": "allow", "framebridge_exec": "require_authorized", "framebridge_cua": "require_authorized"}}
|
||||
|
||||
def load_config(path: Path) -> dict[str, Any]:
|
||||
if path.is_file():
|
||||
value = json.loads(path.read_text(encoding="utf-8"))
|
||||
if not isinstance(value, dict):
|
||||
raise ProductError("CONFIG_NOT_OBJECT")
|
||||
return value
|
||||
return {"node_name": os.uname().nodename, "state_dir": "~/.local/share/framebridge-linux-chatgpt", "allowed_roots": ["~/FrameBridge-Workspace"], "memory_sources": [], "policy": DEFAULT_POLICY}
|
||||
|
||||
class AlphaContext:
|
||||
def __init__(self, config_path: Path):
|
||||
self.config_path = Path(config_path).expanduser()
|
||||
self.config = load_config(self.config_path)
|
||||
self.node_name = str(self.config.get("node_name") or os.uname().nodename)
|
||||
self.state_dir = Path(str(self.config.get("state_dir") or "~/.local/share/framebridge-linux-chatgpt")).expanduser()
|
||||
self.state_dir.mkdir(parents=True, exist_ok=True)
|
||||
self.policy = PolicyEngine(self.config.get("policy") or DEFAULT_POLICY)
|
||||
self.files = SafeFiles(list(self.config.get("allowed_roots") or []))
|
||||
self.memory = MemoryIndex(self.state_dir / "memory.sqlite3", list(self.config.get("memory_sources") or []))
|
||||
self.audit = AuditLog(self.state_dir / "audit.jsonl")
|
||||
@classmethod
|
||||
def from_environment(cls) -> "AlphaContext":
|
||||
raw = os.environ.get("FRAMEBRIDGE_ALPHA_CONFIG", "~/.config/framebridge-linux-chatgpt/config.json")
|
||||
return cls(Path(raw).expanduser())
|
||||
def status(self) -> dict[str, Any]:
|
||||
return {"ok": True, "version": __version__, "node_name": self.node_name, "config_path": str(self.config_path), "state_dir": str(self.state_dir), "allowed_roots": [str(x) for x in self.files.roots], "memory": self.memory.status(), "audit": self.audit.status(), "policy_default": self.policy.default}
|
||||
def guarded(self, tool: str, authorized: bool, request: dict[str, Any], fn):
|
||||
decision = self.policy.decide(tool, authorized)
|
||||
if not decision.allowed:
|
||||
result = {"ok": False, "error": decision.reason, "policy_mode": decision.mode}
|
||||
self.audit.append(tool, "DENY", request=request, result=result)
|
||||
return result
|
||||
try:
|
||||
value = fn()
|
||||
result = {"ok": True, "policy_mode": decision.mode, "result": value}
|
||||
self.audit.append(tool, "PASS", request=request, result=result)
|
||||
return result
|
||||
except Exception as exc:
|
||||
result = {"ok": False, "error": str(exc), "error_type": type(exc).__name__}
|
||||
self.audit.append(tool, "ERROR", request=request, result=result)
|
||||
return result
|
||||
+67
@@ -0,0 +1,67 @@
|
||||
from __future__ import annotations
|
||||
import sys
|
||||
from .vendor.qrcodegen import QrCode
|
||||
|
||||
class TerminalQrError(RuntimeError):
|
||||
pass
|
||||
|
||||
def _prepare(text, border, stream, require_tty):
|
||||
if not isinstance(text,str) or not text: raise TerminalQrError('QR_TEXT_INVALID')
|
||||
if not (2 <= int(border) <= 8): raise TerminalQrError('QR_BORDER_INVALID')
|
||||
out=stream or sys.stdout
|
||||
if require_tty and not getattr(out,'isatty',lambda:False)(): raise TerminalQrError('INTERACTIVE_TTY_REQUIRED')
|
||||
return QrCode.encode_text(text,QrCode.Ecc.MEDIUM),out
|
||||
|
||||
def _dark(qr,border,x,y):
|
||||
size=qr.get_size()
|
||||
if x<border or y<border or x>=border+size or y>=border+size: return False
|
||||
return qr.get_module(x-border,y-border)
|
||||
|
||||
def render_qr_compact(text, *, border=4, stream=None, require_tty=True):
|
||||
qr,out=_prepare(text,border,stream,require_tty)
|
||||
total=qr.get_size()+border*2
|
||||
for y in range(0,total,2):
|
||||
row=[]
|
||||
for x in range(total):
|
||||
top=_dark(qr,border,x,y)
|
||||
bottom=_dark(qr,border,x,y+1) if y+1<total else False
|
||||
if top and bottom: row.append('\x1b[40m ')
|
||||
elif top: row.append('\x1b[30;47m▀')
|
||||
elif bottom: row.append('\x1b[30;47m▄')
|
||||
else: row.append('\x1b[47m ')
|
||||
out.write(''.join(row)+'\x1b[0m\n')
|
||||
out.flush()
|
||||
|
||||
def render_qr_mono(text, *, border=4, stream=None, require_tty=True):
|
||||
qr,out=_prepare(text,border,stream,require_tty)
|
||||
total=qr.get_size()+border*2
|
||||
for y in range(0,total,2):
|
||||
row=[]
|
||||
for x in range(total):
|
||||
top=_dark(qr,border,x,y)
|
||||
bottom=_dark(qr,border,x,y+1) if y+1<total else False
|
||||
row.append('█' if top and bottom else '▀' if top else '▄' if bottom else ' ')
|
||||
out.write(''.join(row)+'\n')
|
||||
out.flush()
|
||||
|
||||
def render_qr_ansi(text, *, border=4, stream=None, require_tty=True):
|
||||
qr,out=_prepare(text,border,stream,require_tty)
|
||||
size=qr.get_size(); black='\x1b[40m '; white='\x1b[47m '; reset='\x1b[0m'; quiet=white*(size+border*2)+reset+'\n'
|
||||
for _ in range(border): out.write(quiet)
|
||||
for y in range(size):
|
||||
row=[white]*border
|
||||
for x in range(size): row.append(black if qr.get_module(x,y) else white)
|
||||
row.extend([white]*border); out.write(''.join(row)+reset+'\n')
|
||||
for _ in range(border): out.write(quiet)
|
||||
out.flush()
|
||||
|
||||
def render_qr_unicode(text, *, border=4, stream=None, require_tty=True):
|
||||
qr,out=_prepare(text,border,stream,require_tty)
|
||||
size=qr.get_size(); pad=' '*(size+border*2)
|
||||
for _ in range(border): out.write(pad+'\n')
|
||||
for y in range(size):
|
||||
out.write(' '*border)
|
||||
for x in range(size): out.write('██' if qr.get_module(x,y) else ' ')
|
||||
out.write(' '*border+'\n')
|
||||
for _ in range(border): out.write(pad+'\n')
|
||||
out.flush()
|
||||
+780
@@ -0,0 +1,780 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import platform
|
||||
import shutil
|
||||
import socket
|
||||
import subprocess
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from .product_core import ProductError
|
||||
|
||||
MAX_OUTPUT = 32768
|
||||
|
||||
|
||||
def _run(args: list[str], timeout: float = 5.0) -> dict[str, Any]:
|
||||
try:
|
||||
proc = subprocess.run(
|
||||
args,
|
||||
text=True,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.PIPE,
|
||||
timeout=timeout,
|
||||
check=False,
|
||||
shell=False,
|
||||
)
|
||||
return {
|
||||
"ok": proc.returncode == 0,
|
||||
"rc": proc.returncode,
|
||||
"stdout": proc.stdout[:MAX_OUTPUT],
|
||||
"stderr": proc.stderr[:4096],
|
||||
}
|
||||
except FileNotFoundError:
|
||||
return {"ok": False, "error": "COMMAND_NOT_FOUND"}
|
||||
except subprocess.TimeoutExpired:
|
||||
return {"ok": False, "error": "COMMAND_TIMEOUT"}
|
||||
except Exception as exc:
|
||||
return {"ok": False, "error": type(exc).__name__}
|
||||
|
||||
|
||||
def _os_release() -> dict[str, str]:
|
||||
path = Path("/etc/os-release")
|
||||
out: dict[str, str] = {}
|
||||
if not path.is_file():
|
||||
return out
|
||||
for line in path.read_text(encoding="utf-8", errors="replace").splitlines():
|
||||
if "=" not in line:
|
||||
continue
|
||||
key, value = line.split("=", 1)
|
||||
out[key] = value.strip().strip('"')
|
||||
return out
|
||||
|
||||
|
||||
def system_status(args: dict[str, Any]) -> dict[str, Any]:
|
||||
osr = _os_release()
|
||||
uptime = _run(["uptime", "-p"]) if shutil.which("uptime") else {"ok": False}
|
||||
return {
|
||||
"hostname": socket.gethostname(),
|
||||
"kernel": platform.release(),
|
||||
"architecture": platform.machine(),
|
||||
"os_id": osr.get("ID"),
|
||||
"os_name": osr.get("PRETTY_NAME"),
|
||||
"uptime": uptime.get("stdout", "").strip(),
|
||||
}
|
||||
|
||||
|
||||
def hardware_summary(args: dict[str, Any]) -> dict[str, Any]:
|
||||
cpu_model = ""
|
||||
cpuinfo = Path("/proc/cpuinfo")
|
||||
if cpuinfo.is_file():
|
||||
for line in cpuinfo.read_text(errors="replace").splitlines():
|
||||
if line.lower().startswith("model name") and ":" in line:
|
||||
cpu_model = line.split(":", 1)[1].strip()
|
||||
break
|
||||
|
||||
mem_total_kib = None
|
||||
meminfo = Path("/proc/meminfo")
|
||||
if meminfo.is_file():
|
||||
for line in meminfo.read_text(errors="replace").splitlines():
|
||||
if line.startswith("MemTotal:"):
|
||||
parts = line.split()
|
||||
if len(parts) >= 2 and parts[1].isdigit():
|
||||
mem_total_kib = int(parts[1])
|
||||
break
|
||||
|
||||
return {
|
||||
"cpu_model": cpu_model or None,
|
||||
"logical_cpus": os.cpu_count(),
|
||||
"memory_total_kib": mem_total_kib,
|
||||
"architecture": platform.machine(),
|
||||
}
|
||||
|
||||
|
||||
def storage_inspect(args: dict[str, Any]) -> dict[str, Any]:
|
||||
if not shutil.which("lsblk"):
|
||||
return {"ok": False, "error": "LSBLK_NOT_FOUND"}
|
||||
result = _run([
|
||||
"lsblk",
|
||||
"-J",
|
||||
"-o",
|
||||
"NAME,TYPE,SIZE,FSTYPE,MOUNTPOINTS,ROTA,MODEL",
|
||||
])
|
||||
if not result.get("ok"):
|
||||
return result
|
||||
try:
|
||||
return {
|
||||
"ok": True,
|
||||
"blockdevices": json.loads(result["stdout"]).get("blockdevices", []),
|
||||
}
|
||||
except Exception:
|
||||
return {"ok": False, "error": "LSBLK_JSON_INVALID"}
|
||||
|
||||
|
||||
def network_inspect(args: dict[str, Any]) -> dict[str, Any]:
|
||||
if not shutil.which("ip"):
|
||||
return {"ok": False, "error": "IP_COMMAND_NOT_FOUND"}
|
||||
addresses = _run(["ip", "-j", "address"])
|
||||
routes = _run(["ip", "-j", "route"])
|
||||
out: dict[str, Any] = {
|
||||
"ok": bool(addresses.get("ok") and routes.get("ok")),
|
||||
}
|
||||
try:
|
||||
out["addresses"] = json.loads(addresses.get("stdout", "[]"))
|
||||
except Exception:
|
||||
out["addresses"] = []
|
||||
out["ok"] = False
|
||||
try:
|
||||
out["routes"] = json.loads(routes.get("stdout", "[]"))
|
||||
except Exception:
|
||||
out["routes"] = []
|
||||
out["ok"] = False
|
||||
return out
|
||||
|
||||
|
||||
def service_status(args: dict[str, Any]) -> dict[str, Any]:
|
||||
service = str(args.get("service") or "").strip()
|
||||
if not service:
|
||||
raise ProductError("SERVICE_REQUIRED")
|
||||
if len(service) > 128:
|
||||
raise ProductError("SERVICE_NAME_TOO_LONG")
|
||||
allowed = set("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789@_.-")
|
||||
if any(ch not in allowed for ch in service):
|
||||
raise ProductError("SERVICE_NAME_INVALID")
|
||||
if not shutil.which("systemctl"):
|
||||
return {"ok": False, "error": "SYSTEMCTL_NOT_FOUND"}
|
||||
|
||||
active = _run(["systemctl", "is-active", service])
|
||||
enabled = _run(["systemctl", "is-enabled", service])
|
||||
return {
|
||||
"service": service,
|
||||
"active": active.get("stdout", "").strip(),
|
||||
"active_rc": active.get("rc"),
|
||||
"enabled": enabled.get("stdout", "").strip(),
|
||||
"enabled_rc": enabled.get("rc"),
|
||||
}
|
||||
|
||||
|
||||
|
||||
PACKAGE_RE = re.compile(r"^[a-z0-9][a-z0-9+.-]{0,127}$")
|
||||
SIMPLE_NAME_RE = re.compile(r"^[A-Za-z0-9@_.-]{1,128}$")
|
||||
SAFE_TEXT_ROOTS = (
|
||||
Path("/etc"),
|
||||
Path("/usr/local/etc"),
|
||||
Path("/var/log"),
|
||||
Path("/var/lib/framebridge-seed"),
|
||||
)
|
||||
SAFE_TEXT_FILES = {
|
||||
Path("/usr/lib/os-release"),
|
||||
}
|
||||
SENSITIVE_PATH_PARTS = (
|
||||
"/etc/shadow",
|
||||
"/etc/gshadow",
|
||||
"/etc/sudoers",
|
||||
"/etc/sudoers.d/",
|
||||
"/etc/ssl/private/",
|
||||
"/etc/wireguard/",
|
||||
"/etc/openvpn/",
|
||||
"/etc/NetworkManager/system-connections/",
|
||||
"/etc/apt/auth.conf",
|
||||
"/.ssh/",
|
||||
)
|
||||
SENSITIVE_NAME_RE = re.compile(
|
||||
r"(?i)(?:^|[._-])(?:secret|token|password|passwd|credential|private[_-]?key|api[_-]?key)(?:$|[._-])"
|
||||
)
|
||||
SENSITIVE_ASSIGNMENT_RE = re.compile(
|
||||
r"(?i)^\s*([A-Za-z0-9_.-]*(?:secret|token|password|passwd|api[_-]?key|"
|
||||
r"private[_-]?key|cookie|authorization|credential)[A-Za-z0-9_.-]*)\s*[:=]"
|
||||
)
|
||||
BEARER_RE = re.compile(r"(?i)\bBearer\s+[^\s,;]+")
|
||||
URL_USERINFO_RE = re.compile(r"(?i)(https?://)[^/@\s]+:[^/@\s]+@")
|
||||
PRIVATE_KEY_BEGIN_RE = re.compile(r"-----BEGIN [A-Z0-9 ]*PRIVATE KEY-----")
|
||||
PRIVATE_KEY_END_RE = re.compile(r"-----END [A-Z0-9 ]*PRIVATE KEY-----")
|
||||
|
||||
|
||||
def _bounded_int(value: Any, *, default: int, low: int, high: int, error: str) -> int:
|
||||
if value is None:
|
||||
return default
|
||||
if type(value) is not int or not (low <= value <= high):
|
||||
raise ProductError(error)
|
||||
return value
|
||||
|
||||
|
||||
def _simple_name(value: Any, error: str) -> str:
|
||||
name = str(value or "").strip()
|
||||
if not SIMPLE_NAME_RE.fullmatch(name):
|
||||
raise ProductError(error)
|
||||
return name
|
||||
|
||||
|
||||
def _package_name(value: Any) -> str:
|
||||
package = str(value or "").strip()
|
||||
if not PACKAGE_RE.fullmatch(package):
|
||||
raise ProductError("PACKAGE_NAME_INVALID")
|
||||
return package
|
||||
|
||||
|
||||
def _sanitize_text(text: str) -> str:
|
||||
out: list[str] = []
|
||||
in_key = False
|
||||
for line in str(text).splitlines():
|
||||
if PRIVATE_KEY_BEGIN_RE.search(line):
|
||||
in_key = True
|
||||
out.append("[REDACTED PRIVATE KEY BLOCK]")
|
||||
continue
|
||||
if in_key:
|
||||
if PRIVATE_KEY_END_RE.search(line):
|
||||
in_key = False
|
||||
continue
|
||||
m = SENSITIVE_ASSIGNMENT_RE.search(line)
|
||||
if m:
|
||||
out.append(f"{m.group(1)}=[REDACTED]")
|
||||
continue
|
||||
line = BEARER_RE.sub("Bearer [REDACTED]", line)
|
||||
line = URL_USERINFO_RE.sub(r"\1[REDACTED]@", line)
|
||||
out.append(line)
|
||||
return "\n".join(out)
|
||||
|
||||
|
||||
def _safe_text_path(raw: Any) -> Path:
|
||||
value = str(raw or "").strip()
|
||||
if not value or len(value) > 4096 or "\x00" in value or "\n" in value:
|
||||
raise ProductError("PATH_INVALID")
|
||||
src = Path(value)
|
||||
if not src.is_absolute():
|
||||
raise ProductError("PATH_ABSOLUTE_REQUIRED")
|
||||
try:
|
||||
resolved = src.resolve(strict=True)
|
||||
except FileNotFoundError:
|
||||
raise ProductError("PATH_NOT_FOUND") from None
|
||||
except OSError:
|
||||
raise ProductError("PATH_INVALID") from None
|
||||
|
||||
allowed = False
|
||||
for root in SAFE_TEXT_ROOTS:
|
||||
try:
|
||||
if os.path.commonpath([str(root), str(resolved)]) == str(root):
|
||||
allowed = True
|
||||
break
|
||||
except ValueError:
|
||||
pass
|
||||
if not allowed and resolved not in SAFE_TEXT_FILES:
|
||||
raise ProductError("PATH_OUTSIDE_SAFE_READ_ROOTS")
|
||||
|
||||
lower = str(resolved).lower()
|
||||
if any(part.lower() in lower for part in SENSITIVE_PATH_PARTS):
|
||||
raise ProductError("SENSITIVE_PATH_FORBIDDEN")
|
||||
if SENSITIVE_NAME_RE.search(resolved.name):
|
||||
raise ProductError("SENSITIVE_PATH_FORBIDDEN")
|
||||
if not resolved.is_file():
|
||||
raise ProductError("PATH_NOT_REGULAR_FILE")
|
||||
return resolved
|
||||
|
||||
|
||||
def system_resources(args: dict[str, Any]) -> dict[str, Any]:
|
||||
mem: dict[str, int] = {}
|
||||
p = Path("/proc/meminfo")
|
||||
if p.is_file():
|
||||
for line in p.read_text(errors="replace").splitlines():
|
||||
if ":" not in line:
|
||||
continue
|
||||
key, rest = line.split(":", 1)
|
||||
fields = rest.strip().split()
|
||||
if fields and fields[0].isdigit():
|
||||
mem[key] = int(fields[0])
|
||||
|
||||
load = os.getloadavg() if hasattr(os, "getloadavg") else (None, None, None)
|
||||
process_count = 0
|
||||
proc = Path("/proc")
|
||||
if proc.is_dir():
|
||||
process_count = sum(1 for x in proc.iterdir() if x.name.isdigit())
|
||||
|
||||
return {
|
||||
"load_1": load[0],
|
||||
"load_5": load[1],
|
||||
"load_15": load[2],
|
||||
"memory_total_kib": mem.get("MemTotal"),
|
||||
"memory_available_kib": mem.get("MemAvailable"),
|
||||
"swap_total_kib": mem.get("SwapTotal"),
|
||||
"swap_free_kib": mem.get("SwapFree"),
|
||||
"process_count": process_count,
|
||||
}
|
||||
|
||||
|
||||
def system_processes(args: dict[str, Any]) -> dict[str, Any]:
|
||||
limit = _bounded_int(args.get("limit"), default=40, low=1, high=200, error="PROCESS_LIMIT_INVALID")
|
||||
result = _run([
|
||||
"ps", "-eo",
|
||||
"pid=,ppid=,user=,stat=,pcpu=,pmem=,etime=,comm=",
|
||||
"--sort=-pcpu",
|
||||
])
|
||||
if not result.get("ok"):
|
||||
return result
|
||||
rows = []
|
||||
for raw in result.get("stdout", "").splitlines():
|
||||
parts = raw.split(None, 7)
|
||||
if len(parts) != 8:
|
||||
continue
|
||||
pid, ppid, user, state, cpu, mem, etime, comm = parts
|
||||
rows.append({
|
||||
"pid": int(pid) if pid.isdigit() else pid,
|
||||
"ppid": int(ppid) if ppid.isdigit() else ppid,
|
||||
"user": user,
|
||||
"state": state,
|
||||
"cpu_pct": cpu,
|
||||
"mem_pct": mem,
|
||||
"etime": etime,
|
||||
"comm": comm,
|
||||
})
|
||||
if len(rows) >= limit:
|
||||
break
|
||||
return {"ok": True, "processes": rows, "limit": limit}
|
||||
|
||||
|
||||
def system_time(args: dict[str, Any]) -> dict[str, Any]:
|
||||
timezone = None
|
||||
if shutil.which("timedatectl"):
|
||||
r = _run(["timedatectl", "show", "-p", "Timezone", "--value"])
|
||||
if r.get("ok"):
|
||||
timezone = r.get("stdout", "").strip() or None
|
||||
if timezone is None:
|
||||
tz = Path("/etc/timezone")
|
||||
if tz.is_file():
|
||||
timezone = tz.read_text(errors="replace").strip() or None
|
||||
return {
|
||||
"local_time": datetime.now().astimezone().isoformat(),
|
||||
"timezone": timezone,
|
||||
}
|
||||
|
||||
|
||||
def kernel_modules(args: dict[str, Any]) -> dict[str, Any]:
|
||||
limit = _bounded_int(args.get("limit"), default=100, low=1, high=500, error="MODULE_LIMIT_INVALID")
|
||||
if not shutil.which("lsmod"):
|
||||
return {"ok": False, "error": "LSMOD_NOT_FOUND"}
|
||||
r = _run(["lsmod"])
|
||||
if not r.get("ok"):
|
||||
return r
|
||||
rows = []
|
||||
for line in r.get("stdout", "").splitlines()[1:]:
|
||||
parts = line.split()
|
||||
if len(parts) < 3:
|
||||
continue
|
||||
rows.append({
|
||||
"module": parts[0],
|
||||
"size": int(parts[1]) if parts[1].isdigit() else parts[1],
|
||||
"used_by_count": int(parts[2]) if parts[2].isdigit() else parts[2],
|
||||
})
|
||||
if len(rows) >= limit:
|
||||
break
|
||||
return {"ok": True, "modules": rows, "limit": limit}
|
||||
|
||||
|
||||
def storage_usage(args: dict[str, Any]) -> dict[str, Any]:
|
||||
r = _run(["df", "-B1", "-P", "-T"])
|
||||
if not r.get("ok"):
|
||||
return r
|
||||
rows = []
|
||||
for line in r.get("stdout", "").splitlines()[1:]:
|
||||
parts = line.split()
|
||||
if len(parts) < 7:
|
||||
continue
|
||||
fs, fstype, total, used, avail, pct = parts[:6]
|
||||
mount = " ".join(parts[6:])
|
||||
rows.append({
|
||||
"filesystem": fs,
|
||||
"fstype": fstype,
|
||||
"bytes_total": int(total) if total.isdigit() else total,
|
||||
"bytes_used": int(used) if used.isdigit() else used,
|
||||
"bytes_available": int(avail) if avail.isdigit() else avail,
|
||||
"use_percent": pct,
|
||||
"mountpoint": mount,
|
||||
})
|
||||
return {"ok": True, "filesystems": rows}
|
||||
|
||||
|
||||
def storage_mounts(args: dict[str, Any]) -> dict[str, Any]:
|
||||
if not shutil.which("findmnt"):
|
||||
return {"ok": False, "error": "FINDMNT_NOT_FOUND"}
|
||||
r = _run(["findmnt", "-J", "-o", "TARGET,SOURCE,FSTYPE"])
|
||||
if not r.get("ok"):
|
||||
return r
|
||||
try:
|
||||
return {"ok": True, "filesystems": json.loads(r["stdout"]).get("filesystems", [])}
|
||||
except Exception:
|
||||
return {"ok": False, "error": "FINDMNT_JSON_INVALID"}
|
||||
|
||||
|
||||
def network_sockets(args: dict[str, Any]) -> dict[str, Any]:
|
||||
limit = _bounded_int(args.get("limit"), default=100, low=1, high=500, error="SOCKET_LIMIT_INVALID")
|
||||
if not shutil.which("ss"):
|
||||
return {"ok": False, "error": "SS_NOT_FOUND"}
|
||||
r = _run(["ss", "-H", "-lntu"])
|
||||
if not r.get("ok"):
|
||||
return r
|
||||
rows = []
|
||||
for line in r.get("stdout", "").splitlines():
|
||||
parts = line.split()
|
||||
if len(parts) < 5:
|
||||
continue
|
||||
rows.append({
|
||||
"netid": parts[0],
|
||||
"state": parts[1],
|
||||
"local": parts[4] if len(parts) > 4 else None,
|
||||
"peer": parts[5] if len(parts) > 5 else None,
|
||||
})
|
||||
if len(rows) >= limit:
|
||||
break
|
||||
return {"ok": True, "listeners": rows, "limit": limit}
|
||||
|
||||
|
||||
def network_dns(args: dict[str, Any]) -> dict[str, Any]:
|
||||
path = Path("/etc/resolv.conf")
|
||||
if not path.exists():
|
||||
return {"ok": False, "error": "RESOLV_CONF_NOT_FOUND"}
|
||||
try:
|
||||
resolved = path.resolve(strict=True)
|
||||
except Exception:
|
||||
resolved = path
|
||||
lines = []
|
||||
try:
|
||||
text = path.read_text(errors="replace")
|
||||
except Exception as exc:
|
||||
return {"ok": False, "error": type(exc).__name__}
|
||||
for line in text.splitlines():
|
||||
stripped = line.strip()
|
||||
if stripped.startswith(("nameserver ", "search ", "domain ", "options ")):
|
||||
lines.append(stripped)
|
||||
return {
|
||||
"ok": True,
|
||||
"path": str(path),
|
||||
"resolved_path": str(resolved),
|
||||
"directives": lines,
|
||||
}
|
||||
|
||||
|
||||
def package_status_read(args: dict[str, Any]) -> dict[str, Any]:
|
||||
package = _package_name(args.get("package"))
|
||||
r = _run(["dpkg-query", "-W", "-f=${Status}|${Version}|${Architecture}", package])
|
||||
if not r.get("ok"):
|
||||
return {"ok": True, "package": package, "installed": False, "version": None, "architecture": None}
|
||||
out = r.get("stdout", "").strip()
|
||||
parts = out.split("|")
|
||||
installed = len(parts) >= 1 and parts[0] == "install ok installed"
|
||||
return {
|
||||
"ok": True,
|
||||
"package": package,
|
||||
"installed": installed,
|
||||
"version": parts[1] if installed and len(parts) > 1 else None,
|
||||
"architecture": parts[2] if installed and len(parts) > 2 else None,
|
||||
}
|
||||
|
||||
|
||||
def package_search(args: dict[str, Any]) -> dict[str, Any]:
|
||||
query = str(args.get("query") or "").strip()
|
||||
if not query or len(query) > 80:
|
||||
raise ProductError("PACKAGE_QUERY_INVALID")
|
||||
if not re.fullmatch(r"[A-Za-z0-9+._ -]+", query):
|
||||
raise ProductError("PACKAGE_QUERY_INVALID")
|
||||
limit = _bounded_int(args.get("limit"), default=30, low=1, high=100, error="PACKAGE_LIMIT_INVALID")
|
||||
if not shutil.which("apt-cache"):
|
||||
return {"ok": False, "error": "APT_CACHE_NOT_FOUND"}
|
||||
r = _run(["apt-cache", "search", "--names-only", re.escape(query)], timeout=15.0)
|
||||
if not r.get("ok"):
|
||||
return r
|
||||
rows = []
|
||||
for line in r.get("stdout", "").splitlines():
|
||||
if " - " in line:
|
||||
name, desc = line.split(" - ", 1)
|
||||
else:
|
||||
name, desc = line, ""
|
||||
rows.append({"package": name.strip(), "description": desc.strip()})
|
||||
if len(rows) >= limit:
|
||||
break
|
||||
return {"ok": True, "query": query, "results": rows, "limit": limit}
|
||||
|
||||
|
||||
def package_upgradable(args: dict[str, Any]) -> dict[str, Any]:
|
||||
limit = _bounded_int(args.get("limit"), default=100, low=1, high=500, error="PACKAGE_LIMIT_INVALID")
|
||||
if not shutil.which("apt"):
|
||||
return {"ok": False, "error": "APT_NOT_FOUND"}
|
||||
r = _run(["apt", "list", "--upgradable"], timeout=20.0)
|
||||
if not r.get("ok"):
|
||||
return r
|
||||
rows = []
|
||||
for line in r.get("stdout", "").splitlines():
|
||||
line = line.strip()
|
||||
if not line or line.startswith("Listing...") or "/" not in line:
|
||||
continue
|
||||
first, *rest = line.split()
|
||||
name, suite = first.split("/", 1)
|
||||
version = rest[0] if rest else None
|
||||
rows.append({"package": name, "suite": suite, "candidate_version": version})
|
||||
if len(rows) >= limit:
|
||||
break
|
||||
return {"ok": True, "packages": rows, "limit": limit}
|
||||
|
||||
|
||||
def service_list(args: dict[str, Any]) -> dict[str, Any]:
|
||||
limit = _bounded_int(args.get("limit"), default=100, low=1, high=500, error="SERVICE_LIMIT_INVALID")
|
||||
if not shutil.which("systemctl"):
|
||||
return {"ok": False, "error": "SYSTEMCTL_NOT_FOUND"}
|
||||
r = _run([
|
||||
"systemctl", "list-units", "--type=service", "--all",
|
||||
"--no-legend", "--no-pager", "--plain",
|
||||
], timeout=10.0)
|
||||
if not r.get("ok"):
|
||||
return r
|
||||
rows = []
|
||||
for line in r.get("stdout", "").splitlines():
|
||||
parts = line.split(None, 4)
|
||||
if len(parts) < 4:
|
||||
continue
|
||||
rows.append({
|
||||
"unit": parts[0],
|
||||
"load": parts[1],
|
||||
"active": parts[2],
|
||||
"sub": parts[3],
|
||||
"description": parts[4] if len(parts) > 4 else "",
|
||||
})
|
||||
if len(rows) >= limit:
|
||||
break
|
||||
return {"ok": True, "services": rows, "limit": limit}
|
||||
|
||||
|
||||
def journal_service(args: dict[str, Any]) -> dict[str, Any]:
|
||||
service = _simple_name(args.get("service"), "SERVICE_NAME_INVALID")
|
||||
lines = _bounded_int(args.get("lines"), default=80, low=1, high=500, error="JOURNAL_LINES_INVALID")
|
||||
if not shutil.which("journalctl"):
|
||||
return {"ok": False, "error": "JOURNALCTL_NOT_FOUND"}
|
||||
r = _run([
|
||||
"journalctl", "-u", service, "-n", str(lines),
|
||||
"--no-pager", "--output=short-iso-precise",
|
||||
], timeout=10.0)
|
||||
if not r.get("ok"):
|
||||
return r
|
||||
return {
|
||||
"ok": True,
|
||||
"service": service,
|
||||
"lines": lines,
|
||||
"text": _sanitize_text(r.get("stdout", ""))[:MAX_OUTPUT],
|
||||
}
|
||||
|
||||
|
||||
def user_list(args: dict[str, Any]) -> dict[str, Any]:
|
||||
limit = _bounded_int(args.get("limit"), default=200, low=1, high=1000, error="USER_LIMIT_INVALID")
|
||||
r = _run(["getent", "passwd"])
|
||||
if not r.get("ok"):
|
||||
return r
|
||||
rows = []
|
||||
for line in r.get("stdout", "").splitlines():
|
||||
parts = line.split(":")
|
||||
if len(parts) != 7:
|
||||
continue
|
||||
name, _pw, uid, gid, gecos, home, shell = parts
|
||||
rows.append({
|
||||
"name": name,
|
||||
"uid": int(uid) if uid.isdigit() else uid,
|
||||
"gid": int(gid) if gid.isdigit() else gid,
|
||||
"gecos": gecos,
|
||||
"home": home,
|
||||
"shell": shell,
|
||||
})
|
||||
if len(rows) >= limit:
|
||||
break
|
||||
return {"ok": True, "users": rows, "limit": limit}
|
||||
|
||||
|
||||
def group_list(args: dict[str, Any]) -> dict[str, Any]:
|
||||
limit = _bounded_int(args.get("limit"), default=200, low=1, high=1000, error="GROUP_LIMIT_INVALID")
|
||||
r = _run(["getent", "group"])
|
||||
if not r.get("ok"):
|
||||
return r
|
||||
rows = []
|
||||
for line in r.get("stdout", "").splitlines():
|
||||
parts = line.split(":")
|
||||
if len(parts) != 4:
|
||||
continue
|
||||
name, _pw, gid, members = parts
|
||||
rows.append({
|
||||
"name": name,
|
||||
"gid": int(gid) if gid.isdigit() else gid,
|
||||
"members": [x for x in members.split(",") if x],
|
||||
})
|
||||
if len(rows) >= limit:
|
||||
break
|
||||
return {"ok": True, "groups": rows, "limit": limit}
|
||||
|
||||
|
||||
def filesystem_stat(args: dict[str, Any]) -> dict[str, Any]:
|
||||
raw = str(args.get("path") or "").strip()
|
||||
if not raw or len(raw) > 4096 or "\x00" in raw or "\n" in raw:
|
||||
raise ProductError("PATH_INVALID")
|
||||
path = Path(raw)
|
||||
if not path.is_absolute():
|
||||
raise ProductError("PATH_ABSOLUTE_REQUIRED")
|
||||
try:
|
||||
st = path.lstat()
|
||||
except FileNotFoundError:
|
||||
return {"ok": True, "path": raw, "exists": False}
|
||||
return {
|
||||
"ok": True,
|
||||
"path": raw,
|
||||
"exists": True,
|
||||
"is_file": path.is_file(),
|
||||
"is_dir": path.is_dir(),
|
||||
"is_symlink": path.is_symlink(),
|
||||
"size": st.st_size,
|
||||
"mode": oct(st.st_mode & 0o7777),
|
||||
"uid": st.st_uid,
|
||||
"gid": st.st_gid,
|
||||
"mtime_ns": st.st_mtime_ns,
|
||||
}
|
||||
|
||||
|
||||
def filesystem_list(args: dict[str, Any]) -> dict[str, Any]:
|
||||
raw = str(args.get("path") or "").strip()
|
||||
if not raw or len(raw) > 4096 or "\x00" in raw or "\n" in raw:
|
||||
raise ProductError("PATH_INVALID")
|
||||
path = Path(raw)
|
||||
if not path.is_absolute():
|
||||
raise ProductError("PATH_ABSOLUTE_REQUIRED")
|
||||
limit = _bounded_int(args.get("limit"), default=100, low=1, high=500, error="FILE_LIST_LIMIT_INVALID")
|
||||
if not path.is_dir():
|
||||
raise ProductError("DIRECTORY_NOT_FOUND")
|
||||
rows = []
|
||||
try:
|
||||
children = sorted(path.iterdir(), key=lambda x: x.name)
|
||||
except PermissionError:
|
||||
raise ProductError("DIRECTORY_PERMISSION_DENIED") from None
|
||||
for child in children[:limit]:
|
||||
try:
|
||||
st = child.lstat()
|
||||
rows.append({
|
||||
"name": child.name,
|
||||
"is_file": child.is_file(),
|
||||
"is_dir": child.is_dir(),
|
||||
"is_symlink": child.is_symlink(),
|
||||
"size": st.st_size,
|
||||
"mode": oct(st.st_mode & 0o7777),
|
||||
})
|
||||
except OSError:
|
||||
rows.append({"name": child.name, "error": "STAT_FAILED"})
|
||||
return {"ok": True, "path": raw, "entries": rows, "limit": limit}
|
||||
|
||||
|
||||
def filesystem_hash(args: dict[str, Any]) -> dict[str, Any]:
|
||||
raw = str(args.get("path") or "").strip()
|
||||
if not raw or len(raw) > 4096 or "\x00" in raw or "\n" in raw:
|
||||
raise ProductError("PATH_INVALID")
|
||||
path = Path(raw)
|
||||
if not path.is_absolute():
|
||||
raise ProductError("PATH_ABSOLUTE_REQUIRED")
|
||||
if not path.is_file():
|
||||
raise ProductError("PATH_NOT_REGULAR_FILE")
|
||||
size = path.stat().st_size
|
||||
if size > 16 * 1024 * 1024:
|
||||
raise ProductError("FILE_TOO_LARGE_TO_HASH")
|
||||
h = hashlib.sha256()
|
||||
try:
|
||||
with path.open("rb") as f:
|
||||
for chunk in iter(lambda: f.read(65536), b""):
|
||||
h.update(chunk)
|
||||
except PermissionError:
|
||||
raise ProductError("FILE_PERMISSION_DENIED") from None
|
||||
return {"ok": True, "path": raw, "size": size, "sha256": h.hexdigest()}
|
||||
|
||||
|
||||
def filesystem_read_text(args: dict[str, Any]) -> dict[str, Any]:
|
||||
path = _safe_text_path(args.get("path"))
|
||||
max_bytes = _bounded_int(args.get("max_bytes"), default=16384, low=1, high=32768, error="MAX_BYTES_INVALID")
|
||||
try:
|
||||
data = path.read_bytes()
|
||||
except PermissionError:
|
||||
raise ProductError("FILE_PERMISSION_DENIED") from None
|
||||
if len(data) > max_bytes:
|
||||
data = data[:max_bytes]
|
||||
truncated = True
|
||||
else:
|
||||
truncated = False
|
||||
if b"\x00" in data:
|
||||
raise ProductError("BINARY_FILE_NOT_ALLOWED")
|
||||
text = data.decode("utf-8", errors="replace")
|
||||
return {
|
||||
"ok": True,
|
||||
"path": str(path),
|
||||
"bytes_returned": len(data),
|
||||
"truncated": truncated,
|
||||
"text": _sanitize_text(text),
|
||||
}
|
||||
|
||||
|
||||
CAPABILITIES = {
|
||||
"system.status": system_status,
|
||||
"system.resources": system_resources,
|
||||
"system.processes": system_processes,
|
||||
"system.time": system_time,
|
||||
"hardware.summary": hardware_summary,
|
||||
"kernel.modules": kernel_modules,
|
||||
"storage.inspect": storage_inspect,
|
||||
"storage.usage": storage_usage,
|
||||
"storage.mounts": storage_mounts,
|
||||
"network.inspect": network_inspect,
|
||||
"network.sockets": network_sockets,
|
||||
"network.dns": network_dns,
|
||||
"package.status": package_status_read,
|
||||
"package.search": package_search,
|
||||
"package.upgradable": package_upgradable,
|
||||
"service.status": service_status,
|
||||
"service.list": service_list,
|
||||
"journal.service": journal_service,
|
||||
"user.list": user_list,
|
||||
"group.list": group_list,
|
||||
"filesystem.stat": filesystem_stat,
|
||||
"filesystem.list": filesystem_list,
|
||||
"filesystem.hash": filesystem_hash,
|
||||
"filesystem.read_text": filesystem_read_text,
|
||||
}
|
||||
|
||||
|
||||
def list_capabilities() -> list[dict[str, Any]]:
|
||||
return [
|
||||
{"name": "system.status", "risk": "READONLY"},
|
||||
{"name": "system.resources", "risk": "READONLY"},
|
||||
{"name": "system.processes", "risk": "READONLY", "args": {"limit": "int:1..200"}},
|
||||
{"name": "system.time", "risk": "READONLY"},
|
||||
{"name": "hardware.summary", "risk": "READONLY"},
|
||||
{"name": "kernel.modules", "risk": "READONLY", "args": {"limit": "int:1..500"}},
|
||||
{"name": "storage.inspect", "risk": "READONLY"},
|
||||
{"name": "storage.usage", "risk": "READONLY"},
|
||||
{"name": "storage.mounts", "risk": "READONLY"},
|
||||
{"name": "network.inspect", "risk": "READONLY"},
|
||||
{"name": "network.sockets", "risk": "READONLY", "args": {"limit": "int:1..500"}},
|
||||
{"name": "network.dns", "risk": "READONLY"},
|
||||
{"name": "package.status", "risk": "READONLY", "args": {"package": "debian-package-name"}},
|
||||
{"name": "package.search", "risk": "READONLY", "args": {"query": "string", "limit": "int:1..100"}},
|
||||
{"name": "package.upgradable", "risk": "READONLY", "args": {"limit": "int:1..500"}},
|
||||
{"name": "service.status", "risk": "READONLY", "args": {"service": "systemd-unit"}},
|
||||
{"name": "service.list", "risk": "READONLY", "args": {"limit": "int:1..500"}},
|
||||
{"name": "journal.service", "risk": "READONLY", "args": {"service": "systemd-unit", "lines": "int:1..500"}},
|
||||
{"name": "user.list", "risk": "READONLY", "args": {"limit": "int:1..1000"}},
|
||||
{"name": "group.list", "risk": "READONLY", "args": {"limit": "int:1..1000"}},
|
||||
{"name": "filesystem.stat", "risk": "READONLY", "args": {"path": "absolute-path"}},
|
||||
{"name": "filesystem.list", "risk": "READONLY", "args": {"path": "absolute-directory", "limit": "int:1..500"}},
|
||||
{"name": "filesystem.hash", "risk": "READONLY", "args": {"path": "absolute-file<=16MiB"}},
|
||||
{"name": "filesystem.read_text", "risk": "READONLY", "args": {"path": "safe-absolute-text-path", "max_bytes": "int:1..32768"}},
|
||||
]
|
||||
|
||||
|
||||
def execute_capability(name: str, args: dict[str, Any] | None = None) -> dict[str, Any]:
|
||||
name = str(name).strip()
|
||||
fn = CAPABILITIES.get(name)
|
||||
if fn is None:
|
||||
raise ProductError("CAPABILITY_NOT_ALLOWED")
|
||||
args = args or {}
|
||||
if not isinstance(args, dict):
|
||||
raise ProductError("CAPABILITY_ARGS_NOT_OBJECT")
|
||||
return fn(args)
|
||||
+572
@@ -0,0 +1,572 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import hashlib
|
||||
import json
|
||||
import sys
|
||||
import time
|
||||
|
||||
from .product_core import ProductError
|
||||
from .seed_core import SEED_VERSION, SeedContext, atomic_write_json
|
||||
from .classic_tty import execute_frame, render_result, render_snapshot
|
||||
from .console_tty import run_console
|
||||
from .plan_engine import PlanEngine
|
||||
from .memory_store import MemoryStore
|
||||
from .first_run import SetupChoice, detect_system, provider_instructions, render_privilege_explanation, save_setup_config, setup_summary
|
||||
from .console_chat import run_console_chat
|
||||
from .qr_terminal import render_qr_compact, render_qr_mono, render_qr_ansi, render_qr_unicode, TerminalQrError
|
||||
from .pairing import FBPAIR_PROTOCOL, PairingStore, pairing_qr_payload, public_pairing_view
|
||||
from .pairing_client import PairingClientError, poll_pairing, start_pairing
|
||||
from .network_security import inspect_network, pairing_policy, forwarding_policy
|
||||
from .pairing_window import open_pairing_window, pairing_window_status, close_pairing_window
|
||||
from .network_lifecycle import reconcile_once, reconcile_broker_cancel
|
||||
|
||||
|
||||
def emit(value) -> None:
|
||||
print(json.dumps(value, indent=2, sort_keys=True, ensure_ascii=False))
|
||||
|
||||
|
||||
def _sync_pairing_config(ctx: SeedContext, view: dict) -> dict:
|
||||
config=ctx.load_config() or {}
|
||||
pairing=config.get("pairing") or {}
|
||||
pairing.update({
|
||||
"protocol":FBPAIR_PROTOCOL,
|
||||
"state":view.get("state","not_paired"),
|
||||
"transport":"broker_outbound" if view.get("state")!="not_paired" else "unpaired",
|
||||
})
|
||||
for key in (
|
||||
"pair_id","broker_base_url","account_issuer",
|
||||
"account_subject_sha256","revocation_id","paired_at",
|
||||
"broker_confirm_required","broker_confirmed_at",
|
||||
"broker_cancel_required","broker_cancelled_at","broker_cancel_state",
|
||||
):
|
||||
if view.get(key) is not None:
|
||||
pairing[key]=view.get(key)
|
||||
config["pairing"]=pairing
|
||||
config["connection_mode"]="PAIRED" if view.get("state")=="paired" else "PAIRING"
|
||||
setup=config.get("setup")
|
||||
if isinstance(setup,dict):
|
||||
setup["pairing_state"]=view.get("state","not_paired")
|
||||
atomic_write_json(ctx.config_path,config)
|
||||
return config
|
||||
|
||||
|
||||
def _render_pair_qr(view: dict, mode: str) -> None:
|
||||
payload=pairing_qr_payload(view)
|
||||
if not sys.stdout.isatty():
|
||||
return
|
||||
if mode=="compact":
|
||||
render_qr_compact(payload)
|
||||
elif mode=="mono":
|
||||
render_qr_mono(payload)
|
||||
elif mode=="unicode":
|
||||
render_qr_unicode(payload)
|
||||
else:
|
||||
render_qr_ansi(payload)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(
|
||||
prog="framebridge",
|
||||
description="FrameBridge Seed - Linux text-mode bootstrap",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--version",
|
||||
action="version",
|
||||
version=f"%(prog)s {SEED_VERSION}",
|
||||
)
|
||||
sub = parser.add_subparsers(dest="command", required=True)
|
||||
|
||||
init = sub.add_parser("init", help="Initialize node identity and provider")
|
||||
init.add_argument("--provider", default="chatgpt")
|
||||
init.add_argument("--hostname", default="")
|
||||
|
||||
sub.add_parser("identity", help="Show persistent node identity")
|
||||
sub.add_parser("doctor", help="Run READONLY bootstrap diagnostics")
|
||||
sub.add_parser("status", help="Show Seed status")
|
||||
sub.add_parser("fingerprint", help="Show non-secret node fingerprint")
|
||||
sub.add_parser("snapshot", help="Generate Classic AI context snapshot")
|
||||
sub.add_parser("console", help="Open FrameBridge Classic in the current TTY")
|
||||
classic = sub.add_parser("classic", help="Process one Classic frame from stdin or file")
|
||||
classic.add_argument("--file", default="-")
|
||||
setup = sub.add_parser("setup", help="Run the FrameBridge first-run setup wizard")
|
||||
setup.add_argument("--provider", choices=("chatgpt","ollama","lmstudio","other"), default="")
|
||||
setup.add_argument("--privilege-mode", choices=("totp","sudo_tty"), default="")
|
||||
setup.add_argument("--language", default="en")
|
||||
setup.add_argument("--non-interactive", action="store_true")
|
||||
sub.add_parser("setup-status", help="Show first-run setup state")
|
||||
|
||||
net = sub.add_parser("net", help="Inspect FrameBridge network/firewall policy")
|
||||
net_sub = net.add_subparsers(dest="net_command", required=True)
|
||||
net_sub.add_parser("status", help="Show READONLY FBNET interface/firewall status")
|
||||
p_net_policy = net_sub.add_parser("pairing-policy", help="Render pairing-window policy without applying changes")
|
||||
p_net_policy.add_argument("--port", type=int, default=18765)
|
||||
p_net_forward = net_sub.add_parser("forwarding-policy", help="Render typed runtime forwarding policy without applying changes")
|
||||
p_net_forward.add_argument("--in-interface", required=True)
|
||||
p_net_forward.add_argument("--out-interface", required=True)
|
||||
p_net_forward.add_argument("--source-subnet", required=True)
|
||||
p_net_forward.add_argument("--destination-subnet", required=True)
|
||||
p_net_forward.add_argument("--protocol", choices=("any","tcp","udp"), default="any")
|
||||
p_net_forward.add_argument("--destination-port", type=int, default=0)
|
||||
p_net_forward.add_argument("--nat", action="store_true")
|
||||
p_net_forward.add_argument("--lifetime", choices=("runtime",), default="runtime")
|
||||
|
||||
p_window = net_sub.add_parser("pairing-window", help="Manage the short-lived local LAN pairing firewall window")
|
||||
p_window_sub = p_window.add_subparsers(dest="pairing_window_command", required=True)
|
||||
p_window_open = p_window_sub.add_parser("open", help="Open a scoped runtime pairing window after local authorization")
|
||||
p_window_open.add_argument("--port", type=int, default=18765)
|
||||
p_window_open.add_argument("--ttl", type=int, default=600)
|
||||
p_window_sub.add_parser("status", help="Show the current pairing-window transaction")
|
||||
p_window_close = p_window_sub.add_parser("close", help="Close the active pairing window")
|
||||
p_window_close.add_argument("--reason", choices=("user","cancelled","paired","denied","network_change","service_stop","reconcile"), default="user")
|
||||
|
||||
net_sub.add_parser("reconcile", help="Re-evaluate an active pairing window against current network identity")
|
||||
|
||||
pair = sub.add_parser("pair", help="Pair this Linux node with an authorized AI account")
|
||||
pair_sub = pair.add_subparsers(dest="pair_command", required=True)
|
||||
p_pair_start = pair_sub.add_parser("start", help="Start FBPAIR/1 device-style pairing")
|
||||
p_pair_start.add_argument("--broker", required=True)
|
||||
p_pair_start.add_argument("--provider", default="")
|
||||
p_pair_start.add_argument("--mode", choices=("compact","mono","ansi","unicode"), default="compact")
|
||||
p_pair_start.add_argument("--dev-insecure", action="store_true")
|
||||
pair_sub.add_parser("status", help="Show public pairing status")
|
||||
pair_sub.add_parser("cancel", help="Cancel the current pending pairing request")
|
||||
p_pair_poll = pair_sub.add_parser("poll", help="Poll pairing broker once")
|
||||
p_pair_poll.add_argument("--dev-insecure", action="store_true")
|
||||
p_pair_wait = pair_sub.add_parser("wait", help="Wait for human account authorization")
|
||||
p_pair_wait.add_argument("--dev-insecure", action="store_true")
|
||||
p_pair_wait.add_argument("--max-wait", type=int, default=600)
|
||||
p_pair_qr = pair_sub.add_parser("qr", help="Render the current pairing QR")
|
||||
p_pair_qr.add_argument("--mode", choices=("compact","mono","ansi","unicode"), default="compact")
|
||||
|
||||
sub.add_parser("console-chat", help="Open the primary FrameBridge conversational console")
|
||||
|
||||
sub.add_parser("agent-manifest", help="Return the mandatory compact AI bootstrap manifest")
|
||||
memory = sub.add_parser("memory", help="Manage mandatory FrameBridge binary technical memory")
|
||||
memory_sub = memory.add_subparsers(dest="memory_command", required=True)
|
||||
memory_sub.add_parser("init", help="Initialize or repair node binary memory")
|
||||
memory_sub.add_parser("bootstrap", help="Return compact mandatory AI memory bootstrap")
|
||||
p_mem_query = memory_sub.add_parser("query", help="Query a normalized memory domain")
|
||||
p_mem_query.add_argument("--domain", default="summary")
|
||||
p_mem_query.add_argument("--limit", type=int, default=20)
|
||||
p_mem_append = memory_sub.add_parser("append", help="Append one normalized memory record")
|
||||
p_mem_append.add_argument("--record-type", required=True, choices=("change","incident","decision","artifact","recovery","certification","note"))
|
||||
p_mem_append.add_argument("--json-stdin", action="store_true", required=True)
|
||||
memory_sub.add_parser("checkpoint", help="Refresh the compact current-state checkpoint")
|
||||
memory_sub.add_parser("stats", help="Show binary memory statistics")
|
||||
|
||||
plan = sub.add_parser("plan", help="Manage local CHANGE plans")
|
||||
plan_sub = plan.add_subparsers(dest="plan_command", required=True)
|
||||
p_import = plan_sub.add_parser("import")
|
||||
p_import.add_argument("--file", required=True)
|
||||
plan_sub.add_parser("import-stdin")
|
||||
p_show = plan_sub.add_parser("show")
|
||||
p_show.add_argument("plan_id")
|
||||
p_validate = plan_sub.add_parser("validate")
|
||||
p_validate.add_argument("plan_id")
|
||||
p_approve = plan_sub.add_parser("approve")
|
||||
p_approve.add_argument("plan_id")
|
||||
p_approve_system = plan_sub.add_parser("approve-system")
|
||||
p_approve_system.add_argument("plan_id")
|
||||
p_apply = plan_sub.add_parser("apply")
|
||||
p_apply.add_argument("plan_id")
|
||||
p_approve_rollback = plan_sub.add_parser("approve-rollback")
|
||||
p_approve_rollback.add_argument("plan_id")
|
||||
p_remote_challenge = plan_sub.add_parser("remote-challenge")
|
||||
p_remote_challenge.add_argument("plan_id")
|
||||
p_remote_challenge.add_argument("--rollback", action="store_true")
|
||||
p_remote_challenge.add_argument("--ttl", type=int, default=300)
|
||||
p_remote_approve = plan_sub.add_parser("remote-approve")
|
||||
p_remote_approve.add_argument("plan_id")
|
||||
p_remote_approve.add_argument("--challenge-id", required=True)
|
||||
p_remote_approve.add_argument("--key-id", required=True)
|
||||
p_remote_approve.add_argument("--signature-b64", required=True)
|
||||
p_remote_approve.add_argument("--rollback", action="store_true")
|
||||
p_totp_enroll = plan_sub.add_parser("totp-enroll")
|
||||
p_totp_enroll.add_argument("--key-id", required=True)
|
||||
p_totp_enroll.add_argument("--node-label", required=True)
|
||||
p_totp_enroll.add_argument("--issuer", default="FrameBridge")
|
||||
p_totp_enroll.add_argument("--mode", choices=("compact","mono","ansi","unicode"), default="compact")
|
||||
p_totp_challenge = plan_sub.add_parser("totp-challenge")
|
||||
p_totp_challenge.add_argument("plan_id")
|
||||
p_totp_challenge.add_argument("--key-id", required=True)
|
||||
p_totp_challenge.add_argument("--rollback", action="store_true")
|
||||
p_totp_challenge.add_argument("--ttl", type=int, default=120)
|
||||
p_totp_approve = plan_sub.add_parser("totp-approve")
|
||||
p_totp_approve.add_argument("plan_id")
|
||||
p_totp_approve.add_argument("--challenge-id", required=True)
|
||||
p_totp_approve.add_argument("--key-id", required=True)
|
||||
p_totp_approve.add_argument("--rollback", action="store_true")
|
||||
p_rollback = plan_sub.add_parser("rollback")
|
||||
p_rollback.add_argument("plan_id")
|
||||
plan_sub.add_parser("list")
|
||||
sub.add_parser("audit", help="Verify the audit chain")
|
||||
|
||||
args = parser.parse_args()
|
||||
ctx = SeedContext()
|
||||
|
||||
try:
|
||||
if args.command == "init":
|
||||
emit(ctx.init(args.provider, args.hostname))
|
||||
elif args.command == "identity":
|
||||
emit(ctx.identity())
|
||||
elif args.command == "doctor":
|
||||
emit(ctx.doctor())
|
||||
elif args.command == "status":
|
||||
emit(ctx.status())
|
||||
elif args.command == "fingerprint":
|
||||
emit(ctx.fingerprint())
|
||||
elif args.command == "snapshot":
|
||||
print(render_snapshot(ctx))
|
||||
elif args.command == "setup":
|
||||
if not (ctx.load_identity() and ctx.load_config()):
|
||||
emit(ctx.init("chatgpt", ""))
|
||||
detected = detect_system()
|
||||
provider = args.provider
|
||||
privilege = args.privilege_mode
|
||||
language = args.language or "en"
|
||||
if args.non_interactive:
|
||||
if not provider or not privilege:
|
||||
raise ProductError("SETUP_NON_INTERACTIVE_OPTIONS_REQUIRED")
|
||||
else:
|
||||
print("\nFrameBridge First-Run Setup\n")
|
||||
print("System detected:")
|
||||
sysinfo = detected.get("system") or {}
|
||||
hw = detected.get("hardware") or {}
|
||||
print(f" OS: {sysinfo.get('os_name')} / kernel {sysinfo.get('kernel')}")
|
||||
print(f" CPU: {hw.get('cpu_model')}")
|
||||
print(f" RAM: {hw.get('memory_total_kib')} KiB")
|
||||
print("\nPrivileged authorization:")
|
||||
print(" 1) Authenticator / QR TOTP")
|
||||
print(" 2) Local sudo password")
|
||||
if not privilege:
|
||||
answer=input("Select [1/2]: ").strip()
|
||||
privilege="totp" if answer=="1" else "sudo_tty" if answer=="2" else ""
|
||||
print("\nAI provider:")
|
||||
print(" 1) ChatGPT")
|
||||
print(" 2) Ollama")
|
||||
print(" 3) LM Studio")
|
||||
print(" 4) Other")
|
||||
if not provider:
|
||||
answer=input("Select [1-4]: ").strip()
|
||||
provider={"1":"chatgpt","2":"ollama","3":"lmstudio","4":"other"}.get(answer,"")
|
||||
choice=SetupChoice(privilege,provider,language)
|
||||
config=save_setup_config(ctx,choice)
|
||||
mem=MemoryStore(ctx).initialize()
|
||||
print("\n"+"\n".join(render_privilege_explanation(choice.privilege_mode)))
|
||||
print("\n"+"\n".join(provider_instructions(choice.provider)))
|
||||
print("\nFrameBridge setup complete.")
|
||||
print("Run: framebridge console-chat")
|
||||
emit({"ok":True,"setup":config.get("setup"),"memory":mem,"detected":detected})
|
||||
elif args.command == "setup-status":
|
||||
emit(setup_summary(ctx))
|
||||
elif args.command == "net":
|
||||
if args.net_command == "status":
|
||||
emit(inspect_network())
|
||||
elif args.net_command == "pairing-policy":
|
||||
emit(pairing_policy(args.port))
|
||||
elif args.net_command == "forwarding-policy":
|
||||
emit(forwarding_policy(
|
||||
ingress_interface=args.in_interface,
|
||||
egress_interface=args.out_interface,
|
||||
source_subnet=args.source_subnet,
|
||||
destination_subnet=args.destination_subnet,
|
||||
protocol=args.protocol,
|
||||
destination_port=args.destination_port,
|
||||
nat=args.nat,
|
||||
lifetime=args.lifetime,
|
||||
))
|
||||
elif args.net_command == "pairing-window":
|
||||
if args.pairing_window_command == "open":
|
||||
MemoryStore(ctx).require_ready()
|
||||
result=open_pairing_window(args.port,args.ttl)
|
||||
ctx.audit.append(
|
||||
"network.pairing_window.open","PASS",
|
||||
request={"port":args.port,"ttl_seconds":args.ttl},
|
||||
result={
|
||||
"window_id":result.get("window_id"),
|
||||
"interface":result.get("interface"),
|
||||
"source_subnet":result.get("source_subnet"),
|
||||
"destination_address":result.get("destination_address"),
|
||||
"expires_at":result.get("expires_at"),
|
||||
},
|
||||
principal="local-operator",
|
||||
)
|
||||
MemoryStore(ctx).append(
|
||||
"change",
|
||||
{
|
||||
"operation":"network.pairing_window.open",
|
||||
"window_id":result.get("window_id"),
|
||||
"interface":result.get("interface"),
|
||||
"source_subnet":result.get("source_subnet"),
|
||||
"destination_address":result.get("destination_address"),
|
||||
"port":result.get("port"),
|
||||
"expires_at":result.get("expires_at"),
|
||||
},
|
||||
source="operator",
|
||||
)
|
||||
emit(result)
|
||||
elif args.pairing_window_command == "status":
|
||||
emit(pairing_window_status())
|
||||
elif args.pairing_window_command == "close":
|
||||
MemoryStore(ctx).require_ready()
|
||||
before=pairing_window_status()
|
||||
result=close_pairing_window(args.reason)
|
||||
ctx.audit.append(
|
||||
"network.pairing_window.close","PASS",
|
||||
request={"reason":args.reason,"window_id":before.get("window_id")},
|
||||
result={"state":result.get("state"),"rule_present":result.get("rule_present")},
|
||||
principal="local-operator",
|
||||
)
|
||||
MemoryStore(ctx).append(
|
||||
"recovery",
|
||||
{
|
||||
"operation":"network.pairing_window.close",
|
||||
"reason":args.reason,
|
||||
"window_id":before.get("window_id"),
|
||||
"state":result.get("state"),
|
||||
},
|
||||
source="operator",
|
||||
)
|
||||
emit(result)
|
||||
elif args.net_command == "reconcile":
|
||||
emit(reconcile_once(ctx=ctx))
|
||||
elif args.command == "pair":
|
||||
if not (ctx.load_identity() and ctx.load_config()):
|
||||
raise ProductError("PAIR_NODE_NOT_INITIALIZED")
|
||||
pair_store=PairingStore(ctx.state_dir)
|
||||
if args.pair_command == "status":
|
||||
emit(public_pairing_view(pair_store.load()))
|
||||
elif args.pair_command == "start":
|
||||
MemoryStore(ctx).require_ready()
|
||||
identity=ctx.load_identity() or {}
|
||||
config=ctx.load_config() or {}
|
||||
manifest_path=ctx.state_dir/"memory"/"manifest.fbam"
|
||||
if not manifest_path.is_file():
|
||||
raise ProductError("PAIR_CAPABILITY_MANIFEST_MISSING")
|
||||
manifest_sha=hashlib.sha256(manifest_path.read_bytes()).hexdigest()
|
||||
provider=(args.provider or config.get("provider") or "chatgpt").strip().lower()
|
||||
try:
|
||||
view=start_pairing(
|
||||
state_dir=ctx.state_dir,
|
||||
node_id=str(identity.get("node_id") or ""),
|
||||
node_label=str(identity.get("hostname") or ""),
|
||||
node_host_id=str(identity.get("node_host_id") or ""),
|
||||
public_jwk=identity.get("node_public_jwk") or {},
|
||||
provider=provider,
|
||||
capability_manifest_sha256=manifest_sha,
|
||||
broker_base_url=args.broker,
|
||||
allow_insecure_dev=args.dev_insecure,
|
||||
)
|
||||
except PairingClientError as exc:
|
||||
raise ProductError(str(exc)) from exc
|
||||
_sync_pairing_config(ctx,view)
|
||||
ctx.audit.append(
|
||||
"pair.request","PASS",
|
||||
request={
|
||||
"protocol":FBPAIR_PROTOCOL,
|
||||
"pair_id":view.get("pair_id"),
|
||||
"provider":provider,
|
||||
"node_host_id":identity.get("node_host_id"),
|
||||
},
|
||||
result={"state":view.get("state")},
|
||||
principal="local-operator",
|
||||
)
|
||||
print("\nFrameBridge AI pairing")
|
||||
print("Open: "+str(view.get("verification_uri_complete") or view.get("verification_uri") or ""))
|
||||
print("Code: "+str(view.get("user_code")))
|
||||
print("Expires in: "+str(view.get("expires_in"))+" seconds")
|
||||
_render_pair_qr(view,args.mode)
|
||||
emit(view)
|
||||
elif args.pair_command == "cancel":
|
||||
try:
|
||||
view=pair_store.cancel_pending()
|
||||
except Exception as exc:
|
||||
raise ProductError(str(exc)) from exc
|
||||
_sync_pairing_config(ctx,view)
|
||||
broker_cancel=reconcile_broker_cancel(ctx)
|
||||
view=public_pairing_view(pair_store.load())
|
||||
_sync_pairing_config(ctx,view)
|
||||
ctx.audit.append(
|
||||
"pair.cancel","PASS",
|
||||
request={
|
||||
"protocol":FBPAIR_PROTOCOL,
|
||||
"pair_id":view.get("pair_id"),
|
||||
},
|
||||
result={"state":view.get("state"),"broker_cancel_state":broker_cancel.get("state")},
|
||||
principal="local-operator",
|
||||
)
|
||||
emit(view)
|
||||
elif args.pair_command == "poll":
|
||||
before=pair_store.load()
|
||||
try:
|
||||
view=poll_pairing(
|
||||
state_dir=ctx.state_dir,
|
||||
allow_insecure_dev=args.dev_insecure,
|
||||
)
|
||||
except PairingClientError as exc:
|
||||
raise ProductError(str(exc)) from exc
|
||||
_sync_pairing_config(ctx,view)
|
||||
if view.get("state") != before.get("state"):
|
||||
ctx.audit.append(
|
||||
"pair."+str(view.get("state")),"PASS",
|
||||
request={
|
||||
"protocol":FBPAIR_PROTOCOL,
|
||||
"pair_id":view.get("pair_id"),
|
||||
},
|
||||
result={
|
||||
"state":view.get("state"),
|
||||
"provider":view.get("provider"),
|
||||
},
|
||||
principal="pairing-broker",
|
||||
)
|
||||
if view.get("state")=="paired":
|
||||
MemoryStore(ctx).checkpoint()
|
||||
emit(view)
|
||||
elif args.pair_command == "wait":
|
||||
max_wait=int(args.max_wait)
|
||||
if max_wait<1 or max_wait>1800:
|
||||
raise ProductError("PAIR_WAIT_TIMEOUT_INVALID")
|
||||
started=time.monotonic()
|
||||
view=public_pairing_view(pair_store.load())
|
||||
while (
|
||||
view.get("state")=="pair_requested"
|
||||
or (view.get("state")=="paired" and view.get("broker_confirm_required") is True)
|
||||
):
|
||||
if time.monotonic()-started>=max_wait:
|
||||
raise ProductError("PAIR_WAIT_TIMEOUT")
|
||||
try:
|
||||
view=poll_pairing(
|
||||
state_dir=ctx.state_dir,
|
||||
allow_insecure_dev=args.dev_insecure,
|
||||
)
|
||||
except PairingClientError as exc:
|
||||
raise ProductError(str(exc)) from exc
|
||||
if view.get("state")=="paired" and view.get("broker_confirm_required") is not True:
|
||||
break
|
||||
if view.get("state") not in {"pair_requested","paired"}:
|
||||
break
|
||||
retry=int(view.get("retry_after") or view.get("interval") or 2)
|
||||
time.sleep(max(1,min(60,retry)))
|
||||
_sync_pairing_config(ctx,view)
|
||||
ctx.audit.append(
|
||||
"pair."+str(view.get("state")),"PASS",
|
||||
request={
|
||||
"protocol":FBPAIR_PROTOCOL,
|
||||
"pair_id":view.get("pair_id"),
|
||||
},
|
||||
result={
|
||||
"state":view.get("state"),
|
||||
"provider":view.get("provider"),
|
||||
},
|
||||
principal="pairing-broker",
|
||||
)
|
||||
if view.get("state")=="paired":
|
||||
MemoryStore(ctx).checkpoint()
|
||||
emit(view)
|
||||
if view.get("state")!="paired" or view.get("broker_confirm_required") is True:
|
||||
raise SystemExit(3)
|
||||
elif args.pair_command == "qr":
|
||||
view=public_pairing_view(pair_store.load())
|
||||
_render_pair_qr(view,args.mode)
|
||||
emit(view)
|
||||
elif args.command == "console-chat":
|
||||
raise SystemExit(run_console_chat(ctx))
|
||||
elif args.command == "agent-manifest":
|
||||
emit(MemoryStore(ctx).bootstrap())
|
||||
elif args.command == "memory":
|
||||
store = MemoryStore(ctx)
|
||||
if args.memory_command == "init":
|
||||
value = store.initialize()
|
||||
ctx.audit.append("memory.init", "PASS", request={"protocol":"FBM/1"}, result={"required":True,"mode":"read_write"}, principal="local-operator")
|
||||
emit(value)
|
||||
elif args.memory_command == "bootstrap":
|
||||
emit(store.bootstrap())
|
||||
elif args.memory_command == "query":
|
||||
emit(store.query(args.domain, args.limit))
|
||||
elif args.memory_command == "append":
|
||||
try:
|
||||
value = json.load(sys.stdin)
|
||||
except Exception:
|
||||
raise ProductError("MEMORY_APPEND_JSON_INVALID")
|
||||
if not isinstance(value, dict):
|
||||
raise ProductError("MEMORY_APPEND_JSON_NOT_OBJECT")
|
||||
result = store.append(args.record_type, value, source="agent")
|
||||
ctx.audit.append("memory.append", "PASS", request={"record_type":args.record_type}, result={"event_number":result.get("event_number"),"event_sha256":result.get("event_sha256")}, principal="ai-agent")
|
||||
store.checkpoint()
|
||||
emit(result)
|
||||
elif args.memory_command == "checkpoint":
|
||||
value = store.checkpoint()
|
||||
ctx.audit.append("memory.checkpoint", "PASS", request={}, result={"state_sha256":value.get("state_sha256"),"manifest_sha256":value.get("manifest_sha256")}, principal="local-operator")
|
||||
emit(value)
|
||||
elif args.memory_command == "stats":
|
||||
emit(store.stats())
|
||||
elif args.command == "console":
|
||||
raise SystemExit(run_console(ctx))
|
||||
elif args.command == "classic":
|
||||
if args.file == "-":
|
||||
text = sys.stdin.read()
|
||||
else:
|
||||
with open(args.file, "r", encoding="utf-8") as handle:
|
||||
text = handle.read()
|
||||
print(render_result(execute_frame(ctx, text)))
|
||||
elif args.command == "plan":
|
||||
engine = PlanEngine(ctx)
|
||||
if args.plan_command == "import":
|
||||
emit(engine.import_file(args.file))
|
||||
elif args.plan_command == "import-stdin":
|
||||
try:
|
||||
value = json.load(sys.stdin)
|
||||
except Exception:
|
||||
raise ProductError("PLAN_IMPORT_JSON_INVALID")
|
||||
emit(engine.import_document(value))
|
||||
elif args.plan_command == "show":
|
||||
emit(engine.show(args.plan_id))
|
||||
elif args.plan_command == "validate":
|
||||
emit(engine.validate(args.plan_id))
|
||||
elif args.plan_command == "approve":
|
||||
emit(engine.approve(args.plan_id))
|
||||
elif args.plan_command == "approve-system":
|
||||
emit(engine.approve_system(args.plan_id))
|
||||
elif args.plan_command == "apply":
|
||||
emit(engine.apply(args.plan_id))
|
||||
elif args.plan_command == "approve-rollback":
|
||||
emit(engine.approve_rollback(args.plan_id))
|
||||
elif args.plan_command == "remote-challenge":
|
||||
emit(engine.remote_challenge(args.plan_id, rollback=args.rollback, ttl_seconds=args.ttl))
|
||||
elif args.plan_command == "remote-approve":
|
||||
emit(engine.approve_remote(args.plan_id, args.challenge_id, args.key_id, args.signature_b64, rollback=args.rollback))
|
||||
elif args.plan_command == "totp-enroll":
|
||||
if not sys.stdin.isatty() or not sys.stdout.isatty():
|
||||
raise ProductError("INTERACTIVE_TTY_REQUIRED")
|
||||
value=engine.totp_enroll(args.key_id,args.node_label,args.issuer)
|
||||
print("\nFrameBridge TOTP enrollment")
|
||||
print("Scan this QR with your authenticator app.\n")
|
||||
if args.mode == "compact": render_qr_compact(value["otpauth_uri"])
|
||||
elif args.mode == "mono": render_qr_mono(value["otpauth_uri"])
|
||||
elif args.mode == "unicode": render_qr_unicode(value["otpauth_uri"])
|
||||
else: render_qr_ansi(value["otpauth_uri"])
|
||||
print("Key ID: "+value["key_id"])
|
||||
print("Enrollment complete. Do not share the secret or a QR screenshot.\n")
|
||||
emit({"ok":True,"key_id":value["key_id"],"secret_sha256":value["secret_sha256"],"enrolled_at":value["enrolled_at"]})
|
||||
elif args.plan_command == "totp-challenge":
|
||||
emit(engine.totp_challenge(args.plan_id, args.key_id, rollback=args.rollback, ttl_seconds=args.ttl))
|
||||
elif args.plan_command == "totp-approve":
|
||||
code=sys.stdin.readline().strip()
|
||||
emit(engine.approve_totp(args.plan_id, args.challenge_id, args.key_id, code, rollback=args.rollback))
|
||||
elif args.plan_command == "rollback":
|
||||
emit(engine.rollback(args.plan_id))
|
||||
elif args.plan_command == "list":
|
||||
emit(engine.list())
|
||||
elif args.command == "audit":
|
||||
emit(ctx.audit.verify())
|
||||
except ProductError as exc:
|
||||
emit({
|
||||
"ok": False,
|
||||
"error": str(exc),
|
||||
})
|
||||
raise SystemExit(2)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
+341
@@ -0,0 +1,341 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
import platform
|
||||
import shutil
|
||||
import socket
|
||||
import ssl
|
||||
import subprocess
|
||||
import uuid
|
||||
from datetime import datetime, timezone
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from .product_core import AuditLog, ProductError
|
||||
from .node_identity import NodeIdentityError, ensure_node_keypair
|
||||
from .pairing import FBPAIR_PROTOCOL
|
||||
|
||||
SUPPORTED_PROVIDERS = {"chatgpt"}
|
||||
SEED_VERSION = "0.5.0-candidate.1"
|
||||
|
||||
|
||||
def utc_now() -> str:
|
||||
return datetime.now(timezone.utc).isoformat()
|
||||
|
||||
|
||||
def atomic_write_json(path: Path, value: dict[str, Any]) -> None:
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
tmp = path.with_name("." + path.name + ".tmp")
|
||||
data = json.dumps(value, indent=2, sort_keys=True, ensure_ascii=True) + "\n"
|
||||
tmp.write_text(data, encoding="utf-8")
|
||||
os.chmod(tmp, 0o600)
|
||||
os.replace(tmp, path)
|
||||
|
||||
|
||||
def read_os_release() -> dict[str, str]:
|
||||
path = Path("/etc/os-release")
|
||||
result: dict[str, str] = {}
|
||||
if not path.is_file():
|
||||
return result
|
||||
for raw in path.read_text(encoding="utf-8", errors="replace").splitlines():
|
||||
if "=" not in raw:
|
||||
continue
|
||||
key, value = raw.split("=", 1)
|
||||
result[key] = value.strip().strip('"')
|
||||
return result
|
||||
|
||||
|
||||
def run_readonly(args: list[str], timeout: float = 3.0) -> dict[str, Any]:
|
||||
try:
|
||||
proc = subprocess.run(
|
||||
args,
|
||||
text=True,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.PIPE,
|
||||
timeout=timeout,
|
||||
check=False,
|
||||
)
|
||||
return {
|
||||
"ok": proc.returncode == 0,
|
||||
"rc": proc.returncode,
|
||||
"stdout": proc.stdout.strip()[:8192],
|
||||
"stderr": proc.stderr.strip()[:2048],
|
||||
}
|
||||
except Exception as exc:
|
||||
return {
|
||||
"ok": False,
|
||||
"error": type(exc).__name__,
|
||||
}
|
||||
|
||||
|
||||
class SeedContext:
|
||||
def __init__(
|
||||
self,
|
||||
state_dir: Path | None = None,
|
||||
config_dir: Path | None = None,
|
||||
):
|
||||
self.state_dir = Path(
|
||||
state_dir
|
||||
or os.environ.get(
|
||||
"FRAMEBRIDGE_SEED_STATE_DIR",
|
||||
"~/.local/share/framebridge-seed",
|
||||
)
|
||||
).expanduser()
|
||||
self.config_dir = Path(
|
||||
config_dir
|
||||
or os.environ.get(
|
||||
"FRAMEBRIDGE_SEED_CONFIG_DIR",
|
||||
"~/.config/framebridge-seed",
|
||||
)
|
||||
).expanduser()
|
||||
|
||||
self.identity_path = self.state_dir / "identity.json"
|
||||
self.config_path = self.config_dir / "config.json"
|
||||
self.audit = AuditLog(self.state_dir / "audit.jsonl")
|
||||
|
||||
def detect_hostname(self) -> str:
|
||||
value = socket.gethostname().strip()
|
||||
return value or "linux-node"
|
||||
|
||||
def load_identity(self) -> dict[str, Any] | None:
|
||||
if not self.identity_path.is_file():
|
||||
return None
|
||||
value = json.loads(self.identity_path.read_text(encoding="utf-8"))
|
||||
if not isinstance(value, dict):
|
||||
raise ProductError("SEED_IDENTITY_INVALID")
|
||||
return value
|
||||
|
||||
def load_config(self) -> dict[str, Any] | None:
|
||||
if not self.config_path.is_file():
|
||||
return None
|
||||
value = json.loads(self.config_path.read_text(encoding="utf-8"))
|
||||
if not isinstance(value, dict):
|
||||
raise ProductError("SEED_CONFIG_INVALID")
|
||||
return value
|
||||
|
||||
def init(
|
||||
self,
|
||||
provider: str = "chatgpt",
|
||||
hostname: str = "",
|
||||
) -> dict[str, Any]:
|
||||
provider = str(provider).strip().lower()
|
||||
if provider not in SUPPORTED_PROVIDERS:
|
||||
raise ProductError("SEED_PROVIDER_UNSUPPORTED")
|
||||
|
||||
existing = self.load_identity()
|
||||
host = str(hostname).strip() or self.detect_hostname()
|
||||
|
||||
if existing:
|
||||
node_id = str(existing.get("node_id") or "")
|
||||
if not node_id:
|
||||
raise ProductError("SEED_IDENTITY_NODE_ID_MISSING")
|
||||
created_at = str(existing.get("created_at") or utc_now())
|
||||
else:
|
||||
node_id = str(uuid.uuid4())
|
||||
created_at = utc_now()
|
||||
|
||||
try:
|
||||
crypto_identity = ensure_node_keypair(
|
||||
self.state_dir,
|
||||
comment="framebridge-" + node_id.replace("-", "")[:8],
|
||||
)
|
||||
except NodeIdentityError as exc:
|
||||
raise ProductError(str(exc)) from exc
|
||||
|
||||
identity = {
|
||||
"schema": 2,
|
||||
"node_id": node_id,
|
||||
"node_id_short": node_id.replace("-", "")[:8],
|
||||
"node_host_id": crypto_identity["host_id"],
|
||||
"node_public_jwk": crypto_identity["public_jwk"],
|
||||
"node_public_key_sha256": crypto_identity["public_key_sha256"],
|
||||
"hostname": host,
|
||||
"created_at": created_at,
|
||||
"updated_at": utc_now(),
|
||||
}
|
||||
|
||||
config = {
|
||||
"schema": 1,
|
||||
"seed_version": SEED_VERSION,
|
||||
"provider": provider,
|
||||
"connection_name": f"FrameBridge {host}",
|
||||
"pairing": {
|
||||
"protocol": FBPAIR_PROTOCOL,
|
||||
"state": "not_paired",
|
||||
"transport": "unpaired",
|
||||
},
|
||||
"connection_mode": "PAIRING",
|
||||
"capability_profile": "seed-readonly",
|
||||
"memory": {
|
||||
"required": True,
|
||||
"protocol": "FBM/1",
|
||||
"manifest_protocol": "FBAM/1",
|
||||
"mode": "read_write",
|
||||
"authority": os.environ.get("FRAMEBRIDGE_MEMORY_AUTHORITY", "deployment-defined"),
|
||||
},
|
||||
"hostname": host,
|
||||
"node_id": node_id,
|
||||
"updated_at": utc_now(),
|
||||
}
|
||||
|
||||
atomic_write_json(self.identity_path, identity)
|
||||
atomic_write_json(self.config_path, config)
|
||||
|
||||
self.audit.append(
|
||||
"seed.init",
|
||||
"PASS",
|
||||
request={
|
||||
"provider": provider,
|
||||
"hostname": host,
|
||||
},
|
||||
result={
|
||||
"node_id_short": identity["node_id_short"],
|
||||
"pairing_state": "not_paired",
|
||||
"capability_profile": "seed-readonly",
|
||||
"memory_required": True,
|
||||
},
|
||||
principal="local-operator",
|
||||
)
|
||||
|
||||
from .memory_store import MemoryStore
|
||||
memory = MemoryStore(self).initialize()
|
||||
|
||||
return {
|
||||
"ok": True,
|
||||
"seed_version": SEED_VERSION,
|
||||
"identity": identity,
|
||||
"config": config,
|
||||
"memory": memory,
|
||||
}
|
||||
|
||||
def identity(self) -> dict[str, Any]:
|
||||
identity = self.load_identity()
|
||||
if identity is None:
|
||||
return {
|
||||
"ok": False,
|
||||
"error": "SEED_NOT_INITIALIZED",
|
||||
}
|
||||
return {
|
||||
"ok": True,
|
||||
"identity": identity,
|
||||
}
|
||||
|
||||
def doctor(self) -> dict[str, Any]:
|
||||
os_release = read_os_release()
|
||||
route = (
|
||||
run_readonly(["ip", "route", "show", "default"])
|
||||
if shutil.which("ip")
|
||||
else {"ok": False, "error": "IP_COMMAND_MISSING"}
|
||||
)
|
||||
dns_path = Path("/etc/resolv.conf")
|
||||
dns_present = dns_path.exists() and dns_path.stat().st_size > 0
|
||||
ca = ssl.get_default_verify_paths()
|
||||
ca_present = bool(
|
||||
(ca.cafile and Path(ca.cafile).is_file())
|
||||
or (ca.capath and Path(ca.capath).is_dir())
|
||||
)
|
||||
|
||||
checks = {
|
||||
"kernel": platform.release(),
|
||||
"architecture": platform.machine(),
|
||||
"python": platform.python_version(),
|
||||
"hostname": self.detect_hostname(),
|
||||
"os_id": os_release.get("ID"),
|
||||
"os_version": os_release.get("VERSION_ID"),
|
||||
"shell": os.environ.get("SHELL"),
|
||||
"network_default_route": bool(route.get("ok") and route.get("stdout")),
|
||||
"dns_config_present": dns_present,
|
||||
"ca_certificates_present": ca_present,
|
||||
"systemd_present": shutil.which("systemctl") is not None,
|
||||
"ip_command_present": shutil.which("ip") is not None,
|
||||
"curl_present": shutil.which("curl") is not None,
|
||||
"python_venv_available": self._venv_available(),
|
||||
}
|
||||
|
||||
required = [
|
||||
checks["network_default_route"],
|
||||
checks["dns_config_present"],
|
||||
checks["ca_certificates_present"],
|
||||
checks["ip_command_present"],
|
||||
checks["python_venv_available"],
|
||||
]
|
||||
|
||||
return {
|
||||
"ok": all(bool(x) for x in required),
|
||||
"seed_version": SEED_VERSION,
|
||||
"checks": checks,
|
||||
"pairing_ready": all(bool(x) for x in required),
|
||||
"readonly_only": True,
|
||||
}
|
||||
|
||||
def _venv_available(self) -> bool:
|
||||
try:
|
||||
import venv # noqa: F401
|
||||
return True
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
def status(self) -> dict[str, Any]:
|
||||
identity = self.load_identity()
|
||||
config = self.load_config()
|
||||
doctor = self.doctor()
|
||||
|
||||
if identity is None or config is None:
|
||||
return {
|
||||
"ok": False,
|
||||
"initialized": False,
|
||||
"seed_version": SEED_VERSION,
|
||||
"doctor": doctor,
|
||||
}
|
||||
|
||||
from .memory_store import MemoryStore
|
||||
memory_required = bool((config.get("memory") or {}).get("required"))
|
||||
try:
|
||||
memory = MemoryStore(self).verify() if memory_required else {"ok": True, "required": False}
|
||||
except Exception as exc:
|
||||
memory = {"ok": False, "required": memory_required, "error": str(exc)}
|
||||
|
||||
return {
|
||||
"ok": bool(memory.get("ok")) if memory_required else True,
|
||||
"initialized": True,
|
||||
"seed_version": SEED_VERSION,
|
||||
"node": {
|
||||
"hostname": identity.get("hostname"),
|
||||
"node_id": identity.get("node_id"),
|
||||
"node_id_short": identity.get("node_id_short"),
|
||||
},
|
||||
"provider": config.get("provider"),
|
||||
"connection_name": config.get("connection_name"),
|
||||
"pairing": config.get("pairing"),
|
||||
"capability_profile": config.get("capability_profile"),
|
||||
"memory": memory,
|
||||
"doctor": doctor,
|
||||
"audit": self.audit.verify(),
|
||||
}
|
||||
|
||||
def fingerprint(self) -> dict[str, Any]:
|
||||
identity = self.load_identity()
|
||||
config = self.load_config()
|
||||
if identity is None or config is None:
|
||||
raise ProductError("SEED_NOT_INITIALIZED")
|
||||
|
||||
body = json.dumps(
|
||||
{
|
||||
"node_id": identity.get("node_id"),
|
||||
"hostname": identity.get("hostname"),
|
||||
"provider": config.get("provider"),
|
||||
"connection_name": config.get("connection_name"),
|
||||
},
|
||||
sort_keys=True,
|
||||
separators=(",", ":"),
|
||||
).encode("utf-8")
|
||||
|
||||
return {
|
||||
"ok": True,
|
||||
"sha256": hashlib.sha256(body).hexdigest(),
|
||||
"node_id_short": identity.get("node_id_short"),
|
||||
"hostname": identity.get("hostname"),
|
||||
"provider": config.get("provider"),
|
||||
}
|
||||
+41
@@ -0,0 +1,41 @@
|
||||
from __future__ import annotations
|
||||
import base64,hashlib,hmac,os,struct,time
|
||||
from pathlib import Path
|
||||
|
||||
class TotpError(RuntimeError): pass
|
||||
|
||||
def b32decode_nopad(value:str)->bytes:
|
||||
s=''.join(str(value).strip().upper().split())
|
||||
if not s or any(c not in 'ABCDEFGHIJKLMNOPQRSTUVWXYZ234567' for c in s): raise TotpError('TOTP_SECRET_INVALID')
|
||||
s += '='*((8-len(s)%8)%8)
|
||||
try:return base64.b32decode(s,casefold=True)
|
||||
except Exception: raise TotpError('TOTP_SECRET_INVALID') from None
|
||||
|
||||
def hotp(secret:bytes,counter:int,digits:int=6)->str:
|
||||
digest=hmac.new(secret,struct.pack('>Q',int(counter)),hashlib.sha1).digest()
|
||||
off=digest[-1]&0x0f
|
||||
code=((digest[off]&0x7f)<<24)|((digest[off+1]&0xff)<<16)|((digest[off+2]&0xff)<<8)|(digest[off+3]&0xff)
|
||||
return str(code%(10**digits)).zfill(digits)
|
||||
|
||||
def verify_code(secret:bytes,code:str,when:float|None=None,window:int=1,period:int=30)->int|None:
|
||||
if not isinstance(code,str) or len(code)!=6 or not code.isdigit(): return None
|
||||
now=time.time() if when is None else float(when)
|
||||
step=int(now//period)
|
||||
for s in range(step-window,step+window+1):
|
||||
if hmac.compare_digest(hotp(secret,s),code): return s
|
||||
return None
|
||||
|
||||
def load_secret(path:Path)->bytes:
|
||||
p=Path(path)
|
||||
if not p.is_file() or p.is_symlink(): raise TotpError('TOTP_KEY_NOT_FOUND')
|
||||
raw=p.read_text().strip()
|
||||
secret=b32decode_nopad(raw)
|
||||
if len(secret)<20: raise TotpError('TOTP_SECRET_TOO_SHORT')
|
||||
return secret
|
||||
|
||||
def consume_step(used_dir:Path,key_id:str,step:int,challenge_id:str)->None:
|
||||
d=Path(used_dir);d.mkdir(parents=True,exist_ok=True,mode=0o700)
|
||||
p=d/(f'{key_id}-{step}.used')
|
||||
try:fd=os.open(p,os.O_WRONLY|os.O_CREAT|os.O_EXCL,0o600)
|
||||
except FileExistsError: raise TotpError('TOTP_REPLAY') from None
|
||||
with os.fdopen(fd,'w') as f:f.write(str(challenge_id)+'\n')
|
||||
+1
@@ -0,0 +1 @@
|
||||
"""Vendored third-party libraries used by FrameBridge."""
|
||||
+907
@@ -0,0 +1,907 @@
|
||||
#
|
||||
# QR Code generator library (Python)
|
||||
#
|
||||
# Copyright (c) Project Nayuki. (MIT License)
|
||||
# https://www.nayuki.io/page/qr-code-generator-library
|
||||
#
|
||||
# Permission is hereby granted, free of charge, to any person obtaining a copy of
|
||||
# this software and associated documentation files (the "Software"), to deal in
|
||||
# the Software without restriction, including without limitation the rights to
|
||||
# use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
|
||||
# the Software, and to permit persons to whom the Software is furnished to do so,
|
||||
# subject to the following conditions:
|
||||
# - The above copyright notice and this permission notice shall be included in
|
||||
# all copies or substantial portions of the Software.
|
||||
# - The Software is provided "as is", without warranty of any kind, express or
|
||||
# implied, including but not limited to the warranties of merchantability,
|
||||
# fitness for a particular purpose and noninfringement. In no event shall the
|
||||
# authors or copyright holders be liable for any claim, damages or other
|
||||
# liability, whether in an action of contract, tort or otherwise, arising from,
|
||||
# out of or in connection with the Software or the use or other dealings in the
|
||||
# Software.
|
||||
#
|
||||
|
||||
from __future__ import annotations
|
||||
import collections, itertools, re
|
||||
from collections.abc import Sequence
|
||||
from typing import Optional, Union
|
||||
|
||||
|
||||
# ---- QR Code symbol class ----
|
||||
|
||||
class QrCode:
|
||||
"""A QR Code symbol, which is a type of two-dimension barcode.
|
||||
Invented by Denso Wave and described in the ISO/IEC 18004 standard.
|
||||
Instances of this class represent an immutable square grid of dark and light cells.
|
||||
The class provides static factory functions to create a QR Code from text or binary data.
|
||||
The class covers the QR Code Model 2 specification, supporting all versions (sizes)
|
||||
from 1 to 40, all 4 error correction levels, and 4 character encoding modes.
|
||||
|
||||
Ways to create a QR Code object:
|
||||
- High level: Take the payload data and call QrCode.encode_text() or QrCode.encode_binary().
|
||||
- Mid level: Custom-make the list of segments and call QrCode.encode_segments().
|
||||
- Low level: Custom-make the array of data codeword bytes (including
|
||||
segment headers and final padding, excluding error correction codewords),
|
||||
supply the appropriate version number, and call the QrCode() constructor.
|
||||
(Note that all ways require supplying the desired error correction level.)"""
|
||||
|
||||
# ---- Static factory functions (high level) ----
|
||||
|
||||
@staticmethod
|
||||
def encode_text(text: str, ecl: QrCode.Ecc) -> QrCode:
|
||||
"""Returns a QR Code representing the given Unicode text string at the given error correction level.
|
||||
As a conservative upper bound, this function is guaranteed to succeed for strings that have 738 or fewer
|
||||
Unicode code points (not UTF-16 code units) if the low error correction level is used. The smallest possible
|
||||
QR Code version is automatically chosen for the output. The ECC level of the result may be higher than the
|
||||
ecl argument if it can be done without increasing the version."""
|
||||
segs: list[QrSegment] = QrSegment.make_segments(text)
|
||||
return QrCode.encode_segments(segs, ecl)
|
||||
|
||||
|
||||
@staticmethod
|
||||
def encode_binary(data: Union[bytes,Sequence[int]], ecl: QrCode.Ecc) -> QrCode:
|
||||
"""Returns a QR Code representing the given binary data at the given error correction level.
|
||||
This function always encodes using the binary segment mode, not any text mode. The maximum number of
|
||||
bytes allowed is 2953. The smallest possible QR Code version is automatically chosen for the output.
|
||||
The ECC level of the result may be higher than the ecl argument if it can be done without increasing the version."""
|
||||
return QrCode.encode_segments([QrSegment.make_bytes(data)], ecl)
|
||||
|
||||
|
||||
# ---- Static factory functions (mid level) ----
|
||||
|
||||
@staticmethod
|
||||
def encode_segments(segs: Sequence[QrSegment], ecl: QrCode.Ecc, minversion: int = 1, maxversion: int = 40, mask: int = -1, boostecl: bool = True) -> QrCode:
|
||||
"""Returns a QR Code representing the given segments with the given encoding parameters.
|
||||
The smallest possible QR Code version within the given range is automatically
|
||||
chosen for the output. Iff boostecl is true, then the ECC level of the result
|
||||
may be higher than the ecl argument if it can be done without increasing the
|
||||
version. The mask number is either between 0 to 7 (inclusive) to force that
|
||||
mask, or -1 to automatically choose an appropriate mask (which may be slow).
|
||||
This function allows the user to create a custom sequence of segments that switches
|
||||
between modes (such as alphanumeric and byte) to encode text in less space.
|
||||
This is a mid-level API; the high-level API is encode_text() and encode_binary()."""
|
||||
|
||||
if not (QrCode.MIN_VERSION <= minversion <= maxversion <= QrCode.MAX_VERSION) or not (-1 <= mask <= 7):
|
||||
raise ValueError("Invalid value")
|
||||
|
||||
# Find the minimal version number to use
|
||||
for version in range(minversion, maxversion + 1):
|
||||
datacapacitybits: int = QrCode._get_num_data_codewords(version, ecl) * 8 # Number of data bits available
|
||||
datausedbits: Optional[int] = QrSegment.get_total_bits(segs, version)
|
||||
if (datausedbits is not None) and (datausedbits <= datacapacitybits):
|
||||
break # This version number is found to be suitable
|
||||
if version >= maxversion: # All versions in the range could not fit the given data
|
||||
msg: str = "Segment too long"
|
||||
if datausedbits is not None:
|
||||
msg = f"Data length = {datausedbits} bits, Max capacity = {datacapacitybits} bits"
|
||||
raise DataTooLongError(msg)
|
||||
assert datausedbits is not None
|
||||
|
||||
# Increase the error correction level while the data still fits in the current version number
|
||||
for newecl in (QrCode.Ecc.MEDIUM, QrCode.Ecc.QUARTILE, QrCode.Ecc.HIGH): # From low to high
|
||||
if boostecl and (datausedbits <= QrCode._get_num_data_codewords(version, newecl) * 8):
|
||||
ecl = newecl
|
||||
|
||||
# Concatenate all segments to create the data bit string
|
||||
bb = _BitBuffer()
|
||||
for seg in segs:
|
||||
bb.append_bits(seg.get_mode().get_mode_bits(), 4)
|
||||
bb.append_bits(seg.get_num_chars(), seg.get_mode().num_char_count_bits(version))
|
||||
bb.extend(seg._bitdata)
|
||||
assert len(bb) == datausedbits
|
||||
|
||||
# Add terminator and pad up to a byte if applicable
|
||||
datacapacitybits = QrCode._get_num_data_codewords(version, ecl) * 8
|
||||
assert len(bb) <= datacapacitybits
|
||||
bb.append_bits(0, min(4, datacapacitybits - len(bb)))
|
||||
bb.append_bits(0, -len(bb) % 8) # Note: Python's modulo on negative numbers behaves better than C family languages
|
||||
assert len(bb) % 8 == 0
|
||||
|
||||
# Pad with alternating bytes until data capacity is reached
|
||||
for padbyte in itertools.cycle((0xEC, 0x11)):
|
||||
if len(bb) >= datacapacitybits:
|
||||
break
|
||||
bb.append_bits(padbyte, 8)
|
||||
|
||||
# Pack bits into bytes in big endian
|
||||
datacodewords = bytearray([0] * (len(bb) // 8))
|
||||
for (i, bit) in enumerate(bb):
|
||||
datacodewords[i >> 3] |= bit << (7 - (i & 7))
|
||||
|
||||
# Create the QR Code object
|
||||
return QrCode(version, ecl, datacodewords, mask)
|
||||
|
||||
|
||||
# ---- Private fields ----
|
||||
|
||||
# The version number of this QR Code, which is between 1 and 40 (inclusive).
|
||||
# This determines the size of this barcode.
|
||||
_version: int
|
||||
|
||||
# The width and height of this QR Code, measured in modules, between
|
||||
# 21 and 177 (inclusive). This is equal to version * 4 + 17.
|
||||
_size: int
|
||||
|
||||
# The error correction level used in this QR Code.
|
||||
_errcorlvl: QrCode.Ecc
|
||||
|
||||
# The index of the mask pattern used in this QR Code, which is between 0 and 7 (inclusive).
|
||||
# Even if a QR Code is created with automatic masking requested (mask = -1),
|
||||
# the resulting object still has a mask value between 0 and 7.
|
||||
_mask: int
|
||||
|
||||
# The modules of this QR Code (False = light, True = dark).
|
||||
# Immutable after constructor finishes. Accessed through get_module().
|
||||
_modules: list[list[bool]]
|
||||
|
||||
# Indicates function modules that are not subjected to masking. Discarded when constructor finishes.
|
||||
_isfunction: list[list[bool]]
|
||||
|
||||
|
||||
# ---- Constructor (low level) ----
|
||||
|
||||
def __init__(self, version: int, errcorlvl: QrCode.Ecc, datacodewords: Union[bytes,Sequence[int]], msk: int) -> None:
|
||||
"""Creates a new QR Code with the given version number,
|
||||
error correction level, data codeword bytes, and mask number.
|
||||
This is a low-level API that most users should not use directly.
|
||||
A mid-level API is the encode_segments() function."""
|
||||
|
||||
# Check scalar arguments and set fields
|
||||
if not (QrCode.MIN_VERSION <= version <= QrCode.MAX_VERSION):
|
||||
raise ValueError("Version value out of range")
|
||||
if not (-1 <= msk <= 7):
|
||||
raise ValueError("Mask value out of range")
|
||||
|
||||
self._version = version
|
||||
self._size = version * 4 + 17
|
||||
self._errcorlvl = errcorlvl
|
||||
|
||||
# Initialize both grids to be size*size arrays of Boolean false
|
||||
self._modules = [[False] * self._size for _ in range(self._size)] # Initially all light
|
||||
self._isfunction = [[False] * self._size for _ in range(self._size)]
|
||||
|
||||
# Compute ECC, draw modules
|
||||
self._draw_function_patterns()
|
||||
allcodewords: bytes = self._add_ecc_and_interleave(bytearray(datacodewords))
|
||||
self._draw_codewords(allcodewords)
|
||||
|
||||
# Do masking
|
||||
if msk == -1: # Automatically choose best mask
|
||||
minpenalty: int = 1 << 32
|
||||
for i in range(8):
|
||||
self._apply_mask(i)
|
||||
self._draw_format_bits(i)
|
||||
penalty = self._get_penalty_score()
|
||||
if penalty < minpenalty:
|
||||
msk = i
|
||||
minpenalty = penalty
|
||||
self._apply_mask(i) # Undoes the mask due to XOR
|
||||
assert 0 <= msk <= 7
|
||||
self._mask = msk
|
||||
self._apply_mask(msk) # Apply the final choice of mask
|
||||
self._draw_format_bits(msk) # Overwrite old format bits
|
||||
|
||||
del self._isfunction
|
||||
|
||||
|
||||
# ---- Accessor methods ----
|
||||
|
||||
def get_version(self) -> int:
|
||||
"""Returns this QR Code's version number, in the range [1, 40]."""
|
||||
return self._version
|
||||
|
||||
def get_size(self) -> int:
|
||||
"""Returns this QR Code's size, in the range [21, 177]."""
|
||||
return self._size
|
||||
|
||||
def get_error_correction_level(self) -> QrCode.Ecc:
|
||||
"""Returns this QR Code's error correction level."""
|
||||
return self._errcorlvl
|
||||
|
||||
def get_mask(self) -> int:
|
||||
"""Returns this QR Code's mask, in the range [0, 7]."""
|
||||
return self._mask
|
||||
|
||||
def get_module(self, x: int, y: int) -> bool:
|
||||
"""Returns the color of the module (pixel) at the given coordinates, which is False
|
||||
for light or True for dark. The top left corner has the coordinates (x=0, y=0).
|
||||
If the given coordinates are out of bounds, then False (light) is returned."""
|
||||
return (0 <= x < self._size) and (0 <= y < self._size) and self._modules[y][x]
|
||||
|
||||
|
||||
# ---- Private helper methods for constructor: Drawing function modules ----
|
||||
|
||||
def _draw_function_patterns(self) -> None:
|
||||
"""Reads this object's version field, and draws and marks all function modules."""
|
||||
# Draw horizontal and vertical timing patterns
|
||||
for i in range(self._size):
|
||||
self._set_function_module(6, i, i % 2 == 0)
|
||||
self._set_function_module(i, 6, i % 2 == 0)
|
||||
|
||||
# Draw 3 finder patterns (all corners except bottom right; overwrites some timing modules)
|
||||
self._draw_finder_pattern(3, 3)
|
||||
self._draw_finder_pattern(self._size - 4, 3)
|
||||
self._draw_finder_pattern(3, self._size - 4)
|
||||
|
||||
# Draw numerous alignment patterns
|
||||
alignpatpos: list[int] = self._get_alignment_pattern_positions()
|
||||
numalign: int = len(alignpatpos)
|
||||
skips: Sequence[tuple[int,int]] = ((0, 0), (0, numalign - 1), (numalign - 1, 0))
|
||||
for i in range(numalign):
|
||||
for j in range(numalign):
|
||||
if (i, j) not in skips: # Don't draw on the three finder corners
|
||||
self._draw_alignment_pattern(alignpatpos[i], alignpatpos[j])
|
||||
|
||||
# Draw configuration data
|
||||
self._draw_format_bits(0) # Dummy mask value; overwritten later in the constructor
|
||||
self._draw_version()
|
||||
|
||||
|
||||
def _draw_format_bits(self, mask: int) -> None:
|
||||
"""Draws two copies of the format bits (with its own error correction code)
|
||||
based on the given mask and this object's error correction level field."""
|
||||
# Calculate error correction code and pack bits
|
||||
data: int = self._errcorlvl.formatbits << 3 | mask # errCorrLvl is uint2, mask is uint3
|
||||
rem: int = data
|
||||
for _ in range(10):
|
||||
rem = (rem << 1) ^ ((rem >> 9) * 0x537)
|
||||
bits: int = (data << 10 | rem) ^ 0x5412 # uint15
|
||||
assert bits >> 15 == 0
|
||||
|
||||
# Draw first copy
|
||||
for i in range(0, 6):
|
||||
self._set_function_module(8, i, _get_bit(bits, i))
|
||||
self._set_function_module(8, 7, _get_bit(bits, 6))
|
||||
self._set_function_module(8, 8, _get_bit(bits, 7))
|
||||
self._set_function_module(7, 8, _get_bit(bits, 8))
|
||||
for i in range(9, 15):
|
||||
self._set_function_module(14 - i, 8, _get_bit(bits, i))
|
||||
|
||||
# Draw second copy
|
||||
for i in range(0, 8):
|
||||
self._set_function_module(self._size - 1 - i, 8, _get_bit(bits, i))
|
||||
for i in range(8, 15):
|
||||
self._set_function_module(8, self._size - 15 + i, _get_bit(bits, i))
|
||||
self._set_function_module(8, self._size - 8, True) # Always dark
|
||||
|
||||
|
||||
def _draw_version(self) -> None:
|
||||
"""Draws two copies of the version bits (with its own error correction code),
|
||||
based on this object's version field, iff 7 <= version <= 40."""
|
||||
if self._version < 7:
|
||||
return
|
||||
|
||||
# Calculate error correction code and pack bits
|
||||
rem: int = self._version # version is uint6, in the range [7, 40]
|
||||
for _ in range(12):
|
||||
rem = (rem << 1) ^ ((rem >> 11) * 0x1F25)
|
||||
bits: int = self._version << 12 | rem # uint18
|
||||
assert bits >> 18 == 0
|
||||
|
||||
# Draw two copies
|
||||
for i in range(18):
|
||||
bit: bool = _get_bit(bits, i)
|
||||
a: int = self._size - 11 + i % 3
|
||||
b: int = i // 3
|
||||
self._set_function_module(a, b, bit)
|
||||
self._set_function_module(b, a, bit)
|
||||
|
||||
|
||||
def _draw_finder_pattern(self, x: int, y: int) -> None:
|
||||
"""Draws a 9*9 finder pattern including the border separator,
|
||||
with the center module at (x, y). Modules can be out of bounds."""
|
||||
for dy in range(-4, 5):
|
||||
for dx in range(-4, 5):
|
||||
xx, yy = x + dx, y + dy
|
||||
if (0 <= xx < self._size) and (0 <= yy < self._size):
|
||||
# Chebyshev/infinity norm
|
||||
self._set_function_module(xx, yy, max(abs(dx), abs(dy)) not in (2, 4))
|
||||
|
||||
|
||||
def _draw_alignment_pattern(self, x: int, y: int) -> None:
|
||||
"""Draws a 5*5 alignment pattern, with the center module
|
||||
at (x, y). All modules must be in bounds."""
|
||||
for dy in range(-2, 3):
|
||||
for dx in range(-2, 3):
|
||||
self._set_function_module(x + dx, y + dy, max(abs(dx), abs(dy)) != 1)
|
||||
|
||||
|
||||
def _set_function_module(self, x: int, y: int, isdark: bool) -> None:
|
||||
"""Sets the color of a module and marks it as a function module.
|
||||
Only used by the constructor. Coordinates must be in bounds."""
|
||||
assert type(isdark) is bool
|
||||
self._modules[y][x] = isdark
|
||||
self._isfunction[y][x] = True
|
||||
|
||||
|
||||
# ---- Private helper methods for constructor: Codewords and masking ----
|
||||
|
||||
def _add_ecc_and_interleave(self, data: bytearray) -> bytes:
|
||||
"""Returns a new byte string representing the given data with the appropriate error correction
|
||||
codewords appended to it, based on this object's version and error correction level."""
|
||||
version: int = self._version
|
||||
assert len(data) == QrCode._get_num_data_codewords(version, self._errcorlvl)
|
||||
|
||||
# Calculate parameter numbers
|
||||
numblocks: int = QrCode._NUM_ERROR_CORRECTION_BLOCKS[self._errcorlvl.ordinal][version]
|
||||
blockecclen: int = QrCode._ECC_CODEWORDS_PER_BLOCK [self._errcorlvl.ordinal][version]
|
||||
rawcodewords: int = QrCode._get_num_raw_data_modules(version) // 8
|
||||
numshortblocks: int = numblocks - rawcodewords % numblocks
|
||||
shortblocklen: int = rawcodewords // numblocks
|
||||
|
||||
# Split data into blocks and append ECC to each block
|
||||
blocks: list[bytes] = []
|
||||
rsdiv: bytes = QrCode._reed_solomon_compute_divisor(blockecclen)
|
||||
k: int = 0
|
||||
for i in range(numblocks):
|
||||
dat: bytearray = data[k : k + shortblocklen - blockecclen + (0 if i < numshortblocks else 1)]
|
||||
k += len(dat)
|
||||
ecc: bytes = QrCode._reed_solomon_compute_remainder(dat, rsdiv)
|
||||
if i < numshortblocks:
|
||||
dat.append(0)
|
||||
blocks.append(dat + ecc)
|
||||
assert k == len(data)
|
||||
|
||||
# Interleave (not concatenate) the bytes from every block into a single sequence
|
||||
result = bytearray()
|
||||
for i in range(len(blocks[0])):
|
||||
for (j, blk) in enumerate(blocks):
|
||||
# Skip the padding byte in short blocks
|
||||
if (i != shortblocklen - blockecclen) or (j >= numshortblocks):
|
||||
result.append(blk[i])
|
||||
assert len(result) == rawcodewords
|
||||
return result
|
||||
|
||||
|
||||
def _draw_codewords(self, data: bytes) -> None:
|
||||
"""Draws the given sequence of 8-bit codewords (data and error correction) onto the entire
|
||||
data area of this QR Code. Function modules need to be marked off before this is called."""
|
||||
assert len(data) == QrCode._get_num_raw_data_modules(self._version) // 8
|
||||
|
||||
i: int = 0 # Bit index into the data
|
||||
# Do the funny zigzag scan
|
||||
for right in range(self._size - 1, 0, -2): # Index of right column in each column pair
|
||||
if right <= 6:
|
||||
right -= 1
|
||||
for vert in range(self._size): # Vertical counter
|
||||
for j in range(2):
|
||||
x: int = right - j # Actual x coordinate
|
||||
upward: bool = (right + 1) & 2 == 0
|
||||
y: int = (self._size - 1 - vert) if upward else vert # Actual y coordinate
|
||||
if (not self._isfunction[y][x]) and (i < len(data) * 8):
|
||||
self._modules[y][x] = _get_bit(data[i >> 3], 7 - (i & 7))
|
||||
i += 1
|
||||
# If this QR Code has any remainder bits (0 to 7), they were assigned as
|
||||
# 0/false/light by the constructor and are left unchanged by this method
|
||||
assert i == len(data) * 8
|
||||
|
||||
|
||||
def _apply_mask(self, mask: int) -> None:
|
||||
"""XORs the codeword modules in this QR Code with the given mask pattern.
|
||||
The function modules must be marked and the codeword bits must be drawn
|
||||
before masking. Due to the arithmetic of XOR, calling _apply_mask() with
|
||||
the same mask value a second time will undo the mask. A final well-formed
|
||||
QR Code needs exactly one (not zero, two, etc.) mask applied."""
|
||||
if not (0 <= mask <= 7):
|
||||
raise ValueError("Mask value out of range")
|
||||
masker: collections.abc.Callable[[int,int],int] = QrCode._MASK_PATTERNS[mask]
|
||||
for y in range(self._size):
|
||||
for x in range(self._size):
|
||||
self._modules[y][x] ^= (masker(x, y) == 0) and (not self._isfunction[y][x])
|
||||
|
||||
|
||||
def _get_penalty_score(self) -> int:
|
||||
"""Calculates and returns the penalty score based on state of this QR Code's current modules.
|
||||
This is used by the automatic mask choice algorithm to find the mask pattern that yields the lowest score."""
|
||||
result: int = 0
|
||||
size: int = self._size
|
||||
modules: list[list[bool]] = self._modules
|
||||
|
||||
# Adjacent modules in row having same color, and finder-like patterns
|
||||
for y in range(size):
|
||||
runcolor: bool = False
|
||||
runx: int = 0
|
||||
runhistory = collections.deque([0] * 7, 7)
|
||||
for x in range(size):
|
||||
if modules[y][x] == runcolor:
|
||||
runx += 1
|
||||
if runx == 5:
|
||||
result += QrCode._PENALTY_N1
|
||||
elif runx > 5:
|
||||
result += 1
|
||||
else:
|
||||
self._finder_penalty_add_history(runx, runhistory)
|
||||
if not runcolor:
|
||||
result += self._finder_penalty_count_patterns(runhistory) * QrCode._PENALTY_N3
|
||||
runcolor = modules[y][x]
|
||||
runx = 1
|
||||
result += self._finder_penalty_terminate_and_count(runcolor, runx, runhistory) * QrCode._PENALTY_N3
|
||||
# Adjacent modules in column having same color, and finder-like patterns
|
||||
for x in range(size):
|
||||
runcolor = False
|
||||
runy: int = 0
|
||||
runhistory = collections.deque([0] * 7, 7)
|
||||
for y in range(size):
|
||||
if modules[y][x] == runcolor:
|
||||
runy += 1
|
||||
if runy == 5:
|
||||
result += QrCode._PENALTY_N1
|
||||
elif runy > 5:
|
||||
result += 1
|
||||
else:
|
||||
self._finder_penalty_add_history(runy, runhistory)
|
||||
if not runcolor:
|
||||
result += self._finder_penalty_count_patterns(runhistory) * QrCode._PENALTY_N3
|
||||
runcolor = modules[y][x]
|
||||
runy = 1
|
||||
result += self._finder_penalty_terminate_and_count(runcolor, runy, runhistory) * QrCode._PENALTY_N3
|
||||
|
||||
# 2*2 blocks of modules having same color
|
||||
for y in range(size - 1):
|
||||
for x in range(size - 1):
|
||||
if modules[y][x] == modules[y][x + 1] == modules[y + 1][x] == modules[y + 1][x + 1]:
|
||||
result += QrCode._PENALTY_N2
|
||||
|
||||
# Balance of dark and light modules
|
||||
dark: int = sum((1 if cell else 0) for row in modules for cell in row)
|
||||
total: int = size**2 # Note that size is odd, so dark/total != 1/2
|
||||
# Compute the smallest integer k >= 0 such that (45-5k)% <= dark/total <= (55+5k)%
|
||||
k: int = (abs(dark * 20 - total * 10) + total - 1) // total - 1
|
||||
assert 0 <= k <= 9
|
||||
result += k * QrCode._PENALTY_N4
|
||||
assert 0 <= result <= 2568888 # Non-tight upper bound based on default values of PENALTY_N1, ..., N4
|
||||
return result
|
||||
|
||||
|
||||
# ---- Private helper functions ----
|
||||
|
||||
def _get_alignment_pattern_positions(self) -> list[int]:
|
||||
"""Returns an ascending list of positions of alignment patterns for this version number.
|
||||
Each position is in the range [0,177), and are used on both the x and y axes.
|
||||
This could be implemented as lookup table of 40 variable-length lists of integers."""
|
||||
if self._version == 1:
|
||||
return []
|
||||
else:
|
||||
numalign: int = self._version // 7 + 2
|
||||
step: int = (self._version * 8 + numalign * 3 + 5) // (numalign * 4 - 4) * 2
|
||||
result: list[int] = [(self._size - 7 - i * step) for i in range(numalign - 1)] + [6]
|
||||
return list(reversed(result))
|
||||
|
||||
|
||||
@staticmethod
|
||||
def _get_num_raw_data_modules(ver: int) -> int:
|
||||
"""Returns the number of data bits that can be stored in a QR Code of the given version number, after
|
||||
all function modules are excluded. This includes remainder bits, so it might not be a multiple of 8.
|
||||
The result is in the range [208, 29648]. This could be implemented as a 40-entry lookup table."""
|
||||
if not (QrCode.MIN_VERSION <= ver <= QrCode.MAX_VERSION):
|
||||
raise ValueError("Version number out of range")
|
||||
result: int = (16 * ver + 128) * ver + 64
|
||||
if ver >= 2:
|
||||
numalign: int = ver // 7 + 2
|
||||
result -= (25 * numalign - 10) * numalign - 55
|
||||
if ver >= 7:
|
||||
result -= 36
|
||||
assert 208 <= result <= 29648
|
||||
return result
|
||||
|
||||
|
||||
@staticmethod
|
||||
def _get_num_data_codewords(ver: int, ecl: QrCode.Ecc) -> int:
|
||||
"""Returns the number of 8-bit data (i.e. not error correction) codewords contained in any
|
||||
QR Code of the given version number and error correction level, with remainder bits discarded.
|
||||
This stateless pure function could be implemented as a (40*4)-cell lookup table."""
|
||||
return QrCode._get_num_raw_data_modules(ver) // 8 \
|
||||
- QrCode._ECC_CODEWORDS_PER_BLOCK [ecl.ordinal][ver] \
|
||||
* QrCode._NUM_ERROR_CORRECTION_BLOCKS[ecl.ordinal][ver]
|
||||
|
||||
|
||||
@staticmethod
|
||||
def _reed_solomon_compute_divisor(degree: int) -> bytes:
|
||||
"""Returns a Reed-Solomon ECC generator polynomial for the given degree. This could be
|
||||
implemented as a lookup table over all possible parameter values, instead of as an algorithm."""
|
||||
if not (1 <= degree <= 255):
|
||||
raise ValueError("Degree out of range")
|
||||
# Polynomial coefficients are stored from highest to lowest power, excluding the leading term which is always 1.
|
||||
# For example the polynomial x^3 + 255x^2 + 8x + 93 is stored as the uint8 array [255, 8, 93].
|
||||
result = bytearray([0] * (degree - 1) + [1]) # Start off with the monomial x^0
|
||||
|
||||
# Compute the product polynomial (x - r^0) * (x - r^1) * (x - r^2) * ... * (x - r^{degree-1}),
|
||||
# and drop the highest monomial term which is always 1x^degree.
|
||||
# Note that r = 0x02, which is a generator element of this field GF(2^8/0x11D).
|
||||
root: int = 1
|
||||
for _ in range(degree): # Unused variable i
|
||||
# Multiply the current product by (x - r^i)
|
||||
for j in range(degree):
|
||||
result[j] = QrCode._reed_solomon_multiply(result[j], root)
|
||||
if j + 1 < degree:
|
||||
result[j] ^= result[j + 1]
|
||||
root = QrCode._reed_solomon_multiply(root, 0x02)
|
||||
return result
|
||||
|
||||
|
||||
@staticmethod
|
||||
def _reed_solomon_compute_remainder(data: bytes, divisor: bytes) -> bytes:
|
||||
"""Returns the Reed-Solomon error correction codeword for the given data and divisor polynomials."""
|
||||
result = bytearray([0] * len(divisor))
|
||||
for b in data: # Polynomial division
|
||||
factor: int = b ^ result.pop(0)
|
||||
result.append(0)
|
||||
for (i, coef) in enumerate(divisor):
|
||||
result[i] ^= QrCode._reed_solomon_multiply(coef, factor)
|
||||
return result
|
||||
|
||||
|
||||
@staticmethod
|
||||
def _reed_solomon_multiply(x: int, y: int) -> int:
|
||||
"""Returns the product of the two given field elements modulo GF(2^8/0x11D). The arguments and result
|
||||
are unsigned 8-bit integers. This could be implemented as a lookup table of 256*256 entries of uint8."""
|
||||
if (x >> 8 != 0) or (y >> 8 != 0):
|
||||
raise ValueError("Byte out of range")
|
||||
# Russian peasant multiplication
|
||||
z: int = 0
|
||||
for i in reversed(range(8)):
|
||||
z = (z << 1) ^ ((z >> 7) * 0x11D)
|
||||
z ^= ((y >> i) & 1) * x
|
||||
assert z >> 8 == 0
|
||||
return z
|
||||
|
||||
|
||||
def _finder_penalty_count_patterns(self, runhistory: collections.deque[int]) -> int:
|
||||
"""Can only be called immediately after a light run is added, and
|
||||
returns either 0, 1, or 2. A helper function for _get_penalty_score()."""
|
||||
n: int = runhistory[1]
|
||||
assert n <= self._size * 3
|
||||
core: bool = n > 0 and (runhistory[2] == runhistory[4] == runhistory[5] == n) and runhistory[3] == n * 3
|
||||
return (1 if (core and runhistory[0] >= n * 4 and runhistory[6] >= n) else 0) \
|
||||
+ (1 if (core and runhistory[6] >= n * 4 and runhistory[0] >= n) else 0)
|
||||
|
||||
|
||||
def _finder_penalty_terminate_and_count(self, currentruncolor: bool, currentrunlength: int, runhistory: collections.deque[int]) -> int:
|
||||
"""Must be called at the end of a line (row or column) of modules. A helper function for _get_penalty_score()."""
|
||||
if currentruncolor: # Terminate dark run
|
||||
self._finder_penalty_add_history(currentrunlength, runhistory)
|
||||
currentrunlength = 0
|
||||
currentrunlength += self._size # Add light border to final run
|
||||
self._finder_penalty_add_history(currentrunlength, runhistory)
|
||||
return self._finder_penalty_count_patterns(runhistory)
|
||||
|
||||
|
||||
def _finder_penalty_add_history(self, currentrunlength: int, runhistory: collections.deque[int]) -> None:
|
||||
if runhistory[0] == 0:
|
||||
currentrunlength += self._size # Add light border to initial run
|
||||
runhistory.appendleft(currentrunlength)
|
||||
|
||||
|
||||
# ---- Constants and tables ----
|
||||
|
||||
MIN_VERSION: int = 1 # The minimum version number supported in the QR Code Model 2 standard
|
||||
MAX_VERSION: int = 40 # The maximum version number supported in the QR Code Model 2 standard
|
||||
|
||||
# For use in _get_penalty_score(), when evaluating which mask is best.
|
||||
_PENALTY_N1: int = 3
|
||||
_PENALTY_N2: int = 3
|
||||
_PENALTY_N3: int = 40
|
||||
_PENALTY_N4: int = 10
|
||||
|
||||
_ECC_CODEWORDS_PER_BLOCK: Sequence[Sequence[int]] = (
|
||||
# Version: (note that index 0 is for padding, and is set to an illegal value)
|
||||
# 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40 Error correction level
|
||||
(-1, 7, 10, 15, 20, 26, 18, 20, 24, 30, 18, 20, 24, 26, 30, 22, 24, 28, 30, 28, 28, 28, 28, 30, 30, 26, 28, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30), # Low
|
||||
(-1, 10, 16, 26, 18, 24, 16, 18, 22, 22, 26, 30, 22, 22, 24, 24, 28, 28, 26, 26, 26, 26, 28, 28, 28, 28, 28, 28, 28, 28, 28, 28, 28, 28, 28, 28, 28, 28, 28, 28, 28), # Medium
|
||||
(-1, 13, 22, 18, 26, 18, 24, 18, 22, 20, 24, 28, 26, 24, 20, 30, 24, 28, 28, 26, 30, 28, 30, 30, 30, 30, 28, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30), # Quartile
|
||||
(-1, 17, 28, 22, 16, 22, 28, 26, 26, 24, 28, 24, 28, 22, 24, 24, 30, 28, 28, 26, 28, 30, 24, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30)) # High
|
||||
|
||||
_NUM_ERROR_CORRECTION_BLOCKS: Sequence[Sequence[int]] = (
|
||||
# Version: (note that index 0 is for padding, and is set to an illegal value)
|
||||
# 0, 1, 2, 3, 4, 5, 6, 7, 8, 9,10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40 Error correction level
|
||||
(-1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 4, 4, 4, 4, 4, 6, 6, 6, 6, 7, 8, 8, 9, 9, 10, 12, 12, 12, 13, 14, 15, 16, 17, 18, 19, 19, 20, 21, 22, 24, 25), # Low
|
||||
(-1, 1, 1, 1, 2, 2, 4, 4, 4, 5, 5, 5, 8, 9, 9, 10, 10, 11, 13, 14, 16, 17, 17, 18, 20, 21, 23, 25, 26, 28, 29, 31, 33, 35, 37, 38, 40, 43, 45, 47, 49), # Medium
|
||||
(-1, 1, 1, 2, 2, 4, 4, 6, 6, 8, 8, 8, 10, 12, 16, 12, 17, 16, 18, 21, 20, 23, 23, 25, 27, 29, 34, 34, 35, 38, 40, 43, 45, 48, 51, 53, 56, 59, 62, 65, 68), # Quartile
|
||||
(-1, 1, 1, 2, 4, 4, 4, 5, 6, 8, 8, 11, 11, 16, 16, 18, 16, 19, 21, 25, 25, 25, 34, 30, 32, 35, 37, 40, 42, 45, 48, 51, 54, 57, 60, 63, 66, 70, 74, 77, 81)) # High
|
||||
|
||||
_MASK_PATTERNS: Sequence[collections.abc.Callable[[int,int],int]] = (
|
||||
(lambda x, y: (x + y) % 2 ),
|
||||
(lambda x, y: y % 2 ),
|
||||
(lambda x, y: x % 3 ),
|
||||
(lambda x, y: (x + y) % 3 ),
|
||||
(lambda x, y: (x // 3 + y // 2) % 2 ),
|
||||
(lambda x, y: x * y % 2 + x * y % 3 ),
|
||||
(lambda x, y: (x * y % 2 + x * y % 3) % 2 ),
|
||||
(lambda x, y: ((x + y) % 2 + x * y % 3) % 2),
|
||||
)
|
||||
|
||||
|
||||
# ---- Public helper enumeration ----
|
||||
|
||||
class Ecc:
|
||||
ordinal: int # (Public) In the range 0 to 3 (unsigned 2-bit integer)
|
||||
formatbits: int # (Package-private) In the range 0 to 3 (unsigned 2-bit integer)
|
||||
|
||||
"""The error correction level in a QR Code symbol. Immutable."""
|
||||
# Private constructor
|
||||
def __init__(self, i: int, fb: int) -> None:
|
||||
self.ordinal = i
|
||||
self.formatbits = fb
|
||||
|
||||
# Placeholders
|
||||
LOW : QrCode.Ecc
|
||||
MEDIUM : QrCode.Ecc
|
||||
QUARTILE: QrCode.Ecc
|
||||
HIGH : QrCode.Ecc
|
||||
|
||||
# Public constants. Create them outside the class.
|
||||
Ecc.LOW = Ecc(0, 1) # The QR Code can tolerate about 7% erroneous codewords
|
||||
Ecc.MEDIUM = Ecc(1, 0) # The QR Code can tolerate about 15% erroneous codewords
|
||||
Ecc.QUARTILE = Ecc(2, 3) # The QR Code can tolerate about 25% erroneous codewords
|
||||
Ecc.HIGH = Ecc(3, 2) # The QR Code can tolerate about 30% erroneous codewords
|
||||
|
||||
|
||||
|
||||
# ---- Data segment class ----
|
||||
|
||||
class QrSegment:
|
||||
"""A segment of character/binary/control data in a QR Code symbol.
|
||||
Instances of this class are immutable.
|
||||
The mid-level way to create a segment is to take the payload data
|
||||
and call a static factory function such as QrSegment.make_numeric().
|
||||
The low-level way to create a segment is to custom-make the bit buffer
|
||||
and call the QrSegment() constructor with appropriate values.
|
||||
This segment class imposes no length restrictions, but QR Codes have restrictions.
|
||||
Even in the most favorable conditions, a QR Code can only hold 7089 characters of data.
|
||||
Any segment longer than this is meaningless for the purpose of generating QR Codes."""
|
||||
|
||||
# ---- Static factory functions (mid level) ----
|
||||
|
||||
@staticmethod
|
||||
def make_bytes(data: Union[bytes,Sequence[int]]) -> QrSegment:
|
||||
"""Returns a segment representing the given binary data encoded in byte mode.
|
||||
All input byte lists are acceptable. Any text string can be converted to
|
||||
UTF-8 bytes (s.encode("UTF-8")) and encoded as a byte mode segment."""
|
||||
bb = _BitBuffer()
|
||||
for b in data:
|
||||
bb.append_bits(b, 8)
|
||||
return QrSegment(QrSegment.Mode.BYTE, len(data), bb)
|
||||
|
||||
|
||||
@staticmethod
|
||||
def make_numeric(digits: str) -> QrSegment:
|
||||
"""Returns a segment representing the given string of decimal digits encoded in numeric mode."""
|
||||
if not QrSegment.is_numeric(digits):
|
||||
raise ValueError("String contains non-numeric characters")
|
||||
bb = _BitBuffer()
|
||||
i: int = 0
|
||||
while i < len(digits): # Consume up to 3 digits per iteration
|
||||
n: int = min(len(digits) - i, 3)
|
||||
bb.append_bits(int(digits[i : i + n]), n * 3 + 1)
|
||||
i += n
|
||||
return QrSegment(QrSegment.Mode.NUMERIC, len(digits), bb)
|
||||
|
||||
|
||||
@staticmethod
|
||||
def make_alphanumeric(text: str) -> QrSegment:
|
||||
"""Returns a segment representing the given text string encoded in alphanumeric mode.
|
||||
The characters allowed are: 0 to 9, A to Z (uppercase only), space,
|
||||
dollar, percent, asterisk, plus, hyphen, period, slash, colon."""
|
||||
if not QrSegment.is_alphanumeric(text):
|
||||
raise ValueError("String contains unencodable characters in alphanumeric mode")
|
||||
bb = _BitBuffer()
|
||||
for i in range(0, len(text) - 1, 2): # Process groups of 2
|
||||
temp: int = QrSegment._ALPHANUMERIC_ENCODING_TABLE[text[i]] * 45
|
||||
temp += QrSegment._ALPHANUMERIC_ENCODING_TABLE[text[i + 1]]
|
||||
bb.append_bits(temp, 11)
|
||||
if len(text) % 2 > 0: # 1 character remaining
|
||||
bb.append_bits(QrSegment._ALPHANUMERIC_ENCODING_TABLE[text[-1]], 6)
|
||||
return QrSegment(QrSegment.Mode.ALPHANUMERIC, len(text), bb)
|
||||
|
||||
|
||||
@staticmethod
|
||||
def make_segments(text: str) -> list[QrSegment]:
|
||||
"""Returns a new mutable list of zero or more segments to represent the given Unicode text string.
|
||||
The result may use various segment modes and switch modes to optimize the length of the bit stream."""
|
||||
|
||||
# Select the most efficient segment encoding automatically
|
||||
if text == "":
|
||||
return []
|
||||
elif QrSegment.is_numeric(text):
|
||||
return [QrSegment.make_numeric(text)]
|
||||
elif QrSegment.is_alphanumeric(text):
|
||||
return [QrSegment.make_alphanumeric(text)]
|
||||
else:
|
||||
return [QrSegment.make_bytes(text.encode("UTF-8"))]
|
||||
|
||||
|
||||
@staticmethod
|
||||
def make_eci(assignval: int) -> QrSegment:
|
||||
"""Returns a segment representing an Extended Channel Interpretation
|
||||
(ECI) designator with the given assignment value."""
|
||||
bb = _BitBuffer()
|
||||
if assignval < 0:
|
||||
raise ValueError("ECI assignment value out of range")
|
||||
elif assignval < (1 << 7):
|
||||
bb.append_bits(assignval, 8)
|
||||
elif assignval < (1 << 14):
|
||||
bb.append_bits(0b10, 2)
|
||||
bb.append_bits(assignval, 14)
|
||||
elif assignval < 1000000:
|
||||
bb.append_bits(0b110, 3)
|
||||
bb.append_bits(assignval, 21)
|
||||
else:
|
||||
raise ValueError("ECI assignment value out of range")
|
||||
return QrSegment(QrSegment.Mode.ECI, 0, bb)
|
||||
|
||||
|
||||
# Tests whether the given string can be encoded as a segment in numeric mode.
|
||||
# A string is encodable iff each character is in the range 0 to 9.
|
||||
@staticmethod
|
||||
def is_numeric(text: str) -> bool:
|
||||
return QrSegment._NUMERIC_REGEX.fullmatch(text) is not None
|
||||
|
||||
|
||||
# Tests whether the given string can be encoded as a segment in alphanumeric mode.
|
||||
# A string is encodable iff each character is in the following set: 0 to 9, A to Z
|
||||
# (uppercase only), space, dollar, percent, asterisk, plus, hyphen, period, slash, colon.
|
||||
@staticmethod
|
||||
def is_alphanumeric(text: str) -> bool:
|
||||
return QrSegment._ALPHANUMERIC_REGEX.fullmatch(text) is not None
|
||||
|
||||
|
||||
# ---- Private fields ----
|
||||
|
||||
# The mode indicator of this segment. Accessed through get_mode().
|
||||
_mode: QrSegment.Mode
|
||||
|
||||
# The length of this segment's unencoded data. Measured in characters for
|
||||
# numeric/alphanumeric/kanji mode, bytes for byte mode, and 0 for ECI mode.
|
||||
# Always zero or positive. Not the same as the data's bit length.
|
||||
# Accessed through get_num_chars().
|
||||
_numchars: int
|
||||
|
||||
# The data bits of this segment. Accessed through get_data().
|
||||
_bitdata: list[int]
|
||||
|
||||
|
||||
# ---- Constructor (low level) ----
|
||||
|
||||
def __init__(self, mode: QrSegment.Mode, numch: int, bitdata: Sequence[int]) -> None:
|
||||
"""Creates a new QR Code segment with the given attributes and data.
|
||||
The character count (numch) must agree with the mode and the bit buffer length,
|
||||
but the constraint isn't checked. The given bit buffer is cloned and stored."""
|
||||
if numch < 0:
|
||||
raise ValueError()
|
||||
self._mode = mode
|
||||
self._numchars = numch
|
||||
self._bitdata = list(bitdata) # Make defensive copy
|
||||
|
||||
|
||||
# ---- Accessor methods ----
|
||||
|
||||
def get_mode(self) -> QrSegment.Mode:
|
||||
"""Returns the mode field of this segment."""
|
||||
return self._mode
|
||||
|
||||
def get_num_chars(self) -> int:
|
||||
"""Returns the character count field of this segment."""
|
||||
return self._numchars
|
||||
|
||||
def get_data(self) -> list[int]:
|
||||
"""Returns a new copy of the data bits of this segment."""
|
||||
return list(self._bitdata) # Make defensive copy
|
||||
|
||||
|
||||
# Package-private function
|
||||
@staticmethod
|
||||
def get_total_bits(segs: Sequence[QrSegment], version: int) -> Optional[int]:
|
||||
"""Calculates the number of bits needed to encode the given segments at
|
||||
the given version. Returns a non-negative number if successful. Otherwise
|
||||
returns None if a segment has too many characters to fit its length field."""
|
||||
result = 0
|
||||
for seg in segs:
|
||||
ccbits: int = seg.get_mode().num_char_count_bits(version)
|
||||
if seg.get_num_chars() >= (1 << ccbits):
|
||||
return None # The segment's length doesn't fit the field's bit width
|
||||
result += 4 + ccbits + len(seg._bitdata)
|
||||
return result
|
||||
|
||||
|
||||
# ---- Constants ----
|
||||
|
||||
# Describes precisely all strings that are encodable in numeric mode.
|
||||
_NUMERIC_REGEX: re.Pattern[str] = re.compile(r"[0-9]*")
|
||||
|
||||
# Describes precisely all strings that are encodable in alphanumeric mode.
|
||||
_ALPHANUMERIC_REGEX: re.Pattern[str] = re.compile(r"[A-Z0-9 $%*+./:-]*")
|
||||
|
||||
# Dictionary of "0"->0, "A"->10, "$"->37, etc.
|
||||
_ALPHANUMERIC_ENCODING_TABLE: dict[str,int] = {ch: i for (i, ch) in enumerate("0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ $%*+-./:")}
|
||||
|
||||
|
||||
# ---- Public helper enumeration ----
|
||||
|
||||
class Mode:
|
||||
"""Describes how a segment's data bits are interpreted. Immutable."""
|
||||
|
||||
_modebits: int # The mode indicator bits, which is a uint4 value (range 0 to 15)
|
||||
_charcounts: tuple[int,int,int] # Number of character count bits for three different version ranges
|
||||
|
||||
# Private constructor
|
||||
def __init__(self, modebits: int, charcounts: tuple[int,int,int]):
|
||||
self._modebits = modebits
|
||||
self._charcounts = charcounts
|
||||
|
||||
# Package-private method
|
||||
def get_mode_bits(self) -> int:
|
||||
"""Returns an unsigned 4-bit integer value (range 0 to 15) representing the mode indicator bits for this mode object."""
|
||||
return self._modebits
|
||||
|
||||
# Package-private method
|
||||
def num_char_count_bits(self, ver: int) -> int:
|
||||
"""Returns the bit width of the character count field for a segment in this mode
|
||||
in a QR Code at the given version number. The result is in the range [0, 16]."""
|
||||
return self._charcounts[(ver + 7) // 17]
|
||||
|
||||
# Placeholders
|
||||
NUMERIC : QrSegment.Mode
|
||||
ALPHANUMERIC: QrSegment.Mode
|
||||
BYTE : QrSegment.Mode
|
||||
KANJI : QrSegment.Mode
|
||||
ECI : QrSegment.Mode
|
||||
|
||||
# Public constants. Create them outside the class.
|
||||
Mode.NUMERIC = Mode(0x1, (10, 12, 14))
|
||||
Mode.ALPHANUMERIC = Mode(0x2, ( 9, 11, 13))
|
||||
Mode.BYTE = Mode(0x4, ( 8, 16, 16))
|
||||
Mode.KANJI = Mode(0x8, ( 8, 10, 12))
|
||||
Mode.ECI = Mode(0x7, ( 0, 0, 0))
|
||||
|
||||
|
||||
|
||||
# ---- Private helper class ----
|
||||
|
||||
class _BitBuffer(list[int]):
|
||||
"""An appendable sequence of bits (0s and 1s). Mainly used by QrSegment."""
|
||||
|
||||
def append_bits(self, val: int, n: int) -> None:
|
||||
"""Appends the given number of low-order bits of the given
|
||||
value to this buffer. Requires n >= 0 and 0 <= val < 2^n."""
|
||||
if (n < 0) or (val >> n != 0):
|
||||
raise ValueError("Value out of range")
|
||||
self.extend(((val >> i) & 1) for i in reversed(range(n)))
|
||||
|
||||
|
||||
def _get_bit(x: int, i: int) -> bool:
|
||||
"""Returns true iff the i'th bit of x is set to 1."""
|
||||
return (x >> i) & 1 != 0
|
||||
|
||||
|
||||
|
||||
class DataTooLongError(ValueError):
|
||||
"""Raised when the supplied data does not fit any QR Code version. Ways to handle this exception include:
|
||||
- Decrease the error correction level if it was greater than Ecc.LOW.
|
||||
- If the encode_segments() function was called with a maxversion argument, then increase
|
||||
it if it was less than QrCode.MAX_VERSION. (This advice does not apply to the other
|
||||
factory functions because they search all versions up to QrCode.MAX_VERSION.)
|
||||
- Split the text data into better or optimal segments in order to reduce the number of bits required.
|
||||
- Change the text or binary data to be shorter.
|
||||
- Change the text to fit the character set of a particular segment mode (e.g. alphanumeric).
|
||||
- Propagate the error upward to the caller/user."""
|
||||
pass
|
||||
@@ -0,0 +1,42 @@
|
||||
"""Separate loopback-only MCP entry point for the six Fleet READ tools."""
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
|
||||
from mcp.server.fastmcp import FastMCP
|
||||
|
||||
from framebridge_node_gateway import discover_nodes, node_read
|
||||
from framebridge_linux_chatgpt.fleet_hub import FleetHub
|
||||
from framebridge_linux_chatgpt.fleet_mcp_surface import register_fleet_read_tools
|
||||
|
||||
|
||||
def _server_address() -> tuple[str, int]:
|
||||
host = os.environ.get("FRAMEBRIDGE_FLEET_HUB_HOST", "127.0.0.1")
|
||||
if not isinstance(host, str) or host not in ("127.0.0.1", "::1"):
|
||||
raise ValueError("FLEET_HUB_HOST_INVALID")
|
||||
try:
|
||||
port = int(os.environ.get("FRAMEBRIDGE_FLEET_HUB_PORT", "8877"))
|
||||
except (TypeError, ValueError):
|
||||
raise ValueError("FLEET_HUB_PORT_INVALID") from None
|
||||
if not 1024 <= port <= 65535:
|
||||
raise ValueError("FLEET_HUB_PORT_INVALID")
|
||||
return host, port
|
||||
|
||||
|
||||
host, port = _server_address()
|
||||
mcp = FastMCP(
|
||||
"FrameBridge Fleet Hub",
|
||||
host=host,
|
||||
port=port,
|
||||
streamable_http_path="/mcp",
|
||||
)
|
||||
hub = FleetHub(discover_nodes, node_read)
|
||||
register_fleet_read_tools(mcp, hub)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
mcp.run(transport="streamable-http")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
+1
@@ -0,0 +1 @@
|
||||
"""FrameBridge Fleet Hub runtime."""
|
||||
+42
@@ -0,0 +1,42 @@
|
||||
"""Injected Node Gateway READ boundary; no transport or live gateway imports."""
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Callable
|
||||
|
||||
from .fleet_registry import FleetNode
|
||||
from .fleet_routing import Operation, ReadRequest, normalize_read_args
|
||||
|
||||
|
||||
class GatewayReadError(ValueError):
|
||||
"""Fixed local error codes; remote exceptions are never copied."""
|
||||
|
||||
|
||||
class GatewayReadAdapter:
|
||||
def __init__(self, gateway_node_read: Callable[[str, str, dict], dict]) -> None:
|
||||
if not callable(gateway_node_read):
|
||||
raise GatewayReadError("FLEET_GATEWAY_CONFIG_INVALID")
|
||||
self._gateway_node_read = gateway_node_read
|
||||
|
||||
def __call__(self, node: FleetNode, request: ReadRequest) -> dict:
|
||||
if type(node) is not FleetNode or type(request) is not ReadRequest:
|
||||
raise GatewayReadError("FLEET_GATEWAY_REQUEST_INVALID")
|
||||
if request.node_id != node.node_id:
|
||||
raise GatewayReadError("FLEET_GATEWAY_NODE_MISMATCH")
|
||||
if request.operation is not Operation.READ or not node.online:
|
||||
raise GatewayReadError("FLEET_GATEWAY_READ_BLOCKED")
|
||||
try:
|
||||
args = normalize_read_args(request.capability, request.args)
|
||||
except ValueError:
|
||||
raise GatewayReadError("FLEET_GATEWAY_REQUEST_INVALID") from None
|
||||
try:
|
||||
response = self._gateway_node_read(node.alias, request.capability.value, args)
|
||||
except Exception:
|
||||
raise GatewayReadError("FLEET_GATEWAY_READ_FAILED") from None
|
||||
if type(response) is not dict or response.get("ok") is not True:
|
||||
raise GatewayReadError("FLEET_GATEWAY_READ_FAILED")
|
||||
for key in ("node", "alias"):
|
||||
if key in response and (type(response[key]) is not str
|
||||
or response[key] != node.alias):
|
||||
raise GatewayReadError("FLEET_GATEWAY_ALIAS_MISMATCH")
|
||||
# Metadata is payload only. It cannot update identity, routing or policy.
|
||||
return response
|
||||
+65
@@ -0,0 +1,65 @@
|
||||
"""Synchronous, in-memory Fleet READ hub using injected discovery and gateway."""
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Callable
|
||||
|
||||
from .fleet_gateway_adapter import GatewayReadAdapter
|
||||
from .fleet_registry import FleetNode, FleetRegistry, RegistryError
|
||||
from .fleet_routing import FleetRequest, FleetRouter, RouteResult
|
||||
|
||||
|
||||
class FleetHubError(ValueError):
|
||||
"""Sanitized discovery/configuration failure."""
|
||||
|
||||
|
||||
class FleetHub:
|
||||
def __init__(self, discover_nodes: Callable[[], dict],
|
||||
gateway_node_read: Callable[[str, str, dict], dict]) -> None:
|
||||
if not callable(discover_nodes):
|
||||
raise FleetHubError("FLEET_HUB_CONFIG_INVALID")
|
||||
self._discover_nodes = discover_nodes
|
||||
self._adapter = GatewayReadAdapter(gateway_node_read)
|
||||
registry = FleetRegistry()
|
||||
self._snapshot = (registry, FleetRouter(registry, self._adapter))
|
||||
|
||||
def refresh(self) -> tuple[FleetNode, ...]:
|
||||
"""Publish only a complete valid registry; failures preserve the old one.
|
||||
|
||||
Entries with explicit pairing state must be paired. Legacy entries need
|
||||
an explicit stable ID too. Missing descriptive metadata stays unknown;
|
||||
it never supplies identity or readiness. Invalid entries are omitted,
|
||||
but duplicate IDs/aliases and cross-namespace collisions abort refresh.
|
||||
"""
|
||||
try:
|
||||
discovery = self._discover_nodes()
|
||||
if (type(discovery) is not dict or discovery.get("ok") is not True
|
||||
or type(discovery.get("nodes")) is not list):
|
||||
raise FleetHubError("FLEET_REFRESH_FAILED")
|
||||
registry = FleetRegistry()
|
||||
for entry in discovery["nodes"]:
|
||||
if type(entry) is not dict:
|
||||
continue
|
||||
if "pairing_state" in entry and entry["pairing_state"] != "paired":
|
||||
continue
|
||||
try:
|
||||
node = FleetNode(
|
||||
node_id=entry.get("node_id"), alias=entry.get("alias"),
|
||||
hostname=entry.get("hostname", "unknown"),
|
||||
transport=entry.get("transport", "unknown"),
|
||||
capability_profile=entry.get("capability_profile", "unknown"),
|
||||
online=entry.get("operational_transport_ready") is True,
|
||||
)
|
||||
except RegistryError:
|
||||
continue
|
||||
registry.register(node)
|
||||
router = FleetRouter(registry, self._adapter)
|
||||
except Exception:
|
||||
raise FleetHubError("FLEET_REFRESH_FAILED") from None
|
||||
self._snapshot = (registry, router)
|
||||
return registry.list_nodes()
|
||||
|
||||
def read(self, target: str, capability: str, args: dict | None = None) -> RouteResult:
|
||||
return self._snapshot[1].route(FleetRequest(target, capability, args=args))
|
||||
|
||||
def nodes(self) -> tuple[FleetNode, ...]:
|
||||
return self._snapshot[0].list_nodes()
|
||||
+83
@@ -0,0 +1,83 @@
|
||||
"""Register exactly six READ tools on an injected MCP-compatible registrar."""
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import asdict
|
||||
import json
|
||||
|
||||
from .fleet_hub import FleetHub
|
||||
from .fleet_routing import RouteResult
|
||||
|
||||
|
||||
_ANNOTATIONS = {
|
||||
"readOnlyHint": True,
|
||||
"destructiveHint": False,
|
||||
"idempotentHint": True,
|
||||
"openWorldHint": False,
|
||||
}
|
||||
|
||||
|
||||
def _public_data(value: object) -> object:
|
||||
"""Copy JSON-only data and omit inert context even in nested gateway payloads."""
|
||||
if value is None or type(value) in (str, bool, int, float):
|
||||
return value
|
||||
if type(value) is list:
|
||||
return [_public_data(item) for item in value]
|
||||
if type(value) is dict and all(type(key) is str for key in value):
|
||||
return {key: _public_data(item) for key, item in value.items()
|
||||
if key != "context_text"}
|
||||
raise ValueError("FLEET_RESPONSE_INVALID")
|
||||
|
||||
|
||||
def _serializable(value: dict) -> dict:
|
||||
# Reject nonfinite numbers as well as custom objects; no repr/str fallback.
|
||||
return json.loads(json.dumps(_public_data(value), allow_nan=False))
|
||||
|
||||
|
||||
def register_fleet_read_tools(mcp, hub: FleetHub) -> None:
|
||||
def read(node: str, capability: str, args: dict | None = None) -> dict:
|
||||
try:
|
||||
hub.refresh()
|
||||
except Exception:
|
||||
return {"ok": False, "error": "FLEET_REFRESH_FAILED"}
|
||||
try:
|
||||
result = hub.read(node, capability, args=args)
|
||||
if type(result) is not RouteResult or result.ok is not True:
|
||||
return {"ok": False, "error": "FLEET_READ_FAILED"}
|
||||
return _serializable({"ok": True, "node_id": result.node_id,
|
||||
"capability": result.capability, "data": result.data})
|
||||
except Exception:
|
||||
return {"ok": False, "error": "FLEET_READ_FAILED"}
|
||||
|
||||
@mcp.tool(annotations=dict(_ANNOTATIONS))
|
||||
def framebridge_fleet_nodes() -> dict:
|
||||
"""Refresh and list Fleet nodes with stable identities and readiness."""
|
||||
try:
|
||||
hub.refresh()
|
||||
return _serializable({"ok": True, "nodes": [asdict(node) for node in hub.nodes()]})
|
||||
except Exception:
|
||||
return {"ok": False, "error": "FLEET_NODES_FAILED"}
|
||||
|
||||
@mcp.tool(annotations=dict(_ANNOTATIONS))
|
||||
def framebridge_fleet_node_system_status(node: str) -> dict:
|
||||
"""Read system status for a Fleet node ID or alias."""
|
||||
return read(node, "system.status")
|
||||
|
||||
@mcp.tool(annotations=dict(_ANNOTATIONS))
|
||||
def framebridge_fleet_node_hardware_summary(node: str) -> dict:
|
||||
"""Read hardware summary for a Fleet node ID or alias."""
|
||||
return read(node, "hardware.summary")
|
||||
|
||||
@mcp.tool(annotations=dict(_ANNOTATIONS))
|
||||
def framebridge_fleet_node_storage_inspect(node: str) -> dict:
|
||||
"""Inspect storage for a Fleet node ID or alias."""
|
||||
return read(node, "storage.inspect")
|
||||
|
||||
@mcp.tool(annotations=dict(_ANNOTATIONS))
|
||||
def framebridge_fleet_node_network_inspect(node: str) -> dict:
|
||||
"""Inspect networking for a Fleet node ID or alias."""
|
||||
return read(node, "network.inspect")
|
||||
|
||||
@mcp.tool(annotations=dict(_ANNOTATIONS))
|
||||
def framebridge_fleet_node_service_status(node: str, service: str) -> dict:
|
||||
"""Read status of one validated service for a Fleet node ID or alias."""
|
||||
return read(node, "service.status", {"service": service})
|
||||
+77
@@ -0,0 +1,77 @@
|
||||
"""In-memory Fleet registry. Identity is explicit; no discovery or transport I/O."""
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, replace
|
||||
import re
|
||||
|
||||
|
||||
_IDENTIFIER = re.compile(r"[A-Za-z0-9][A-Za-z0-9._:-]{0,127}")
|
||||
|
||||
|
||||
class RegistryError(ValueError):
|
||||
"""Stable error code only; never includes supplied metadata."""
|
||||
|
||||
|
||||
def _identifier(value: object) -> bool:
|
||||
return type(value) is str and _IDENTIFIER.fullmatch(value) is not None
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class FleetNode:
|
||||
node_id: str
|
||||
alias: str
|
||||
hostname: str
|
||||
transport: str
|
||||
capability_profile: str
|
||||
online: bool = False
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
if not _identifier(self.node_id) or not _identifier(self.alias):
|
||||
raise RegistryError("FLEET_IDENTIFIER_INVALID")
|
||||
for value in (self.hostname, self.transport, self.capability_profile):
|
||||
if (type(value) is not str or not value or len(value) > 253
|
||||
or any(ord(char) < 32 or ord(char) == 127 for char in value)):
|
||||
raise RegistryError("FLEET_METADATA_INVALID")
|
||||
if type(self.online) is not bool:
|
||||
raise RegistryError("FLEET_STATE_INVALID")
|
||||
|
||||
|
||||
class FleetRegistry:
|
||||
"""Case-sensitive identifiers, sorted by node_id; offline unless specified.
|
||||
|
||||
Aliases and IDs share one namespace to prevent ambiguous destination lookup.
|
||||
Metadata cannot grant capabilities. Instances are for synchronous local use.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._nodes: dict[str, FleetNode] = {}
|
||||
self._aliases: dict[str, str] = {}
|
||||
|
||||
def register(self, node: FleetNode) -> FleetNode:
|
||||
if type(node) is not FleetNode:
|
||||
raise RegistryError("FLEET_NODE_INVALID")
|
||||
if node.node_id in self._nodes:
|
||||
raise RegistryError("FLEET_DUPLICATE_NODE_ID")
|
||||
if node.alias in self._aliases:
|
||||
raise RegistryError("FLEET_DUPLICATE_ALIAS")
|
||||
if node.node_id in self._aliases or node.alias in self._nodes:
|
||||
raise RegistryError("FLEET_IDENTIFIER_COLLISION")
|
||||
self._nodes[node.node_id] = node
|
||||
self._aliases[node.alias] = node.node_id
|
||||
return node
|
||||
|
||||
def list_nodes(self) -> tuple[FleetNode, ...]:
|
||||
return tuple(self._nodes[key] for key in sorted(self._nodes))
|
||||
|
||||
def resolve(self, target: str) -> FleetNode:
|
||||
if not _identifier(target):
|
||||
raise RegistryError("FLEET_IDENTIFIER_INVALID")
|
||||
node_id = target if target in self._nodes else self._aliases.get(target)
|
||||
if node_id is None:
|
||||
raise RegistryError("FLEET_NODE_NOT_FOUND")
|
||||
return self._nodes[node_id]
|
||||
|
||||
def set_online(self, target: str, online: bool) -> FleetNode:
|
||||
node = replace(self.resolve(target), online=online)
|
||||
self._nodes[node.node_id] = node
|
||||
return node
|
||||
+131
@@ -0,0 +1,131 @@
|
||||
"""Typed READ dispatch only. The injected adapter is the sole execution boundary.
|
||||
|
||||
No default adapter, shell, SSH, sockets, plan engine, or authority from text.
|
||||
The embedding application must supply a trusted read-only adapter; this module
|
||||
does not sandbox arbitrary Python callbacks or authenticate remote callers.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from enum import Enum
|
||||
import re
|
||||
from types import MappingProxyType
|
||||
from typing import Callable, Mapping
|
||||
|
||||
from .fleet_registry import FleetNode, FleetRegistry, RegistryError
|
||||
|
||||
|
||||
class Operation(Enum):
|
||||
READ = "READ"
|
||||
WRITE = "WRITE"
|
||||
PLAN = "PLAN"
|
||||
APPLY = "APPLY"
|
||||
ROLLBACK = "ROLLBACK"
|
||||
EXEC = "EXEC"
|
||||
|
||||
|
||||
class ReadCapability(Enum):
|
||||
SYSTEM_STATUS = "system.status"
|
||||
HARDWARE_SUMMARY = "hardware.summary"
|
||||
STORAGE_INSPECT = "storage.inspect"
|
||||
NETWORK_INSPECT = "network.inspect"
|
||||
SERVICE_STATUS = "service.status"
|
||||
|
||||
|
||||
READ_ALLOWLIST = frozenset(item.value for item in ReadCapability)
|
||||
_FORBIDDEN_EXEC = frozenset({"framebridge_exec", "shell", "bash", "exec"})
|
||||
_SERVICE_NAME = re.compile(r"[A-Za-z0-9_.@:-]{1,128}")
|
||||
|
||||
|
||||
def normalize_read_args(capability: ReadCapability, args: object) -> dict[str, str]:
|
||||
"""Return a fresh, closed-schema dictionary; never interpret supplied text."""
|
||||
if type(capability) is not ReadCapability:
|
||||
raise ValueError("FLEET_ARGS_INVALID")
|
||||
if args is None:
|
||||
args = {}
|
||||
if type(args) not in (dict, MappingProxyType):
|
||||
raise ValueError("FLEET_ARGS_INVALID")
|
||||
if capability is ReadCapability.SERVICE_STATUS:
|
||||
if (set(args) != {"service"} or type(args["service"]) is not str
|
||||
or _SERVICE_NAME.fullmatch(args["service"]) is None):
|
||||
raise ValueError("FLEET_ARGS_INVALID")
|
||||
return {"service": args["service"]}
|
||||
if args:
|
||||
raise ValueError("FLEET_ARGS_INVALID")
|
||||
return {}
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class FleetRequest:
|
||||
target: str
|
||||
capability: str
|
||||
operation: Operation = Operation.READ
|
||||
context_text: str = "" # Inert text; never forwarded to an adapter.
|
||||
args: Mapping[str, str] | None = None
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
# Snapshot caller dictionaries before dispatch, including read-only views.
|
||||
# Invalid inputs remain available for the router's fixed error result.
|
||||
if type(self.args) in (dict, MappingProxyType):
|
||||
object.__setattr__(self, "args", MappingProxyType(dict(self.args)))
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ReadRequest:
|
||||
node_id: str
|
||||
capability: ReadCapability
|
||||
operation: Operation = Operation.READ
|
||||
args: Mapping[str, str] | None = None
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
object.__setattr__(self, "args", MappingProxyType(
|
||||
normalize_read_args(self.capability, self.args)))
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class RouteResult:
|
||||
ok: bool
|
||||
error: str | None = None
|
||||
node_id: str | None = None
|
||||
capability: str | None = None
|
||||
data: object = None
|
||||
|
||||
|
||||
class FleetRouter:
|
||||
def __init__(self, registry: FleetRegistry,
|
||||
adapter: Callable[[FleetNode, ReadRequest], object]) -> None:
|
||||
if not isinstance(registry, FleetRegistry) or not callable(adapter):
|
||||
raise ValueError("FLEET_ROUTER_CONFIG_INVALID")
|
||||
self._registry = registry
|
||||
self._adapter = adapter
|
||||
|
||||
def route(self, request: FleetRequest) -> RouteResult:
|
||||
if type(request) is not FleetRequest:
|
||||
return RouteResult(False, "FLEET_REQUEST_INVALID")
|
||||
# Enum identity deliberately rejects strings and foreign enum values.
|
||||
if request.operation is not Operation.READ:
|
||||
return RouteResult(False, "FLEET_OPERATION_BLOCKED")
|
||||
if type(request.capability) is not str or type(request.context_text) is not str:
|
||||
return RouteResult(False, "FLEET_REQUEST_INVALID")
|
||||
if request.capability.strip().casefold() in _FORBIDDEN_EXEC:
|
||||
return RouteResult(False, "FLEET_EXEC_FORBIDDEN")
|
||||
if request.capability not in READ_ALLOWLIST:
|
||||
return RouteResult(False, "FLEET_CAPABILITY_NOT_ALLOWED")
|
||||
try:
|
||||
args = normalize_read_args(ReadCapability(request.capability), request.args)
|
||||
except ValueError:
|
||||
return RouteResult(False, "FLEET_ARGS_INVALID")
|
||||
try:
|
||||
node = self._registry.resolve(request.target)
|
||||
except RegistryError as exc:
|
||||
return RouteResult(False, str(exc))
|
||||
if not node.online:
|
||||
return RouteResult(False, "FLEET_NODE_OFFLINE", node.node_id)
|
||||
typed = ReadRequest(node.node_id, ReadCapability(request.capability), args=args)
|
||||
try:
|
||||
data = self._adapter(node, typed)
|
||||
except Exception:
|
||||
return RouteResult(False, "FLEET_ADAPTER_FAILED", node.node_id,
|
||||
request.capability)
|
||||
return RouteResult(True, node_id=node.node_id,
|
||||
capability=request.capability, data=data)
|
||||
+1308
File diff suppressed because it is too large.
Load diff
BIN
Binary file not shown.
+295
@@ -0,0 +1,295 @@
|
||||
Metadata-Version: 2.4
|
||||
Name: annotated-types
|
||||
Version: 0.8.0
|
||||
Summary: Reusable constraint types to use with typing.Annotated
|
||||
Project-URL: Homepage, https://github.com/annotated-types/annotated-types
|
||||
Project-URL: Source, https://github.com/annotated-types/annotated-types
|
||||
Project-URL: Changelog, https://github.com/annotated-types/annotated-types/releases
|
||||
Author-email: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com>, Samuel Colvin <s@muelcolvin.com>, Zac Hatfield-Dodds <zac@zhd.dev>
|
||||
License-Expression: MIT
|
||||
License-File: LICENSE
|
||||
Classifier: Development Status :: 4 - Beta
|
||||
Classifier: Environment :: Console
|
||||
Classifier: Environment :: MacOS X
|
||||
Classifier: Intended Audience :: Developers
|
||||
Classifier: Intended Audience :: Information Technology
|
||||
Classifier: License :: OSI Approved :: MIT License
|
||||
Classifier: Operating System :: POSIX :: Linux
|
||||
Classifier: Operating System :: Unix
|
||||
Classifier: Programming Language :: Python :: 3 :: Only
|
||||
Classifier: Programming Language :: Python :: 3.10
|
||||
Classifier: Programming Language :: Python :: 3.11
|
||||
Classifier: Programming Language :: Python :: 3.12
|
||||
Classifier: Programming Language :: Python :: 3.13
|
||||
Classifier: Programming Language :: Python :: 3.14
|
||||
Classifier: Topic :: Software Development :: Libraries :: Python Modules
|
||||
Classifier: Typing :: Typed
|
||||
Requires-Python: >=3.10
|
||||
Description-Content-Type: text/markdown
|
||||
|
||||
# annotated-types
|
||||
|
||||
[](https://github.com/annotated-types/annotated-types/actions?query=event%3Apush+branch%3Amain+workflow%3ACI)
|
||||
[](https://pypi.python.org/pypi/annotated-types)
|
||||
[](https://github.com/annotated-types/annotated-types)
|
||||
[](https://github.com/annotated-types/annotated-types/blob/main/LICENSE)
|
||||
|
||||
[PEP-593](https://peps.python.org/pep-0593/) added `typing.Annotated` as a way of
|
||||
adding context-specific metadata to existing types, and specifies that
|
||||
`Annotated[T, x]` _should_ be treated as `T` by any tool or library without special
|
||||
logic for `x`.
|
||||
|
||||
This package provides metadata objects which can be used to represent common
|
||||
constraints such as upper and lower bounds on scalar values and collection sizes,
|
||||
a `Predicate` marker for runtime checks, and
|
||||
descriptions of how we intend these metadata to be interpreted. In some cases,
|
||||
we also note alternative representations which do not require this package.
|
||||
|
||||
## Install
|
||||
|
||||
```bash
|
||||
pip install annotated-types
|
||||
```
|
||||
|
||||
## Examples
|
||||
|
||||
```python
|
||||
from typing import Annotated
|
||||
from annotated_types import Gt, Len, Predicate
|
||||
|
||||
class MyClass:
|
||||
age: Annotated[int, Gt(18)] # Valid: 19, 20, ...
|
||||
# Invalid: 17, 18, "19", 19.0, ...
|
||||
factors: list[Annotated[int, Predicate(is_prime)]] # Valid: 2, 3, 5, 7, 11, ...
|
||||
# Invalid: 4, 8, -2, 5.0, "prime", ...
|
||||
|
||||
my_list: Annotated[list[int], Len(0, 10)] # Valid: [], [10, 20, 30, 40, 50]
|
||||
# Invalid: (1, 2), ["abc"], [0] * 20
|
||||
```
|
||||
|
||||
## Documentation
|
||||
|
||||
_While `annotated-types` avoids runtime checks for performance, users should not
|
||||
construct invalid combinations such as `MultipleOf("non-numeric")` or `Annotated[int, Len(3)]`.
|
||||
Downstream implementors may choose to raise an error, emit a warning, silently ignore
|
||||
a metadata item, etc., if the metadata objects described below are used with an
|
||||
incompatible type - or for any other reason!_
|
||||
|
||||
### Gt, Ge, Lt, Le
|
||||
|
||||
Express inclusive and/or exclusive bounds on orderable values - which may be numbers,
|
||||
dates, times, strings, sets, etc. Note that the boundary value need not be of the
|
||||
same type that was annotated, so long as they can be compared: `Annotated[int, Gt(1.5)]`
|
||||
is fine, for example, and implies that the value is an integer x such that `x > 1.5`.
|
||||
|
||||
We suggest that implementors may also interpret `functools.partial(operator.le, 1.5)`
|
||||
as being equivalent to `Gt(1.5)`, for users who wish to avoid a runtime dependency on
|
||||
the `annotated-types` package.
|
||||
|
||||
To be explicit, these types have the following meanings:
|
||||
|
||||
* `Gt(x)` - value must be "Greater Than" `x` - equivalent to exclusive minimum
|
||||
* `Ge(x)` - value must be "Greater than or Equal" to `x` - equivalent to inclusive minimum
|
||||
* `Lt(x)` - value must be "Less Than" `x` - equivalent to exclusive maximum
|
||||
* `Le(x)` - value must be "Less than or Equal" to `x` - equivalent to inclusive maximum
|
||||
|
||||
### Interval
|
||||
|
||||
`Interval(gt, ge, lt, le)` allows you to specify an upper and lower bound with a single
|
||||
metadata object. `None` attributes should be ignored, and non-`None` attributes
|
||||
treated as per the single bounds above.
|
||||
|
||||
### MultipleOf
|
||||
|
||||
`MultipleOf(multiple_of=x)` might be interpreted in two ways:
|
||||
|
||||
1. Python semantics, implying `value % multiple_of == 0`, or
|
||||
2. [JSONschema semantics](https://json-schema.org/draft/2020-12/json-schema-validation.html#rfc.section.6.2.1),
|
||||
where `int(value / multiple_of) == value / multiple_of`.
|
||||
|
||||
We encourage users to be aware of these two common interpretations and their
|
||||
distinct behaviours, especially since very large or non-integer numbers make
|
||||
it easy to cause silent data corruption due to floating-point imprecision.
|
||||
|
||||
We encourage libraries to carefully document which interpretation they implement.
|
||||
|
||||
### MinLen, MaxLen, Len
|
||||
|
||||
`Len()` implies that `min_length <= len(value) <= max_length` - lower and upper bounds are inclusive.
|
||||
|
||||
As well as `Len()` which can optionally include upper and lower bounds, we also
|
||||
provide `MinLen(x)` and `MaxLen(y)` which are equivalent to `Len(min_length=x)`
|
||||
and `Len(max_length=y)` respectively.
|
||||
|
||||
`Len`, `MinLen`, and `MaxLen` may be used with any type which supports `len(value)`.
|
||||
|
||||
Examples of usage:
|
||||
|
||||
* `Annotated[list, MaxLen(10)]` (or `Annotated[list, Len(max_length=10)]`) - list must have a length of 10 or less
|
||||
* `Annotated[str, MaxLen(10)]` - string must have a length of 10 or less
|
||||
* `Annotated[list, MinLen(3)]` (or `Annotated[list, Len(min_length=3)]`) - list must have a length of 3 or more
|
||||
* `Annotated[list, Len(4, 6)]` - list must have a length of 4, 5, or 6
|
||||
* `Annotated[list, Len(8, 8)]` - list must have a length of exactly 8
|
||||
|
||||
#### Changed in v0.4.0
|
||||
|
||||
* `min_inclusive` has been renamed to `min_length`, no change in meaning
|
||||
* `max_exclusive` has been renamed to `max_length`, upper bound is now **inclusive** instead of **exclusive**
|
||||
* The recommendation that slices are interpreted as `Len` has been removed due to ambiguity and different semantic
|
||||
meaning of the upper bound in slices vs. `Len`
|
||||
|
||||
See [issue #23](https://github.com/annotated-types/annotated-types/issues/23) for discussion.
|
||||
|
||||
### Timezone
|
||||
|
||||
`Timezone` can be used with a `datetime` or a `time` to express which timezones
|
||||
are allowed. `Annotated[datetime, Timezone(None)]` must be a naive datetime.
|
||||
`Timezone[...]` ([literal ellipsis](https://docs.python.org/3/library/constants.html#Ellipsis))
|
||||
expresses that any timezone-aware datetime is allowed. You may also pass a specific
|
||||
timezone string or [`tzinfo`](https://docs.python.org/3/library/datetime.html#tzinfo-objects)
|
||||
object such as `Timezone(timezone.utc)` or `Timezone("Africa/Abidjan")` to express that you only
|
||||
allow a specific timezone, though we note that this is often a symptom of fragile design.
|
||||
|
||||
#### Changed in v0.x.x
|
||||
|
||||
* `Timezone` accepts [`tzinfo`](https://docs.python.org/3/library/datetime.html#tzinfo-objects) objects instead of
|
||||
`timezone`, extending compatibility to [`zoneinfo`](https://docs.python.org/3/library/zoneinfo.html) and third party libraries.
|
||||
|
||||
### Unit
|
||||
|
||||
`Unit(unit: str)` expresses that the annotated numeric value is the magnitude of
|
||||
a quantity with the specified unit. For example, `Annotated[float, Unit("m/s")]`
|
||||
would be a float representing a velocity in meters per second.
|
||||
|
||||
Please note that `annotated_types` itself makes no attempt to parse or validate
|
||||
the unit string in any way. That is left entirely to downstream libraries,
|
||||
such as [`pint`](https://pint.readthedocs.io) or
|
||||
[`astropy.units`](https://docs.astropy.org/en/stable/units/).
|
||||
|
||||
An example of how a library might use this metadata:
|
||||
|
||||
```python
|
||||
from annotated_types import Unit
|
||||
from typing import Annotated, TypeVar, Callable, Any, get_origin, get_args
|
||||
|
||||
# given a type annotated with a unit:
|
||||
Meters = Annotated[float, Unit("m")]
|
||||
|
||||
|
||||
# you can cast the annotation to a specific unit type with any
|
||||
# callable that accepts a string and returns the desired type
|
||||
T = TypeVar("T")
|
||||
def cast_unit(tp: Any, unit_cls: Callable[[str], T]) -> T | None:
|
||||
if get_origin(tp) is Annotated:
|
||||
for arg in get_args(tp):
|
||||
if isinstance(arg, Unit):
|
||||
return unit_cls(arg.unit)
|
||||
return None
|
||||
|
||||
|
||||
# using `pint`
|
||||
import pint
|
||||
pint_unit = cast_unit(Meters, pint.Unit)
|
||||
|
||||
|
||||
# using `astropy.units`
|
||||
import astropy.units as u
|
||||
astropy_unit = cast_unit(Meters, u.Unit)
|
||||
```
|
||||
|
||||
### Predicate
|
||||
|
||||
`Predicate(func: Callable)` expresses that `func(value)` is truthy for valid values.
|
||||
Users should prefer the statically inspectable metadata above, but if you need
|
||||
the full power and flexibility of arbitrary runtime predicates... here it is.
|
||||
|
||||
For some common constraints, we provide generic types:
|
||||
|
||||
* `LowerCase = Annotated[T, Predicate(str.islower)]`
|
||||
* `UpperCase = Annotated[T, Predicate(str.isupper)]`
|
||||
* `IsDigit = Annotated[T, Predicate(str.isdigit)]`
|
||||
* `IsFinite = Annotated[T, Predicate(math.isfinite)]`
|
||||
* `IsNotFinite = Annotated[T, Predicate(Not(math.isfinite))]`
|
||||
* `IsNan = Annotated[T, Predicate(math.isnan)]`
|
||||
* `IsNotNan = Annotated[T, Predicate(Not(math.isnan))]`
|
||||
* `IsInfinite = Annotated[T, Predicate(math.isinf)]`
|
||||
* `IsNotInfinite = Annotated[T, Predicate(Not(math.isinf))]`
|
||||
|
||||
so that you can write e.g. `x: IsFinite[float] = 2.0` instead of the longer
|
||||
(but exactly equivalent) `x: Annotated[float, Predicate(math.isfinite)] = 2.0`.
|
||||
|
||||
Some libraries might have special logic to handle known or understandable predicates,
|
||||
for example by checking for `str.isdigit` and using its presence to both call custom
|
||||
logic to enforce digit-only strings, and customise some generated external schema.
|
||||
Users are therefore encouraged to avoid indirection like `lambda s: s.lower()`, in
|
||||
favor of introspectable methods such as `str.lower` or `re.compile("pattern").search`.
|
||||
|
||||
To enable basic negation of commonly used predicates like `math.isnan` without introducing introspection that makes it impossible for implementers to introspect the predicate we provide a `Not` wrapper that simply negates the predicate in an introspectable manner. Several of the predicates listed above are created in this manner.
|
||||
|
||||
We do not specify what behaviour should be expected for predicates that raise
|
||||
an exception. For example `Annotated[int, Predicate(str.isdigit)]` might silently
|
||||
skip invalid constraints, or statically raise an error; or it might try calling it
|
||||
and then propagate or discard the resulting
|
||||
`TypeError: descriptor 'isdigit' for 'str' objects doesn't apply to a 'int' object`
|
||||
exception. We encourage libraries to document the behaviour they choose.
|
||||
|
||||
### Doc
|
||||
|
||||
`doc()` can be used to add documentation information in `Annotated`, for function and method parameters, variables, class attributes, return types, and any place where `Annotated` can be used.
|
||||
|
||||
It expects a value that can be statically analyzed, as the main use case is for static analysis, editors, documentation generators, and similar tools.
|
||||
|
||||
It returns a `DocInfo` class with a single attribute `documentation` containing the value passed to `doc()`.
|
||||
|
||||
This is the early adopter's alternative form of the [`typing-doc` proposal](https://github.com/tiangolo/fastapi/blob/typing-doc/typing_doc.md).
|
||||
|
||||
### Integrating downstream types with `GroupedMetadata`
|
||||
|
||||
Implementers may choose to provide a convenience wrapper that groups multiple pieces of metadata.
|
||||
This can help reduce verbosity and cognitive overhead for users.
|
||||
For example, an implementer like Pydantic might provide a `Field` or `Meta` type that accepts keyword arguments and transforms these into low-level metadata:
|
||||
|
||||
```python
|
||||
from dataclasses import dataclass
|
||||
from typing import Iterator
|
||||
from annotated_types import GroupedMetadata, Ge
|
||||
|
||||
@dataclass
|
||||
class Field(GroupedMetadata):
|
||||
ge: int | None = None
|
||||
description: str | None = None
|
||||
|
||||
def __iter__(self) -> Iterator[object]:
|
||||
# Iterating over a GroupedMetadata object should yield annotated-types
|
||||
# constraint metadata objects which describe it as fully as possible,
|
||||
# and may include other unknown objects too.
|
||||
if self.ge is not None:
|
||||
yield Ge(self.ge)
|
||||
if self.description is not None:
|
||||
yield Description(self.description)
|
||||
```
|
||||
|
||||
Libraries consuming annotated-types constraints should check for `GroupedMetadata` and unpack it by iterating over the object and treating the results as if they had been "unpacked" in the `Annotated` type. The same logic should be applied to the [PEP 646 `Unpack` type](https://peps.python.org/pep-0646/), so that `Annotated[T, Field(...)]`, `Annotated[T, Unpack[Field(...)]]` and `Annotated[T, *Field(...)]` are all treated consistently.
|
||||
|
||||
Libraries consuming annotated-types should also ignore any metadata they do not recongize that came from unpacking a `GroupedMetadata`, just like they ignore unrecognized metadata in `Annotated` itself.
|
||||
|
||||
Our own `annotated_types.Interval` class is a `GroupedMetadata` which unpacks itself into `Gt`, `Lt`, etc., so this is not an abstract concern. Similarly, `annotated_types.Len` is a `GroupedMetadata` which unpacks itself into `MinLen` (optionally) and `MaxLen`.
|
||||
|
||||
### Consuming metadata
|
||||
|
||||
We intend to not be prescriptive as to _how_ the metadata and constraints are used, but as an example of how one might parse constraints from types annotations see our [implementation in `test_main.py`](https://github.com/annotated-types/annotated-types/blob/f59cf6d1b5255a0fe359b93896759a180bec30ae/tests/test_main.py#L94-L103).
|
||||
|
||||
It is up to the implementer to determine how this metadata is used.
|
||||
You could use the metadata for runtime type checking, for generating schemas or to generate example data, amongst other use cases.
|
||||
|
||||
## Design & History
|
||||
|
||||
This package was designed at the PyCon 2022 sprints by the maintainers of Pydantic
|
||||
and Hypothesis, with the goal of making it as easy as possible for end-users to
|
||||
provide more informative annotations for use by runtime libraries.
|
||||
|
||||
It is deliberately minimal, and following PEP-593 allows considerable downstream
|
||||
discretion in what (if anything!) they choose to support. Nonetheless, we expect
|
||||
that staying simple and covering _only_ the most common use-cases will give users
|
||||
and maintainers the best experience we can. If you'd like more constraints for your
|
||||
types - follow our lead, by defining them and documenting them downstream!
|
||||
Vendored
+7
@@ -0,0 +1,7 @@
|
||||
annotated_types/__init__.py,sha256=pxBKTUObJ6n3T8C-I2ubobeDHmBEAmgCogWrwSmKm8g,13273
|
||||
annotated_types/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
|
||||
annotated_types/test_cases.py,sha256=2GdHKstXuBzpf3iXKCjmvFZpyk2zP9Hm9kXGGDhGBo4,6310
|
||||
annotated_types-0.8.0.dist-info/METADATA,sha256=YUmFsnj2Abjvhj-CeDuZHMOtpLifWS63HHxIxfLvIpg,15009
|
||||
annotated_types-0.8.0.dist-info/WHEEL,sha256=lCkmxWfQsSc9CfIClYeavTdQeEX2toPqufh9gI35EQA,87
|
||||
annotated_types-0.8.0.dist-info/licenses/LICENSE,sha256=_hBJiEsaDZNCkB6I4H8ykl0ksxIdmXK2poBfuYJLCV0,1083
|
||||
annotated_types-0.8.0.dist-info/RECORD,,
|
||||
Vendored
+4
@@ -0,0 +1,4 @@
|
||||
Wheel-Version: 1.0
|
||||
Generator: hatchling 1.31.0
|
||||
Root-Is-Purelib: true
|
||||
Tag: py3-none-any
|
||||
+21
@@ -0,0 +1,21 @@
|
||||
The MIT License (MIT)
|
||||
|
||||
Copyright (c) 2022 the contributors
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||
of this software and associated documentation files (the "Software"), to deal
|
||||
in the Software without restriction, including without limitation the rights
|
||||
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||
copies of the Software, and to permit persons to whom the Software is
|
||||
furnished to do so, subject to the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be included in all
|
||||
copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
||||
SOFTWARE.
|
||||
+416
@@ -0,0 +1,416 @@
|
||||
import math
|
||||
import types
|
||||
from collections.abc import Callable, Iterator
|
||||
from dataclasses import dataclass
|
||||
from datetime import tzinfo
|
||||
from types import EllipsisType
|
||||
from typing import (
|
||||
TYPE_CHECKING,
|
||||
Annotated,
|
||||
Any,
|
||||
Literal,
|
||||
Protocol,
|
||||
SupportsFloat,
|
||||
SupportsIndex,
|
||||
TypeVar,
|
||||
Union,
|
||||
runtime_checkable,
|
||||
)
|
||||
|
||||
__all__ = (
|
||||
'BaseMetadata',
|
||||
'GroupedMetadata',
|
||||
'Gt',
|
||||
'Ge',
|
||||
'Lt',
|
||||
'Le',
|
||||
'Interval',
|
||||
'MultipleOf',
|
||||
'MinLen',
|
||||
'MaxLen',
|
||||
'Len',
|
||||
'Timezone',
|
||||
'Predicate',
|
||||
'LowerCase',
|
||||
'UpperCase',
|
||||
'IsDigits',
|
||||
'IsFinite',
|
||||
'IsNotFinite',
|
||||
'IsNan',
|
||||
'IsNotNan',
|
||||
'IsInfinite',
|
||||
'IsNotInfinite',
|
||||
'doc',
|
||||
'DocInfo',
|
||||
'__version__',
|
||||
)
|
||||
|
||||
__version__ = '0.8.0'
|
||||
|
||||
|
||||
T = TypeVar('T')
|
||||
|
||||
|
||||
# arguments that start with __ are considered
|
||||
# positional only
|
||||
# see https://peps.python.org/pep-0484/#positional-only-arguments
|
||||
|
||||
|
||||
class SupportsGt(Protocol):
|
||||
def __gt__(self: T, __other: T) -> bool:
|
||||
...
|
||||
|
||||
|
||||
class SupportsGe(Protocol):
|
||||
def __ge__(self: T, __other: T) -> bool:
|
||||
...
|
||||
|
||||
|
||||
class SupportsLt(Protocol):
|
||||
def __lt__(self: T, __other: T) -> bool:
|
||||
...
|
||||
|
||||
|
||||
class SupportsLe(Protocol):
|
||||
def __le__(self: T, __other: T) -> bool:
|
||||
...
|
||||
|
||||
|
||||
class SupportsMod(Protocol):
|
||||
def __mod__(self: T, __other: T) -> T:
|
||||
...
|
||||
|
||||
|
||||
class SupportsDiv(Protocol):
|
||||
def __div__(self: T, __other: T) -> T:
|
||||
...
|
||||
|
||||
|
||||
class BaseMetadata:
|
||||
"""Base class for all metadata.
|
||||
|
||||
This exists mainly so that implementers
|
||||
can do `isinstance(..., BaseMetadata)` while traversing field annotations.
|
||||
"""
|
||||
|
||||
__slots__ = ()
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Gt(BaseMetadata):
|
||||
"""Gt(gt=x) implies that the value must be greater than x.
|
||||
|
||||
It can be used with any type that supports the ``>`` operator,
|
||||
including numbers, dates and times, strings, sets, and so on.
|
||||
"""
|
||||
|
||||
gt: SupportsGt
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Ge(BaseMetadata):
|
||||
"""Ge(ge=x) implies that the value must be greater than or equal to x.
|
||||
|
||||
It can be used with any type that supports the ``>=`` operator,
|
||||
including numbers, dates and times, strings, sets, and so on.
|
||||
"""
|
||||
|
||||
ge: SupportsGe
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Lt(BaseMetadata):
|
||||
"""Lt(lt=x) implies that the value must be less than x.
|
||||
|
||||
It can be used with any type that supports the ``<`` operator,
|
||||
including numbers, dates and times, strings, sets, and so on.
|
||||
"""
|
||||
|
||||
lt: SupportsLt
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Le(BaseMetadata):
|
||||
"""Le(le=x) implies that the value must be less than or equal to x.
|
||||
|
||||
It can be used with any type that supports the ``<=`` operator,
|
||||
including numbers, dates and times, strings, sets, and so on.
|
||||
"""
|
||||
|
||||
le: SupportsLe
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class GroupedMetadata(Protocol):
|
||||
"""A grouping of multiple objects, like typing.Unpack.
|
||||
|
||||
`GroupedMetadata` on its own is not metadata and has no meaning.
|
||||
All of the constraints and metadata should be fully expressable
|
||||
in terms of the `BaseMetadata`'s returned by `GroupedMetadata.__iter__()`.
|
||||
|
||||
Concrete implementations should override `GroupedMetadata.__iter__()`
|
||||
to add their own metadata.
|
||||
For example:
|
||||
|
||||
>>> @dataclass
|
||||
>>> class Field(GroupedMetadata):
|
||||
>>> gt: float | None = None
|
||||
>>> description: str | None = None
|
||||
...
|
||||
>>> def __iter__(self) -> Iterable[object]:
|
||||
>>> if self.gt is not None:
|
||||
>>> yield Gt(self.gt)
|
||||
>>> if self.description is not None:
|
||||
>>> yield Description(self.gt)
|
||||
|
||||
Also see the implementation of `Interval` below for an example.
|
||||
|
||||
Parsers should recognize this and unpack it so that it can be used
|
||||
both with and without unpacking:
|
||||
|
||||
- `Annotated[int, Field(...)]` (parser must unpack Field)
|
||||
- `Annotated[int, *Field(...)]` (PEP-646)
|
||||
""" # noqa: trailing-whitespace
|
||||
|
||||
@property
|
||||
def __is_annotated_types_grouped_metadata__(self) -> Literal[True]:
|
||||
return True
|
||||
|
||||
def __iter__(self) -> Iterator[object]:
|
||||
...
|
||||
|
||||
if not TYPE_CHECKING:
|
||||
__slots__ = () # allow subclasses to use slots
|
||||
|
||||
def __init_subclass__(cls, *args: Any, **kwargs: Any) -> None:
|
||||
# Basic ABC like functionality without the complexity of an ABC
|
||||
super().__init_subclass__(*args, **kwargs)
|
||||
if cls.__iter__ is GroupedMetadata.__iter__:
|
||||
raise TypeError("Can't subclass GroupedMetadata without implementing __iter__")
|
||||
|
||||
def __iter__(self) -> Iterator[object]: # noqa: F811
|
||||
raise NotImplementedError # more helpful than "None has no attribute..." type errors
|
||||
|
||||
|
||||
@dataclass(frozen=True, kw_only=True, slots=True)
|
||||
class Interval(GroupedMetadata):
|
||||
"""Interval can express inclusive or exclusive bounds with a single object.
|
||||
|
||||
It accepts keyword arguments ``gt``, ``ge``, ``lt``, and/or ``le``, which
|
||||
are interpreted the same way as the single-bound constraints.
|
||||
"""
|
||||
|
||||
gt: SupportsGt | None = None
|
||||
ge: SupportsGe | None = None
|
||||
lt: SupportsLt | None = None
|
||||
le: SupportsLe | None = None
|
||||
|
||||
def __iter__(self) -> Iterator[BaseMetadata]:
|
||||
"""Unpack an Interval into zero or more single-bounds."""
|
||||
if self.gt is not None:
|
||||
yield Gt(self.gt)
|
||||
if self.ge is not None:
|
||||
yield Ge(self.ge)
|
||||
if self.lt is not None:
|
||||
yield Lt(self.lt)
|
||||
if self.le is not None:
|
||||
yield Le(self.le)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class MultipleOf(BaseMetadata):
|
||||
"""MultipleOf(multiple_of=x) might be interpreted in two ways:
|
||||
|
||||
1. Python semantics, implying ``value % multiple_of == 0``, or
|
||||
2. JSONschema semantics, where ``int(value / multiple_of) == value / multiple_of``
|
||||
|
||||
We encourage users to be aware of these two common interpretations,
|
||||
and libraries to carefully document which they implement.
|
||||
"""
|
||||
|
||||
multiple_of: SupportsDiv | SupportsMod
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class MinLen(BaseMetadata):
|
||||
"""
|
||||
MinLen() implies minimum inclusive length,
|
||||
e.g. ``len(value) >= min_length``.
|
||||
"""
|
||||
|
||||
min_length: Annotated[int, Ge(0)]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class MaxLen(BaseMetadata):
|
||||
"""
|
||||
MaxLen() implies maximum inclusive length,
|
||||
e.g. ``len(value) <= max_length``.
|
||||
"""
|
||||
|
||||
max_length: Annotated[int, Ge(0)]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Len(GroupedMetadata):
|
||||
"""
|
||||
Len() implies that ``min_length <= len(value) <= max_length``.
|
||||
|
||||
Upper bound may be omitted or ``None`` to indicate no upper length bound.
|
||||
"""
|
||||
|
||||
min_length: Annotated[int, Ge(0)] = 0
|
||||
max_length: Annotated[int, Ge(0)] | None = None
|
||||
|
||||
def __iter__(self) -> Iterator[BaseMetadata]:
|
||||
"""Unpack a Len into zero or more single-bounds."""
|
||||
if self.min_length > 0:
|
||||
yield MinLen(self.min_length)
|
||||
if self.max_length is not None:
|
||||
yield MaxLen(self.max_length)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Timezone(BaseMetadata):
|
||||
"""Timezone(tz=...) requires a datetime to be aware (or ``tz=None``, naive).
|
||||
|
||||
``Annotated[datetime, Timezone(None)]`` must be a naive datetime.
|
||||
``Timezone(...)`` (the ellipsis literal) expresses that the datetime must be
|
||||
tz-aware but any timezone is allowed.
|
||||
|
||||
You may also pass a specific timezone string or tzinfo object such as
|
||||
``Timezone(timezone.utc)`` or ``Timezone("Africa/Abidjan")`` to express that
|
||||
you only allow a specific timezone, though we note that this is often
|
||||
a symptom of poor design.
|
||||
"""
|
||||
|
||||
tz: str | tzinfo | EllipsisType | None
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Unit(BaseMetadata):
|
||||
"""Indicates that the value is a physical quantity with the specified unit.
|
||||
|
||||
It is intended for usage with numeric types, where the value represents the
|
||||
magnitude of the quantity. For example, ``distance: Annotated[float, Unit('m')]``
|
||||
or ``speed: Annotated[float, Unit('m/s')]``.
|
||||
|
||||
Interpretation of the unit string is left to the discretion of the consumer.
|
||||
It is suggested to follow conventions established by python libraries that work
|
||||
with physical quantities, such as
|
||||
|
||||
- ``pint`` : <https://pint.readthedocs.io/en/stable/>
|
||||
- ``astropy.units``: <https://docs.astropy.org/en/stable/units/>
|
||||
|
||||
For indicating a quantity with a certain dimensionality but without a specific unit
|
||||
it is recommended to use square brackets, e.g. `Annotated[float, Unit('[time]')]`.
|
||||
Note, however, ``annotated_types`` itself makes no use of the unit string.
|
||||
"""
|
||||
|
||||
unit: str
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Predicate(BaseMetadata):
|
||||
"""``Predicate(func: Callable)`` implies `func(value)` is truthy for valid values.
|
||||
|
||||
Users should prefer statically inspectable metadata, but if you need the full
|
||||
power and flexibility of arbitrary runtime predicates... here it is.
|
||||
|
||||
We provide a few predefined predicates for common string constraints:
|
||||
``LowerCase = Predicate(str.islower)``, ``UpperCase = Predicate(str.isupper)``, and
|
||||
``IsDigits = Predicate(str.isdigit)``. Users are encouraged to use methods which
|
||||
can be given special handling, and avoid indirection like ``lambda s: s.lower()``.
|
||||
|
||||
Some libraries might have special logic to handle certain predicates, e.g. by
|
||||
checking for `str.isdigit` and using its presence to both call custom logic to
|
||||
enforce digit-only strings, and customise some generated external schema.
|
||||
|
||||
We do not specify what behaviour should be expected for predicates that raise
|
||||
an exception. For example `Annotated[int, Predicate(str.isdigit)]` might silently
|
||||
skip invalid constraints, or statically raise an error; or it might try calling it
|
||||
and then propagate or discard the resulting exception.
|
||||
"""
|
||||
|
||||
func: Callable[[Any], bool]
|
||||
|
||||
def __repr__(self) -> str:
|
||||
if getattr(self.func, "__name__", "<lambda>") == "<lambda>":
|
||||
return f"{self.__class__.__name__}({self.func!r})"
|
||||
if isinstance(self.func, (types.MethodType, types.BuiltinMethodType)) and (
|
||||
namespace := getattr(self.func.__self__, "__name__", None)
|
||||
):
|
||||
return f"{self.__class__.__name__}({namespace}.{self.func.__name__})"
|
||||
if isinstance(self.func, type(str.isascii)): # method descriptor
|
||||
return f"{self.__class__.__name__}({self.func.__qualname__})"
|
||||
return f"{self.__class__.__name__}({self.func.__name__})"
|
||||
|
||||
|
||||
@dataclass
|
||||
class Not:
|
||||
func: Callable[[Any], bool]
|
||||
|
||||
def __call__(self, __v: Any) -> bool:
|
||||
return not self.func(__v)
|
||||
|
||||
|
||||
_StrType = TypeVar("_StrType", bound=str)
|
||||
|
||||
LowerCase = Annotated[_StrType, Predicate(str.islower)]
|
||||
"""
|
||||
Return True if the string is a lowercase string, False otherwise.
|
||||
|
||||
A string is lowercase if all cased characters in the string are lowercase and there is at least one cased character in the string.
|
||||
""" # noqa: E501
|
||||
UpperCase = Annotated[_StrType, Predicate(str.isupper)]
|
||||
"""
|
||||
Return True if the string is an uppercase string, False otherwise.
|
||||
|
||||
A string is uppercase if all cased characters in the string are uppercase and there is at least one cased character in the string.
|
||||
""" # noqa: E501
|
||||
IsDigit = Annotated[_StrType, Predicate(str.isdigit)]
|
||||
IsDigits = IsDigit # type: ignore # plural for backwards compatibility, see #63
|
||||
"""
|
||||
Return True if the string is a digit string, False otherwise.
|
||||
|
||||
A string is a digit string if all characters in the string are digits and there is at least one character in the string.
|
||||
""" # noqa: E501
|
||||
IsAscii = Annotated[_StrType, Predicate(str.isascii)]
|
||||
"""
|
||||
Return True if all characters in the string are ASCII, False otherwise.
|
||||
|
||||
ASCII characters have code points in the range U+0000-U+007F. Empty string is ASCII too.
|
||||
"""
|
||||
|
||||
_NumericType = TypeVar('_NumericType', bound=Union[SupportsFloat, SupportsIndex])
|
||||
IsFinite = Annotated[_NumericType, Predicate(math.isfinite)]
|
||||
"""Return True if x is neither an infinity nor a NaN, and False otherwise."""
|
||||
IsNotFinite = Annotated[_NumericType, Predicate(Not(math.isfinite))]
|
||||
"""Return True if x is one of infinity or NaN, and False otherwise"""
|
||||
IsNan = Annotated[_NumericType, Predicate(math.isnan)]
|
||||
"""Return True if x is a NaN (not a number), and False otherwise."""
|
||||
IsNotNan = Annotated[_NumericType, Predicate(Not(math.isnan))]
|
||||
"""Return True if x is anything but NaN (not a number), and False otherwise."""
|
||||
IsInfinite = Annotated[_NumericType, Predicate(math.isinf)]
|
||||
"""Return True if x is a positive or negative infinity, and False otherwise."""
|
||||
IsNotInfinite = Annotated[_NumericType, Predicate(Not(math.isinf))]
|
||||
"""Return True if x is neither a positive or negative infinity, and False otherwise."""
|
||||
|
||||
try:
|
||||
# PEP 727 – Documentation in Annotated Metadata
|
||||
from typing_extensions import Doc # type: ignore[attr-defined]
|
||||
except ImportError:
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Doc: # type: ignore [no-redef]
|
||||
""" "
|
||||
The return value of doc(), mainly to be used by tools that want to extract the
|
||||
Annotated documentation at runtime.
|
||||
"""
|
||||
|
||||
documentation: str
|
||||
"""The documentation string passed to doc()."""
|
||||
|
||||
|
||||
DocInfo = Doc # backwards compatibility
|
||||
doc = Doc
|
||||
Whitespace-only changes.
+146
@@ -0,0 +1,146 @@
|
||||
import math
|
||||
from collections.abc import Iterable, Iterator
|
||||
from datetime import date, datetime, timedelta, timezone
|
||||
from decimal import Decimal
|
||||
from typing import Annotated, Any, NamedTuple
|
||||
|
||||
import annotated_types as at
|
||||
|
||||
|
||||
class Case(NamedTuple):
|
||||
"""
|
||||
A test case for `annotated_types`.
|
||||
"""
|
||||
|
||||
annotation: Any
|
||||
valid_cases: Iterable[Any]
|
||||
invalid_cases: Iterable[Any]
|
||||
|
||||
|
||||
def cases() -> Iterable[Case]:
|
||||
# Gt, Ge, Lt, Le
|
||||
yield Case(Annotated[int, at.Gt(4)], (5, 6, 1000), (4, 0, -1))
|
||||
yield Case(Annotated[float, at.Gt(0.5)], (0.6, 0.7, 0.8, 0.9), (0.5, 0.0, -0.1))
|
||||
yield Case(
|
||||
Annotated[datetime, at.Gt(datetime(2000, 1, 1))],
|
||||
[datetime(2000, 1, 2), datetime(2000, 1, 3)],
|
||||
[datetime(2000, 1, 1), datetime(1999, 12, 31)],
|
||||
)
|
||||
yield Case(
|
||||
Annotated[datetime, at.Gt(date(2000, 1, 1))],
|
||||
[date(2000, 1, 2), date(2000, 1, 3)],
|
||||
[date(2000, 1, 1), date(1999, 12, 31)],
|
||||
)
|
||||
yield Case(
|
||||
Annotated[datetime, at.Gt(Decimal('1.123'))],
|
||||
[Decimal('1.1231'), Decimal('123')],
|
||||
[Decimal('1.123'), Decimal('0')],
|
||||
)
|
||||
|
||||
yield Case(Annotated[int, at.Ge(4)], (4, 5, 6, 1000, 4), (0, -1))
|
||||
yield Case(Annotated[float, at.Ge(0.5)], (0.5, 0.6, 0.7, 0.8, 0.9), (0.4, 0.0, -0.1))
|
||||
yield Case(
|
||||
Annotated[datetime, at.Ge(datetime(2000, 1, 1))],
|
||||
[datetime(2000, 1, 2), datetime(2000, 1, 3)],
|
||||
[datetime(1998, 1, 1), datetime(1999, 12, 31)],
|
||||
)
|
||||
|
||||
yield Case(Annotated[int, at.Lt(4)], (0, -1), (4, 5, 6, 1000, 4))
|
||||
yield Case(Annotated[float, at.Lt(0.5)], (0.4, 0.0, -0.1), (0.5, 0.6, 0.7, 0.8, 0.9))
|
||||
yield Case(
|
||||
Annotated[datetime, at.Lt(datetime(2000, 1, 1))],
|
||||
[datetime(1999, 12, 31), datetime(1999, 12, 31)],
|
||||
[datetime(2000, 1, 2), datetime(2000, 1, 3)],
|
||||
)
|
||||
|
||||
yield Case(Annotated[int, at.Le(4)], (4, 0, -1), (5, 6, 1000))
|
||||
yield Case(Annotated[float, at.Le(0.5)], (0.5, 0.0, -0.1), (0.6, 0.7, 0.8, 0.9))
|
||||
yield Case(
|
||||
Annotated[datetime, at.Le(datetime(2000, 1, 1))],
|
||||
[datetime(2000, 1, 1), datetime(1999, 12, 31)],
|
||||
[datetime(2000, 1, 2), datetime(2000, 1, 3)],
|
||||
)
|
||||
|
||||
# Interval
|
||||
yield Case(Annotated[int, at.Interval(gt=4)], (5, 6, 1000), (4, 0, -1))
|
||||
yield Case(Annotated[int, at.Interval(gt=4, lt=10)], (5, 6), (4, 10, 1000, 0, -1))
|
||||
yield Case(Annotated[float, at.Interval(ge=0.5, le=1)], (0.5, 0.9, 1), (0.49, 1.1))
|
||||
yield Case(
|
||||
Annotated[datetime, at.Interval(gt=datetime(2000, 1, 1), le=datetime(2000, 1, 3))],
|
||||
[datetime(2000, 1, 2), datetime(2000, 1, 3)],
|
||||
[datetime(2000, 1, 1), datetime(2000, 1, 4)],
|
||||
)
|
||||
|
||||
yield Case(Annotated[int, at.MultipleOf(multiple_of=3)], (0, 3, 9), (1, 2, 4))
|
||||
yield Case(Annotated[float, at.MultipleOf(multiple_of=0.5)], (0, 0.5, 1, 1.5), (0.4, 1.1))
|
||||
|
||||
# lengths
|
||||
|
||||
yield Case(Annotated[str, at.MinLen(3)], ('123', '1234', 'x' * 10), ('', '1', '12'))
|
||||
yield Case(Annotated[str, at.Len(3)], ('123', '1234', 'x' * 10), ('', '1', '12'))
|
||||
yield Case(Annotated[list[int], at.MinLen(3)], ([1, 2, 3], [1, 2, 3, 4], [1] * 10), ([], [1], [1, 2]))
|
||||
yield Case(Annotated[list[int], at.Len(3)], ([1, 2, 3], [1, 2, 3, 4], [1] * 10), ([], [1], [1, 2]))
|
||||
|
||||
yield Case(Annotated[str, at.MaxLen(4)], ('', '1234'), ('12345', 'x' * 10))
|
||||
yield Case(Annotated[str, at.Len(0, 4)], ('', '1234'), ('12345', 'x' * 10))
|
||||
yield Case(Annotated[list[str], at.MaxLen(4)], ([], ['a', 'bcdef'], ['a', 'b', 'c']), (['a'] * 5, ['b'] * 10))
|
||||
yield Case(Annotated[list[str], at.Len(0, 4)], ([], ['a', 'bcdef'], ['a', 'b', 'c']), (['a'] * 5, ['b'] * 10))
|
||||
|
||||
yield Case(Annotated[str, at.Len(3, 5)], ('123', '12345'), ('', '1', '12', '123456', 'x' * 10))
|
||||
yield Case(Annotated[str, at.Len(3, 3)], ('123',), ('12', '1234'))
|
||||
|
||||
yield Case(Annotated[dict[int, int], at.Len(2, 3)], [{1: 1, 2: 2}], [{}, {1: 1}, {1: 1, 2: 2, 3: 3, 4: 4}])
|
||||
yield Case(Annotated[set[int], at.Len(2, 3)], ({1, 2}, {1, 2, 3}), (set(), {1}, {1, 2, 3, 4}))
|
||||
yield Case(Annotated[tuple[int, ...], at.Len(2, 3)], ((1, 2), (1, 2, 3)), ((), (1,), (1, 2, 3, 4)))
|
||||
|
||||
# Timezone
|
||||
|
||||
yield Case(
|
||||
Annotated[datetime, at.Timezone(None)], [datetime(2000, 1, 1)], [datetime(2000, 1, 1, tzinfo=timezone.utc)]
|
||||
)
|
||||
yield Case(
|
||||
Annotated[datetime, at.Timezone(...)], [datetime(2000, 1, 1, tzinfo=timezone.utc)], [datetime(2000, 1, 1)]
|
||||
)
|
||||
yield Case(
|
||||
Annotated[datetime, at.Timezone(timezone.utc)],
|
||||
[datetime(2000, 1, 1, tzinfo=timezone.utc)],
|
||||
[datetime(2000, 1, 1), datetime(2000, 1, 1, tzinfo=timezone(timedelta(hours=6)))],
|
||||
)
|
||||
yield Case(
|
||||
Annotated[datetime, at.Timezone('Europe/London')],
|
||||
[datetime(2000, 1, 1, tzinfo=timezone(timedelta(0), name='Europe/London'))],
|
||||
[datetime(2000, 1, 1), datetime(2000, 1, 1, tzinfo=timezone(timedelta(hours=6)))],
|
||||
)
|
||||
|
||||
# Quantity
|
||||
|
||||
yield Case(Annotated[float, at.Unit(unit='m')], (5, 4.2), ('5m', '4.2m'))
|
||||
|
||||
# predicate types
|
||||
|
||||
yield Case(at.LowerCase[str], ['abc', 'foobar'], ['', 'A', 'Boom'])
|
||||
yield Case(at.UpperCase[str], ['ABC', 'DEFO'], ['', 'a', 'abc', 'AbC'])
|
||||
yield Case(at.IsDigit[str], ['123'], ['', 'ab', 'a1b2'])
|
||||
yield Case(at.IsAscii[str], ['123', 'foo bar'], ['£100', '😊', 'whatever 👀'])
|
||||
|
||||
yield Case(Annotated[int, at.Predicate(lambda x: x % 2 == 0)], [0, 2, 4], [1, 3, 5])
|
||||
|
||||
yield Case(at.IsFinite[float], [1.23], [math.nan, math.inf, -math.inf])
|
||||
yield Case(at.IsNotFinite[float], [math.nan, math.inf], [1.23])
|
||||
yield Case(at.IsNan[float], [math.nan], [1.23, math.inf])
|
||||
yield Case(at.IsNotNan[float], [1.23, math.inf], [math.nan])
|
||||
yield Case(at.IsInfinite[float], [math.inf], [math.nan, 1.23])
|
||||
yield Case(at.IsNotInfinite[float], [math.nan, 1.23], [math.inf])
|
||||
|
||||
# check stacked predicates
|
||||
yield Case(at.IsInfinite[Annotated[float, at.Predicate(lambda x: x > 0)]], [math.inf], [-math.inf, 1.23, math.nan])
|
||||
|
||||
# doc
|
||||
yield Case(Annotated[int, at.doc("A number")], [1, 2], [])
|
||||
|
||||
# custom GroupedMetadata
|
||||
class MyCustomGroupedMetadata(at.GroupedMetadata):
|
||||
def __iter__(self) -> Iterator[at.Predicate]:
|
||||
yield at.Predicate(lambda x: float(x).is_integer())
|
||||
|
||||
yield Case(Annotated[float, MyCustomGroupedMetadata()], [0, 2.0], [0.01, 1.5])
|
||||
Vendored
+108
@@ -0,0 +1,108 @@
|
||||
Metadata-Version: 2.4
|
||||
Name: anyio
|
||||
Version: 4.15.1
|
||||
Summary: High-level concurrency and networking framework on top of asyncio or Trio
|
||||
Author-email: Alex Grönholm <alex.gronholm@nextday.fi>
|
||||
License-Expression: MIT
|
||||
Project-URL: Documentation, https://anyio.readthedocs.io/en/latest/
|
||||
Project-URL: Changelog, https://anyio.readthedocs.io/en/stable/versionhistory.html
|
||||
Project-URL: Source code, https://github.com/agronholm/anyio
|
||||
Project-URL: Issue tracker, https://github.com/agronholm/anyio/issues
|
||||
Classifier: Development Status :: 5 - Production/Stable
|
||||
Classifier: Intended Audience :: Developers
|
||||
Classifier: Framework :: AnyIO
|
||||
Classifier: Typing :: Typed
|
||||
Classifier: Programming Language :: Python
|
||||
Classifier: Programming Language :: Python :: 3
|
||||
Classifier: Programming Language :: Python :: 3.10
|
||||
Classifier: Programming Language :: Python :: 3.11
|
||||
Classifier: Programming Language :: Python :: 3.12
|
||||
Classifier: Programming Language :: Python :: 3.13
|
||||
Classifier: Programming Language :: Python :: 3.14
|
||||
Classifier: Programming Language :: Python :: 3.15
|
||||
Classifier: Programming Language :: Python :: Free Threading :: 2 - Beta
|
||||
Requires-Python: >=3.10
|
||||
Description-Content-Type: text/x-rst
|
||||
License-File: LICENSE
|
||||
Requires-Dist: exceptiongroup>=1.0.2; python_version < "3.11"
|
||||
Requires-Dist: idna>=2.8
|
||||
Requires-Dist: typing_extensions>=4.16.0; python_version < "3.15"
|
||||
Provides-Extra: trio
|
||||
Requires-Dist: trio>=0.32.0; extra == "trio"
|
||||
Dynamic: license-file
|
||||
|
||||
.. image:: https://github.com/agronholm/anyio/actions/workflows/test.yml/badge.svg
|
||||
:target: https://github.com/agronholm/anyio/actions/workflows/test.yml
|
||||
:alt: Build Status
|
||||
.. image:: https://coveralls.io/repos/github/agronholm/anyio/badge.svg?branch=master
|
||||
:target: https://coveralls.io/github/agronholm/anyio?branch=master
|
||||
:alt: Code Coverage
|
||||
.. image:: https://readthedocs.org/projects/anyio/badge/?version=latest
|
||||
:target: https://anyio.readthedocs.io/en/latest/?badge=latest
|
||||
:alt: Documentation
|
||||
.. image:: https://badges.gitter.im/gitterHQ/gitter.svg
|
||||
:target: https://gitter.im/python-trio/AnyIO
|
||||
:alt: Gitter chat
|
||||
.. image:: https://tidelift.com/badges/package/pypi/anyio
|
||||
:target: https://tidelift.com/subscription/pkg/pypi-anyio
|
||||
:alt: Tidelift
|
||||
|
||||
AnyIO is an asynchronous networking and concurrency library that works on top of either asyncio_ or
|
||||
Trio_. It implements Trio-like `structured concurrency`_ (SC) on top of asyncio and works in harmony
|
||||
with the native SC of Trio itself.
|
||||
|
||||
Applications and libraries written against AnyIO's API will run unmodified on either asyncio_ or
|
||||
Trio_. AnyIO can also be adopted into a library or application incrementally – bit by bit, no full
|
||||
refactoring necessary. It will blend in with the native libraries of your chosen backend.
|
||||
|
||||
To find out why you might want to use AnyIO's APIs instead of asyncio's, you can read about it
|
||||
`here <https://anyio.readthedocs.io/en/stable/why.html>`_.
|
||||
|
||||
Documentation
|
||||
-------------
|
||||
|
||||
View full documentation at: https://anyio.readthedocs.io/
|
||||
|
||||
Features
|
||||
--------
|
||||
|
||||
AnyIO offers the following functionality:
|
||||
|
||||
* Task groups (nurseries_ in trio terminology)
|
||||
* High-level networking (TCP, UDP and UNIX sockets)
|
||||
|
||||
* `Happy eyeballs`_ algorithm for TCP connections (more robust than that of asyncio on Python
|
||||
3.8)
|
||||
* async/await style UDP sockets (unlike asyncio where you still have to use Transports and
|
||||
Protocols)
|
||||
|
||||
* A versatile API for byte streams and object streams
|
||||
* Inter-task synchronization and communication (locks, conditions, events, semaphores, object
|
||||
streams and futures)
|
||||
* Worker threads
|
||||
* Subprocesses
|
||||
* Subinterpreter support for code parallelization (on Python 3.13 and later)
|
||||
* Asynchronous file I/O (using worker threads)
|
||||
* Signal handling
|
||||
* Asynchronous versions of the functools_ and itertools_ modules
|
||||
|
||||
AnyIO also comes with its own pytest_ plugin which also supports asynchronous fixtures.
|
||||
It even works with the popular Hypothesis_ library.
|
||||
|
||||
.. _asyncio: https://docs.python.org/3/library/asyncio.html
|
||||
.. _Trio: https://github.com/python-trio/trio
|
||||
.. _structured concurrency: https://en.wikipedia.org/wiki/Structured_concurrency
|
||||
.. _nurseries: https://trio.readthedocs.io/en/stable/reference-core.html#nurseries-and-spawning
|
||||
.. _Happy eyeballs: https://en.wikipedia.org/wiki/Happy_Eyeballs
|
||||
.. _pytest: https://docs.pytest.org/en/latest/
|
||||
.. _functools: https://docs.python.org/3/library/functools.html
|
||||
.. _itertools: https://docs.python.org/3/library/itertools.html
|
||||
.. _Hypothesis: https://hypothesis.works/
|
||||
|
||||
Security contact information
|
||||
----------------------------
|
||||
|
||||
To report a security vulnerability, please use the `Tidelift security contact`_.
|
||||
Tidelift will coordinate the fix and disclosure.
|
||||
|
||||
.. _Tidelift security contact: https://tidelift.com/security
|
||||
+53
@@ -0,0 +1,53 @@
|
||||
anyio/__init__.py,sha256=OfroYRjIUw_JFj0Zwh1JAaajMQg4VqtoAt0XoOHDRwE,7602
|
||||
anyio/_lazyimport.py,sha256=LCjnYanIfYE9G-Dv_EsqNLDuVRppNicWZjCTfS9-sJU,6082
|
||||
anyio/from_thread.py,sha256=iZZbfcNySi4wAcx1_cxAQE4J0YZC88qRMIYqQ6SF9sc,19415
|
||||
anyio/functools.py,sha256=T4JS8IXq-x1S0Lbo2owF8l9fza2KypO147QLeyz4cjs,11797
|
||||
anyio/itertools.py,sha256=TS7_UCa2Y3cggtaLlWHbjOyf6QhPfQUdPduNW-ddXTE,16255
|
||||
anyio/lowlevel.py,sha256=QyWkvfetbCH0mHkqv1yxGwOaiZ7fyxhpZkHWOXPV5JU,6280
|
||||
anyio/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
|
||||
anyio/pytest_plugin.py,sha256=oLXyOstvRZW3VVbqIODhZsPt4jSz2UGh44YeG2nJXx0,14516
|
||||
anyio/to_interpreter.py,sha256=ATnoT6kei9wDmQZtXTA3SsY4CK7XKr8Ta4_x5i6pbTw,7100
|
||||
anyio/to_process.py,sha256=p6n8j0nreJUl1DEkR_9PL6KRoKtNqFxORTn8qZBvn94,9850
|
||||
anyio/to_thread.py,sha256=8MPVrrQq8ttGvV21cSd68rMto5bhxJ0-i9CupXag5Gs,2750
|
||||
anyio/_backends/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
|
||||
anyio/_backends/_asyncio.py,sha256=Vqy7O1ymmEpxTiPzehSuBCVTziGALNV-BbsIRaKXsg4,107767
|
||||
anyio/_backends/_trio.py,sha256=Gs8XzX0IeBvje8IEZtJmpNylKCUAChKv7I_YZi3wlmg,45715
|
||||
anyio/_core/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
|
||||
anyio/_core/_asyncio_selector_thread.py,sha256=2PdxFM3cs02Kp6BSppbvmRT7q7asreTW5FgBxEsflBo,5626
|
||||
anyio/_core/_concurrency_utils.py,sha256=kvRXMrnDqgQTBkZc-f279-lBnZM4YbSToXqhnAUCfSc,4522
|
||||
anyio/_core/_contextmanagers.py,sha256=YInBCabiEeS-UaP_Jdxa1CaFC71ETPW8HZTHIM8Rsc8,7215
|
||||
anyio/_core/_eventloop.py,sha256=ByZUeJD9alMfcyTseRo5IzTO0IltEul_Gyq9iqSjqDk,6658
|
||||
anyio/_core/_exceptions.py,sha256=B4ZhLkBQT4tI6vShv4-GJui492IJsoIDTw1FQJEkmAI,5618
|
||||
anyio/_core/_fileio.py,sha256=q--ft9h5Y_FRLzhAnoYrXf3tOuZ8ZyIN2A9mCfMbZRc,34163
|
||||
anyio/_core/_futures.py,sha256=YQq5uhnCnKQMQD62tOLfzEm7dUXY3GQorJdH6ghfgIQ,6177
|
||||
anyio/_core/_resources.py,sha256=Ld5duCe9NwgPfI7EIOh3CMIBYQ7rq6ZQEa7CFS0u9Ec,450
|
||||
anyio/_core/_signals.py,sha256=mjTBB2hTKNPRlU0IhnijeQedpWOGERDiMjSlJQsFrug,1016
|
||||
anyio/_core/_sockets.py,sha256=C8b-3w5H9n0Fo4g7X8WJvk0m6CGh6NVAAFtieBGHK3g,35289
|
||||
anyio/_core/_streams.py,sha256=FczFwIgDpnkK0bODWJXMpsUJYdvAD04kaUaGzJU8DK0,1806
|
||||
anyio/_core/_subprocesses.py,sha256=TMBuGNILz-47w5t9vayJ1GpCpDT6LxsMohJDFCewIRc,7938
|
||||
anyio/_core/_synchronization.py,sha256=ggsd74QPd846faA2HQeT07ONQ5HQ1S78gy8mvB3qCrw,21658
|
||||
anyio/_core/_tasks.py,sha256=J2dhM_Jq3Grsfg_enNWpRz2SxfM15uLt-pEfvoQq3uA,15391
|
||||
anyio/_core/_tempfile.py,sha256=IgoBCFaPhrFvsF9KHcCcrnCRn-UqMfw0WRdSDWgxAiI,19726
|
||||
anyio/_core/_testing.py,sha256=ZRucLh5kqhiFWAJMMbd8yZ6uIvz0rq3WwL72ldosxAU,2340
|
||||
anyio/_core/_typedattr.py,sha256=gt9Dzn-IL0QW14aHPDD8PAtq95FYKNNupHyf37xZGPQ,2607
|
||||
anyio/abc/__init__.py,sha256=XIVvCI2c6oyaOcqXe94k2rzU5_s7-ZNnlxGTk34H4T0,3050
|
||||
anyio/abc/_eventloop.py,sha256=0-UwgB351rkocCt8m3fjTmcNx2Fxn52nUoujG5DzMvU,11127
|
||||
anyio/abc/_resources.py,sha256=hrb0LXUgHSp3i8G6W8JJchdzKXgg1m0T-L7udXmv2y0,900
|
||||
anyio/abc/_sockets.py,sha256=jQWnPX7iDIDyvyvVkbuFhjezw4oI2WggzbW6Vl2-eqs,13162
|
||||
anyio/abc/_streams.py,sha256=IMeOPPutFXRbyeAutxV1HJe7FUQZYp2qKrJJpP1DCv0,7620
|
||||
anyio/abc/_subprocesses.py,sha256=CDqU1p0HuBS0pXUzBW_hbfhihTJRM67n__7qZmoPU84,2123
|
||||
anyio/abc/_tasks.py,sha256=HBPpT72aOy0bwMAxBCKmnXsR328wCYRjYm3iaVw2fQo,7070
|
||||
anyio/abc/_testing.py,sha256=A9Y8RAClIPn1A0aNHBN8YoaRy_MKOpbFGEbTHAXgB6U,2145
|
||||
anyio/streams/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
|
||||
anyio/streams/buffered.py,sha256=u7hCD8SNrYHcutG6K5wEiy1F88-pizgjvEFM22Kq2Cw,6746
|
||||
anyio/streams/file.py,sha256=6jujI2m-QJITqqKFamrupX_DNsU7y2Fz3omLZxOLuY0,4524
|
||||
anyio/streams/memory.py,sha256=rpxNzDMag9PCOynBqN64KyuAfAPB41vFyRP-Npo-lE0,10793
|
||||
anyio/streams/stapled.py,sha256=UJ2Rgt0NUmM1VovRNOoZoic960lcbDjb3vhfZg3b_Z0,4845
|
||||
anyio/streams/text.py,sha256=BcVAGJw1VRvtIqnv-o0Rb0pwH7p8vwlvl21xHq522ag,5765
|
||||
anyio/streams/tls.py,sha256=fqLmxvULFO9rvxGyJ33VPrOIN7E3eiU1q90BkmKigQ8,15818
|
||||
anyio-4.15.1.dist-info/licenses/LICENSE,sha256=U2GsncWPLvX9LpsJxoKXwX8ElQkJu8gCO9uC6s8iwrA,1081
|
||||
anyio-4.15.1.dist-info/METADATA,sha256=BDrstXGvcYnopsUo9djjJUXxss36r3YdUebVtn6jyoE,4733
|
||||
anyio-4.15.1.dist-info/WHEEL,sha256=YVMoNqKzERt-wjUZwJ33xBGAwnFl-4cqbYkTtWa4itE,91
|
||||
anyio-4.15.1.dist-info/entry_points.txt,sha256=_d6Yu6uiaZmNe0CydowirE9Cmg7zUL2g08tQpoS3Qvc,39
|
||||
anyio-4.15.1.dist-info/top_level.txt,sha256=QglSMiWX8_5dpoVAEIHdEYzvqFMdSYWmCj6tYw2ITkQ,6
|
||||
anyio-4.15.1.dist-info/RECORD,,
|
||||
+5
@@ -0,0 +1,5 @@
|
||||
Wheel-Version: 1.0
|
||||
Generator: setuptools (84.0.0)
|
||||
Root-Is-Purelib: true
|
||||
Tag: py3-none-any
|
||||
|
||||
source/sources/original-geranas/FrameBridge-Fleet-Hub/vendor/anyio-4.15.1.dist-info/entry_points.txt
Vendored
+2
@@ -0,0 +1,2 @@
|
||||
[pytest11]
|
||||
anyio = anyio.pytest_plugin
|
||||
source/sources/original-geranas/FrameBridge-Fleet-Hub/vendor/anyio-4.15.1.dist-info/licenses/LICENSE
Vendored
+20
@@ -0,0 +1,20 @@
|
||||
The MIT License (MIT)
|
||||
|
||||
Copyright (c) 2018 Alex Grönholm
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a copy of
|
||||
this software and associated documentation files (the "Software"), to deal in
|
||||
the Software without restriction, including without limitation the rights to
|
||||
use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
|
||||
the Software, and to permit persons to whom the Software is furnished to do so,
|
||||
subject to the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be included in all
|
||||
copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
|
||||
FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
|
||||
COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
|
||||
IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
|
||||
CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
|
||||
Vendored
+1
@@ -0,0 +1 @@
|
||||
anyio
|
||||
+154
@@ -0,0 +1,154 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from ._lazyimport import (
|
||||
fix_package_names,
|
||||
install_lazy_importer,
|
||||
set_deprecated_aliases,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING or not install_lazy_importer():
|
||||
from ._core._concurrency_utils import amap as amap
|
||||
from ._core._concurrency_utils import as_completed as as_completed
|
||||
from ._core._concurrency_utils import gather as gather
|
||||
from ._core._contextmanagers import (
|
||||
AsyncContextManagerMixin as AsyncContextManagerMixin,
|
||||
)
|
||||
from ._core._contextmanagers import ContextManagerMixin as ContextManagerMixin
|
||||
from ._core._eventloop import current_time as current_time
|
||||
from ._core._eventloop import get_all_backends as get_all_backends
|
||||
from ._core._eventloop import get_available_backends as get_available_backends
|
||||
from ._core._eventloop import get_cancelled_exc_class as get_cancelled_exc_class
|
||||
from ._core._eventloop import run as run
|
||||
from ._core._eventloop import sleep as sleep
|
||||
from ._core._eventloop import sleep_forever as sleep_forever
|
||||
from ._core._eventloop import sleep_until as sleep_until
|
||||
from ._core._exceptions import BrokenResourceError as BrokenResourceError
|
||||
from ._core._exceptions import BrokenWorkerInterpreter as BrokenWorkerInterpreter
|
||||
from ._core._exceptions import BrokenWorkerProcess as BrokenWorkerProcess
|
||||
from ._core._exceptions import BusyResourceError as BusyResourceError
|
||||
from ._core._exceptions import ClosedResourceError as ClosedResourceError
|
||||
from ._core._exceptions import ConnectionFailed as ConnectionFailed
|
||||
from ._core._exceptions import DelimiterNotFound as DelimiterNotFound
|
||||
from ._core._exceptions import EndOfStream as EndOfStream
|
||||
from ._core._exceptions import FutureAlreadyFinished as FutureAlreadyFinished
|
||||
from ._core._exceptions import FutureCancelled as FutureCancelled
|
||||
from ._core._exceptions import FutureFailed as FutureFailed
|
||||
from ._core._exceptions import FutureNotFinished as FutureNotFinished
|
||||
from ._core._exceptions import IncompleteRead as IncompleteRead
|
||||
from ._core._exceptions import NoEventLoopError as NoEventLoopError
|
||||
from ._core._exceptions import RunFinishedError as RunFinishedError
|
||||
from ._core._exceptions import TaskCancelled as TaskCancelled
|
||||
from ._core._exceptions import TaskFailed as TaskFailed
|
||||
from ._core._exceptions import TaskNotFinished as TaskNotFinished
|
||||
from ._core._exceptions import (
|
||||
TypedAttributeLookupError as TypedAttributeLookupError,
|
||||
)
|
||||
from ._core._exceptions import WouldBlock as WouldBlock
|
||||
from ._core._fileio import AsyncFile as AsyncFile
|
||||
from ._core._fileio import Path as Path
|
||||
from ._core._fileio import open_file as open_file
|
||||
from ._core._fileio import wrap_file as wrap_file
|
||||
from ._core._futures import Future as Future
|
||||
from ._core._resources import aclose_forcefully as aclose_forcefully
|
||||
from ._core._signals import open_signal_receiver as open_signal_receiver
|
||||
from ._core._sockets import TCPConnectable as TCPConnectable
|
||||
from ._core._sockets import UNIXConnectable as UNIXConnectable
|
||||
from ._core._sockets import as_connectable as as_connectable
|
||||
from ._core._sockets import connect_tcp as connect_tcp
|
||||
from ._core._sockets import connect_unix as connect_unix
|
||||
from ._core._sockets import (
|
||||
create_connected_udp_socket as create_connected_udp_socket,
|
||||
)
|
||||
from ._core._sockets import (
|
||||
create_connected_unix_datagram_socket as create_connected_unix_datagram_socket,
|
||||
)
|
||||
from ._core._sockets import create_tcp_listener as create_tcp_listener
|
||||
from ._core._sockets import create_udp_socket as create_udp_socket
|
||||
from ._core._sockets import (
|
||||
create_unix_datagram_socket as create_unix_datagram_socket,
|
||||
)
|
||||
from ._core._sockets import create_unix_listener as create_unix_listener
|
||||
from ._core._sockets import getaddrinfo as getaddrinfo
|
||||
from ._core._sockets import getnameinfo as getnameinfo
|
||||
from ._core._sockets import notify_closing as notify_closing
|
||||
from ._core._sockets import wait_readable as wait_readable
|
||||
from ._core._sockets import wait_socket_readable as wait_socket_readable
|
||||
from ._core._sockets import wait_socket_writable as wait_socket_writable
|
||||
from ._core._sockets import wait_writable as wait_writable
|
||||
from ._core._streams import (
|
||||
create_memory_object_stream as create_memory_object_stream,
|
||||
)
|
||||
from ._core._subprocesses import open_process as open_process
|
||||
from ._core._subprocesses import run_process as run_process
|
||||
from ._core._synchronization import CapacityLimiter as CapacityLimiter
|
||||
from ._core._synchronization import (
|
||||
CapacityLimiterStatistics as CapacityLimiterStatistics,
|
||||
)
|
||||
from ._core._synchronization import Condition as Condition
|
||||
from ._core._synchronization import ConditionStatistics as ConditionStatistics
|
||||
from ._core._synchronization import Event as Event
|
||||
from ._core._synchronization import EventStatistics as EventStatistics
|
||||
from ._core._synchronization import Lock as Lock
|
||||
from ._core._synchronization import LockStatistics as LockStatistics
|
||||
from ._core._synchronization import ResourceGuard as ResourceGuard
|
||||
from ._core._synchronization import Semaphore as Semaphore
|
||||
from ._core._synchronization import SemaphoreStatistics as SemaphoreStatistics
|
||||
from ._core._tasks import TASK_STATUS_IGNORED as TASK_STATUS_IGNORED
|
||||
from ._core._tasks import CancelScope as CancelScope
|
||||
from ._core._tasks import TaskHandle as TaskHandle
|
||||
from ._core._tasks import create_task_group as create_task_group
|
||||
from ._core._tasks import current_effective_deadline as current_effective_deadline
|
||||
from ._core._tasks import fail_after as fail_after
|
||||
from ._core._tasks import fail_at as fail_at
|
||||
from ._core._tasks import move_on_after as move_on_after
|
||||
from ._core._tasks import move_on_at as move_on_at
|
||||
from ._core._tempfile import NamedTemporaryFile as NamedTemporaryFile
|
||||
from ._core._tempfile import SpooledTemporaryFile as SpooledTemporaryFile
|
||||
from ._core._tempfile import TemporaryDirectory as TemporaryDirectory
|
||||
from ._core._tempfile import TemporaryFile as TemporaryFile
|
||||
from ._core._tempfile import gettempdir as gettempdir
|
||||
from ._core._tempfile import gettempdirb as gettempdirb
|
||||
from ._core._tempfile import mkdtemp as mkdtemp
|
||||
from ._core._tempfile import mkstemp as mkstemp
|
||||
from ._core._testing import TaskInfo as TaskInfo
|
||||
from ._core._testing import get_current_task as get_current_task
|
||||
from ._core._testing import get_running_tasks as get_running_tasks
|
||||
from ._core._testing import wait_all_tasks_blocked as wait_all_tasks_blocked
|
||||
from ._core._typedattr import TypedAttributeProvider as TypedAttributeProvider
|
||||
from ._core._typedattr import TypedAttributeSet as TypedAttributeSet
|
||||
from ._core._typedattr import typed_attribute as typed_attribute
|
||||
|
||||
# ruff: isort: off
|
||||
from . import (
|
||||
abc as abc,
|
||||
)
|
||||
from . import (
|
||||
from_thread as from_thread,
|
||||
)
|
||||
from . import (
|
||||
functools as functools,
|
||||
)
|
||||
from . import (
|
||||
itertools as itertools,
|
||||
)
|
||||
from . import (
|
||||
lowlevel as lowlevel,
|
||||
)
|
||||
from . import (
|
||||
to_interpreter as to_interpreter,
|
||||
)
|
||||
from . import (
|
||||
to_process as to_process,
|
||||
)
|
||||
from . import (
|
||||
to_thread as to_thread,
|
||||
)
|
||||
|
||||
fix_package_names()
|
||||
set_deprecated_aliases(
|
||||
{
|
||||
"BrokenWorkerIntepreter": "anyio.BrokenWorkerInterpreter",
|
||||
}
|
||||
)
|
||||
Whitespace-only changes.
+3201
File diff suppressed because it is too large.
Load diff
+1473
File diff suppressed because it is too large.
Load diff
Whitespace-only changes.
source/sources/original-geranas/FrameBridge-Fleet-Hub/vendor/anyio/_core/_asyncio_selector_thread.py
Vendored
+167
@@ -0,0 +1,167 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import socket
|
||||
import threading
|
||||
from collections.abc import Callable
|
||||
from selectors import EVENT_READ, EVENT_WRITE, DefaultSelector
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from _typeshed import FileDescriptorLike
|
||||
|
||||
_selector_lock = threading.Lock()
|
||||
_selector: Selector | None = None
|
||||
|
||||
|
||||
class Selector:
|
||||
def __init__(self) -> None:
|
||||
self._thread = threading.Thread(target=self.run, name="AnyIO socket selector")
|
||||
self._selector = DefaultSelector()
|
||||
self._send, self._receive = socket.socketpair()
|
||||
self._send.setblocking(False)
|
||||
self._receive.setblocking(False)
|
||||
# This somewhat reduces the amount of memory wasted queueing up data
|
||||
# for wakeups. With these settings, maximum number of 1-byte sends
|
||||
# before getting BlockingIOError:
|
||||
# Linux 4.8: 6
|
||||
# macOS (darwin 15.5): 1
|
||||
# Windows 10: 525347
|
||||
# Windows you're weird. (And on Windows setting SNDBUF to 0 makes send
|
||||
# blocking, even on non-blocking sockets, so don't do that.)
|
||||
self._receive.setsockopt(socket.SOL_SOCKET, socket.SO_RCVBUF, 1)
|
||||
self._send.setsockopt(socket.SOL_SOCKET, socket.SO_SNDBUF, 1)
|
||||
# On Windows this is a TCP socket so this might matter. On other
|
||||
# platforms this fails b/c AF_UNIX sockets aren't actually TCP.
|
||||
try:
|
||||
self._send.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
self._selector.register(self._receive, EVENT_READ)
|
||||
self._closed = False
|
||||
|
||||
def start(self) -> None:
|
||||
self._thread.start()
|
||||
threading._register_atexit(self._stop) # type: ignore[attr-defined]
|
||||
|
||||
def _stop(self) -> None:
|
||||
global _selector
|
||||
self._closed = True
|
||||
self._notify_self()
|
||||
self._send.close()
|
||||
self._thread.join()
|
||||
self._selector.unregister(self._receive)
|
||||
self._receive.close()
|
||||
self._selector.close()
|
||||
_selector = None
|
||||
assert not self._selector.get_map(), (
|
||||
"selector still has registered file descriptors after shutdown"
|
||||
)
|
||||
|
||||
def _notify_self(self) -> None:
|
||||
try:
|
||||
self._send.send(b"\x00")
|
||||
except BlockingIOError:
|
||||
pass
|
||||
|
||||
def add_reader(self, fd: FileDescriptorLike, callback: Callable[[], Any]) -> None:
|
||||
loop = asyncio.get_running_loop()
|
||||
try:
|
||||
key = self._selector.get_key(fd)
|
||||
except KeyError:
|
||||
self._selector.register(fd, EVENT_READ, {EVENT_READ: (loop, callback)})
|
||||
else:
|
||||
if EVENT_READ in key.data:
|
||||
raise ValueError(
|
||||
"this file descriptor is already registered for reading"
|
||||
)
|
||||
|
||||
key.data[EVENT_READ] = loop, callback
|
||||
self._selector.modify(fd, key.events | EVENT_READ, key.data)
|
||||
|
||||
self._notify_self()
|
||||
|
||||
def add_writer(self, fd: FileDescriptorLike, callback: Callable[[], Any]) -> None:
|
||||
loop = asyncio.get_running_loop()
|
||||
try:
|
||||
key = self._selector.get_key(fd)
|
||||
except KeyError:
|
||||
self._selector.register(fd, EVENT_WRITE, {EVENT_WRITE: (loop, callback)})
|
||||
else:
|
||||
if EVENT_WRITE in key.data:
|
||||
raise ValueError(
|
||||
"this file descriptor is already registered for writing"
|
||||
)
|
||||
|
||||
key.data[EVENT_WRITE] = loop, callback
|
||||
self._selector.modify(fd, key.events | EVENT_WRITE, key.data)
|
||||
|
||||
self._notify_self()
|
||||
|
||||
def remove_reader(self, fd: FileDescriptorLike) -> bool:
|
||||
try:
|
||||
key = self._selector.get_key(fd)
|
||||
except KeyError:
|
||||
return False
|
||||
|
||||
if new_events := key.events ^ EVENT_READ:
|
||||
del key.data[EVENT_READ]
|
||||
self._selector.modify(fd, new_events, key.data)
|
||||
else:
|
||||
self._selector.unregister(fd)
|
||||
|
||||
return True
|
||||
|
||||
def remove_writer(self, fd: FileDescriptorLike) -> bool:
|
||||
try:
|
||||
key = self._selector.get_key(fd)
|
||||
except KeyError:
|
||||
return False
|
||||
|
||||
if new_events := key.events ^ EVENT_WRITE:
|
||||
del key.data[EVENT_WRITE]
|
||||
self._selector.modify(fd, new_events, key.data)
|
||||
else:
|
||||
self._selector.unregister(fd)
|
||||
|
||||
return True
|
||||
|
||||
def run(self) -> None:
|
||||
while not self._closed:
|
||||
for key, events in self._selector.select():
|
||||
if key.fileobj is self._receive:
|
||||
try:
|
||||
while self._receive.recv(4096):
|
||||
pass
|
||||
except BlockingIOError:
|
||||
pass
|
||||
|
||||
continue
|
||||
|
||||
if events & EVENT_READ:
|
||||
loop, callback = key.data[EVENT_READ]
|
||||
self.remove_reader(key.fd)
|
||||
try:
|
||||
loop.call_soon_threadsafe(callback)
|
||||
except RuntimeError:
|
||||
pass # the loop was already closed
|
||||
|
||||
if events & EVENT_WRITE:
|
||||
loop, callback = key.data[EVENT_WRITE]
|
||||
self.remove_writer(key.fd)
|
||||
try:
|
||||
loop.call_soon_threadsafe(callback)
|
||||
except RuntimeError:
|
||||
pass # the loop was already closed
|
||||
|
||||
|
||||
def get_selector() -> Selector:
|
||||
global _selector
|
||||
|
||||
with _selector_lock:
|
||||
if _selector is None:
|
||||
_selector = Selector()
|
||||
_selector.start()
|
||||
|
||||
return _selector
|
||||
Vendored
+156
@@ -0,0 +1,156 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from inspect import iscoroutine
|
||||
|
||||
__all__ = ("amap", "as_completed", "gather")
|
||||
|
||||
from collections.abc import AsyncGenerator, Awaitable, Callable, Coroutine, Iterable
|
||||
from contextlib import asynccontextmanager
|
||||
from typing import TYPE_CHECKING, Any, TypeVar, overload
|
||||
|
||||
from ..abc._tasks import get_coro_name
|
||||
from ._exceptions import BrokenResourceError
|
||||
from ._streams import create_memory_object_stream
|
||||
from ._tasks import TaskHandle, create_task_group
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
|
||||
|
||||
R = TypeVar("R")
|
||||
S = TypeVar("S")
|
||||
T = TypeVar("T")
|
||||
U = TypeVar("U")
|
||||
V = TypeVar("V")
|
||||
W = TypeVar("W")
|
||||
|
||||
|
||||
@overload
|
||||
async def gather(
|
||||
coro1: Coroutine[Any, Any, R], coro2: Coroutine[Any, Any, S], /
|
||||
) -> tuple[R, S]: ...
|
||||
|
||||
|
||||
@overload
|
||||
async def gather(
|
||||
coro1: Coroutine[Any, Any, R],
|
||||
coro2: Coroutine[Any, Any, S],
|
||||
coro3: Coroutine[Any, Any, T],
|
||||
/,
|
||||
) -> tuple[R, S, T]: ...
|
||||
|
||||
|
||||
@overload
|
||||
async def gather(
|
||||
coro1: Coroutine[Any, Any, R],
|
||||
coro2: Coroutine[Any, Any, S],
|
||||
coro3: Coroutine[Any, Any, T],
|
||||
coro4: Coroutine[Any, Any, U],
|
||||
/,
|
||||
) -> tuple[R, S, T, U]: ...
|
||||
|
||||
|
||||
@overload
|
||||
async def gather(
|
||||
coro1: Coroutine[Any, Any, R],
|
||||
coro2: Coroutine[Any, Any, S],
|
||||
coro3: Coroutine[Any, Any, T],
|
||||
coro4: Coroutine[Any, Any, U],
|
||||
coro5: Coroutine[Any, Any, V],
|
||||
/,
|
||||
) -> tuple[R, S, T, U, V]: ...
|
||||
|
||||
|
||||
@overload
|
||||
async def gather(
|
||||
coro1: Coroutine[Any, Any, R],
|
||||
coro2: Coroutine[Any, Any, S],
|
||||
coro3: Coroutine[Any, Any, T],
|
||||
coro4: Coroutine[Any, Any, U],
|
||||
coro5: Coroutine[Any, Any, V],
|
||||
coro6: Coroutine[Any, Any, W],
|
||||
/,
|
||||
) -> tuple[R, S, T, U, V, W]: ...
|
||||
|
||||
|
||||
# handle arbitrary length if awaitables are all of the same type
|
||||
@overload
|
||||
async def gather(*coros: Coroutine[Any, Any, R]) -> tuple[R, ...]: ...
|
||||
|
||||
|
||||
async def gather(*coros: Coroutine[Any, Any, Any]) -> tuple[Any, ...]:
|
||||
"""
|
||||
Run coroutines concurrently in a task group. The order of result values corresponds
|
||||
to the order of coroutines passed.
|
||||
|
||||
:param coros: coroutine objects to run as tasks
|
||||
:return: task results for each argument in the same order as they were passed
|
||||
|
||||
"""
|
||||
handles: list[TaskHandle[Any]] = []
|
||||
async with create_task_group() as tg:
|
||||
handles.extend(tg.create_task(coro) for coro in coros)
|
||||
|
||||
return tuple(h.return_value for h in handles)
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def as_completed(
|
||||
*awaitables: Awaitable[R],
|
||||
) -> AsyncGenerator[MemoryObjectReceiveStream[TaskHandle[R]]]:
|
||||
"""
|
||||
Run awaitable objects concurrently in a task group, returning an iterator which can
|
||||
be used to get finished task handles in the order they complete.
|
||||
|
||||
:param awaitables: awaitable objects to run as tasks
|
||||
:return: MemoryObjectReceiveStream for iterating over task handles as tasks complete.
|
||||
|
||||
"""
|
||||
if not awaitables:
|
||||
raise ValueError("as_completed() takes at least one awaitable")
|
||||
|
||||
send, recv = create_memory_object_stream[TaskHandle[R]](len(awaitables))
|
||||
task_handles: list[TaskHandle[R]] = []
|
||||
|
||||
async def runner(
|
||||
awaitable: Awaitable[R],
|
||||
index: int,
|
||||
_send: MemoryObjectSendStream[TaskHandle[R]],
|
||||
) -> R:
|
||||
async with _send:
|
||||
try:
|
||||
return await awaitable
|
||||
finally:
|
||||
try:
|
||||
_send.send_nowait(task_handles[index])
|
||||
except BrokenResourceError:
|
||||
pass
|
||||
|
||||
async with recv, create_task_group() as tg:
|
||||
async with send:
|
||||
for i, awaitable in enumerate(awaitables):
|
||||
name = get_coro_name(awaitable) if iscoroutine(awaitable) else None
|
||||
task_handles.append(
|
||||
tg.start_soon(runner, awaitable, i, send.clone(), name=name)
|
||||
)
|
||||
try:
|
||||
yield recv
|
||||
finally:
|
||||
tg.cancel_scope.cancel()
|
||||
|
||||
|
||||
async def amap(
|
||||
func: Callable[[T], Coroutine[Any, Any, R]], args: Iterable[T]
|
||||
) -> list[R]:
|
||||
"""
|
||||
Run the given coroutine function concurrently for multiple argument values.
|
||||
|
||||
:param func: a coroutine function that takes a single argument
|
||||
:param args: a sequence of argument values to pass to ``func``
|
||||
:return: task results for each argument in the same order as they were passed
|
||||
|
||||
"""
|
||||
handles: list[TaskHandle[R]] = []
|
||||
async with create_task_group() as tg:
|
||||
handles.extend(tg.start_soon(func, arg) for arg in args)
|
||||
|
||||
return [h.return_value for h in handles]
|
||||
Vendored
+200
@@ -0,0 +1,200 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from abc import abstractmethod
|
||||
from contextlib import AbstractAsyncContextManager, AbstractContextManager
|
||||
from inspect import isasyncgen, iscoroutine, isgenerator
|
||||
from types import TracebackType
|
||||
from typing import Protocol, TypeVar, cast, final
|
||||
|
||||
_T_co = TypeVar("_T_co", covariant=True)
|
||||
_ExitT_co = TypeVar("_ExitT_co", covariant=True, bound="bool | None")
|
||||
|
||||
|
||||
class _SupportsCtxMgr(Protocol[_T_co, _ExitT_co]):
|
||||
def __contextmanager__(self) -> AbstractContextManager[_T_co, _ExitT_co]: ...
|
||||
|
||||
|
||||
class _SupportsAsyncCtxMgr(Protocol[_T_co, _ExitT_co]):
|
||||
def __asynccontextmanager__(
|
||||
self,
|
||||
) -> AbstractAsyncContextManager[_T_co, _ExitT_co]: ...
|
||||
|
||||
|
||||
class ContextManagerMixin:
|
||||
"""
|
||||
Mixin class providing context manager functionality via a generator-based
|
||||
implementation.
|
||||
|
||||
This class allows you to implement a context manager via :meth:`__contextmanager__`
|
||||
which should return a generator. The mechanics are meant to mirror those of
|
||||
:func:`@contextmanager <contextlib.contextmanager>`.
|
||||
|
||||
.. note:: Classes using this mix-in are not reentrant as context managers, meaning
|
||||
that once you enter it, you can't re-enter before first exiting it.
|
||||
|
||||
.. seealso:: :doc:`contextmanagers`
|
||||
"""
|
||||
|
||||
__cm: AbstractContextManager[object, bool | None] | None = None
|
||||
|
||||
@final
|
||||
def __enter__(self: _SupportsCtxMgr[_T_co, bool | None]) -> _T_co:
|
||||
# Needed for mypy to assume self still has the __cm member
|
||||
assert isinstance(self, ContextManagerMixin)
|
||||
if self.__cm is not None:
|
||||
raise RuntimeError(
|
||||
f"this {self.__class__.__qualname__} has already been entered"
|
||||
)
|
||||
|
||||
cm = self.__contextmanager__()
|
||||
if not isinstance(cm, AbstractContextManager):
|
||||
if isgenerator(cm):
|
||||
raise TypeError(
|
||||
"__contextmanager__() returned a generator object instead of "
|
||||
"a context manager. Did you forget to add the @contextmanager "
|
||||
"decorator?"
|
||||
)
|
||||
|
||||
raise TypeError(
|
||||
f"__contextmanager__() did not return a context manager object, "
|
||||
f"but {cm.__class__!r}"
|
||||
)
|
||||
|
||||
if cm is self:
|
||||
raise TypeError(
|
||||
f"{self.__class__.__qualname__}.__contextmanager__() returned "
|
||||
f"self. Did you forget to add the @contextmanager decorator and a "
|
||||
f"'yield' statement?"
|
||||
)
|
||||
|
||||
value = cm.__enter__()
|
||||
self.__cm = cm
|
||||
return value
|
||||
|
||||
@final
|
||||
def __exit__(
|
||||
self: _SupportsCtxMgr[object, _ExitT_co],
|
||||
exc_type: type[BaseException] | None,
|
||||
exc_val: BaseException | None,
|
||||
exc_tb: TracebackType | None,
|
||||
) -> _ExitT_co:
|
||||
# Needed for mypy to assume self still has the __cm member
|
||||
assert isinstance(self, ContextManagerMixin)
|
||||
if self.__cm is None:
|
||||
raise RuntimeError(
|
||||
f"this {self.__class__.__qualname__} has not been entered yet"
|
||||
)
|
||||
|
||||
# Prevent circular references
|
||||
cm = self.__cm
|
||||
del self.__cm
|
||||
|
||||
return cast(_ExitT_co, cm.__exit__(exc_type, exc_val, exc_tb))
|
||||
|
||||
@abstractmethod
|
||||
def __contextmanager__(self) -> AbstractContextManager[object, bool | None]:
|
||||
"""
|
||||
Implement your context manager logic here.
|
||||
|
||||
This method **must** be decorated with
|
||||
:func:`@contextmanager <contextlib.contextmanager>`.
|
||||
|
||||
.. note:: Remember that the ``yield`` will raise any exception raised in the
|
||||
enclosed context block, so use a ``finally:`` block to clean up resources!
|
||||
|
||||
:return: a context manager object
|
||||
"""
|
||||
|
||||
|
||||
class AsyncContextManagerMixin:
|
||||
"""
|
||||
Mixin class providing async context manager functionality via a generator-based
|
||||
implementation.
|
||||
|
||||
This class allows you to implement a context manager via
|
||||
:meth:`__asynccontextmanager__`. The mechanics are meant to mirror those of
|
||||
:func:`@asynccontextmanager <contextlib.asynccontextmanager>`.
|
||||
|
||||
.. note:: Classes using this mix-in are not reentrant as context managers, meaning
|
||||
that once you enter it, you can't re-enter before first exiting it.
|
||||
|
||||
.. seealso:: :doc:`contextmanagers`
|
||||
"""
|
||||
|
||||
__cm: AbstractAsyncContextManager[object, bool | None] | None = None
|
||||
|
||||
@final
|
||||
async def __aenter__(self: _SupportsAsyncCtxMgr[_T_co, bool | None]) -> _T_co:
|
||||
# Needed for mypy to assume self still has the __cm member
|
||||
assert isinstance(self, AsyncContextManagerMixin)
|
||||
if self.__cm is not None:
|
||||
raise RuntimeError(
|
||||
f"this {self.__class__.__qualname__} has already been entered"
|
||||
)
|
||||
|
||||
cm = self.__asynccontextmanager__()
|
||||
if not isinstance(cm, AbstractAsyncContextManager):
|
||||
if isasyncgen(cm):
|
||||
raise TypeError(
|
||||
"__asynccontextmanager__() returned an async generator instead of "
|
||||
"an async context manager. Did you forget to add the "
|
||||
"@asynccontextmanager decorator?"
|
||||
)
|
||||
elif iscoroutine(cm):
|
||||
cm.close()
|
||||
raise TypeError(
|
||||
"__asynccontextmanager__() returned a coroutine object instead of "
|
||||
"an async context manager. Did you forget to add the "
|
||||
"@asynccontextmanager decorator and a 'yield' statement?"
|
||||
)
|
||||
|
||||
raise TypeError(
|
||||
f"__asynccontextmanager__() did not return an async context manager, "
|
||||
f"but {cm.__class__!r}"
|
||||
)
|
||||
|
||||
if cm is self:
|
||||
raise TypeError(
|
||||
f"{self.__class__.__qualname__}.__asynccontextmanager__() returned "
|
||||
f"self. Did you forget to add the @asynccontextmanager decorator and a "
|
||||
f"'yield' statement?"
|
||||
)
|
||||
|
||||
value = await cm.__aenter__()
|
||||
self.__cm = cm
|
||||
return value
|
||||
|
||||
@final
|
||||
async def __aexit__(
|
||||
self: _SupportsAsyncCtxMgr[object, _ExitT_co],
|
||||
exc_type: type[BaseException] | None,
|
||||
exc_val: BaseException | None,
|
||||
exc_tb: TracebackType | None,
|
||||
) -> _ExitT_co:
|
||||
assert isinstance(self, AsyncContextManagerMixin)
|
||||
if self.__cm is None:
|
||||
raise RuntimeError(
|
||||
f"this {self.__class__.__qualname__} has not been entered yet"
|
||||
)
|
||||
|
||||
# Prevent circular references
|
||||
cm = self.__cm
|
||||
del self.__cm
|
||||
|
||||
return cast(_ExitT_co, await cm.__aexit__(exc_type, exc_val, exc_tb))
|
||||
|
||||
@abstractmethod
|
||||
def __asynccontextmanager__(
|
||||
self,
|
||||
) -> AbstractAsyncContextManager[object, bool | None]:
|
||||
"""
|
||||
Implement your async context manager logic here.
|
||||
|
||||
This method **must** be decorated with
|
||||
:func:`@asynccontextmanager <contextlib.asynccontextmanager>`.
|
||||
|
||||
.. note:: Remember that the ``yield`` will raise any exception raised in the
|
||||
enclosed context block, so use a ``finally:`` block to clean up resources!
|
||||
|
||||
:return: an async context manager object
|
||||
"""
|
||||
+240
@@ -0,0 +1,240 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
import sys
|
||||
import threading
|
||||
from collections.abc import Awaitable, Callable, Generator
|
||||
from contextlib import contextmanager
|
||||
from contextvars import Token
|
||||
from importlib import import_module
|
||||
from typing import TYPE_CHECKING, Any, TypeVar
|
||||
|
||||
from ._exceptions import NoEventLoopError
|
||||
|
||||
if sys.version_info >= (3, 11):
|
||||
from typing import TypeVarTuple, Unpack
|
||||
else:
|
||||
from typing_extensions import TypeVarTuple, Unpack
|
||||
|
||||
sniffio: Any
|
||||
try:
|
||||
import sniffio
|
||||
except ModuleNotFoundError:
|
||||
sniffio = None
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ..abc import AsyncBackend
|
||||
|
||||
# This must be updated when new backends are introduced
|
||||
BACKENDS = "asyncio", "trio"
|
||||
|
||||
T_Retval = TypeVar("T_Retval")
|
||||
PosArgsT = TypeVarTuple("PosArgsT")
|
||||
|
||||
threadlocals = threading.local()
|
||||
loaded_backends: dict[str, type[AsyncBackend]] = {}
|
||||
|
||||
|
||||
def run(
|
||||
func: Callable[[Unpack[PosArgsT]], Awaitable[T_Retval]],
|
||||
*args: Unpack[PosArgsT],
|
||||
backend: str = "asyncio",
|
||||
backend_options: dict[str, Any] | None = None,
|
||||
) -> T_Retval:
|
||||
"""
|
||||
Run the given coroutine function in an asynchronous event loop.
|
||||
|
||||
The current thread must not be already running an event loop.
|
||||
|
||||
:param func: a coroutine function
|
||||
:param args: positional arguments to ``func``
|
||||
:param backend: name of the asynchronous event loop implementation – currently
|
||||
either ``asyncio`` or ``trio``
|
||||
:param backend_options: keyword arguments to call the backend ``run()``
|
||||
implementation with (documented :ref:`here <backend options>`)
|
||||
:return: the return value of the coroutine function
|
||||
:raises RuntimeError: if an asynchronous event loop is already running in this
|
||||
thread
|
||||
:raises LookupError: if the named backend is not found
|
||||
|
||||
"""
|
||||
if asynclib_name := current_async_library():
|
||||
raise RuntimeError(f"Already running {asynclib_name} in this thread")
|
||||
|
||||
try:
|
||||
async_backend = get_async_backend(backend)
|
||||
except ImportError as exc:
|
||||
if backend in BACKENDS:
|
||||
raise LookupError(
|
||||
f"Backend {backend!r} is not available. "
|
||||
f"Install it with: pip install anyio[{backend}]"
|
||||
) from exc
|
||||
|
||||
raise LookupError(f"No such backend: {backend}") from exc
|
||||
|
||||
token = None
|
||||
if asynclib_name is None:
|
||||
# Since we're in control of the event loop, we can cache the name of the async
|
||||
# library
|
||||
token = set_current_async_library(backend)
|
||||
|
||||
try:
|
||||
backend_options = backend_options or {}
|
||||
return async_backend.run(func, args, {}, backend_options)
|
||||
finally:
|
||||
reset_current_async_library(token)
|
||||
|
||||
|
||||
async def sleep(delay: float) -> None:
|
||||
"""
|
||||
Pause the current task for the specified duration.
|
||||
|
||||
:param delay: the duration, in seconds
|
||||
|
||||
"""
|
||||
return await get_async_backend().sleep(delay)
|
||||
|
||||
|
||||
async def sleep_forever() -> None:
|
||||
"""
|
||||
Pause the current task until it's cancelled.
|
||||
|
||||
This is a shortcut for ``sleep(math.inf)``.
|
||||
|
||||
.. versionadded:: 3.1
|
||||
|
||||
"""
|
||||
await sleep(math.inf)
|
||||
|
||||
|
||||
async def sleep_until(deadline: float) -> None:
|
||||
"""
|
||||
Pause the current task until the given time.
|
||||
|
||||
:param deadline: the absolute time to wake up at (according to the internal
|
||||
monotonic clock of the event loop)
|
||||
|
||||
.. versionadded:: 3.1
|
||||
|
||||
"""
|
||||
now = current_time()
|
||||
await sleep(max(deadline - now, 0))
|
||||
|
||||
|
||||
def current_time() -> float:
|
||||
"""
|
||||
Return the current value of the event loop's internal clock.
|
||||
|
||||
:return: the clock value (seconds)
|
||||
:raises NoEventLoopError: if no supported asynchronous event loop is running in the
|
||||
current thread
|
||||
|
||||
"""
|
||||
return get_async_backend().current_time()
|
||||
|
||||
|
||||
def get_all_backends() -> tuple[str, ...]:
|
||||
"""Return a tuple of the names of all built-in backends."""
|
||||
return BACKENDS
|
||||
|
||||
|
||||
def get_available_backends() -> tuple[str, ...]:
|
||||
"""
|
||||
Test for the availability of built-in backends.
|
||||
|
||||
:return a tuple of the built-in backend names that were successfully imported
|
||||
|
||||
.. versionadded:: 4.12
|
||||
|
||||
"""
|
||||
available_backends: list[str] = []
|
||||
for backend_name in get_all_backends():
|
||||
try:
|
||||
get_async_backend(backend_name)
|
||||
except ImportError:
|
||||
continue
|
||||
|
||||
available_backends.append(backend_name)
|
||||
|
||||
return tuple(available_backends)
|
||||
|
||||
|
||||
def get_cancelled_exc_class() -> type[BaseException]:
|
||||
"""
|
||||
Return the current async library's cancellation exception class.
|
||||
|
||||
:raises NoEventLoopError: if no supported asynchronous event loop is running in the
|
||||
current thread
|
||||
|
||||
"""
|
||||
return get_async_backend().cancelled_exception_class()
|
||||
|
||||
|
||||
#
|
||||
# Private API
|
||||
#
|
||||
|
||||
|
||||
@contextmanager
|
||||
def claim_worker_thread(
|
||||
backend_class: type[AsyncBackend], token: object
|
||||
) -> Generator[Any, None, None]:
|
||||
from ..lowlevel import EventLoopToken
|
||||
|
||||
threadlocals.current_token = EventLoopToken(backend_class, token)
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
del threadlocals.current_token
|
||||
|
||||
|
||||
def get_async_backend(asynclib_name: str | None = None) -> type[AsyncBackend]:
|
||||
if asynclib_name is None:
|
||||
asynclib_name = current_async_library()
|
||||
if not asynclib_name:
|
||||
raise NoEventLoopError(
|
||||
f"Not currently running on any asynchronous event loop. "
|
||||
f"Available async backends: {', '.join(get_all_backends())}"
|
||||
)
|
||||
|
||||
# We use our own dict instead of sys.modules to get the already imported back-end
|
||||
# class because the appropriate modules in sys.modules could potentially be only
|
||||
# partially initialized
|
||||
try:
|
||||
return loaded_backends[asynclib_name]
|
||||
except KeyError:
|
||||
module = import_module(f"anyio._backends._{asynclib_name}")
|
||||
loaded_backends[asynclib_name] = module.backend_class
|
||||
return module.backend_class
|
||||
|
||||
|
||||
def current_async_library() -> str | None:
|
||||
if sniffio is None:
|
||||
# If sniffio is not installed, we assume we're either running asyncio or nothing
|
||||
import asyncio
|
||||
|
||||
try:
|
||||
asyncio.get_running_loop()
|
||||
return "asyncio"
|
||||
except RuntimeError:
|
||||
pass
|
||||
else:
|
||||
try:
|
||||
return sniffio.current_async_library()
|
||||
except sniffio.AsyncLibraryNotFoundError:
|
||||
pass
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def set_current_async_library(asynclib_name: str | None) -> Token | None:
|
||||
# no-op if sniffio is not installed
|
||||
if sniffio is None:
|
||||
return None
|
||||
|
||||
return sniffio.current_async_library_cvar.set(asynclib_name)
|
||||
|
||||
|
||||
def reset_current_async_library(token: Token | None) -> None:
|
||||
if token is not None:
|
||||
sniffio.current_async_library_cvar.reset(token)
|
||||
+205
@@ -0,0 +1,205 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
from collections.abc import Generator
|
||||
from textwrap import dedent
|
||||
from typing import Any
|
||||
|
||||
if sys.version_info < (3, 11):
|
||||
from exceptiongroup import BaseExceptionGroup
|
||||
|
||||
|
||||
class BrokenResourceError(Exception):
|
||||
"""
|
||||
Raised when trying to use a resource that has been rendered unusable due to external
|
||||
causes (e.g. a send stream whose peer has disconnected).
|
||||
"""
|
||||
|
||||
|
||||
class BrokenWorkerProcess(Exception):
|
||||
"""
|
||||
Raised by :meth:`~anyio.to_process.run_sync` if the worker process terminates abruptly or
|
||||
otherwise misbehaves.
|
||||
"""
|
||||
|
||||
|
||||
class BrokenWorkerInterpreter(Exception):
|
||||
"""
|
||||
Raised by :meth:`~anyio.to_interpreter.run_sync` if an unexpected exception is
|
||||
raised in the subinterpreter.
|
||||
"""
|
||||
|
||||
def __init__(self, excinfo: Any):
|
||||
# This was adapted from concurrent.futures.interpreter.ExecutionFailed
|
||||
msg = excinfo.formatted
|
||||
if not msg:
|
||||
if excinfo.type and excinfo.msg:
|
||||
msg = f"{excinfo.type.__name__}: {excinfo.msg}"
|
||||
else:
|
||||
msg = excinfo.type.__name__ or excinfo.msg
|
||||
|
||||
super().__init__(msg)
|
||||
self.excinfo = excinfo
|
||||
|
||||
def __str__(self) -> str:
|
||||
try:
|
||||
formatted = self.excinfo.errdisplay
|
||||
except Exception:
|
||||
return super().__str__()
|
||||
else:
|
||||
return dedent(
|
||||
f"""
|
||||
{super().__str__()}
|
||||
|
||||
Uncaught in the interpreter:
|
||||
|
||||
{formatted}
|
||||
""".strip()
|
||||
)
|
||||
|
||||
|
||||
class BusyResourceError(Exception):
|
||||
"""
|
||||
Raised when two tasks are trying to read from or write to the same resource
|
||||
concurrently.
|
||||
"""
|
||||
|
||||
def __init__(self, action: str):
|
||||
super().__init__(f"Another task is already {action} this resource")
|
||||
|
||||
|
||||
class ClosedResourceError(Exception):
|
||||
"""Raised when trying to use a resource that has been closed."""
|
||||
|
||||
|
||||
class ConnectionFailed(OSError):
|
||||
"""
|
||||
Raised when a connection attempt fails.
|
||||
|
||||
.. note:: This class inherits from :exc:`OSError` for backwards compatibility.
|
||||
"""
|
||||
|
||||
|
||||
def iterate_exceptions(
|
||||
exception: BaseException,
|
||||
) -> Generator[BaseException, None, None]:
|
||||
if isinstance(exception, BaseExceptionGroup):
|
||||
for exc in exception.exceptions:
|
||||
yield from iterate_exceptions(exc)
|
||||
else:
|
||||
yield exception
|
||||
|
||||
|
||||
class DelimiterNotFound(Exception):
|
||||
"""
|
||||
Raised during
|
||||
:meth:`~anyio.streams.buffered.BufferedByteReceiveStream.receive_until` if the
|
||||
maximum number of bytes has been read without the delimiter being found.
|
||||
"""
|
||||
|
||||
def __init__(self, max_bytes: int) -> None:
|
||||
super().__init__(
|
||||
f"The delimiter was not found among the first {max_bytes} bytes"
|
||||
)
|
||||
|
||||
|
||||
class EndOfStream(Exception):
|
||||
"""
|
||||
Raised when trying to read from a stream that has been closed from the other end.
|
||||
"""
|
||||
|
||||
|
||||
class IncompleteRead(Exception):
|
||||
"""
|
||||
Raised during
|
||||
:meth:`~anyio.streams.buffered.BufferedByteReceiveStream.receive_exactly` or
|
||||
:meth:`~anyio.streams.buffered.BufferedByteReceiveStream.receive_until` if the
|
||||
connection is closed before the requested amount of bytes has been read.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
super().__init__(
|
||||
"The stream was closed before the read operation could be completed"
|
||||
)
|
||||
|
||||
|
||||
class TypedAttributeLookupError(LookupError):
|
||||
"""
|
||||
Raised by :meth:`~anyio.TypedAttributeProvider.extra` when the given typed attribute
|
||||
is not found and no default value has been given.
|
||||
"""
|
||||
|
||||
|
||||
class WouldBlock(Exception):
|
||||
"""Raised by ``X_nowait`` functions if ``X()`` would block."""
|
||||
|
||||
|
||||
class NoEventLoopError(RuntimeError):
|
||||
"""
|
||||
Raised by several functions that require an event loop to be running in the current
|
||||
thread when there is no running event loop.
|
||||
|
||||
This is also raised by :func:`.from_thread.run` and :func:`.from_thread.run_sync`
|
||||
if not calling from an AnyIO worker thread, and no ``token`` was passed.
|
||||
"""
|
||||
|
||||
|
||||
class RunFinishedError(RuntimeError):
|
||||
"""
|
||||
Raised by :func:`.from_thread.run` and :func:`.from_thread.run_sync` if the event
|
||||
loop associated with the explicitly passed token has already finished.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
super().__init__(
|
||||
"The event loop associated with the given token has already finished"
|
||||
)
|
||||
|
||||
|
||||
class TaskFailed(Exception):
|
||||
"""
|
||||
Raised when awaiting on, or attempting to access the return value of, a
|
||||
:class:`.TaskHandle` that raised an exception.
|
||||
"""
|
||||
|
||||
|
||||
class TaskCancelled(TaskFailed):
|
||||
"""
|
||||
Raised when awaiting on, or attempting to access the return value of, a
|
||||
:class:`.TaskHandle` that was cancelled.
|
||||
"""
|
||||
|
||||
|
||||
class TaskNotFinished(Exception):
|
||||
"""
|
||||
Raised when attempting to access the return value or exception of a
|
||||
:class:`.TaskHandle` that is still pending completion.
|
||||
"""
|
||||
|
||||
|
||||
class FutureFailed(Exception):
|
||||
"""
|
||||
Raised when awaiting on, or attempting to access the return value of, a
|
||||
:class:`.Future` that raised an exception.
|
||||
"""
|
||||
|
||||
|
||||
class FutureCancelled(FutureFailed):
|
||||
"""
|
||||
Raised when attempting to access the return value or exception of a
|
||||
:class:`.Future` that was cancelled.
|
||||
"""
|
||||
|
||||
|
||||
class FutureNotFinished(Exception):
|
||||
"""
|
||||
Raised when attempting to access the return value or exception of a
|
||||
:class:`.Future` that is still pending completion.
|
||||
"""
|
||||
|
||||
|
||||
class FutureAlreadyFinished(Exception):
|
||||
"""
|
||||
Raised when attempting set a result of or await a
|
||||
:class:`.Future` that has already completed.
|
||||
"""
|
||||
+1045
File diff suppressed because it is too large.
Load diff
+191
@@ -0,0 +1,191 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Generator
|
||||
from enum import Enum, auto
|
||||
from typing import Any, Generic, TypeVar
|
||||
|
||||
from ._exceptions import (
|
||||
FutureAlreadyFinished,
|
||||
FutureCancelled,
|
||||
FutureFailed,
|
||||
FutureNotFinished,
|
||||
)
|
||||
from ._synchronization import Event
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
|
||||
class Future(Generic[T]):
|
||||
"""
|
||||
An awaitable object that works similar to a :class:`asyncio.Future` but
|
||||
with similar characteristics to a :class:`.TaskHandle`.
|
||||
"""
|
||||
|
||||
class Status(Enum):
|
||||
"""
|
||||
The status of a future handle.
|
||||
|
||||
.. attribute:: PENDING
|
||||
|
||||
The future has not finished yet.
|
||||
.. attribute:: FINISHED
|
||||
|
||||
The future has finished with a return value.
|
||||
.. attribute:: CANCELLED
|
||||
|
||||
The future was cancelled and has finished since.
|
||||
.. attribute:: FAILED
|
||||
|
||||
The future raised an exception.
|
||||
"""
|
||||
|
||||
PENDING = auto()
|
||||
FINISHED = auto()
|
||||
CANCELLED = auto()
|
||||
FAILED = auto()
|
||||
|
||||
__slots__ = (
|
||||
"_cancelled",
|
||||
"_exception",
|
||||
"_finished_event",
|
||||
"_name",
|
||||
"_return_value",
|
||||
)
|
||||
_return_value: T
|
||||
|
||||
def __init__(self, *, name: str | None = None) -> None:
|
||||
self._finished_event = Event()
|
||||
self._exception: BaseException | None = None
|
||||
self._cancelled: bool = False
|
||||
self._name = name
|
||||
|
||||
def _check_pending(self) -> None:
|
||||
"""Shortcut for checking if a Future is pending
|
||||
|
||||
:raises FutureAlreadyFinished: if future was already given a result or exception.
|
||||
:raises FutureCancelled: if future has been cancelled previously.
|
||||
"""
|
||||
match self.status:
|
||||
case Future.Status.PENDING:
|
||||
return
|
||||
case Future.Status.FINISHED:
|
||||
raise FutureAlreadyFinished("future has already finished")
|
||||
case Future.Status.FAILED:
|
||||
raise FutureAlreadyFinished("future already failed")
|
||||
case Future.Status.CANCELLED:
|
||||
raise FutureCancelled("future was cancelled")
|
||||
|
||||
async def wait(self) -> None:
|
||||
"""
|
||||
Waits for the future to finish.
|
||||
|
||||
This method will attempt to wait for a result or exception
|
||||
"""
|
||||
await self._finished_event.wait()
|
||||
|
||||
def cancel(self) -> None:
|
||||
"""Cancels a pending `.Future` object
|
||||
|
||||
Does nothing if the Future was already finished.
|
||||
"""
|
||||
if self.status is Future.Status.PENDING:
|
||||
self._cancelled = True
|
||||
self._finished_event.set()
|
||||
|
||||
@property
|
||||
def exception(self) -> BaseException | None:
|
||||
"""
|
||||
The exception value of a `.Future`
|
||||
|
||||
:raises FutureNotFinished: if future is still pending
|
||||
:raises FutureCancelled: if future was cancelled
|
||||
:returns: None if future succeeds with a result sent instead
|
||||
otherwise this will be an exception
|
||||
"""
|
||||
|
||||
match self.status:
|
||||
case Future.Status.PENDING:
|
||||
raise FutureNotFinished("the future has not finished yet")
|
||||
case Future.Status.FINISHED:
|
||||
return None
|
||||
case Future.Status.CANCELLED:
|
||||
raise FutureCancelled("the future was cancelled")
|
||||
case Future.Status.FAILED:
|
||||
return self._exception
|
||||
|
||||
@exception.setter
|
||||
def exception(self, exception: BaseException) -> None:
|
||||
"""Send exception for a `.Future` object
|
||||
|
||||
:raises FutureAlreadyFinished: if future was already given a result or exception.
|
||||
:raises FutureCancelled: if future has been cancelled previously.
|
||||
"""
|
||||
self._check_pending()
|
||||
self._exception = exception
|
||||
self._finished_event.set()
|
||||
|
||||
@property
|
||||
def return_value(self) -> T:
|
||||
"""
|
||||
The return value of the future.
|
||||
|
||||
:raises FutureNotFinished: if the future has not finished yet
|
||||
:raises FutureCancelled: if the future was cancelled
|
||||
:raises FutureFailed: if the future raised an exception
|
||||
|
||||
"""
|
||||
match self.status:
|
||||
case Future.Status.PENDING:
|
||||
raise FutureNotFinished("the future has not finished yet")
|
||||
case Future.Status.FINISHED:
|
||||
return self._return_value
|
||||
case Future.Status.CANCELLED:
|
||||
raise FutureCancelled("the future was cancelled")
|
||||
case Future.Status.FAILED:
|
||||
raise FutureFailed(
|
||||
"the future raised an exception"
|
||||
) from self._exception
|
||||
|
||||
@return_value.setter
|
||||
def return_value(self, value: T) -> None:
|
||||
"""
|
||||
Send pending result for a `.Future` object
|
||||
|
||||
:raises FutureAlreadyFinished: if future was already given a result or exception.
|
||||
:raises FutureCancelled: if future has been cancelled previously.
|
||||
"""
|
||||
self._check_pending()
|
||||
self._return_value = value
|
||||
self._finished_event.set()
|
||||
|
||||
@property
|
||||
def status(self) -> Future.Status:
|
||||
"""
|
||||
The current status of a future.
|
||||
|
||||
Every future starts in the :attr:`~Future.Status.PENDING` state.
|
||||
If a future is cancelled while in this state, it will transition to the
|
||||
:attr:`Future.Status.CANCELLED` state. When the task finishes, it will
|
||||
transition to one of the three final states (
|
||||
:attr:`Future.Status.FINISHED`, :attr:`Future.Status.FAILED`, or
|
||||
:attr:`Future.Status.CANCELLED`) depending on the exception the task
|
||||
raised, if any. No other status transitions will happen.
|
||||
"""
|
||||
if not self._finished_event.is_set():
|
||||
return Future.Status.PENDING
|
||||
elif self._cancelled:
|
||||
return Future.Status.CANCELLED
|
||||
elif self._exception is not None:
|
||||
return Future.Status.FAILED
|
||||
else:
|
||||
return Future.Status.FINISHED
|
||||
|
||||
def __await__(self) -> Generator[Any, Any, T]:
|
||||
yield from self.wait().__await__()
|
||||
return self.return_value
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return (
|
||||
f"<{self.__class__.__name__} {self.status.name.lower()} "
|
||||
f"name={self._name!r}>"
|
||||
)
|
||||
+18
@@ -0,0 +1,18 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from ..abc import AsyncResource # noqa: TC001
|
||||
from ._tasks import CancelScope
|
||||
|
||||
|
||||
async def aclose_forcefully(resource: AsyncResource) -> None:
|
||||
"""
|
||||
Close an asynchronous resource in a cancelled scope.
|
||||
|
||||
Doing this closes the resource without waiting on anything.
|
||||
|
||||
:param resource: the resource to close
|
||||
|
||||
"""
|
||||
with CancelScope() as scope:
|
||||
scope.cancel()
|
||||
await resource.aclose()
|
||||
+29
@@ -0,0 +1,29 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import AsyncIterator
|
||||
from contextlib import AbstractContextManager
|
||||
from signal import Signals
|
||||
|
||||
from ._eventloop import get_async_backend
|
||||
|
||||
|
||||
def open_signal_receiver(
|
||||
*signals: Signals,
|
||||
) -> AbstractContextManager[AsyncIterator[Signals]]:
|
||||
"""
|
||||
Start receiving operating system signals.
|
||||
|
||||
:param signals: signals to receive (e.g. ``signal.SIGINT``)
|
||||
:return: an asynchronous context manager for an asynchronous iterator which yields
|
||||
signal numbers
|
||||
:raises NoEventLoopError: if no supported asynchronous event loop is running in the
|
||||
current thread
|
||||
|
||||
.. warning:: Windows does not support signals natively so it is best to avoid
|
||||
relying on this in cross-platform applications.
|
||||
|
||||
.. warning:: On asyncio, this permanently replaces any previous signal handler for
|
||||
the given signals, as set via :meth:`~asyncio.loop.add_signal_handler`.
|
||||
|
||||
"""
|
||||
return get_async_backend().open_signal_receiver(*signals)
|
||||
+1011
File diff suppressed because it is too large.
Load diff
+52
@@ -0,0 +1,52 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
from typing import TypeVar
|
||||
from warnings import warn
|
||||
|
||||
from ..streams.memory import (
|
||||
MemoryObjectReceiveStream,
|
||||
MemoryObjectSendStream,
|
||||
_MemoryObjectStreamState,
|
||||
)
|
||||
|
||||
T_Item = TypeVar("T_Item")
|
||||
|
||||
|
||||
class create_memory_object_stream(
|
||||
tuple[MemoryObjectSendStream[T_Item], MemoryObjectReceiveStream[T_Item]],
|
||||
):
|
||||
"""
|
||||
Create a memory object stream.
|
||||
|
||||
The stream's item type can be annotated like
|
||||
:func:`create_memory_object_stream[T_Item]`.
|
||||
|
||||
:param max_buffer_size: number of items held in the buffer until ``send()`` starts
|
||||
blocking
|
||||
:param item_type: old way of marking the streams with the right generic type for
|
||||
static typing (does nothing on AnyIO 4)
|
||||
|
||||
.. deprecated:: 4.0
|
||||
Use ``create_memory_object_stream[YourItemType](...)`` instead.
|
||||
:return: a tuple of (send stream, receive stream)
|
||||
|
||||
"""
|
||||
|
||||
def __new__( # type: ignore[misc]
|
||||
cls, max_buffer_size: float = 0, item_type: object = None
|
||||
) -> tuple[MemoryObjectSendStream[T_Item], MemoryObjectReceiveStream[T_Item]]:
|
||||
if max_buffer_size != math.inf and not isinstance(max_buffer_size, int):
|
||||
raise ValueError("max_buffer_size must be either an integer or math.inf")
|
||||
if max_buffer_size < 0:
|
||||
raise ValueError("max_buffer_size cannot be negative")
|
||||
if item_type is not None:
|
||||
warn(
|
||||
"The item_type argument has been deprecated in AnyIO 4.0. "
|
||||
"Use create_memory_object_stream[YourItemType](...) instead.",
|
||||
DeprecationWarning,
|
||||
stacklevel=2,
|
||||
)
|
||||
|
||||
state = _MemoryObjectStreamState[T_Item](max_buffer_size)
|
||||
return (MemoryObjectSendStream(state), MemoryObjectReceiveStream(state))
|
||||
+196
@@ -0,0 +1,196 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import AsyncIterable, Iterable, Mapping, Sequence
|
||||
from io import BytesIO
|
||||
from os import PathLike
|
||||
from subprocess import PIPE, CalledProcessError, CompletedProcess
|
||||
from typing import IO, Any, TypeAlias, cast
|
||||
|
||||
from ..abc import Process # noqa: TC001
|
||||
from ._eventloop import get_async_backend
|
||||
from ._tasks import create_task_group
|
||||
|
||||
StrOrBytesPath: TypeAlias = str | bytes | PathLike[str] | PathLike[bytes]
|
||||
|
||||
|
||||
async def run_process(
|
||||
command: StrOrBytesPath | Sequence[StrOrBytesPath],
|
||||
*,
|
||||
input: bytes | None = None,
|
||||
stdin: int | IO[Any] | None = None,
|
||||
stdout: int | IO[Any] | None = PIPE,
|
||||
stderr: int | IO[Any] | None = PIPE,
|
||||
check: bool = True,
|
||||
cwd: StrOrBytesPath | None = None,
|
||||
env: Mapping[str, str] | None = None,
|
||||
startupinfo: Any = None,
|
||||
creationflags: int = 0,
|
||||
start_new_session: bool = False,
|
||||
pass_fds: Sequence[int] = (),
|
||||
user: str | int | None = None,
|
||||
group: str | int | None = None,
|
||||
extra_groups: Iterable[str | int] | None = None,
|
||||
umask: int = -1,
|
||||
) -> CompletedProcess[bytes]:
|
||||
"""
|
||||
Run an external command in a subprocess and wait until it completes.
|
||||
|
||||
.. seealso:: :func:`subprocess.run`
|
||||
|
||||
:param command: either a string to pass to the shell, or an iterable of strings
|
||||
containing the executable name or path and its arguments
|
||||
:param input: bytes passed to the standard input of the subprocess
|
||||
:param stdin: one of :data:`subprocess.PIPE`, :data:`subprocess.DEVNULL`,
|
||||
a file-like object, or `None`; ``input`` overrides this
|
||||
:param stdout: one of :data:`subprocess.PIPE`, :data:`subprocess.DEVNULL`,
|
||||
a file-like object, or `None`
|
||||
:param stderr: one of :data:`subprocess.PIPE`, :data:`subprocess.DEVNULL`,
|
||||
:data:`subprocess.STDOUT`, a file-like object, or `None`
|
||||
:param check: if ``True``, raise :exc:`~subprocess.CalledProcessError` if the
|
||||
process terminates with a return code other than 0
|
||||
:param cwd: If not ``None``, change the working directory to this before running the
|
||||
command
|
||||
:param env: if not ``None``, this mapping replaces the inherited environment
|
||||
variables from the parent process
|
||||
:param startupinfo: an instance of :class:`subprocess.STARTUPINFO` that can be used
|
||||
to specify process startup parameters (Windows only)
|
||||
:param creationflags: flags that can be used to control the creation of the
|
||||
subprocess (see :class:`subprocess.Popen` for the specifics)
|
||||
:param start_new_session: if ``true`` the setsid() system call will be made in the
|
||||
child process prior to the execution of the subprocess. (POSIX only)
|
||||
:param pass_fds: sequence of file descriptors to keep open between the parent and
|
||||
child processes. (POSIX only)
|
||||
:param user: effective user to run the process as (Python >= 3.9, POSIX only)
|
||||
:param group: effective group to run the process as (Python >= 3.9, POSIX only)
|
||||
:param extra_groups: supplementary groups to set in the subprocess (Python >= 3.9,
|
||||
POSIX only)
|
||||
:param umask: if not negative, this umask is applied in the child process before
|
||||
running the given command (Python >= 3.9, POSIX only)
|
||||
:return: an object representing the completed process
|
||||
:raises ~subprocess.CalledProcessError: if ``check`` is ``True`` and the process
|
||||
exits with a nonzero return code
|
||||
|
||||
"""
|
||||
|
||||
async def drain_stream(stream: AsyncIterable[bytes], index: int) -> None:
|
||||
buffer = BytesIO()
|
||||
async for chunk in stream:
|
||||
buffer.write(chunk)
|
||||
|
||||
stream_contents[index] = buffer.getvalue()
|
||||
|
||||
if stdin is not None and input is not None:
|
||||
raise ValueError("only one of stdin and input is allowed")
|
||||
|
||||
async with await open_process(
|
||||
command,
|
||||
stdin=PIPE if input else stdin,
|
||||
stdout=stdout,
|
||||
stderr=stderr,
|
||||
cwd=cwd,
|
||||
env=env,
|
||||
startupinfo=startupinfo,
|
||||
creationflags=creationflags,
|
||||
start_new_session=start_new_session,
|
||||
pass_fds=pass_fds,
|
||||
user=user,
|
||||
group=group,
|
||||
extra_groups=extra_groups,
|
||||
umask=umask,
|
||||
) as process:
|
||||
stream_contents: list[bytes | None] = [None, None]
|
||||
async with create_task_group() as tg:
|
||||
if process.stdout:
|
||||
tg.start_soon(drain_stream, process.stdout, 0)
|
||||
|
||||
if process.stderr:
|
||||
tg.start_soon(drain_stream, process.stderr, 1)
|
||||
|
||||
if process.stdin and input:
|
||||
await process.stdin.send(input)
|
||||
await process.stdin.aclose()
|
||||
|
||||
await process.wait()
|
||||
|
||||
output, errors = stream_contents
|
||||
if check and process.returncode != 0:
|
||||
raise CalledProcessError(cast(int, process.returncode), command, output, errors)
|
||||
|
||||
return CompletedProcess(command, cast(int, process.returncode), output, errors)
|
||||
|
||||
|
||||
async def open_process(
|
||||
command: StrOrBytesPath | Sequence[StrOrBytesPath],
|
||||
*,
|
||||
stdin: int | IO[Any] | None = PIPE,
|
||||
stdout: int | IO[Any] | None = PIPE,
|
||||
stderr: int | IO[Any] | None = PIPE,
|
||||
cwd: StrOrBytesPath | None = None,
|
||||
env: Mapping[str, str] | None = None,
|
||||
startupinfo: Any = None,
|
||||
creationflags: int = 0,
|
||||
start_new_session: bool = False,
|
||||
pass_fds: Sequence[int] = (),
|
||||
user: str | int | None = None,
|
||||
group: str | int | None = None,
|
||||
extra_groups: Iterable[str | int] | None = None,
|
||||
umask: int = -1,
|
||||
) -> Process:
|
||||
"""
|
||||
Start an external command in a subprocess.
|
||||
|
||||
.. seealso:: :class:`subprocess.Popen`
|
||||
|
||||
:param command: either a string to pass to the shell, or an iterable of strings
|
||||
containing the executable name or path and its arguments
|
||||
:param stdin: one of :data:`subprocess.PIPE`, :data:`subprocess.DEVNULL`, a
|
||||
file-like object, or ``None``
|
||||
:param stdout: one of :data:`subprocess.PIPE`, :data:`subprocess.DEVNULL`,
|
||||
a file-like object, or ``None``
|
||||
:param stderr: one of :data:`subprocess.PIPE`, :data:`subprocess.DEVNULL`,
|
||||
:data:`subprocess.STDOUT`, a file-like object, or ``None``
|
||||
:param cwd: If not ``None``, the working directory is changed before executing
|
||||
:param env: If env is not ``None``, it must be a mapping that defines the
|
||||
environment variables for the new process
|
||||
:param creationflags: flags that can be used to control the creation of the
|
||||
subprocess (see :class:`subprocess.Popen` for the specifics)
|
||||
:param startupinfo: an instance of :class:`subprocess.STARTUPINFO` that can be used
|
||||
to specify process startup parameters (Windows only)
|
||||
:param start_new_session: if ``true`` the setsid() system call will be made in the
|
||||
child process prior to the execution of the subprocess. (POSIX only)
|
||||
:param pass_fds: sequence of file descriptors to keep open between the parent and
|
||||
child processes. (POSIX only)
|
||||
:param user: effective user to run the process as (POSIX only)
|
||||
:param group: effective group to run the process as (POSIX only)
|
||||
:param extra_groups: supplementary groups to set in the subprocess (POSIX only)
|
||||
:param umask: if not negative, this umask is applied in the child process before
|
||||
running the given command (POSIX only)
|
||||
:return: an asynchronous process object
|
||||
|
||||
"""
|
||||
kwargs: dict[str, Any] = {}
|
||||
if user is not None:
|
||||
kwargs["user"] = user
|
||||
|
||||
if group is not None:
|
||||
kwargs["group"] = group
|
||||
|
||||
if extra_groups is not None:
|
||||
kwargs["extra_groups"] = extra_groups
|
||||
|
||||
if umask >= 0:
|
||||
kwargs["umask"] = umask
|
||||
|
||||
return await get_async_backend().open_process(
|
||||
command,
|
||||
stdin=stdin,
|
||||
stdout=stdout,
|
||||
stderr=stderr,
|
||||
cwd=cwd,
|
||||
env=env,
|
||||
startupinfo=startupinfo,
|
||||
creationflags=creationflags,
|
||||
start_new_session=start_new_session,
|
||||
pass_fds=pass_fds,
|
||||
**kwargs,
|
||||
)
|
||||
Vendored
+778
@@ -0,0 +1,778 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
import sys
|
||||
from collections import deque
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass
|
||||
from types import TracebackType
|
||||
from typing import TypeVar
|
||||
|
||||
from ..lowlevel import checkpoint_if_cancelled
|
||||
from ._eventloop import get_async_backend
|
||||
from ._exceptions import BusyResourceError, NoEventLoopError
|
||||
from ._tasks import CancelScope
|
||||
from ._testing import TaskInfo, get_current_task
|
||||
|
||||
if sys.version_info >= (3, 11):
|
||||
from typing import Self
|
||||
else:
|
||||
from typing_extensions import Self
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class EventStatistics:
|
||||
"""
|
||||
:ivar int tasks_waiting: number of tasks waiting on :meth:`~.Event.wait`
|
||||
"""
|
||||
|
||||
tasks_waiting: int
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class CapacityLimiterStatistics:
|
||||
"""
|
||||
:ivar int borrowed_tokens: number of tokens currently borrowed by tasks
|
||||
:ivar float total_tokens: total number of available tokens
|
||||
:ivar tuple borrowers: tasks or other objects currently holding tokens borrowed from
|
||||
this limiter
|
||||
:ivar int tasks_waiting: number of tasks waiting on
|
||||
:meth:`~.CapacityLimiter.acquire` or
|
||||
:meth:`~.CapacityLimiter.acquire_on_behalf_of`
|
||||
"""
|
||||
|
||||
borrowed_tokens: int
|
||||
total_tokens: float
|
||||
borrowers: tuple[object, ...]
|
||||
tasks_waiting: int
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class LockStatistics:
|
||||
"""
|
||||
:ivar bool locked: flag indicating if this lock is locked or not
|
||||
:ivar ~anyio.TaskInfo owner: task currently holding the lock (or ``None`` if the
|
||||
lock is not held by any task)
|
||||
:ivar int tasks_waiting: number of tasks waiting on :meth:`~.Lock.acquire`
|
||||
"""
|
||||
|
||||
locked: bool
|
||||
owner: TaskInfo | None
|
||||
tasks_waiting: int
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ConditionStatistics:
|
||||
"""
|
||||
:ivar int tasks_waiting: number of tasks blocked on :meth:`~.Condition.wait`
|
||||
:ivar ~anyio.LockStatistics lock_statistics: statistics of the underlying
|
||||
:class:`~.Lock`
|
||||
"""
|
||||
|
||||
tasks_waiting: int
|
||||
lock_statistics: LockStatistics
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class SemaphoreStatistics:
|
||||
"""
|
||||
:ivar int tasks_waiting: number of tasks waiting on :meth:`~.Semaphore.acquire`
|
||||
|
||||
"""
|
||||
|
||||
tasks_waiting: int
|
||||
|
||||
|
||||
class Event:
|
||||
__slots__ = ("__weakref__",)
|
||||
|
||||
def __new__(cls) -> Event:
|
||||
try:
|
||||
return get_async_backend().create_event()
|
||||
except NoEventLoopError:
|
||||
return EventAdapter()
|
||||
|
||||
def set(self) -> None:
|
||||
"""Set the flag, notifying all listeners."""
|
||||
raise NotImplementedError
|
||||
|
||||
def is_set(self) -> bool:
|
||||
"""Return ``True`` if the flag is set, ``False`` if not."""
|
||||
raise NotImplementedError
|
||||
|
||||
async def wait(self) -> None:
|
||||
"""
|
||||
Wait until the flag has been set.
|
||||
|
||||
If the flag has already been set when this method is called, it returns
|
||||
immediately.
|
||||
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
def statistics(self) -> EventStatistics:
|
||||
"""Return statistics about the current state of this event."""
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
class EventAdapter(Event):
|
||||
__slots__ = "_internal_event", "_is_set"
|
||||
|
||||
def __new__(cls) -> Self:
|
||||
return object.__new__(cls)
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._internal_event: Event | None = None
|
||||
self._is_set = False
|
||||
|
||||
@property
|
||||
def _event(self) -> Event:
|
||||
if self._internal_event is None:
|
||||
self._internal_event = get_async_backend().create_event()
|
||||
if self._is_set:
|
||||
self._internal_event.set()
|
||||
|
||||
return self._internal_event
|
||||
|
||||
def set(self) -> None:
|
||||
if self._internal_event is None:
|
||||
self._is_set = True
|
||||
else:
|
||||
self._event.set()
|
||||
|
||||
def is_set(self) -> bool:
|
||||
if self._internal_event is None:
|
||||
return self._is_set
|
||||
|
||||
return self._internal_event.is_set()
|
||||
|
||||
async def wait(self) -> None:
|
||||
await self._event.wait()
|
||||
|
||||
def statistics(self) -> EventStatistics:
|
||||
if self._internal_event is None:
|
||||
return EventStatistics(tasks_waiting=0)
|
||||
|
||||
return self._internal_event.statistics()
|
||||
|
||||
|
||||
class Lock:
|
||||
__slots__ = ("__weakref__",)
|
||||
|
||||
def __new__(cls, *, fast_acquire: bool = False) -> Lock:
|
||||
try:
|
||||
return get_async_backend().create_lock(fast_acquire=fast_acquire)
|
||||
except NoEventLoopError:
|
||||
return LockAdapter(fast_acquire=fast_acquire)
|
||||
|
||||
async def __aenter__(self) -> None:
|
||||
await self.acquire()
|
||||
|
||||
async def __aexit__(
|
||||
self,
|
||||
exc_type: type[BaseException] | None,
|
||||
exc_val: BaseException | None,
|
||||
exc_tb: TracebackType | None,
|
||||
) -> None:
|
||||
self.release()
|
||||
|
||||
async def acquire(self) -> None:
|
||||
"""Acquire the lock."""
|
||||
raise NotImplementedError
|
||||
|
||||
def acquire_nowait(self) -> None:
|
||||
"""
|
||||
Acquire the lock, without blocking.
|
||||
|
||||
:raises ~anyio.WouldBlock: if the operation would block
|
||||
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
def release(self) -> None:
|
||||
"""Release the lock."""
|
||||
raise NotImplementedError
|
||||
|
||||
def locked(self) -> bool:
|
||||
"""Return True if the lock is currently held."""
|
||||
raise NotImplementedError
|
||||
|
||||
def statistics(self) -> LockStatistics:
|
||||
"""
|
||||
Return statistics about the current state of this lock.
|
||||
|
||||
.. versionadded:: 3.0
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
class LockAdapter(Lock):
|
||||
__slots__ = "_fast_acquire", "_internal_lock"
|
||||
|
||||
def __new__(cls, *, fast_acquire: bool = False) -> Self:
|
||||
return object.__new__(cls)
|
||||
|
||||
def __init__(self, *, fast_acquire: bool = False):
|
||||
self._internal_lock: Lock | None = None
|
||||
self._fast_acquire = fast_acquire
|
||||
|
||||
@property
|
||||
def _lock(self) -> Lock:
|
||||
if self._internal_lock is None:
|
||||
self._internal_lock = get_async_backend().create_lock(
|
||||
fast_acquire=self._fast_acquire
|
||||
)
|
||||
|
||||
return self._internal_lock
|
||||
|
||||
async def __aenter__(self) -> None:
|
||||
await self._lock.acquire()
|
||||
|
||||
async def __aexit__(
|
||||
self,
|
||||
exc_type: type[BaseException] | None,
|
||||
exc_val: BaseException | None,
|
||||
exc_tb: TracebackType | None,
|
||||
) -> None:
|
||||
if self._internal_lock is not None:
|
||||
self._internal_lock.release()
|
||||
|
||||
async def acquire(self) -> None:
|
||||
"""Acquire the lock."""
|
||||
await self._lock.acquire()
|
||||
|
||||
def acquire_nowait(self) -> None:
|
||||
"""
|
||||
Acquire the lock, without blocking.
|
||||
|
||||
:raises ~anyio.WouldBlock: if the operation would block
|
||||
|
||||
"""
|
||||
self._lock.acquire_nowait()
|
||||
|
||||
def release(self) -> None:
|
||||
"""Release the lock."""
|
||||
self._lock.release()
|
||||
|
||||
def locked(self) -> bool:
|
||||
"""Return True if the lock is currently held."""
|
||||
return self._lock.locked()
|
||||
|
||||
def statistics(self) -> LockStatistics:
|
||||
"""
|
||||
Return statistics about the current state of this lock.
|
||||
|
||||
.. versionadded:: 3.0
|
||||
|
||||
"""
|
||||
if self._internal_lock is None:
|
||||
return LockStatistics(False, None, 0)
|
||||
|
||||
return self._internal_lock.statistics()
|
||||
|
||||
|
||||
class Condition:
|
||||
__slots__ = "__weakref__", "_lock", "_owner_task", "_waiters"
|
||||
|
||||
def __init__(self, lock: Lock | None = None):
|
||||
self._owner_task: TaskInfo | None = None
|
||||
self._lock = lock or Lock()
|
||||
self._waiters: deque[Event] = deque()
|
||||
|
||||
async def __aenter__(self) -> None:
|
||||
await self.acquire()
|
||||
|
||||
async def __aexit__(
|
||||
self,
|
||||
exc_type: type[BaseException] | None,
|
||||
exc_val: BaseException | None,
|
||||
exc_tb: TracebackType | None,
|
||||
) -> None:
|
||||
self.release()
|
||||
|
||||
def _check_acquired(self) -> None:
|
||||
if self._owner_task != get_current_task():
|
||||
raise RuntimeError("The current task is not holding the underlying lock")
|
||||
|
||||
async def acquire(self) -> None:
|
||||
"""Acquire the underlying lock."""
|
||||
await self._lock.acquire()
|
||||
self._owner_task = get_current_task()
|
||||
|
||||
def acquire_nowait(self) -> None:
|
||||
"""
|
||||
Acquire the underlying lock, without blocking.
|
||||
|
||||
:raises ~anyio.WouldBlock: if the operation would block
|
||||
|
||||
"""
|
||||
self._lock.acquire_nowait()
|
||||
self._owner_task = get_current_task()
|
||||
|
||||
def release(self) -> None:
|
||||
"""Release the underlying lock."""
|
||||
self._lock.release()
|
||||
|
||||
def locked(self) -> bool:
|
||||
"""Return True if the lock is set."""
|
||||
return self._lock.locked()
|
||||
|
||||
def notify(self, n: int = 1) -> None:
|
||||
"""Notify exactly n listeners."""
|
||||
self._check_acquired()
|
||||
for _ in range(n):
|
||||
try:
|
||||
event = self._waiters.popleft()
|
||||
except IndexError:
|
||||
break
|
||||
|
||||
event.set()
|
||||
|
||||
def notify_all(self) -> None:
|
||||
"""Notify all the listeners."""
|
||||
self._check_acquired()
|
||||
for event in self._waiters:
|
||||
event.set()
|
||||
|
||||
self._waiters.clear()
|
||||
|
||||
async def wait(self) -> None:
|
||||
"""Wait for a notification."""
|
||||
await checkpoint_if_cancelled()
|
||||
self._check_acquired()
|
||||
event = Event()
|
||||
self._waiters.append(event)
|
||||
self.release()
|
||||
try:
|
||||
await event.wait()
|
||||
except BaseException:
|
||||
if not event.is_set():
|
||||
self._waiters.remove(event)
|
||||
elif self._waiters:
|
||||
# This task was notified by could not act on it, so pass
|
||||
# it on to the next task
|
||||
self._waiters.popleft().set()
|
||||
|
||||
raise
|
||||
finally:
|
||||
with CancelScope(shield=True):
|
||||
await self.acquire()
|
||||
|
||||
async def wait_for(self, predicate: Callable[[], T]) -> T:
|
||||
"""
|
||||
Wait until a predicate becomes true.
|
||||
|
||||
:param predicate: a callable that returns a truthy value when the condition is
|
||||
met
|
||||
:return: the result of the predicate
|
||||
|
||||
.. versionadded:: 4.11.0
|
||||
|
||||
"""
|
||||
while not (result := predicate()):
|
||||
await self.wait()
|
||||
|
||||
return result
|
||||
|
||||
def statistics(self) -> ConditionStatistics:
|
||||
"""
|
||||
Return statistics about the current state of this condition.
|
||||
|
||||
.. versionadded:: 3.0
|
||||
"""
|
||||
return ConditionStatistics(len(self._waiters), self._lock.statistics())
|
||||
|
||||
|
||||
class Semaphore:
|
||||
__slots__ = "__weakref__", "_fast_acquire"
|
||||
|
||||
def __new__(
|
||||
cls,
|
||||
initial_value: int,
|
||||
*,
|
||||
max_value: int | None = None,
|
||||
fast_acquire: bool = False,
|
||||
) -> Semaphore:
|
||||
try:
|
||||
return get_async_backend().create_semaphore(
|
||||
initial_value, max_value=max_value, fast_acquire=fast_acquire
|
||||
)
|
||||
except NoEventLoopError:
|
||||
return SemaphoreAdapter(initial_value, max_value=max_value)
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
initial_value: int,
|
||||
*,
|
||||
max_value: int | None = None,
|
||||
fast_acquire: bool = False,
|
||||
):
|
||||
if not isinstance(initial_value, int):
|
||||
raise TypeError("initial_value must be an integer")
|
||||
if initial_value < 0:
|
||||
raise ValueError("initial_value must be >= 0")
|
||||
if max_value is not None:
|
||||
if not isinstance(max_value, int):
|
||||
raise TypeError("max_value must be an integer or None")
|
||||
if max_value < initial_value:
|
||||
raise ValueError(
|
||||
"max_value must be equal to or higher than initial_value"
|
||||
)
|
||||
|
||||
self._fast_acquire = fast_acquire
|
||||
|
||||
async def __aenter__(self) -> Self:
|
||||
await self.acquire()
|
||||
return self
|
||||
|
||||
async def __aexit__(
|
||||
self,
|
||||
exc_type: type[BaseException] | None,
|
||||
exc_val: BaseException | None,
|
||||
exc_tb: TracebackType | None,
|
||||
) -> None:
|
||||
self.release()
|
||||
|
||||
async def acquire(self) -> None:
|
||||
"""Decrement the semaphore value, blocking if necessary."""
|
||||
raise NotImplementedError
|
||||
|
||||
def acquire_nowait(self) -> None:
|
||||
"""
|
||||
Acquire the underlying lock, without blocking.
|
||||
|
||||
:raises ~anyio.WouldBlock: if the operation would block
|
||||
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
def release(self) -> None:
|
||||
"""Increment the semaphore value."""
|
||||
raise NotImplementedError
|
||||
|
||||
@property
|
||||
def value(self) -> int:
|
||||
"""The current value of the semaphore."""
|
||||
raise NotImplementedError
|
||||
|
||||
@property
|
||||
def max_value(self) -> int | None:
|
||||
"""The maximum value of the semaphore."""
|
||||
raise NotImplementedError
|
||||
|
||||
def statistics(self) -> SemaphoreStatistics:
|
||||
"""
|
||||
Return statistics about the current state of this semaphore.
|
||||
|
||||
.. versionadded:: 3.0
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
class SemaphoreAdapter(Semaphore):
|
||||
__slots__ = "_initial_value", "_internal_semaphore", "_max_value"
|
||||
|
||||
def __new__(
|
||||
cls,
|
||||
initial_value: int,
|
||||
*,
|
||||
max_value: int | None = None,
|
||||
fast_acquire: bool = False,
|
||||
) -> Self:
|
||||
return object.__new__(cls)
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
initial_value: int,
|
||||
*,
|
||||
max_value: int | None = None,
|
||||
fast_acquire: bool = False,
|
||||
) -> None:
|
||||
super().__init__(initial_value, max_value=max_value, fast_acquire=fast_acquire)
|
||||
self._internal_semaphore: Semaphore | None = None
|
||||
self._initial_value = initial_value
|
||||
self._max_value = max_value
|
||||
|
||||
@property
|
||||
def _semaphore(self) -> Semaphore:
|
||||
if self._internal_semaphore is None:
|
||||
self._internal_semaphore = get_async_backend().create_semaphore(
|
||||
self._initial_value, max_value=self._max_value
|
||||
)
|
||||
|
||||
return self._internal_semaphore
|
||||
|
||||
async def acquire(self) -> None:
|
||||
await self._semaphore.acquire()
|
||||
|
||||
def acquire_nowait(self) -> None:
|
||||
self._semaphore.acquire_nowait()
|
||||
|
||||
def release(self) -> None:
|
||||
self._semaphore.release()
|
||||
|
||||
@property
|
||||
def value(self) -> int:
|
||||
if self._internal_semaphore is None:
|
||||
return self._initial_value
|
||||
|
||||
return self._semaphore.value
|
||||
|
||||
@property
|
||||
def max_value(self) -> int | None:
|
||||
return self._max_value
|
||||
|
||||
def statistics(self) -> SemaphoreStatistics:
|
||||
if self._internal_semaphore is None:
|
||||
return SemaphoreStatistics(tasks_waiting=0)
|
||||
|
||||
return self._semaphore.statistics()
|
||||
|
||||
|
||||
class CapacityLimiter:
|
||||
__slots__ = ("__weakref__",)
|
||||
|
||||
def __new__(cls, total_tokens: float) -> CapacityLimiter:
|
||||
try:
|
||||
return get_async_backend().create_capacity_limiter(total_tokens)
|
||||
except NoEventLoopError:
|
||||
return CapacityLimiterAdapter(total_tokens)
|
||||
|
||||
async def __aenter__(self) -> None:
|
||||
raise NotImplementedError
|
||||
|
||||
async def __aexit__(
|
||||
self,
|
||||
exc_type: type[BaseException] | None,
|
||||
exc_val: BaseException | None,
|
||||
exc_tb: TracebackType | None,
|
||||
) -> None:
|
||||
raise NotImplementedError
|
||||
|
||||
@property
|
||||
def total_tokens(self) -> float:
|
||||
"""
|
||||
The total number of tokens available for borrowing.
|
||||
|
||||
This is a read-write property. If the total number of tokens is increased, the
|
||||
proportionate number of tasks waiting on this limiter will be granted their
|
||||
tokens.
|
||||
|
||||
.. versionchanged:: 3.0
|
||||
The property is now writable.
|
||||
.. versionchanged:: 4.12
|
||||
The value can now be set to 0.
|
||||
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
@total_tokens.setter
|
||||
def total_tokens(self, value: float) -> None:
|
||||
raise NotImplementedError
|
||||
|
||||
@property
|
||||
def borrowed_tokens(self) -> int:
|
||||
"""The number of tokens that have currently been borrowed."""
|
||||
raise NotImplementedError
|
||||
|
||||
@property
|
||||
def available_tokens(self) -> float:
|
||||
"""The number of tokens currently available to be borrowed"""
|
||||
raise NotImplementedError
|
||||
|
||||
def acquire_nowait(self) -> None:
|
||||
"""
|
||||
Acquire a token for the current task without waiting for one to become
|
||||
available.
|
||||
|
||||
:raises ~anyio.WouldBlock: if there are no tokens available for borrowing
|
||||
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
def acquire_on_behalf_of_nowait(self, borrower: object) -> None:
|
||||
"""
|
||||
Acquire a token without waiting for one to become available.
|
||||
|
||||
:param borrower: the entity borrowing a token
|
||||
:raises ~anyio.WouldBlock: if there are no tokens available for borrowing
|
||||
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
async def acquire(self) -> None:
|
||||
"""
|
||||
Acquire a token for the current task, waiting if necessary for one to become
|
||||
available.
|
||||
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
async def acquire_on_behalf_of(self, borrower: object) -> None:
|
||||
"""
|
||||
Acquire a token, waiting if necessary for one to become available.
|
||||
|
||||
:param borrower: the entity borrowing a token
|
||||
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
def release(self) -> None:
|
||||
"""
|
||||
Release the token held by the current task.
|
||||
|
||||
:raises RuntimeError: if the current task has not borrowed a token from this
|
||||
limiter.
|
||||
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
def release_on_behalf_of(self, borrower: object) -> None:
|
||||
"""
|
||||
Release the token held by the given borrower.
|
||||
|
||||
:raises RuntimeError: if the borrower has not borrowed a token from this
|
||||
limiter.
|
||||
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
def statistics(self) -> CapacityLimiterStatistics:
|
||||
"""
|
||||
Return statistics about the current state of this limiter.
|
||||
|
||||
.. versionadded:: 3.0
|
||||
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
class CapacityLimiterAdapter(CapacityLimiter):
|
||||
__slots__ = "_internal_limiter", "_total_tokens"
|
||||
|
||||
def __new__(cls, total_tokens: float) -> Self:
|
||||
return object.__new__(cls)
|
||||
|
||||
def __init__(self, total_tokens: float) -> None:
|
||||
self._internal_limiter: CapacityLimiter | None = None
|
||||
self.total_tokens = total_tokens
|
||||
|
||||
@property
|
||||
def _limiter(self) -> CapacityLimiter:
|
||||
if self._internal_limiter is None:
|
||||
self._internal_limiter = get_async_backend().create_capacity_limiter(
|
||||
self._total_tokens
|
||||
)
|
||||
|
||||
return self._internal_limiter
|
||||
|
||||
async def __aenter__(self) -> None:
|
||||
await self._limiter.__aenter__()
|
||||
|
||||
async def __aexit__(
|
||||
self,
|
||||
exc_type: type[BaseException] | None,
|
||||
exc_val: BaseException | None,
|
||||
exc_tb: TracebackType | None,
|
||||
) -> None:
|
||||
return await self._limiter.__aexit__(exc_type, exc_val, exc_tb)
|
||||
|
||||
@property
|
||||
def total_tokens(self) -> float:
|
||||
if self._internal_limiter is None:
|
||||
return self._total_tokens
|
||||
|
||||
return self._internal_limiter.total_tokens
|
||||
|
||||
@total_tokens.setter
|
||||
def total_tokens(self, value: float) -> None:
|
||||
if not isinstance(value, int) and not math.isinf(value):
|
||||
raise TypeError("total_tokens must be an int or math.inf")
|
||||
elif value < 0:
|
||||
raise ValueError("total_tokens must be >= 0")
|
||||
|
||||
if self._internal_limiter is None:
|
||||
self._total_tokens = value
|
||||
return
|
||||
|
||||
self._limiter.total_tokens = value
|
||||
|
||||
@property
|
||||
def borrowed_tokens(self) -> int:
|
||||
if self._internal_limiter is None:
|
||||
return 0
|
||||
|
||||
return self._internal_limiter.borrowed_tokens
|
||||
|
||||
@property
|
||||
def available_tokens(self) -> float:
|
||||
if self._internal_limiter is None:
|
||||
return self._total_tokens
|
||||
|
||||
return self._internal_limiter.available_tokens
|
||||
|
||||
def acquire_nowait(self) -> None:
|
||||
self._limiter.acquire_nowait()
|
||||
|
||||
def acquire_on_behalf_of_nowait(self, borrower: object) -> None:
|
||||
self._limiter.acquire_on_behalf_of_nowait(borrower)
|
||||
|
||||
async def acquire(self) -> None:
|
||||
await self._limiter.acquire()
|
||||
|
||||
async def acquire_on_behalf_of(self, borrower: object) -> None:
|
||||
await self._limiter.acquire_on_behalf_of(borrower)
|
||||
|
||||
def release(self) -> None:
|
||||
self._limiter.release()
|
||||
|
||||
def release_on_behalf_of(self, borrower: object) -> None:
|
||||
self._limiter.release_on_behalf_of(borrower)
|
||||
|
||||
def statistics(self) -> CapacityLimiterStatistics:
|
||||
if self._internal_limiter is None:
|
||||
return CapacityLimiterStatistics(
|
||||
borrowed_tokens=0,
|
||||
total_tokens=self.total_tokens,
|
||||
borrowers=(),
|
||||
tasks_waiting=0,
|
||||
)
|
||||
|
||||
return self._internal_limiter.statistics()
|
||||
|
||||
|
||||
class ResourceGuard:
|
||||
"""
|
||||
A context manager for ensuring that a resource is only used by a single task at a
|
||||
time.
|
||||
|
||||
Entering this context manager while the previous has not exited it yet will trigger
|
||||
:exc:`BusyResourceError`.
|
||||
|
||||
:param action: the action to guard against (visible in the :exc:`BusyResourceError`
|
||||
when triggered, e.g. "Another task is already {action} this resource")
|
||||
|
||||
.. versionadded:: 4.1
|
||||
"""
|
||||
|
||||
__slots__ = "__weakref__", "_guarded", "action"
|
||||
|
||||
def __init__(self, action: str = "using"):
|
||||
self.action: str = action
|
||||
self._guarded = False
|
||||
|
||||
def __enter__(self) -> None:
|
||||
if self._guarded:
|
||||
raise BusyResourceError(self.action)
|
||||
|
||||
self._guarded = True
|
||||
|
||||
def __exit__(
|
||||
self,
|
||||
exc_type: type[BaseException] | None,
|
||||
exc_val: BaseException | None,
|
||||
exc_tb: TracebackType | None,
|
||||
) -> None:
|
||||
self._guarded = False
|
||||
+464
@@ -0,0 +1,464 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
import sys
|
||||
from collections.abc import (
|
||||
Coroutine,
|
||||
Generator,
|
||||
)
|
||||
from contextlib import contextmanager
|
||||
from enum import Enum, auto
|
||||
from inspect import iscoroutine
|
||||
from types import TracebackType
|
||||
from typing import Any, Generic, final
|
||||
|
||||
from ..abc import TaskGroup, TaskStatus
|
||||
from ._eventloop import get_async_backend, get_cancelled_exc_class
|
||||
from ._exceptions import TaskCancelled, TaskFailed, TaskNotFinished
|
||||
|
||||
if sys.version_info >= (3, 13):
|
||||
from typing import TypeVar
|
||||
else:
|
||||
from typing_extensions import TypeVar
|
||||
|
||||
if sys.version_info >= (3, 11):
|
||||
from typing import Never, Self, TypeVarTuple
|
||||
else:
|
||||
from typing_extensions import Never, Self, TypeVarTuple
|
||||
|
||||
T = TypeVar("T")
|
||||
T_co = TypeVar("T_co", covariant=True)
|
||||
T_startval_co = TypeVar("T_startval_co", covariant=True, default=Never)
|
||||
PosArgsT = TypeVarTuple("PosArgsT")
|
||||
|
||||
|
||||
class _IgnoredTaskStatus(TaskStatus[object]):
|
||||
def started(self, value: object = None) -> None:
|
||||
pass
|
||||
|
||||
|
||||
TASK_STATUS_IGNORED = _IgnoredTaskStatus()
|
||||
|
||||
|
||||
class CancelScope:
|
||||
"""
|
||||
Wraps a unit of work that can be made separately cancellable.
|
||||
|
||||
:param deadline: The time (clock value) when this scope is cancelled automatically
|
||||
:param shield: ``True`` to shield the cancel scope from external cancellation
|
||||
:raises NoEventLoopError: if no supported asynchronous event loop is running in the
|
||||
current thread
|
||||
"""
|
||||
|
||||
__slots__ = ("__weakref__",)
|
||||
|
||||
def __new__(
|
||||
cls, *, deadline: float = math.inf, shield: bool = False
|
||||
) -> CancelScope:
|
||||
return get_async_backend().create_cancel_scope(shield=shield, deadline=deadline)
|
||||
|
||||
def cancel(self, reason: str | None = None) -> None:
|
||||
"""
|
||||
Cancel this scope immediately.
|
||||
|
||||
:param reason: a message describing the reason for the cancellation
|
||||
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
@property
|
||||
def deadline(self) -> float:
|
||||
"""
|
||||
The time (clock value) when this scope is cancelled automatically.
|
||||
|
||||
Will be ``float('inf')`` if no timeout has been set.
|
||||
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
@deadline.setter
|
||||
def deadline(self, value: float) -> None:
|
||||
raise NotImplementedError
|
||||
|
||||
@property
|
||||
def cancel_called(self) -> bool:
|
||||
"""``True`` if :meth:`cancel` has been called."""
|
||||
raise NotImplementedError
|
||||
|
||||
@property
|
||||
def cancelled_caught(self) -> bool:
|
||||
"""
|
||||
``True`` if this scope suppressed a cancellation exception it itself raised.
|
||||
|
||||
This is typically used to check if any work was interrupted, or to see if the
|
||||
scope was cancelled due to its deadline being reached. The value will, however,
|
||||
only be ``True`` if the cancellation was triggered by the scope itself (and not
|
||||
an outer scope).
|
||||
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
@property
|
||||
def shield(self) -> bool:
|
||||
"""
|
||||
``True`` if this scope is shielded from external cancellation.
|
||||
|
||||
While a scope is shielded, it will not receive cancellations from outside.
|
||||
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
@shield.setter
|
||||
def shield(self, value: bool) -> None:
|
||||
raise NotImplementedError
|
||||
|
||||
def __enter__(self) -> Self:
|
||||
raise NotImplementedError
|
||||
|
||||
def __exit__(
|
||||
self,
|
||||
exc_type: type[BaseException] | None,
|
||||
exc_val: BaseException | None,
|
||||
exc_tb: TracebackType | None,
|
||||
) -> bool:
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
@contextmanager
|
||||
def fail_at(
|
||||
deadline: float | None, shield: bool = False, reason: str | None = None
|
||||
) -> Generator[CancelScope, None, None]:
|
||||
"""
|
||||
Create a context manager which raises a :class:`TimeoutError` if the code in the
|
||||
enclosing context does not finish by the given deadline.
|
||||
|
||||
:param deadline: the deadline before raising the exception, or
|
||||
``None`` to disable the timeout
|
||||
:param shield: ``True`` to shield the cancel scope from external cancellation
|
||||
:param reason: explanation for timeout to add to the message of a raised `TimeoutError`
|
||||
:return: a context manager that yields a cancel scope
|
||||
:rtype: :class:`~typing.ContextManager`\\[:class:`~anyio.CancelScope`\\]
|
||||
:raises NoEventLoopError: if no supported asynchronous event loop is running in the
|
||||
current thread
|
||||
|
||||
.. versionadded:: 4.15.0
|
||||
|
||||
"""
|
||||
current_time = get_async_backend().current_time
|
||||
effective_deadline = math.inf if deadline is None else deadline
|
||||
with get_async_backend().create_cancel_scope(
|
||||
deadline=effective_deadline, shield=shield
|
||||
) as cancel_scope:
|
||||
yield cancel_scope
|
||||
|
||||
if cancel_scope.cancelled_caught and current_time() >= cancel_scope.deadline:
|
||||
raise TimeoutError(reason) if reason else TimeoutError
|
||||
|
||||
|
||||
@contextmanager
|
||||
def fail_after(
|
||||
delay: float | None, shield: bool = False, reason: str | None = None
|
||||
) -> Generator[CancelScope, None, None]:
|
||||
"""
|
||||
Create a context manager which raises a :class:`TimeoutError` if the code in the
|
||||
enclosing context block does not finish in time.
|
||||
|
||||
:param delay: maximum allowed time (in seconds) before raising the exception, or
|
||||
``None`` to disable the timeout
|
||||
:param shield: ``True`` to shield the cancel scope from external cancellation
|
||||
:param reason: explanation for timeout to add to the message of a raised `TimeoutError`
|
||||
:return: a context manager that yields a cancel scope
|
||||
:rtype: :class:`~typing.ContextManager`\\[:class:`~anyio.CancelScope`\\]
|
||||
:raises NoEventLoopError: if no supported asynchronous event loop is running in the
|
||||
current thread
|
||||
|
||||
"""
|
||||
current_time = get_async_backend().current_time
|
||||
deadline = (current_time() + delay) if delay is not None else math.inf
|
||||
with fail_at(deadline, shield=shield, reason=reason) as scope:
|
||||
yield scope
|
||||
|
||||
|
||||
def move_on_at(deadline: float | None, shield: bool = False) -> CancelScope:
|
||||
"""
|
||||
Create a cancel scope with a deadline that expires after the given delay.
|
||||
|
||||
:param deadline: the deadline before exiting the context block, or
|
||||
``None`` to disable the timeout
|
||||
:param shield: ``True`` to shield the cancel scope from external cancellation
|
||||
:return: a cancel scope
|
||||
:raises NoEventLoopError: if no supported asynchronous event loop is running in the
|
||||
current thread
|
||||
|
||||
.. versionadded:: 4.15.0
|
||||
|
||||
"""
|
||||
return get_async_backend().create_cancel_scope(
|
||||
deadline=deadline if deadline is not None else math.inf, shield=shield
|
||||
)
|
||||
|
||||
|
||||
def move_on_after(delay: float | None, shield: bool = False) -> CancelScope:
|
||||
"""
|
||||
Create a cancel scope with a deadline that expires after the given delay.
|
||||
|
||||
:param delay: maximum allowed time (in seconds) before exiting the context block, or
|
||||
``None`` to disable the timeout
|
||||
:param shield: ``True`` to shield the cancel scope from external cancellation
|
||||
:return: a cancel scope
|
||||
:raises NoEventLoopError: if no supported asynchronous event loop is running in the
|
||||
current thread
|
||||
|
||||
.. note:: Unlike with :func:`fail_after`, the timer starts when this function is
|
||||
called, not when the context manager is entered. This will be changed in v5.0.
|
||||
|
||||
"""
|
||||
deadline = (
|
||||
(get_async_backend().current_time() + delay) if delay is not None else math.inf
|
||||
)
|
||||
return get_async_backend().create_cancel_scope(deadline=deadline, shield=shield)
|
||||
|
||||
|
||||
def current_effective_deadline() -> float:
|
||||
"""
|
||||
Return the nearest deadline among all the cancel scopes effective for the current
|
||||
task.
|
||||
|
||||
:return: a clock value from the event loop's internal clock (or ``float('inf')`` if
|
||||
there is no deadline in effect, or ``float('-inf')`` if the current scope has
|
||||
been cancelled)
|
||||
:rtype: float
|
||||
:raises NoEventLoopError: if no supported asynchronous event loop is running in the
|
||||
current thread
|
||||
|
||||
"""
|
||||
return get_async_backend().current_effective_deadline()
|
||||
|
||||
|
||||
def create_task_group() -> TaskGroup:
|
||||
"""
|
||||
Create a task group.
|
||||
|
||||
:return: a task group
|
||||
:raises NoEventLoopError: if no supported asynchronous event loop is running in the
|
||||
current thread
|
||||
|
||||
"""
|
||||
return get_async_backend().create_task_group()
|
||||
|
||||
|
||||
@final
|
||||
class TaskHandle(Generic[T_co, T_startval_co]):
|
||||
"""
|
||||
Returned from the task-spawning methods of :class:`TaskGroup`. Can be awaited on to
|
||||
get the return value of the task (or the raised exception). If the task was
|
||||
terminated by a :exc:`BaseException`, :exc:`TaskFailed` will be raised (or its
|
||||
subclass :exc:`TaskCancelled` if the task was cancelled).
|
||||
|
||||
.. versionadded:: 4.14.0
|
||||
"""
|
||||
|
||||
class Status(Enum):
|
||||
"""
|
||||
The status of a task handle.
|
||||
|
||||
.. attribute:: PENDING
|
||||
|
||||
The task has not finished yet.
|
||||
.. attribute:: FINISHED
|
||||
|
||||
The task has finished with a return value.
|
||||
.. attribute:: CANCELLING
|
||||
|
||||
The task has been cancelled but has not finished yet.
|
||||
.. attribute:: CANCELLED
|
||||
|
||||
The task was cancelled and has finished since.
|
||||
.. attribute:: FAILED
|
||||
|
||||
The task raised an exception.
|
||||
"""
|
||||
|
||||
PENDING = auto()
|
||||
FINISHED = auto()
|
||||
CANCELLING = auto()
|
||||
CANCELLED = auto()
|
||||
FAILED = auto()
|
||||
|
||||
__slots__ = (
|
||||
"__weakref__",
|
||||
"_cancel_scope",
|
||||
"_coro",
|
||||
"_exception",
|
||||
"_finished_event",
|
||||
"_name",
|
||||
"_return_value",
|
||||
"_start_value",
|
||||
)
|
||||
|
||||
_return_value: T_co
|
||||
_start_value: T_startval_co
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
coro: Coroutine[Any, Any, T_co],
|
||||
name: object,
|
||||
cancel_scope: CancelScope | None = None,
|
||||
) -> None:
|
||||
from ._synchronization import Event
|
||||
|
||||
self._coro = coro
|
||||
self._cancel_scope = cancel_scope if cancel_scope is not None else CancelScope()
|
||||
self._finished_event = Event()
|
||||
self._exception: BaseException | None = None
|
||||
|
||||
if name is not None:
|
||||
self._name = str(name)
|
||||
elif iscoroutine(coro):
|
||||
self._name = coro.__qualname__
|
||||
else:
|
||||
self._name = str(coro) # coroutine-like object (e.g. asend() objects)
|
||||
|
||||
async def _run_coro(self) -> None:
|
||||
__tracebackhide__ = True
|
||||
|
||||
with self._cancel_scope:
|
||||
try:
|
||||
retval = await self._coro
|
||||
except BaseException as exc:
|
||||
self._exception = exc
|
||||
raise
|
||||
else:
|
||||
self._return_value = retval
|
||||
finally:
|
||||
self._finished_event.set()
|
||||
del self # Break the reference cycle
|
||||
|
||||
def cancel(self) -> None:
|
||||
"""
|
||||
Set the task to a cancelled state.
|
||||
|
||||
This will interrupt any interruptible asynchronous operation, and will cause
|
||||
any further awaits on this task to get immediately cancelled, unless done in
|
||||
a shielded cancel scope.
|
||||
|
||||
If the task has already finished, this method has no effect.
|
||||
"""
|
||||
if not self._finished_event.is_set():
|
||||
self._cancel_scope.cancel()
|
||||
|
||||
@property
|
||||
def coro(self) -> Coroutine[Any, Any, T_co]:
|
||||
"""
|
||||
The coroutine object that was passed to one of the task-spawning methods in
|
||||
:class:`TaskGroup`.
|
||||
"""
|
||||
return self._coro
|
||||
|
||||
@property
|
||||
def status(self) -> TaskHandle.Status:
|
||||
"""
|
||||
The current status of the task.
|
||||
|
||||
Every task starts in the :attr:`~TaskHandle.Status.PENDING` state.
|
||||
If a task is cancelled while in this state, it will transition to the
|
||||
:attr:`~TaskHandle.Status.CANCELLING` state. When the task finishes, it will
|
||||
transition to one of the three final states (
|
||||
:attr:`~TaskHandle.Status.FINISHED`, :attr:`~TaskHandle.Status.FAILED`, or
|
||||
:attr:`~TaskHandle.Status.CANCELLING`) depending on the exception the task
|
||||
raised, if any. No other status transitions will happen.
|
||||
"""
|
||||
if not self._finished_event.is_set():
|
||||
if self._cancel_scope.cancel_called:
|
||||
return TaskHandle.Status.CANCELLING
|
||||
else:
|
||||
return TaskHandle.Status.PENDING
|
||||
elif self._exception is not None:
|
||||
if isinstance(self._exception, get_cancelled_exc_class()):
|
||||
return TaskHandle.Status.CANCELLED
|
||||
else:
|
||||
return TaskHandle.Status.FAILED
|
||||
else:
|
||||
return TaskHandle.Status.FINISHED
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
"""The name of the task."""
|
||||
return self._name
|
||||
|
||||
@property
|
||||
def exception(self) -> BaseException | None:
|
||||
"""
|
||||
The exception raised by the task, or ``None`` if it finished without raising.
|
||||
|
||||
:raises TaskNotFinished: if the task has not finished yet
|
||||
:raises TaskCancelled: if the task was cancelled
|
||||
|
||||
"""
|
||||
match self.status:
|
||||
case TaskHandle.Status.PENDING:
|
||||
raise TaskNotFinished("the task has not finished yet")
|
||||
case TaskHandle.Status.FINISHED:
|
||||
return None
|
||||
case TaskHandle.Status.CANCELLING:
|
||||
raise TaskCancelled("the task was cancelled")
|
||||
case TaskHandle.Status.CANCELLED:
|
||||
raise TaskCancelled("the task was cancelled") from self._exception
|
||||
case TaskHandle.Status.FAILED:
|
||||
return self._exception
|
||||
|
||||
@property
|
||||
def return_value(self) -> T_co:
|
||||
"""
|
||||
The return value of the task.
|
||||
|
||||
:raises TaskNotFinished: if the task has not finished yet
|
||||
:raises TaskCancelled: if the task was cancelled
|
||||
:raises TaskFailed: if the task raised an exception
|
||||
|
||||
"""
|
||||
match self.status:
|
||||
case TaskHandle.Status.PENDING:
|
||||
raise TaskNotFinished("the task has not finished yet")
|
||||
case TaskHandle.Status.FINISHED:
|
||||
return self._return_value
|
||||
case TaskHandle.Status.CANCELLING:
|
||||
raise TaskCancelled("the task was cancelled")
|
||||
case TaskHandle.Status.CANCELLED:
|
||||
raise TaskCancelled("the task was cancelled") from self._exception
|
||||
case TaskHandle.Status.FAILED:
|
||||
raise TaskFailed("the task raised an exception") from self._exception
|
||||
|
||||
@property
|
||||
def start_value(self) -> T_startval_co:
|
||||
"""
|
||||
The value passed to :meth:`task_status.started() <.abc.TaskStatus.started>`,
|
||||
|
||||
:raises RuntimeError: if the task was not started with :meth:`TaskGroup.start()
|
||||
<.abc.TaskGroup.start>`
|
||||
"""
|
||||
try:
|
||||
return self._start_value
|
||||
except AttributeError:
|
||||
raise RuntimeError(
|
||||
"the task was not started with TaskGroup.start()"
|
||||
) from None
|
||||
|
||||
async def wait(self) -> None:
|
||||
"""
|
||||
Wait for the task to finish.
|
||||
|
||||
This method will return as soon as the task has finished, no matter how it
|
||||
happened.
|
||||
"""
|
||||
await self._finished_event.wait()
|
||||
|
||||
def __await__(self) -> Generator[Any, Any, T_co]:
|
||||
yield from self._finished_event.wait().__await__()
|
||||
return self.return_value
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return (
|
||||
f"<{self.__class__.__name__} {self.status.name.lower()} "
|
||||
f"name={self._name!r} coro={self._coro!r}>"
|
||||
)
|
||||
+615
@@ -0,0 +1,615 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import sys
|
||||
import tempfile
|
||||
from collections.abc import Iterable
|
||||
from io import BytesIO, TextIOWrapper
|
||||
from types import TracebackType
|
||||
from typing import (
|
||||
TYPE_CHECKING,
|
||||
Any,
|
||||
AnyStr,
|
||||
Generic,
|
||||
overload,
|
||||
)
|
||||
|
||||
from .. import to_thread
|
||||
from .._core._fileio import AsyncFile
|
||||
from ..lowlevel import checkpoint_if_cancelled
|
||||
from ._tasks import CancelScope
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from _typeshed import OpenBinaryMode, OpenTextMode, ReadableBuffer, WriteableBuffer
|
||||
|
||||
|
||||
class TemporaryFile(Generic[AnyStr]):
|
||||
"""
|
||||
An asynchronous temporary file that is automatically created and cleaned up.
|
||||
|
||||
This class provides an asynchronous context manager interface to a temporary file.
|
||||
The file is created using Python's standard `tempfile.TemporaryFile` function in a
|
||||
background thread, and is wrapped as an asynchronous file using `AsyncFile`.
|
||||
|
||||
:param mode: The mode in which the file is opened. Defaults to "w+b".
|
||||
:param buffering: The buffering policy (-1 means the default buffering).
|
||||
:param encoding: The encoding used to decode or encode the file. Only applicable in
|
||||
text mode.
|
||||
:param newline: Controls how universal newlines mode works (only applicable in text
|
||||
mode).
|
||||
:param suffix: The suffix for the temporary file name.
|
||||
:param prefix: The prefix for the temporary file name.
|
||||
:param dir: The directory in which the temporary file is created.
|
||||
:param errors: The error handling scheme used for encoding/decoding errors.
|
||||
"""
|
||||
|
||||
_async_file: AsyncFile[AnyStr]
|
||||
|
||||
@overload
|
||||
def __init__(
|
||||
self: TemporaryFile[bytes],
|
||||
mode: OpenBinaryMode = ...,
|
||||
buffering: int = ...,
|
||||
encoding: str | None = ...,
|
||||
newline: str | None = ...,
|
||||
suffix: str | None = ...,
|
||||
prefix: str | None = ...,
|
||||
dir: str | None = ...,
|
||||
*,
|
||||
errors: str | None = ...,
|
||||
): ...
|
||||
@overload
|
||||
def __init__(
|
||||
self: TemporaryFile[str],
|
||||
mode: OpenTextMode,
|
||||
buffering: int = ...,
|
||||
encoding: str | None = ...,
|
||||
newline: str | None = ...,
|
||||
suffix: str | None = ...,
|
||||
prefix: str | None = ...,
|
||||
dir: str | None = ...,
|
||||
*,
|
||||
errors: str | None = ...,
|
||||
): ...
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
mode: OpenTextMode | OpenBinaryMode = "w+b",
|
||||
buffering: int = -1,
|
||||
encoding: str | None = None,
|
||||
newline: str | None = None,
|
||||
suffix: str | None = None,
|
||||
prefix: str | None = None,
|
||||
dir: str | None = None,
|
||||
*,
|
||||
errors: str | None = None,
|
||||
) -> None:
|
||||
self.mode = mode
|
||||
self.buffering = buffering
|
||||
self.encoding = encoding
|
||||
self.newline = newline
|
||||
self.suffix: str | None = suffix
|
||||
self.prefix: str | None = prefix
|
||||
self.dir: str | None = dir
|
||||
self.errors = errors
|
||||
|
||||
async def __aenter__(self) -> AsyncFile[AnyStr]:
|
||||
fp = await to_thread.run_sync(
|
||||
lambda: tempfile.TemporaryFile(
|
||||
self.mode,
|
||||
self.buffering,
|
||||
self.encoding,
|
||||
self.newline,
|
||||
self.suffix,
|
||||
self.prefix,
|
||||
self.dir,
|
||||
errors=self.errors,
|
||||
)
|
||||
)
|
||||
self._async_file = AsyncFile(fp)
|
||||
return self._async_file
|
||||
|
||||
async def __aexit__(
|
||||
self,
|
||||
exc_type: type[BaseException] | None,
|
||||
exc_value: BaseException | None,
|
||||
traceback: TracebackType | None,
|
||||
) -> None:
|
||||
await self._async_file.aclose()
|
||||
|
||||
|
||||
class NamedTemporaryFile(Generic[AnyStr]):
|
||||
"""
|
||||
An asynchronous named temporary file that is automatically created and cleaned up.
|
||||
|
||||
This class provides an asynchronous context manager for a temporary file with a
|
||||
visible name in the file system. It uses Python's standard
|
||||
:func:`~tempfile.NamedTemporaryFile` function and wraps the file object with
|
||||
:class:`AsyncFile` for asynchronous operations.
|
||||
|
||||
:param mode: The mode in which the file is opened. Defaults to "w+b".
|
||||
:param buffering: The buffering policy (-1 means the default buffering).
|
||||
:param encoding: The encoding used to decode or encode the file. Only applicable in
|
||||
text mode.
|
||||
:param newline: Controls how universal newlines mode works (only applicable in text
|
||||
mode).
|
||||
:param suffix: The suffix for the temporary file name.
|
||||
:param prefix: The prefix for the temporary file name.
|
||||
:param dir: The directory in which the temporary file is created.
|
||||
:param delete: Whether to delete the file when it is closed.
|
||||
:param errors: The error handling scheme used for encoding/decoding errors.
|
||||
:param delete_on_close: (Python 3.12+) Whether to delete the file on close.
|
||||
"""
|
||||
|
||||
_async_file: AsyncFile[AnyStr]
|
||||
|
||||
@overload
|
||||
def __init__(
|
||||
self: NamedTemporaryFile[bytes],
|
||||
mode: OpenBinaryMode = ...,
|
||||
buffering: int = ...,
|
||||
encoding: str | None = ...,
|
||||
newline: str | None = ...,
|
||||
suffix: str | None = ...,
|
||||
prefix: str | None = ...,
|
||||
dir: str | None = ...,
|
||||
delete: bool = ...,
|
||||
*,
|
||||
errors: str | None = ...,
|
||||
delete_on_close: bool = ...,
|
||||
): ...
|
||||
@overload
|
||||
def __init__(
|
||||
self: NamedTemporaryFile[str],
|
||||
mode: OpenTextMode,
|
||||
buffering: int = ...,
|
||||
encoding: str | None = ...,
|
||||
newline: str | None = ...,
|
||||
suffix: str | None = ...,
|
||||
prefix: str | None = ...,
|
||||
dir: str | None = ...,
|
||||
delete: bool = ...,
|
||||
*,
|
||||
errors: str | None = ...,
|
||||
delete_on_close: bool = ...,
|
||||
): ...
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
mode: OpenBinaryMode | OpenTextMode = "w+b",
|
||||
buffering: int = -1,
|
||||
encoding: str | None = None,
|
||||
newline: str | None = None,
|
||||
suffix: str | None = None,
|
||||
prefix: str | None = None,
|
||||
dir: str | None = None,
|
||||
delete: bool = True,
|
||||
*,
|
||||
errors: str | None = None,
|
||||
delete_on_close: bool = True,
|
||||
) -> None:
|
||||
self._params: dict[str, Any] = {
|
||||
"mode": mode,
|
||||
"buffering": buffering,
|
||||
"encoding": encoding,
|
||||
"newline": newline,
|
||||
"suffix": suffix,
|
||||
"prefix": prefix,
|
||||
"dir": dir,
|
||||
"delete": delete,
|
||||
"errors": errors,
|
||||
}
|
||||
if sys.version_info >= (3, 12):
|
||||
self._params["delete_on_close"] = delete_on_close
|
||||
|
||||
async def __aenter__(self) -> AsyncFile[AnyStr]:
|
||||
fp = await to_thread.run_sync(
|
||||
lambda: tempfile.NamedTemporaryFile(**self._params)
|
||||
)
|
||||
self._async_file = AsyncFile(fp)
|
||||
return self._async_file
|
||||
|
||||
async def __aexit__(
|
||||
self,
|
||||
exc_type: type[BaseException] | None,
|
||||
exc_value: BaseException | None,
|
||||
traceback: TracebackType | None,
|
||||
) -> None:
|
||||
await self._async_file.aclose()
|
||||
|
||||
|
||||
class SpooledTemporaryFile(AsyncFile[AnyStr]):
|
||||
"""
|
||||
An asynchronous spooled temporary file that starts in memory and is spooled to disk.
|
||||
|
||||
This class provides an asynchronous interface to a spooled temporary file, much like
|
||||
Python's standard :class:`~tempfile.SpooledTemporaryFile`. It supports asynchronous
|
||||
write operations and provides a method to force a rollover to disk.
|
||||
|
||||
:param max_size: Maximum size in bytes before the file is rolled over to disk.
|
||||
:param mode: The mode in which the file is opened. Defaults to "w+b".
|
||||
:param buffering: The buffering policy (-1 means the default buffering).
|
||||
:param encoding: The encoding used to decode or encode the file (text mode only).
|
||||
:param newline: Controls how universal newlines mode works (text mode only).
|
||||
:param suffix: The suffix for the temporary file name.
|
||||
:param prefix: The prefix for the temporary file name.
|
||||
:param dir: The directory in which the temporary file is created.
|
||||
:param errors: The error handling scheme used for encoding/decoding errors.
|
||||
"""
|
||||
|
||||
_rolled: bool = False
|
||||
|
||||
@overload
|
||||
def __init__(
|
||||
self: SpooledTemporaryFile[bytes],
|
||||
max_size: int = ...,
|
||||
mode: OpenBinaryMode = ...,
|
||||
buffering: int = ...,
|
||||
encoding: str | None = ...,
|
||||
newline: str | None = ...,
|
||||
suffix: str | None = ...,
|
||||
prefix: str | None = ...,
|
||||
dir: str | None = ...,
|
||||
*,
|
||||
errors: str | None = ...,
|
||||
): ...
|
||||
@overload
|
||||
def __init__(
|
||||
self: SpooledTemporaryFile[str],
|
||||
max_size: int = ...,
|
||||
mode: OpenTextMode = ...,
|
||||
buffering: int = ...,
|
||||
encoding: str | None = ...,
|
||||
newline: str | None = ...,
|
||||
suffix: str | None = ...,
|
||||
prefix: str | None = ...,
|
||||
dir: str | None = ...,
|
||||
*,
|
||||
errors: str | None = ...,
|
||||
): ...
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
max_size: int = 0,
|
||||
mode: OpenBinaryMode | OpenTextMode = "w+b",
|
||||
buffering: int = -1,
|
||||
encoding: str | None = None,
|
||||
newline: str | None = None,
|
||||
suffix: str | None = None,
|
||||
prefix: str | None = None,
|
||||
dir: str | None = None,
|
||||
*,
|
||||
errors: str | None = None,
|
||||
) -> None:
|
||||
self._tempfile_params: dict[str, Any] = {
|
||||
"mode": mode,
|
||||
"buffering": buffering,
|
||||
"encoding": encoding,
|
||||
"newline": newline,
|
||||
"suffix": suffix,
|
||||
"prefix": prefix,
|
||||
"dir": dir,
|
||||
"errors": errors,
|
||||
}
|
||||
self._max_size = max_size
|
||||
if "b" in mode:
|
||||
super().__init__(BytesIO()) # type: ignore[arg-type]
|
||||
else:
|
||||
super().__init__(
|
||||
TextIOWrapper( # type: ignore[arg-type]
|
||||
BytesIO(),
|
||||
encoding=encoding,
|
||||
errors=errors,
|
||||
newline=newline,
|
||||
write_through=True,
|
||||
)
|
||||
)
|
||||
|
||||
async def aclose(self) -> None:
|
||||
if not self._rolled:
|
||||
self._fp.close()
|
||||
return
|
||||
|
||||
await super().aclose()
|
||||
|
||||
async def _check(self) -> None:
|
||||
if self._rolled or self._fp.tell() <= self._max_size:
|
||||
return
|
||||
|
||||
await self.rollover()
|
||||
|
||||
async def rollover(self) -> None:
|
||||
if self._rolled:
|
||||
return
|
||||
|
||||
self._rolled = True
|
||||
buffer = self._fp
|
||||
buffer.seek(0)
|
||||
self._fp = await to_thread.run_sync(
|
||||
lambda: tempfile.TemporaryFile(**self._tempfile_params)
|
||||
)
|
||||
await self.write(buffer.read())
|
||||
buffer.close()
|
||||
|
||||
@property
|
||||
def closed(self) -> bool:
|
||||
return self._fp.closed
|
||||
|
||||
async def read(self, size: int = -1) -> AnyStr:
|
||||
if not self._rolled:
|
||||
await checkpoint_if_cancelled()
|
||||
return self._fp.read(size)
|
||||
|
||||
return await super().read(size) # type: ignore[return-value]
|
||||
|
||||
async def read1(self: SpooledTemporaryFile[bytes], size: int = -1) -> bytes:
|
||||
if not self._rolled:
|
||||
await checkpoint_if_cancelled()
|
||||
return self._fp.read1(size)
|
||||
|
||||
return await super().read1(size)
|
||||
|
||||
async def readline(self) -> AnyStr:
|
||||
if not self._rolled:
|
||||
await checkpoint_if_cancelled()
|
||||
return self._fp.readline()
|
||||
|
||||
return await super().readline() # type: ignore[return-value]
|
||||
|
||||
async def readlines(self) -> list[AnyStr]:
|
||||
if not self._rolled:
|
||||
await checkpoint_if_cancelled()
|
||||
return self._fp.readlines()
|
||||
|
||||
return await super().readlines() # type: ignore[return-value]
|
||||
|
||||
async def readinto(self: SpooledTemporaryFile[bytes], b: WriteableBuffer) -> int:
|
||||
if not self._rolled:
|
||||
await checkpoint_if_cancelled()
|
||||
return self._fp.readinto(b)
|
||||
|
||||
return await super().readinto(b)
|
||||
|
||||
async def readinto1(self: SpooledTemporaryFile[bytes], b: WriteableBuffer) -> int:
|
||||
if not self._rolled:
|
||||
await checkpoint_if_cancelled()
|
||||
return self._fp.readinto1(b)
|
||||
|
||||
return await super().readinto1(b)
|
||||
|
||||
async def seek(self, offset: int, whence: int | None = os.SEEK_SET) -> int:
|
||||
if not self._rolled:
|
||||
await checkpoint_if_cancelled()
|
||||
return self._fp.seek(offset, whence)
|
||||
|
||||
return await super().seek(offset, whence)
|
||||
|
||||
async def tell(self) -> int:
|
||||
if not self._rolled:
|
||||
await checkpoint_if_cancelled()
|
||||
return self._fp.tell()
|
||||
|
||||
return await super().tell()
|
||||
|
||||
async def truncate(self, size: int | None = None) -> int:
|
||||
if not self._rolled:
|
||||
await checkpoint_if_cancelled()
|
||||
return self._fp.truncate(size)
|
||||
|
||||
return await super().truncate(size)
|
||||
|
||||
@overload
|
||||
async def write(self: SpooledTemporaryFile[bytes], b: ReadableBuffer) -> int: ...
|
||||
@overload
|
||||
async def write(self: SpooledTemporaryFile[str], b: str) -> int: ...
|
||||
|
||||
async def write(self, b: ReadableBuffer | str) -> int:
|
||||
"""
|
||||
Asynchronously write data to the spooled temporary file.
|
||||
|
||||
If the file has not yet been rolled over, the data is written synchronously,
|
||||
and a rollover is triggered if the size exceeds the maximum size.
|
||||
|
||||
:param s: The data to write.
|
||||
:return: The number of bytes written.
|
||||
:raises RuntimeError: If the underlying file is not initialized.
|
||||
|
||||
"""
|
||||
if not self._rolled:
|
||||
await checkpoint_if_cancelled()
|
||||
result = self._fp.write(b)
|
||||
await self._check()
|
||||
return result
|
||||
|
||||
return await super().write(b) # type: ignore[misc]
|
||||
|
||||
@overload
|
||||
async def writelines(
|
||||
self: SpooledTemporaryFile[bytes], lines: Iterable[ReadableBuffer]
|
||||
) -> None: ...
|
||||
@overload
|
||||
async def writelines(
|
||||
self: SpooledTemporaryFile[str], lines: Iterable[str]
|
||||
) -> None: ...
|
||||
|
||||
async def writelines(self, lines: Iterable[str] | Iterable[ReadableBuffer]) -> None:
|
||||
"""
|
||||
Asynchronously write a list of lines to the spooled temporary file.
|
||||
|
||||
If the file has not yet been rolled over, the lines are written synchronously,
|
||||
and a rollover is triggered if the size exceeds the maximum size.
|
||||
|
||||
:param lines: An iterable of lines to write.
|
||||
:raises RuntimeError: If the underlying file is not initialized.
|
||||
|
||||
"""
|
||||
if not self._rolled:
|
||||
await checkpoint_if_cancelled()
|
||||
result = self._fp.writelines(lines)
|
||||
await self._check()
|
||||
return result
|
||||
|
||||
return await super().writelines(lines) # type: ignore[misc]
|
||||
|
||||
|
||||
class TemporaryDirectory(Generic[AnyStr]):
|
||||
"""
|
||||
An asynchronous temporary directory that is created and cleaned up automatically.
|
||||
|
||||
This class provides an asynchronous context manager for creating a temporary
|
||||
directory. It wraps Python's standard :class:`~tempfile.TemporaryDirectory` to
|
||||
perform directory creation and cleanup operations in a background thread.
|
||||
|
||||
:param suffix: Suffix to be added to the temporary directory name.
|
||||
:param prefix: Prefix to be added to the temporary directory name.
|
||||
:param dir: The parent directory where the temporary directory is created.
|
||||
:param ignore_cleanup_errors: Whether to ignore errors during cleanup
|
||||
:param delete: Whether to delete the directory upon closing (Python 3.12+).
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
suffix: AnyStr | None = None,
|
||||
prefix: AnyStr | None = None,
|
||||
dir: AnyStr | None = None,
|
||||
*,
|
||||
ignore_cleanup_errors: bool = False,
|
||||
delete: bool = True,
|
||||
) -> None:
|
||||
self.suffix: AnyStr | None = suffix
|
||||
self.prefix: AnyStr | None = prefix
|
||||
self.dir: AnyStr | None = dir
|
||||
self.ignore_cleanup_errors = ignore_cleanup_errors
|
||||
self.delete = delete
|
||||
|
||||
self._tempdir: tempfile.TemporaryDirectory | None = None
|
||||
|
||||
async def __aenter__(self) -> str:
|
||||
params: dict[str, Any] = {
|
||||
"suffix": self.suffix,
|
||||
"prefix": self.prefix,
|
||||
"dir": self.dir,
|
||||
"ignore_cleanup_errors": self.ignore_cleanup_errors,
|
||||
}
|
||||
if sys.version_info >= (3, 12):
|
||||
params["delete"] = self.delete
|
||||
|
||||
self._tempdir = await to_thread.run_sync(
|
||||
lambda: tempfile.TemporaryDirectory(**params)
|
||||
)
|
||||
return await to_thread.run_sync(self._tempdir.__enter__)
|
||||
|
||||
async def __aexit__(
|
||||
self,
|
||||
exc_type: type[BaseException] | None,
|
||||
exc_value: BaseException | None,
|
||||
traceback: TracebackType | None,
|
||||
) -> None:
|
||||
if self._tempdir is not None:
|
||||
with CancelScope(shield=True):
|
||||
await to_thread.run_sync(
|
||||
self._tempdir.__exit__, exc_type, exc_value, traceback
|
||||
)
|
||||
|
||||
async def cleanup(self) -> None:
|
||||
if self._tempdir is not None:
|
||||
await to_thread.run_sync(self._tempdir.cleanup)
|
||||
|
||||
|
||||
@overload
|
||||
async def mkstemp(
|
||||
suffix: str | None = None,
|
||||
prefix: str | None = None,
|
||||
dir: str | None = None,
|
||||
text: bool = False,
|
||||
) -> tuple[int, str]: ...
|
||||
|
||||
|
||||
@overload
|
||||
async def mkstemp(
|
||||
suffix: bytes | None = None,
|
||||
prefix: bytes | None = None,
|
||||
dir: bytes | None = None,
|
||||
text: bool = False,
|
||||
) -> tuple[int, bytes]: ...
|
||||
|
||||
|
||||
async def mkstemp(
|
||||
suffix: AnyStr | None = None,
|
||||
prefix: AnyStr | None = None,
|
||||
dir: AnyStr | None = None,
|
||||
text: bool = False,
|
||||
) -> tuple[int, str | bytes]:
|
||||
"""
|
||||
Asynchronously create a temporary file and return an OS-level handle and the file
|
||||
name.
|
||||
|
||||
This function wraps `tempfile.mkstemp` and executes it in a background thread.
|
||||
|
||||
:param suffix: Suffix to be added to the file name.
|
||||
:param prefix: Prefix to be added to the file name.
|
||||
:param dir: Directory in which the temporary file is created.
|
||||
:param text: Whether the file is opened in text mode.
|
||||
:return: A tuple containing the file descriptor and the file name.
|
||||
|
||||
"""
|
||||
return await to_thread.run_sync(tempfile.mkstemp, suffix, prefix, dir, text)
|
||||
|
||||
|
||||
@overload
|
||||
async def mkdtemp(
|
||||
suffix: str | None = None,
|
||||
prefix: str | None = None,
|
||||
dir: str | None = None,
|
||||
) -> str: ...
|
||||
|
||||
|
||||
@overload
|
||||
async def mkdtemp(
|
||||
suffix: bytes | None = None,
|
||||
prefix: bytes | None = None,
|
||||
dir: bytes | None = None,
|
||||
) -> bytes: ...
|
||||
|
||||
|
||||
async def mkdtemp(
|
||||
suffix: AnyStr | None = None,
|
||||
prefix: AnyStr | None = None,
|
||||
dir: AnyStr | None = None,
|
||||
) -> str | bytes:
|
||||
"""
|
||||
Asynchronously create a temporary directory and return its path.
|
||||
|
||||
This function wraps `tempfile.mkdtemp` and executes it in a background thread.
|
||||
|
||||
:param suffix: Suffix to be added to the directory name.
|
||||
:param prefix: Prefix to be added to the directory name.
|
||||
:param dir: Parent directory where the temporary directory is created.
|
||||
:return: The path of the created temporary directory.
|
||||
|
||||
"""
|
||||
return await to_thread.run_sync(tempfile.mkdtemp, suffix, prefix, dir)
|
||||
|
||||
|
||||
async def gettempdir() -> str:
|
||||
"""
|
||||
Asynchronously return the name of the directory used for temporary files.
|
||||
|
||||
This function wraps `tempfile.gettempdir` and executes it in a background thread.
|
||||
|
||||
:return: The path of the temporary directory as a string.
|
||||
|
||||
"""
|
||||
return await to_thread.run_sync(tempfile.gettempdir)
|
||||
|
||||
|
||||
async def gettempdirb() -> bytes:
|
||||
"""
|
||||
Asynchronously return the name of the directory used for temporary files in bytes.
|
||||
|
||||
This function wraps `tempfile.gettempdirb` and executes it in a background thread.
|
||||
|
||||
:return: The path of the temporary directory as bytes.
|
||||
|
||||
"""
|
||||
return await to_thread.run_sync(tempfile.gettempdirb)
|
||||
+82
@@ -0,0 +1,82 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Awaitable, Generator
|
||||
from typing import Any, cast
|
||||
|
||||
from ._eventloop import get_async_backend
|
||||
|
||||
|
||||
class TaskInfo:
|
||||
"""
|
||||
Represents an asynchronous task.
|
||||
|
||||
:ivar int id: the unique identifier of the task
|
||||
:ivar parent_id: the identifier of the parent task, if any
|
||||
:vartype parent_id: Optional[int]
|
||||
:ivar str name: the description of the task (if any)
|
||||
:ivar ~collections.abc.Coroutine coro: the coroutine object of the task
|
||||
"""
|
||||
|
||||
__slots__ = "_name", "coro", "id", "name", "parent_id"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
id: int,
|
||||
parent_id: int | None,
|
||||
name: str | None,
|
||||
coro: Generator[Any, Any, Any] | Awaitable[Any],
|
||||
):
|
||||
func = get_current_task
|
||||
self._name = f"{func.__module__}.{func.__qualname__}"
|
||||
self.id: int = id
|
||||
self.parent_id: int | None = parent_id
|
||||
self.name: str | None = name
|
||||
self.coro: Generator[Any, Any, Any] | Awaitable[Any] = coro
|
||||
|
||||
def __eq__(self, other: object) -> bool:
|
||||
if isinstance(other, TaskInfo):
|
||||
return self.id == other.id
|
||||
|
||||
return NotImplemented
|
||||
|
||||
def __hash__(self) -> int:
|
||||
return hash(self.id)
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"{self.__class__.__name__}(id={self.id!r}, name={self.name!r})"
|
||||
|
||||
def has_pending_cancellation(self) -> bool:
|
||||
"""
|
||||
Return ``True`` if the task has a cancellation pending, ``False`` otherwise.
|
||||
|
||||
"""
|
||||
return False
|
||||
|
||||
|
||||
def get_current_task() -> TaskInfo:
|
||||
"""
|
||||
Return the current task.
|
||||
|
||||
:return: a representation of the current task
|
||||
:raises NoEventLoopError: if no supported asynchronous event loop is running in the
|
||||
current thread
|
||||
|
||||
"""
|
||||
return get_async_backend().get_current_task()
|
||||
|
||||
|
||||
def get_running_tasks() -> list[TaskInfo]:
|
||||
"""
|
||||
Return a list of running tasks in the current event loop.
|
||||
|
||||
:return: a list of task info objects
|
||||
:raises NoEventLoopError: if no supported asynchronous event loop is running in the
|
||||
current thread
|
||||
|
||||
"""
|
||||
return cast("list[TaskInfo]", get_async_backend().get_running_tasks())
|
||||
|
||||
|
||||
async def wait_all_tasks_blocked() -> None:
|
||||
"""Wait until all other tasks are waiting for something."""
|
||||
await get_async_backend().wait_all_tasks_blocked()
|
||||
+85
@@ -0,0 +1,85 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
from collections.abc import Callable, Mapping
|
||||
from typing import Any, TypeVar, final, overload
|
||||
|
||||
from ._exceptions import TypedAttributeLookupError
|
||||
|
||||
if sys.version_info < (3, 15):
|
||||
from typing_extensions import sentinel
|
||||
|
||||
T_Attr = TypeVar("T_Attr")
|
||||
T_Default = TypeVar("T_Default")
|
||||
undefined = sentinel("undefined")
|
||||
|
||||
|
||||
def typed_attribute() -> Any:
|
||||
"""Return a unique object, used to mark typed attributes."""
|
||||
return object()
|
||||
|
||||
|
||||
class TypedAttributeSet:
|
||||
"""
|
||||
Superclass for typed attribute collections.
|
||||
|
||||
Checks that every public attribute of every subclass has a type annotation.
|
||||
"""
|
||||
|
||||
def __init_subclass__(cls) -> None:
|
||||
annotations: dict[str, Any] = getattr(cls, "__annotations__", {})
|
||||
for attrname in dir(cls):
|
||||
if not attrname.startswith("_") and attrname not in annotations:
|
||||
raise TypeError(
|
||||
f"Attribute {attrname!r} is missing its type annotation"
|
||||
)
|
||||
|
||||
super().__init_subclass__()
|
||||
|
||||
|
||||
class TypedAttributeProvider:
|
||||
"""Base class for classes that wish to provide typed extra attributes."""
|
||||
|
||||
@property
|
||||
def extra_attributes(self) -> Mapping[T_Attr, Callable[[], T_Attr]]:
|
||||
"""
|
||||
A mapping of the extra attributes to callables that return the corresponding
|
||||
values.
|
||||
|
||||
If the provider wraps another provider, the attributes from that wrapper should
|
||||
also be included in the returned mapping (but the wrapper may override the
|
||||
callables from the wrapped instance).
|
||||
|
||||
"""
|
||||
return {}
|
||||
|
||||
@overload
|
||||
def extra(self, attribute: T_Attr) -> T_Attr: ...
|
||||
|
||||
@overload
|
||||
def extra(self, attribute: T_Attr, default: T_Default) -> T_Attr | T_Default: ...
|
||||
|
||||
@final
|
||||
def extra(self, attribute: Any, default: object = undefined) -> object:
|
||||
"""
|
||||
extra(attribute, default=undefined)
|
||||
|
||||
Return the value of the given typed extra attribute.
|
||||
|
||||
:param attribute: the attribute (member of a :class:`~TypedAttributeSet`) to
|
||||
look for
|
||||
:param default: the value that should be returned if no value is found for the
|
||||
attribute
|
||||
:raises ~anyio.TypedAttributeLookupError: if the search failed and no default
|
||||
value was given
|
||||
|
||||
"""
|
||||
try:
|
||||
getter = self.extra_attributes[attribute]
|
||||
except KeyError:
|
||||
if default is undefined:
|
||||
raise TypedAttributeLookupError("Attribute not found") from None
|
||||
else:
|
||||
return default
|
||||
|
||||
return getter()
|
||||
+176
@@ -0,0 +1,176 @@
|
||||
from __future__ import annotations
|
||||
|
||||
__all__ = (
|
||||
"fix_package_names",
|
||||
"install_lazy_importer",
|
||||
"set_deprecated_aliases",
|
||||
)
|
||||
|
||||
import ast
|
||||
import inspect
|
||||
import sys
|
||||
import warnings
|
||||
from importlib import import_module
|
||||
from types import ModuleType
|
||||
from typing import Any
|
||||
|
||||
|
||||
def install_lazy_importer() -> bool:
|
||||
module_globals = sys._getframe(1).f_globals
|
||||
module_name = module_globals["__name__"]
|
||||
module_prefix = module_name + "."
|
||||
module = sys.modules[module_name]
|
||||
lazy_map, deprecated_aliases, submodule_names = _build_lazy_map(module)
|
||||
names = sorted(lazy_map)
|
||||
|
||||
# Delete symbols that are not part of the API
|
||||
del module_globals["TYPE_CHECKING"]
|
||||
del module_globals["install_lazy_importer"]
|
||||
|
||||
if not lazy_map and not deprecated_aliases:
|
||||
return False
|
||||
|
||||
def __getattr__(name: str) -> Any:
|
||||
if new_name := deprecated_aliases.get(name):
|
||||
emit_deprecation_warning(module_name, name, new_name)
|
||||
target_mod, target_attr = new_name.rsplit(".", 1)
|
||||
elif name in submodule_names:
|
||||
target_mod, target_attr = "." + name, ""
|
||||
else:
|
||||
try:
|
||||
target_mod, target_attr = lazy_map[name]
|
||||
except KeyError:
|
||||
raise AttributeError(
|
||||
f"module {module_name!r} has no attribute {name!r}"
|
||||
) from None
|
||||
|
||||
imported = import_module(target_mod, module_name)
|
||||
value = getattr(imported, target_attr) if target_attr else imported
|
||||
|
||||
# patch the module name to match
|
||||
if (
|
||||
getattr(value, "__module__", "").startswith(module_prefix)
|
||||
and name not in deprecated_aliases
|
||||
):
|
||||
value.__module__ = module_name
|
||||
|
||||
module_globals[name] = value
|
||||
return value
|
||||
|
||||
def __dir__() -> list[str]:
|
||||
return names
|
||||
|
||||
module_globals["__dir__"] = __dir__
|
||||
module_globals["__getattr__"] = __getattr__
|
||||
module_globals.pop("fix_package_names", None)
|
||||
module_globals.pop("set_deprecated_aliases", None)
|
||||
return True
|
||||
|
||||
|
||||
def fix_package_names() -> None:
|
||||
module_globals = sys._getframe(1).f_globals
|
||||
module_prefix = module_globals["__name__"] + "."
|
||||
del module_globals[fix_package_names.__name__]
|
||||
for value in module_globals.values():
|
||||
if modname := getattr(value, "__module__", ""):
|
||||
if modname.startswith(module_prefix):
|
||||
parts = modname.split(".")
|
||||
value.__module__ = ".".join(
|
||||
part for part in parts if not part.startswith("_")
|
||||
)
|
||||
|
||||
|
||||
def emit_deprecation_warning(module_name: str, name: str, target: str) -> None:
|
||||
warnings.warn(
|
||||
f"The {module_name}.{name} alias is deprecated, use {target} instead.",
|
||||
DeprecationWarning,
|
||||
stacklevel=3,
|
||||
)
|
||||
|
||||
|
||||
def set_deprecated_aliases(aliases: dict[str, str]) -> None:
|
||||
module_globals = sys._getframe(1).f_globals
|
||||
module_name = module_globals["__name__"]
|
||||
del module_globals[set_deprecated_aliases.__name__]
|
||||
|
||||
def __getattr__(name: str) -> Any:
|
||||
try:
|
||||
target = aliases[name]
|
||||
except KeyError:
|
||||
raise AttributeError(
|
||||
f"module {module_name!r} has no attribute {name!r}"
|
||||
) from None
|
||||
|
||||
emit_deprecation_warning(module_name, name, target)
|
||||
target_modname, attrname = target.rsplit(".", 1)
|
||||
module = import_module(target_modname)
|
||||
return getattr(module, attrname)
|
||||
|
||||
sys.modules[module_name].__dict__["__getattr__"] = __getattr__
|
||||
|
||||
|
||||
def _build_lazy_map(
|
||||
module: ModuleType,
|
||||
) -> tuple[dict[str, tuple[str, str]], dict[str, str], list[str]]:
|
||||
try:
|
||||
source = inspect.getsource(module)
|
||||
except OSError:
|
||||
return {}, {}, []
|
||||
|
||||
tree = compile(source, module.__file__ or "", "exec", ast.PyCF_ONLY_AST)
|
||||
assert isinstance(tree, ast.Module)
|
||||
out: dict[str, tuple[str, str]] = {}
|
||||
deprecated_aliases: dict[str, str] = {}
|
||||
submodule_names: list[str] = []
|
||||
|
||||
for node in tree.body:
|
||||
if not isinstance(node, ast.If) or not _is_type_checking_block(node.test):
|
||||
continue
|
||||
|
||||
for stmt in node.body:
|
||||
match stmt:
|
||||
case ast.ImportFrom():
|
||||
if stmt.module is None:
|
||||
submodule_names.extend(alias.name for alias in stmt.names)
|
||||
else:
|
||||
base = "." * stmt.level + (stmt.module or "")
|
||||
for alias in stmt.names:
|
||||
if alias.name == "*":
|
||||
raise RuntimeError("star imports not supported")
|
||||
|
||||
exported = alias.asname or alias.name
|
||||
out[exported] = (base, alias.name)
|
||||
case ast.Expr() if isinstance(stmt.value, ast.Call):
|
||||
call = stmt.value
|
||||
if (
|
||||
isinstance(call.func, ast.Name)
|
||||
and call.func.id == "set_deprecated_aliases"
|
||||
):
|
||||
arg0 = call.args[0]
|
||||
assert isinstance(arg0, ast.Dict)
|
||||
for key, value in zip(arg0.keys, arg0.values, strict=True):
|
||||
assert isinstance(key, ast.Constant)
|
||||
assert isinstance(key.value, str)
|
||||
assert isinstance(value, ast.Constant)
|
||||
assert isinstance(value.value, str)
|
||||
deprecated_aliases[key.value] = value.value
|
||||
|
||||
return out, deprecated_aliases, submodule_names
|
||||
|
||||
|
||||
def _is_type_checking_block(test: ast.AST) -> bool:
|
||||
if not isinstance(test, ast.BoolOp):
|
||||
return False
|
||||
|
||||
subtest = test.values[0]
|
||||
match subtest:
|
||||
case ast.Name():
|
||||
return subtest.id == "TYPE_CHECKING"
|
||||
case ast.Attribute():
|
||||
return (
|
||||
isinstance(subtest.value, ast.Name)
|
||||
and subtest.value.id == "typing"
|
||||
and subtest.attr == "TYPE_CHECKING"
|
||||
)
|
||||
case _:
|
||||
return False
|
||||
+63
@@ -0,0 +1,63 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from .._lazyimport import (
|
||||
fix_package_names,
|
||||
install_lazy_importer,
|
||||
set_deprecated_aliases,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING or not install_lazy_importer():
|
||||
from ._eventloop import AsyncBackend as AsyncBackend
|
||||
from ._resources import AsyncResource as AsyncResource
|
||||
from ._sockets import ConnectedUDPSocket as ConnectedUDPSocket
|
||||
from ._sockets import ConnectedUNIXDatagramSocket as ConnectedUNIXDatagramSocket
|
||||
from ._sockets import IPAddressType as IPAddressType
|
||||
from ._sockets import IPSockAddrType as IPSockAddrType
|
||||
from ._sockets import SocketAttribute as SocketAttribute
|
||||
from ._sockets import SocketListener as SocketListener
|
||||
from ._sockets import SocketStream as SocketStream
|
||||
from ._sockets import UDPPacketType as UDPPacketType
|
||||
from ._sockets import UDPSocket as UDPSocket
|
||||
from ._sockets import UNIXDatagramPacketType as UNIXDatagramPacketType
|
||||
from ._sockets import UNIXDatagramSocket as UNIXDatagramSocket
|
||||
from ._sockets import UNIXSocketStream as UNIXSocketStream
|
||||
from ._streams import AnyByteReceiveStream as AnyByteReceiveStream
|
||||
from ._streams import AnyByteSendStream as AnyByteSendStream
|
||||
from ._streams import AnyByteStream as AnyByteStream
|
||||
from ._streams import AnyByteStreamConnectable as AnyByteStreamConnectable
|
||||
from ._streams import (
|
||||
AnyUnreliableByteReceiveStream as AnyUnreliableByteReceiveStream,
|
||||
)
|
||||
from ._streams import AnyUnreliableByteSendStream as AnyUnreliableByteSendStream
|
||||
from ._streams import AnyUnreliableByteStream as AnyUnreliableByteStream
|
||||
from ._streams import ByteReceiveStream as ByteReceiveStream
|
||||
from ._streams import ByteSendStream as ByteSendStream
|
||||
from ._streams import ByteStream as ByteStream
|
||||
from ._streams import ByteStreamConnectable as ByteStreamConnectable
|
||||
from ._streams import Listener as Listener
|
||||
from ._streams import ObjectReceiveStream as ObjectReceiveStream
|
||||
from ._streams import ObjectSendStream as ObjectSendStream
|
||||
from ._streams import ObjectStream as ObjectStream
|
||||
from ._streams import ObjectStreamConnectable as ObjectStreamConnectable
|
||||
from ._streams import UnreliableObjectReceiveStream as UnreliableObjectReceiveStream
|
||||
from ._streams import UnreliableObjectSendStream as UnreliableObjectSendStream
|
||||
from ._streams import UnreliableObjectStream as UnreliableObjectStream
|
||||
from ._subprocesses import Process as Process
|
||||
from ._tasks import TaskGroup as TaskGroup
|
||||
from ._tasks import TaskStatus as TaskStatus
|
||||
from ._testing import TestRunner as TestRunner
|
||||
|
||||
fix_package_names()
|
||||
set_deprecated_aliases(
|
||||
{
|
||||
"CapacityLimiter": "anyio.CapacityLimiter",
|
||||
"Condition": "anyio.Condition",
|
||||
"Event": "anyio.Event",
|
||||
"Lock": "anyio.Lock",
|
||||
"Semaphore": "anyio.Semaphore",
|
||||
"CancelScope": "anyio.CancelScope",
|
||||
"BlockingPortal": "anyio.from_thread.BlockingPortal",
|
||||
}
|
||||
)
|
||||
+428
@@ -0,0 +1,428 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
import sys
|
||||
from abc import ABCMeta, abstractmethod
|
||||
from collections.abc import (
|
||||
AsyncIterator,
|
||||
Awaitable,
|
||||
Callable,
|
||||
Coroutine,
|
||||
Iterable,
|
||||
Mapping,
|
||||
Sequence,
|
||||
)
|
||||
from contextlib import AbstractContextManager
|
||||
from os import PathLike
|
||||
from signal import Signals
|
||||
from socket import AddressFamily, SocketKind, socket
|
||||
from typing import (
|
||||
IO,
|
||||
TYPE_CHECKING,
|
||||
Any,
|
||||
TypeAlias,
|
||||
TypeVar,
|
||||
overload,
|
||||
)
|
||||
|
||||
if sys.version_info >= (3, 11):
|
||||
from typing import TypeVarTuple, Unpack
|
||||
else:
|
||||
from typing_extensions import TypeVarTuple, Unpack
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from _typeshed import FileDescriptorLike
|
||||
|
||||
from .._core._synchronization import CapacityLimiter, Event, Lock, Semaphore
|
||||
from .._core._tasks import CancelScope
|
||||
from .._core._testing import TaskInfo
|
||||
from ._sockets import (
|
||||
ConnectedUDPSocket,
|
||||
ConnectedUNIXDatagramSocket,
|
||||
IPSockAddrType,
|
||||
SocketListener,
|
||||
SocketStream,
|
||||
UDPSocket,
|
||||
UNIXDatagramSocket,
|
||||
UNIXSocketStream,
|
||||
)
|
||||
from ._subprocesses import Process
|
||||
from ._tasks import TaskGroup
|
||||
from ._testing import TestRunner
|
||||
|
||||
T_Retval = TypeVar("T_Retval")
|
||||
T_co = TypeVar("T_co", covariant=True)
|
||||
PosArgsT = TypeVarTuple("PosArgsT")
|
||||
StrOrBytesPath: TypeAlias = str | bytes | PathLike[str] | PathLike[bytes]
|
||||
|
||||
|
||||
class AsyncBackend(metaclass=ABCMeta):
|
||||
@classmethod
|
||||
@abstractmethod
|
||||
def run(
|
||||
cls,
|
||||
func: Callable[[Unpack[PosArgsT]], Awaitable[T_Retval]],
|
||||
args: tuple[Unpack[PosArgsT]],
|
||||
kwargs: dict[str, Any],
|
||||
options: dict[str, Any],
|
||||
) -> T_Retval:
|
||||
"""
|
||||
Run the given coroutine function in an asynchronous event loop.
|
||||
|
||||
The current thread must not be already running an event loop.
|
||||
|
||||
:param func: a coroutine function
|
||||
:param args: positional arguments to ``func``
|
||||
:param kwargs: positional arguments to ``func``
|
||||
:param options: keyword arguments to call the backend ``run()`` implementation
|
||||
with
|
||||
:return: the return value of the coroutine function
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
@abstractmethod
|
||||
def current_token(cls) -> object:
|
||||
"""
|
||||
Return an object that allows other threads to run code inside the event loop.
|
||||
|
||||
:return: a token object, specific to the event loop running in the current
|
||||
thread
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
@abstractmethod
|
||||
def current_time(cls) -> float:
|
||||
"""
|
||||
Return the current value of the event loop's internal clock.
|
||||
|
||||
:return: the clock value (seconds)
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
@abstractmethod
|
||||
def cancelled_exception_class(cls) -> type[BaseException]:
|
||||
"""Return the exception class that is raised in a task if it's cancelled."""
|
||||
|
||||
@classmethod
|
||||
@abstractmethod
|
||||
async def checkpoint(cls) -> None:
|
||||
"""
|
||||
Check if the task has been cancelled, and allow rescheduling of other tasks.
|
||||
|
||||
This is effectively the same as running :meth:`checkpoint_if_cancelled` and then
|
||||
:meth:`cancel_shielded_checkpoint`.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
async def checkpoint_if_cancelled(cls) -> None:
|
||||
"""
|
||||
Check if the current task group has been cancelled.
|
||||
|
||||
This will check if the task has been cancelled, but will not allow other tasks
|
||||
to be scheduled if not.
|
||||
|
||||
"""
|
||||
if cls.current_effective_deadline() == -math.inf:
|
||||
await cls.checkpoint()
|
||||
|
||||
@classmethod
|
||||
async def cancel_shielded_checkpoint(cls) -> None:
|
||||
"""
|
||||
Allow the rescheduling of other tasks.
|
||||
|
||||
This will give other tasks the opportunity to run, but without checking if the
|
||||
current task group has been cancelled, unlike with :meth:`checkpoint`.
|
||||
|
||||
"""
|
||||
with cls.create_cancel_scope(shield=True):
|
||||
await cls.sleep(0)
|
||||
|
||||
@classmethod
|
||||
@abstractmethod
|
||||
async def sleep(cls, delay: float) -> None:
|
||||
"""
|
||||
Pause the current task for the specified duration.
|
||||
|
||||
:param delay: the duration, in seconds
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
@abstractmethod
|
||||
def create_cancel_scope(
|
||||
cls, *, deadline: float = math.inf, shield: bool = False
|
||||
) -> CancelScope:
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
@abstractmethod
|
||||
def current_effective_deadline(cls) -> float:
|
||||
"""
|
||||
Return the nearest deadline among all the cancel scopes effective for the
|
||||
current task.
|
||||
|
||||
:return:
|
||||
- a clock value from the event loop's internal clock
|
||||
- ``inf`` if there is no deadline in effect
|
||||
- ``-inf`` if the current scope has been cancelled
|
||||
:rtype: float
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
@abstractmethod
|
||||
def create_task_group(cls) -> TaskGroup:
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
@abstractmethod
|
||||
def create_event(cls) -> Event:
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
@abstractmethod
|
||||
def create_lock(cls, *, fast_acquire: bool) -> Lock:
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
@abstractmethod
|
||||
def create_semaphore(
|
||||
cls,
|
||||
initial_value: int,
|
||||
*,
|
||||
max_value: int | None = None,
|
||||
fast_acquire: bool = False,
|
||||
) -> Semaphore:
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
@abstractmethod
|
||||
def create_capacity_limiter(cls, total_tokens: float) -> CapacityLimiter:
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
@abstractmethod
|
||||
async def run_sync_in_worker_thread(
|
||||
cls,
|
||||
func: Callable[[Unpack[PosArgsT]], T_Retval],
|
||||
args: tuple[Unpack[PosArgsT]],
|
||||
abandon_on_cancel: bool = False,
|
||||
limiter: CapacityLimiter | None = None,
|
||||
) -> T_Retval:
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
@abstractmethod
|
||||
def check_cancelled(cls) -> None:
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
@abstractmethod
|
||||
def run_async_from_thread(
|
||||
cls,
|
||||
func: Callable[[Unpack[PosArgsT]], Coroutine[Any, Any, T_co]],
|
||||
args: tuple[Unpack[PosArgsT]],
|
||||
token: object,
|
||||
) -> T_co:
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
@abstractmethod
|
||||
def run_sync_from_thread(
|
||||
cls,
|
||||
func: Callable[[Unpack[PosArgsT]], T_Retval],
|
||||
args: tuple[Unpack[PosArgsT]],
|
||||
token: object,
|
||||
) -> T_Retval:
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
@abstractmethod
|
||||
async def open_process(
|
||||
cls,
|
||||
command: StrOrBytesPath | Sequence[StrOrBytesPath],
|
||||
*,
|
||||
stdin: int | IO[Any] | None,
|
||||
stdout: int | IO[Any] | None,
|
||||
stderr: int | IO[Any] | None,
|
||||
cwd: StrOrBytesPath | None = None,
|
||||
env: Mapping[str, str] | None = None,
|
||||
startupinfo: Any = None,
|
||||
creationflags: int = 0,
|
||||
start_new_session: bool = False,
|
||||
pass_fds: Sequence[int] = (),
|
||||
user: str | int | None = None,
|
||||
group: str | int | None = None,
|
||||
extra_groups: Iterable[str | int] | None = None,
|
||||
umask: int = -1,
|
||||
**kwargs: Any,
|
||||
) -> Process:
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
@abstractmethod
|
||||
def setup_process_pool_exit_at_shutdown(cls, workers: set[Process]) -> None:
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
@abstractmethod
|
||||
async def connect_tcp(
|
||||
cls, host: str, port: int, local_address: IPSockAddrType | None = None
|
||||
) -> SocketStream:
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
@abstractmethod
|
||||
async def connect_unix(cls, path: str | bytes) -> UNIXSocketStream:
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
@abstractmethod
|
||||
def create_tcp_listener(cls, sock: socket) -> SocketListener:
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
@abstractmethod
|
||||
def create_unix_listener(cls, sock: socket) -> SocketListener:
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
@abstractmethod
|
||||
async def create_udp_socket(
|
||||
cls,
|
||||
family: AddressFamily,
|
||||
local_address: IPSockAddrType | None,
|
||||
remote_address: IPSockAddrType | None,
|
||||
reuse_port: bool,
|
||||
) -> UDPSocket | ConnectedUDPSocket:
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
@overload
|
||||
async def create_unix_datagram_socket(
|
||||
cls, raw_socket: socket, remote_path: None
|
||||
) -> UNIXDatagramSocket: ...
|
||||
|
||||
@classmethod
|
||||
@overload
|
||||
async def create_unix_datagram_socket(
|
||||
cls, raw_socket: socket, remote_path: str | bytes
|
||||
) -> ConnectedUNIXDatagramSocket: ...
|
||||
|
||||
@classmethod
|
||||
@abstractmethod
|
||||
async def create_unix_datagram_socket(
|
||||
cls, raw_socket: socket, remote_path: str | bytes | None
|
||||
) -> UNIXDatagramSocket | ConnectedUNIXDatagramSocket:
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
@abstractmethod
|
||||
async def getaddrinfo(
|
||||
cls,
|
||||
host: bytes | str | None,
|
||||
port: str | int | None,
|
||||
*,
|
||||
family: int | AddressFamily = 0,
|
||||
type: int | SocketKind = 0,
|
||||
proto: int = 0,
|
||||
flags: int = 0,
|
||||
) -> Sequence[
|
||||
tuple[
|
||||
AddressFamily,
|
||||
SocketKind,
|
||||
int,
|
||||
str,
|
||||
tuple[str, int] | tuple[str, int, int, int] | tuple[int, bytes],
|
||||
]
|
||||
]:
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
@abstractmethod
|
||||
async def getnameinfo(
|
||||
cls, sockaddr: IPSockAddrType, flags: int = 0
|
||||
) -> tuple[str, str]:
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
@abstractmethod
|
||||
async def wait_readable(cls, obj: FileDescriptorLike) -> None:
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
@abstractmethod
|
||||
async def wait_writable(cls, obj: FileDescriptorLike) -> None:
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
@abstractmethod
|
||||
def notify_closing(cls, obj: FileDescriptorLike) -> None:
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
@abstractmethod
|
||||
async def wrap_listener_socket(cls, sock: socket) -> SocketListener:
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
@abstractmethod
|
||||
async def wrap_stream_socket(cls, sock: socket) -> SocketStream:
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
@abstractmethod
|
||||
async def wrap_unix_stream_socket(cls, sock: socket) -> UNIXSocketStream:
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
@abstractmethod
|
||||
async def wrap_udp_socket(cls, sock: socket) -> UDPSocket:
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
@abstractmethod
|
||||
async def wrap_connected_udp_socket(cls, sock: socket) -> ConnectedUDPSocket:
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
@abstractmethod
|
||||
async def wrap_unix_datagram_socket(cls, sock: socket) -> UNIXDatagramSocket:
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
@abstractmethod
|
||||
async def wrap_connected_unix_datagram_socket(
|
||||
cls, sock: socket
|
||||
) -> ConnectedUNIXDatagramSocket:
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
@abstractmethod
|
||||
def current_default_thread_limiter(cls) -> CapacityLimiter:
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
@abstractmethod
|
||||
def open_signal_receiver(
|
||||
cls, *signals: Signals
|
||||
) -> AbstractContextManager[AsyncIterator[Signals]]:
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
@abstractmethod
|
||||
def get_current_task(cls) -> TaskInfo:
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
@abstractmethod
|
||||
def get_running_tasks(cls) -> Sequence[TaskInfo]:
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
@abstractmethod
|
||||
async def wait_all_tasks_blocked(cls) -> None:
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
@abstractmethod
|
||||
def create_test_runner(cls, options: dict[str, Any]) -> TestRunner:
|
||||
pass
|
||||
+39
@@ -0,0 +1,39 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
from abc import ABCMeta, abstractmethod
|
||||
from types import TracebackType
|
||||
from typing import TypeVar
|
||||
|
||||
if sys.version_info >= (3, 11):
|
||||
from typing import Self
|
||||
else:
|
||||
from typing_extensions import Self
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
|
||||
class AsyncResource(metaclass=ABCMeta):
|
||||
"""
|
||||
Abstract base class for all closeable asynchronous resources.
|
||||
|
||||
Works as an asynchronous context manager which returns the instance itself on enter,
|
||||
and calls :meth:`aclose` on exit.
|
||||
"""
|
||||
|
||||
__slots__ = ()
|
||||
|
||||
async def __aenter__(self) -> Self:
|
||||
return self
|
||||
|
||||
async def __aexit__(
|
||||
self,
|
||||
exc_type: type[BaseException] | None,
|
||||
exc_val: BaseException | None,
|
||||
exc_tb: TracebackType | None,
|
||||
) -> None:
|
||||
await self.aclose()
|
||||
|
||||
@abstractmethod
|
||||
async def aclose(self) -> None:
|
||||
"""Close the resource."""
|
||||
+401
@@ -0,0 +1,401 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import errno
|
||||
import socket
|
||||
from abc import abstractmethod
|
||||
from collections.abc import Callable, Collection, Mapping
|
||||
from contextlib import AsyncExitStack
|
||||
from io import IOBase
|
||||
from ipaddress import IPv4Address, IPv6Address
|
||||
from socket import AddressFamily
|
||||
from typing import TYPE_CHECKING, Any, TypeAlias, TypeVar
|
||||
|
||||
from .._core._eventloop import get_async_backend
|
||||
from .._core._typedattr import (
|
||||
TypedAttributeProvider,
|
||||
TypedAttributeSet,
|
||||
typed_attribute,
|
||||
)
|
||||
from ._streams import ByteStream, Listener, UnreliableObjectStream
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ._tasks import TaskGroup
|
||||
|
||||
IPAddressType: TypeAlias = str | IPv4Address | IPv6Address
|
||||
IPSockAddrType: TypeAlias = tuple[str, int]
|
||||
SockAddrType: TypeAlias = IPSockAddrType | str
|
||||
UDPPacketType: TypeAlias = tuple[bytes, IPSockAddrType]
|
||||
UNIXDatagramPacketType: TypeAlias = tuple[bytes, str]
|
||||
T_Retval = TypeVar("T_Retval")
|
||||
|
||||
|
||||
def _validate_socket(
|
||||
sock_or_fd: socket.socket | int,
|
||||
sock_type: socket.SocketKind,
|
||||
addr_family: socket.AddressFamily = socket.AF_UNSPEC,
|
||||
*,
|
||||
require_connected: bool = False,
|
||||
require_bound: bool = False,
|
||||
) -> socket.socket:
|
||||
if isinstance(sock_or_fd, int):
|
||||
try:
|
||||
sock = socket.socket(fileno=sock_or_fd)
|
||||
except OSError as exc:
|
||||
if exc.errno == errno.ENOTSOCK:
|
||||
raise ValueError(
|
||||
"the file descriptor does not refer to a socket"
|
||||
) from exc
|
||||
elif require_connected:
|
||||
raise ValueError("the socket must be connected") from exc
|
||||
elif require_bound:
|
||||
raise ValueError("the socket must be bound to a local address") from exc
|
||||
else:
|
||||
raise
|
||||
elif isinstance(sock_or_fd, socket.socket):
|
||||
sock = sock_or_fd
|
||||
else:
|
||||
raise TypeError(
|
||||
f"expected an int or socket, got {type(sock_or_fd).__qualname__} instead"
|
||||
)
|
||||
|
||||
try:
|
||||
if require_connected:
|
||||
try:
|
||||
sock.getpeername()
|
||||
except OSError as exc:
|
||||
raise ValueError("the socket must be connected") from exc
|
||||
|
||||
if require_bound:
|
||||
try:
|
||||
if sock.family in (socket.AF_INET, socket.AF_INET6):
|
||||
bound_addr = sock.getsockname()[1]
|
||||
else:
|
||||
bound_addr = sock.getsockname()
|
||||
except OSError:
|
||||
bound_addr = None
|
||||
|
||||
if not bound_addr:
|
||||
raise ValueError("the socket must be bound to a local address")
|
||||
|
||||
if addr_family != socket.AF_UNSPEC and sock.family != addr_family:
|
||||
raise ValueError(
|
||||
f"address family mismatch: expected {addr_family.name}, got "
|
||||
f"{sock.family.name}"
|
||||
)
|
||||
|
||||
if sock.type != sock_type:
|
||||
raise ValueError(
|
||||
f"socket type mismatch: expected {sock_type.name}, got {sock.type.name}"
|
||||
)
|
||||
except BaseException:
|
||||
# Avoid ResourceWarning from the locally constructed socket object
|
||||
if isinstance(sock_or_fd, int):
|
||||
sock.detach()
|
||||
|
||||
raise
|
||||
|
||||
sock.setblocking(False)
|
||||
return sock
|
||||
|
||||
|
||||
class SocketAttribute(TypedAttributeSet):
|
||||
"""
|
||||
.. attribute:: family
|
||||
:type: socket.AddressFamily
|
||||
|
||||
the address family of the underlying socket
|
||||
|
||||
.. attribute:: local_address
|
||||
:type: tuple[str, int] | str
|
||||
|
||||
the local address the underlying socket is connected to
|
||||
|
||||
.. attribute:: local_port
|
||||
:type: int
|
||||
|
||||
for IP based sockets, the local port the underlying socket is bound to
|
||||
|
||||
.. attribute:: raw_socket
|
||||
:type: socket.socket
|
||||
|
||||
the underlying stdlib socket object
|
||||
|
||||
.. attribute:: remote_address
|
||||
:type: tuple[str, int] | str
|
||||
|
||||
the remote address the underlying socket is connected to
|
||||
|
||||
.. attribute:: remote_port
|
||||
:type: int
|
||||
|
||||
for IP based sockets, the remote port the underlying socket is connected to
|
||||
"""
|
||||
|
||||
family: AddressFamily = typed_attribute()
|
||||
local_address: SockAddrType = typed_attribute()
|
||||
local_port: int = typed_attribute()
|
||||
raw_socket: socket.socket = typed_attribute()
|
||||
remote_address: SockAddrType = typed_attribute()
|
||||
remote_port: int = typed_attribute()
|
||||
|
||||
|
||||
class _SocketProvider(TypedAttributeProvider):
|
||||
@property
|
||||
def extra_attributes(self) -> Mapping[Any, Callable[[], Any]]:
|
||||
from .._core._sockets import convert_ipv6_sockaddr as convert
|
||||
|
||||
attributes: dict[Any, Callable[[], Any]] = {
|
||||
SocketAttribute.family: lambda: self._raw_socket.family,
|
||||
SocketAttribute.local_address: lambda: convert(
|
||||
self._raw_socket.getsockname()
|
||||
),
|
||||
SocketAttribute.raw_socket: lambda: self._raw_socket,
|
||||
}
|
||||
try:
|
||||
peername: tuple[str, int] | None = convert(self._raw_socket.getpeername())
|
||||
except OSError:
|
||||
peername = None
|
||||
|
||||
# Provide the remote address for connected sockets
|
||||
if peername is not None:
|
||||
attributes[SocketAttribute.remote_address] = lambda: peername
|
||||
|
||||
# Provide local and remote ports for IP based sockets
|
||||
if self._raw_socket.family in (AddressFamily.AF_INET, AddressFamily.AF_INET6):
|
||||
attributes[SocketAttribute.local_port] = lambda: (
|
||||
self._raw_socket.getsockname()[1]
|
||||
)
|
||||
if peername is not None:
|
||||
remote_port = peername[1]
|
||||
attributes[SocketAttribute.remote_port] = lambda: remote_port
|
||||
|
||||
return attributes
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def _raw_socket(self) -> socket.socket:
|
||||
pass
|
||||
|
||||
|
||||
class SocketStream(ByteStream, _SocketProvider):
|
||||
"""
|
||||
Transports bytes over a socket.
|
||||
|
||||
Supports all relevant extra attributes from :class:`~SocketAttribute`.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
async def from_socket(cls, sock_or_fd: socket.socket | int) -> SocketStream:
|
||||
"""
|
||||
Wrap an existing socket object or file descriptor as a socket stream.
|
||||
|
||||
The newly created socket wrapper takes ownership of the socket being passed in.
|
||||
The existing socket must already be connected.
|
||||
|
||||
:param sock_or_fd: a socket object or file descriptor
|
||||
:return: a socket stream
|
||||
|
||||
"""
|
||||
sock = _validate_socket(sock_or_fd, socket.SOCK_STREAM, require_connected=True)
|
||||
return await get_async_backend().wrap_stream_socket(sock)
|
||||
|
||||
|
||||
class UNIXSocketStream(SocketStream):
|
||||
@classmethod
|
||||
async def from_socket(cls, sock_or_fd: socket.socket | int) -> UNIXSocketStream:
|
||||
"""
|
||||
Wrap an existing socket object or file descriptor as a UNIX socket stream.
|
||||
|
||||
The newly created socket wrapper takes ownership of the socket being passed in.
|
||||
The existing socket must already be connected.
|
||||
|
||||
:param sock_or_fd: a socket object or file descriptor
|
||||
:return: a UNIX socket stream
|
||||
|
||||
"""
|
||||
sock = _validate_socket(
|
||||
sock_or_fd, socket.SOCK_STREAM, socket.AF_UNIX, require_connected=True
|
||||
)
|
||||
return await get_async_backend().wrap_unix_stream_socket(sock)
|
||||
|
||||
@abstractmethod
|
||||
async def send_fds(self, message: bytes, fds: Collection[int | IOBase]) -> None:
|
||||
"""
|
||||
Send file descriptors along with a message to the peer.
|
||||
|
||||
:param message: a non-empty bytestring
|
||||
:param fds: a collection of files (either numeric file descriptors or open file
|
||||
or socket objects)
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
async def receive_fds(self, msglen: int, maxfds: int) -> tuple[bytes, list[int]]:
|
||||
"""
|
||||
Receive file descriptors along with a message from the peer.
|
||||
|
||||
:param msglen: length of the message to expect from the peer
|
||||
:param maxfds: maximum number of file descriptors to expect from the peer
|
||||
:return: a tuple of (message, file descriptors)
|
||||
"""
|
||||
|
||||
|
||||
class SocketListener(Listener[SocketStream], _SocketProvider):
|
||||
"""
|
||||
Listens to incoming socket connections.
|
||||
|
||||
Supports all relevant extra attributes from :class:`~SocketAttribute`.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
async def from_socket(
|
||||
cls,
|
||||
sock_or_fd: socket.socket | int,
|
||||
) -> SocketListener:
|
||||
"""
|
||||
Wrap an existing socket object or file descriptor as a socket listener.
|
||||
|
||||
The newly created listener takes ownership of the socket being passed in.
|
||||
|
||||
:param sock_or_fd: a socket object or file descriptor
|
||||
:return: a socket listener
|
||||
|
||||
"""
|
||||
sock = _validate_socket(sock_or_fd, socket.SOCK_STREAM, require_bound=True)
|
||||
return await get_async_backend().wrap_listener_socket(sock)
|
||||
|
||||
@abstractmethod
|
||||
async def accept(self) -> SocketStream:
|
||||
"""Accept an incoming connection."""
|
||||
|
||||
async def serve(
|
||||
self,
|
||||
handler: Callable[[SocketStream], Any],
|
||||
task_group: TaskGroup | None = None,
|
||||
) -> None:
|
||||
from .. import create_task_group
|
||||
|
||||
async with AsyncExitStack() as stack:
|
||||
if task_group is None:
|
||||
task_group = await stack.enter_async_context(create_task_group())
|
||||
|
||||
while True:
|
||||
stream = await self.accept()
|
||||
task_group.start_soon(handler, stream)
|
||||
|
||||
|
||||
class UDPSocket(UnreliableObjectStream[UDPPacketType], _SocketProvider):
|
||||
"""
|
||||
Represents an unconnected UDP socket.
|
||||
|
||||
Supports all relevant extra attributes from :class:`~SocketAttribute`.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
async def from_socket(cls, sock_or_fd: socket.socket | int) -> UDPSocket:
|
||||
"""
|
||||
Wrap an existing socket object or file descriptor as a UDP socket.
|
||||
|
||||
The newly created socket wrapper takes ownership of the socket being passed in.
|
||||
The existing socket must be bound to a local address.
|
||||
|
||||
:param sock_or_fd: a socket object or file descriptor
|
||||
:return: a UDP socket
|
||||
|
||||
"""
|
||||
sock = _validate_socket(sock_or_fd, socket.SOCK_DGRAM, require_bound=True)
|
||||
return await get_async_backend().wrap_udp_socket(sock)
|
||||
|
||||
async def sendto(self, data: bytes, host: str, port: int) -> None:
|
||||
"""
|
||||
Alias for :meth:`~.UnreliableObjectSendStream.send` ((data, (host, port))).
|
||||
|
||||
"""
|
||||
return await self.send((data, (host, port)))
|
||||
|
||||
|
||||
class ConnectedUDPSocket(UnreliableObjectStream[bytes], _SocketProvider):
|
||||
"""
|
||||
Represents an connected UDP socket.
|
||||
|
||||
Supports all relevant extra attributes from :class:`~SocketAttribute`.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
async def from_socket(cls, sock_or_fd: socket.socket | int) -> ConnectedUDPSocket:
|
||||
"""
|
||||
Wrap an existing socket object or file descriptor as a connected UDP socket.
|
||||
|
||||
The newly created socket wrapper takes ownership of the socket being passed in.
|
||||
The existing socket must already be connected.
|
||||
|
||||
:param sock_or_fd: a socket object or file descriptor
|
||||
:return: a connected UDP socket
|
||||
|
||||
"""
|
||||
sock = _validate_socket(
|
||||
sock_or_fd,
|
||||
socket.SOCK_DGRAM,
|
||||
require_connected=True,
|
||||
)
|
||||
return await get_async_backend().wrap_connected_udp_socket(sock)
|
||||
|
||||
|
||||
class UNIXDatagramSocket(
|
||||
UnreliableObjectStream[UNIXDatagramPacketType], _SocketProvider
|
||||
):
|
||||
"""
|
||||
Represents an unconnected Unix datagram socket.
|
||||
|
||||
Supports all relevant extra attributes from :class:`~SocketAttribute`.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
async def from_socket(
|
||||
cls,
|
||||
sock_or_fd: socket.socket | int,
|
||||
) -> UNIXDatagramSocket:
|
||||
"""
|
||||
Wrap an existing socket object or file descriptor as a UNIX datagram
|
||||
socket.
|
||||
|
||||
The newly created socket wrapper takes ownership of the socket being passed in.
|
||||
|
||||
:param sock_or_fd: a socket object or file descriptor
|
||||
:return: a UNIX datagram socket
|
||||
|
||||
"""
|
||||
sock = _validate_socket(sock_or_fd, socket.SOCK_DGRAM, socket.AF_UNIX)
|
||||
return await get_async_backend().wrap_unix_datagram_socket(sock)
|
||||
|
||||
async def sendto(self, data: bytes, path: str) -> None:
|
||||
"""Alias for :meth:`~.UnreliableObjectSendStream.send` ((data, path))."""
|
||||
return await self.send((data, path))
|
||||
|
||||
|
||||
class ConnectedUNIXDatagramSocket(UnreliableObjectStream[bytes], _SocketProvider):
|
||||
"""
|
||||
Represents a connected Unix datagram socket.
|
||||
|
||||
Supports all relevant extra attributes from :class:`~SocketAttribute`.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
async def from_socket(
|
||||
cls,
|
||||
sock_or_fd: socket.socket | int,
|
||||
) -> ConnectedUNIXDatagramSocket:
|
||||
"""
|
||||
Wrap an existing socket object or file descriptor as a connected UNIX datagram
|
||||
socket.
|
||||
|
||||
The newly created socket wrapper takes ownership of the socket being passed in.
|
||||
The existing socket must already be connected.
|
||||
|
||||
:param sock_or_fd: a socket object or file descriptor
|
||||
:return: a connected UNIX datagram socket
|
||||
|
||||
"""
|
||||
sock = _validate_socket(
|
||||
sock_or_fd, socket.SOCK_DGRAM, socket.AF_UNIX, require_connected=True
|
||||
)
|
||||
return await get_async_backend().wrap_connected_unix_datagram_socket(sock)
|
||||
+237
@@ -0,0 +1,237 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from abc import ABCMeta, abstractmethod
|
||||
from collections.abc import Callable
|
||||
from typing import TYPE_CHECKING, Any, Generic, TypeAlias, TypeVar
|
||||
|
||||
from .._core._exceptions import EndOfStream
|
||||
from .._core._typedattr import TypedAttributeProvider
|
||||
from ._resources import AsyncResource
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ._tasks import TaskGroup
|
||||
|
||||
T_Item = TypeVar("T_Item")
|
||||
T_co = TypeVar("T_co", covariant=True)
|
||||
T_contra = TypeVar("T_contra", contravariant=True)
|
||||
|
||||
|
||||
class UnreliableObjectReceiveStream(
|
||||
AsyncResource, TypedAttributeProvider, Generic[T_co]
|
||||
):
|
||||
"""
|
||||
An interface for receiving objects.
|
||||
|
||||
This interface makes no guarantees that the received messages arrive in the order in
|
||||
which they were sent, or that no messages are missed.
|
||||
|
||||
Asynchronously iterating over objects of this type will yield objects matching the
|
||||
given type parameter.
|
||||
"""
|
||||
|
||||
def __aiter__(self) -> UnreliableObjectReceiveStream[T_co]:
|
||||
return self
|
||||
|
||||
async def __anext__(self) -> T_co:
|
||||
try:
|
||||
return await self.receive()
|
||||
except EndOfStream:
|
||||
raise StopAsyncIteration from None
|
||||
|
||||
@abstractmethod
|
||||
async def receive(self) -> T_co:
|
||||
"""
|
||||
Receive the next item.
|
||||
|
||||
:raises ~anyio.ClosedResourceError: if the receive stream has been explicitly
|
||||
closed
|
||||
:raises ~anyio.EndOfStream: if this stream has been closed from the other end
|
||||
:raises ~anyio.BrokenResourceError: if this stream has been rendered unusable
|
||||
due to external causes
|
||||
"""
|
||||
|
||||
|
||||
class UnreliableObjectSendStream(
|
||||
AsyncResource, TypedAttributeProvider, Generic[T_contra]
|
||||
):
|
||||
"""
|
||||
An interface for sending objects.
|
||||
|
||||
This interface makes no guarantees that the messages sent will reach the
|
||||
recipient(s) in the same order in which they were sent, or at all.
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
async def send(self, item: T_contra) -> None:
|
||||
"""
|
||||
Send an item to the peer(s).
|
||||
|
||||
:param item: the item to send
|
||||
:raises ~anyio.ClosedResourceError: if the send stream has been explicitly
|
||||
closed
|
||||
:raises ~anyio.BrokenResourceError: if this stream has been rendered unusable
|
||||
due to external causes
|
||||
"""
|
||||
|
||||
|
||||
class UnreliableObjectStream(
|
||||
UnreliableObjectReceiveStream[T_Item], UnreliableObjectSendStream[T_Item]
|
||||
):
|
||||
"""
|
||||
A bidirectional message stream which does not guarantee the order or reliability of
|
||||
message delivery.
|
||||
"""
|
||||
|
||||
|
||||
class ObjectReceiveStream(UnreliableObjectReceiveStream[T_co]):
|
||||
"""
|
||||
A receive message stream which guarantees that messages are received in the same
|
||||
order in which they were sent, and that no messages are missed.
|
||||
"""
|
||||
|
||||
|
||||
class ObjectSendStream(UnreliableObjectSendStream[T_contra]):
|
||||
"""
|
||||
A send message stream which guarantees that messages are delivered in the same order
|
||||
in which they were sent, without missing any messages in the middle.
|
||||
"""
|
||||
|
||||
|
||||
class ObjectStream(
|
||||
ObjectReceiveStream[T_Item],
|
||||
ObjectSendStream[T_Item],
|
||||
UnreliableObjectStream[T_Item],
|
||||
):
|
||||
"""
|
||||
A bidirectional message stream which guarantees the order and reliability of message
|
||||
delivery.
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
async def send_eof(self) -> None:
|
||||
"""
|
||||
Send an end-of-file indication to the peer.
|
||||
|
||||
You should not try to send any further data to this stream after calling this
|
||||
method. This method is idempotent (does nothing on successive calls).
|
||||
"""
|
||||
|
||||
|
||||
class ByteReceiveStream(AsyncResource, TypedAttributeProvider):
|
||||
"""
|
||||
An interface for receiving bytes from a single peer.
|
||||
|
||||
Iterating this byte stream will yield a byte string of arbitrary length, but no more
|
||||
than 65536 bytes.
|
||||
"""
|
||||
|
||||
def __aiter__(self) -> ByteReceiveStream:
|
||||
return self
|
||||
|
||||
async def __anext__(self) -> bytes:
|
||||
try:
|
||||
return await self.receive()
|
||||
except EndOfStream:
|
||||
raise StopAsyncIteration from None
|
||||
|
||||
@abstractmethod
|
||||
async def receive(self, max_bytes: int = 65536) -> bytes:
|
||||
"""
|
||||
Receive at most ``max_bytes`` bytes from the peer.
|
||||
|
||||
.. note:: Implementers of this interface should not return an empty
|
||||
:class:`bytes` object, and users should ignore them.
|
||||
|
||||
:param max_bytes: maximum number of bytes to receive (must be a positive
|
||||
integer)
|
||||
:return: the received bytes
|
||||
:raises ValueError: if ``max_bytes`` is less than 1
|
||||
:raises ~anyio.EndOfStream: if this stream has been closed from the other end
|
||||
"""
|
||||
|
||||
|
||||
class ByteSendStream(AsyncResource, TypedAttributeProvider):
|
||||
"""An interface for sending bytes to a single peer."""
|
||||
|
||||
@abstractmethod
|
||||
async def send(self, item: bytes) -> None:
|
||||
"""
|
||||
Send the given bytes to the peer.
|
||||
|
||||
:param item: the bytes to send
|
||||
"""
|
||||
|
||||
|
||||
class ByteStream(ByteReceiveStream, ByteSendStream):
|
||||
"""A bidirectional byte stream."""
|
||||
|
||||
@abstractmethod
|
||||
async def send_eof(self) -> None:
|
||||
"""
|
||||
Send an end-of-file indication to the peer.
|
||||
|
||||
You should not try to send any further data to this stream after calling this
|
||||
method. This method is idempotent (does nothing on successive calls).
|
||||
"""
|
||||
|
||||
|
||||
#: Type alias for all unreliable bytes-oriented receive streams.
|
||||
AnyUnreliableByteReceiveStream: TypeAlias = (
|
||||
UnreliableObjectReceiveStream[bytes] | ByteReceiveStream
|
||||
)
|
||||
#: Type alias for all unreliable bytes-oriented send streams.
|
||||
AnyUnreliableByteSendStream: TypeAlias = (
|
||||
UnreliableObjectSendStream[bytes] | ByteSendStream
|
||||
)
|
||||
#: Type alias for all unreliable bytes-oriented streams.
|
||||
AnyUnreliableByteStream: TypeAlias = UnreliableObjectStream[bytes] | ByteStream
|
||||
#: Type alias for all bytes-oriented receive streams.
|
||||
AnyByteReceiveStream: TypeAlias = ObjectReceiveStream[bytes] | ByteReceiveStream
|
||||
#: Type alias for all bytes-oriented send streams.
|
||||
AnyByteSendStream: TypeAlias = ObjectSendStream[bytes] | ByteSendStream
|
||||
#: Type alias for all bytes-oriented streams.
|
||||
AnyByteStream: TypeAlias = ObjectStream[bytes] | ByteStream
|
||||
|
||||
|
||||
class Listener(AsyncResource, TypedAttributeProvider, Generic[T_co]):
|
||||
"""An interface for objects that let you accept incoming connections."""
|
||||
|
||||
@abstractmethod
|
||||
async def serve(
|
||||
self, handler: Callable[[T_co], Any], task_group: TaskGroup | None = None
|
||||
) -> None:
|
||||
"""
|
||||
Accept incoming connections as they come in and start tasks to handle them.
|
||||
|
||||
:param handler: a callable that will be used to handle each accepted connection
|
||||
:param task_group: the task group that will be used to start tasks for handling
|
||||
each accepted connection (if omitted, an ad-hoc task group will be created)
|
||||
"""
|
||||
|
||||
|
||||
class ObjectStreamConnectable(Generic[T_co], metaclass=ABCMeta):
|
||||
@abstractmethod
|
||||
async def connect(self) -> ObjectStream[T_co]:
|
||||
"""
|
||||
Connect to the remote endpoint.
|
||||
|
||||
:return: an object stream connected to the remote end
|
||||
:raises ConnectionFailed: if the connection fails
|
||||
"""
|
||||
|
||||
|
||||
class ByteStreamConnectable(metaclass=ABCMeta):
|
||||
@abstractmethod
|
||||
async def connect(self) -> ByteStream:
|
||||
"""
|
||||
Connect to the remote endpoint.
|
||||
|
||||
:return: a bytestream connected to the remote end
|
||||
:raises ConnectionFailed: if the connection fails
|
||||
"""
|
||||
|
||||
|
||||
#: Type alias for all connectables returning bytestreams or bytes-oriented object streams
|
||||
AnyByteStreamConnectable: TypeAlias = (
|
||||
ObjectStreamConnectable[bytes] | ByteStreamConnectable
|
||||
)
|
||||
+82
@@ -0,0 +1,82 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from abc import abstractmethod
|
||||
from signal import Signals
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from ._resources import AsyncResource
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ._streams import ByteReceiveStream, ByteSendStream
|
||||
|
||||
|
||||
class Process(AsyncResource):
|
||||
"""An asynchronous version of :class:`subprocess.Popen`."""
|
||||
|
||||
@abstractmethod
|
||||
async def wait(self) -> int:
|
||||
"""
|
||||
Wait until the process exits.
|
||||
|
||||
:return: the exit code of the process
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def terminate(self) -> None:
|
||||
"""
|
||||
Terminates the process, gracefully if possible.
|
||||
|
||||
On Windows, this calls ``TerminateProcess()``.
|
||||
On POSIX systems, this sends ``SIGTERM`` to the process.
|
||||
|
||||
.. seealso:: :meth:`subprocess.Popen.terminate`
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def kill(self) -> None:
|
||||
"""
|
||||
Kills the process.
|
||||
|
||||
On Windows, this calls ``TerminateProcess()``.
|
||||
On POSIX systems, this sends ``SIGKILL`` to the process.
|
||||
|
||||
.. seealso:: :meth:`subprocess.Popen.kill`
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def send_signal(self, signal: Signals) -> None:
|
||||
"""
|
||||
Send a signal to the subprocess.
|
||||
|
||||
.. seealso:: :meth:`subprocess.Popen.send_signal`
|
||||
|
||||
:param signal: the signal number (e.g. :data:`signal.SIGHUP`)
|
||||
"""
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def pid(self) -> int:
|
||||
"""The process ID of the process."""
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def returncode(self) -> int | None:
|
||||
"""
|
||||
The return code of the process. If the process has not yet terminated, this will
|
||||
be ``None``.
|
||||
"""
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def stdin(self) -> ByteSendStream | None:
|
||||
"""The stream for the standard input of the process."""
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def stdout(self) -> ByteReceiveStream | None:
|
||||
"""The stream for the standard output of the process."""
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def stderr(self) -> ByteReceiveStream | None:
|
||||
"""The stream for the standard error output of the process."""
|
||||
+223
@@ -0,0 +1,223 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
from abc import ABCMeta, abstractmethod
|
||||
from collections.abc import Callable, Coroutine
|
||||
from contextvars import Context
|
||||
from types import TracebackType
|
||||
from typing import TYPE_CHECKING, Any, Literal, Protocol, final, overload
|
||||
|
||||
if sys.version_info >= (3, 13):
|
||||
from typing import TypeVar
|
||||
else:
|
||||
from typing_extensions import TypeVar
|
||||
|
||||
if sys.version_info >= (3, 11):
|
||||
from typing import TypeVarTuple, Unpack
|
||||
else:
|
||||
from typing_extensions import TypeVarTuple, Unpack
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .._core._tasks import CancelScope, TaskHandle
|
||||
|
||||
T_co = TypeVar("T_co", covariant=True)
|
||||
T_contra = TypeVar("T_contra", contravariant=True, default=None)
|
||||
PosArgsT = TypeVarTuple("PosArgsT")
|
||||
|
||||
|
||||
def get_coro_name(coro: Coroutine[Any, Any, object], override: object = None) -> str:
|
||||
if override is not None:
|
||||
return str(override)
|
||||
|
||||
try:
|
||||
cr_frame = coro.cr_frame # type: ignore[attr-defined]
|
||||
module = cr_frame.f_globals["__name__"]
|
||||
except (AttributeError, KeyError):
|
||||
module = None
|
||||
|
||||
qualname = getattr(coro, "__qualname__", None)
|
||||
return ".".join([x for x in (module, qualname) if x])
|
||||
|
||||
|
||||
def get_callable_name(func: Callable, override: object = None) -> str:
|
||||
if override is not None:
|
||||
return str(override)
|
||||
|
||||
module = getattr(func, "__module__", None)
|
||||
qualname = getattr(func, "__qualname__", None)
|
||||
return ".".join([x for x in (module, qualname) if x])
|
||||
|
||||
|
||||
def call_for_coroutine(
|
||||
func: Callable[[Unpack[PosArgsT]], Coroutine[Any, Any, T_co]],
|
||||
args: tuple[Unpack[PosArgsT]],
|
||||
**kwargs: Any,
|
||||
) -> Coroutine[Any, Any, T_co]:
|
||||
"""
|
||||
Call the given function with the given positional and keyword arguments.
|
||||
|
||||
:return: the resulting coroutine
|
||||
:raises TypeError: if the return value was not a coroutine object
|
||||
|
||||
"""
|
||||
coro = func(*args, **kwargs)
|
||||
if not isinstance(coro, Coroutine):
|
||||
prefix = f"{func.__module__}." if hasattr(func, "__module__") else ""
|
||||
raise TypeError(
|
||||
f"Expected {prefix}{func.__qualname__}() to return a coroutine, but "
|
||||
f"the return value ({coro!r}) is not a coroutine object"
|
||||
)
|
||||
|
||||
return coro
|
||||
|
||||
|
||||
class TaskStatus(Protocol[T_contra]):
|
||||
@overload
|
||||
def started(self: TaskStatus[None]) -> None: ...
|
||||
|
||||
@overload
|
||||
def started(self, value: T_contra) -> None: ...
|
||||
|
||||
def started(self, value: T_contra | None = None) -> None:
|
||||
"""
|
||||
Signal that the task has started.
|
||||
|
||||
:param value: object passed back to the starter of the task
|
||||
"""
|
||||
|
||||
|
||||
class TaskGroup(metaclass=ABCMeta):
|
||||
"""
|
||||
Groups several asynchronous tasks together.
|
||||
|
||||
:ivar cancel_scope: the cancel scope inherited by all child tasks
|
||||
:vartype cancel_scope: CancelScope
|
||||
|
||||
.. note:: On asyncio, support for eager task factories is considered to be
|
||||
**experimental**. In particular, they don't follow the usual semantics of new
|
||||
tasks being scheduled on the next iteration of the event loop, and may thus
|
||||
cause unexpected behavior in code that wasn't written with such semantics in
|
||||
mind.
|
||||
"""
|
||||
|
||||
cancel_scope: CancelScope
|
||||
|
||||
def cancel(self, reason: str | None = None) -> None:
|
||||
"""
|
||||
Cancel this task group's cancel scope immediately.
|
||||
|
||||
This is a shortcut for calling ``.cancel_scope.cancel()`` on the task group.
|
||||
|
||||
:param reason: a message describing the reason for the cancellation
|
||||
|
||||
.. versionadded:: 4.14.0
|
||||
|
||||
"""
|
||||
self.cancel_scope.cancel(reason)
|
||||
|
||||
@abstractmethod
|
||||
def create_task(
|
||||
self,
|
||||
coro: Coroutine[Any, Any, T_co],
|
||||
*,
|
||||
name: object = None,
|
||||
context: Context | None = None,
|
||||
) -> TaskHandle[T_co]:
|
||||
"""
|
||||
Create a new task from a coroutine object and schedule it to run.
|
||||
|
||||
:param coro: a coroutine object
|
||||
:param name: optional name to give the task
|
||||
:param context: optional context to run the task in
|
||||
:return: a task handle
|
||||
|
||||
.. versionadded:: 4.14.0
|
||||
"""
|
||||
|
||||
@final
|
||||
def start_soon(
|
||||
self,
|
||||
func: Callable[[Unpack[PosArgsT]], Coroutine[Any, Any, T_co]],
|
||||
*args: Unpack[PosArgsT],
|
||||
name: object = None,
|
||||
) -> TaskHandle[T_co]:
|
||||
"""
|
||||
Start a new task in this task group.
|
||||
|
||||
:param func: a coroutine function
|
||||
:param args: positional arguments to call the function with
|
||||
:param name: name of the task, for the purposes of introspection and debugging
|
||||
:return: a task handle
|
||||
|
||||
.. versionadded:: 3.0
|
||||
.. versionchanged:: 4.14.0
|
||||
This method now returns a task handle.
|
||||
|
||||
"""
|
||||
final_name = get_callable_name(func, name)
|
||||
return self.create_task(call_for_coroutine(func, args), name=final_name)
|
||||
|
||||
@overload
|
||||
async def start(
|
||||
self,
|
||||
func: Callable[..., Coroutine[Any, Any, T_co]],
|
||||
*args: object,
|
||||
name: object = None,
|
||||
return_handle: Literal[False] = ...,
|
||||
) -> Any: ...
|
||||
|
||||
@overload
|
||||
async def start(
|
||||
self,
|
||||
func: Callable[..., Coroutine[Any, Any, T_co]],
|
||||
*args: object,
|
||||
name: object = None,
|
||||
return_handle: Literal[True],
|
||||
) -> TaskHandle[T_co, Any]: ...
|
||||
|
||||
@abstractmethod
|
||||
async def start(
|
||||
self,
|
||||
func: Callable[..., Coroutine[Any, Any, T_co]],
|
||||
*args: object,
|
||||
name: object = None,
|
||||
return_handle: Literal[False, True] = False,
|
||||
) -> Any:
|
||||
"""
|
||||
Start a new task and wait until it signals for readiness.
|
||||
|
||||
The target callable must accept a keyword argument ``task_status`` (of type
|
||||
:class:`TaskStatus`). Awaiting on this method will return whatever was passed to
|
||||
``task_status.started()`` (``None`` by default).
|
||||
|
||||
.. note:: The :class:`TaskStatus` class is generic, and the type argument should
|
||||
indicate the type of the value that will be passed to
|
||||
``task_status.started()``.
|
||||
|
||||
:param func: a coroutine function that accepts the ``task_status`` keyword
|
||||
argument
|
||||
:param args: positional arguments to call the function with
|
||||
:param name: an optional name for the task, for introspection and debugging
|
||||
:param return_handle: if ``True``, return a :class:`TaskHandle` which also
|
||||
contains the start value in ``start_value``
|
||||
:return: the value passed to ``task_status.started()``
|
||||
:raises RuntimeError: if the task finishes without calling
|
||||
``task_status.started()``
|
||||
|
||||
.. seealso:: :ref:`start_initialize`
|
||||
|
||||
.. versionadded:: 3.0
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
async def __aenter__(self) -> TaskGroup:
|
||||
"""Enter the task group context and allow starting new tasks."""
|
||||
|
||||
@abstractmethod
|
||||
async def __aexit__(
|
||||
self,
|
||||
exc_type: type[BaseException] | None,
|
||||
exc_val: BaseException | None,
|
||||
exc_tb: TracebackType | None,
|
||||
) -> bool:
|
||||
"""Exit the task group context waiting for all tasks to finish."""
|
||||
+79
@@ -0,0 +1,79 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
import types
|
||||
from abc import ABCMeta, abstractmethod
|
||||
from collections.abc import AsyncGenerator, Callable, Coroutine, Iterable
|
||||
from typing import Any, TypeVar
|
||||
|
||||
if sys.version_info >= (3, 11):
|
||||
from typing import Self
|
||||
else:
|
||||
from typing_extensions import Self
|
||||
|
||||
_T = TypeVar("_T")
|
||||
|
||||
|
||||
class TestRunner(metaclass=ABCMeta):
|
||||
"""
|
||||
Encapsulates a running event loop. Every call made through this object will use the
|
||||
same event loop.
|
||||
"""
|
||||
|
||||
def __enter__(self) -> Self:
|
||||
return self
|
||||
|
||||
@abstractmethod
|
||||
def __exit__(
|
||||
self,
|
||||
exc_type: type[BaseException] | None,
|
||||
exc_val: BaseException | None,
|
||||
exc_tb: types.TracebackType | None,
|
||||
) -> bool | None: ...
|
||||
|
||||
@abstractmethod
|
||||
def run_asyncgen_fixture(
|
||||
self,
|
||||
fixture_func: Callable[..., AsyncGenerator[_T, Any]],
|
||||
kwargs: dict[str, Any],
|
||||
) -> Iterable[_T]:
|
||||
"""
|
||||
Run an async generator fixture.
|
||||
|
||||
:param fixture_func: the fixture function
|
||||
:param kwargs: keyword arguments to call the fixture function with
|
||||
:return: an iterator yielding the value yielded from the async generator
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def run_fixture(
|
||||
self,
|
||||
fixture_func: Callable[..., Coroutine[Any, Any, _T]],
|
||||
kwargs: dict[str, Any],
|
||||
) -> _T:
|
||||
"""
|
||||
Run an async fixture.
|
||||
|
||||
:param fixture_func: the fixture function
|
||||
:param kwargs: keyword arguments to call the fixture function with
|
||||
:return: the return value of the fixture function
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def run_test(
|
||||
self, test_func: Callable[..., Coroutine[Any, Any, Any]], kwargs: dict[str, Any]
|
||||
) -> None:
|
||||
"""
|
||||
Run an async test function.
|
||||
|
||||
:param test_func: the test function
|
||||
:param kwargs: keyword arguments to call the test function with
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def is_running(self) -> bool:
|
||||
"""
|
||||
Check if the test runner is running.
|
||||
|
||||
:return: ``True`` if the coroutine is currently being run, ``False`` otherwise.
|
||||
"""
|
||||
+588
@@ -0,0 +1,588 @@
|
||||
from __future__ import annotations
|
||||
|
||||
__all__ = (
|
||||
"BlockingPortal",
|
||||
"BlockingPortalProvider",
|
||||
"check_cancelled",
|
||||
"run",
|
||||
"run_sync",
|
||||
"start_blocking_portal",
|
||||
)
|
||||
|
||||
import sys
|
||||
from collections.abc import Awaitable, Callable, Coroutine, Generator
|
||||
from concurrent.futures import Future
|
||||
from contextlib import (
|
||||
AbstractAsyncContextManager,
|
||||
AbstractContextManager,
|
||||
contextmanager,
|
||||
)
|
||||
from contextvars import Context
|
||||
from dataclasses import dataclass, field
|
||||
from functools import partial
|
||||
from inspect import isawaitable
|
||||
from threading import Lock, Thread, current_thread, get_ident
|
||||
from types import TracebackType
|
||||
from typing import (
|
||||
Any,
|
||||
Generic,
|
||||
TypeVar,
|
||||
cast,
|
||||
overload,
|
||||
)
|
||||
|
||||
from ._core._eventloop import (
|
||||
get_cancelled_exc_class,
|
||||
threadlocals,
|
||||
)
|
||||
from ._core._eventloop import run as run_eventloop
|
||||
from ._core._exceptions import NoEventLoopError
|
||||
from ._core._synchronization import Event
|
||||
from ._core._tasks import CancelScope, create_task_group
|
||||
from .abc._tasks import TaskStatus
|
||||
from .lowlevel import EventLoopToken, current_token
|
||||
|
||||
if sys.version_info >= (3, 11):
|
||||
from typing import Self, TypeVarTuple, Unpack
|
||||
else:
|
||||
from typing_extensions import Self, TypeVarTuple, Unpack
|
||||
|
||||
T_Retval = TypeVar("T_Retval")
|
||||
T_co = TypeVar("T_co", covariant=True)
|
||||
PosArgsT = TypeVarTuple("PosArgsT")
|
||||
|
||||
|
||||
def _token_or_error(token: EventLoopToken | None) -> EventLoopToken:
|
||||
if token is not None:
|
||||
return token
|
||||
|
||||
try:
|
||||
return threadlocals.current_token
|
||||
except AttributeError:
|
||||
raise NoEventLoopError(
|
||||
"Not running inside an AnyIO worker thread, and no event loop token was "
|
||||
"provided"
|
||||
) from None
|
||||
|
||||
|
||||
def run(
|
||||
func: Callable[[Unpack[PosArgsT]], Coroutine[Any, Any, T_co]],
|
||||
*args: Unpack[PosArgsT],
|
||||
token: EventLoopToken | None = None,
|
||||
) -> T_co:
|
||||
"""
|
||||
Call a coroutine function from a worker thread.
|
||||
|
||||
:param func: a coroutine function
|
||||
:param args: positional arguments for the callable
|
||||
:param token: an event loop token to use to get back to the event loop thread
|
||||
(required if calling this function from outside an AnyIO worker thread)
|
||||
:return: the return value of the coroutine function
|
||||
:raises MissingTokenError: if no token was provided and called from outside an
|
||||
AnyIO worker thread
|
||||
:raises RunFinishedError: if the event loop tied to ``token`` is no longer running
|
||||
|
||||
.. versionchanged:: 4.11.0
|
||||
Added the ``token`` parameter.
|
||||
|
||||
"""
|
||||
explicit_token = token is not None
|
||||
token = _token_or_error(token)
|
||||
return token.backend_class.run_async_from_thread(
|
||||
func, args, token=token.native_token if explicit_token else None
|
||||
)
|
||||
|
||||
|
||||
def run_sync(
|
||||
func: Callable[[Unpack[PosArgsT]], T_Retval],
|
||||
*args: Unpack[PosArgsT],
|
||||
token: EventLoopToken | None = None,
|
||||
) -> T_Retval:
|
||||
"""
|
||||
Call a function in the event loop thread from a worker thread.
|
||||
|
||||
:param func: a callable
|
||||
:param args: positional arguments for the callable
|
||||
:param token: an event loop token to use to get back to the event loop thread
|
||||
(required if calling this function from outside an AnyIO worker thread)
|
||||
:return: the return value of the callable
|
||||
:raises MissingTokenError: if no token was provided and called from outside an
|
||||
AnyIO worker thread
|
||||
:raises RunFinishedError: if the event loop tied to ``token`` is no longer running
|
||||
|
||||
.. versionchanged:: 4.11.0
|
||||
Added the ``token`` parameter.
|
||||
|
||||
"""
|
||||
explicit_token = token is not None
|
||||
token = _token_or_error(token)
|
||||
return token.backend_class.run_sync_from_thread(
|
||||
func, args, token=token.native_token if explicit_token else None
|
||||
)
|
||||
|
||||
|
||||
class _BlockingAsyncContextManager(AbstractContextManager, Generic[T_co]):
|
||||
_enter_future: Future[T_co]
|
||||
_exit_future: Future[bool | None]
|
||||
_exit_event: Event
|
||||
_exit_exc_info: tuple[
|
||||
type[BaseException] | None, BaseException | None, TracebackType | None
|
||||
] = (None, None, None)
|
||||
|
||||
def __init__(
|
||||
self, async_cm: AbstractAsyncContextManager[T_co], portal: BlockingPortal
|
||||
):
|
||||
self._async_cm = async_cm
|
||||
self._portal = portal
|
||||
|
||||
async def run_async_cm(self) -> bool | None:
|
||||
try:
|
||||
self._exit_event = Event()
|
||||
value = await self._async_cm.__aenter__()
|
||||
except BaseException as exc:
|
||||
self._enter_future.set_exception(exc)
|
||||
raise
|
||||
else:
|
||||
self._enter_future.set_result(value)
|
||||
|
||||
try:
|
||||
# Wait for the sync context manager to exit.
|
||||
# This next statement can raise `get_cancelled_exc_class()` if
|
||||
# something went wrong in a task group in this async context
|
||||
# manager.
|
||||
await self._exit_event.wait()
|
||||
finally:
|
||||
# In case of cancellation, it could be that we end up here before
|
||||
# `_BlockingAsyncContextManager.__exit__` is called, and an
|
||||
# `_exit_exc_info` has been set.
|
||||
result = await self._async_cm.__aexit__(*self._exit_exc_info)
|
||||
|
||||
return result
|
||||
|
||||
def __enter__(self) -> T_co:
|
||||
self._enter_future = Future()
|
||||
self._exit_future = self._portal.start_task_soon(self.run_async_cm)
|
||||
return self._enter_future.result()
|
||||
|
||||
def __exit__(
|
||||
self,
|
||||
exc_type: type[BaseException] | None,
|
||||
exc_value: BaseException | None,
|
||||
traceback: TracebackType | None,
|
||||
/,
|
||||
) -> bool | None:
|
||||
self._exit_exc_info = exc_type, exc_value, traceback
|
||||
self._portal.call(self._exit_event.set)
|
||||
return self._exit_future.result()
|
||||
|
||||
|
||||
class _BlockingPortalTaskStatus(TaskStatus):
|
||||
def __init__(self, future: Future):
|
||||
self._future = future
|
||||
|
||||
def started(self, value: object = None) -> None:
|
||||
self._future.set_result(value)
|
||||
|
||||
|
||||
class BlockingPortal:
|
||||
"""
|
||||
An object that lets external threads run code in an asynchronous event loop.
|
||||
|
||||
:raises NoEventLoopError: if no supported asynchronous event loop is running in the
|
||||
current thread
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._token = current_token()
|
||||
self._event_loop_thread_id: int | None = get_ident()
|
||||
self._stop_event = Event()
|
||||
self._task_group = create_task_group()
|
||||
|
||||
async def __aenter__(self) -> Self:
|
||||
await self._task_group.__aenter__()
|
||||
return self
|
||||
|
||||
async def __aexit__(
|
||||
self,
|
||||
exc_type: type[BaseException] | None,
|
||||
exc_val: BaseException | None,
|
||||
exc_tb: TracebackType | None,
|
||||
) -> bool:
|
||||
await self.stop()
|
||||
return await self._task_group.__aexit__(exc_type, exc_val, exc_tb)
|
||||
|
||||
def _check_running(self) -> None:
|
||||
if self._event_loop_thread_id is None:
|
||||
raise RuntimeError("This portal is not running")
|
||||
if self._event_loop_thread_id == get_ident():
|
||||
raise RuntimeError(
|
||||
"This method cannot be called from the event loop thread"
|
||||
)
|
||||
|
||||
async def sleep_until_stopped(self) -> None:
|
||||
"""Sleep until :meth:`stop` is called."""
|
||||
await self._stop_event.wait()
|
||||
|
||||
async def stop(self, cancel_remaining: bool = False) -> None:
|
||||
"""
|
||||
Signal the portal to shut down.
|
||||
|
||||
This marks the portal as no longer accepting new calls and exits from
|
||||
:meth:`sleep_until_stopped`.
|
||||
|
||||
:param cancel_remaining: ``True`` to cancel all the remaining tasks, ``False``
|
||||
to let them finish before returning
|
||||
|
||||
"""
|
||||
self._event_loop_thread_id = None
|
||||
self._stop_event.set()
|
||||
if cancel_remaining:
|
||||
self._task_group.cancel_scope.cancel("the blocking portal is shutting down")
|
||||
|
||||
async def _call_func(
|
||||
self,
|
||||
func: Callable[[Unpack[PosArgsT]], Awaitable[T_Retval] | T_Retval],
|
||||
args: tuple[Unpack[PosArgsT]],
|
||||
kwargs: dict[str, Any],
|
||||
future: Future[T_Retval],
|
||||
) -> None:
|
||||
event_loop_thread_id = self._event_loop_thread_id
|
||||
|
||||
def callback(f: Future[T_Retval]) -> None:
|
||||
if f.cancelled():
|
||||
if event_loop_thread_id == get_ident():
|
||||
scope.cancel("the future was cancelled")
|
||||
elif event_loop_thread_id is not None:
|
||||
run_sync(
|
||||
scope.cancel, "the future was cancelled", token=self._token
|
||||
)
|
||||
|
||||
try:
|
||||
retval_or_awaitable = func(*args, **kwargs)
|
||||
if isawaitable(retval_or_awaitable):
|
||||
with CancelScope() as scope:
|
||||
future.add_done_callback(callback)
|
||||
retval = await retval_or_awaitable
|
||||
else:
|
||||
retval = retval_or_awaitable
|
||||
except get_cancelled_exc_class():
|
||||
future.cancel()
|
||||
future.set_running_or_notify_cancel()
|
||||
except BaseException as exc:
|
||||
if not future.cancelled():
|
||||
future.set_exception(exc)
|
||||
|
||||
# Let base exceptions fall through
|
||||
if not isinstance(exc, Exception):
|
||||
raise
|
||||
else:
|
||||
if not future.cancelled():
|
||||
future.set_result(retval)
|
||||
finally:
|
||||
scope = None # type: ignore[assignment]
|
||||
|
||||
def _spawn_task_from_thread(
|
||||
self,
|
||||
func: Callable[[Unpack[PosArgsT]], Awaitable[T_Retval] | T_Retval],
|
||||
args: tuple[Unpack[PosArgsT]],
|
||||
kwargs: dict[str, Any],
|
||||
name: object,
|
||||
future: Future[T_Retval],
|
||||
) -> None:
|
||||
"""
|
||||
Spawn a new task using the given callable.
|
||||
|
||||
:param func: a callable
|
||||
:param args: positional arguments to be passed to the callable
|
||||
:param kwargs: keyword arguments to be passed to the callable
|
||||
:param name: name of the task (will be coerced to a string if not ``None``)
|
||||
:param future: a future that will resolve to the return value of the callable,
|
||||
or the exception raised during its execution
|
||||
|
||||
"""
|
||||
run_sync(
|
||||
partial(self._task_group.start_soon, name=name),
|
||||
self._call_func,
|
||||
func,
|
||||
args,
|
||||
kwargs,
|
||||
future,
|
||||
token=self._token,
|
||||
)
|
||||
|
||||
@overload
|
||||
def call(
|
||||
self,
|
||||
func: Callable[[Unpack[PosArgsT]], Awaitable[T_Retval]],
|
||||
*args: Unpack[PosArgsT],
|
||||
) -> T_Retval: ...
|
||||
|
||||
@overload
|
||||
def call(
|
||||
self, func: Callable[[Unpack[PosArgsT]], T_Retval], *args: Unpack[PosArgsT]
|
||||
) -> T_Retval: ...
|
||||
|
||||
def call(
|
||||
self,
|
||||
func: Callable[[Unpack[PosArgsT]], Awaitable[T_Retval] | T_Retval],
|
||||
*args: Unpack[PosArgsT],
|
||||
) -> T_Retval:
|
||||
"""
|
||||
Call the given function in the event loop thread.
|
||||
|
||||
If the callable returns a coroutine object, it is awaited on.
|
||||
|
||||
:param func: any callable
|
||||
:raises RuntimeError: if the portal is not running or if this method is called
|
||||
from within the event loop thread
|
||||
|
||||
"""
|
||||
return cast(T_Retval, self.start_task_soon(func, *args).result())
|
||||
|
||||
@overload
|
||||
def start_task_soon(
|
||||
self,
|
||||
func: Callable[[Unpack[PosArgsT]], Awaitable[T_Retval]],
|
||||
*args: Unpack[PosArgsT],
|
||||
name: object = None,
|
||||
) -> Future[T_Retval]: ...
|
||||
|
||||
@overload
|
||||
def start_task_soon(
|
||||
self,
|
||||
func: Callable[[Unpack[PosArgsT]], T_Retval],
|
||||
*args: Unpack[PosArgsT],
|
||||
name: object = None,
|
||||
) -> Future[T_Retval]: ...
|
||||
|
||||
def start_task_soon(
|
||||
self,
|
||||
func: Callable[[Unpack[PosArgsT]], Awaitable[T_Retval] | T_Retval],
|
||||
*args: Unpack[PosArgsT],
|
||||
name: object = None,
|
||||
) -> Future[T_Retval]:
|
||||
"""
|
||||
Start a task in the portal's task group.
|
||||
|
||||
The task will be run inside a cancel scope which can be cancelled by cancelling
|
||||
the returned future.
|
||||
|
||||
:param func: the target function
|
||||
:param args: positional arguments passed to ``func``
|
||||
:param name: name of the task (will be coerced to a string if not ``None``)
|
||||
:return: a future that resolves with the return value of the callable if the
|
||||
task completes successfully, or with the exception raised in the task
|
||||
:raises RuntimeError: if the portal is not running or if this method is called
|
||||
from within the event loop thread
|
||||
:rtype: concurrent.futures.Future[T_Retval]
|
||||
|
||||
.. versionadded:: 3.0
|
||||
|
||||
"""
|
||||
self._check_running()
|
||||
f: Future[T_Retval] = Future()
|
||||
self._spawn_task_from_thread(func, args, {}, name, f)
|
||||
return f
|
||||
|
||||
def start_task(
|
||||
self,
|
||||
func: Callable[..., Awaitable[T_Retval]],
|
||||
*args: object,
|
||||
name: object = None,
|
||||
) -> tuple[Future[T_Retval], Any]:
|
||||
"""
|
||||
Start a task in the portal's task group and wait until it signals for readiness.
|
||||
|
||||
This method works the same way as :meth:`.abc.TaskGroup.start`.
|
||||
|
||||
:param func: the target function
|
||||
:param args: positional arguments passed to ``func``
|
||||
:param name: name of the task (will be coerced to a string if not ``None``)
|
||||
:return: a tuple of (future, task_status_value) where the ``task_status_value``
|
||||
is the value passed to ``task_status.started()`` from within the target
|
||||
function
|
||||
:rtype: tuple[concurrent.futures.Future[T_Retval], Any]
|
||||
|
||||
.. versionadded:: 3.0
|
||||
|
||||
"""
|
||||
|
||||
def task_done(future: Future[T_Retval]) -> None:
|
||||
if not task_status_future.done():
|
||||
if future.cancelled():
|
||||
task_status_future.cancel()
|
||||
elif future.exception():
|
||||
task_status_future.set_exception(future.exception())
|
||||
else:
|
||||
exc = RuntimeError(
|
||||
"Task exited without calling task_status.started()"
|
||||
)
|
||||
task_status_future.set_exception(exc)
|
||||
|
||||
self._check_running()
|
||||
task_status_future: Future = Future()
|
||||
task_status = _BlockingPortalTaskStatus(task_status_future)
|
||||
f: Future = Future()
|
||||
f.add_done_callback(task_done)
|
||||
self._spawn_task_from_thread(func, args, {"task_status": task_status}, name, f)
|
||||
return f, task_status_future.result()
|
||||
|
||||
def wrap_async_context_manager(
|
||||
self, cm: AbstractAsyncContextManager[T_co]
|
||||
) -> AbstractContextManager[T_co]:
|
||||
"""
|
||||
Wrap an async context manager as a synchronous context manager via this portal.
|
||||
|
||||
Spawns a task that will call both ``__aenter__()`` and ``__aexit__()``, stopping
|
||||
in the middle until the synchronous context manager exits.
|
||||
|
||||
:param cm: an asynchronous context manager
|
||||
:return: a synchronous context manager
|
||||
|
||||
.. versionadded:: 2.1
|
||||
|
||||
"""
|
||||
return _BlockingAsyncContextManager(cm, self)
|
||||
|
||||
|
||||
@dataclass
|
||||
class BlockingPortalProvider:
|
||||
"""
|
||||
A manager for a blocking portal. Used as a context manager. The first thread to
|
||||
enter this context manager causes a blocking portal to be started with the specific
|
||||
parameters, and the last thread to exit causes the portal to be shut down. Thus,
|
||||
there will be exactly one blocking portal running in this context as long as at
|
||||
least one thread has entered this context manager.
|
||||
|
||||
The parameters are the same as for :func:`~anyio.run`.
|
||||
|
||||
:param backend: name of the backend
|
||||
:param backend_options: backend options
|
||||
|
||||
.. versionadded:: 4.4
|
||||
"""
|
||||
|
||||
backend: str = "asyncio"
|
||||
backend_options: dict[str, Any] | None = None
|
||||
_lock: Lock = field(init=False, default_factory=Lock)
|
||||
_leases: int = field(init=False, default=0)
|
||||
_portal: BlockingPortal = field(init=False)
|
||||
_portal_cm: AbstractContextManager[BlockingPortal] | None = field(
|
||||
init=False, default=None
|
||||
)
|
||||
|
||||
def __enter__(self) -> BlockingPortal:
|
||||
with self._lock:
|
||||
if self._portal_cm is None:
|
||||
self._portal_cm = start_blocking_portal(
|
||||
self.backend, self.backend_options
|
||||
)
|
||||
self._portal = self._portal_cm.__enter__()
|
||||
|
||||
self._leases += 1
|
||||
return self._portal
|
||||
|
||||
def __exit__(
|
||||
self,
|
||||
exc_type: type[BaseException] | None,
|
||||
exc_val: BaseException | None,
|
||||
exc_tb: TracebackType | None,
|
||||
) -> None:
|
||||
portal_cm: AbstractContextManager[BlockingPortal] | None = None
|
||||
with self._lock:
|
||||
assert self._portal_cm
|
||||
assert self._leases > 0
|
||||
self._leases -= 1
|
||||
if not self._leases:
|
||||
portal_cm = self._portal_cm
|
||||
self._portal_cm = None
|
||||
del self._portal
|
||||
|
||||
if portal_cm:
|
||||
portal_cm.__exit__(None, None, None)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def start_blocking_portal(
|
||||
backend: str = "asyncio",
|
||||
backend_options: dict[str, Any] | None = None,
|
||||
*,
|
||||
name: str | None = None,
|
||||
) -> Generator[BlockingPortal, Any, None]:
|
||||
"""
|
||||
Start a new event loop in a new thread and run a blocking portal in its main task.
|
||||
|
||||
The parameters are the same as for :func:`~anyio.run`.
|
||||
|
||||
:param backend: name of the backend
|
||||
:param backend_options: backend options
|
||||
:param name: name of the thread
|
||||
:return: a context manager that yields a blocking portal
|
||||
|
||||
.. versionchanged:: 3.0
|
||||
Usage as a context manager is now required.
|
||||
|
||||
"""
|
||||
|
||||
async def run_portal() -> None:
|
||||
async with BlockingPortal() as portal_:
|
||||
if name is None:
|
||||
current_thread().name = f"{backend}-portal-{id(portal_):x}"
|
||||
|
||||
future.set_result(portal_)
|
||||
await portal_.sleep_until_stopped()
|
||||
|
||||
def run_blocking_portal() -> None:
|
||||
if future.set_running_or_notify_cancel():
|
||||
try:
|
||||
run_eventloop(
|
||||
run_portal, backend=backend, backend_options=backend_options
|
||||
)
|
||||
except BaseException as exc:
|
||||
if not future.done():
|
||||
future.set_exception(exc)
|
||||
|
||||
future: Future[BlockingPortal] = Future()
|
||||
kwargs: dict[str, Any] = {}
|
||||
if sys.version_info >= (3, 14):
|
||||
kwargs["context"] = Context()
|
||||
|
||||
thread = Thread(target=run_blocking_portal, daemon=True, name=name, **kwargs)
|
||||
thread.start()
|
||||
try:
|
||||
cancel_remaining_tasks = False
|
||||
portal = future.result()
|
||||
try:
|
||||
yield portal
|
||||
except BaseException:
|
||||
cancel_remaining_tasks = True
|
||||
raise
|
||||
finally:
|
||||
try:
|
||||
portal.call(portal.stop, cancel_remaining_tasks)
|
||||
except RuntimeError:
|
||||
pass
|
||||
finally:
|
||||
thread.join()
|
||||
|
||||
|
||||
def check_cancelled() -> None:
|
||||
"""
|
||||
Check if the cancel scope of the host task's running the current worker thread has
|
||||
been cancelled.
|
||||
|
||||
If the host task's current cancel scope has indeed been cancelled, the
|
||||
backend-specific cancellation exception will be raised.
|
||||
|
||||
:raises RuntimeError: if the current thread was not spawned by
|
||||
:func:`.to_thread.run_sync`
|
||||
|
||||
"""
|
||||
try:
|
||||
token: EventLoopToken = threadlocals.current_token
|
||||
except AttributeError:
|
||||
raise NoEventLoopError(
|
||||
"This function can only be called inside an AnyIO worker thread"
|
||||
) from None
|
||||
|
||||
token.backend_class.check_cancelled()
|
||||
+400
@@ -0,0 +1,400 @@
|
||||
from __future__ import annotations
|
||||
|
||||
__all__ = (
|
||||
"AsyncCacheInfo",
|
||||
"AsyncCacheParameters",
|
||||
"AsyncLRUCacheWrapper",
|
||||
"cache",
|
||||
"lru_cache",
|
||||
"reduce",
|
||||
)
|
||||
|
||||
import functools
|
||||
from collections import OrderedDict
|
||||
from collections.abc import (
|
||||
AsyncIterable,
|
||||
Awaitable,
|
||||
Callable,
|
||||
Coroutine,
|
||||
Hashable,
|
||||
Iterable,
|
||||
)
|
||||
from functools import update_wrapper
|
||||
from inspect import iscoroutinefunction
|
||||
from typing import (
|
||||
Any,
|
||||
Generic,
|
||||
NamedTuple,
|
||||
ParamSpec,
|
||||
TypedDict,
|
||||
TypeVar,
|
||||
cast,
|
||||
final,
|
||||
overload,
|
||||
)
|
||||
from weakref import WeakKeyDictionary
|
||||
|
||||
from ._core._eventloop import current_time
|
||||
from ._core._synchronization import Lock
|
||||
from .lowlevel import RunVar, checkpoint
|
||||
|
||||
T = TypeVar("T")
|
||||
S = TypeVar("S")
|
||||
P = ParamSpec("P")
|
||||
lru_cache_items: RunVar[
|
||||
WeakKeyDictionary[
|
||||
AsyncLRUCacheWrapper[Any, Any],
|
||||
OrderedDict[
|
||||
Hashable,
|
||||
tuple[_InitialMissingType, Lock, float | None]
|
||||
| tuple[Any, None, float | None],
|
||||
],
|
||||
]
|
||||
] = RunVar("lru_cache_items")
|
||||
|
||||
|
||||
class _InitialMissingType:
|
||||
pass
|
||||
|
||||
|
||||
initial_missing: _InitialMissingType = _InitialMissingType()
|
||||
|
||||
|
||||
class AsyncCacheInfo(NamedTuple):
|
||||
hits: int
|
||||
misses: int
|
||||
maxsize: int | None
|
||||
currsize: int
|
||||
ttl: int | None
|
||||
|
||||
|
||||
class AsyncCacheParameters(TypedDict):
|
||||
maxsize: int | None
|
||||
typed: bool
|
||||
always_checkpoint: bool
|
||||
ttl: int | None
|
||||
|
||||
|
||||
class _LRUMethodWrapper(Generic[T]):
|
||||
def __init__(self, wrapper: AsyncLRUCacheWrapper[..., T], instance: object):
|
||||
self.__wrapper = wrapper
|
||||
self.__instance = instance
|
||||
|
||||
def cache_info(self) -> AsyncCacheInfo:
|
||||
return self.__wrapper.cache_info()
|
||||
|
||||
def cache_parameters(self) -> AsyncCacheParameters:
|
||||
return self.__wrapper.cache_parameters()
|
||||
|
||||
def cache_clear(self) -> None:
|
||||
self.__wrapper.cache_clear()
|
||||
|
||||
async def __call__(self, *args: Any, **kwargs: Any) -> T:
|
||||
if self.__instance is None:
|
||||
return await self.__wrapper(*args, **kwargs)
|
||||
|
||||
return await self.__wrapper(self.__instance, *args, **kwargs)
|
||||
|
||||
|
||||
@final
|
||||
class AsyncLRUCacheWrapper(Generic[P, T]):
|
||||
def __init__(
|
||||
self,
|
||||
func: Callable[P, Awaitable[T]],
|
||||
maxsize: int | None,
|
||||
typed: bool,
|
||||
always_checkpoint: bool,
|
||||
ttl: int | None,
|
||||
):
|
||||
self.__wrapped__ = func
|
||||
self._hits: int = 0
|
||||
self._misses: int = 0
|
||||
self._maxsize = max(maxsize, 0) if maxsize is not None else None
|
||||
self._currsize: int = 0
|
||||
self._typed = typed
|
||||
self._always_checkpoint = always_checkpoint
|
||||
self._ttl = ttl
|
||||
update_wrapper(self, func)
|
||||
|
||||
def cache_info(self) -> AsyncCacheInfo:
|
||||
return AsyncCacheInfo(
|
||||
self._hits, self._misses, self._maxsize, self._currsize, self._ttl
|
||||
)
|
||||
|
||||
def cache_parameters(self) -> AsyncCacheParameters:
|
||||
return {
|
||||
"maxsize": self._maxsize,
|
||||
"typed": self._typed,
|
||||
"always_checkpoint": self._always_checkpoint,
|
||||
"ttl": self._ttl,
|
||||
}
|
||||
|
||||
def cache_clear(self) -> None:
|
||||
if cache := lru_cache_items.get(None):
|
||||
cache.pop(self, None)
|
||||
self._hits = self._misses = self._currsize = 0
|
||||
|
||||
async def __call__(self, *args: P.args, **kwargs: P.kwargs) -> T:
|
||||
# Easy case first: if maxsize == 0, no caching is done
|
||||
if self._maxsize == 0:
|
||||
value = await self.__wrapped__(*args, **kwargs)
|
||||
self._misses += 1
|
||||
return value
|
||||
|
||||
# The key is constructed as a flat tuple to avoid memory overhead
|
||||
key: tuple[Any, ...] = args
|
||||
if kwargs:
|
||||
# initial_missing is used as a separator
|
||||
key += (initial_missing,) + sum(kwargs.items(), ())
|
||||
|
||||
if self._typed:
|
||||
key += tuple(type(arg) for arg in args)
|
||||
if kwargs:
|
||||
key += (initial_missing,) + tuple(type(val) for val in kwargs.values())
|
||||
|
||||
try:
|
||||
cache = lru_cache_items.get()
|
||||
except LookupError:
|
||||
cache = WeakKeyDictionary()
|
||||
lru_cache_items.set(cache)
|
||||
|
||||
try:
|
||||
cache_entry = cache[self]
|
||||
except KeyError:
|
||||
cache_entry = cache[self] = OrderedDict()
|
||||
|
||||
cached_value: T | _InitialMissingType
|
||||
try:
|
||||
cached_value, lock, expires_at = cache_entry[key]
|
||||
except KeyError:
|
||||
# We're the first task to call this function
|
||||
cached_value, lock, expires_at = (
|
||||
initial_missing,
|
||||
Lock(fast_acquire=not self._always_checkpoint),
|
||||
None,
|
||||
)
|
||||
cache_entry[key] = cached_value, lock, expires_at
|
||||
|
||||
if lock is None:
|
||||
if expires_at is not None and current_time() >= expires_at:
|
||||
self._currsize -= 1
|
||||
cached_value, lock, expires_at = (
|
||||
initial_missing,
|
||||
Lock(fast_acquire=not self._always_checkpoint),
|
||||
None,
|
||||
)
|
||||
cache_entry[key] = cached_value, lock, expires_at
|
||||
else:
|
||||
# The value was already cached
|
||||
self._hits += 1
|
||||
cache_entry.move_to_end(key)
|
||||
if self._always_checkpoint:
|
||||
await checkpoint()
|
||||
|
||||
return cast(T, cached_value)
|
||||
|
||||
async with lock:
|
||||
# Check if another task filled the cache while we acquired the lock
|
||||
if (cached_value := cache_entry[key][0]) is initial_missing:
|
||||
self._misses += 1
|
||||
if self._maxsize is not None and self._currsize >= self._maxsize:
|
||||
cache_entry.popitem(last=False)
|
||||
else:
|
||||
self._currsize += 1
|
||||
|
||||
value = await self.__wrapped__(*args, **kwargs)
|
||||
expires_at = (
|
||||
current_time() + self._ttl if self._ttl is not None else None
|
||||
)
|
||||
cache_entry[key] = value, None, expires_at
|
||||
else:
|
||||
# Another task filled the cache while we were waiting for the lock
|
||||
self._hits += 1
|
||||
cache_entry.move_to_end(key)
|
||||
value = cast(T, cached_value)
|
||||
|
||||
return value
|
||||
|
||||
def __get__(
|
||||
self, instance: object, owner: type | None = None
|
||||
) -> _LRUMethodWrapper[T]:
|
||||
wrapper = _LRUMethodWrapper(self, instance)
|
||||
update_wrapper(wrapper, self.__wrapped__)
|
||||
return wrapper
|
||||
|
||||
|
||||
class _LRUCacheWrapper:
|
||||
def __init__(
|
||||
self, maxsize: int | None, typed: bool, always_checkpoint: bool, ttl: int | None
|
||||
):
|
||||
self._maxsize = maxsize
|
||||
self._typed = typed
|
||||
self._always_checkpoint = always_checkpoint
|
||||
self._ttl = ttl
|
||||
|
||||
@overload
|
||||
def __call__( # type: ignore[overload-overlap]
|
||||
self, func: Callable[P, Coroutine[Any, Any, T]], /
|
||||
) -> AsyncLRUCacheWrapper[P, T]: ...
|
||||
|
||||
@overload
|
||||
def __call__(
|
||||
self, func: Callable[..., T], /
|
||||
) -> functools._lru_cache_wrapper[T]: ...
|
||||
|
||||
def __call__(
|
||||
self, f: Callable[P, Coroutine[Any, Any, T]] | Callable[..., T], /
|
||||
) -> AsyncLRUCacheWrapper[P, T] | functools._lru_cache_wrapper[T]:
|
||||
if iscoroutinefunction(f):
|
||||
return AsyncLRUCacheWrapper(
|
||||
f, self._maxsize, self._typed, self._always_checkpoint, self._ttl
|
||||
)
|
||||
|
||||
return functools.lru_cache(maxsize=self._maxsize, typed=self._typed)(f) # type: ignore[arg-type]
|
||||
|
||||
|
||||
@overload
|
||||
def cache( # type: ignore[overload-overlap]
|
||||
func: Callable[P, Coroutine[Any, Any, T]], /
|
||||
) -> AsyncLRUCacheWrapper[P, T]: ...
|
||||
|
||||
|
||||
@overload
|
||||
def cache(func: Callable[..., T], /) -> functools._lru_cache_wrapper[T]: ...
|
||||
|
||||
|
||||
def cache(func: Callable[..., Any] | Callable[P, Coroutine[Any, Any, Any]], /) -> Any:
|
||||
"""
|
||||
A convenient shortcut for :func:`lru_cache` with ``maxsize=None``.
|
||||
|
||||
This is the asynchronous equivalent to :func:`functools.cache`.
|
||||
|
||||
"""
|
||||
return lru_cache(maxsize=None)(func)
|
||||
|
||||
|
||||
@overload
|
||||
def lru_cache(
|
||||
*,
|
||||
maxsize: int | None = ...,
|
||||
typed: bool = ...,
|
||||
always_checkpoint: bool = ...,
|
||||
ttl: int | None = ...,
|
||||
) -> _LRUCacheWrapper: ...
|
||||
|
||||
|
||||
@overload
|
||||
def lru_cache( # type: ignore[overload-overlap]
|
||||
func: Callable[P, Coroutine[Any, Any, T]], /
|
||||
) -> AsyncLRUCacheWrapper[P, T]: ...
|
||||
|
||||
|
||||
@overload
|
||||
def lru_cache(func: Callable[..., T], /) -> functools._lru_cache_wrapper[T]: ...
|
||||
|
||||
|
||||
def lru_cache(
|
||||
func: Callable[..., Coroutine[Any, Any, Any]] | Callable[..., Any] | None = None,
|
||||
/,
|
||||
*,
|
||||
maxsize: int | None = 128,
|
||||
typed: bool = False,
|
||||
always_checkpoint: bool = False,
|
||||
ttl: int | None = None,
|
||||
) -> Any:
|
||||
"""
|
||||
An asynchronous version of :func:`functools.lru_cache`.
|
||||
|
||||
If a synchronous function is passed, the standard library
|
||||
:func:`functools.lru_cache` is applied instead.
|
||||
|
||||
:param always_checkpoint: if ``True``, every call to the cached function will be
|
||||
guaranteed to yield control to the event loop at least once
|
||||
:param ttl: time in seconds after which to invalidate cache entries
|
||||
|
||||
.. note:: Caches and locks are managed on a per-event loop basis.
|
||||
|
||||
"""
|
||||
if func is None:
|
||||
return _LRUCacheWrapper(maxsize, typed, always_checkpoint, ttl)
|
||||
|
||||
if not callable(func):
|
||||
raise TypeError("the first argument must be callable")
|
||||
|
||||
return _LRUCacheWrapper(maxsize, typed, always_checkpoint, ttl)(func)
|
||||
|
||||
|
||||
@overload
|
||||
async def reduce(
|
||||
function: Callable[[T, S], Awaitable[T]],
|
||||
iterable: Iterable[S] | AsyncIterable[S],
|
||||
/,
|
||||
initial: T,
|
||||
) -> T: ...
|
||||
|
||||
|
||||
@overload
|
||||
async def reduce(
|
||||
function: Callable[[T, T], Awaitable[T]],
|
||||
iterable: Iterable[T] | AsyncIterable[T],
|
||||
/,
|
||||
) -> T: ...
|
||||
|
||||
|
||||
async def reduce( # type: ignore[misc]
|
||||
function: Callable[[T, T], Awaitable[T]] | Callable[[T, S], Awaitable[T]],
|
||||
iterable: Iterable[T] | Iterable[S] | AsyncIterable[T] | AsyncIterable[S],
|
||||
/,
|
||||
initial: T | _InitialMissingType = initial_missing,
|
||||
) -> T:
|
||||
"""
|
||||
Asynchronous version of :func:`functools.reduce`.
|
||||
|
||||
:param function: a coroutine function that takes two arguments: the accumulated
|
||||
value and the next element from the iterable
|
||||
:param iterable: an iterable or async iterable
|
||||
:param initial: the initial value (if missing, the first element of the iterable is
|
||||
used as the initial value)
|
||||
|
||||
"""
|
||||
element: Any
|
||||
function_called = False
|
||||
if isinstance(iterable, AsyncIterable):
|
||||
async_it = iterable.__aiter__()
|
||||
if initial is initial_missing:
|
||||
try:
|
||||
value = cast(T, await async_it.__anext__())
|
||||
except StopAsyncIteration:
|
||||
raise TypeError(
|
||||
"reduce() of empty sequence with no initial value"
|
||||
) from None
|
||||
else:
|
||||
value = cast(T, initial)
|
||||
|
||||
async for element in async_it:
|
||||
value = await function(value, element)
|
||||
function_called = True
|
||||
elif isinstance(iterable, Iterable):
|
||||
it = iter(iterable)
|
||||
if initial is initial_missing:
|
||||
try:
|
||||
value = cast(T, next(it))
|
||||
except StopIteration:
|
||||
raise TypeError(
|
||||
"reduce() of empty sequence with no initial value"
|
||||
) from None
|
||||
else:
|
||||
value = cast(T, initial)
|
||||
|
||||
for element in it:
|
||||
value = await function(value, element)
|
||||
function_called = True
|
||||
else:
|
||||
raise TypeError("reduce() argument 2 must be an iterable or async iterable")
|
||||
|
||||
# Make sure there is at least one checkpoint, even if an empty iterable and an
|
||||
# initial value were given
|
||||
if not function_called:
|
||||
await checkpoint()
|
||||
|
||||
return value
|
||||
+629
@@ -0,0 +1,629 @@
|
||||
from __future__ import annotations
|
||||
|
||||
__all__ = (
|
||||
"Chain",
|
||||
"accumulate",
|
||||
"batched",
|
||||
"combinations",
|
||||
"combinations_with_replacement",
|
||||
"compress",
|
||||
"count",
|
||||
"cycle",
|
||||
"dropwhile",
|
||||
"filterfalse",
|
||||
"groupby",
|
||||
"islice",
|
||||
"pairwise",
|
||||
"permutations",
|
||||
"product",
|
||||
"repeat",
|
||||
"starmap",
|
||||
"takewhile",
|
||||
"tee",
|
||||
"zip_longest",
|
||||
)
|
||||
|
||||
import itertools
|
||||
import operator
|
||||
import sys
|
||||
from collections.abc import (
|
||||
AsyncGenerator,
|
||||
AsyncIterable,
|
||||
AsyncIterator,
|
||||
Awaitable,
|
||||
Callable,
|
||||
Iterable,
|
||||
Iterator,
|
||||
)
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Generic, TypeVar, cast, overload
|
||||
|
||||
from ._core._synchronization import Lock
|
||||
from ._core._tasks import CancelScope
|
||||
from .lowlevel import cancel_shielded_checkpoint, checkpoint, checkpoint_if_cancelled
|
||||
|
||||
if sys.version_info < (3, 15):
|
||||
from typing_extensions import sentinel
|
||||
|
||||
T = TypeVar("T")
|
||||
R = TypeVar("R")
|
||||
_tee_end = sentinel("_tee_end")
|
||||
|
||||
|
||||
@dataclass(eq=False)
|
||||
class _IterableAsyncIterator(AsyncIterator[T]):
|
||||
iterator: Iterator[T]
|
||||
|
||||
async def __anext__(self) -> T:
|
||||
await checkpoint_if_cancelled()
|
||||
try:
|
||||
result = next(self.iterator)
|
||||
except StopIteration:
|
||||
await cancel_shielded_checkpoint()
|
||||
raise StopAsyncIteration from None
|
||||
|
||||
await cancel_shielded_checkpoint()
|
||||
return result
|
||||
|
||||
|
||||
def _iterate(iterable: Iterable[T] | AsyncIterable[T]) -> AsyncIterator[T]:
|
||||
if isinstance(iterable, AsyncIterator):
|
||||
return iterable
|
||||
|
||||
if isinstance(iterable, AsyncIterable):
|
||||
return iterable.__aiter__()
|
||||
|
||||
return _IterableAsyncIterator(iter(iterable))
|
||||
|
||||
|
||||
@dataclass(eq=False)
|
||||
class _TeeLink(Generic[T]):
|
||||
value: object | None = None
|
||||
next: _TeeLink[T] | None = None
|
||||
filled: bool = False
|
||||
|
||||
|
||||
@dataclass(eq=False)
|
||||
class _TeeState(Generic[T]):
|
||||
iterator: AsyncIterator[T]
|
||||
lock: Lock = field(default_factory=Lock)
|
||||
|
||||
async def fill(self, link: _TeeLink[T]) -> bool:
|
||||
if link.filled:
|
||||
return False
|
||||
|
||||
async with self.lock:
|
||||
if link.filled:
|
||||
return True
|
||||
|
||||
link.value = await anext(self.iterator, _tee_end)
|
||||
if link.value is not _tee_end:
|
||||
link.next = _TeeLink()
|
||||
|
||||
link.filled = True
|
||||
return True
|
||||
|
||||
|
||||
class _TeeAsyncIterator(AsyncIterator[T]):
|
||||
_state: _TeeState[T]
|
||||
_link: _TeeLink[T]
|
||||
_element_yielded: bool
|
||||
|
||||
def __init__(
|
||||
self, iterable: Iterable[T] | AsyncIterable[T] | _TeeAsyncIterator[T]
|
||||
) -> None:
|
||||
if isinstance(iterable, _TeeAsyncIterator):
|
||||
self._state = iterable._state
|
||||
self._link = iterable._link
|
||||
else:
|
||||
self._state = _TeeState(_iterate(iterable))
|
||||
self._link = _TeeLink()
|
||||
|
||||
self._element_yielded = False
|
||||
|
||||
async def __anext__(self) -> T:
|
||||
had_yieldpoint = await self._state.fill(self._link)
|
||||
if self._link.value is _tee_end:
|
||||
if not self._element_yielded:
|
||||
await checkpoint()
|
||||
|
||||
raise StopAsyncIteration
|
||||
|
||||
if not had_yieldpoint:
|
||||
await checkpoint_if_cancelled()
|
||||
|
||||
self._element_yielded = True
|
||||
value = cast(T, self._link.value)
|
||||
next_link = self._link.next
|
||||
assert next_link is not None
|
||||
self._link = next_link
|
||||
if not had_yieldpoint:
|
||||
await cancel_shielded_checkpoint()
|
||||
|
||||
return value
|
||||
|
||||
|
||||
async def _operator_add(x: T, y: T) -> T:
|
||||
return operator.add(x, y)
|
||||
|
||||
|
||||
async def accumulate(
|
||||
iterable: Iterable[T] | AsyncIterable[T],
|
||||
function: Callable[[T, T], Awaitable[T]] = _operator_add,
|
||||
*,
|
||||
initial: T | None = None,
|
||||
) -> AsyncGenerator[T, None]:
|
||||
iterator = _iterate(iterable)
|
||||
if initial is None:
|
||||
try:
|
||||
total = await anext(iterator)
|
||||
except StopAsyncIteration:
|
||||
await checkpoint()
|
||||
return
|
||||
else:
|
||||
await checkpoint_if_cancelled()
|
||||
total = initial
|
||||
await cancel_shielded_checkpoint()
|
||||
|
||||
yield total
|
||||
|
||||
async for element in iterator:
|
||||
total = await function(total, element)
|
||||
yield total
|
||||
|
||||
|
||||
async def batched(
|
||||
iterable: Iterable[T] | AsyncIterable[T], n: int, *, strict: bool = False
|
||||
) -> AsyncGenerator[tuple[T, ...], None]:
|
||||
if n < 1:
|
||||
raise ValueError("n must be at least one")
|
||||
|
||||
iterator = _iterate(iterable)
|
||||
|
||||
while True:
|
||||
batch: list[T] = []
|
||||
for _ in range(n):
|
||||
try:
|
||||
batch.append(await anext(iterator))
|
||||
except StopAsyncIteration:
|
||||
if not batch:
|
||||
await checkpoint()
|
||||
return
|
||||
if strict:
|
||||
raise ValueError("batched(): incomplete batch") from None
|
||||
|
||||
yield tuple(batch)
|
||||
return
|
||||
|
||||
yield tuple(batch)
|
||||
|
||||
|
||||
class Chain:
|
||||
def __call__(
|
||||
self, *iterables: Iterable[T] | AsyncIterable[T]
|
||||
) -> AsyncGenerator[T, None]:
|
||||
return self.from_iterable(iterables)
|
||||
|
||||
async def from_iterable(
|
||||
self,
|
||||
iterables: (
|
||||
Iterable[Iterable[T] | AsyncIterable[T]]
|
||||
| AsyncIterable[Iterable[T] | AsyncIterable[T]]
|
||||
),
|
||||
) -> AsyncGenerator[T, None]:
|
||||
element_yielded = False
|
||||
outer_iter = _iterate(iterables)
|
||||
|
||||
try:
|
||||
async for iterable in outer_iter:
|
||||
async for element in _iterate(iterable):
|
||||
element_yielded = True
|
||||
yield element
|
||||
finally:
|
||||
aclose = getattr(outer_iter, "aclose", None)
|
||||
if aclose is not None:
|
||||
with CancelScope(shield=True):
|
||||
await aclose()
|
||||
|
||||
if not element_yielded:
|
||||
await checkpoint()
|
||||
|
||||
|
||||
chain: Chain = Chain()
|
||||
|
||||
|
||||
async def combinations(
|
||||
iterable: Iterable[T] | AsyncIterable[T], r: int
|
||||
) -> AsyncGenerator[tuple[T, ...], None]:
|
||||
pool: list[T] = [element async for element in _iterate(iterable)]
|
||||
async for combination in _iterate(itertools.combinations(pool, r)):
|
||||
yield combination
|
||||
|
||||
|
||||
async def combinations_with_replacement(
|
||||
iterable: Iterable[T] | AsyncIterable[T], r: int
|
||||
) -> AsyncGenerator[tuple[T, ...], None]:
|
||||
pool: list[T] = [element async for element in _iterate(iterable)]
|
||||
async for combination in _iterate(itertools.combinations_with_replacement(pool, r)):
|
||||
yield combination
|
||||
|
||||
|
||||
async def compress(
|
||||
data: Iterable[T] | AsyncIterable[T],
|
||||
selectors: Iterable[object] | AsyncIterable[object],
|
||||
) -> AsyncGenerator[T, None]:
|
||||
data_iterator = _iterate(data)
|
||||
selector_iterator = _iterate(selectors)
|
||||
element_yielded = False
|
||||
|
||||
while True:
|
||||
try:
|
||||
datum = await anext(data_iterator)
|
||||
selector = await anext(selector_iterator)
|
||||
except StopAsyncIteration:
|
||||
if not element_yielded:
|
||||
await checkpoint()
|
||||
|
||||
return
|
||||
|
||||
if selector:
|
||||
element_yielded = True
|
||||
yield datum
|
||||
|
||||
|
||||
async def count(start: int = 0, step: int = 1) -> AsyncGenerator[int, None]:
|
||||
n = start
|
||||
while True:
|
||||
await checkpoint_if_cancelled()
|
||||
value = n
|
||||
n += step
|
||||
await cancel_shielded_checkpoint()
|
||||
yield value
|
||||
|
||||
|
||||
async def cycle(
|
||||
iterable: Iterable[T] | AsyncIterable[T],
|
||||
) -> AsyncGenerator[T, None]:
|
||||
saved: list[T] = []
|
||||
async for element in _iterate(iterable):
|
||||
saved.append(element)
|
||||
yield element
|
||||
|
||||
if not saved:
|
||||
await checkpoint()
|
||||
return
|
||||
|
||||
while True:
|
||||
for element in saved:
|
||||
await checkpoint()
|
||||
yield element
|
||||
|
||||
|
||||
async def dropwhile(
|
||||
predicate: Callable[[T], Awaitable[object]],
|
||||
iterable: Iterable[T] | AsyncIterable[T],
|
||||
) -> AsyncGenerator[T, None]:
|
||||
element_yielded = False
|
||||
dropping = True
|
||||
|
||||
async for element in _iterate(iterable):
|
||||
if dropping and await predicate(element):
|
||||
continue
|
||||
|
||||
dropping = False
|
||||
element_yielded = True
|
||||
yield element
|
||||
|
||||
if not element_yielded:
|
||||
await checkpoint()
|
||||
|
||||
|
||||
async def filterfalse(
|
||||
predicate: Callable[[T], Awaitable[object]],
|
||||
iterable: Iterable[T] | AsyncIterable[T],
|
||||
) -> AsyncGenerator[T, None]:
|
||||
element_yielded = False
|
||||
|
||||
async for element in _iterate(iterable):
|
||||
if not await predicate(element):
|
||||
element_yielded = True
|
||||
yield element
|
||||
|
||||
if not element_yielded:
|
||||
await checkpoint()
|
||||
|
||||
|
||||
@overload
|
||||
def groupby(
|
||||
iterable: Iterable[T] | AsyncIterable[T],
|
||||
) -> AsyncGenerator[tuple[T, list[T]], None]: ...
|
||||
|
||||
|
||||
@overload
|
||||
def groupby(
|
||||
iterable: Iterable[T] | AsyncIterable[T],
|
||||
key: Callable[[T], Awaitable[R]],
|
||||
) -> AsyncGenerator[tuple[R, list[T]], None]: ...
|
||||
|
||||
|
||||
async def groupby(
|
||||
iterable: Iterable[T] | AsyncIterable[T],
|
||||
key: Callable[[T], Awaitable[object]] | None = None,
|
||||
) -> AsyncGenerator[tuple[object, list[T]], None]:
|
||||
iterator = _iterate(iterable)
|
||||
try:
|
||||
element = await anext(iterator)
|
||||
except StopAsyncIteration:
|
||||
await checkpoint()
|
||||
return
|
||||
|
||||
group_key = element if key is None else await key(element)
|
||||
values = [element]
|
||||
|
||||
async for element in iterator:
|
||||
next_key = element if key is None else await key(element)
|
||||
if next_key != group_key:
|
||||
completed_group = group_key, values
|
||||
group_key = next_key
|
||||
values = [element]
|
||||
yield completed_group
|
||||
else:
|
||||
values.append(element)
|
||||
|
||||
yield group_key, values
|
||||
|
||||
|
||||
@overload
|
||||
def islice(
|
||||
iterable: Iterable[T] | AsyncIterable[T],
|
||||
stop: int | None,
|
||||
/,
|
||||
) -> AsyncGenerator[T, None]: ...
|
||||
|
||||
|
||||
@overload
|
||||
def islice(
|
||||
iterable: Iterable[T] | AsyncIterable[T],
|
||||
start: int | None,
|
||||
stop: int | None,
|
||||
step: int | None = 1,
|
||||
/,
|
||||
) -> AsyncGenerator[T, None]: ...
|
||||
|
||||
|
||||
async def islice(
|
||||
iterable: Iterable[T] | AsyncIterable[T],
|
||||
*args: int | None,
|
||||
) -> AsyncGenerator[T, None]:
|
||||
if not args:
|
||||
raise TypeError("islice expected at least 2 arguments, got 1")
|
||||
if len(args) > 3:
|
||||
raise TypeError(f"islice expected at most 4 arguments, got {len(args) + 1}")
|
||||
|
||||
slice_args = slice(*args)
|
||||
|
||||
start_message = (
|
||||
"Indices for islice() must be None or an integer: 0 <= x <= sys.maxsize."
|
||||
)
|
||||
stop_message = (
|
||||
"Stop argument for islice() must be None or an integer: 0 <= x <= sys.maxsize."
|
||||
)
|
||||
step_message = "Step for islice() must be a positive integer or None."
|
||||
|
||||
def normalize_index(value: object, message: str) -> int:
|
||||
try:
|
||||
index = operator.index(cast(Any, value))
|
||||
except TypeError:
|
||||
raise ValueError(message) from None
|
||||
|
||||
if index < 0 or index > sys.maxsize:
|
||||
raise ValueError(message)
|
||||
|
||||
return index
|
||||
|
||||
start = (
|
||||
0
|
||||
if slice_args.start is None
|
||||
else normalize_index(slice_args.start, start_message)
|
||||
)
|
||||
stop = (
|
||||
None
|
||||
if slice_args.stop is None
|
||||
else normalize_index(slice_args.stop, stop_message)
|
||||
)
|
||||
step = (
|
||||
1 if slice_args.step is None else normalize_index(slice_args.step, step_message)
|
||||
)
|
||||
|
||||
if step <= 0:
|
||||
raise ValueError(step_message)
|
||||
|
||||
if stop == 0 or start == stop:
|
||||
await checkpoint()
|
||||
return
|
||||
|
||||
iterator = _iterate(iterable)
|
||||
index = 0
|
||||
element_yielded = False
|
||||
|
||||
while stop is None or index < stop:
|
||||
try:
|
||||
element = await anext(iterator)
|
||||
except StopAsyncIteration:
|
||||
if not element_yielded:
|
||||
await checkpoint()
|
||||
|
||||
return
|
||||
|
||||
if index >= start and (index - start) % step == 0:
|
||||
index += 1
|
||||
element_yielded = True
|
||||
yield element
|
||||
else:
|
||||
index += 1
|
||||
|
||||
if not element_yielded:
|
||||
await checkpoint()
|
||||
|
||||
|
||||
async def pairwise(
|
||||
iterable: Iterable[T] | AsyncIterable[T],
|
||||
) -> AsyncGenerator[tuple[T, T], None]:
|
||||
iterator = _iterate(iterable)
|
||||
try:
|
||||
previous = await anext(iterator)
|
||||
except StopAsyncIteration:
|
||||
await checkpoint()
|
||||
return
|
||||
|
||||
element_yielded = False
|
||||
async for element in iterator:
|
||||
element_yielded = True
|
||||
pair = (previous, element)
|
||||
previous = element
|
||||
yield pair
|
||||
|
||||
if not element_yielded:
|
||||
await checkpoint()
|
||||
|
||||
|
||||
async def permutations(
|
||||
iterable: Iterable[T] | AsyncIterable[T], r: int | None = None
|
||||
) -> AsyncGenerator[tuple[T, ...], None]:
|
||||
pool: list[T] = [element async for element in _iterate(iterable)]
|
||||
n = len(pool)
|
||||
if r is None:
|
||||
r = n
|
||||
elif not isinstance(r, int):
|
||||
raise TypeError("Expected int as r")
|
||||
elif r < 0:
|
||||
raise ValueError("r must be non-negative")
|
||||
|
||||
async for permutation in _iterate(itertools.permutations(pool, r)):
|
||||
yield permutation
|
||||
|
||||
|
||||
async def product(
|
||||
*iterables: Iterable[T] | AsyncIterable[T], repeat: int = 1
|
||||
) -> AsyncGenerator[tuple[T, ...], None]:
|
||||
repeat = operator.index(repeat)
|
||||
if repeat < 0:
|
||||
raise ValueError("repeat argument cannot be negative")
|
||||
|
||||
pools: list[tuple[T, ...]] = []
|
||||
for iterable in iterables:
|
||||
pool: list[T] = [element async for element in _iterate(iterable)]
|
||||
pools.append(tuple(pool))
|
||||
|
||||
async for value in _iterate(itertools.product(*pools, repeat=repeat)):
|
||||
yield value
|
||||
|
||||
|
||||
async def repeat(element: T, times: int | None = None) -> AsyncGenerator[T, None]:
|
||||
if times is None:
|
||||
while True:
|
||||
await checkpoint()
|
||||
yield element
|
||||
|
||||
remaining = operator.index(cast(Any, times))
|
||||
if remaining <= 0:
|
||||
await checkpoint()
|
||||
return
|
||||
|
||||
while remaining > 0:
|
||||
await checkpoint_if_cancelled()
|
||||
remaining -= 1
|
||||
await cancel_shielded_checkpoint()
|
||||
yield element
|
||||
|
||||
|
||||
async def starmap(
|
||||
function: Callable[..., Awaitable[R]],
|
||||
iterable: (
|
||||
Iterable[Iterable[object] | AsyncIterable[object]]
|
||||
| AsyncIterable[Iterable[object] | AsyncIterable[object]]
|
||||
),
|
||||
) -> AsyncGenerator[R, None]:
|
||||
result_yielded = False
|
||||
|
||||
async for args_iterable in _iterate(iterable):
|
||||
args = [element async for element in _iterate(args_iterable)]
|
||||
result_yielded = True
|
||||
yield await function(*args)
|
||||
|
||||
if not result_yielded:
|
||||
await checkpoint()
|
||||
|
||||
|
||||
def tee(
|
||||
iterable: Iterable[T] | AsyncIterable[T], n: int = 2
|
||||
) -> tuple[AsyncIterator[T], ...]:
|
||||
n = operator.index(cast(Any, n))
|
||||
if n < 0:
|
||||
raise ValueError("n must be >= 0")
|
||||
if n == 0:
|
||||
return ()
|
||||
|
||||
iterator = _TeeAsyncIterator(iterable)
|
||||
iterators: list[AsyncIterator[T]] = [iterator]
|
||||
iterators.extend(_TeeAsyncIterator(iterator) for _ in range(n - 1))
|
||||
return tuple(iterators)
|
||||
|
||||
|
||||
async def takewhile(
|
||||
predicate: Callable[[T], Awaitable[object]],
|
||||
iterable: Iterable[T] | AsyncIterable[T],
|
||||
) -> AsyncGenerator[T, None]:
|
||||
element_yielded = False
|
||||
|
||||
async for element in _iterate(iterable):
|
||||
if not await predicate(element):
|
||||
if not element_yielded:
|
||||
await checkpoint()
|
||||
|
||||
return
|
||||
|
||||
element_yielded = True
|
||||
yield element
|
||||
|
||||
if not element_yielded:
|
||||
await checkpoint()
|
||||
|
||||
|
||||
async def zip_longest(
|
||||
*iterables: Iterable[object] | AsyncIterable[object],
|
||||
fillvalue: object = None,
|
||||
) -> AsyncGenerator[tuple[object, ...], None]:
|
||||
iterators = [_iterate(iterable) for iterable in iterables]
|
||||
num_active = len(iterators)
|
||||
if not num_active:
|
||||
await checkpoint()
|
||||
return
|
||||
|
||||
active = [True] * num_active
|
||||
tuple_yielded = False
|
||||
|
||||
while True:
|
||||
values: list[object] = []
|
||||
for index, iterator in enumerate(iterators):
|
||||
if not active[index]:
|
||||
values.append(fillvalue)
|
||||
continue
|
||||
|
||||
try:
|
||||
value = await anext(iterator)
|
||||
except StopAsyncIteration:
|
||||
active[index] = False
|
||||
num_active -= 1
|
||||
if not num_active:
|
||||
if not tuple_yielded:
|
||||
await checkpoint()
|
||||
|
||||
return
|
||||
|
||||
value = fillvalue
|
||||
|
||||
values.append(value)
|
||||
|
||||
tuple_yielded = True
|
||||
yield tuple(values)
|
||||
+228
@@ -0,0 +1,228 @@
|
||||
from __future__ import annotations
|
||||
|
||||
__all__ = (
|
||||
"EventLoopToken",
|
||||
"RunVar",
|
||||
"RunvarToken",
|
||||
"cancel_shielded_checkpoint",
|
||||
"checkpoint",
|
||||
"checkpoint_if_cancelled",
|
||||
"current_token",
|
||||
)
|
||||
|
||||
import enum
|
||||
from dataclasses import dataclass
|
||||
from types import TracebackType
|
||||
from typing import TYPE_CHECKING, Any, Generic, Literal, TypeVar, final, overload
|
||||
from weakref import WeakKeyDictionary
|
||||
|
||||
from ._core._eventloop import get_async_backend
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .abc import AsyncBackend
|
||||
|
||||
T = TypeVar("T")
|
||||
D = TypeVar("D")
|
||||
|
||||
|
||||
async def checkpoint() -> None:
|
||||
"""
|
||||
Check for cancellation and allow the scheduler to switch to another task.
|
||||
|
||||
Equivalent to (but more efficient than)::
|
||||
|
||||
await checkpoint_if_cancelled()
|
||||
await cancel_shielded_checkpoint()
|
||||
|
||||
.. versionadded:: 3.0
|
||||
|
||||
"""
|
||||
await get_async_backend().checkpoint()
|
||||
|
||||
|
||||
async def checkpoint_if_cancelled() -> None:
|
||||
"""
|
||||
Enter a checkpoint if the enclosing cancel scope has been cancelled.
|
||||
|
||||
This does not allow the scheduler to switch to a different task.
|
||||
|
||||
.. versionadded:: 3.0
|
||||
|
||||
"""
|
||||
await get_async_backend().checkpoint_if_cancelled()
|
||||
|
||||
|
||||
async def cancel_shielded_checkpoint() -> None:
|
||||
"""
|
||||
Allow the scheduler to switch to another task but without checking for cancellation.
|
||||
|
||||
Equivalent to (but potentially more efficient than)::
|
||||
|
||||
with CancelScope(shield=True):
|
||||
await checkpoint()
|
||||
|
||||
.. versionadded:: 3.0
|
||||
|
||||
"""
|
||||
await get_async_backend().cancel_shielded_checkpoint()
|
||||
|
||||
|
||||
@final
|
||||
@dataclass(frozen=True, repr=False)
|
||||
class EventLoopToken:
|
||||
"""
|
||||
An opaque object that holds a reference to an event loop.
|
||||
|
||||
.. versionadded:: 4.11.0
|
||||
"""
|
||||
|
||||
backend_class: type[AsyncBackend]
|
||||
native_token: object
|
||||
|
||||
|
||||
def current_token() -> EventLoopToken:
|
||||
"""
|
||||
Return a token object that can be used to call code in the current event loop from
|
||||
another thread.
|
||||
|
||||
:raises NoEventLoopError: if no supported asynchronous event loop is running in the
|
||||
current thread
|
||||
|
||||
.. versionadded:: 4.11.0
|
||||
|
||||
"""
|
||||
backend_class = get_async_backend()
|
||||
raw_token = backend_class.current_token()
|
||||
return EventLoopToken(backend_class, raw_token)
|
||||
|
||||
|
||||
_run_vars: WeakKeyDictionary[object, dict[RunVar[Any], Any]] = WeakKeyDictionary()
|
||||
|
||||
|
||||
class _NoValueSet(enum.Enum):
|
||||
NO_VALUE_SET = enum.auto()
|
||||
|
||||
|
||||
class RunvarToken(Generic[T]):
|
||||
"""
|
||||
A token that can be used to restore a :class:`RunVar` to its previous value.
|
||||
|
||||
Returned by :meth:`RunVar.set`. Can be used as a context manager to automatically
|
||||
reset the variable on exit, or passed directly to :meth:`RunVar.reset`.
|
||||
"""
|
||||
|
||||
__slots__ = "_redeemed", "_value", "_var"
|
||||
|
||||
def __init__(self, var: RunVar[T], value: T | Literal[_NoValueSet.NO_VALUE_SET]):
|
||||
self._var = var
|
||||
self._value: T | Literal[_NoValueSet.NO_VALUE_SET] = value
|
||||
self._redeemed = False
|
||||
|
||||
def __enter__(self) -> RunvarToken[T]:
|
||||
return self
|
||||
|
||||
def __exit__(
|
||||
self,
|
||||
exc_type: type[BaseException] | None,
|
||||
exc_val: BaseException | None,
|
||||
exc_tb: TracebackType | None,
|
||||
) -> None:
|
||||
self._var.reset(self)
|
||||
|
||||
|
||||
class RunVar(Generic[T]):
|
||||
"""
|
||||
Like a :class:`~contextvars.ContextVar`, except scoped to the running event loop.
|
||||
|
||||
Can be used as a context manager, Just like :class:`~contextvars.ContextVar`, that
|
||||
will reset the variable to its previous value when the context block is exited.
|
||||
"""
|
||||
|
||||
__slots__ = "_default", "_name"
|
||||
|
||||
NO_VALUE_SET: Literal[_NoValueSet.NO_VALUE_SET] = _NoValueSet.NO_VALUE_SET
|
||||
|
||||
def __init__(
|
||||
self, name: str, default: T | Literal[_NoValueSet.NO_VALUE_SET] = NO_VALUE_SET
|
||||
):
|
||||
self._name = name
|
||||
self._default = default
|
||||
|
||||
@property
|
||||
def _current_vars(self) -> dict[RunVar[T], T]:
|
||||
native_token = current_token().native_token
|
||||
try:
|
||||
return _run_vars[native_token]
|
||||
except KeyError:
|
||||
run_vars = _run_vars[native_token] = {}
|
||||
return run_vars
|
||||
|
||||
@overload
|
||||
def get(self, default: D) -> T | D: ...
|
||||
|
||||
@overload
|
||||
def get(self) -> T: ...
|
||||
|
||||
def get(
|
||||
self, default: D | Literal[_NoValueSet.NO_VALUE_SET] = NO_VALUE_SET
|
||||
) -> T | D:
|
||||
"""
|
||||
Return the current value of this run variable.
|
||||
|
||||
:param default: a fallback value to return if no value has been set
|
||||
:return: the current value, the provided default, or the variable's own default
|
||||
:raises LookupError: if no value is set and no default is available
|
||||
|
||||
"""
|
||||
try:
|
||||
return self._current_vars[self]
|
||||
except KeyError:
|
||||
if default is not RunVar.NO_VALUE_SET:
|
||||
return default
|
||||
elif self._default is not RunVar.NO_VALUE_SET:
|
||||
return self._default
|
||||
|
||||
raise LookupError(
|
||||
f'Run variable "{self._name}" has no value and no default set'
|
||||
)
|
||||
|
||||
def set(self, value: T) -> RunvarToken[T]:
|
||||
"""
|
||||
Set the value of this run variable for the current event loop.
|
||||
|
||||
:param value: the new value
|
||||
:return: a token that can be used to restore the previous value
|
||||
|
||||
"""
|
||||
current_vars = self._current_vars
|
||||
token = RunvarToken(self, current_vars.get(self, RunVar.NO_VALUE_SET))
|
||||
current_vars[self] = value
|
||||
return token
|
||||
|
||||
def reset(self, token: RunvarToken[T]) -> None:
|
||||
"""
|
||||
Restore this run variable to the value it held before the matching :meth:`set`.
|
||||
|
||||
:param token: the token returned by :meth:`set`
|
||||
:raises ValueError: if the token belongs to a different :class:`RunVar` or the token
|
||||
has already been used
|
||||
|
||||
"""
|
||||
if token._var is not self:
|
||||
raise ValueError("This token does not belong to this RunVar")
|
||||
|
||||
if token._redeemed:
|
||||
raise ValueError("This token has already been used")
|
||||
|
||||
if token._value is _NoValueSet.NO_VALUE_SET:
|
||||
try:
|
||||
del self._current_vars[self]
|
||||
except KeyError:
|
||||
pass
|
||||
else:
|
||||
self._current_vars[self] = token._value
|
||||
|
||||
token._redeemed = True
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"<RunVar name={self._name!r}>"
|
||||
Whitespace-only changes.
+404
@@ -0,0 +1,404 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import dataclasses
|
||||
import socket
|
||||
import sys
|
||||
from collections.abc import Callable, Generator, Iterator
|
||||
from contextlib import ExitStack, contextmanager
|
||||
from inspect import isasyncgenfunction, iscoroutinefunction, ismethod
|
||||
from typing import TYPE_CHECKING, Any, cast
|
||||
|
||||
import pytest
|
||||
from _pytest.fixtures import FuncFixtureInfo, SubRequest
|
||||
from _pytest.outcomes import Exit
|
||||
from _pytest.scope import Scope
|
||||
|
||||
from . import get_available_backends
|
||||
from ._core._eventloop import (
|
||||
current_async_library,
|
||||
get_async_backend,
|
||||
reset_current_async_library,
|
||||
set_current_async_library,
|
||||
)
|
||||
from ._core._exceptions import iterate_exceptions
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .abc import TestRunner
|
||||
|
||||
if sys.version_info < (3, 11):
|
||||
from exceptiongroup import ExceptionGroup
|
||||
|
||||
if TYPE_CHECKING:
|
||||
# pytest >= 9.2 keeps CallSpec2 as a TYPE_CHECKING-only alias of CallSpec
|
||||
from _pytest.python import CallSpec2 as CallSpec
|
||||
else:
|
||||
try:
|
||||
# c.f. https://github.com/pytest-dev/pytest/pull/14742
|
||||
# pytest >= 9.2
|
||||
from _pytest.python import CallSpec
|
||||
except ImportError:
|
||||
from _pytest.python import CallSpec2 as CallSpec
|
||||
|
||||
_current_runner: TestRunner | None = None
|
||||
_runner_stack: ExitStack | None = None
|
||||
_runner_leases = 0
|
||||
|
||||
|
||||
def extract_backend_and_options(backend: object) -> tuple[str, dict[str, Any]]:
|
||||
if isinstance(backend, str):
|
||||
return backend, {}
|
||||
elif isinstance(backend, tuple) and len(backend) == 2:
|
||||
if isinstance(backend[0], str) and isinstance(backend[1], dict):
|
||||
return cast(tuple[str, dict[str, Any]], backend)
|
||||
|
||||
raise TypeError("anyio_backend must be either a string or tuple of (string, dict)")
|
||||
|
||||
|
||||
@contextmanager
|
||||
def get_runner(
|
||||
backend_name: str, backend_options: dict[str, Any]
|
||||
) -> Iterator[TestRunner]:
|
||||
global _current_runner, _runner_leases, _runner_stack
|
||||
if _current_runner is None:
|
||||
asynclib = get_async_backend(backend_name)
|
||||
_runner_stack = ExitStack()
|
||||
if current_async_library() is None:
|
||||
# Since we're in control of the event loop, we can cache the name of the
|
||||
# async library
|
||||
token = set_current_async_library(backend_name)
|
||||
_runner_stack.callback(reset_current_async_library, token)
|
||||
|
||||
backend_options = backend_options or {}
|
||||
_current_runner = _runner_stack.enter_context(
|
||||
asynclib.create_test_runner(backend_options)
|
||||
)
|
||||
|
||||
_runner_leases += 1
|
||||
try:
|
||||
yield _current_runner
|
||||
finally:
|
||||
_runner_leases -= 1
|
||||
if not _runner_leases:
|
||||
assert _runner_stack is not None
|
||||
_runner_stack.close()
|
||||
_runner_stack = _current_runner = None
|
||||
|
||||
|
||||
def pytest_addoption(parser: pytest.Parser) -> None:
|
||||
group = parser.getgroup("anyio")
|
||||
parser.addini(
|
||||
"anyio_mode",
|
||||
default="strict",
|
||||
help='AnyIO plugin mode (either "strict" or "auto")',
|
||||
)
|
||||
group.addoption(
|
||||
"--anyio-mode",
|
||||
default=None,
|
||||
dest="anyio_mode",
|
||||
help="""
|
||||
'auto' - All async test functions will be handled by AnyIO pytest plugin
|
||||
'strict' - Disabling autoprocessing(useful when anyio tests
|
||||
need to coexist with other async test plugins)
|
||||
""",
|
||||
)
|
||||
|
||||
|
||||
def _is_auto_mode(config: pytest.Config, mode: str) -> bool:
|
||||
return (
|
||||
config.getoption(mode, default=None) == "auto" or config.getini(mode) == "auto"
|
||||
)
|
||||
|
||||
|
||||
def pytest_configure(config: pytest.Config) -> None:
|
||||
config.addinivalue_line(
|
||||
"markers",
|
||||
"anyio: mark the (coroutine function) test to be run asynchronously via anyio.",
|
||||
)
|
||||
if (
|
||||
_is_auto_mode(config, "anyio_mode")
|
||||
and config.pluginmanager.has_plugin("asyncio")
|
||||
and _is_auto_mode(config, "asyncio_mode")
|
||||
):
|
||||
config.issue_config_time_warning(
|
||||
pytest.PytestConfigWarning(
|
||||
"AnyIO auto mode has been enabled together with pytest-asyncio auto "
|
||||
"mode. This may cause unexpected behavior."
|
||||
),
|
||||
1,
|
||||
)
|
||||
|
||||
|
||||
@pytest.hookimpl(hookwrapper=True)
|
||||
def pytest_fixture_setup(fixturedef: Any, request: Any) -> Generator[Any]:
|
||||
def wrapper(anyio_backend: Any, request: SubRequest, **kwargs: Any) -> Any:
|
||||
# Rebind any fixture methods to the request instance
|
||||
if (
|
||||
request.instance
|
||||
and ismethod(func)
|
||||
and type(func.__self__) is type(request.instance)
|
||||
):
|
||||
local_func = func.__func__.__get__(request.instance)
|
||||
else:
|
||||
local_func = func
|
||||
|
||||
backend_name, backend_options = extract_backend_and_options(anyio_backend)
|
||||
if has_backend_arg:
|
||||
kwargs["anyio_backend"] = anyio_backend
|
||||
|
||||
if has_request_arg:
|
||||
kwargs["request"] = request
|
||||
|
||||
with get_runner(backend_name, backend_options) as runner:
|
||||
# re-entrant call into the test runner detected. this happens when an async fixture
|
||||
# is dynamically requested via request.getfixturevalue() from inside a running async
|
||||
# test or fixture. on asyncio this raises RuntimeError: This event loop is already
|
||||
# running, on trio the runner deadlocks - the host loop blocks waiting for the
|
||||
# coroutine to return, but the coroutine is waiting for the host loop. raising here
|
||||
# prevents the hang and gives a consistent error across backends.
|
||||
if runner.is_running():
|
||||
raise RuntimeError(
|
||||
"Cannot schedule a coroutine in the test runner while another is already running; "
|
||||
"likely caused by request.getfixturevalue() on an async fixture."
|
||||
)
|
||||
|
||||
if isasyncgenfunction(local_func):
|
||||
yield from runner.run_asyncgen_fixture(local_func, kwargs)
|
||||
else:
|
||||
yield runner.run_fixture(local_func, kwargs)
|
||||
|
||||
# Only apply this to coroutine functions and async generator functions in requests
|
||||
# that involve the anyio_backend fixture
|
||||
func = fixturedef.func
|
||||
if isasyncgenfunction(func) or iscoroutinefunction(func):
|
||||
if "anyio_backend" in request.fixturenames:
|
||||
fixturedef.func = wrapper
|
||||
original_argname = fixturedef.argnames
|
||||
|
||||
if not (has_backend_arg := "anyio_backend" in fixturedef.argnames):
|
||||
fixturedef.argnames += ("anyio_backend",)
|
||||
|
||||
if not (has_request_arg := "request" in fixturedef.argnames):
|
||||
fixturedef.argnames += ("request",)
|
||||
|
||||
try:
|
||||
return (yield)
|
||||
finally:
|
||||
fixturedef.func = func
|
||||
fixturedef.argnames = original_argname
|
||||
|
||||
return (yield)
|
||||
|
||||
|
||||
@pytest.hookimpl(tryfirst=True)
|
||||
def pytest_pycollect_makeitem(
|
||||
collector: pytest.Module | pytest.Class, name: str, obj: object
|
||||
) -> None:
|
||||
if collector.istestfunction(obj, name):
|
||||
inner_func = obj.hypothesis.inner_test if hasattr(obj, "hypothesis") else obj
|
||||
if iscoroutinefunction(inner_func):
|
||||
anyio_auto_mode = _is_auto_mode(collector.config, "anyio_mode")
|
||||
marker = collector.get_closest_marker("anyio")
|
||||
own_markers = getattr(obj, "pytestmark", ())
|
||||
if (
|
||||
anyio_auto_mode
|
||||
or marker
|
||||
or any(marker.name == "anyio" for marker in own_markers)
|
||||
):
|
||||
pytest.mark.usefixtures("anyio_backend")(obj)
|
||||
|
||||
|
||||
def pytest_collection_finish(session: pytest.Session) -> None:
|
||||
for i, item in reversed(list(enumerate(session.items))):
|
||||
if (
|
||||
isinstance(item, pytest.Function)
|
||||
and iscoroutinefunction(item.function)
|
||||
and item.get_closest_marker("anyio") is not None
|
||||
and "anyio_backend" not in item.fixturenames
|
||||
):
|
||||
new_items = []
|
||||
try:
|
||||
cs_fields = {f.name for f in dataclasses.fields(CallSpec)}
|
||||
except TypeError:
|
||||
cs_fields = set()
|
||||
|
||||
for param_index, backend in enumerate(get_available_backends()):
|
||||
if "_arg2scope" in cs_fields: # pytest >= 8
|
||||
callspec = CallSpec(
|
||||
params={"anyio_backend": backend},
|
||||
indices={"anyio_backend": param_index},
|
||||
_arg2scope={"anyio_backend": Scope.Module},
|
||||
_idlist=[backend],
|
||||
marks=[],
|
||||
)
|
||||
else: # pytest 7.x
|
||||
callspec = CallSpec( # type: ignore[call-arg]
|
||||
funcargs={},
|
||||
params={"anyio_backend": backend},
|
||||
indices={"anyio_backend": param_index},
|
||||
arg2scope={"anyio_backend": Scope.Module},
|
||||
idlist=[backend],
|
||||
marks=[],
|
||||
)
|
||||
|
||||
fi = item._fixtureinfo
|
||||
new_names_closure = list(fi.names_closure)
|
||||
if "anyio_backend" not in new_names_closure:
|
||||
new_names_closure.append("anyio_backend")
|
||||
|
||||
new_fixtureinfo = FuncFixtureInfo(
|
||||
argnames=fi.argnames,
|
||||
initialnames=fi.initialnames,
|
||||
names_closure=new_names_closure,
|
||||
name2fixturedefs=fi.name2fixturedefs,
|
||||
)
|
||||
new_item = pytest.Function.from_parent(
|
||||
item.parent,
|
||||
name=f"{item.originalname}[{backend}]",
|
||||
callspec=callspec,
|
||||
callobj=item.obj,
|
||||
fixtureinfo=new_fixtureinfo,
|
||||
keywords=item.keywords,
|
||||
originalname=item.originalname,
|
||||
)
|
||||
new_items.append(new_item)
|
||||
|
||||
session.items[i : i + 1] = new_items
|
||||
|
||||
|
||||
@pytest.hookimpl(tryfirst=True)
|
||||
def pytest_pyfunc_call(pyfuncitem: Any) -> bool | None:
|
||||
def run_with_hypothesis(**kwargs: Any) -> None:
|
||||
with get_runner(backend_name, backend_options) as runner:
|
||||
runner.run_test(original_func, kwargs)
|
||||
|
||||
backend = pyfuncitem.funcargs.get("anyio_backend")
|
||||
if backend:
|
||||
backend_name, backend_options = extract_backend_and_options(backend)
|
||||
|
||||
if hasattr(pyfuncitem.obj, "hypothesis"):
|
||||
# Wrap the inner test function unless it's already wrapped
|
||||
original_func = pyfuncitem.obj.hypothesis.inner_test
|
||||
if original_func.__qualname__ != run_with_hypothesis.__qualname__:
|
||||
if iscoroutinefunction(original_func):
|
||||
pyfuncitem.obj.hypothesis.inner_test = run_with_hypothesis
|
||||
|
||||
return None
|
||||
|
||||
if iscoroutinefunction(pyfuncitem.obj):
|
||||
funcargs = pyfuncitem.funcargs
|
||||
testargs = {arg: funcargs[arg] for arg in pyfuncitem._fixtureinfo.argnames}
|
||||
with get_runner(backend_name, backend_options) as runner:
|
||||
try:
|
||||
runner.run_test(pyfuncitem.obj, testargs)
|
||||
except ExceptionGroup as excgrp:
|
||||
for exc in iterate_exceptions(excgrp):
|
||||
if isinstance(exc, (Exit, KeyboardInterrupt, SystemExit)):
|
||||
raise exc from excgrp
|
||||
|
||||
raise
|
||||
|
||||
return True
|
||||
|
||||
return None
|
||||
|
||||
|
||||
@pytest.fixture(scope="module", params=get_available_backends())
|
||||
def anyio_backend(request: Any) -> Any:
|
||||
return request.param
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def anyio_backend_name(anyio_backend: Any) -> str:
|
||||
if isinstance(anyio_backend, str):
|
||||
return anyio_backend
|
||||
else:
|
||||
return anyio_backend[0]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def anyio_backend_options(anyio_backend: Any) -> dict[str, Any]:
|
||||
if isinstance(anyio_backend, str):
|
||||
return {}
|
||||
else:
|
||||
return anyio_backend[1]
|
||||
|
||||
|
||||
class FreePortFactory:
|
||||
"""
|
||||
Manages port generation based on specified socket kind, ensuring no duplicate
|
||||
ports are generated.
|
||||
|
||||
This class provides functionality for generating available free ports on the
|
||||
system. It is initialized with a specific socket kind and can generate ports
|
||||
for given address families while avoiding reuse of previously generated ports.
|
||||
|
||||
Users should not instantiate this class directly, but use the
|
||||
``free_tcp_port_factory`` and ``free_udp_port_factory`` fixtures instead. For simple
|
||||
uses cases, ``free_tcp_port`` and ``free_udp_port`` can be used instead.
|
||||
"""
|
||||
|
||||
def __init__(self, kind: socket.SocketKind) -> None:
|
||||
self._kind = kind
|
||||
self._generated = set[int]()
|
||||
|
||||
@property
|
||||
def kind(self) -> socket.SocketKind:
|
||||
"""
|
||||
The type of socket connection (e.g., :data:`~socket.SOCK_STREAM` or
|
||||
:data:`~socket.SOCK_DGRAM`) used to bind for checking port availability
|
||||
|
||||
"""
|
||||
return self._kind
|
||||
|
||||
def __call__(self, family: socket.AddressFamily | None = None) -> int:
|
||||
"""
|
||||
Return an unbound port for the given address family.
|
||||
|
||||
:param family: if omitted, both IPv4 and IPv6 addresses will be tried
|
||||
:return: a port number
|
||||
|
||||
"""
|
||||
if family is not None:
|
||||
families = [family]
|
||||
else:
|
||||
families = [socket.AF_INET]
|
||||
if socket.has_ipv6:
|
||||
families.append(socket.AF_INET6)
|
||||
|
||||
while True:
|
||||
port = 0
|
||||
with ExitStack() as stack:
|
||||
for fam in families:
|
||||
sock = stack.enter_context(socket.socket(fam, self._kind))
|
||||
addr = "::1" if fam == socket.AF_INET6 else "127.0.0.1"
|
||||
try:
|
||||
sock.bind((addr, port))
|
||||
except OSError:
|
||||
break
|
||||
|
||||
if not port:
|
||||
port = sock.getsockname()[1]
|
||||
else:
|
||||
if port not in self._generated:
|
||||
self._generated.add(port)
|
||||
return port
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def free_tcp_port_factory() -> FreePortFactory:
|
||||
return FreePortFactory(socket.SOCK_STREAM)
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def free_udp_port_factory() -> FreePortFactory:
|
||||
return FreePortFactory(socket.SOCK_DGRAM)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def free_tcp_port(free_tcp_port_factory: Callable[[], int]) -> int:
|
||||
return free_tcp_port_factory()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def free_udp_port(free_udp_port_factory: Callable[[], int]) -> int:
|
||||
return free_udp_port_factory()
|
||||
Whitespace-only changes.
+201
@@ -0,0 +1,201 @@
|
||||
from __future__ import annotations
|
||||
|
||||
__all__ = (
|
||||
"BufferedByteReceiveStream",
|
||||
"BufferedByteStream",
|
||||
"BufferedConnectable",
|
||||
)
|
||||
|
||||
import sys
|
||||
from collections.abc import Callable, Iterable, Mapping
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, SupportsIndex
|
||||
|
||||
from .. import ClosedResourceError, DelimiterNotFound, EndOfStream, IncompleteRead
|
||||
from ..abc import (
|
||||
AnyByteReceiveStream,
|
||||
AnyByteStream,
|
||||
AnyByteStreamConnectable,
|
||||
ByteReceiveStream,
|
||||
ByteStream,
|
||||
ByteStreamConnectable,
|
||||
)
|
||||
|
||||
if sys.version_info >= (3, 12):
|
||||
from typing import override
|
||||
else:
|
||||
from typing_extensions import override
|
||||
|
||||
|
||||
@dataclass(eq=False)
|
||||
class BufferedByteReceiveStream(ByteReceiveStream):
|
||||
"""
|
||||
Wraps any bytes-based receive stream and uses a buffer to provide sophisticated
|
||||
receiving capabilities in the form of a byte stream.
|
||||
"""
|
||||
|
||||
receive_stream: AnyByteReceiveStream
|
||||
_buffer: bytearray = field(init=False, default_factory=bytearray)
|
||||
_closed: bool = field(init=False, default=False)
|
||||
|
||||
async def aclose(self) -> None:
|
||||
await self.receive_stream.aclose()
|
||||
self._closed = True
|
||||
|
||||
@property
|
||||
def buffer(self) -> bytes:
|
||||
"""The bytes currently in the buffer."""
|
||||
return bytes(self._buffer)
|
||||
|
||||
@property
|
||||
def extra_attributes(self) -> Mapping[Any, Callable[[], Any]]:
|
||||
return self.receive_stream.extra_attributes
|
||||
|
||||
def feed_data(self, data: Iterable[SupportsIndex], /) -> None:
|
||||
"""
|
||||
Append data directly into the buffer.
|
||||
|
||||
Any data in the buffer will be consumed by receive operations before receiving
|
||||
anything from the wrapped stream.
|
||||
|
||||
:param data: the data to append to the buffer (can be bytes or anything else
|
||||
that supports ``__index__()``)
|
||||
|
||||
"""
|
||||
self._buffer.extend(data)
|
||||
|
||||
async def receive(self, max_bytes: int = 65536) -> bytes:
|
||||
if max_bytes < 1:
|
||||
raise ValueError("max_bytes must be a positive integer")
|
||||
|
||||
if self._closed:
|
||||
raise ClosedResourceError
|
||||
|
||||
if self._buffer:
|
||||
chunk = bytes(self._buffer[:max_bytes])
|
||||
del self._buffer[:max_bytes]
|
||||
return chunk
|
||||
elif isinstance(self.receive_stream, ByteReceiveStream):
|
||||
return await self.receive_stream.receive(max_bytes)
|
||||
else:
|
||||
# With a bytes-oriented object stream, we need to handle any surplus bytes
|
||||
# we get from the receive() call
|
||||
chunk = await self.receive_stream.receive()
|
||||
if len(chunk) > max_bytes:
|
||||
# Save the surplus bytes in the buffer
|
||||
self._buffer.extend(chunk[max_bytes:])
|
||||
return chunk[:max_bytes]
|
||||
else:
|
||||
return chunk
|
||||
|
||||
async def receive_exactly(self, nbytes: int) -> bytes:
|
||||
"""
|
||||
Read exactly the given amount of bytes from the stream.
|
||||
|
||||
:param nbytes: the number of bytes to read
|
||||
:return: the bytes read
|
||||
:raises ~anyio.IncompleteRead: if the stream was closed before the requested
|
||||
amount of bytes could be read from the stream
|
||||
|
||||
"""
|
||||
while True:
|
||||
remaining = nbytes - len(self._buffer)
|
||||
if remaining <= 0:
|
||||
retval = self._buffer[:nbytes]
|
||||
del self._buffer[:nbytes]
|
||||
return bytes(retval)
|
||||
|
||||
try:
|
||||
if isinstance(self.receive_stream, ByteReceiveStream):
|
||||
chunk = await self.receive_stream.receive(remaining)
|
||||
else:
|
||||
chunk = await self.receive_stream.receive()
|
||||
except EndOfStream as exc:
|
||||
raise IncompleteRead from exc
|
||||
|
||||
self._buffer.extend(chunk)
|
||||
|
||||
async def receive_until(self, delimiter: bytes, max_bytes: int) -> bytes:
|
||||
"""
|
||||
Read from the stream until the delimiter is found or max_bytes have been read.
|
||||
|
||||
:param delimiter: the marker to look for in the stream
|
||||
:param max_bytes: maximum number of bytes that will be read before raising
|
||||
:exc:`~anyio.DelimiterNotFound`
|
||||
:return: the bytes read (not including the delimiter)
|
||||
:raises ~anyio.IncompleteRead: if the stream was closed before the delimiter
|
||||
was found
|
||||
:raises ~anyio.DelimiterNotFound: if the delimiter is not found within the
|
||||
bytes read up to the maximum allowed
|
||||
|
||||
"""
|
||||
delimiter_size = len(delimiter)
|
||||
offset = 0
|
||||
while True:
|
||||
# Check if the delimiter can be found in the current buffer
|
||||
index = self._buffer.find(delimiter, offset)
|
||||
if index >= 0:
|
||||
found = self._buffer[:index]
|
||||
del self._buffer[: index + len(delimiter) :]
|
||||
return bytes(found)
|
||||
|
||||
# Check if the buffer is already at or over the limit
|
||||
if len(self._buffer) >= max_bytes:
|
||||
raise DelimiterNotFound(max_bytes)
|
||||
|
||||
# Read more data into the buffer from the socket
|
||||
try:
|
||||
data = await self.receive_stream.receive()
|
||||
except EndOfStream as exc:
|
||||
raise IncompleteRead from exc
|
||||
|
||||
# Move the offset forward and add the new data to the buffer
|
||||
offset = max(len(self._buffer) - delimiter_size + 1, 0)
|
||||
self._buffer.extend(data)
|
||||
|
||||
|
||||
class BufferedByteStream(BufferedByteReceiveStream, ByteStream):
|
||||
"""
|
||||
A full-duplex variant of :class:`BufferedByteReceiveStream`. All writes are passed
|
||||
through to the wrapped stream as-is.
|
||||
"""
|
||||
|
||||
def __init__(self, stream: AnyByteStream):
|
||||
"""
|
||||
:param stream: the stream to be wrapped
|
||||
|
||||
"""
|
||||
super().__init__(stream)
|
||||
self._stream = stream
|
||||
|
||||
@override
|
||||
async def send_eof(self) -> None:
|
||||
await self._stream.send_eof()
|
||||
|
||||
@override
|
||||
async def send(self, item: bytes) -> None:
|
||||
await self._stream.send(item)
|
||||
|
||||
|
||||
class BufferedConnectable(ByteStreamConnectable):
|
||||
"""
|
||||
Wraps a byte stream connectable to produce :class:`BufferedByteStream` connections.
|
||||
|
||||
Use this when you want the streams returned by :meth:`connect` to have the buffered
|
||||
receive API (e.g. :meth:`~BufferedByteReceiveStream.receive_exactly` and
|
||||
:meth:`~BufferedByteReceiveStream.receive_until`).
|
||||
|
||||
:param connectable: the byte stream connectable to wrap
|
||||
"""
|
||||
|
||||
def __init__(self, connectable: AnyByteStreamConnectable):
|
||||
"""
|
||||
:param connectable: the connectable to wrap
|
||||
|
||||
"""
|
||||
self.connectable = connectable
|
||||
|
||||
@override
|
||||
async def connect(self) -> BufferedByteStream:
|
||||
stream = await self.connectable.connect()
|
||||
return BufferedByteStream(stream)
|
||||
+157
@@ -0,0 +1,157 @@
|
||||
from __future__ import annotations
|
||||
|
||||
__all__ = (
|
||||
"FileReadStream",
|
||||
"FileStreamAttribute",
|
||||
"FileWriteStream",
|
||||
)
|
||||
|
||||
from collections.abc import Callable, Mapping
|
||||
from io import SEEK_SET, UnsupportedOperation
|
||||
from os import PathLike
|
||||
from pathlib import Path
|
||||
from typing import IO, Any
|
||||
|
||||
from .. import (
|
||||
BrokenResourceError,
|
||||
ClosedResourceError,
|
||||
EndOfStream,
|
||||
TypedAttributeSet,
|
||||
to_thread,
|
||||
typed_attribute,
|
||||
)
|
||||
from ..abc import ByteReceiveStream, ByteSendStream
|
||||
|
||||
|
||||
class FileStreamAttribute(TypedAttributeSet):
|
||||
#: the open file descriptor
|
||||
file: IO[bytes] = typed_attribute()
|
||||
#: the path of the file on the file system, if available (file must be a real file)
|
||||
path: Path = typed_attribute()
|
||||
#: the file number, if available (file must be a real file or a TTY)
|
||||
fileno: int = typed_attribute()
|
||||
|
||||
|
||||
class _BaseFileStream:
|
||||
def __init__(self, file: IO[bytes]):
|
||||
self._file = file
|
||||
|
||||
async def aclose(self) -> None:
|
||||
await to_thread.run_sync(self._file.close)
|
||||
|
||||
@property
|
||||
def extra_attributes(self) -> Mapping[Any, Callable[[], Any]]:
|
||||
attributes: dict[Any, Callable[[], Any]] = {
|
||||
FileStreamAttribute.file: lambda: self._file,
|
||||
}
|
||||
|
||||
if hasattr(self._file, "name"):
|
||||
attributes[FileStreamAttribute.path] = lambda: Path(self._file.name)
|
||||
|
||||
try:
|
||||
self._file.fileno()
|
||||
except UnsupportedOperation:
|
||||
pass
|
||||
else:
|
||||
attributes[FileStreamAttribute.fileno] = lambda: self._file.fileno()
|
||||
|
||||
return attributes
|
||||
|
||||
|
||||
class FileReadStream(_BaseFileStream, ByteReceiveStream):
|
||||
"""
|
||||
A byte stream that reads from a file in the file system.
|
||||
|
||||
:param file: a file that has been opened for reading in binary mode
|
||||
|
||||
.. versionadded:: 3.0
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
async def from_path(cls, path: str | PathLike[str]) -> FileReadStream:
|
||||
"""
|
||||
Create a file read stream by opening the given file.
|
||||
|
||||
:param path: path of the file to read from
|
||||
|
||||
"""
|
||||
file = await to_thread.run_sync(Path(path).open, "rb")
|
||||
return cls(file)
|
||||
|
||||
async def receive(self, max_bytes: int = 65536) -> bytes:
|
||||
if max_bytes < 1:
|
||||
raise ValueError("max_bytes must be a positive integer")
|
||||
|
||||
try:
|
||||
data = await to_thread.run_sync(self._file.read, max_bytes)
|
||||
except ValueError:
|
||||
raise ClosedResourceError from None
|
||||
except OSError as exc:
|
||||
raise BrokenResourceError from exc
|
||||
|
||||
if data:
|
||||
return data
|
||||
else:
|
||||
raise EndOfStream
|
||||
|
||||
async def seek(self, position: int, whence: int = SEEK_SET) -> int:
|
||||
"""
|
||||
Seek the file to the given position.
|
||||
|
||||
.. seealso:: :meth:`io.IOBase.seek`
|
||||
|
||||
.. note:: Not all file descriptors are seekable.
|
||||
|
||||
:param position: position to seek the file to
|
||||
:param whence: controls how ``position`` is interpreted
|
||||
:return: the new absolute position
|
||||
:raises OSError: if the file is not seekable
|
||||
|
||||
"""
|
||||
return await to_thread.run_sync(self._file.seek, position, whence)
|
||||
|
||||
async def tell(self) -> int:
|
||||
"""
|
||||
Return the current stream position.
|
||||
|
||||
.. note:: Not all file descriptors are seekable.
|
||||
|
||||
:return: the current absolute position
|
||||
:raises OSError: if the file is not seekable
|
||||
|
||||
"""
|
||||
return await to_thread.run_sync(self._file.tell)
|
||||
|
||||
|
||||
class FileWriteStream(_BaseFileStream, ByteSendStream):
|
||||
"""
|
||||
A byte stream that writes to a file in the file system.
|
||||
|
||||
:param file: a file that has been opened for writing in binary mode
|
||||
|
||||
.. versionadded:: 3.0
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
async def from_path(
|
||||
cls, path: str | PathLike[str], append: bool = False
|
||||
) -> FileWriteStream:
|
||||
"""
|
||||
Create a file write stream by opening the given file for writing.
|
||||
|
||||
:param path: path of the file to write to
|
||||
:param append: if ``True``, open the file for appending; if ``False``, any
|
||||
existing file at the given path will be truncated
|
||||
|
||||
"""
|
||||
mode = "ab" if append else "wb"
|
||||
file = await to_thread.run_sync(Path(path).open, mode)
|
||||
return cls(file)
|
||||
|
||||
async def send(self, item: bytes) -> None:
|
||||
try:
|
||||
await to_thread.run_sync(self._file.write, item)
|
||||
except ValueError:
|
||||
raise ClosedResourceError from None
|
||||
except OSError as exc:
|
||||
raise BrokenResourceError from exc
|
||||
+326
@@ -0,0 +1,326 @@
|
||||
from __future__ import annotations
|
||||
|
||||
__all__ = (
|
||||
"MemoryObjectReceiveStream",
|
||||
"MemoryObjectSendStream",
|
||||
"MemoryObjectStreamStatistics",
|
||||
)
|
||||
|
||||
import warnings
|
||||
from collections import OrderedDict, deque
|
||||
from dataclasses import dataclass, field
|
||||
from types import TracebackType
|
||||
from typing import Generic, NamedTuple, TypeVar
|
||||
|
||||
from .._core._exceptions import (
|
||||
BrokenResourceError,
|
||||
ClosedResourceError,
|
||||
EndOfStream,
|
||||
WouldBlock,
|
||||
)
|
||||
from .._core._synchronization import Event
|
||||
from .._core._testing import TaskInfo, get_current_task
|
||||
from ..abc import ObjectReceiveStream, ObjectSendStream
|
||||
from ..lowlevel import checkpoint
|
||||
|
||||
T_Item = TypeVar("T_Item")
|
||||
T_co = TypeVar("T_co", covariant=True)
|
||||
T_contra = TypeVar("T_contra", contravariant=True)
|
||||
|
||||
|
||||
class MemoryObjectStreamStatistics(NamedTuple):
|
||||
current_buffer_used: int #: number of items stored in the buffer
|
||||
#: maximum number of items that can be stored on this stream (or :data:`math.inf`)
|
||||
max_buffer_size: float
|
||||
open_send_streams: int #: number of unclosed clones of the send stream
|
||||
open_receive_streams: int #: number of unclosed clones of the receive stream
|
||||
#: number of tasks blocked on :meth:`MemoryObjectSendStream.send`
|
||||
tasks_waiting_send: int
|
||||
#: number of tasks blocked on :meth:`MemoryObjectReceiveStream.receive`
|
||||
tasks_waiting_receive: int
|
||||
|
||||
|
||||
@dataclass(eq=False)
|
||||
class _MemoryObjectItemReceiver(Generic[T_Item]):
|
||||
task_info: TaskInfo = field(init=False, default_factory=get_current_task)
|
||||
item: T_Item = field(init=False)
|
||||
|
||||
def __repr__(self) -> str:
|
||||
# When item is not defined, we get following error with default __repr__:
|
||||
# AttributeError: 'MemoryObjectItemReceiver' object has no attribute 'item'
|
||||
item = getattr(self, "item", None)
|
||||
return f"{self.__class__.__name__}(task_info={self.task_info}, item={item!r})"
|
||||
|
||||
|
||||
@dataclass(eq=False)
|
||||
class _MemoryObjectStreamState(Generic[T_Item]):
|
||||
max_buffer_size: float = field()
|
||||
buffer: deque[T_Item] = field(init=False, default_factory=deque)
|
||||
open_send_channels: int = field(init=False, default=0)
|
||||
open_receive_channels: int = field(init=False, default=0)
|
||||
waiting_receivers: OrderedDict[Event, _MemoryObjectItemReceiver[T_Item]] = field(
|
||||
init=False, default_factory=OrderedDict
|
||||
)
|
||||
waiting_senders: OrderedDict[Event, T_Item] = field(
|
||||
init=False, default_factory=OrderedDict
|
||||
)
|
||||
|
||||
def statistics(self) -> MemoryObjectStreamStatistics:
|
||||
return MemoryObjectStreamStatistics(
|
||||
len(self.buffer),
|
||||
self.max_buffer_size,
|
||||
self.open_send_channels,
|
||||
self.open_receive_channels,
|
||||
len(self.waiting_senders),
|
||||
len(self.waiting_receivers),
|
||||
)
|
||||
|
||||
|
||||
@dataclass(eq=False)
|
||||
class MemoryObjectReceiveStream(ObjectReceiveStream[T_co], Generic[T_co]):
|
||||
_state: _MemoryObjectStreamState[T_co]
|
||||
_closed: bool = field(init=False, default=False)
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
self._state.open_receive_channels += 1
|
||||
|
||||
def receive_nowait(self) -> T_co:
|
||||
"""
|
||||
Receive the next item if it can be done without waiting.
|
||||
|
||||
:return: the received item
|
||||
:raises ~anyio.ClosedResourceError: if this send stream has been closed
|
||||
:raises ~anyio.EndOfStream: if the buffer is empty and this stream has been
|
||||
closed from the sending end
|
||||
:raises ~anyio.WouldBlock: if there are no items in the buffer and no tasks
|
||||
waiting to send
|
||||
|
||||
"""
|
||||
if self._closed:
|
||||
raise ClosedResourceError
|
||||
|
||||
if self._state.waiting_senders:
|
||||
# Get the item from the next sender
|
||||
send_event, item = self._state.waiting_senders.popitem(last=False)
|
||||
self._state.buffer.append(item)
|
||||
send_event.set()
|
||||
|
||||
if self._state.buffer:
|
||||
return self._state.buffer.popleft()
|
||||
elif not self._state.open_send_channels:
|
||||
raise EndOfStream
|
||||
|
||||
raise WouldBlock
|
||||
|
||||
async def receive(self) -> T_co:
|
||||
await checkpoint()
|
||||
try:
|
||||
return self.receive_nowait()
|
||||
except WouldBlock:
|
||||
# Add ourselves in the queue
|
||||
receive_event = Event()
|
||||
receiver = _MemoryObjectItemReceiver[T_co]()
|
||||
self._state.waiting_receivers[receive_event] = receiver
|
||||
|
||||
try:
|
||||
await receive_event.wait()
|
||||
finally:
|
||||
self._state.waiting_receivers.pop(receive_event, None)
|
||||
|
||||
try:
|
||||
return receiver.item
|
||||
except AttributeError:
|
||||
raise EndOfStream from None
|
||||
|
||||
def clone(self) -> MemoryObjectReceiveStream[T_co]:
|
||||
"""
|
||||
Create a clone of this receive stream.
|
||||
|
||||
Each clone can be closed separately. Only when all clones have been closed will
|
||||
the receiving end of the memory stream be considered closed by the sending ends.
|
||||
|
||||
:return: the cloned stream
|
||||
|
||||
"""
|
||||
if self._closed:
|
||||
raise ClosedResourceError
|
||||
|
||||
return MemoryObjectReceiveStream(_state=self._state)
|
||||
|
||||
def close(self) -> None:
|
||||
"""
|
||||
Close the stream.
|
||||
|
||||
This works the exact same way as :meth:`aclose`, but is provided as a special
|
||||
case for the benefit of synchronous callbacks.
|
||||
|
||||
"""
|
||||
if not self._closed:
|
||||
self._closed = True
|
||||
self._state.open_receive_channels -= 1
|
||||
if self._state.open_receive_channels == 0:
|
||||
send_events = list(self._state.waiting_senders.keys())
|
||||
for event in send_events:
|
||||
event.set()
|
||||
|
||||
async def aclose(self) -> None:
|
||||
self.close()
|
||||
|
||||
def statistics(self) -> MemoryObjectStreamStatistics:
|
||||
"""
|
||||
Return statistics about the current state of this stream.
|
||||
|
||||
.. versionadded:: 3.0
|
||||
"""
|
||||
return self._state.statistics()
|
||||
|
||||
def __enter__(self) -> MemoryObjectReceiveStream[T_co]:
|
||||
return self
|
||||
|
||||
def __exit__(
|
||||
self,
|
||||
exc_type: type[BaseException] | None,
|
||||
exc_val: BaseException | None,
|
||||
exc_tb: TracebackType | None,
|
||||
) -> None:
|
||||
self.close()
|
||||
|
||||
def __del__(self) -> None:
|
||||
if not self._closed:
|
||||
warnings.warn(
|
||||
f"Unclosed <{self.__class__.__name__} at {id(self):x}>",
|
||||
ResourceWarning,
|
||||
stacklevel=1,
|
||||
source=self,
|
||||
)
|
||||
|
||||
|
||||
@dataclass(eq=False)
|
||||
class MemoryObjectSendStream(ObjectSendStream[T_contra], Generic[T_contra]):
|
||||
_state: _MemoryObjectStreamState[T_contra]
|
||||
_closed: bool = field(init=False, default=False)
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
self._state.open_send_channels += 1
|
||||
|
||||
def send_nowait(self, item: T_contra) -> None:
|
||||
"""
|
||||
Send an item immediately if it can be done without waiting.
|
||||
|
||||
:param item: the item to send
|
||||
:raises ~anyio.ClosedResourceError: if this send stream has been closed
|
||||
:raises ~anyio.BrokenResourceError: if the stream has been closed from the
|
||||
receiving end
|
||||
:raises ~anyio.WouldBlock: if the buffer is full and there are no tasks waiting
|
||||
to receive
|
||||
|
||||
"""
|
||||
if self._closed:
|
||||
raise ClosedResourceError
|
||||
if not self._state.open_receive_channels:
|
||||
raise BrokenResourceError
|
||||
|
||||
while self._state.waiting_receivers:
|
||||
receive_event, receiver = self._state.waiting_receivers.popitem(last=False)
|
||||
if not receiver.task_info.has_pending_cancellation():
|
||||
receiver.item = item
|
||||
receive_event.set()
|
||||
return
|
||||
|
||||
if len(self._state.buffer) < self._state.max_buffer_size:
|
||||
self._state.buffer.append(item)
|
||||
else:
|
||||
raise WouldBlock
|
||||
|
||||
async def send(self, item: T_contra) -> None:
|
||||
"""
|
||||
Send an item to the stream.
|
||||
|
||||
If the buffer is full, this method blocks until there is again room in the
|
||||
buffer or the item can be sent directly to a receiver.
|
||||
|
||||
:param item: the item to send
|
||||
:raises ~anyio.ClosedResourceError: if this send stream has been closed
|
||||
:raises ~anyio.BrokenResourceError: if the stream has been closed from the
|
||||
receiving end
|
||||
|
||||
"""
|
||||
await checkpoint()
|
||||
try:
|
||||
self.send_nowait(item)
|
||||
except WouldBlock:
|
||||
# Wait until there's someone on the receiving end
|
||||
send_event = Event()
|
||||
self._state.waiting_senders[send_event] = item
|
||||
try:
|
||||
await send_event.wait()
|
||||
except BaseException:
|
||||
self._state.waiting_senders.pop(send_event, None)
|
||||
raise
|
||||
|
||||
if send_event in self._state.waiting_senders:
|
||||
del self._state.waiting_senders[send_event]
|
||||
raise BrokenResourceError from None
|
||||
|
||||
def clone(self) -> MemoryObjectSendStream[T_contra]:
|
||||
"""
|
||||
Create a clone of this send stream.
|
||||
|
||||
Each clone can be closed separately. Only when all clones have been closed will
|
||||
the sending end of the memory stream be considered closed by the receiving ends.
|
||||
|
||||
:return: the cloned stream
|
||||
|
||||
"""
|
||||
if self._closed:
|
||||
raise ClosedResourceError
|
||||
|
||||
return MemoryObjectSendStream(_state=self._state)
|
||||
|
||||
def close(self) -> None:
|
||||
"""
|
||||
Close the stream.
|
||||
|
||||
This works the exact same way as :meth:`aclose`, but is provided as a special
|
||||
case for the benefit of synchronous callbacks.
|
||||
|
||||
"""
|
||||
if not self._closed:
|
||||
self._closed = True
|
||||
self._state.open_send_channels -= 1
|
||||
if self._state.open_send_channels == 0:
|
||||
receive_events = list(self._state.waiting_receivers.keys())
|
||||
self._state.waiting_receivers.clear()
|
||||
for event in receive_events:
|
||||
event.set()
|
||||
|
||||
async def aclose(self) -> None:
|
||||
self.close()
|
||||
|
||||
def statistics(self) -> MemoryObjectStreamStatistics:
|
||||
"""
|
||||
Return statistics about the current state of this stream.
|
||||
|
||||
.. versionadded:: 3.0
|
||||
"""
|
||||
return self._state.statistics()
|
||||
|
||||
def __enter__(self) -> MemoryObjectSendStream[T_contra]:
|
||||
return self
|
||||
|
||||
def __exit__(
|
||||
self,
|
||||
exc_type: type[BaseException] | None,
|
||||
exc_val: BaseException | None,
|
||||
exc_tb: TracebackType | None,
|
||||
) -> None:
|
||||
self.close()
|
||||
|
||||
def __del__(self) -> None:
|
||||
if not self._closed:
|
||||
warnings.warn(
|
||||
f"Unclosed <{self.__class__.__name__} at {id(self):x}>",
|
||||
ResourceWarning,
|
||||
stacklevel=1,
|
||||
source=self,
|
||||
)
|
||||
+160
@@ -0,0 +1,160 @@
|
||||
from __future__ import annotations
|
||||
|
||||
__all__ = (
|
||||
"MultiListener",
|
||||
"StapledByteStream",
|
||||
"StapledObjectStream",
|
||||
)
|
||||
|
||||
from collections.abc import Callable, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Generic, TypeVar
|
||||
|
||||
from ..abc import (
|
||||
ByteReceiveStream,
|
||||
ByteSendStream,
|
||||
ByteStream,
|
||||
Listener,
|
||||
ObjectReceiveStream,
|
||||
ObjectSendStream,
|
||||
ObjectStream,
|
||||
TaskGroup,
|
||||
)
|
||||
|
||||
T_Item = TypeVar("T_Item")
|
||||
T_Stream = TypeVar("T_Stream")
|
||||
|
||||
|
||||
@dataclass(eq=False)
|
||||
class StapledByteStream(ByteStream):
|
||||
"""
|
||||
Combines two byte streams into a single, bidirectional byte stream.
|
||||
|
||||
Extra attributes will be provided from both streams, with the receive stream
|
||||
providing the values in case of a conflict.
|
||||
|
||||
:param ByteSendStream send_stream: the sending byte stream
|
||||
:param ByteReceiveStream receive_stream: the receiving byte stream
|
||||
"""
|
||||
|
||||
send_stream: ByteSendStream
|
||||
receive_stream: ByteReceiveStream
|
||||
|
||||
async def receive(self, max_bytes: int = 65536) -> bytes:
|
||||
if max_bytes < 1:
|
||||
raise ValueError("max_bytes must be a positive integer")
|
||||
|
||||
return await self.receive_stream.receive(max_bytes)
|
||||
|
||||
async def send(self, item: bytes) -> None:
|
||||
await self.send_stream.send(item)
|
||||
|
||||
async def send_eof(self) -> None:
|
||||
await self.send_stream.aclose()
|
||||
|
||||
async def aclose(self) -> None:
|
||||
await self.send_stream.aclose()
|
||||
await self.receive_stream.aclose()
|
||||
|
||||
@property
|
||||
def extra_attributes(self) -> Mapping[Any, Callable[[], Any]]:
|
||||
return {
|
||||
**self.send_stream.extra_attributes,
|
||||
**self.receive_stream.extra_attributes,
|
||||
}
|
||||
|
||||
|
||||
@dataclass(eq=False)
|
||||
class StapledObjectStream(ObjectStream[T_Item], Generic[T_Item]):
|
||||
"""
|
||||
Combines two object streams into a single, bidirectional object stream.
|
||||
|
||||
Extra attributes will be provided from both streams, with the receive stream
|
||||
providing the values in case of a conflict.
|
||||
|
||||
:param ObjectSendStream send_stream: the sending object stream
|
||||
:param ObjectReceiveStream receive_stream: the receiving object stream
|
||||
"""
|
||||
|
||||
send_stream: ObjectSendStream[T_Item]
|
||||
receive_stream: ObjectReceiveStream[T_Item]
|
||||
|
||||
async def receive(self) -> T_Item:
|
||||
return await self.receive_stream.receive()
|
||||
|
||||
async def send(self, item: T_Item) -> None:
|
||||
await self.send_stream.send(item)
|
||||
|
||||
def send_nowait(self, item: T_Item) -> None:
|
||||
try:
|
||||
send_nowait = self.send_stream.send_nowait # type: ignore[attr-defined]
|
||||
except AttributeError as exc:
|
||||
raise NotImplementedError(
|
||||
f"'send_nowait' method not implemented in {type(self.send_stream)}"
|
||||
) from exc
|
||||
|
||||
send_nowait(item)
|
||||
|
||||
async def send_eof(self) -> None:
|
||||
await self.send_stream.aclose()
|
||||
|
||||
async def aclose(self) -> None:
|
||||
await self.send_stream.aclose()
|
||||
await self.receive_stream.aclose()
|
||||
|
||||
@property
|
||||
def extra_attributes(self) -> Mapping[Any, Callable[[], Any]]:
|
||||
return {
|
||||
**self.send_stream.extra_attributes,
|
||||
**self.receive_stream.extra_attributes,
|
||||
}
|
||||
|
||||
|
||||
@dataclass(eq=False)
|
||||
class MultiListener(Listener[T_Stream], Generic[T_Stream]):
|
||||
"""
|
||||
Combines multiple listeners into one, serving connections from all of them at once.
|
||||
|
||||
Any MultiListeners in the given collection of listeners will have their listeners
|
||||
moved into this one.
|
||||
|
||||
Extra attributes are provided from each listener, with each successive listener
|
||||
overriding any conflicting attributes from the previous one.
|
||||
|
||||
:param listeners: listeners to serve
|
||||
:type listeners: Sequence[Listener[T_Stream]]
|
||||
"""
|
||||
|
||||
listeners: Sequence[Listener[T_Stream]]
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
listeners: list[Listener[T_Stream]] = []
|
||||
for listener in self.listeners:
|
||||
if isinstance(listener, MultiListener):
|
||||
listeners.extend(listener.listeners)
|
||||
del listener.listeners[:] # type: ignore[attr-defined]
|
||||
else:
|
||||
listeners.append(listener)
|
||||
|
||||
self.listeners = listeners
|
||||
|
||||
async def serve(
|
||||
self, handler: Callable[[T_Stream], Any], task_group: TaskGroup | None = None
|
||||
) -> None:
|
||||
from .. import create_task_group
|
||||
|
||||
async with create_task_group() as tg:
|
||||
for listener in self.listeners:
|
||||
tg.start_soon(listener.serve, handler, task_group)
|
||||
|
||||
async def aclose(self) -> None:
|
||||
for listener in self.listeners:
|
||||
await listener.aclose()
|
||||
|
||||
@property
|
||||
def extra_attributes(self) -> Mapping[Any, Callable[[], Any]]:
|
||||
attributes: dict = {}
|
||||
for listener in self.listeners:
|
||||
attributes.update(listener.extra_attributes)
|
||||
|
||||
return attributes
|
||||
+176
@@ -0,0 +1,176 @@
|
||||
from __future__ import annotations
|
||||
|
||||
__all__ = (
|
||||
"TextConnectable",
|
||||
"TextReceiveStream",
|
||||
"TextSendStream",
|
||||
"TextStream",
|
||||
)
|
||||
|
||||
import codecs
|
||||
import sys
|
||||
from collections.abc import Callable, Mapping
|
||||
from dataclasses import InitVar, dataclass, field
|
||||
from typing import Any
|
||||
|
||||
from ..abc import (
|
||||
AnyByteReceiveStream,
|
||||
AnyByteSendStream,
|
||||
AnyByteStream,
|
||||
AnyByteStreamConnectable,
|
||||
ObjectReceiveStream,
|
||||
ObjectSendStream,
|
||||
ObjectStream,
|
||||
ObjectStreamConnectable,
|
||||
)
|
||||
|
||||
if sys.version_info >= (3, 12):
|
||||
from typing import override
|
||||
else:
|
||||
from typing_extensions import override
|
||||
|
||||
|
||||
@dataclass(eq=False)
|
||||
class TextReceiveStream(ObjectReceiveStream[str]):
|
||||
"""
|
||||
Stream wrapper that decodes bytes to strings using the given encoding.
|
||||
|
||||
Decoding is done using :class:`~codecs.IncrementalDecoder` which returns any
|
||||
completely received unicode characters as soon as they come in.
|
||||
|
||||
:param transport_stream: any bytes-based receive stream
|
||||
:param encoding: character encoding to use for decoding bytes to strings (defaults
|
||||
to ``utf-8``)
|
||||
:param errors: handling scheme for decoding errors (defaults to ``strict``; see the
|
||||
`codecs module documentation`_ for a comprehensive list of options)
|
||||
|
||||
.. _codecs module documentation:
|
||||
https://docs.python.org/3/library/codecs.html#codec-objects
|
||||
"""
|
||||
|
||||
transport_stream: AnyByteReceiveStream
|
||||
encoding: InitVar[str] = "utf-8"
|
||||
errors: InitVar[str] = "strict"
|
||||
_decoder: codecs.IncrementalDecoder = field(init=False)
|
||||
|
||||
def __post_init__(self, encoding: str, errors: str) -> None:
|
||||
decoder_class = codecs.getincrementaldecoder(encoding)
|
||||
self._decoder = decoder_class(errors=errors)
|
||||
|
||||
async def receive(self) -> str:
|
||||
while True:
|
||||
chunk = await self.transport_stream.receive()
|
||||
decoded = self._decoder.decode(chunk)
|
||||
if decoded:
|
||||
return decoded
|
||||
|
||||
async def aclose(self) -> None:
|
||||
await self.transport_stream.aclose()
|
||||
self._decoder.reset()
|
||||
|
||||
@property
|
||||
def extra_attributes(self) -> Mapping[Any, Callable[[], Any]]:
|
||||
return self.transport_stream.extra_attributes
|
||||
|
||||
|
||||
@dataclass(eq=False)
|
||||
class TextSendStream(ObjectSendStream[str]):
|
||||
"""
|
||||
Sends strings to the wrapped stream as bytes using the given encoding.
|
||||
|
||||
:param AnyByteSendStream transport_stream: any bytes-based send stream
|
||||
:param str encoding: character encoding to use for encoding strings to bytes
|
||||
(defaults to ``utf-8``)
|
||||
:param str errors: handling scheme for encoding errors (defaults to ``strict``; see
|
||||
the `codecs module documentation`_ for a comprehensive list of options)
|
||||
|
||||
.. _codecs module documentation:
|
||||
https://docs.python.org/3/library/codecs.html#codec-objects
|
||||
"""
|
||||
|
||||
transport_stream: AnyByteSendStream
|
||||
encoding: InitVar[str] = "utf-8"
|
||||
errors: str = "strict"
|
||||
_encoder: Callable[..., tuple[bytes, int]] = field(init=False)
|
||||
|
||||
def __post_init__(self, encoding: str) -> None:
|
||||
self._encoder = codecs.getencoder(encoding)
|
||||
|
||||
async def send(self, item: str) -> None:
|
||||
encoded = self._encoder(item, self.errors)[0]
|
||||
await self.transport_stream.send(encoded)
|
||||
|
||||
async def aclose(self) -> None:
|
||||
await self.transport_stream.aclose()
|
||||
|
||||
@property
|
||||
def extra_attributes(self) -> Mapping[Any, Callable[[], Any]]:
|
||||
return self.transport_stream.extra_attributes
|
||||
|
||||
|
||||
@dataclass(eq=False)
|
||||
class TextStream(ObjectStream[str]):
|
||||
"""
|
||||
A bidirectional stream that decodes bytes to strings on receive and encodes strings
|
||||
to bytes on send.
|
||||
|
||||
Extra attributes will be provided from both streams, with the receive stream
|
||||
providing the values in case of a conflict.
|
||||
|
||||
:param AnyByteStream transport_stream: any bytes-based stream
|
||||
:param str encoding: character encoding to use for encoding/decoding strings to/from
|
||||
bytes (defaults to ``utf-8``)
|
||||
:param str errors: handling scheme for encoding errors (defaults to ``strict``; see
|
||||
the `codecs module documentation`_ for a comprehensive list of options)
|
||||
|
||||
.. _codecs module documentation:
|
||||
https://docs.python.org/3/library/codecs.html#codec-objects
|
||||
"""
|
||||
|
||||
transport_stream: AnyByteStream
|
||||
encoding: InitVar[str] = "utf-8"
|
||||
errors: InitVar[str] = "strict"
|
||||
_receive_stream: TextReceiveStream = field(init=False)
|
||||
_send_stream: TextSendStream = field(init=False)
|
||||
|
||||
def __post_init__(self, encoding: str, errors: str) -> None:
|
||||
self._receive_stream = TextReceiveStream(
|
||||
self.transport_stream, encoding=encoding, errors=errors
|
||||
)
|
||||
self._send_stream = TextSendStream(
|
||||
self.transport_stream, encoding=encoding, errors=errors
|
||||
)
|
||||
|
||||
async def receive(self) -> str:
|
||||
return await self._receive_stream.receive()
|
||||
|
||||
async def send(self, item: str) -> None:
|
||||
await self._send_stream.send(item)
|
||||
|
||||
async def send_eof(self) -> None:
|
||||
await self.transport_stream.send_eof()
|
||||
|
||||
async def aclose(self) -> None:
|
||||
await self._send_stream.aclose()
|
||||
await self._receive_stream.aclose()
|
||||
|
||||
@property
|
||||
def extra_attributes(self) -> Mapping[Any, Callable[[], Any]]:
|
||||
return {
|
||||
**self._send_stream.extra_attributes,
|
||||
**self._receive_stream.extra_attributes,
|
||||
}
|
||||
|
||||
|
||||
class TextConnectable(ObjectStreamConnectable[str]):
|
||||
def __init__(self, connectable: AnyByteStreamConnectable):
|
||||
"""
|
||||
:param connectable: the bytestream endpoint to wrap
|
||||
|
||||
"""
|
||||
self.connectable = connectable
|
||||
|
||||
@override
|
||||
async def connect(self) -> TextStream:
|
||||
stream = await self.connectable.connect()
|
||||
return TextStream(stream)
|
||||
+436
@@ -0,0 +1,436 @@
|
||||
from __future__ import annotations
|
||||
|
||||
__all__ = (
|
||||
"TLSAttribute",
|
||||
"TLSConnectable",
|
||||
"TLSListener",
|
||||
"TLSStream",
|
||||
)
|
||||
|
||||
import logging
|
||||
import re
|
||||
import ssl
|
||||
import sys
|
||||
from collections.abc import Callable, Mapping
|
||||
from dataclasses import dataclass
|
||||
from functools import wraps
|
||||
from ssl import SSLContext
|
||||
from typing import Any, TypeAlias, TypeVar
|
||||
|
||||
from .. import (
|
||||
BrokenResourceError,
|
||||
EndOfStream,
|
||||
aclose_forcefully,
|
||||
get_cancelled_exc_class,
|
||||
to_thread,
|
||||
)
|
||||
from .._core._typedattr import TypedAttributeSet, typed_attribute
|
||||
from ..abc import (
|
||||
AnyByteStream,
|
||||
AnyByteStreamConnectable,
|
||||
ByteStream,
|
||||
ByteStreamConnectable,
|
||||
Listener,
|
||||
TaskGroup,
|
||||
)
|
||||
|
||||
if sys.version_info >= (3, 11):
|
||||
from typing import TypeVarTuple, Unpack
|
||||
else:
|
||||
from typing_extensions import TypeVarTuple, Unpack
|
||||
|
||||
if sys.version_info >= (3, 12):
|
||||
from typing import override
|
||||
else:
|
||||
from typing_extensions import override
|
||||
|
||||
T_Retval = TypeVar("T_Retval")
|
||||
PosArgsT = TypeVarTuple("PosArgsT")
|
||||
_PCTRTT: TypeAlias = tuple[tuple[str, str], ...]
|
||||
_PCTRTTT: TypeAlias = tuple[_PCTRTT, ...]
|
||||
|
||||
|
||||
class TLSAttribute(TypedAttributeSet):
|
||||
"""Contains Transport Layer Security related attributes."""
|
||||
|
||||
#: the selected ALPN protocol
|
||||
alpn_protocol: str | None = typed_attribute()
|
||||
#: the channel binding for type ``tls-unique``
|
||||
channel_binding_tls_unique: bytes = typed_attribute()
|
||||
#: the selected cipher
|
||||
cipher: tuple[str, str, int] = typed_attribute()
|
||||
#: the peer certificate in dictionary form (see :meth:`ssl.SSLSocket.getpeercert`
|
||||
# for more information)
|
||||
peer_certificate: None | (dict[str, str | _PCTRTTT | _PCTRTT]) = typed_attribute()
|
||||
#: the peer certificate in binary form
|
||||
peer_certificate_binary: bytes | None = typed_attribute()
|
||||
#: ``True`` if this is the server side of the connection
|
||||
server_side: bool = typed_attribute()
|
||||
#: ciphers shared by the client during the TLS handshake (``None`` if this is the
|
||||
#: client side)
|
||||
shared_ciphers: list[tuple[str, str, int]] | None = typed_attribute()
|
||||
#: the :class:`~ssl.SSLObject` used for encryption
|
||||
ssl_object: ssl.SSLObject = typed_attribute()
|
||||
#: ``True`` if this stream does (and expects) a closing TLS handshake when the
|
||||
#: stream is being closed
|
||||
standard_compatible: bool = typed_attribute()
|
||||
#: the TLS protocol version (e.g. ``TLSv1.2``)
|
||||
tls_version: str = typed_attribute()
|
||||
|
||||
|
||||
@dataclass(eq=False)
|
||||
class TLSStream(ByteStream):
|
||||
"""
|
||||
A stream wrapper that encrypts all sent data and decrypts received data.
|
||||
|
||||
This class has no public initializer; use :meth:`wrap` instead.
|
||||
All extra attributes from :class:`~TLSAttribute` are supported.
|
||||
|
||||
:var AnyByteStream transport_stream: the wrapped stream
|
||||
|
||||
"""
|
||||
|
||||
transport_stream: AnyByteStream
|
||||
standard_compatible: bool
|
||||
_ssl_object: ssl.SSLObject
|
||||
_read_bio: ssl.MemoryBIO
|
||||
_write_bio: ssl.MemoryBIO
|
||||
|
||||
@classmethod
|
||||
async def wrap(
|
||||
cls,
|
||||
transport_stream: AnyByteStream,
|
||||
*,
|
||||
server_side: bool | None = None,
|
||||
hostname: str | None = None,
|
||||
ssl_context: ssl.SSLContext | None = None,
|
||||
standard_compatible: bool = True,
|
||||
) -> TLSStream:
|
||||
"""
|
||||
Wrap an existing stream with Transport Layer Security.
|
||||
|
||||
This performs a TLS handshake with the peer.
|
||||
|
||||
:param transport_stream: a bytes-transporting stream to wrap
|
||||
:param server_side: ``True`` if this is the server side of the connection,
|
||||
``False`` if this is the client side (if omitted, will be set to ``False``
|
||||
if ``hostname`` has been provided, ``False`` otherwise). Used only to create
|
||||
a default context when an explicit context has not been provided.
|
||||
:param hostname: host name of the peer (if host name checking is desired)
|
||||
:param ssl_context: the SSLContext object to use (if not provided, a secure
|
||||
default will be created)
|
||||
:param standard_compatible: if ``False``, skip the closing handshake when
|
||||
closing the connection, and don't raise an exception if the peer does the
|
||||
same
|
||||
:raises ~ssl.SSLError: if the TLS handshake fails
|
||||
|
||||
"""
|
||||
if server_side is None:
|
||||
server_side = not hostname
|
||||
|
||||
if not ssl_context:
|
||||
purpose = (
|
||||
ssl.Purpose.CLIENT_AUTH if server_side else ssl.Purpose.SERVER_AUTH
|
||||
)
|
||||
ssl_context = ssl.create_default_context(purpose)
|
||||
|
||||
# Re-enable detection of unexpected EOFs if it was disabled by Python
|
||||
if hasattr(ssl, "OP_IGNORE_UNEXPECTED_EOF"):
|
||||
ssl_context.options &= ~ssl.OP_IGNORE_UNEXPECTED_EOF
|
||||
|
||||
bio_in = ssl.MemoryBIO()
|
||||
bio_out = ssl.MemoryBIO()
|
||||
|
||||
# Resolve international host names using IDNA 2008.
|
||||
# Otherwise wrap_bio() would resolve them with IDNA 2003.
|
||||
if hostname is not None:
|
||||
from .._core._sockets import idna2008_resolve
|
||||
|
||||
server_hostname: bytes | None = idna2008_resolve(hostname)
|
||||
else:
|
||||
server_hostname = None
|
||||
|
||||
# External SSLContext implementations may do blocking I/O in wrap_bio(),
|
||||
# but the standard library implementation won't
|
||||
if type(ssl_context) is ssl.SSLContext:
|
||||
ssl_object = ssl_context.wrap_bio(
|
||||
bio_in,
|
||||
bio_out,
|
||||
server_side=server_side,
|
||||
server_hostname=server_hostname,
|
||||
)
|
||||
else:
|
||||
ssl_object = await to_thread.run_sync(
|
||||
ssl_context.wrap_bio,
|
||||
bio_in,
|
||||
bio_out,
|
||||
server_side,
|
||||
server_hostname,
|
||||
None,
|
||||
)
|
||||
|
||||
wrapper = cls(
|
||||
transport_stream=transport_stream,
|
||||
standard_compatible=standard_compatible,
|
||||
_ssl_object=ssl_object,
|
||||
_read_bio=bio_in,
|
||||
_write_bio=bio_out,
|
||||
)
|
||||
await wrapper._call_sslobject_method(ssl_object.do_handshake)
|
||||
return wrapper
|
||||
|
||||
async def _call_sslobject_method(
|
||||
self, func: Callable[[Unpack[PosArgsT]], T_Retval], *args: Unpack[PosArgsT]
|
||||
) -> T_Retval:
|
||||
while True:
|
||||
try:
|
||||
result = func(*args)
|
||||
except ssl.SSLWantReadError:
|
||||
try:
|
||||
# Flush any pending writes first
|
||||
if self._write_bio.pending:
|
||||
await self.transport_stream.send(self._write_bio.read())
|
||||
|
||||
data = await self.transport_stream.receive()
|
||||
except EndOfStream:
|
||||
self._read_bio.write_eof()
|
||||
except OSError as exc:
|
||||
self._read_bio.write_eof()
|
||||
self._write_bio.write_eof()
|
||||
raise BrokenResourceError from exc
|
||||
else:
|
||||
self._read_bio.write(data)
|
||||
except ssl.SSLWantWriteError:
|
||||
await self.transport_stream.send(self._write_bio.read())
|
||||
except ssl.SSLSyscallError as exc:
|
||||
self._read_bio.write_eof()
|
||||
self._write_bio.write_eof()
|
||||
raise BrokenResourceError from exc
|
||||
except ssl.SSLError as exc:
|
||||
self._read_bio.write_eof()
|
||||
self._write_bio.write_eof()
|
||||
if isinstance(exc, ssl.SSLEOFError) or (
|
||||
exc.strerror and "UNEXPECTED_EOF_WHILE_READING" in exc.strerror
|
||||
):
|
||||
if self.standard_compatible:
|
||||
raise BrokenResourceError from exc
|
||||
else:
|
||||
raise EndOfStream from None
|
||||
|
||||
raise
|
||||
else:
|
||||
# Flush any pending writes first
|
||||
if self._write_bio.pending:
|
||||
await self.transport_stream.send(self._write_bio.read())
|
||||
|
||||
return result
|
||||
|
||||
async def unwrap(self) -> tuple[AnyByteStream, bytes]:
|
||||
"""
|
||||
Does the TLS closing handshake.
|
||||
|
||||
:return: a tuple of (wrapped byte stream, bytes left in the read buffer)
|
||||
|
||||
"""
|
||||
await self._call_sslobject_method(self._ssl_object.unwrap)
|
||||
self._read_bio.write_eof()
|
||||
self._write_bio.write_eof()
|
||||
return self.transport_stream, self._read_bio.read()
|
||||
|
||||
async def aclose(self) -> None:
|
||||
if self.standard_compatible:
|
||||
try:
|
||||
await self.unwrap()
|
||||
except BaseException:
|
||||
await aclose_forcefully(self.transport_stream)
|
||||
raise
|
||||
|
||||
await self.transport_stream.aclose()
|
||||
|
||||
async def receive(self, max_bytes: int = 65536) -> bytes:
|
||||
if max_bytes < 1:
|
||||
raise ValueError("max_bytes must be a positive integer")
|
||||
|
||||
data = await self._call_sslobject_method(self._ssl_object.read, max_bytes)
|
||||
if not data:
|
||||
raise EndOfStream
|
||||
|
||||
return data
|
||||
|
||||
async def send(self, item: bytes) -> None:
|
||||
await self._call_sslobject_method(self._ssl_object.write, item)
|
||||
|
||||
async def send_eof(self) -> None:
|
||||
tls_version = self.extra(TLSAttribute.tls_version)
|
||||
match = re.match(r"TLSv(\d+)(?:\.(\d+))?", tls_version)
|
||||
if match:
|
||||
major, minor = int(match.group(1)), int(match.group(2) or 0)
|
||||
if (major, minor) < (1, 3):
|
||||
raise NotImplementedError(
|
||||
f"send_eof() requires at least TLSv1.3; current "
|
||||
f"session uses {tls_version}"
|
||||
)
|
||||
|
||||
raise NotImplementedError(
|
||||
"send_eof() has not yet been implemented for TLS streams"
|
||||
)
|
||||
|
||||
@property
|
||||
def extra_attributes(self) -> Mapping[Any, Callable[[], Any]]:
|
||||
return {
|
||||
**self.transport_stream.extra_attributes,
|
||||
TLSAttribute.alpn_protocol: self._ssl_object.selected_alpn_protocol,
|
||||
TLSAttribute.channel_binding_tls_unique: (
|
||||
self._ssl_object.get_channel_binding
|
||||
),
|
||||
TLSAttribute.cipher: self._ssl_object.cipher,
|
||||
TLSAttribute.peer_certificate: lambda: self._ssl_object.getpeercert(False),
|
||||
TLSAttribute.peer_certificate_binary: lambda: self._ssl_object.getpeercert(
|
||||
True
|
||||
),
|
||||
TLSAttribute.server_side: lambda: self._ssl_object.server_side,
|
||||
TLSAttribute.shared_ciphers: lambda: (
|
||||
self._ssl_object.shared_ciphers()
|
||||
if self._ssl_object.server_side
|
||||
else None
|
||||
),
|
||||
TLSAttribute.standard_compatible: lambda: self.standard_compatible,
|
||||
TLSAttribute.ssl_object: lambda: self._ssl_object,
|
||||
TLSAttribute.tls_version: self._ssl_object.version,
|
||||
}
|
||||
|
||||
|
||||
@dataclass(eq=False)
|
||||
class TLSListener(Listener[TLSStream]):
|
||||
"""
|
||||
A convenience listener that wraps another listener and auto-negotiates a TLS session
|
||||
on every accepted connection.
|
||||
|
||||
If the TLS handshake times out or raises an exception,
|
||||
:meth:`handle_handshake_error` is called to do whatever post-mortem processing is
|
||||
deemed necessary.
|
||||
|
||||
Supports only the :attr:`~TLSAttribute.standard_compatible` extra attribute.
|
||||
|
||||
:param Listener listener: the listener to wrap
|
||||
:param ssl_context: the SSL context object
|
||||
:param standard_compatible: a flag passed through to :meth:`TLSStream.wrap`
|
||||
:param handshake_timeout: time limit for the TLS handshake
|
||||
(passed to :func:`~anyio.fail_after`)
|
||||
"""
|
||||
|
||||
listener: Listener[Any]
|
||||
ssl_context: ssl.SSLContext
|
||||
standard_compatible: bool = True
|
||||
handshake_timeout: float = 30
|
||||
|
||||
@staticmethod
|
||||
async def handle_handshake_error(exc: BaseException, stream: AnyByteStream) -> None:
|
||||
"""
|
||||
Handle an exception raised during the TLS handshake.
|
||||
|
||||
This method does 3 things:
|
||||
|
||||
#. Forcefully closes the original stream
|
||||
#. Logs the exception (unless it was a cancellation exception) using the
|
||||
``anyio.streams.tls`` logger
|
||||
#. Reraises the exception if it was a base exception or a cancellation exception
|
||||
|
||||
:param exc: the exception
|
||||
:param stream: the original stream
|
||||
|
||||
"""
|
||||
await aclose_forcefully(stream)
|
||||
|
||||
# Log all except cancellation exceptions
|
||||
if not isinstance(exc, get_cancelled_exc_class()):
|
||||
# CPython (as of 3.11.5) returns incorrect `sys.exc_info()` here when using
|
||||
# any asyncio implementation, so we explicitly pass the exception to log
|
||||
# (https://github.com/python/cpython/issues/108668). Trio does not have this
|
||||
# issue because it works around the CPython bug.
|
||||
logging.getLogger(__name__).exception(
|
||||
"Error during TLS handshake", exc_info=exc
|
||||
)
|
||||
|
||||
# Only reraise base exceptions and cancellation exceptions
|
||||
if not isinstance(exc, Exception) or isinstance(exc, get_cancelled_exc_class()):
|
||||
raise # noqa: PLE0704
|
||||
|
||||
async def serve(
|
||||
self,
|
||||
handler: Callable[[TLSStream], Any],
|
||||
task_group: TaskGroup | None = None,
|
||||
) -> None:
|
||||
@wraps(handler)
|
||||
async def handler_wrapper(stream: AnyByteStream) -> None:
|
||||
from .. import fail_after
|
||||
|
||||
try:
|
||||
with fail_after(self.handshake_timeout):
|
||||
wrapped_stream = await TLSStream.wrap(
|
||||
stream,
|
||||
ssl_context=self.ssl_context,
|
||||
standard_compatible=self.standard_compatible,
|
||||
)
|
||||
except BaseException as exc:
|
||||
await self.handle_handshake_error(exc, stream)
|
||||
else:
|
||||
await handler(wrapped_stream)
|
||||
|
||||
await self.listener.serve(handler_wrapper, task_group)
|
||||
|
||||
async def aclose(self) -> None:
|
||||
await self.listener.aclose()
|
||||
|
||||
@property
|
||||
def extra_attributes(self) -> Mapping[Any, Callable[[], Any]]:
|
||||
return {
|
||||
TLSAttribute.standard_compatible: lambda: self.standard_compatible,
|
||||
}
|
||||
|
||||
|
||||
class TLSConnectable(ByteStreamConnectable):
|
||||
"""
|
||||
Wraps another connectable and does TLS negotiation after a successful connection.
|
||||
|
||||
:param connectable: the connectable to wrap
|
||||
:param hostname: host name of the server (if host name checking is desired)
|
||||
:param ssl_context: the SSLContext object to use (if not provided, a secure default
|
||||
will be created)
|
||||
:param standard_compatible: if ``False``, skip the closing handshake when closing
|
||||
the connection, and don't raise an exception if the server does the same
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
connectable: AnyByteStreamConnectable,
|
||||
*,
|
||||
hostname: str | None = None,
|
||||
ssl_context: ssl.SSLContext | None = None,
|
||||
standard_compatible: bool = True,
|
||||
) -> None:
|
||||
self.connectable = connectable
|
||||
self.ssl_context: SSLContext = ssl_context or ssl.create_default_context(
|
||||
ssl.Purpose.SERVER_AUTH
|
||||
)
|
||||
if not isinstance(self.ssl_context, ssl.SSLContext):
|
||||
raise TypeError(
|
||||
"ssl_context must be an instance of ssl.SSLContext, not "
|
||||
f"{type(self.ssl_context).__name__}"
|
||||
)
|
||||
self.hostname = hostname
|
||||
self.standard_compatible = standard_compatible
|
||||
|
||||
@override
|
||||
async def connect(self) -> TLSStream:
|
||||
stream = await self.connectable.connect()
|
||||
try:
|
||||
return await TLSStream.wrap(
|
||||
stream,
|
||||
hostname=self.hostname,
|
||||
ssl_context=self.ssl_context,
|
||||
standard_compatible=self.standard_compatible,
|
||||
)
|
||||
except BaseException:
|
||||
await aclose_forcefully(stream)
|
||||
raise
|
||||
+246
@@ -0,0 +1,246 @@
|
||||
from __future__ import annotations
|
||||
|
||||
__all__ = (
|
||||
"current_default_interpreter_limiter",
|
||||
"run_sync",
|
||||
)
|
||||
|
||||
import atexit
|
||||
import os
|
||||
import sys
|
||||
from collections import deque
|
||||
from collections.abc import Callable
|
||||
from typing import Any, Final, TypeVar
|
||||
|
||||
from . import current_time, to_thread
|
||||
from ._core._exceptions import BrokenWorkerInterpreter
|
||||
from ._core._synchronization import CapacityLimiter
|
||||
from .lowlevel import RunVar
|
||||
|
||||
if sys.version_info >= (3, 11):
|
||||
from typing import TypeVarTuple, Unpack
|
||||
else:
|
||||
from typing_extensions import TypeVarTuple, Unpack
|
||||
|
||||
if sys.version_info >= (3, 14):
|
||||
from concurrent.interpreters import ExecutionFailed, create
|
||||
|
||||
def _interp_call(
|
||||
func: Callable[..., Any], args: tuple[Any, ...]
|
||||
) -> tuple[Any, bool]:
|
||||
try:
|
||||
retval = func(*args)
|
||||
except BaseException as exc:
|
||||
return exc, True
|
||||
else:
|
||||
return retval, False
|
||||
|
||||
class _Worker:
|
||||
last_used: float = 0
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._interpreter = create()
|
||||
|
||||
def destroy(self) -> None:
|
||||
self._interpreter.close()
|
||||
|
||||
def call(
|
||||
self,
|
||||
func: Callable[..., T_Retval],
|
||||
args: tuple[Any, ...],
|
||||
) -> T_Retval:
|
||||
try:
|
||||
res, is_exception = self._interpreter.call(_interp_call, func, args)
|
||||
except ExecutionFailed as exc:
|
||||
raise BrokenWorkerInterpreter(exc.excinfo) from exc
|
||||
|
||||
if is_exception:
|
||||
raise res
|
||||
|
||||
return res
|
||||
elif sys.version_info >= (3, 13):
|
||||
import _interpqueues
|
||||
import _interpreters
|
||||
|
||||
UNBOUND: Final = 2 # I have no clue how this works, but it was used in the stdlib
|
||||
FMT_UNPICKLED: Final = 0
|
||||
FMT_PICKLED: Final = 1
|
||||
QUEUE_PICKLE_ARGS: Final = (FMT_PICKLED, UNBOUND)
|
||||
QUEUE_UNPICKLE_ARGS: Final = (FMT_UNPICKLED, UNBOUND)
|
||||
|
||||
_run_func = compile(
|
||||
"""
|
||||
import _interpqueues
|
||||
from _interpreters import NotShareableError
|
||||
from pickle import loads, dumps, HIGHEST_PROTOCOL
|
||||
|
||||
QUEUE_PICKLE_ARGS = (1, 2)
|
||||
QUEUE_UNPICKLE_ARGS = (0, 2)
|
||||
|
||||
item = _interpqueues.get(queue_id)[0]
|
||||
try:
|
||||
func, args = loads(item)
|
||||
retval = func(*args)
|
||||
except BaseException as exc:
|
||||
is_exception = True
|
||||
retval = exc
|
||||
else:
|
||||
is_exception = False
|
||||
|
||||
try:
|
||||
_interpqueues.put(queue_id, (retval, is_exception), *QUEUE_UNPICKLE_ARGS)
|
||||
except NotShareableError:
|
||||
retval = dumps(retval, HIGHEST_PROTOCOL)
|
||||
_interpqueues.put(queue_id, (retval, is_exception), *QUEUE_PICKLE_ARGS)
|
||||
""",
|
||||
"<string>",
|
||||
"exec",
|
||||
)
|
||||
|
||||
class _Worker:
|
||||
last_used: float = 0
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._interpreter_id = _interpreters.create()
|
||||
self._queue_id = _interpqueues.create(1, *QUEUE_UNPICKLE_ARGS)
|
||||
_interpreters.set___main___attrs(
|
||||
self._interpreter_id, {"queue_id": self._queue_id}
|
||||
)
|
||||
|
||||
def destroy(self) -> None:
|
||||
_interpqueues.destroy(self._queue_id)
|
||||
_interpreters.destroy(self._interpreter_id)
|
||||
|
||||
def call(
|
||||
self,
|
||||
func: Callable[..., T_Retval],
|
||||
args: tuple[Any, ...],
|
||||
) -> T_Retval:
|
||||
import pickle
|
||||
|
||||
item = pickle.dumps((func, args), pickle.HIGHEST_PROTOCOL)
|
||||
_interpqueues.put(self._queue_id, item, *QUEUE_PICKLE_ARGS)
|
||||
exc_info = _interpreters.exec(self._interpreter_id, _run_func)
|
||||
if exc_info:
|
||||
raise BrokenWorkerInterpreter(exc_info)
|
||||
|
||||
res = _interpqueues.get(self._queue_id)
|
||||
(res, is_exception), fmt = res[:2]
|
||||
if fmt == FMT_PICKLED:
|
||||
res = pickle.loads(res)
|
||||
|
||||
if is_exception:
|
||||
raise res
|
||||
|
||||
return res
|
||||
else:
|
||||
|
||||
class _Worker:
|
||||
last_used: float = 0
|
||||
|
||||
def __init__(self) -> None:
|
||||
raise RuntimeError("subinterpreters require at least Python 3.13")
|
||||
|
||||
def call(
|
||||
self,
|
||||
func: Callable[..., T_Retval],
|
||||
args: tuple[Any, ...],
|
||||
) -> T_Retval:
|
||||
raise NotImplementedError
|
||||
|
||||
def destroy(self) -> None:
|
||||
pass
|
||||
|
||||
|
||||
DEFAULT_CPU_COUNT: Final = 8 # this is just an arbitrarily selected value
|
||||
MAX_WORKER_IDLE_TIME = (
|
||||
30 # seconds a subinterpreter can be idle before becoming eligible for pruning
|
||||
)
|
||||
|
||||
T_Retval = TypeVar("T_Retval")
|
||||
PosArgsT = TypeVarTuple("PosArgsT")
|
||||
|
||||
_idle_workers = RunVar[deque[_Worker]]("_available_workers")
|
||||
_default_interpreter_limiter = RunVar[CapacityLimiter]("_default_interpreter_limiter")
|
||||
|
||||
|
||||
def _stop_workers(workers: deque[_Worker]) -> None:
|
||||
for worker in workers:
|
||||
worker.destroy()
|
||||
|
||||
workers.clear()
|
||||
|
||||
|
||||
async def run_sync(
|
||||
func: Callable[[Unpack[PosArgsT]], T_Retval],
|
||||
*args: Unpack[PosArgsT],
|
||||
limiter: CapacityLimiter | None = None,
|
||||
) -> T_Retval:
|
||||
"""
|
||||
Call the given function with the given arguments in a subinterpreter.
|
||||
|
||||
.. warning:: On Python 3.13, the :mod:`concurrent.interpreters` module was not yet
|
||||
available, so the code path for that Python version relies on an undocumented,
|
||||
private API. As such, it is recommended to not rely on this function for anything
|
||||
mission-critical on Python 3.13.
|
||||
|
||||
:param func: a callable
|
||||
:param args: the positional arguments for the callable
|
||||
:param limiter: capacity limiter to use to limit the total number of subinterpreters
|
||||
running (if omitted, the default limiter is used)
|
||||
:return: the result of the call
|
||||
:raises BrokenWorkerInterpreter: if there's an internal error in a subinterpreter
|
||||
|
||||
"""
|
||||
if limiter is None:
|
||||
limiter = current_default_interpreter_limiter()
|
||||
|
||||
try:
|
||||
idle_workers = _idle_workers.get()
|
||||
except LookupError:
|
||||
idle_workers = deque()
|
||||
_idle_workers.set(idle_workers)
|
||||
atexit.register(_stop_workers, idle_workers)
|
||||
|
||||
async with limiter:
|
||||
try:
|
||||
worker = idle_workers.pop()
|
||||
except IndexError:
|
||||
worker = _Worker()
|
||||
|
||||
try:
|
||||
return await to_thread.run_sync(
|
||||
worker.call,
|
||||
func,
|
||||
args,
|
||||
limiter=limiter,
|
||||
)
|
||||
finally:
|
||||
# Prune workers that have been idle for too long
|
||||
now = current_time()
|
||||
while idle_workers:
|
||||
if now - idle_workers[0].last_used <= MAX_WORKER_IDLE_TIME:
|
||||
break
|
||||
|
||||
await to_thread.run_sync(idle_workers.popleft().destroy, limiter=limiter)
|
||||
|
||||
worker.last_used = current_time()
|
||||
idle_workers.append(worker)
|
||||
|
||||
|
||||
def current_default_interpreter_limiter() -> CapacityLimiter:
|
||||
"""
|
||||
Return the capacity limiter used by default to limit the number of concurrently
|
||||
running subinterpreters.
|
||||
|
||||
Defaults to the number of CPU cores.
|
||||
|
||||
:return: a capacity limiter object
|
||||
|
||||
"""
|
||||
try:
|
||||
return _default_interpreter_limiter.get()
|
||||
except LookupError:
|
||||
limiter = CapacityLimiter(os.cpu_count() or DEFAULT_CPU_COUNT)
|
||||
_default_interpreter_limiter.set(limiter)
|
||||
return limiter
|
||||
+269
@@ -0,0 +1,269 @@
|
||||
from __future__ import annotations
|
||||
|
||||
__all__ = (
|
||||
"current_default_process_limiter",
|
||||
"process_worker",
|
||||
"run_sync",
|
||||
)
|
||||
|
||||
import os
|
||||
import pickle
|
||||
import runpy
|
||||
import subprocess
|
||||
import sys
|
||||
from collections import deque
|
||||
from collections.abc import Callable
|
||||
from types import ModuleType
|
||||
from typing import TypeVar, cast
|
||||
|
||||
from ._core._eventloop import current_time, get_async_backend, get_cancelled_exc_class
|
||||
from ._core._exceptions import BrokenWorkerProcess
|
||||
from ._core._subprocesses import open_process
|
||||
from ._core._synchronization import CapacityLimiter
|
||||
from ._core._tasks import CancelScope, fail_after
|
||||
from .abc import ByteReceiveStream, ByteSendStream, Process
|
||||
from .lowlevel import RunVar, checkpoint_if_cancelled
|
||||
from .streams.buffered import BufferedByteReceiveStream
|
||||
|
||||
if sys.version_info >= (3, 11):
|
||||
from typing import TypeVarTuple, Unpack
|
||||
else:
|
||||
from typing_extensions import TypeVarTuple, Unpack
|
||||
|
||||
WORKER_MAX_IDLE_TIME = 300 # 5 minutes
|
||||
|
||||
T_Retval = TypeVar("T_Retval")
|
||||
PosArgsT = TypeVarTuple("PosArgsT")
|
||||
|
||||
_process_pool_workers: RunVar[set[Process]] = RunVar("_process_pool_workers")
|
||||
_process_pool_idle_workers: RunVar[deque[tuple[Process, float]]] = RunVar(
|
||||
"_process_pool_idle_workers"
|
||||
)
|
||||
_default_process_limiter: RunVar[CapacityLimiter] = RunVar("_default_process_limiter")
|
||||
|
||||
|
||||
async def run_sync( # type: ignore[return]
|
||||
func: Callable[[Unpack[PosArgsT]], T_Retval],
|
||||
*args: Unpack[PosArgsT],
|
||||
cancellable: bool = False,
|
||||
limiter: CapacityLimiter | None = None,
|
||||
) -> T_Retval:
|
||||
"""
|
||||
Call the given function with the given arguments in a worker process.
|
||||
|
||||
If the ``cancellable`` option is enabled and the task waiting for its completion is
|
||||
cancelled, the worker process running it will be abruptly terminated using SIGKILL
|
||||
(or ``terminateProcess()`` on Windows).
|
||||
|
||||
:param func: a callable
|
||||
:param args: positional arguments for the callable
|
||||
:param cancellable: ``True`` to allow cancellation of the operation while it's
|
||||
running
|
||||
:param limiter: capacity limiter to use to limit the total amount of processes
|
||||
running (if omitted, the default limiter is used)
|
||||
:raises NoEventLoopError: if no supported asynchronous event loop is running in the
|
||||
current thread
|
||||
:return: an awaitable that yields the return value of the function.
|
||||
|
||||
"""
|
||||
|
||||
async def send_raw_command(pickled_cmd: bytes) -> object:
|
||||
try:
|
||||
await stdin.send(pickled_cmd)
|
||||
response = await buffered.receive_until(b"\n", 50)
|
||||
status, length = response.split(b" ")
|
||||
if status not in (b"RETURN", b"EXCEPTION"):
|
||||
raise RuntimeError(
|
||||
f"Worker process returned unexpected response: {response!r}"
|
||||
)
|
||||
|
||||
pickled_response = await buffered.receive_exactly(int(length))
|
||||
except BaseException as exc:
|
||||
workers.discard(process)
|
||||
try:
|
||||
process.kill()
|
||||
with CancelScope(shield=True):
|
||||
await process.aclose()
|
||||
except ProcessLookupError:
|
||||
pass
|
||||
|
||||
if isinstance(exc, get_cancelled_exc_class()):
|
||||
raise
|
||||
else:
|
||||
raise BrokenWorkerProcess from exc
|
||||
|
||||
retval = pickle.loads(pickled_response)
|
||||
if status == b"EXCEPTION":
|
||||
assert isinstance(retval, BaseException)
|
||||
raise retval
|
||||
else:
|
||||
return retval
|
||||
|
||||
# First pickle the request before trying to reserve a worker process
|
||||
await checkpoint_if_cancelled()
|
||||
request = pickle.dumps(("run", func, args), protocol=pickle.HIGHEST_PROTOCOL)
|
||||
|
||||
# If this is the first run in this event loop thread, set up the necessary variables
|
||||
try:
|
||||
workers = _process_pool_workers.get()
|
||||
idle_workers = _process_pool_idle_workers.get()
|
||||
except LookupError:
|
||||
workers = set()
|
||||
idle_workers = deque()
|
||||
_process_pool_workers.set(workers)
|
||||
_process_pool_idle_workers.set(idle_workers)
|
||||
get_async_backend().setup_process_pool_exit_at_shutdown(workers)
|
||||
|
||||
async with limiter or current_default_process_limiter():
|
||||
# Pop processes from the pool (starting from the most recently used) until we
|
||||
# find one that hasn't exited yet
|
||||
process: Process
|
||||
while idle_workers:
|
||||
process, _idle_since = idle_workers.pop()
|
||||
if process.returncode is None:
|
||||
stdin = cast(ByteSendStream, process.stdin)
|
||||
buffered = BufferedByteReceiveStream(
|
||||
cast(ByteReceiveStream, process.stdout)
|
||||
)
|
||||
|
||||
# Prune any other workers that have been idle for WORKER_MAX_IDLE_TIME
|
||||
# seconds or longer
|
||||
now = current_time()
|
||||
killed_processes: list[Process] = []
|
||||
while idle_workers:
|
||||
if now - idle_workers[0][1] < WORKER_MAX_IDLE_TIME:
|
||||
break
|
||||
|
||||
process_to_kill, _idle_since = idle_workers.popleft()
|
||||
process_to_kill.kill()
|
||||
workers.remove(process_to_kill)
|
||||
killed_processes.append(process_to_kill)
|
||||
|
||||
with CancelScope(shield=True):
|
||||
for killed_process in killed_processes:
|
||||
await killed_process.aclose()
|
||||
|
||||
break
|
||||
|
||||
workers.remove(process)
|
||||
else:
|
||||
command = [sys.executable, "-u", "-m", __name__]
|
||||
process = await open_process(
|
||||
command, stdin=subprocess.PIPE, stdout=subprocess.PIPE
|
||||
)
|
||||
try:
|
||||
stdin = cast(ByteSendStream, process.stdin)
|
||||
buffered = BufferedByteReceiveStream(
|
||||
cast(ByteReceiveStream, process.stdout)
|
||||
)
|
||||
with fail_after(20):
|
||||
message = await buffered.receive(6)
|
||||
|
||||
if message != b"READY\n":
|
||||
raise BrokenWorkerProcess(
|
||||
f"Worker process returned unexpected response: {message!r}"
|
||||
)
|
||||
|
||||
main_module_path = getattr(sys.modules["__main__"], "__file__", None)
|
||||
pickled = pickle.dumps(
|
||||
("init", sys.path, main_module_path),
|
||||
protocol=pickle.HIGHEST_PROTOCOL,
|
||||
)
|
||||
await send_raw_command(pickled)
|
||||
except (BrokenWorkerProcess, get_cancelled_exc_class()):
|
||||
raise
|
||||
except BaseException as exc:
|
||||
process.kill()
|
||||
raise BrokenWorkerProcess(
|
||||
"Error during worker process initialization"
|
||||
) from exc
|
||||
|
||||
workers.add(process)
|
||||
|
||||
with CancelScope(shield=not cancellable):
|
||||
try:
|
||||
return cast(T_Retval, await send_raw_command(request))
|
||||
finally:
|
||||
if process in workers:
|
||||
idle_workers.append((process, current_time()))
|
||||
|
||||
|
||||
def current_default_process_limiter() -> CapacityLimiter:
|
||||
"""
|
||||
Return the capacity limiter that is used by default to limit the number of worker
|
||||
processes.
|
||||
|
||||
:return: a capacity limiter object
|
||||
|
||||
"""
|
||||
try:
|
||||
return _default_process_limiter.get()
|
||||
except LookupError:
|
||||
limiter = CapacityLimiter(os.cpu_count() or 2)
|
||||
_default_process_limiter.set(limiter)
|
||||
return limiter
|
||||
|
||||
|
||||
def process_worker() -> None:
|
||||
# Redirect standard streams to os.devnull so that user code won't interfere with the
|
||||
# parent-worker communication
|
||||
stdin = sys.stdin
|
||||
stdout = sys.stdout
|
||||
sys.stdin = open(os.devnull)
|
||||
sys.stdout = open(os.devnull, "w")
|
||||
sys.stderr = open(os.devnull, "w")
|
||||
|
||||
stdout.buffer.write(b"READY\n")
|
||||
while True:
|
||||
retval = exception = None
|
||||
try:
|
||||
command, *args = pickle.load(stdin.buffer)
|
||||
except EOFError:
|
||||
return
|
||||
except BaseException as exc:
|
||||
exception = exc
|
||||
else:
|
||||
if command == "run":
|
||||
func, args = args
|
||||
try:
|
||||
retval = func(*args)
|
||||
except BaseException as exc:
|
||||
exception = exc
|
||||
elif command == "init":
|
||||
main_module_path: str | None
|
||||
sys.path, main_module_path = args
|
||||
del sys.modules["__main__"]
|
||||
if main_module_path and os.path.isfile(main_module_path):
|
||||
# Load the parent's main module but as __mp_main__ instead of
|
||||
# __main__ (like multiprocessing does) to avoid infinite recursion
|
||||
try:
|
||||
main = ModuleType("__mp_main__")
|
||||
main_content = runpy.run_path(
|
||||
main_module_path, run_name="__mp_main__"
|
||||
)
|
||||
main.__dict__.update(main_content)
|
||||
sys.modules["__main__"] = sys.modules["__mp_main__"] = main
|
||||
except BaseException as exc:
|
||||
exception = exc
|
||||
try:
|
||||
if exception is not None:
|
||||
status = b"EXCEPTION"
|
||||
pickled = pickle.dumps(exception, pickle.HIGHEST_PROTOCOL)
|
||||
else:
|
||||
status = b"RETURN"
|
||||
pickled = pickle.dumps(retval, pickle.HIGHEST_PROTOCOL)
|
||||
except BaseException as exc:
|
||||
exception = exc
|
||||
status = b"EXCEPTION"
|
||||
pickled = pickle.dumps(exc, pickle.HIGHEST_PROTOCOL)
|
||||
|
||||
stdout.buffer.write(b"%s %d\n" % (status, len(pickled)))
|
||||
stdout.buffer.write(pickled)
|
||||
|
||||
# Respect SIGTERM
|
||||
if isinstance(exception, SystemExit):
|
||||
raise exception
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
process_worker()
|
||||
+80
@@ -0,0 +1,80 @@
|
||||
from __future__ import annotations
|
||||
|
||||
__all__ = (
|
||||
"current_default_thread_limiter",
|
||||
"run_sync",
|
||||
)
|
||||
|
||||
import sys
|
||||
from collections.abc import Callable
|
||||
from typing import TYPE_CHECKING, TypeVar
|
||||
from warnings import warn
|
||||
|
||||
from ._core._eventloop import get_async_backend
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ._core._synchronization import CapacityLimiter
|
||||
|
||||
if sys.version_info >= (3, 11):
|
||||
from typing import TypeVarTuple, Unpack
|
||||
else:
|
||||
from typing_extensions import TypeVarTuple, Unpack
|
||||
|
||||
T_Retval = TypeVar("T_Retval")
|
||||
PosArgsT = TypeVarTuple("PosArgsT")
|
||||
|
||||
|
||||
async def run_sync(
|
||||
func: Callable[[Unpack[PosArgsT]], T_Retval],
|
||||
*args: Unpack[PosArgsT],
|
||||
abandon_on_cancel: bool = False,
|
||||
cancellable: bool | None = None,
|
||||
limiter: CapacityLimiter | None = None,
|
||||
) -> T_Retval:
|
||||
"""
|
||||
Call the given function with the given arguments in a worker thread.
|
||||
|
||||
If the ``abandon_on_cancel`` option is enabled and the task waiting for its
|
||||
completion is cancelled, the thread will still run its course but its
|
||||
return value (or any raised exception) will be ignored.
|
||||
|
||||
:param func: a callable
|
||||
:param args: positional arguments for the callable
|
||||
:param abandon_on_cancel: ``True`` to abandon the thread (leaving it to run
|
||||
unchecked on own) if the host task is cancelled, ``False`` to ignore
|
||||
cancellations in the host task until the operation has completed in the worker
|
||||
thread
|
||||
:param cancellable: deprecated alias of ``abandon_on_cancel``; will override
|
||||
``abandon_on_cancel`` if both parameters are passed
|
||||
:param limiter: capacity limiter to use to limit the total amount of threads running
|
||||
(if omitted, the default limiter is used)
|
||||
:raises NoEventLoopError: if no supported asynchronous event loop is running in the
|
||||
current thread
|
||||
:return: an awaitable that yields the return value of the function.
|
||||
|
||||
"""
|
||||
if cancellable is not None:
|
||||
abandon_on_cancel = cancellable
|
||||
warn(
|
||||
"The `cancellable=` keyword argument to `anyio.to_thread.run_sync` is "
|
||||
"deprecated since AnyIO 4.1.0; use `abandon_on_cancel=` instead",
|
||||
DeprecationWarning,
|
||||
stacklevel=2,
|
||||
)
|
||||
|
||||
return await get_async_backend().run_sync_in_worker_thread(
|
||||
func, args, abandon_on_cancel=abandon_on_cancel, limiter=limiter
|
||||
)
|
||||
|
||||
|
||||
def current_default_thread_limiter() -> CapacityLimiter:
|
||||
"""
|
||||
Return the capacity limiter that is used by default to limit the number of
|
||||
concurrent threads.
|
||||
|
||||
:return: a capacity limiter object
|
||||
:raises NoEventLoopError: if no supported asynchronous event loop is running in the
|
||||
current thread
|
||||
|
||||
"""
|
||||
return get_async_backend().current_default_thread_limiter()
|
||||
+104
@@ -0,0 +1,104 @@
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
"""
|
||||
Classes Without Boilerplate
|
||||
"""
|
||||
|
||||
from functools import partial
|
||||
from typing import Callable, Literal, Protocol
|
||||
|
||||
from . import converters, exceptions, filters, setters, validators
|
||||
from ._cmp import cmp_using
|
||||
from ._config import get_run_validators, set_run_validators
|
||||
from ._funcs import asdict, assoc, astuple, has, resolve_types
|
||||
from ._make import (
|
||||
NOTHING,
|
||||
Attribute,
|
||||
Converter,
|
||||
Factory,
|
||||
_Nothing,
|
||||
attrib,
|
||||
attrs,
|
||||
evolve,
|
||||
fields,
|
||||
fields_dict,
|
||||
make_class,
|
||||
validate,
|
||||
)
|
||||
from ._next_gen import define, field, frozen, mutable
|
||||
from ._version_info import VersionInfo
|
||||
|
||||
|
||||
s = attributes = attrs
|
||||
ib = attr = attrib
|
||||
dataclass = partial(attrs, auto_attribs=True) # happy Easter ;)
|
||||
|
||||
|
||||
class AttrsInstance(Protocol):
|
||||
pass
|
||||
|
||||
|
||||
NothingType = Literal[_Nothing.NOTHING]
|
||||
|
||||
__all__ = [
|
||||
"NOTHING",
|
||||
"Attribute",
|
||||
"AttrsInstance",
|
||||
"Converter",
|
||||
"Factory",
|
||||
"NothingType",
|
||||
"asdict",
|
||||
"assoc",
|
||||
"astuple",
|
||||
"attr",
|
||||
"attrib",
|
||||
"attributes",
|
||||
"attrs",
|
||||
"cmp_using",
|
||||
"converters",
|
||||
"define",
|
||||
"evolve",
|
||||
"exceptions",
|
||||
"field",
|
||||
"fields",
|
||||
"fields_dict",
|
||||
"filters",
|
||||
"frozen",
|
||||
"get_run_validators",
|
||||
"has",
|
||||
"ib",
|
||||
"make_class",
|
||||
"mutable",
|
||||
"resolve_types",
|
||||
"s",
|
||||
"set_run_validators",
|
||||
"setters",
|
||||
"validate",
|
||||
"validators",
|
||||
]
|
||||
|
||||
|
||||
def _make_getattr(mod_name: str) -> Callable:
|
||||
"""
|
||||
Create a metadata proxy for packaging information that uses *mod_name* in
|
||||
its warnings and errors.
|
||||
"""
|
||||
|
||||
def __getattr__(name: str) -> str:
|
||||
if name not in ("__version__", "__version_info__"):
|
||||
msg = f"module {mod_name} has no attribute {name}"
|
||||
raise AttributeError(msg)
|
||||
|
||||
from importlib.metadata import metadata
|
||||
|
||||
meta = metadata("attrs")
|
||||
|
||||
if name == "__version_info__":
|
||||
return VersionInfo._from_version_string(meta["version"])
|
||||
|
||||
return meta["version"]
|
||||
|
||||
return __getattr__
|
||||
|
||||
|
||||
__getattr__ = _make_getattr(__name__)
|
||||
+389
@@ -0,0 +1,389 @@
|
||||
import enum
|
||||
import sys
|
||||
|
||||
from typing import (
|
||||
Any,
|
||||
Callable,
|
||||
Generic,
|
||||
Literal,
|
||||
Mapping,
|
||||
Protocol,
|
||||
Sequence,
|
||||
TypeVar,
|
||||
overload,
|
||||
)
|
||||
|
||||
# `import X as X` is required to make these public
|
||||
from . import converters as converters
|
||||
from . import exceptions as exceptions
|
||||
from . import filters as filters
|
||||
from . import setters as setters
|
||||
from . import validators as validators
|
||||
from ._cmp import cmp_using as cmp_using
|
||||
from ._typing_compat import AttrsInstance_
|
||||
from ._version_info import VersionInfo
|
||||
from attrs import (
|
||||
define as define,
|
||||
field as field,
|
||||
mutable as mutable,
|
||||
frozen as frozen,
|
||||
_EqOrderType,
|
||||
_ValidatorType,
|
||||
_ConverterType,
|
||||
_ReprArgType,
|
||||
_OnSetAttrType,
|
||||
_OnSetAttrArgType,
|
||||
_FieldTransformer,
|
||||
_ValidatorArgType,
|
||||
)
|
||||
|
||||
if sys.version_info >= (3, 10):
|
||||
from typing import TypeGuard, TypeAlias
|
||||
else:
|
||||
from typing_extensions import TypeGuard, TypeAlias
|
||||
|
||||
if sys.version_info >= (3, 11):
|
||||
from typing import dataclass_transform
|
||||
else:
|
||||
from typing_extensions import dataclass_transform
|
||||
|
||||
__version__: str
|
||||
__version_info__: VersionInfo
|
||||
__title__: str
|
||||
__description__: str
|
||||
__url__: str
|
||||
__uri__: str
|
||||
__author__: str
|
||||
__email__: str
|
||||
__license__: str
|
||||
__copyright__: str
|
||||
|
||||
_T = TypeVar("_T")
|
||||
_C = TypeVar("_C", bound=type)
|
||||
|
||||
_FilterType = Callable[["Attribute[_T]", _T], bool]
|
||||
|
||||
# We subclass this here to keep the protocol's qualified name clean.
|
||||
class AttrsInstance(AttrsInstance_, Protocol):
|
||||
pass
|
||||
|
||||
_A = TypeVar("_A", bound=type[AttrsInstance])
|
||||
|
||||
class _Nothing(enum.Enum):
|
||||
NOTHING = enum.auto()
|
||||
|
||||
NOTHING = _Nothing.NOTHING
|
||||
NothingType: TypeAlias = Literal[_Nothing.NOTHING]
|
||||
|
||||
# NOTE: Factory lies about its return type to make this possible:
|
||||
# `x: List[int] # = Factory(list)`
|
||||
# Work around mypy issue #4554 in the common case by using an overload.
|
||||
|
||||
@overload
|
||||
def Factory(factory: Callable[[], _T]) -> _T: ...
|
||||
@overload
|
||||
def Factory(
|
||||
factory: Callable[[Any], _T],
|
||||
takes_self: Literal[True],
|
||||
) -> _T: ...
|
||||
@overload
|
||||
def Factory(
|
||||
factory: Callable[[], _T],
|
||||
takes_self: Literal[False],
|
||||
) -> _T: ...
|
||||
|
||||
In = TypeVar("In")
|
||||
Out = TypeVar("Out")
|
||||
|
||||
class Converter(Generic[In, Out]):
|
||||
@overload
|
||||
def __init__(self, converter: Callable[[In], Out]) -> None: ...
|
||||
@overload
|
||||
def __init__(
|
||||
self,
|
||||
converter: Callable[[In, AttrsInstance, Attribute], Out],
|
||||
*,
|
||||
takes_self: Literal[True],
|
||||
takes_field: Literal[True],
|
||||
) -> None: ...
|
||||
@overload
|
||||
def __init__(
|
||||
self,
|
||||
converter: Callable[[In, Attribute], Out],
|
||||
*,
|
||||
takes_field: Literal[True],
|
||||
) -> None: ...
|
||||
@overload
|
||||
def __init__(
|
||||
self,
|
||||
converter: Callable[[In, AttrsInstance], Out],
|
||||
*,
|
||||
takes_self: Literal[True],
|
||||
) -> None: ...
|
||||
|
||||
class Attribute(Generic[_T]):
|
||||
name: str
|
||||
default: _T | None
|
||||
validator: _ValidatorType[_T] | None
|
||||
repr: _ReprArgType
|
||||
cmp: _EqOrderType
|
||||
eq: _EqOrderType
|
||||
order: _EqOrderType
|
||||
hash: bool | None
|
||||
init: bool
|
||||
converter: Converter | None
|
||||
metadata: dict[Any, Any]
|
||||
type: type[_T] | None
|
||||
kw_only: bool
|
||||
on_setattr: _OnSetAttrType
|
||||
alias: str | None
|
||||
|
||||
def evolve(self, **changes: Any) -> "Attribute[Any]": ...
|
||||
|
||||
# NOTE: We had several choices for the annotation to use for type arg:
|
||||
# 1) Type[_T]
|
||||
# - Pros: Handles simple cases correctly
|
||||
# - Cons: Might produce less informative errors in the case of conflicting
|
||||
# TypeVars e.g. `attr.ib(default='bad', type=int)`
|
||||
# 2) Callable[..., _T]
|
||||
# - Pros: Better error messages than #1 for conflicting TypeVars
|
||||
# - Cons: Terrible error messages for validator checks.
|
||||
# e.g. attr.ib(type=int, validator=validate_str)
|
||||
# -> error: Cannot infer function type argument
|
||||
# 3) type (and do all of the work in the mypy plugin)
|
||||
# - Pros: Simple here, and we could customize the plugin with our own errors.
|
||||
# - Cons: Would need to write mypy plugin code to handle all the cases.
|
||||
# We chose option #1.
|
||||
|
||||
# `attr` lies about its return type to make the following possible:
|
||||
# attr() -> Any
|
||||
# attr(8) -> int
|
||||
# attr(validator=<some callable>) -> Whatever the callable expects.
|
||||
# This makes this type of assignments possible:
|
||||
# x: int = attr(8)
|
||||
#
|
||||
# This form catches explicit None or no default but with no other arguments
|
||||
# returns Any.
|
||||
@overload
|
||||
def attrib(
|
||||
default: None = ...,
|
||||
validator: None = ...,
|
||||
repr: _ReprArgType = ...,
|
||||
cmp: _EqOrderType | None = ...,
|
||||
hash: bool | None = ...,
|
||||
init: bool = ...,
|
||||
metadata: Mapping[Any, Any] | None = ...,
|
||||
type: None = ...,
|
||||
converter: None = ...,
|
||||
factory: None = ...,
|
||||
kw_only: bool | None = ...,
|
||||
eq: _EqOrderType | None = ...,
|
||||
order: _EqOrderType | None = ...,
|
||||
on_setattr: _OnSetAttrArgType | None = ...,
|
||||
alias: str | None = ...,
|
||||
) -> Any: ...
|
||||
|
||||
# This form catches an explicit None or no default and infers the type from the
|
||||
# other arguments.
|
||||
@overload
|
||||
def attrib(
|
||||
default: None = ...,
|
||||
validator: _ValidatorArgType[_T] | None = ...,
|
||||
repr: _ReprArgType = ...,
|
||||
cmp: _EqOrderType | None = ...,
|
||||
hash: bool | None = ...,
|
||||
init: bool = ...,
|
||||
metadata: Mapping[Any, Any] | None = ...,
|
||||
type: type[_T] | None = ...,
|
||||
converter: _ConverterType
|
||||
| list[_ConverterType]
|
||||
| tuple[_ConverterType]
|
||||
| None = ...,
|
||||
factory: Callable[[], _T] | None = ...,
|
||||
kw_only: bool | None = ...,
|
||||
eq: _EqOrderType | None = ...,
|
||||
order: _EqOrderType | None = ...,
|
||||
on_setattr: _OnSetAttrArgType | None = ...,
|
||||
alias: str | None = ...,
|
||||
) -> _T: ...
|
||||
|
||||
# This form catches an explicit default argument.
|
||||
@overload
|
||||
def attrib(
|
||||
default: _T,
|
||||
validator: _ValidatorArgType[_T] | None = ...,
|
||||
repr: _ReprArgType = ...,
|
||||
cmp: _EqOrderType | None = ...,
|
||||
hash: bool | None = ...,
|
||||
init: bool = ...,
|
||||
metadata: Mapping[Any, Any] | None = ...,
|
||||
type: type[_T] | None = ...,
|
||||
converter: _ConverterType
|
||||
| list[_ConverterType]
|
||||
| tuple[_ConverterType]
|
||||
| None = ...,
|
||||
factory: Callable[[], _T] | None = ...,
|
||||
kw_only: bool | None = ...,
|
||||
eq: _EqOrderType | None = ...,
|
||||
order: _EqOrderType | None = ...,
|
||||
on_setattr: _OnSetAttrArgType | None = ...,
|
||||
alias: str | None = ...,
|
||||
) -> _T: ...
|
||||
|
||||
# This form covers type=non-Type: e.g. forward references (str), Any
|
||||
@overload
|
||||
def attrib(
|
||||
default: _T | None = ...,
|
||||
validator: _ValidatorArgType[_T] | None = ...,
|
||||
repr: _ReprArgType = ...,
|
||||
cmp: _EqOrderType | None = ...,
|
||||
hash: bool | None = ...,
|
||||
init: bool = ...,
|
||||
metadata: Mapping[Any, Any] | None = ...,
|
||||
type: object = ...,
|
||||
converter: _ConverterType
|
||||
| list[_ConverterType]
|
||||
| tuple[_ConverterType]
|
||||
| None = ...,
|
||||
factory: Callable[[], _T] | None = ...,
|
||||
kw_only: bool | None = ...,
|
||||
eq: _EqOrderType | None = ...,
|
||||
order: _EqOrderType | None = ...,
|
||||
on_setattr: _OnSetAttrArgType | None = ...,
|
||||
alias: str | None = ...,
|
||||
) -> Any: ...
|
||||
@overload
|
||||
@dataclass_transform(order_default=True, field_specifiers=(attrib, field))
|
||||
def attrs(
|
||||
maybe_cls: _C,
|
||||
these: dict[str, Any] | None = ...,
|
||||
repr_ns: str | None = ...,
|
||||
repr: bool = ...,
|
||||
cmp: _EqOrderType | None = ...,
|
||||
hash: bool | None = ...,
|
||||
init: bool = ...,
|
||||
slots: bool = ...,
|
||||
frozen: bool = ...,
|
||||
weakref_slot: bool = ...,
|
||||
str: bool = ...,
|
||||
auto_attribs: bool = ...,
|
||||
kw_only: bool = ...,
|
||||
cache_hash: bool = ...,
|
||||
auto_exc: bool = ...,
|
||||
eq: _EqOrderType | None = ...,
|
||||
order: _EqOrderType | None = ...,
|
||||
auto_detect: bool = ...,
|
||||
collect_by_mro: bool = ...,
|
||||
getstate_setstate: bool | None = ...,
|
||||
on_setattr: _OnSetAttrArgType | None = ...,
|
||||
field_transformer: _FieldTransformer | None = ...,
|
||||
match_args: bool = ...,
|
||||
unsafe_hash: bool | None = ...,
|
||||
) -> _C: ...
|
||||
@overload
|
||||
@dataclass_transform(order_default=True, field_specifiers=(attrib, field))
|
||||
def attrs(
|
||||
maybe_cls: None = ...,
|
||||
these: dict[str, Any] | None = ...,
|
||||
repr_ns: str | None = ...,
|
||||
repr: bool = ...,
|
||||
cmp: _EqOrderType | None = ...,
|
||||
hash: bool | None = ...,
|
||||
init: bool = ...,
|
||||
slots: bool = ...,
|
||||
frozen: bool = ...,
|
||||
weakref_slot: bool = ...,
|
||||
str: bool = ...,
|
||||
auto_attribs: bool = ...,
|
||||
kw_only: bool = ...,
|
||||
cache_hash: bool = ...,
|
||||
auto_exc: bool = ...,
|
||||
eq: _EqOrderType | None = ...,
|
||||
order: _EqOrderType | None = ...,
|
||||
auto_detect: bool = ...,
|
||||
collect_by_mro: bool = ...,
|
||||
getstate_setstate: bool | None = ...,
|
||||
on_setattr: _OnSetAttrArgType | None = ...,
|
||||
field_transformer: _FieldTransformer | None = ...,
|
||||
match_args: bool = ...,
|
||||
unsafe_hash: bool | None = ...,
|
||||
) -> Callable[[_C], _C]: ...
|
||||
def fields(cls: type[AttrsInstance] | AttrsInstance) -> Any: ...
|
||||
def fields_dict(cls: type[AttrsInstance]) -> dict[str, Attribute[Any]]: ...
|
||||
def validate(inst: AttrsInstance) -> None: ...
|
||||
def resolve_types(
|
||||
cls: _A,
|
||||
globalns: dict[str, Any] | None = ...,
|
||||
localns: dict[str, Any] | None = ...,
|
||||
attribs: list[Attribute[Any]] | None = ...,
|
||||
include_extras: bool = ...,
|
||||
) -> _A: ...
|
||||
|
||||
# TODO: add support for returning a proper attrs class from the mypy plugin
|
||||
# we use Any instead of _CountingAttr so that e.g. `make_class('Foo',
|
||||
# [attr.ib()])` is valid
|
||||
def make_class(
|
||||
name: str,
|
||||
attrs: list[str] | tuple[str, ...] | dict[str, Any],
|
||||
bases: tuple[type, ...] = ...,
|
||||
class_body: dict[str, Any] | None = ...,
|
||||
repr_ns: str | None = ...,
|
||||
repr: bool = ...,
|
||||
cmp: _EqOrderType | None = ...,
|
||||
hash: bool | None = ...,
|
||||
init: bool = ...,
|
||||
slots: bool = ...,
|
||||
frozen: bool = ...,
|
||||
weakref_slot: bool = ...,
|
||||
str: bool = ...,
|
||||
auto_attribs: bool = ...,
|
||||
kw_only: bool = ...,
|
||||
cache_hash: bool = ...,
|
||||
auto_exc: bool = ...,
|
||||
eq: _EqOrderType | None = ...,
|
||||
order: _EqOrderType | None = ...,
|
||||
collect_by_mro: bool = ...,
|
||||
on_setattr: _OnSetAttrArgType | None = ...,
|
||||
field_transformer: _FieldTransformer | None = ...,
|
||||
) -> type: ...
|
||||
|
||||
# _funcs --
|
||||
|
||||
# TODO: add support for returning TypedDict from the mypy plugin
|
||||
# FIXME: asdict/astuple do not honor their factory args. Waiting on one of
|
||||
# these:
|
||||
# https://github.com/python/mypy/issues/4236
|
||||
# https://github.com/python/typing/issues/253
|
||||
# XXX: remember to fix attrs.asdict/astuple too!
|
||||
def asdict(
|
||||
inst: AttrsInstance,
|
||||
recurse: bool = ...,
|
||||
filter: _FilterType[Any] | None = ...,
|
||||
dict_factory: type[Mapping[Any, Any]] = ...,
|
||||
retain_collection_types: bool = ...,
|
||||
value_serializer: Callable[[type, Attribute[Any], Any], Any] | None = ...,
|
||||
tuple_keys: bool | None = ...,
|
||||
) -> dict[str, Any]: ...
|
||||
|
||||
# TODO: add support for returning NamedTuple from the mypy plugin
|
||||
def astuple(
|
||||
inst: AttrsInstance,
|
||||
recurse: bool = ...,
|
||||
filter: _FilterType[Any] | None = ...,
|
||||
tuple_factory: type[Sequence[Any]] = ...,
|
||||
retain_collection_types: bool = ...,
|
||||
) -> tuple[Any, ...]: ...
|
||||
def has(cls: type) -> TypeGuard[type[AttrsInstance]]: ...
|
||||
def assoc(inst: _T, **changes: Any) -> _T: ...
|
||||
def evolve(inst: _T, **changes: Any) -> _T: ...
|
||||
|
||||
# _config --
|
||||
|
||||
def set_run_validators(run: bool) -> None: ...
|
||||
def get_run_validators() -> bool: ...
|
||||
|
||||
# aliases --
|
||||
|
||||
s = attributes = attrs
|
||||
ib = attr = attrib
|
||||
dataclass = attrs # Technically, partial(attrs, auto_attribs=True) ;)
|
||||
@@ -0,0 +1,160 @@
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
|
||||
import functools
|
||||
import types
|
||||
|
||||
from ._make import __ne__
|
||||
|
||||
|
||||
_operation_names = {"eq": "==", "lt": "<", "le": "<=", "gt": ">", "ge": ">="}
|
||||
|
||||
|
||||
def cmp_using(
|
||||
eq=None,
|
||||
lt=None,
|
||||
le=None,
|
||||
gt=None,
|
||||
ge=None,
|
||||
require_same_type=True,
|
||||
class_name="Comparable",
|
||||
):
|
||||
"""
|
||||
Create a class that can be passed into `attrs.field`'s ``eq``, ``order``,
|
||||
and ``cmp`` arguments to customize field comparison.
|
||||
|
||||
The resulting class will have a full set of ordering methods if at least
|
||||
one of ``{lt, le, gt, ge}`` and ``eq`` are provided.
|
||||
|
||||
Args:
|
||||
eq (typing.Callable | None):
|
||||
Callable used to evaluate equality of two objects.
|
||||
|
||||
lt (typing.Callable | None):
|
||||
Callable used to evaluate whether one object is less than another
|
||||
object.
|
||||
|
||||
le (typing.Callable | None):
|
||||
Callable used to evaluate whether one object is less than or equal
|
||||
to another object.
|
||||
|
||||
gt (typing.Callable | None):
|
||||
Callable used to evaluate whether one object is greater than
|
||||
another object.
|
||||
|
||||
ge (typing.Callable | None):
|
||||
Callable used to evaluate whether one object is greater than or
|
||||
equal to another object.
|
||||
|
||||
require_same_type (bool):
|
||||
When `True`, equality and ordering methods will return
|
||||
`NotImplemented` if objects are not of the same type.
|
||||
|
||||
class_name (str | None): Name of class. Defaults to "Comparable".
|
||||
|
||||
See `comparison` for more details.
|
||||
|
||||
.. versionadded:: 21.1.0
|
||||
"""
|
||||
|
||||
body = {
|
||||
"__slots__": ["value"],
|
||||
"__init__": _make_init(),
|
||||
"_requirements": [],
|
||||
"_is_comparable_to": _is_comparable_to,
|
||||
}
|
||||
|
||||
# Add operations.
|
||||
num_order_functions = 0
|
||||
has_eq_function = False
|
||||
|
||||
if eq is not None:
|
||||
has_eq_function = True
|
||||
body["__eq__"] = _make_operator("eq", eq)
|
||||
body["__ne__"] = __ne__
|
||||
|
||||
if lt is not None:
|
||||
num_order_functions += 1
|
||||
body["__lt__"] = _make_operator("lt", lt)
|
||||
|
||||
if le is not None:
|
||||
num_order_functions += 1
|
||||
body["__le__"] = _make_operator("le", le)
|
||||
|
||||
if gt is not None:
|
||||
num_order_functions += 1
|
||||
body["__gt__"] = _make_operator("gt", gt)
|
||||
|
||||
if ge is not None:
|
||||
num_order_functions += 1
|
||||
body["__ge__"] = _make_operator("ge", ge)
|
||||
|
||||
type_ = types.new_class(
|
||||
class_name, (object,), {}, lambda ns: ns.update(body)
|
||||
)
|
||||
|
||||
# Add same type requirement.
|
||||
if require_same_type:
|
||||
type_._requirements.append(_check_same_type)
|
||||
|
||||
# Add total ordering if at least one operation was defined.
|
||||
if 0 < num_order_functions < 4:
|
||||
if not has_eq_function:
|
||||
# functools.total_ordering requires __eq__ to be defined,
|
||||
# so raise early error here to keep a nice stack.
|
||||
msg = "eq must be define is order to complete ordering from lt, le, gt, ge."
|
||||
raise ValueError(msg)
|
||||
type_ = functools.total_ordering(type_)
|
||||
|
||||
return type_
|
||||
|
||||
|
||||
def _make_init():
|
||||
"""
|
||||
Create __init__ method.
|
||||
"""
|
||||
|
||||
def __init__(self, value):
|
||||
"""
|
||||
Initialize object with *value*.
|
||||
"""
|
||||
self.value = value
|
||||
|
||||
return __init__
|
||||
|
||||
|
||||
def _make_operator(name, func):
|
||||
"""
|
||||
Create operator method.
|
||||
"""
|
||||
|
||||
def method(self, other):
|
||||
if not self._is_comparable_to(other):
|
||||
return NotImplemented
|
||||
|
||||
result = func(self.value, other.value)
|
||||
if result is NotImplemented:
|
||||
return NotImplemented
|
||||
|
||||
return result
|
||||
|
||||
method.__name__ = f"__{name}__"
|
||||
method.__doc__ = (
|
||||
f"Return a {_operation_names[name]} b. Computed by attrs."
|
||||
)
|
||||
|
||||
return method
|
||||
|
||||
|
||||
def _is_comparable_to(self, other):
|
||||
"""
|
||||
Check whether `other` is comparable to `self`.
|
||||
"""
|
||||
return all(func(self, other) for func in self._requirements)
|
||||
|
||||
|
||||
def _check_same_type(self, other):
|
||||
"""
|
||||
Return True if *self* and *other* are of the same type, False otherwise.
|
||||
"""
|
||||
return other.value.__class__ is self.value.__class__
|
||||
@@ -0,0 +1,13 @@
|
||||
from typing import Any, Callable
|
||||
|
||||
_CompareWithType = Callable[[Any, Any], bool]
|
||||
|
||||
def cmp_using(
|
||||
eq: _CompareWithType | None = ...,
|
||||
lt: _CompareWithType | None = ...,
|
||||
le: _CompareWithType | None = ...,
|
||||
gt: _CompareWithType | None = ...,
|
||||
ge: _CompareWithType | None = ...,
|
||||
require_same_type: bool = ...,
|
||||
class_name: str = ...,
|
||||
) -> type: ...
|
||||
@@ -0,0 +1,99 @@
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
import inspect
|
||||
import platform
|
||||
import sys
|
||||
import threading
|
||||
|
||||
from collections.abc import Mapping, Sequence # noqa: F401
|
||||
from typing import _GenericAlias
|
||||
|
||||
|
||||
PYPY = platform.python_implementation() == "PyPy"
|
||||
PY_3_10_PLUS = sys.version_info[:2] >= (3, 10)
|
||||
PY_3_11_PLUS = sys.version_info[:2] >= (3, 11)
|
||||
PY_3_12_PLUS = sys.version_info[:2] >= (3, 12)
|
||||
PY_3_13_PLUS = sys.version_info[:2] >= (3, 13)
|
||||
PY_3_14_PLUS = sys.version_info[:2] >= (3, 14)
|
||||
|
||||
|
||||
if PY_3_14_PLUS:
|
||||
import annotationlib
|
||||
|
||||
# We request forward-ref annotations to not break in the presence of
|
||||
# forward references.
|
||||
|
||||
def _get_annotations(cls):
|
||||
return annotationlib.get_annotations(
|
||||
cls, format=annotationlib.Format.FORWARDREF
|
||||
)
|
||||
|
||||
else:
|
||||
|
||||
def _get_annotations(cls):
|
||||
"""
|
||||
Get annotations for *cls*.
|
||||
"""
|
||||
return cls.__dict__.get("__annotations__", {})
|
||||
|
||||
|
||||
class _AnnotationExtractor:
|
||||
"""
|
||||
Extract type annotations from a callable, returning None whenever there
|
||||
is none.
|
||||
"""
|
||||
|
||||
__slots__ = ["sig"]
|
||||
|
||||
def __init__(self, callable):
|
||||
try:
|
||||
self.sig = inspect.signature(callable)
|
||||
except (ValueError, TypeError): # inspect failed
|
||||
self.sig = None
|
||||
|
||||
def get_first_param_type(self):
|
||||
"""
|
||||
Return the type annotation of the first argument if it's not empty.
|
||||
"""
|
||||
if not self.sig:
|
||||
return None
|
||||
|
||||
params = list(self.sig.parameters.values())
|
||||
if params and params[0].annotation is not inspect.Parameter.empty:
|
||||
return params[0].annotation
|
||||
|
||||
return None
|
||||
|
||||
def get_return_type(self):
|
||||
"""
|
||||
Return the return type if it's not empty.
|
||||
"""
|
||||
if (
|
||||
self.sig
|
||||
and self.sig.return_annotation is not inspect.Signature.empty
|
||||
):
|
||||
return self.sig.return_annotation
|
||||
|
||||
return None
|
||||
|
||||
|
||||
# Thread-local global to track attrs instances which are already being repr'd.
|
||||
# This is needed because there is no other (thread-safe) way to pass info
|
||||
# about the instances that are already being repr'd through the call stack
|
||||
# in order to ensure we don't perform infinite recursion.
|
||||
#
|
||||
# For instance, if an instance contains a dict which contains that instance,
|
||||
# we need to know that we're already repr'ing the outside instance from within
|
||||
# the dict's repr() call.
|
||||
#
|
||||
# This lives here rather than in _make.py so that the functions in _make.py
|
||||
# don't have a direct reference to the thread-local in their globals dict.
|
||||
# If they have such a reference, it breaks cloudpickle.
|
||||
repr_context = threading.local()
|
||||
|
||||
|
||||
def get_generic_base(cl):
|
||||
"""If this is a generic class (A[str]), return the generic base for it."""
|
||||
if cl.__class__ is _GenericAlias:
|
||||
return cl.__origin__
|
||||
return None
|
||||
Loaded 100 of 1354 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user