scott/maubot-tmdb
1
0
Fork 0
Code Issues Pull requests Packages Projects Releases Wiki Activity

Refactoring:

Browse source 
* Remove the movie poster parsing
* Fix some code issues
* Increase blurb trunctation length
This commit is contained in:
Scott Wallace 2024-11-28 09:03:49 +00:00
parent 0c400fb4d5
commit cc89e95403
Signed by: scott
SSH key fingerprint: SHA256:+LJug6Dj01Jdg86CILGng9r0lJseUrpI0xfRqdW9Uws
8 changed files with 423 additions and 315 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

10
aiohttp_tester.py
View file

@ -1,15 +1,13 @@
from tmdb.tmdb_api import Movie
import asyncio import asyncio
import time
from tmdb.tmdb_api import Movie
async def test(): async def test():
start = time.time()
movie = Movie() movie = Movie()
await movie.load_parameters() await movie.load_parameters()
id = await movie.search_title('Breakfast Club') await movie.search_title("Breakfast Club")
print("ID of Breakfast Club " + str(id))
print(time.time() - start)
await movie.close_session() await movie.close_session()
asyncio.run(test()) asyncio.run(test())

4
create_release.sh
View file

@ -1 +1,3 @@
zip -r lomion.tmdb.$1.mbp maubot.yaml tmdb/__init__.py tmdb/tmdb.py tmdb/tmdb_api.py tmdb/database.py #!/usr/bin/env bash
zip -r lomion.tmdb."$(grep version: maubot.yaml | cut -f2 -d' ')".mbp maubot.yaml tmdb/__init__.py tmdb/tmdb.py tmdb/tmdb_api.py tmdb/database.py

2
maubot.yaml
View file

@ -1,6 +1,6 @@
maubot: 0.1.0 maubot: 0.1.0
id: lomion.tmdb id: lomion.tmdb
version: 1.3.0+scott version: 1.3.0+scott-0.15
license: AGPL 3.0 license: AGPL 3.0
modules: modules:
- tmdb - tmdb

1
requirements.txt
View file

@ -1,2 +1,3 @@
maubot maubot
aiohttp aiohttp
sqlalchemy

211
test_tmdb.py
View file

