mirror of
https://github.com/artiemis/artemis.git
synced 2026-02-14 08:31:55 +00:00
85 lines
2.6 KiB
Python
85 lines
2.6 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import io
|
|
from dataclasses import dataclass
|
|
from typing import TYPE_CHECKING, Any, Literal
|
|
from utils.common import ArtemisError
|
|
|
|
import aiohttp
|
|
|
|
if TYPE_CHECKING:
|
|
from bot import Artemis
|
|
|
|
|
|
@dataclass
|
|
class DeepLResult:
|
|
src: str
|
|
dst: str
|
|
translation: str
|
|
|
|
|
|
class API:
|
|
def __init__(self, bot: Artemis, token: str):
|
|
self.base_url = "http://127.0.0.1:3000"
|
|
self.token = token
|
|
self.session: aiohttp.ClientSession = bot.session
|
|
self.HEADERS = {"User-Agent": bot.real_user_agent}
|
|
self.AUTHED_HEADERS = {**self.HEADERS, "Authorization": f"Bearer {self.token}"}
|
|
|
|
async def _aioread(self, fp):
|
|
return await asyncio.to_thread(fp.read)
|
|
|
|
async def _request(
|
|
self,
|
|
method: str,
|
|
path: str,
|
|
authed: bool = False,
|
|
res_type: Literal["json", "text", "bytes"] = "json",
|
|
**kwargs,
|
|
) -> Any:
|
|
headers = self.AUTHED_HEADERS if authed else self.HEADERS
|
|
async with self.session.request(
|
|
method, self.base_url + path, headers=headers, **kwargs
|
|
) as r:
|
|
match res_type:
|
|
case "json":
|
|
return await r.json()
|
|
case "text":
|
|
return await r.text()
|
|
case "bytes":
|
|
return await r.read()
|
|
|
|
async def screenshot(
|
|
self,
|
|
url: str,
|
|
selector: str | None = None,
|
|
waitForSelector: str | None = None,
|
|
waitForFunction: str | None = None,
|
|
) -> io.BytesIO:
|
|
"""Returns a PNG screenshot of the website at url with optional selector."""
|
|
params = {"url": url}
|
|
if selector:
|
|
params["selector"] = selector
|
|
if waitForSelector:
|
|
params["waitForSelector"] = waitForSelector
|
|
if waitForFunction:
|
|
params["waitForFunction"] = waitForFunction
|
|
|
|
res: bytes = await self._request(
|
|
"GET", "/webdriver/screenshot", authed=True, res_type="bytes", params=params
|
|
)
|
|
return io.BytesIO(res)
|
|
|
|
async def deepl(self, text: str, src: str = "auto", dst: str = "en") -> DeepLResult:
|
|
"""Returns DeepL translated text."""
|
|
data = {"src": src.lower(), "dst": dst.lower(), "text": text}
|
|
|
|
async with self.session.post(
|
|
self.base_url + "/webdriver/deepl", json=data, headers=self.AUTHED_HEADERS
|
|
) as r:
|
|
data = await r.json()
|
|
if not r.ok:
|
|
raise ArtemisError(f"DeepL Error: `{data.get('error', 'Unknown')}`")
|
|
return DeepLResult(**data)
|