maubot-tmdb/tmdb/tmdb.py
lomion 2c5d3607f3 Changed release to 0.1.0
Some flake8 improvements
2020-06-23 20:12:43 +02:00

123 lines
4.5 KiB
Python

'''
This file is part of tmdb-bot.
tmdb-bot is free software: you can redistribute it and/or modify
it under the terms of the GNU Affero General Public License version 3 as published by
the Free Software Foundation.
tmdb-bot is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
GNU Affero General Public License for more details.
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/>.
'''
from html import escape
from mautrix.types import TextMessageEventContent, MediaMessageEventContent, MessageType, Format
from maubot import Plugin, MessageEvent
from maubot.handlers import command
from tmdb.tmdb_api import Movie
from sqlalchemy import (Column, String, Integer, Table, MetaData,
select)
from sqlalchemy.engine.base import Engine
class Database:
db: Engine
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(db)
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))
return rows.fetchone()
class TmdbBot(Plugin):
db: Database
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}'
if len(movie.overview) > 200:
three_dotts = " [...]"
else:
three_dotts = ""
cast = "Acting: "
for actor in movie.cast[:3]:
cast += f'{actor}, '
cast = cast[:-2]
html_message = f"""<p><b>{escape(movie.title)}</b></p>
<p>{escape(movie.overview)[:200]}{three_dotts}</p>
<p>{cast}</p>
<p>taken from www.themoviedb.org</p>"""
content = TextMessageEventContent(
msgtype=MessageType.TEXT, format=Format.HTML,
body=f"{text_message}",
formatted_body=f"{html_message}")
await evt.respond(content)
content = MediaMessageEventContent(
msgtype=MessageType.IMAGE,
body=f"Image {movie.title}",
url=f"{mxc_uri}")
await evt.respond(content)
@command.new("movie-id", help="Movie lookup by id")
@command.argument("message", pass_raw=True, required=True)
async def movie_id(self, evt: MessageEvent, message: str = "") -> None:
movie = Movie()
language = self.db.get_language(evt.sender)
if language:
movie.set_language(language)
movie.query_details(message)
await self.send_movie_info(evt, movie)
@command.new("movie-search", help="Movie lookup by Title")
@command.argument("message", pass_raw=True, required=True)
async def movie_search(self, evt: MessageEvent, message: str = "") -> None:
movie = Movie()
language = self.db.get_language(evt.sender)
if language:
movie.set_language(language)
movie.search_title(message)
if movie.valid:
await self.send_movie_info(evt, movie)
else:
content = TextMessageEventContent(
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.set_language(evt.sender, message)
content = TextMessageEventContent(
msgtype=MessageType.NOTICE, format=Format.HTML,
body=f"Language set to {message}!")
await evt.respond(content)