maubot-tmdb/tmdb/tmdb.py

243 lines
10 KiB
Python
Raw Normal View History

'''
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/>.
'''
2020-06-22 12:15:28 +01:00
from html import escape
2020-06-26 15:41:44 +01:00
import re
2020-06-22 12:15:28 +01:00
from mautrix.types import TextMessageEventContent, MediaMessageEventContent, MessageType, Format
2020-06-22 12:15:28 +01:00
from maubot import Plugin, MessageEvent
from maubot.handlers import command
2020-09-11 10:15:04 +01:00
from tmdb.tmdb_api import TmdbApi, Movie, TvShow
2020-06-23 16:34:34 +01:00
from tmdb.database import Database
2020-06-22 12:15:28 +01:00
2020-06-22 15:21:18 +01:00
class TmdbBot(Plugin):
db: Database
2020-06-22 13:41:31 +01:00
async def start(self) -> None:
await super().start()
self.db = Database(self.database)
2020-09-11 10:15:04 +01:00
# 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>"""
2020-06-26 17:19:21 +01:00
2020-06-24 10:17:19 +01:00
async def send_html_message(self, evt: MessageEvent, text_message, html_message) -> None:
2020-06-22 12:15:28 +01:00
content = TextMessageEventContent(
msgtype=MessageType.TEXT, format=Format.HTML,
body=f"{text_message}",
formatted_body=f"{html_message}")
await evt.respond(content)
2020-07-03 14:06:33 +01:00
async def send_notice(self, evt: MessageEvent, message: str = "") -> None:
content = TextMessageEventContent(
msgtype=MessageType.NOTICE, format=Format.HTML,
body=message)
await evt.respond(content)
async def send_help(self, evt: MessageEvent) -> None:
html = """<p>Use <b>!tmdb id movie {tmdb id}</b> to get movie detail for tmdb-id.</p>
<p>Use <b>!tmdb movie {title} [y:{release year}]</b> to get movie detail based on the given title.</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>
<p>Use <b>!tmdb tvshow {title}</b> to get detail about a tv show based on the given title.</p>"""
content = TextMessageEventContent(
msgtype=MessageType.TEXT, format=Format.HTML,
body=f"Help for TMDB Bot",
formatted_body=f"{html}")
await evt.respond(content)
2020-06-24 10:17:19 +01:00
async def send_image(self, evt: MessageEvent, title, image) -> None:
mxc_uri = await self.client.upload_media(image)
2020-06-22 12:15:28 +01:00
content = MediaMessageEventContent(
msgtype=MessageType.IMAGE,
body=f"Image {title}",
2020-06-22 12:15:28 +01:00
url=f"{mxc_uri}")
await evt.respond(content)
def construct_html_message(self, movie, overview_length = 200, cast_length = 3) -> str:
if len(movie.overview) > overview_length:
three_dotts = " [...]"
else:
three_dotts = ""
cast = "Acting: "
for actor in movie.cast[:cast_length]:
cast+= f'{actor}, '
cast = cast[:-2]
2020-06-24 09:51:42 +01:00
html_message = f"""<p><a href="{movie.web_url}"><b>{escape(movie.title)}</b></a></p>
<p>{escape(movie.overview)[:overview_length]}{three_dotts}</p>
<p>{cast}</p>
2020-06-24 10:17:19 +01:00
<p>Taken from www.themoviedb.org</p>"""
return html_message
2020-06-26 15:41:44 +01:00
def split_title_year(self, message : str) -> (str, int):
m = re.search(r'^(.*) (y:\d\d\d\d)', message)
if m:
title = m.group(1)
year = int(m.group(2)[2:])
return (title, year)
return (message, None)
2020-09-11 10:15:04 +01:00
def set_language(self, evt: MessageEvent, movie : TmdbApi):
language = self.db.get_language(evt.sender)
if language:
movie.set_language(language)
def poster_size(self, evt: MessageEvent, movie : TmdbApi):
size = self.db.get_poster_size(evt.sender)
if size:
movie.set_poster_size(size)
async def send_movie_info(self, evt: MessageEvent, movie) -> None:
html_message = self.construct_html_message(movie)
2020-06-24 10:17:19 +01:00
await self.send_html_message(evt, f'{movie.title}', html_message)
await self.send_image(evt, movie.title, movie.get_image_binary())
2020-06-22 12:15:28 +01:00
async def movie_id(self, evt: MessageEvent, message: str = "") -> None:
movie = Movie()
2020-09-11 10:15:04 +01:00
self.poster_size(evt, movie)
2020-06-22 15:21:18 +01:00
language = self.db.get_language(evt.sender)
if language:
movie.set_language(language)
2020-06-22 12:15:28 +01:00
movie.query_details(message)
await self.send_movie_info(evt, movie)
2020-07-03 14:06:33 +01:00
@command.new("movie-id", help="Movie lookup by id")
2020-06-22 12:15:28 +01:00
@command.argument("message", pass_raw=True, required=True)
2020-07-03 14:06:33 +01:00
async def command_movie_id(self, evt: MessageEvent, message: str = "") -> None:
await self.movie_id(evt, message)
2020-06-22 12:15:28 +01:00
async def movie_search(self, evt: MessageEvent, message: str = "") -> None:
movie = Movie()
2020-09-11 10:15:04 +01:00
self.poster_size(evt, movie)
2020-06-22 13:41:31 +01:00
language = self.db.get_language(evt.sender)
if language:
movie.set_language(language)
2020-06-26 15:41:44 +01:00
title, year = self.split_title_year(message)
movie.search_title(title, year)
2020-06-22 12:15:28 +01:00
if movie.valid:
await self.send_movie_info(evt, movie)
else:
content = TextMessageEventContent(
msgtype=MessageType.NOTICE, format=Format.HTML,
body=f"No movie found!")
2020-09-16 16:00:42 +01:00
await evt.respond(content)
2020-06-22 13:41:31 +01:00
2020-07-03 14:06:33 +01:00
@command.new("movie-search", help="Movie lookup by Title")
2020-06-22 13:41:31 +01:00
@command.argument("message", pass_raw=True, required=True)
2020-07-03 14:06:33 +01:00
async def command_movie_search(self, evt: MessageEvent, message: str = "") -> None:
await self.movie_search(evt, message)
2020-06-22 13:41:31 +01:00
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,
2020-06-22 15:21:18 +01:00
body=f"Language set to {message}!")
await evt.respond(content)
2020-09-11 10:15:04 +01:00
async def set_poster_size(self, evt: MessageEvent, message: str = None) -> None:
movie = Movie()
2020-09-11 10:15:04 +01:00
poster_sizes = ""
for x in movie.poster_sizes:
poster_sizes += x + " "
if message:
size = movie.set_poster_size(message)
2020-09-11 10:15:04 +01:00
if size:
self.db.set_poster_size(evt.sender, size)
await self.send_notice(evt, f"Set default poster size to {size}")
else:
await self.send_notice(evt, f"Failed setting poster size. Valid sizes are {poster_sizes}.")
else:
2020-09-11 10:15:04 +01:00
await self.send_notice(evt, f"Valid sizes are {poster_sizes}.")
2020-06-24 10:17:19 +01:00
2020-07-03 14:06:33 +01:00
@command.new("movie-language", help="Set language for lookup")
@command.argument("message", pass_raw=True, required=True)
async def command_movie_language(self, evt: MessageEvent, message: str = "") -> None:
await self.movie_language(evt, message)
@command.new("movie-help", help="Help for TMDB Bot")
async def movie_help(self, evt: MessageEvent, message: str = "") -> None:
await self.send_help(evt)
@command.new("tvshow-help", help="Help for TMDB Bot")
async def tvshow_help(self, evt: MessageEvent, message: str = "") -> None:
await self.send_help(evt)
2020-06-24 10:17:19 +01:00
async def tvshow_search(self, evt: MessageEvent, message: str = "") -> None:
movie = TvShow()
2020-09-11 10:15:04 +01:00
self.poster_size(evt, movie)
2020-06-24 10:17:19 +01:00
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:
2020-09-11 10:15:04 +01:00
await self.send_notice(evt, "No tv show found!")
2020-07-03 14:06:33 +01:00
@command.new("tvshow-search", help="TV Show lookup by Title")
@command.argument("message", pass_raw=True, required=True)
async def command_tvshow_search(self, evt: MessageEvent, message: str = "") -> None:
await self.tvshow_search(evt, message)
@command.new("tmdb", help="TMDB Bot")
@command.argument("message", pass_raw=True, required=True)
async def command_dispatcher(self, evt: MessageEvent, message: str = "") -> None:
m = re.search(r'^([^\s]*)\s*(.*)', message)
if m:
command = m.group(1)
parameters = m.group(2)
if command.lower() == 'help':
await self.send_help(evt)
elif command.lower() == 'movie':
await self.movie_search(evt, parameters)
elif command.lower() == 'language':
await self.movie_language(evt, parameters)
elif command.lower() == 'poster_size':
2020-09-11 10:15:04 +01:00
await self.set_poster_size(evt, parameters)
2020-07-03 14:06:33 +01:00
elif command.lower() == 'tvshow':
await self.tvshow_search(evt, parameters)
elif command.lower() == 'id':
m = re.search(r'^(movie|tvshow) (\d*)', parameters)
if m:
type = m.group(1)
id = m.group(2)
if type.lower() == 'movie':
await self.movie_id(evt, id)
elif type.lower() == 'tvshow':
await self.send_notice(evt, 'Not yet implemented. Search TV shows by title for now (!tmdb tvshow {title})')
else:
await self.send_notice(evt, 'Syntax wrong: !tmdb id {movie|tvshow} {tmdb id}')
else:
await self.send_notice(evt, 'Syntax wrong: !tmdb id {movie|tvshow} {tmdb id}')
else:
await self.send_help(evt)
2020-07-03 14:06:33 +01:00
else:
await self.send_help(evt)