Add "offers" request

Add function requesting all offers for an entry with given node ID.
Allow specifying multiple countries to look for offers for.

Author: Electronic-Mango

src/simplejustwatchapi/justwatch.py

@@ -1,14 +1,21 @@
 """Main module orchestrating requests to JustWatch GraphQL API.
-Currently only search requests are supported.
+Currently supported requests are:
+ - ``search`` - search for entries via title
+ - ``details`` - get details for given node ID
+ - ``offers_for_countries`` - get all offers for entry with given node ID,
+   can look for offers for multiple countries
 """
 
 from httpx import post
 
 from simplejustwatchapi.query import (
     MediaEntry,
+    Offer,
     parse_details_response,
+    parse_offers_for_countries_response,
     parse_search_response,
     prepare_details_request,
+    prepare_offers_for_countries_request,
     prepare_search_request,
 )
 
@@ -55,9 +62,57 @@ def details(
         best_only: return only best offers if ``True``, return all offers if ``False``
 
     Returns:
-        MediaEntry NamedTuple with data about requested entry.
+        ``MediaEntry`` NamedTuple with data about requested entry.
     """
     request = prepare_details_request(node_id, country, language, best_only)
     response = post(_GRAPHQL_API_URL, json=request)
     response.raise_for_status()
     return parse_details_response(response.json())
+
+
+def offers_for_countries(
+    node_id: str, countries: set[str], language: str = "en", best_only: bool = True
+) -> dict[str, list[Offer]]:
+    """Get offers for entry of given node ID for all countries passed as argument.
+    Language argument only specifies format of price string, e.g. whether ".", or "," is used
+    in decimal fractions.
+
+    Returned dict has keys matching "countries" argument and values are list of found offers.
+    If no countries are passed (an empty set given as argument) empty dict is returned.
+
+    Country codes passed as argument are case-insensitive, however keys in returned dict will match
+    them exactly. E.g. for countries specified as:
+    .. code-block:: python
+
+        {"uK", "Us", "AU", "ca"}
+
+    returned dict will have the following structure:
+    .. code-block:: python
+
+        {
+            "uK": [... offers ...],
+            "Us": [... offers ...],
+            "AU": [... offers ...],
+            "ca": [... offers ...],
+        }
+
+    ``best_only`` allows filtering out redundant offers, e.g. when if provide offers service
+    in 4K, HD and SD, using ``best_only = True`` returns only 4K option, ``best_only = False``
+    returns all three.
+
+    Args:
+        node_id: ID of entry to look up offers for
+        countries: set of country codes to search for offers
+        language: language of responses, ``en`` by default
+        best_only: return only best offers if ``True``, return all offers if ``False``
+
+    Returns:
+        ``dict`` where keys match values in ``countries`` and keys are all found offers for their
+        respective countries
+    """
+    if not countries:
+        return {}
+    request = prepare_offers_for_countries_request(node_id, countries, language, best_only)
+    response = post(_GRAPHQL_API_URL, json=request)
+    response.raise_for_status()
+    return parse_offers_for_countries_response(response.json(), countries)

src/simplejustwatchapi/query.py

@@ -58,6 +58,24 @@
 }
 """
 
+_GRAPHQL_OFFERS_BY_COUNTRY_QUERY = """
+query GetTitleOffers(
+  $nodeId: ID!,
+  $language: Language!,
+  $formatOfferIcon: ImageFormat,
+  $filter: OfferFilter!,
+) {{
+  node(id: $nodeId) {{
+    ... on MovieOrShow {{
+      {country_entries}
+      __typename
+    }}
+    __typename
+  }}
+  __typename
+}}
+"""
+
 _GRAPHQL_DETAILS_FRAGMENT = """
 fragment TitleDetails on MovieOrShow {
   id
@@ -122,6 +140,13 @@
 }
 """
 
+_GRAPHQL_COUNTRY_OFFERS_ENTRY = """
+      {country_code}: offers(country: {country_code}, platform: WEB, filter: $filter) {{
+        ...TitleOffer
+        __typename
+      }}
+"""
+
 
 class OfferPackage(NamedTuple):
     """Parsed single offer package from JustWatch GraphQL API for single entry."""
@@ -270,10 +295,78 @@ def parse_details_response(json: any) -> MediaEntry | None:
     return _parse_entry(json["data"]["node"]) if "errors" not in json else None
 
 
+def prepare_offers_for_countries_request(
+    node_id: str, countries: set[str], language: str, best_only: bool
+) -> dict:
+    """Prepare an offers request for specified node ID and for all specified countries
+    to JustWatch GraphQL API.
+    Creates a ``GetTitleOffers`` GraphQL query.
+    Country codes should be two uppercase letters, however they will be auto-converted to uppercase.
+    ``countries`` argument mustn't be empty.
+
+    Args:
+        node_id: node ID of entry to get details for
+        countries: list of country codes to search for offers
+        language: language of responses
+        best_only: return only best offers if ``True``, return all offers if ``False``
+
+    Returns:
+        JSON/dict with GraphQL POST body
+    """
+    assert countries, "Cannot prepare offers request without specified countries"
+    for country in countries:
+        _assert_country_code_is_valid(country)
+    return {
+        "operationName": "GetTitleOffers",
+        "variables": {
+            "nodeId": node_id,
+            "language": language,
+            "formatPoster": "JPG",
+            "formatOfferIcon": "PNG",
+            "profile": "S718",
+            "backdropProfile": "S1920",
+            "filter": {"bestOnly": best_only},
+        },
+        "query": _prepare_offers_for_countries_entry(countries),
+    }
+
+
+def parse_offers_for_countries_response(json: any, countries: set[str]) -> dict[str, list[Offer]]:
+    """Parse response from offers query from JustWatch GraphQL API.
+    Parses response for ``GetTitleOffers`` query.
+    Response if searched for country codes passed as ``countries`` argument.
+    Countries in JSON response which are not present in ``countries`` set will be ignored.
+    If response doesn't have offers for a country, then that country still will be present
+    in returned dict, just with an empty list as value.
+
+    Args:
+        json: JSON returned by JustWatch GraphQL API
+        countries: set of countries to look for in API response
+
+    Returns:
+        A dict, where keys are matching ``countries`` argument and values are offers for a given
+        country parsed from JSON response.
+    """
+    offers_node = json["data"]["node"]
+    return {
+        country: list(map(_parse_offer, offers_node.get(country.upper(), [])))
+        for country in countries
+    }
+
+
 def _assert_country_code_is_valid(code: str) -> None:
     assert len(code) == 2, f"Invalid country code: {code}, code must be 2 characters long"
 
 
+def _prepare_offers_for_countries_entry(countries: set[str]) -> str:
+    offer_requests = [
+        _GRAPHQL_COUNTRY_OFFERS_ENTRY.format(country_code=country_code.upper())
+        for country_code in countries
+    ]
+    main_body = _GRAPHQL_OFFERS_BY_COUNTRY_QUERY.format(country_entries="\n".join(offer_requests))
+    return main_body + _GRAPHQL_OFFER_FRAGMENT
+
+
 def _parse_entry(json: any) -> MediaEntry:
     entry_id = json.get("id")
     object_id = json.get("objectId")

test/simplejustwatchapi/test_justwatch.py

@@ -1,7 +1,8 @@
 from unittest.mock import MagicMock, patch
+
 from pytest import fixture
 
-from simplejustwatchapi.justwatch import search, details
+from simplejustwatchapi.justwatch import details, search
 
 JUSTWATCH_GRAPHQL_URL = "https://apis.justwatch.com/graphql"
 SEARCH_INPUT = ("TITLE", "COUNTRY", "LANGUAGE", 5, True)
@@ -22,7 +23,7 @@ def httpx_post_mock(mocker):
 
 @patch("simplejustwatchapi.justwatch.parse_search_response", return_value=DUMMY_ENTRIES)
 @patch("simplejustwatchapi.justwatch.prepare_search_request", return_value=DUMMY_REQUEST)
-def test_search(requests_mock, parser_mock, httpx_post_mock) -> None:
+def test_search(requests_mock, parser_mock, httpx_post_mock):
     results = search(*SEARCH_INPUT)
     requests_mock.assert_called_with(*SEARCH_INPUT)
     parser_mock.assert_called_with(DUMMY_RESPONSE)
@@ -31,7 +32,7 @@ def test_search(requests_mock, parser_mock, httpx_post_mock) -> None:
 
 @patch("simplejustwatchapi.justwatch.parse_details_response", return_value=DUMMY_ENTRIES)
 @patch("simplejustwatchapi.justwatch.prepare_details_request", return_value=DUMMY_REQUEST)
-def test_details(requests_mock, parser_mock, httpx_post_mock) -> None:
+def test_details(requests_mock, parser_mock, httpx_post_mock):
     results = details(*DETAILS_INPUT)
     requests_mock.assert_called_with(*DETAILS_INPUT)
     parser_mock.assert_called_with(DUMMY_RESPONSE)

test/simplejustwatchapi/test_parser.py

@@ -1,6 +1,13 @@
 from pytest import mark
 
-from simplejustwatchapi.query import MediaEntry, Offer, parse_search_response, OfferPackage, parse_details_response
+from simplejustwatchapi.query import (
+    MediaEntry,
+    Offer,
+    OfferPackage,
+    parse_details_response,
+    parse_offers_for_countries_response,
+    parse_search_response,
+)
 
 DETAILS_URL = "https://justwatch.com"
 IMAGES_URL = "https://images.justwatch.com"
@@ -282,10 +289,10 @@
     argnames=["response_json", "expected_output"],
     argvalues=[
         (API_SEARCH_RESPONSE_JSON, [PARSED_NODE_1, PARSED_NODE_2, PARSED_NODE_3]),
-        (API_SEARCH_RESPONSE_NO_DATA, [])
-    ]
+        (API_SEARCH_RESPONSE_NO_DATA, []),
+    ],
 )
-def test_parse_search_response(response_json, expected_output) -> None:
+def test_parse_search_response(response_json: dict, expected_output: list[MediaEntry]) -> None:
     parsed_entries = parse_search_response(response_json)
     assert parsed_entries == expected_output
 
@@ -296,9 +303,47 @@ def test_parse_search_response(response_json, expected_output) -> None:
         ({"data": {"node": RESPONSE_NODE_1}}, PARSED_NODE_1),
         ({"data": {"node": RESPONSE_NODE_2}}, PARSED_NODE_2),
         ({"data": {"node": RESPONSE_NODE_3}}, PARSED_NODE_3),
-        ({"errors": [], "data": {"node": None}}, None)
-    ]
+        ({"errors": [], "data": {"node": None}}, None),
+    ],
 )
-def test_parse_details_response(response_json, expected_output) -> None:
+def test_parse_details_response(response_json: dict, expected_output: MediaEntry) -> None:
     parsed_entries = parse_details_response(response_json)
     assert parsed_entries == expected_output
+
+
+@mark.parametrize(
+    argnames=["response_json", "countries", "expected_output"],
+    argvalues=[
+        (
+            {"data": {"node": {"US": RESPONSE_NODE_1["offers"]}}},
+            {"US"},
+            {"US": PARSED_NODE_1.offers},
+        ),
+        (
+            {"data": {"node": {"US": RESPONSE_NODE_1["offers"], "GB": RESPONSE_NODE_2["offers"]}}},
+            {"US", "GB"},
+            {"US": PARSED_NODE_1.offers, "GB": PARSED_NODE_2.offers},
+        ),
+        (
+            {"data": {"node": {"US": RESPONSE_NODE_1["offers"], "GB": RESPONSE_NODE_2["offers"]}}},
+            {"US"},
+            {"US": PARSED_NODE_1.offers},
+        ),
+        (
+            {"data": {"node": {"US": RESPONSE_NODE_1["offers"], "GB": RESPONSE_NODE_2["offers"]}}},
+            {"GB"},
+            {"GB": PARSED_NODE_2.offers},
+        ),
+        (
+            {"data": {"node": {"US": RESPONSE_NODE_1["offers"], "GB": []}}},
+            {"US", "GB"},
+            {"US": PARSED_NODE_1.offers, "GB": []},
+        ),
+        ({"data": {"node": {"US": []}}}, {"US"}, {"US": []}),
+    ],
+)
+def test_parse_offers_for_countries_response(
+    response_json: dict, countries: set[str], expected_output: dict[str, list[Offer]]
+) -> None:
+    parsed_entries = parse_offers_for_countries_response(response_json, countries)
+    assert parsed_entries == expected_output