diff --git a/test_tmdb.py b/test_tmdb.py index ac8f3af..68bc6a3 100644 --- a/test_tmdb.py +++ b/test_tmdb.py @@ -26,7 +26,7 @@ class TestTmdbMethods(unittest.TestCase): def test_overview(self): movie = Movie() movie.search_title('Breakfast Club') - description = 'Samstag morgen in einer amerikanischen High-School' + description = 'Five high school students from different walks of' self.assertEqual(description, movie.overview[:len(description)]) def test_change_language(self): @@ -42,7 +42,7 @@ class TestTmdbMethods(unittest.TestCase): movie.query_details('550') message = tmdb.construct_html_message(movie, overview_length = 10) self.assertEqual(message, """
-Ein Yuppie [...]
+A ticking- [...]
Acting: Edward Norton, Brad Pitt, Helena Bonham Carter
Taken from www.themoviedb.org
""") @@ -80,6 +80,13 @@ class TestTmdbMethods(unittest.TestCase): title, year = tmdb.split_title_year('Dune y:2020 ') self.assertEqual('Dune', title) self.assertEqual(2020, year) + + def test_set_poster_size(self): + movie = Movie() + size = movie.set_poster_size("w500") + self.assertEqual(size, "w500") + size = movie.set_poster_size("w666") + self.assertEqual(size, None) # TV Shows def test_search_tvshow(self): @@ -95,7 +102,7 @@ class TestTmdbMethods(unittest.TestCase): def test_cast(self): movie = TvShow() movie.search_title('The Flash') - self.assertEqual('Grant Gustin', movie.cast[0]) + self.assertEqual('Tom Cavanagh', movie.cast[0]) self.assertEqual('Carlos Valdes', movie.cast[2]) if __name__ == '__main__': diff --git a/tmdb/database.py b/tmdb/database.py index fdba0dd..4dfd103 100644 --- a/tmdb/database.py +++ b/tmdb/database.py @@ -31,6 +31,10 @@ class Database: Column("id", Integer, primary_key=True, autoincrement=True), Column("user_id", String(255), nullable=False), Column("language", String(255), nullable=False),) + self.tmdb_poster_size = Table("tmdb_poster_size", meta, + Column("id", Integer, primary_key=True, autoincrement=True), + Column("user_id", String(255), nullable=False), + Column("size", String(255), nullable=False),) meta.create_all(db) def set_language(self, user_id, language): @@ -41,4 +45,14 @@ class Database: 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 + return rows.fetchone()['language'] + + def set_poster_size(self, user_id, size): + 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(self.tmdb_poster_size.insert().values(user_id=user_id, size=size)) + + def get_poster_size(self, user_id): + rows = self.db.execute(select([self.tmdb_poster_size.c.size]) + .where(self.tmdb_poster_size.c.user_id == user_id)) + return rows.fetchone()['size'] \ No newline at end of file diff --git a/tmdb/tmdb.py b/tmdb/tmdb.py index 11b8cdb..4e12725 100644 --- a/tmdb/tmdb.py +++ b/tmdb/tmdb.py @@ -147,7 +147,18 @@ class TmdbBot(Plugin): content = TextMessageEventContent( msgtype=MessageType.NOTICE, format=Format.HTML, body=f"Language set to {message}!") - await evt.respond(content) + await evt.respond(content) + + async def poster_size(self, evt: MessageEvent, message: str = None) -> None: + movie = Movie() + if message: + size = movie.set_poster_size(message) + else: + 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) @command.new("movie-language", help="Set language for lookup") @command.argument("message", pass_raw=True, required=True) @@ -191,6 +202,8 @@ class TmdbBot(Plugin): await self.movie_search(evt, parameters) elif command.lower() == 'language': await self.movie_language(evt, parameters) + elif command.lower() == 'poster_size': + await self.poster_size(evt, parameters) elif command.lower() == 'tvshow': await self.tvshow_search(evt, parameters) elif command.lower() == 'id': diff --git a/tmdb/tmdb_api.py b/tmdb/tmdb_api.py index f3c0339..0d4a365 100644 --- a/tmdb/tmdb_api.py +++ b/tmdb/tmdb_api.py @@ -34,6 +34,7 @@ class TmdbApi(): 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.poster_sizes = result['images']['poster_sizes'] def get_apikey(self): return { 'api_key' : self.api_key } @@ -54,6 +55,13 @@ class TmdbApi(): return requests.get(self.poster_url).content else: return None + + def set_poster_size(self, size): + for x in self.poster_sizes: + if x == size: + self.base_url_poster = result['images']['base_url'] + x + return x + return None class Movie(TmdbApi):