from __future__ import annotations import json import logging from pathlib import Path from typing import Any import httpx from yuxi.channel.extensions.zulip.errors import ( ZulipAPIError, ZulipAuthError, ZulipQueueExpiredError, ZulipRateLimitError, ) logger = logging.getLogger(__name__) class ZulipAsyncClient: def __init__( self, realm_url: str, bot_email: str, bot_api_key: str, timeout: float = 90.0, connect_timeout: float = 10.0, ): self.realm_url = realm_url.rstrip("/") self.api_url = f"{self.realm_url}/api/v1" self.bot_email = bot_email limits = httpx.Limits(max_keepalive_connections=5, max_connections=20) transport = httpx.AsyncHTTPTransport(limits=limits, retries=0) self._client = httpx.AsyncClient( base_url=self.api_url, auth=(bot_email, bot_api_key), timeout=httpx.Timeout(timeout, connect=connect_timeout), transport=transport, ) async def close(self) -> None: await self._client.aclose() async def _request( self, method: str, path: str, data: dict | None = None, params: dict | None = None, files: dict | None = None, timeout: httpx.Timeout | None = None, ) -> dict: kwargs: dict[str, Any] = {"method": method, "url": path} if data: kwargs["data"] = data if params: kwargs["params"] = params if files: kwargs["files"] = files if timeout: kwargs["timeout"] = timeout resp = await self._client.request(**kwargs) if resp.status_code >= 400: try: body = resp.json() except Exception: body = {"code": "UNKNOWN", "msg": resp.text} if resp.status_code == 401: raise ZulipAuthError(resp.status_code, body) if resp.status_code == 429: raise ZulipRateLimitError(resp.status_code, body) if body.get("code") == "BAD_EVENT_QUEUE_ID": raise ZulipQueueExpiredError(resp.status_code, body) raise ZulipAPIError(resp.status_code, body) return resp.json() async def register_queue( self, event_types: list[str] | None = None, narrow: list | None = None, fetch_event_types: list[str] | None = None, ) -> dict: data: dict[str, str] = {} if event_types: data["event_types"] = json.dumps(event_types) if narrow: data["narrow"] = json.dumps(narrow) if fetch_event_types: data["fetch_event_types"] = json.dumps(fetch_event_types) return await self._request("POST", "/register", data=data) async def get_events( self, queue_id: str, last_event_id: int, dont_block: bool = False, ) -> dict: params: dict[str, Any] = { "queue_id": queue_id, "last_event_id": last_event_id, } if dont_block: params["dont_block"] = "true" return await self._request( "GET", "/events", params=params, timeout=httpx.Timeout(90.0), ) async def delete_queue(self, queue_id: str) -> dict: return await self._request("DELETE", "/events", data={"queue_id": queue_id}) async def send_message(self, payload: dict) -> dict: return await self._request("POST", "/messages", data=payload) async def update_message( self, message_id: int, content: str, propagate_mode: str = "change_one", ) -> dict: data: dict[str, Any] = { "content": content, "propagate_mode": propagate_mode, } return await self._request("PATCH", f"/messages/{message_id}", data=data) async def delete_message(self, message_id: int) -> dict: return await self._request("DELETE", f"/messages/{message_id}") async def upload_file(self, file_path: str) -> dict: path = Path(file_path) with open(path, "rb") as f: files = {"file": (path.name, f)} return await self._request("POST", "/user_uploads", files=files) async def add_reaction(self, message_id: int, emoji_name: str) -> dict: data = { "emoji_name": emoji_name, "reaction_type": "unicode_emoji", } return await self._request("POST", f"/messages/{message_id}/reactions", data=data) async def remove_reaction(self, message_id: int, emoji_name: str) -> dict: params = { "emoji_name": emoji_name, "reaction_type": "unicode_emoji", } return await self._request("DELETE", f"/messages/{message_id}/reactions", params=params) async def get_own_user(self) -> dict: return await self._request("GET", "/users/me") async def get_users(self) -> dict: return await self._request("GET", "/users") async def get_user_by_id(self, user_id: int) -> dict: return await self._request("GET", f"/users/{user_id}") async def get_streams(self, include_public: bool = True) -> dict: params: dict[str, Any] = {} if include_public: params["include_public"] = "true" return await self._request("GET", "/streams", params=params) async def get_stream_id(self, stream_name: str) -> dict: return await self._request("GET", "/get_stream_id", params={"stream": stream_name}) async def get_subscriptions(self) -> dict: return await self._request("GET", "/users/me/subscriptions") async def get_messages( self, narrow: list[dict] | None = None, anchor: str | None = None, num_before: int = 0, num_after: int = 20, include_anchor: bool = True, apply_markdown: bool = True, client_gravatar: bool = True, ) -> dict: params: dict[str, Any] = { "num_before": num_before, "num_after": num_after, "include_anchor": str(include_anchor).lower(), "apply_markdown": str(apply_markdown).lower(), "client_gravatar": str(client_gravatar).lower(), } if anchor: params["anchor"] = anchor if narrow: params["narrow"] = json.dumps(narrow) return await self._request("GET", "/messages", params=params) async def send_typing( self, op: str, to: list[str] | str, topic: str | None = None, ) -> dict: data: dict[str, Any] = {"op": op, "to": to} if topic: data["topic"] = topic return await self._request("POST", "/typing", data=data) async def get_bot_storage(self, key: str | None = None) -> dict: params: dict[str, str] = {} if key: params["key"] = key return await self._request("GET", "/bot_storage", params=params) async def put_bot_storage(self, key: str, value: str) -> dict: return await self._request("PUT", "/bot_storage", data={"key": key, "value": value}) async def delete_bot_storage(self, key: str) -> dict: return await self._request("DELETE", "/bot_storage", data={"key": key}) async def get_message( self, message_id: int, apply_markdown: bool = True, ) -> dict: params = {"apply_markdown": str(apply_markdown).lower()} return await self._request("GET", f"/messages/{message_id}", params=params) async def get_message_history(self, message_id: int) -> dict: return await self._request("GET", f"/messages/{message_id}/history") async def subscribe(self, subscriptions: list[dict]) -> dict: data = {"subscriptions": json.dumps(subscriptions)} return await self._request("POST", "/users/me/subscriptions", data=data) async def unsubscribe(self, stream_names: list[str]) -> dict: return await self._request( "DELETE", "/users/me/subscriptions", data={"subscriptions": json.dumps(stream_names)}, ) async def get_topics(self, stream_id: int) -> dict: return await self._request("GET", f"/users/me/{stream_id}/topics") async def get_stream_members(self, stream_id: int) -> dict: return await self._request("GET", f"/streams/{stream_id}/members") async def get_user_by_email(self, email: str) -> dict: return await self._request("GET", f"/users/{email}") async def update_message_flags( self, messages: list[int], op: str, flag: str, ) -> dict: data = { "messages": json.dumps(messages), "op": op, "flag": flag, } return await self._request("POST", "/messages/flags", data=data) async def mark_all_as_read(self) -> dict: return await self._request("POST", "/mark_all_as_read", data={}) async def mark_stream_as_read(self, stream_id: int) -> dict: return await self._request("POST", "/mark_stream_as_read", data={"stream_id": stream_id}) async def mark_topic_as_read(self, stream_id: int, topic_name: str) -> dict: return await self._request( "POST", "/mark_topic_as_read", data={"stream_id": stream_id, "topic_name": topic_name}, ) async def get_user_status(self, user_id: int) -> dict: return await self._request("GET", f"/users/{user_id}/status") async def update_bot_status( self, status_text: str = "", away: bool = False, emoji_name: str = "", emoji_code: str = "", ) -> dict: data: dict[str, Any] = {} if status_text: data["status_text"] = status_text if away: data["away"] = str(away).lower() if emoji_name: data["emoji_name"] = emoji_name if emoji_code: data["emoji_code"] = emoji_code return await self._request("POST", "/users/me/status", data=data) async def get_realm_emoji(self) -> dict: return await self._request("GET", "/realm/emoji") async def get_server_settings(self) -> dict: return await self._request("GET", "/server_settings") async def get_stream_by_id(self, stream_id: int) -> dict: return await self._request("GET", f"/streams/{stream_id}") def get_file_url(self, realm_id_str: str, filename: str) -> str: return f"{self.realm_url}/user_uploads/{realm_id_str}/{filename}" async def report_message(self, message_id: int, reason: str) -> dict: return await self._request( "POST", "/report/message", data={"message_id": message_id, "reason": reason}, ) async def create_stream( self, name: str, description: str = "", invite_only: bool = False, history_public_to_subscribers: bool = False, ) -> dict: data = { "subscriptions": json.dumps([{"name": name, "description": description}]), "invite_only": str(invite_only).lower(), "history_public_to_subscribers": str(history_public_to_subscribers).lower(), } return await self._request("POST", "/users/me/subscriptions", data=data) async def update_stream( self, stream_id: int, *, name: str | None = None, description: str | None = None, is_private: bool | None = None, ) -> dict: data: dict[str, Any] = {} if name is not None: data["new_name"] = name if description is not None: data["description"] = description if is_private is not None: data["is_private"] = str(is_private).lower() return await self._request("PATCH", f"/streams/{stream_id}", data=data) async def archive_stream(self, stream_id: int) -> dict: return await self._request("DELETE", f"/streams/{stream_id}") async def get_stream_email(self, stream_id: int) -> dict: return await self._request("GET", f"/streams/{stream_id}/email_address") async def mute_topic( self, stream_name: str, topic: str, op: str = "add", ) -> dict: return await self._request( "PATCH", "/users/me/subscriptions/muted_topics", data={"stream": stream_name, "topic": topic, "op": op}, ) async def update_user_topic( self, stream_id: int, topic: str, visibility_policy: int, ) -> dict: return await self._request( "POST", "/user_topics", data={ "stream_id": stream_id, "topic": topic, "visibility_policy": visibility_policy, }, ) async def delete_topic(self, stream_id: int, topic_name: str) -> dict: return await self._request( "DELETE", f"/streams/{stream_id}/delete_topic", data={"topic_name": topic_name}, ) async def get_subscription_status(self, user_id: int, stream_id: int) -> dict: return await self._request("GET", f"/users/{user_id}/subscriptions/{stream_id}")