@ -1,19 +1,26 @@
#!/usr/bin/env python3 """
Unit tests
"""
import unittest import unittest
from html import escape
from tmdb.tmdb_api import Movie, TvShow, MoviePopular
from tmdb.tmdb import TmdbBot
from tmdb.database import Database
from sqlalchemy import create_engine
import aiohttp import aiohttp
from mautrix.util.logging.trace import TraceLogger
from sqlalchemy import create_engine
from tmdb.database import Database
from tmdb.tmdb import TmdbBot
from tmdb.tmdb_api import Movie, MoviePopular, TvShow
# pylint: disable=missing-class-docstring,missing-function-docstring
async def apiRequests(command): async def api_requests(command):
api_key = '51d75c00dc1502dc894b7773ec3e7a15' api_key = "51d75c00dc1502dc894b7773ec3e7a15"
base_url = "https://api.themoviedb.org/3/" base_url = "https://api.themoviedb.org/3/"
url = base_url + command.lstrip('/') url = base_url + command.lstrip("/")
params = {'api_key': api_key} params = {"api_key": api_key}
params.update({'language': 'en'}) params.update({"language": "en"})
async with aiohttp.ClientSession() as client: async with aiohttp.ClientSession() as client:
async with client.get(url, params=params) as resp: async with client.get(url, params=params) as resp:
return await resp.json() return await resp.json()
@ -24,98 +31,104 @@ class TestTmdbMethods(unittest.IsolatedAsyncioTestCase):
async def test_search_item(self): async def test_search_item(self):
movie = Movie() movie = Movie()
await movie.load_parameters() await movie.load_parameters()
id = await movie.search_title('Breakfast Club') movie_id = await movie.search_title("Breakfast Club")
self.assertEqual(id, 2108) self.assertEqual(movie_id, 2108)
self.assertEqual(movie.valid, True) self.assertEqual(movie.valid, True)
await movie.close_session() await movie.close_session()
async def test_vote(self): async def test_vote(self):
movie = Movie() movie = Movie()
await movie.load_parameters() await movie.load_parameters()
await movie.search_title('Dune') await movie.search_title("Dune")
vote = movie.vote_average vote = movie.vote_average
self.assertEqual(vote, 8.0) self.assertGreaterEqual(vote, 7)
await movie.close_session()
async def test_vote(self):
movie = Movie()
await movie.load_parameters()
await movie.search_title('Dune')
vote = movie.vote_average
self.assertEqual(vote, 8.0)
await movie.close_session() await movie.close_session()
async def test_cast(self): async def test_cast(self):
movie = Movie() movie = Movie()
await movie.load_parameters() await movie.load_parameters()
await movie.search_title('Breakfast Club') await movie.search_title("Breakfast Club")
self.assertEqual('Emilio Estevez', movie.cast[0]) self.assertEqual("Emilio Estevez", movie.cast[0])
await movie.close_session() await movie.close_session()
async def test_title(self): async def test_title(self):
movie = Movie() movie = Movie()
await movie.load_parameters() await movie.load_parameters()
await movie.search_title('Breakfast Club') await movie.search_title("Breakfast Club")
self.assertEqual('The Breakfast Club', movie.title) self.assertEqual("The Breakfast Club", movie.title)
await movie.close_session() await movie.close_session()
async def test_overview(self): async def test_overview(self):
movie = Movie() movie = Movie()
await movie.load_parameters() await movie.load_parameters()
await movie.search_title('Breakfast Club') await movie.search_title("Breakfast Club")
description = 'Five high school students from different walks of' description = "Five high school students from different walks of"
self.assertEqual(description, movie.overview[: len(description)]) self.assertEqual(description, movie.overview[: len(description)])
await movie.close_session() await movie.close_session()
async def test_change_language(self): async def test_change_language(self):
movie = Movie() movie = Movie()
await movie.load_parameters() await movie.load_parameters()
movie.set_language('en') movie.set_language("en")
await movie.search_title('Breakfast Club') await movie.search_title("Breakfast Club")
description = 'Five high school students from different walks of life endure a Saturday detention' description = (
"Five high school students from different walks of "
"life endure a Saturday detention"
)
self.assertEqual(description, movie.overview[: len(description)]) self.assertEqual(description, movie.overview[: len(description)])
await movie.close_session() await movie.close_session()
def test_database_language(self): def test_database_language(self):
engine = create_engine('sqlite:///test.db', echo=True) engine = create_engine("sqlite:///test.db", echo=True)
db = Database(engine) db = Database(engine)
db.set_language('@testuser:example.com', 'de') db.set_language("@testuser:example.com", "de")
self.assertEqual(str(db.get_language('@testuser:example.com')), 'de') self.assertEqual(str(db.get_language("@testuser:example.com")), "de")
db.set_language('@testuser:example.com', 'en') db.set_language("@testuser:example.com", "en")
self.assertEqual(str(db.get_language('@testuser:example.com')), 'en') self.assertEqual(str(db.get_language("@testuser:example.com")), "en")
async def test_id_lookup(self): async def test_id_lookup(self):
movie = Movie() movie = Movie()
await movie.load_parameters() await movie.load_parameters()
await movie.query_details('2108') await movie.query_details("2108")
self.assertEqual('The Breakfast Club', movie.title) self.assertEqual("The Breakfast Club", movie.title)
await movie.close_session() await movie.close_session()
async def test_search_fails(self): async def test_search_fails(self):
movie = Movie() movie = Movie()
await movie.load_parameters() await movie.load_parameters()
id = await movie.search_title('Breakfast Club 2019') movie_id = await movie.search_title("Breakfast Club 2019")
self.assertEqual(id, None) self.assertEqual(movie_id, 0)
self.assertEqual(None, movie.title) self.assertEqual("", movie.title)
self.assertEqual(movie.valid, False) self.assertEqual(movie.valid, False)
await movie.close_session() await movie.close_session()
async def test_search_year(self): async def test_search_year(self):
movie = Movie() movie = Movie()
await movie.load_parameters() await movie.load_parameters()
id = await movie.search_title('Dune') movie_id = await movie.search_title("Dune")
self.assertEqual(id, 438631) self.assertEqual(movie_id, 438631)
id = await movie.search_title('Dune', 1984) movie_id = await movie.search_title("Dune", 1984)
self.assertEqual(id, 841) self.assertEqual(movie_id, 841)
await movie.close_session() await movie.close_session()
def test_split_year(self): def test_split_year(self):
tmdb = TmdbBot("", "", "", "", "", "", "", "", "", "") tmdb = TmdbBot(
title, year = tmdb.split_title_year('Dune') client=None, # pyright: ignore[reportArgumentType]
self.assertEqual('Dune', title) loop=None, # pyright: ignore[reportArgumentType]
self.assertEqual(None, year) http=None, # pyright: ignore[reportArgumentType]
title, year = tmdb.split_title_year('Dune y:2020 ') instance_id="",
self.assertEqual('Dune', title) log=TraceLogger(""),
config=None,
database=None,
webapp=None,
webapp_url="",
loader=None, # pyright: ignore[reportArgumentType]
)
title, year = tmdb.split_title_year("Dune")
self.assertEqual("Dune", title)
self.assertEqual(0, year)
title, year = tmdb.split_title_year("Dune y:2020 ")
self.assertEqual("Dune", title)
self.assertEqual(2020, year) self.assertEqual(2020, year)
async def test_set_poster_size(self): async def test_set_poster_size(self):
@ -131,95 +144,97 @@ class TestTmdbMethods(unittest.IsolatedAsyncioTestCase):
async def test_year_no_y(self): async def test_year_no_y(self):
movie = Movie() movie = Movie()
await movie.load_parameters() await movie.load_parameters()
id = await movie.search_title('infinite 2021') movie_id = await movie.search_title("infinite y:2021")
self.assertEqual(id, None) self.assertEqual(movie_id, 0)
self.assertEqual(movie.valid, False) self.assertEqual(movie.valid, False)
# TV Shows # TV Shows
async def test_search_tvshow(self): async def test_search_tvshow(self):
movie = TvShow() movie = TvShow()
await movie.load_parameters() await movie.load_parameters()
id = await movie.search_title('The Flash') movie_id = await movie.search_title("The Flash")
self.assertEqual(id, 60735) self.assertEqual(movie_id, 60735)
await movie.close_session() await movie.close_session()
async def test_tv_title(self): async def test_tv_title(self):
movie = TvShow() movie = TvShow()
await movie.load_parameters() await movie.load_parameters()
await movie.search_title('The Flash') await movie.search_title("The Flash")
self.assertEqual('The Flash', movie.title) self.assertEqual("The Flash", movie.title)
await movie.close_session() await movie.close_session()
async def test_search_tvshow_v(self): async def test_search_tvshow_v(self):
movie = TvShow() show = TvShow()
await movie.load_parameters() await show.load_parameters()
mid = await movie.search_title('V') mid = await show.search_title("V 2009")
self.assertEqual(mid, 21494) self.assertEqual(mid, 21494)
await movie.close_session() await show.close_session()
async def test_tv_title_v(self): async def test_tv_title_v(self):
movie = TvShow() movie = TvShow()
await movie.load_parameters() await movie.load_parameters()
await movie.search_title('V') await movie.search_title("V 1983")
self.assertEqual('V', movie.title) self.assertEqual("V", movie.title)
await movie.close_session() await movie.close_session()
async def test_search_tvshow_v_2009(self): async def test_search_tvshow_v_1983(self):
movie = TvShow() movie = TvShow()
await movie.load_parameters() await movie.load_parameters()
mid = await movie.search_title('V', 1983) mid = await movie.search_title("V", 1983)
self.assertEqual(mid, 14141) self.assertEqual(mid, 14141)
await movie.close_session() await movie.close_session()
async def test_tv_title_v_1983(self): async def test_tv_title_v_1983(self):
movie = TvShow() movie = TvShow()
await movie.load_parameters() await movie.load_parameters()
await movie.search_title('V', 1983) await movie.search_title("V", 1983)
self.assertEqual('V', movie.title) self.assertEqual("V", movie.title)
await movie.close_session() await movie.close_session()
async def test_cast_2(self): async def test_cast_2(self):
movie = TvShow() movie = TvShow()
await movie.load_parameters() await movie.load_parameters()
await movie.search_title('The Flash') await movie.search_title("The Flash")
self.assertEqual('Danielle Panabaker', movie.cast[2]) self.assertEqual("Danielle Panabaker", movie.cast[2])
await movie.close_session() await movie.close_session()
async def test_poster_path(self): # async def test_poster_path(self):
movie = Movie() # movie = Movie()
await movie.load_parameters() # await movie.load_parameters()
await movie.search_title('Dune') # await movie.search_title("Dune")
self.assertEqual(movie.poster_url, "http://image.tmdb.org/t/p/w92/d5NXSklXo0qyIYkgV94XAgMIckC.jpg") # self.assertEqual(
await movie.close_session() # movie.poster_url,
# "http://image.tmdb.org/t/p/w92/d5NXSklXo0qyIYkgV94XAgMIckC.jpg",
# )
# await movie.close_session()
async def test_movie_popular_length(self): async def test_movie_popular_length(self):
results = await apiRequests('/movie/popular') results = await api_requests("/movie/popular")
list = MoviePopular() movie_list = MoviePopular()
await list.load_parameters() await movie_list.load_parameters()
text = await list.query() text = await movie_list.query()
self.assertEqual(text, results['total_results']) self.assertEqual(text, results["total_results"])
await list.close_session() await movie_list.close_session()
async def test_movie_popular_id(self): async def test_movie_popular_id(self):
results = await apiRequests('/movie/popular') results = await api_requests("/movie/popular")
list = MoviePopular() movie_list = MoviePopular()
await list.load_parameters() await movie_list.load_parameters()
await list.query() await movie_list.query()
self.assertEqual(list.list[2]['id'], results['results'][2]['id']) self.assertEqual(movie_list.list[2]["id"], results["results"][2]["id"])
await list.close_session() await movie_list.close_session()
async def test_movie_popular_text(self): async def test_movie_popular_text(self):
results = await apiRequests('/movie/popular') results = await api_requests("/movie/popular")
list = MoviePopular() movie_list = MoviePopular()
await list.load_parameters() await movie_list.load_parameters()
await list.query() await movie_list.query()
test_result = results['results'][-1]['title'] test_result = results["results"][-1]["title"]
tested = list.getListText() tested = movie_list.getListText()
tested = tested[(len(results['results'][-1]['title'])) * -1:] tested = tested[(len(results["results"][-1]["title"])) * -1 :]
self.assertEqual(tested, test_result) self.assertEqual(tested, test_result)
await list.close_session() await movie_list.close_session()
if __name__ == "__main__":
if __name__ == '__main__':
unittest.main() unittest.main()

78
tmdb/database.py
View file

@ -1,4 +1,4 @@
''' """
This file is part of tmdb-bot. This file is part of tmdb-bot.
tmdb-bot is free software: you can redistribute it and/or modify tmdb-bot is free software: you can redistribute it and/or modify
@ -12,10 +12,13 @@ GNU Affero General Public License for more details.
You should have received a copy of the GNU Affero General Public License You should have received a copy of the GNU Affero General Public License
along with tmdb-bot. If not, see <https://www.gnu.org/licenses/>. along with tmdb-bot. If not, see <https://www.gnu.org/licenses/>.
''' """
# pylint: disable=missing-class-docstring,missing-function-docstring
from time import time from time import time
from sqlalchemy import (Column, String, Integer, Table, MetaData, select, Float)
from sqlalchemy import Column, Float, Integer, MetaData, String, Table, select
from sqlalchemy.engine.base import Engine from sqlalchemy.engine.base import Engine
@ -27,58 +30,83 @@ class Database:
meta = MetaData() meta = MetaData()
meta.bind = db meta.bind = db
self.language = Table("tmdb_language", meta, self.language = Table(
"tmdb_language",
meta,
Column("id", Integer, primary_key=True, autoincrement=True), Column("id", Integer, primary_key=True, autoincrement=True),
Column("user_id", String(255), nullable=False), Column("user_id", String(255), nullable=False),
Column("language", String(255), nullable=False),) Column("language", String(255), nullable=False),
self.tmdb_poster_size = Table("tmdb_poster_size", meta, )
self.tmdb_poster_size = Table(
"tmdb_poster_size",
meta,
Column("id", Integer, primary_key=True, autoincrement=True), Column("id", Integer, primary_key=True, autoincrement=True),
Column("user_id", String(255), nullable=False), Column("user_id", String(255), nullable=False),
Column("size", String(255), nullable=False),) Column("size", String(255), nullable=False),
self.tmdb_messages = Table("tmdb_messages", meta, )
self.tmdb_messages = Table(
"tmdb_messages",
meta,
Column("timestamp", Float, primary_key=True), Column("timestamp", Float, primary_key=True),
Column("event_id", String(64), nullable=False), Column("event_id", String(64), nullable=False),
Column("result_json", String(255), nullable=False),) Column("result_json", String(255), nullable=False),
)
meta.create_all(db) meta.create_all(db)
def set_message(self, event_id, result_json): def set_message(self, event_id, result_json):
with self.db.begin() as tx: with self.db.begin() as tx:
timestamp = time() timestamp = time()
tx.execute(self.tmdb_messages.insert().values(timestamp=timestamp, event_id=event_id, result_json=result_json)) tx.execute(
self.tmdb_messages.insert().values(
timestamp=timestamp, event_id=event_id, result_json=result_json
)
)
def get_message(self, event_id): def get_message(self, event_id):
rows = self.db.execute(select([self.tmdb_messages.c.result_json]) rows = self.db.execute(
.where(self.tmdb_messages.c.event_id == event_id)) select([self.tmdb_messages.c.result_json]).where(
self.tmdb_messages.c.event_id == event_id
)
)
row = rows.fetchone() row = rows.fetchone()
if row: if row:
return row['result_json'] return row["result_json"]
else:
return None return None
def set_language(self, user_id, language): def set_language(self, user_id, language):
with self.db.begin() as tx: with self.db.begin() as tx:
tx.execute(self.language.delete().where(self.language.c.user_id == user_id)) tx.execute(self.language.delete().where(self.language.c.user_id == user_id))
tx.execute(self.language.insert().values(user_id=user_id, language=language)) tx.execute(
self.language.insert().values(user_id=user_id, language=language)
)
def get_language(self, user_id): def get_language(self, user_id):
rows = self.db.execute(select([self.language.c.language]) rows = self.db.execute(
.where(self.language.c.user_id == user_id)) select([self.language.c.language]).where(self.language.c.user_id == user_id)
)
row = rows.fetchone() row = rows.fetchone()
if row: if row:
return row['language'] return row["language"]
else:
return None return None
def set_poster_size(self, user_id, size): def set_poster_size(self, user_id, size):
with self.db.begin() as tx: with self.db.begin() as tx:
tx.execute(self.tmdb_poster_size.delete().where(self.tmdb_poster_size.c.user_id == user_id)) tx.execute(
tx.execute(self.tmdb_poster_size.insert().values(user_id=user_id, size=size)) self.tmdb_poster_size.delete().where(
self.tmdb_poster_size.c.user_id == user_id
)
)
tx.execute(
self.tmdb_poster_size.insert().values(user_id=user_id, size=size)
)
def get_poster_size(self, user_id): def get_poster_size(self, user_id):
rows = self.db.execute(select([self.tmdb_poster_size.c.size]) rows = self.db.execute(
.where(self.tmdb_poster_size.c.user_id == user_id)) select([self.tmdb_poster_size.c.size]).where(
self.tmdb_poster_size.c.user_id == user_id
)
)
row = rows.fetchone() row = rows.fetchone()
if row: if row:
return row['size'] return row["size"]
else:
return None return None

170
tmdb/tmdb.py
View file

@ -1,4 +1,4 @@
''' """
This file is part of tmdb-bot. This file is part of tmdb-bot.
tmdb-bot is free software: you can redistribute it and/or modify tmdb-bot is free software: you can redistribute it and/or modify
@ -12,79 +12,107 @@ GNU Affero General Public License for more details.
You should have received a copy of the GNU Affero General Public License You should have received a copy of the GNU Affero General Public License
along with tmdb-bot. If not, see <https://www.gnu.org/licenses/>. along with tmdb-bot. If not, see <https://www.gnu.org/licenses/>.
''' """
from html import escape
import re # pylint: disable=missing-class-docstring,missing-function-docstring
import json import json
import re
from html import escape
from typing import Callable
from mautrix.types import TextMessageEventContent, MediaMessageEventContent, MessageType, Format, ImageInfo, EventType
from maubot import Plugin, MessageEvent
from maubot.handlers import command, event from maubot.handlers import command, event
from maubot.plugin_base import Plugin
from mautrix.types import (
EventType,
Format,
ImageInfo,
MediaMessageEventContent,
MessageType,
TextMessageEventContent,
)
from mautrix.types.event.message import MessageEvent
from tmdb.tmdb_api import TmdbApi, Movie, TvShow, MoviePopular from tmdb.database import Database, Engine
from tmdb.database import Database from tmdb.tmdb_api import Movie, MoviePopular, TmdbApi, TvShow
class MessageConstructor(): class MessageConstructor:
def __init__(self, movie: TmdbApi): def __init__(self, movie: TmdbApi):
self.movie = movie self.movie = movie
self.overview_length = 200 self.overview_length = 1024 * 32
self.cast_length = 3 self.cast_length = 3
def three_dotts(self): def three_dotts(self):
if len(self.movie.overview) > self.overview_length: if len(self.movie.overview) > self.overview_length:
return " [...]" return " [...]"
else:
return "" return ""
def cast(self): def cast(self):
cast = "Acting: " cast = "Acting: "
for actor in self.movie.cast[: self.cast_length]: for actor in self.movie.cast[: self.cast_length]:
cast += f'{actor}, ' cast += f"{actor}, "
return cast[:-2] return cast[:-2]
def construct_html_message(self) -> str: def construct_html_message(self) -> str:
html_message = f"""<p><a href="{self.movie.web_url}"><b>{escape(self.movie.title)}</b></a> - {str(int(self.movie.vote_average*10))}%</p> html_message = (
<p>{escape(self.movie.overview)[:self.overview_length]}{self.three_dotts()}</p> f'<p><a href="{self.movie.web_url}"><b>{escape(self.movie.title)}</b></a> - '
<p>{self.cast()}</p> f"{str(int(self.movie.vote_average*10))}%</p>"
<p>Taken from www.themoviedb.org</p>""" f"<p>{escape(self.movie.overview)[:self.overview_length]}{self.three_dotts()}</p>"
f"<p>{self.cast()}</p>"
f"<p>Taken from www.themoviedb.org</p>"
)
return html_message return html_message
class TmdbBot(Plugin): class TmdbBot(Plugin):
database: Engine
db: Database db: Database
api = None
async def start(self) -> None: async def start(self) -> None:
await super().start() await super().start()
self.db = Database(self.database) self.db = Database(self.database)
self.api = None self.api = None
async def send_html_message(self, evt: MessageEvent, text_message: str, html_message: str, data=None) -> None: async def send_html_message(
self, evt: MessageEvent, text_message: str, html_message: str, data=None
) -> None:
content = TextMessageEventContent( content = TextMessageEventContent(
msgtype=MessageType.TEXT, format=Format.HTML, msgtype=MessageType.TEXT,
format=Format.HTML,
body=f"{text_message}", body=f"{text_message}",
formatted_body=f"{html_message}") formatted_body=f"{html_message}",
)
event_id = await evt.respond(content) event_id = await evt.respond(content)
if data: if data:
self.db.set_message(event_id, json.dumps(data)) self.db.set_message(event_id, json.dumps(data))
async def send_notice(self, evt: MessageEvent, message: str = "") -> None: async def send_notice(self, evt: MessageEvent, message: str = "") -> None:
content = TextMessageEventContent( content = TextMessageEventContent(
msgtype=MessageType.NOTICE, format=Format.HTML, msgtype=MessageType.NOTICE, format=Format.HTML, body=message
body=message) )
await evt.respond(content) await evt.respond(content)
async def send_help(self, evt: MessageEvent) -> None: async def send_help(self, evt: MessageEvent) -> None:
html = """<p>Use <b>!tmdb movie {title} [y:{release year}]</b> to get movie detail based on the given title.</p> html = (
<p>Use <b>!tmdb tvshow {title}</b> to get detail about a tv show based on the given title.</p> "<p>Use <b>!tmdb movie {title} [y:{release year}]</b> to get movie detail based "
<p>Use <b>!tmdb popular [{rating}]</b> to get most popular movies. Get details about any one movie in the list by adding the {rating}.</p> "on the given title.</p>"
<p>Use <b>!tmdb language {language}</b> to set your prefered language.</p> "<p>Use <b>!tmdb tvshow {title}</b> to get detail about a tv show based on the "
<p>Use <b>!tmdb poster_size [{size}]</b> to set your prefered poster size. With empty {size} all available sizes are listed.</p>""" "given title.</p>"
"<p>Use <b>!tmdb popular [{rating}]</b> to get most popular movies. Get details "
"about any one movie in the list by adding the {rating}.</p>"
"<p>Use <b>!tmdb language {language}</b> to set your prefered language.</p>"
"<p>Use <b>!tmdb poster_size [{size}]</b> to set your prefered poster size. With "
"empty {size} all available sizes are listed.</p>"
)
content = TextMessageEventContent( content = TextMessageEventContent(
msgtype=MessageType.TEXT, format=Format.HTML, msgtype=MessageType.TEXT,
format=Format.HTML,
body="Help for TMDB Bot", body="Help for TMDB Bot",
formatted_body=f"{html}") formatted_body=f"{html}",
)
await evt.respond(content) await evt.respond(content)
async def send_image(self, evt: MessageEvent, title, image) -> None: async def send_image(self, evt: MessageEvent, title, image) -> None:
@ -94,16 +122,17 @@ class TmdbBot(Plugin):
msgtype=MessageType.IMAGE, msgtype=MessageType.IMAGE,
body=f"Image {title}", body=f"Image {title}",
url=f"{mxc_uri}", url=f"{mxc_uri}",
info=ImageInfo(mimetype='image/jpg')) info=ImageInfo(mimetype="image/jpg"),
)
await evt.respond(content) await evt.respond(content)
def split_title_year(self, message: str) -> (str, int): def split_title_year(self, message: str) -> tuple[str, int]:
m = re.search(r'^(.*) (y:\d\d\d\d)', message) m = re.search(r"^(.*) (y:\d\d\d\d)", message)
if m: if m:
title = m.group(1) title = m.group(1)
year = int(m.group(2)[2:]) year = int(m.group(2)[2:])
return (title, year) return (title, year)
return (message, None) return (message, 0)
def set_language(self, evt: MessageEvent, movie: TmdbApi): def set_language(self, evt: MessageEvent, movie: TmdbApi):
language = self.db.get_language(evt.sender) language = self.db.get_language(evt.sender)
@ -118,7 +147,7 @@ class TmdbBot(Plugin):
async def send_movie_info(self, evt: MessageEvent, movie) -> None: async def send_movie_info(self, evt: MessageEvent, movie) -> None:
constructor = MessageConstructor(movie) constructor = MessageConstructor(movie)
html_message = constructor.construct_html_message() html_message = constructor.construct_html_message()
await self.send_html_message(evt, f'{movie.title}', html_message) await self.send_html_message(evt, f"{movie.title}", html_message)
if movie.get_image_binary(): if movie.get_image_binary():
await self.send_image(evt, movie.title, movie.get_image_binary()) await self.send_image(evt, movie.title, movie.get_image_binary())
@ -161,7 +190,7 @@ class TmdbBot(Plugin):
await popular.query() await popular.query()
m = re.search(r'([1-5])', message) m = re.search(r"([1-5])", message)
if m: if m:
number = m.group(1) number = m.group(1)
movie = await popular.getMovieByNumber(number) movie = await popular.getMovieByNumber(number)
@ -170,9 +199,15 @@ class TmdbBot(Plugin):
await self.send_movie_info(evt, movie) await self.send_movie_info(evt, movie)
else: else:
text = popular.getListText(length) text = popular.getListText(length)
html = '<p><b>Currently most popular at <a href="https://www.themoviedb.org">www.themoviedb.org</a>:</b></p>' html = (
"<p><b>Currently most popular at "
'<a href="https://www.themoviedb.org">www.themoviedb.org</a>:</b></p>'
)
html += popular.getListHtml(length) html += popular.getListHtml(length)
html += '<p>For details reply to this message with the ranking number from this list</p>' html += (
"<p>For details reply to this message with the ranking number "
"from this list</p>"
)
results = popular.getDict(length) results = popular.getDict(length)
await self.send_html_message(evt, text, html, results) await self.send_html_message(evt, text, html, results)
@ -184,11 +219,13 @@ class TmdbBot(Plugin):
async def movie_language(self, evt: MessageEvent, message: str = "") -> None: async def movie_language(self, evt: MessageEvent, message: str = "") -> None:
self.db.set_language(evt.sender, message) self.db.set_language(evt.sender, message)
content = TextMessageEventContent( content = TextMessageEventContent(
msgtype=MessageType.NOTICE, format=Format.HTML, msgtype=MessageType.NOTICE,
body=f"Language set to {message}!") format=Format.HTML,
body=f"Language set to {message}!",
)
await evt.respond(content) await evt.respond(content)
async def set_poster_size(self, evt: MessageEvent, message: str = None) -> None: async def set_poster_size(self, evt: MessageEvent, message: str = "") -> None:
movie = await self.init_movie() movie = await self.init_movie()
poster_sizes = "" poster_sizes = ""
for x in movie.poster_sizes: for x in movie.poster_sizes:
@ -199,21 +236,28 @@ class TmdbBot(Plugin):
self.db.set_poster_size(evt.sender, size) self.db.set_poster_size(evt.sender, size)
await self.send_notice(evt, f"Set default poster size to {size}") await self.send_notice(evt, f"Set default poster size to {size}")
else: else:
await self.send_notice(evt, f"Failed setting poster size. Valid sizes are {poster_sizes}.") await self.send_notice(
evt,
f"Failed setting poster size. Valid sizes are {poster_sizes}.",
)
else: else:
await self.send_notice(evt, f"Valid sizes are {poster_sizes}.") await self.send_notice(evt, f"Valid sizes are {poster_sizes}.")
@command.new("movie-language", help="Set language for lookup") @command.new("movie-language", help="Set language for lookup")
@command.argument("message", pass_raw=True, required=True) @command.argument("message", pass_raw=True, required=True)
async def command_movie_language(self, evt: MessageEvent, message: str = "") -> None: async def command_movie_language(
self,
evt: MessageEvent,
message: str = "",
) -> None:
await self.movie_language(evt, message) await self.movie_language(evt, message)
@command.new("movie-help", help="Help for TMDB Bot") @command.new("movie-help", help="Help for TMDB Bot")
async def movie_help(self, evt: MessageEvent, message: str = "") -> None: async def movie_help(self, evt: MessageEvent) -> None:
await self.send_help(evt) await self.send_help(evt)
@command.new("tvshow-help", help="Help for TMDB Bot") @command.new("tvshow-help", help="Help for TMDB Bot")
async def tvshow_help(self, evt: MessageEvent, message: str = "") -> None: async def tvshow_help(self, evt: MessageEvent) -> None:
await self.send_help(evt) await self.send_help(evt)
async def tvshow_search(self, evt: MessageEvent, message: str = "") -> None: async def tvshow_search(self, evt: MessageEvent, message: str = "") -> None:
@ -236,6 +280,8 @@ class TmdbBot(Plugin):
async def movie_search(self, evt: MessageEvent, message: str = "") -> None: async def movie_search(self, evt: MessageEvent, message: str = "") -> None:
await self.init_movie() await self.init_movie()
if not self.api:
return
self.poster_size(evt, self.api) self.poster_size(evt, self.api)
language = self.db.get_language(evt.sender) language = self.db.get_language(evt.sender)
if language: if language:
@ -246,48 +292,50 @@ class TmdbBot(Plugin):
await self.send_movie_info(evt, self.api) await self.send_movie_info(evt, self.api)
else: else:
content = TextMessageEventContent( content = TextMessageEventContent(
msgtype=MessageType.NOTICE, format=Format.HTML, msgtype=MessageType.NOTICE, format=Format.HTML, body="No movie found!"
body="No movie found!") )
await evt.respond(content) await evt.respond(content)
@command.new("tmdb", help="TMDB Bot") @command.new("tmdb", help="TMDB Bot")
@command.argument("message", pass_raw=True, required=True) @command.argument("message", pass_raw=True, required=True)
async def command_dispatcher(self, evt: MessageEvent, message: str = "") -> None: async def command_dispatcher(self, evt: MessageEvent, message: str = "") -> None:
m = re.search(r'^([^\s]*)\s*(.*)', message) m = re.search(r"^([^\s]*)\s*(.*)", message)
if m: if m:
command = m.group(1) cmd = m.group(1)
parameters = m.group(2) parameters = m.group(2)
if command.lower() == 'help': match cmd.lower():
case "help":
await self.send_help(evt) await self.send_help(evt)
elif command.lower() == 'movie': case "movie":
await self.movie_search(evt, parameters) await self.movie_search(evt, parameters)
elif command.lower() == 'popular': case "popular":
await self.movie_popular(evt, parameters) await self.movie_popular(evt, parameters)
elif command.lower() == 'language': case "language":
await self.movie_language(evt, parameters) await self.movie_language(evt, parameters)
elif command.lower() == 'poster_size': case "poster_size":
await self.set_poster_size(evt, parameters) await self.set_poster_size(evt, parameters)
elif command.lower() == 'tvshow': case "tvshow":
await self.tvshow_search(evt, parameters) await self.tvshow_search(evt, parameters)
else: case _:
await self.send_help(evt) await self.send_help(evt)
else: else:
await self.send_help(evt) await self.send_help(evt)
if self.api: if self.api:
await self.api.close_session() await self.api.close_session()
@event.on(EventType.ROOM_MESSAGE) @event.on(EventType.ROOM_MESSAGE) # pylint: disable=no-member
async def handle_reply(self, evt: MessageEvent) -> None: async def handle_reply(self, evt: MessageEvent) -> None:
reply_to = evt.content.get_reply_to() reply_to = evt.content.get_reply_to
if reply_to: if isinstance(reply_to, Callable):
self.log.info(f"{evt.event_id} received. Reply to {evt.content.get_reply_to()}") self.log.info(f"{evt.event_id} received. Reply to {reply_to}")
result_json = self.db.get_message(reply_to) result_json = self.db.get_message(reply_to())
if result_json: if result_json:
requ = int(evt.content.body) requ = int(str(evt.content.body))
if requ > 0: if requ > 0:
populars = json.loads(result_json) populars = json.loads(result_json)
if str(requ) in populars: if str(requ) in populars:
await self.init_movie() await self.init_movie()
if self.api:
await self.api.search_id(populars[str(requ)]) await self.api.search_id(populars[str(requ)])
await self.send_movie_info(evt, self.api) await self.send_movie_info(evt, self.api)
else: else:

