Skip to content
This repository has been archived by the owner on Jun 24, 2024. It is now read-only.

Updated compatability for trademoe's new API #69

Open
wants to merge 1 commit into
base: master
Choose a base branch
from
Open

Updated compatability for trademoe's new API #69

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 5 additions & 5 deletions twsaucenao/sauce.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -117,7 +117,7 @@ async def _video_preview(self, sauce: AnimeSource, path_or_fh: typing.Union[str,
return None

try:
tracemoe_sauce = await tracemoe.search(path_or_fh, is_url=is_url)
tracemoe_sauce = await tracemoe.search(path_or_fh, is_url=is_url, anilist_id=sauce.anilist_id)
except ClientResponseError as e:
if e.status == 503:
self._log.warning("Tracemoe is not accepting API queries right now; aborting search query")
Expand All @@ -127,18 +127,18 @@ async def _video_preview(self, sauce: AnimeSource, path_or_fh: typing.Union[str,
except Exception:
self._log.exception("Tracemoe returned an exception; aborting search query")
return None
if not tracemoe_sauce.get('docs'):
if not tracemoe_sauce.get('result'):
self._log.info("Tracemoe returned no results")
return None

# Make sure our search results match
if await sauce.load_ids():
if sauce.anilist_id != tracemoe_sauce['docs'][0]['anilist_id']:
self._log.info(f"saucenao and trace.moe provided mismatched anilist entries: `{sauce.anilist_id}` vs. `{tracemoe_sauce['docs'][0]['anilist_id']}`")
if sauce.anilist_id != tracemoe_sauce['result'][0]['anilist']:
self._log.info(f"saucenao and trace.moe provided mismatched anilist entries: `{sauce.anilist_id}` vs. `{tracemoe_sauce['result'][0]['anilist']}`")
return None

self._log.info(f'Downloading video preview for AniList entry {sauce.anilist_id} from trace.moe')
tracemoe_preview = await tracemoe.video_preview_natural(tracemoe_sauce)
tracemoe_preview = await tracemoe.video_preview(tracemoe_sauce)
return tracemoe_preview

return None
Expand Down
141 changes: 50 additions & 91 deletions twsaucenao/tracemoe.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,27 +1,26 @@
# -*- coding: utf-8 -*-
# authors: Ethosa, FujiMakoto
import io
import re
from base64 import b64encode
from json import loads
from urllib.parse import urlencode

from aiohttp import ClientSession
from PIL import Image

from twsaucenao.config import config


class ATraceMoe:

DISCORD_IMAGE_URL_RE = re.compile(r'^https://cdn\.discordapp\.com/attachments/(\S)+\.(jpg|jpeg|png|webp|gif)$')
VIDEO_LARGE = 'l'
VIDEO_MEDIUM = 'm'
VIDEO_SMALL = 's'

def __init__(self, token=""):
"""
Initialize trace moe API.
"""
self.api_url = "https://trace.moe/api/"
self.main_url = "https://trace.moe/"
self.media_url = "https://media.trace.moe/"
self.api_url = "https://api.trace.moe"
self.token = token
self.session = ClientSession(
headers={
Expand All @@ -37,75 +36,28 @@ async def me(self):
Returns:
dict -- server response
"""
url = "%sme" % (self.api_url)
url = f"{self.api_url}/me"

if self.token:
url += "?token=%s" % (self.token)
url += f"?token={self.token}"

response = await self.session.get(url)

return await response.json()

async def image_preview(self, response, index=0, page="thumbnail.php"):
"""
Gets image preview after server response.

Arguments:
response {dict} -- server response

Returns:
bytes -- content for the write-in file.
"""
response = response["docs"][index]
url = "%s%s?anilist_id=%s&file=%s&t=%s&token=%s" % (
self.main_url, page, response["anilist_id"],
response["filename"], response["at"], response["tokenthumb"]
)
response = await self.session.get(url)

return await response.content.read()

async def video_preview(self, response, index=0):
"""
Gets video preview after server response.

Arguments:
response {dict} -- server response

Returns:
bytes -- content for the write-in file.
"""
return await self.image_preview(response, index, "preview.php")

async def video_preview_natural(self, response, index=0, mute=False):
"""
With trace.moe-media, it can now detect timestamp boundaries of a scene naturally.

Arguments:
response {dict} -- server response

Keyword Arguments:
index {number} -- index from response (default: {0})
mute {bool} -- mute video sound. {default: {False}}
async def video_preview(self, response, mute=False, size=None):
video_url = response['result'][0]['video']

Returns:
bytes -- content for the write-in file.
"""
response = response["docs"][index]
url = "%svideo/%s/%s?t=%s&token=%s" % (
self.media_url, response["anilist_id"],
response["filename"], response["at"],
response["tokenthumb"]
)
# Append the video size
size = size or self.VIDEO_LARGE
video_url += f"&size={size}"

if mute:
url += "&mute"

response = await self.session.get(url)
video_url += "&mute"

response = await self.session.get(video_url)
return await response.content.read()

async def search(self, path, search_filter=0, is_url=False):
async def search(self, path, search_filter=0, is_url=False, anilist_id=None):
"""
Searchs anime by image.

Expand All @@ -118,38 +70,45 @@ async def search(self, path, search_filter=0, is_url=False):
Returns:
dict -- server response
"""
url = "%ssearch" % (self.api_url)
url = f"{self.api_url}/search?"
args = {}

if self.token:
url += "?token=%s" % (self.token)
args['token'] = self.token

# Are we filtering by a specific anilist ID?
if anilist_id:
args['anilistID'] = anilist_id

url += urlencode(args)

if is_url:
# Discord URL's tend to break with trace.moe at the moment
if self.DISCORD_IMAGE_URL_RE.match(path):
# Load the image
response = await self.session.get(path, read_until_eof=False)
data = io.BytesIO(await response.read())

# Verify it's a valid image first
image = Image.open(data)
image.verify()
image = Image.open(data)

# If it's an animated gif, we need to extract a frame
if image.is_animated:
data = io.BytesIO()
image.seek(0)
image.save(data, format='PNG')

del image
encoded = b64encode(data.getvalue()).decode("utf-8")
response = await self.session.post(
url, json={"image": encoded, "filter": search_filter}
)
else:
response = await self.session.get(
url, params={"url": path}
)
# # Discord URL's tend to break with trace.moe at the moment
# if self.DISCORD_IMAGE_URL_RE.match(path):
# # Load the image
# response = await self.session.get(path, read_until_eof=False)
# data = io.BytesIO(await response.read())
#
# # Verify it's a valid image first
# image = Image.open(data)
# image.verify()
# image = Image.open(data)
#
# # If it's an animated gif, we need to extract a frame
# if image.is_animated:
# data = io.BytesIO()
# image.seek(0)
# image.save(data, format='PNG')
#
# del image
# encoded = b64encode(data.getvalue()).decode("utf-8")
# response = await self.session.post(
# url, json={"image": encoded, "filter": search_filter}
# )
# else:
response = await self.session.get(
url, params={"url": path}
)
return loads(await response.text())
elif isinstance(path, io.BufferedIOBase):
encoded = b64encode(path.read()).decode("utf-8")
Expand Down