diff --git a/.gitignore b/.gitignore index 7b8fdf6..892e881 100644 --- a/.gitignore +++ b/.gitignore @@ -130,3 +130,4 @@ dmypy.json .pyre/ lomion.tmdb.mbp +test.db diff --git a/README.md b/README.md index bff2726..c0164cd 100644 --- a/README.md +++ b/README.md @@ -2,11 +2,15 @@ A [maubot](https://github.com/maubot/maubot) to get information about movies from [TheMovieDB.org](https://www.themoviedb.org/). ## Usage -Use `!movie-id ` to get movie detail for tmdb-id. +Use `!movie-help` to get help about the bot usage. -Use `!movie-search ` to get movie detail based on the given title. +Use `!movie-id {tmdb id}` to get movie detail for tmdb-id. -Use `!movie-language <language>` to set your prefered language. +Use `!movie-search {title} [y:{release year}]` to get movie detail based on the given title. + +Use `!movie-language {language}` to set your prefered language. + +Use `!tvshow-search {title}` to get detail about a tv show based on the given title. ## Discussion Matrix room: [#tmdb-bot:matrix.sarkasti.eu](https://matrix.to/#/#tmdb-bot:matrix.sarkasti.eu) \ No newline at end of file diff --git a/maubot.yaml b/maubot.yaml index 5009c62..92fe944 100644 --- a/maubot.yaml +++ b/maubot.yaml @@ -1,6 +1,6 @@ maubot: 0.1.0 id: lomion.tmdb -version: 0.1.0 +version: 0.2.0 license: AGPL 3.0 modules: - tmdb diff --git a/test_tmdb.py b/test_tmdb.py new file mode 100644 index 0000000..ac8f3af --- /dev/null +++ b/test_tmdb.py @@ -0,0 +1,102 @@ +#!/usr/bin/env python3 +import unittest +from tmdb.tmdb_api import Movie, TvShow +from tmdb.tmdb import TmdbBot +from tmdb.database import Database +from sqlalchemy import create_engine + + +class TestTmdbMethods(unittest.TestCase): + ### TMDB API + def test_search_item(self): + movie = Movie() + id = movie.search_title('Breakfast Club') + self.assertEqual(id, 2108) + + def test_cast(self): + movie = Movie() + movie.search_title('Breakfast Club') + self.assertEqual('Anthony Michael Hall', movie.cast[0]) + + def test_title(self): + movie = Movie() + movie.search_title('Breakfast Club') + self.assertEqual('The Breakfast Club', movie.title) + + def test_overview(self): + movie = Movie() + movie.search_title('Breakfast Club') + description = 'Samstag morgen in einer amerikanischen High-School' + self.assertEqual(description, movie.overview[:len(description)]) + + def test_change_language(self): + movie = Movie() + movie.set_language('en') + movie.search_title('Breakfast Club') + description = 'Five high school students from different walks of life endure a Saturday detention' + self.assertEqual(description, movie.overview[:len(description)]) + + def test_html_construction(self): + movie = Movie() + tmdb = TmdbBot("","" ,"" ,"" ,"" ,"" ,"" ,"" ,"" ) + movie.query_details('550') + message = tmdb.construct_html_message(movie, overview_length = 10) + self.assertEqual(message, """<p><a href="https://www.themoviedb.org/movie/550"><b>Fight Club</b></a></p> + <p>Ein Yuppie [...]</p> + <p>Acting: Edward Norton, Brad Pitt, Helena Bonham Carter</p> + <p>Taken from www.themoviedb.org</p>""") + + def test_database_language(self): + engine = create_engine('sqlite:///test.db', echo = True) + db = Database(engine) + db.set_language('@testuser:example.com', 'de') + self.assertEqual(str(db.get_language('@testuser:example.com')), 'de') + db.set_language('@testuser:example.com', 'en') + self.assertEqual(str(db.get_language('@testuser:example.com')), 'en') + + def test_id_lookup(self): + movie = Movie() + movie.query_details('2108') + self.assertEqual('The Breakfast Club', movie.title) + + def test_search_fails(self): + movie = Movie() + id = movie.search_title('Breakfast Club 2019') + self.assertEqual(id, None) + self.assertEqual(None, movie.title) + + def test_search_year(self): + movie = Movie() + id = movie.search_title('Dune') + self.assertEqual(id, 841) + id = movie.search_title('Dune', 2020) + self.assertEqual(id, 438631) + + def test_split_year(self): + tmdb = TmdbBot("","" ,"" ,"" ,"" ,"" ,"" ,"" ,"" ) + title, year = tmdb.split_title_year('Dune') + self.assertEqual('Dune', title) + self.assertEqual(None, year) + title, year = tmdb.split_title_year('Dune y:2020 ') + self.assertEqual('Dune', title) + self.assertEqual(2020, year) + + # TV Shows + def test_search_tvshow(self): + movie = TvShow() + id = movie.search_title('The Flash') + self.assertEqual(id, 60735) + + def test_tv_title(self): + movie = TvShow() + movie.search_title('The Flash') + self.assertEqual('The Flash', movie.title) + + def test_cast(self): + movie = TvShow() + movie.search_title('The Flash') + self.assertEqual('Grant Gustin', movie.cast[0]) + self.assertEqual('Carlos Valdes', movie.cast[2]) + +if __name__ == '__main__': + unittest.main() diff --git a/tmdb/database.py b/tmdb/database.py new file mode 100644 index 0000000..fdba0dd --- /dev/null +++ b/tmdb/database.py @@ -0,0 +1,44 @@ +''' +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 sqlalchemy import (Column, String, Integer, ForeignKey, Table, MetaData, + select, and_) +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()['language'] \ No newline at end of file diff --git a/tmdb/test_tmdb.py b/tmdb/test_tmdb.py deleted file mode 100644 index 4d0cdc2..0000000 --- a/tmdb/test_tmdb.py +++ /dev/null @@ -1,36 +0,0 @@ -#!/usr/bin/env python3 -import unittest -from tmdb_api import Movie - -class TestTmdbMethods(unittest.TestCase): - ### TMDB API - def test_search_item(self): - movie = Movie() - id = movie.search_title('Breakfast Club') - self.assertEqual(id, 2108) - - def test_cast(self): - movie = Movie() - movie.search_title('Breakfast Club') - self.assertEqual('Anthony Michael Hall', movie.cast[0]) - - def test_title(self): - movie = Movie() - movie.search_title('Breakfast Club') - self.assertEqual('The Breakfast Club', movie.title) - - def test_overview(self): - movie = Movie() - movie.search_title('Breakfast Club') - description = 'Samstag morgen in einer amerikanischen High-School' - self.assertEqual(description, movie.overview[:len(description)]) - - def test_change_language(self): - movie = Movie() - movie.set_language('en') - movie.search_title('Breakfast Club') - description = 'Five high school students from different walks of life endure a Saturday detention' - self.assertEqual(description, movie.overview[:len(description)]) - -if __name__ == '__main__': - unittest.main() diff --git a/tmdb/tmdb.py b/tmdb/tmdb.py index f20af64..7b0da4c 100644 --- a/tmdb/tmdb.py +++ b/tmdb/tmdb.py @@ -14,43 +14,15 @@ 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 +import re 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() +from tmdb.tmdb_api import Movie, TvShow +from tmdb.database import Database class TmdbBot(Plugin): @@ -75,17 +47,62 @@ class TmdbBot(Plugin): <p>{escape(movie.overview)[:200]}{three_dotts}</p> <p>{cast}</p> <p>taken from www.themoviedb.org</p>""" + + async def send_html_message(self, evt: MessageEvent, text_message, html_message) -> None: content = TextMessageEventContent( msgtype=MessageType.TEXT, format=Format.HTML, body=f"{text_message}", formatted_body=f"{html_message}") await evt.respond(content) + + async def send_help(self, evt: MessageEvent) -> None: + html = """Use <b>!movie-id {tmdb id}</b> to get movie detail for tmdb-id.</br> + Use <b>!movie-search {title} [y:{release year}]</b> to get movie detail based on the given title.</br> + Use <b>!movie-language {language}</b> to set your prefered language.</br> + Use <b>!tvshow-search {title}</b> to get detail about a tv show based on the given title.</br>""" + content = TextMessageEventContent( + msgtype=MessageType.TEXT, format=Format.HTML, + body=f"Help for TMDB Bot", + formatted_body=f"{html}") + await evt.respond(content) + + async def send_image(self, evt: MessageEvent, title, image) -> None: + mxc_uri = await self.client.upload_media(image) content = MediaMessageEventContent( msgtype=MessageType.IMAGE, - body=f"Image {movie.title}", + body=f"Image {title}", 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] + 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> + <p>Taken from www.themoviedb.org</p>""" + return html_message + + 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) + + async def send_movie_info(self, evt: MessageEvent, movie) -> None: + html_message = self.construct_html_message(movie) + await self.send_html_message(evt, f'{movie.title}', html_message) + await self.send_image(evt, movie.title, movie.get_image_binary()) + @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: @@ -103,7 +120,8 @@ class TmdbBot(Plugin): language = self.db.get_language(evt.sender) if language: movie.set_language(language) - movie.search_title(message) + title, year = self.split_title_year(message) + movie.search_title(title, year) if movie.valid: await self.send_movie_info(evt, movie) else: @@ -119,4 +137,28 @@ class TmdbBot(Plugin): content = TextMessageEventContent( msgtype=MessageType.NOTICE, format=Format.HTML, body=f"Language set to {message}!") + await evt.respond(content) + + @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) + + @command.new("tvshow-search", help="TV Show lookup by Title") + @command.argument("message", pass_raw=True, required=True) + async def tvshow_search(self, evt: MessageEvent, message: str = "") -> None: + movie = TvShow() + 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 tv show found!") await evt.respond(content) diff --git a/tmdb/tmdb_api.py b/tmdb/tmdb_api.py index c6a82a4..cb2ee1f 100644 --- a/tmdb/tmdb_api.py +++ b/tmdb/tmdb_api.py @@ -15,9 +15,18 @@ along with tmdb-bot. If not, see <https://www.gnu.org/licenses/>. ''' import requests -class Connection(): +class TmdbApi(): def __init__(self): self.load_parameters() + self.title = None + self.id = None + self.poster_url = None + self.overview = None + self.web_url = None + self.vote_average = None + self.language = 'de' + + self.valid = False def load_parameters(self): self.api_key = '51d75c00dc1502dc894b7773ec3e7a15' @@ -25,18 +34,15 @@ class Connection(): self.base_url = "https://api.themoviedb.org/3/" 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 } - def request(self, request_uri): - url = self.base_url + request_uri - payload = self.get_apikey() - payload['language'] = self.language - result = requests.get(url, params=payload) + def request(self, request_uri, params : dict = {}): + url = self.base_url + request_uri.lstrip('/') + params.update(self.get_apikey()) + params.update({ 'language' : self.language }) + result = requests.get(url, params=params) self.valid = True return result @@ -44,20 +50,22 @@ class Connection(): self.language = language def get_image_binary(self): - return requests.get(self.poster_url).content + if self.poster_url: + return requests.get(self.poster_url).content + else: + return None -class Movie(Connection): +class Movie(TmdbApi): def __init__(self): - self.load_parameters() - pass + super().__init__() - def search_title(self, title): - url = self.base_url+ 'search/movie' - payload = self.get_apikey() - payload['language'] = self.language + def search_title(self, title : str, year: int = None) -> int: + payload = {} payload['query'] = title - result = requests.get(url, params=payload) + if year: + payload['year'] = year + result = self.request('search/movie', params=payload) json = result.json() if json['total_results'] > 0: movie_id = json['results'][0]['id'] @@ -82,4 +90,37 @@ class Movie(Connection): def get_cast(self, amount): return self.cast[:amount] + + +class TvShow(TmdbApi): + def __init__(self): + super().__init__() + def search_title(self, title): + payload = {} + payload['query'] = title + result = self.request('/search/tv', params=payload) + json = result.json() + if json['total_results'] > 0: + movie_id = json['results'][0]['id'] + self.query_details(movie_id) + return movie_id + + def query_details(self, id): + data = self.request('tv/' + str(id)).json() + self.title = data['name'] + self.id = data['id'] + self.poster_url = self.base_url_poster + data['poster_path'] + self.overview = data['overview'] + self.web_url = 'https://www.themoviedb.org/tv/' + str(self.id) + self.vote_average = str(data['vote_average']) + self.query_cast() + + def query_cast(self): + data = self.request('tv/'+str(self.id)+'/credits').json() + self.cast = [] + for actor in data['cast']: + self.cast.append(actor['name']) + + def get_cast(self, amount): + return self.cast[:amount] \ No newline at end of file