diff --git a/README.md b/README.md index 80526e7..4ada9d9 100644 --- a/README.md +++ b/README.md @@ -4,7 +4,7 @@ A [maubot](https://github.com/maubot/maubot) to get information about movies fro ## Usage Use `!movie-id {tmdb id}` to get movie detail for tmdb-id. -Use `!movie-search {title}` to get movie detail based on the given title. +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. diff --git a/test_tmdb.py b/test_tmdb.py index 123c241..ac8f3af 100644 --- a/test_tmdb.py +++ b/test_tmdb.py @@ -64,7 +64,23 @@ class TestTmdbMethods(unittest.TestCase): 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() diff --git a/tmdb/tmdb.py b/tmdb/tmdb.py index 4d9a8d1..2afd256 100644 --- a/tmdb/tmdb.py +++ b/tmdb/tmdb.py @@ -14,6 +14,7 @@ You should have received a copy of the GNU Affero General Public License along with tmdb-bot. If not, see . ''' from html import escape +import re from mautrix.types import TextMessageEventContent, MediaMessageEventContent, MessageType, Format, RelatesTo, RelationType @@ -62,6 +63,14 @@ class TmdbBot(Plugin):

Taken from www.themoviedb.org

""" 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) @@ -86,7 +95,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: diff --git a/tmdb/tmdb_api.py b/tmdb/tmdb_api.py index b2013b7..cb2ee1f 100644 --- a/tmdb/tmdb_api.py +++ b/tmdb/tmdb_api.py @@ -60,9 +60,11 @@ class Movie(TmdbApi): def __init__(self): super().__init__() - def search_title(self, title): + def search_title(self, title : str, year: int = None) -> int: payload = {} payload['query'] = title + if year: + payload['year'] = year result = self.request('search/movie', params=payload) json = result.json() if json['total_results'] > 0: