Merge pull request #24 from dbatten5/imdb-id

Add support for imdb_id

docs/imdb.md

@@ -15,6 +15,30 @@ from phylm.sources import Imdb
 imdb = Imdb(raw_title="The Matrix")
 ```
 
+## Movie ID
+
+If you know the IMDb movie id you can pass it to `load_sources` to perform a more
+accurate search for the movie:
+
+```python
+phylm.load_source("imdb", movie_id="0133093")
+```
+
+`phylm` will first perform a search based on the ID. If the ID is valid the result will
+be selected, if not then it will fall back to a title search.
+
+Alternatively, instantiate the IMDb source class with the ID:
+
+```python
+from phylm.sources import Imdb
+
+imdb = Imdb(movie_id="0133093")
+```
+
+!!! warning ""
+    If instantiating the class directly you must supply at least one of `movie_id`
+    or `raw_title`, otherwise a `ValueError` will be raised.
+
 # Reference
 
 ::: phylm.sources.imdb.Imdb

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "phylm"
-version = "4.0.3"
+version = "4.0.4"
 description = "Phylm"
 authors = ["Dom Batten <[email protected]>"]
 license = "MIT"

src/phylm/phylm.py

@@ -20,7 +20,7 @@ def __init__(self, title: str) -> None:
         self._rt: Optional[Rt] = None
 
     @property
-    def imdb(self) -> Optional[Imdb]:
+    def imdb(self) -> Imdb:
         """Return the IMDb data.
 
         Returns:
@@ -35,7 +35,7 @@ def imdb(self) -> Optional[Imdb]:
         return self._imdb
 
     @property
-    def mtc(self) -> Optional[Mtc]:
+    def mtc(self) -> Mtc:
         """Return the Metacritic data.
 
         Returns:
@@ -68,11 +68,13 @@ def rt(self) -> Rt:
 
         return self._rt
 
-    def load_source(self, source: str) -> "Phylm":
+    def load_source(self, source: str, imdb_id: Optional[str] = None) -> "Phylm":
         """Load the film data for a source.
 
         Args:
-            source (str): the desired source
+            source: the desired source
+            imdb_id: an optional `IMDb` id which will be used to load the imdb data
+                instead of a basic search on the title
 
         Returns:
             the instance
@@ -82,7 +84,7 @@ def load_source(self, source: str) -> "Phylm":
         """
         if source == "imdb":
             if not self._imdb:
-                self._imdb = Imdb(raw_title=self.title)
+                self._imdb = Imdb(raw_title=self.title, movie_id=imdb_id)
             return self
 
         if source == "mtc":

src/phylm/sources/imdb.py

@@ -11,13 +11,49 @@
 class Imdb:
     """Class to abstract an IMDb movie object."""
 
-    def __init__(self, raw_title: str) -> None:
-        """Initialize the object."""
-        self.raw_title: str = raw_title
+    def __init__(
+        self,
+        raw_title: Optional[str] = None,
+        movie_id: Optional[str] = None,
+    ) -> None:
+        """Initialize the object.
+
+        Note that at least one of `raw_title` or `movie_id` must be given to be used as
+        a search term. `movie_id` is preferred over `raw_title`.
+
+        Args:
+            raw_title: the title of the movie
+            movie_id: the `IMDb` id of the movie
+
+        Raises:
+            ValueError: if neither `raw_title` nor `movie_id` is supplied
+        """
+        if not (raw_title or movie_id):
+            raise ValueError("At least one of raw_title and movie_id must be given")
+
+        self.raw_title: Optional[str] = raw_title
+        self.movie_id: Optional[str] = movie_id
         self.low_confidence: bool = False
         self._imdb_data: Optional[Movie] = self._get_imdb_data()
 
     def _get_imdb_data(self) -> Optional[Movie]:
+        """Fetch the data from IMDb.
+
+        If `self.movie_id` exists, prefer that as a search query, falling back to
+        `self.raw_title` if that exists. If `movie_id` exists but is unrecognised
+        by `IMDb` then we also fall back to the `raw_title`.
+
+        Returns:
+            an optional `IMDb` `Movie` object
+        """
+        if self.movie_id:
+            get_movie_result: Movie = ia.get_movie(self.movie_id)
+            if get_movie_result:
+                return get_movie_result
+
+        if not self.raw_title:
+            return None
+
         results: List[Movie] = ia.search_movie(self.raw_title)
         if not results:
             return None
@@ -30,6 +66,7 @@ def _get_imdb_data(self) -> Optional[Movie]:
             target = results[0]
             self.low_confidence = True
         ia.update(target, info=["main"])
+
         return target
 
     @property

tests/unit/sources/test_imdb.py

@@ -409,3 +409,83 @@ def test_no_movie_results(self, mock_ia: MagicMock) -> None:
         imdb = Imdb("The Movie")
 
         assert imdb.plot is None
+
+
+class TestMovieId:
+    """Tests for working with an IMDb `movie_id`."""
+
+    def test_no_raw_title_or_movie_id(self) -> None:
+        """
+        When an `Imdb` instance is created without a `raw_title` or ` movie_id`,
+        Then a `ValueError` is raised
+        """
+        with pytest.raises(
+            ValueError,
+            match="At least one of raw_title and movie_id must be given",
+        ):
+            Imdb()
+
+    @patch(IMDB_IA_PATH)
+    def test_valid_movie_id(
+        self,
+        mock_ia: MagicMock,
+        the_matrix_full: Movie,
+    ) -> None:
+        """
+        Given a valid `movie_id`,
+        When there is a match from IMDb,
+        Then the match is selected and `low_confidence` remains False
+        """
+        mock_ia.get_movie.return_value = the_matrix_full
+
+        imdb = Imdb(movie_id="0133093")
+
+        assert imdb.title == "The Matrix"
+        assert imdb.low_confidence is False
+
+    def test_invalid_movie_id_with_no_raw_title(
+        self,
+    ) -> None:
+        """
+        Given an unrecognised `movie_id` and no `raw_title` given,
+        When there are no matches from IMDb,
+        Then data remains as `None`
+        """
+        imdb = Imdb(movie_id="9999999999999999999999999999")
+
+        assert imdb.title is None
+
+    @patch(IMDB_IA_PATH)
+    def test_invalid_movie_id_with_raw_title(
+        self,
+        mock_ia: MagicMock,
+        the_matrix: Movie,
+    ) -> None:
+        """
+        Given an unrecognised `movie_id` and a `raw_title` given,
+        When there are no matches from IMDb id search,
+        Then the title is used to perform the search
+        """
+        mock_ia.get_movie.return_value = Movie(movieID="9999999999999999999999999999")
+        mock_ia.search_movie.return_value = [the_matrix]
+
+        imdb = Imdb(raw_title="The Matrix", movie_id="9999999999999999999999999999")
+
+        assert imdb.title == "The Matrix"
+
+    @patch(IMDB_IA_PATH)
+    def test_valid_movie_id_with_raw_title(
+        self,
+        mock_ia: MagicMock,
+        the_matrix: Movie,
+    ) -> None:
+        """
+        Given a recognized `movie_id` and a `raw_title` given,
+        Then the movie_id is used to perform the search
+        """
+        mock_ia.get_movie.return_value = the_matrix
+
+        imdb = Imdb(movie_id="0133093")
+
+        assert imdb.title == "The Matrix"
+        mock_ia.search_movie.assert_not_called()

tests/unit/test_phylm.py

@@ -59,7 +59,22 @@ def test_recognized_source_imdb(self, mock_imdb: MagicMock) -> None:
         phylm.load_source("imdb")
 
         assert phylm.imdb == mock_imdb.return_value
-        mock_imdb.assert_called_once_with(raw_title="bar")
+        mock_imdb.assert_called_once_with(raw_title="bar", movie_id=None)
+
+    @patch(f"{MODULE_PATH}.Imdb", autospec=True)
+    def test_recognized_source_imdb_with_movie_id(self, mock_imdb: MagicMock) -> None:
+        """
+        Given a `Phylm` instance,
+        When the `load_source` method is invoked with the `imdb` source
+            and a `movie_id`,
+        Then the source is loaded with the `movie_id`
+        """
+        phylm = Phylm(title="bar")
+
+        phylm.load_source("imdb", imdb_id="abc")
+
+        assert phylm.imdb == mock_imdb.return_value
+        mock_imdb.assert_called_once_with(raw_title="bar", movie_id="abc")
 
     @patch(f"{MODULE_PATH}.Mtc", autospec=True)
     def test_recognized_source_mtc(self, mock_mtc: MagicMock) -> None: