diff --git a/maubot.yaml b/maubot.yaml index 9dda137..31d9324 100644 --- a/maubot.yaml +++ b/maubot.yaml @@ -1,7 +1,8 @@ maubot: 0.1.0 id: lomion.tmdb -version: 0.0.1 +version: 0.0.2 license: AGPL 3.0 modules: - tmdb main_class: TmdbBot +database: true diff --git a/tmdb/tmdb.py b/tmdb/tmdb.py index 229b4e7..cbba73b 100644 --- a/tmdb/tmdb.py +++ b/tmdb/tmdb.py @@ -7,8 +7,17 @@ from maubot.handlers import command from tmdb.tmdb_api import Movie +from sqlalchemy import (Column, String, Integer, ForeignKey, Table, MetaData, + select, and_) +from sqlalchemy.engine.base import Engine + class TmdbBot(Plugin): + + async def start(self) -> None: + await super().start() + self.db = Database(self.database) + async def send_movie_info(self, evt: MessageEvent, movie) -> None: mxc_uri = await self.client.upload_media(data=movie.get_image_binary()) text_message = f'{movie.title}' @@ -48,6 +57,13 @@ class TmdbBot(Plugin): @command.argument("message", pass_raw=True, required=True) async def movie_search(self, evt: MessageEvent, message: str = "") -> None: movie = Movie() + self.db = Database(self.database) + language = self.db.get_language(evt.sender) + if language: + movie.set_language(language) + else: + self.db.set_language(evt.sender, 'de') + movie.set_language('de') movie.search_title(message) if movie.valid: await self.send_movie_info(evt, movie) @@ -56,3 +72,40 @@ class TmdbBot(Plugin): msgtype=MessageType.NOTICE, format=Format.HTML, body=f"No movie found!") await evt.respond(content) + + @command.new("movie-language", help="Set language for lookup") + @command.argument("message", pass_raw=True, required=True) + async def movie_language(self, evt: MessageEvent, message: str = "") -> None: + self.db = Database(self.database) + self.db.set_language(evt.sender, message) + content = TextMessageEventContent( + msgtype=MessageType.NOTICE, format=Format.HTML, + body=f"Language set to {language}!") + await evt.respond(content) + +class Database: + def __init__(self, db: Engine) -> None: + self.db = db + + meta = MetaData() + meta.bind = db + + self.language = Table("tmdb_language", meta, + Column("id", Integer, primary_key=True, autoincrement=True), + Column("user_id", String(255), nullable=False), + Column("language", String(255), nullable=False),) + meta.create_all() + + def set_language(self, user_id, language): + with self.db.begin() as tx: + 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)) + + def get_language(self, user_id): + rows = self.db.execute(select([self.language.c.language]) + .where(self.language.c.user_id == user_id)) + if len(rows) == 1: + row = rows[0] + return row[2] + else: + return None diff --git a/tmdb/tmdb_api.py b/tmdb/tmdb_api.py index f819671..03e9740 100644 --- a/tmdb/tmdb_api.py +++ b/tmdb/tmdb_api.py @@ -11,6 +11,8 @@ class Connection(): result = requests.get(self.base_url + 'configuration', params = self.get_apikey()).json() self.base_url_poster = result['images']['base_url'] + result['images']['poster_sizes'][0] self.valid = False + + self.language = 'de' def get_apikey(self): return { 'api_key' : self.api_key } @@ -18,11 +20,14 @@ class Connection(): def request(self, request_uri): url = self.base_url + request_uri payload = self.get_apikey() - payload['language'] = 'de' + payload['language'] = self.language result = requests.get(url, params=payload) self.valid = True return result + def set_language(self, language): + self.language = language + def get_image_binary(self): return requests.get(self.poster_url).content @@ -35,7 +40,7 @@ class Movie(Connection): def search_title(self, title): url = self.base_url+ 'search/movie' payload = self.get_apikey() - payload['language'] = 'de' + payload['language'] = self.language payload['query'] = title result = requests.get(url, params=payload) json = result.json()