212
tmdb/tmdb_api.py
View file

@ -1,4 +1,4 @@
''' """
This file is part of tmdb-bot. This file is part of tmdb-bot.
tmdb-bot is free software: you can redistribute it and/or modify tmdb-bot is free software: you can redistribute it and/or modify
@ -12,34 +12,55 @@ GNU Affero General Public License for more details.
You should have received a copy of the GNU Affero General Public License You should have received a copy of the GNU Affero General Public License
along with tmdb-bot. If not, see <https://www.gnu.org/licenses/>. along with tmdb-bot. If not, see <https://www.gnu.org/licenses/>.
''' """
from html import escape
import aiohttp # pylint: disable=missing-class-docstring,missing-function-docstring
import asyncio import asyncio
from html import escape
import aiohttp
import aiohttp.client_exceptions
class TmdbApi(): class TmdbApi:
def __init__(self): def __init__(self):
self.session = aiohttp.ClientSession() self.session = aiohttp.ClientSession()
self.language = 'en' self.language = "en"
self.valid = False self.valid = False
self.overview = ""
self.cast = []
self.web_url = ""
self.title = ""
self.vote_average = 0.0
self.base_url_poster = ""
self.base_url = ""
self.api_key = ""
self.poster_sizes = []
self.base_url_images = ""
async def load_parameters(self): async def load_parameters(self):
self.api_key = '51d75c00dc1502dc894b7773ec3e7a15' self.api_key = "51d75c00dc1502dc894b7773ec3e7a15"
self.base_url = "https://api.themoviedb.org/3/" self.base_url = "https://api.themoviedb.org/3/"
async with self.session.get(self.base_url + 'configuration', params=self.get_apikey()) as resp: async with self.session.get(
self.base_url + "configuration", params=self.get_apikey()
) as resp:
result = await resp.json() result = await resp.json()
self.base_url_images = result['images']['base_url'] self.base_url_images = result["images"]["secure_base_url"]
self.base_url_poster = self.base_url_images + result['images']['poster_sizes'][0] self.base_url_poster = (
self.poster_sizes = result['images']['poster_sizes'] self.base_url_images + result["images"]["poster_sizes"][-4]
)
self.poster_sizes = result["images"]["poster_sizes"]
def get_apikey(self): def get_apikey(self):
return {'api_key': self.api_key} return {"api_key": self.api_key}
async def request(self, request_uri, params: dict = {}): async def request(self, request_uri, params: dict | None = None):
url = self.base_url + request_uri.lstrip('/') if not params:
params = {}
url = self.base_url + request_uri.lstrip("/")
params.update(self.get_apikey()) params.update(self.get_apikey())
params.update({'language': self.language}) params.update({"language": self.language})
result = None result = None
async with self.session.get(url, params=params) as resp: async with self.session.get(url, params=params) as resp:
result = await resp.json() result = await resp.json()
@ -63,18 +84,19 @@ class TmdbApi():
class TmdbApiSingle(TmdbApi): class TmdbApiSingle(TmdbApi):
def __init__(self): def __init__(self):
super().__init__() super().__init__()
self.title = None self.title = ""
self.id = None self.id = None
self.poster_url = None self.poster_url = ""
self.poster_binary = None self.poster_binary = None
self.overview = None self.overview = ""
self.web_url = None self.web_url = ""
self.vote_average = None self.vote_average = None
def get_image_binary(self): def get_image_binary(self):
return self.poster_binary return self.poster_binary
async def query_image_binary(self): async def query_image_binary(self):
self.poster_url = ""
if self.poster_url: if self.poster_url:
try: try:
async with self.session.get(self.poster_url) as resp: async with self.session.get(self.poster_url) as resp:
@ -92,43 +114,47 @@ class MoviePopular(TmdbApi):
self.length = 0 self.length = 0
async def query(self) -> int: async def query(self) -> int:
result = await self.request('/movie/popular') result = await self.request("/movie/popular")
self.length = result['total_results'] self.length = result["total_results"]
self.list = result['results'] self.list = result["results"]
return self.length return self.length
def getListHtml(self, length: int = None) -> str: def getListHtml(self, length: int = 0) -> str:
html = "" html = ""
if length: if length:
loop = length loop = length
else: else:
loop = self.length loop = self.length
id = 1 movie_id = 1
for element in self.list[:loop]: for element in self.list[:loop]:
html += f"""<p>{str(id)} - <a href="https://www.themoviedb.org/movie/{str(element['id'])}">{escape(element['title'])}</a> - {str(int(element['vote_average']*10))}%</p>""" html += (
id += 1 f"<p>{str(movie_id)} - "
f'<a href="https://www.themoviedb.org/movie/{str(element["id"])}">'
f"{escape(element['title'])}</a> - {str(int(element['vote_average']*10))}%</p>"
)
movie_id += 1
return html return html
def getListText(self, length: int = None) -> str: def getListText(self, length: int = 0) -> str:
text = "" text = ""
if length: if length:
loop = length loop = length
else: else:
loop = self.length loop = self.length
for element in self.list[:loop]: for element in self.list[:loop]:
text += element['title'] text += element["title"]
return text return text
def getDict(self, length: int = None) -> str: def getDict(self, length: int = 0) -> dict[int, int]:
result = {} result = {}
if length: if length:
loop = length loop = length
else: else:
loop = self.length loop = self.length
id = 1 movie_id = 1
for element in self.list[:loop]: for element in self.list[:loop]:
result[id] = str(element['id']) result[movie_id] = str(element["id"])
id += 1 movie_id += 1
return result return result
async def getMovieByNumber(self, number): async def getMovieByNumber(self, number):
@ -141,105 +167,95 @@ class MoviePopular(TmdbApi):
class Movie(TmdbApiSingle): class Movie(TmdbApiSingle):
def __init__(self):
super().__init__()
async def setData(self, data): async def setData(self, data):
self.title = data['title'] self.title = data["title"]
if not self.title: if not self.title:
self.valid = False self.valid = False
self.id = data['id'] self.id = data["id"]
self.poster_url = self.base_url_poster + data['poster_path'] self.poster_url = self.base_url_poster + data.get(
self.overview = data['overview'] "poster_path", "__no_poster_path__"
self.web_url = 'https://www.themoviedb.org/movie/' + str(self.id) )
self.vote_average = data['vote_average'] self.overview = data["overview"]
await asyncio.gather( self.web_url = "https://www.themoviedb.org/movie/" + str(self.id)
self.query_cast(self.id), self.vote_average = data["vote_average"]
self.query_image_binary()) await asyncio.gather(self.query_cast(self.id), self.query_image_binary())
return self.id return self.id
async def search_title(self, title: str, year: int = None) -> int: async def search_title(self, title: str, year: int = 0) -> int:
payload = {} payload = {}
payload['query'] = title payload["query"] = title
if year: if year:
payload['year'] = year payload["year"] = year
json = await self.request('search/movie', params=payload) json = await self.request("search/movie", params=payload)
if json['total_results'] > 0: if json["total_results"] > 0:
movie_id = json['results'][0]['id'] movie_id = json["results"][0]["id"]
await self.query_details(movie_id) await self.query_details(movie_id)
await asyncio.gather( await asyncio.gather(self.query_cast(movie_id), self.query_image_binary())
self.query_cast(movie_id),
self.query_image_binary())
return movie_id return movie_id
else:
self.valid = False self.valid = False
return None return 0
async def search_id(self, id): async def search_id(self, movie_id):
await self.query_details(id) await self.query_details(movie_id)
await asyncio.gather( await asyncio.gather(self.query_cast(movie_id), self.query_image_binary())
self.query_cast(id), return movie_id
self.query_image_binary())
return id
async def query_details(self, id): async def query_details(self, movie_id):
data = await self.request('movie/' + str(id)) data = await self.request("movie/" + str(movie_id))
self.title = data['title'] self.title = data["title"]
if not self.title: if not self.title:
self.valid = False self.valid = False
self.id = data['id'] self.id = data["id"]
self.poster_url = self.base_url_poster + data['poster_path'] self.poster_url = self.base_url_poster + data.get(
self.overview = data['overview'] "poster_path", "__no_poster_path__"
self.web_url = 'https://www.themoviedb.org/movie/' + str(self.id) )
self.vote_average = data['vote_average'] self.overview = data["overview"]
self.web_url = "https://www.themoviedb.org/movie/" + str(self.id)
self.vote_average = data["vote_average"]
async def query_cast(self, id): async def query_cast(self, movie_id):
data = await self.request('movie/' + str(id) + '/credits') data = await self.request("movie/" + str(movie_id) + "/credits")
self.cast = [] self.cast = []
for actor in data['cast']: for actor in data["cast"]:
self.cast.append(actor['name']) self.cast.append(actor["name"])
def get_cast(self, amount): def get_cast(self, amount):
return self.cast[:amount] return self.cast[:amount]
class TvShow(TmdbApiSingle): class TvShow(TmdbApiSingle):
def __init__(self): async def search_title(self, title, year: int = 0) -> int:
super().__init__()
async def search_title(self, title, year: int = 0):
payload = {} payload = {}
payload['query'] = title payload["query"] = title
if year: if year:
payload['first_air_date_year'] = str(year) payload["first_air_date_year"] = str(year)
json = await self.request('/search/tv', params=payload) json = await self.request("/search/tv", params=payload)
if json['total_results'] > 0: if json["total_results"] > 0:
movie_id = json['results'][0]['id'] movie_id = json["results"][0]["id"]
await self.query_details(movie_id) await self.query_details(movie_id)
await asyncio.gather( await asyncio.gather(self.query_cast(), self.query_image_binary())
self.query_cast(),
self.query_image_binary())
return movie_id return movie_id
else:
self.valid = False self.valid = False
return None return 0
async def query_details(self, id): async def query_details(self, movie_id):
data = await self.request('tv/' + str(id)) data = await self.request("tv/" + str(movie_id))
self.title = data['name'] self.title = data["name"]
if not self.title: if not self.title:
self.valid = False self.valid = False
self.id = data['id'] self.id = data["id"]
self.poster_url = self.base_url_poster + data['poster_path'] self.poster_url = self.base_url_poster + data.get(
self.overview = data['overview'] "poster_path", "__no_poster_path__"
self.web_url = 'https://www.themoviedb.org/tv/' + str(self.id) )
self.vote_average = data['vote_average'] self.overview = data["overview"]
self.web_url = "https://www.themoviedb.org/tv/" + str(self.id)
self.vote_average = data["vote_average"]
async def query_cast(self): async def query_cast(self):
data = await self.request('tv/' + str(self.id) + '/credits') data = await self.request("tv/" + str(self.id) + "/credits")
self.cast = [] self.cast = []
for actor in data['cast']: for actor in data["cast"]:
self.cast.append(actor['name']) self.cast.append(actor["name"])
def get_cast(self, amount): def get_cast(self, amount):
return self.cast[:amount] return self.cast[:amount]