Add support to search multiple markets at once (#526)

* Add support to search multiple markets. Pass in a list or ALL to search all markets.

* pep8 formatting and verification with flake8

* work on comments from stephanebruckert

* pep8 formatting

* run autopep8 formatting

* fix typo

* allow tuple of markets to be passed. Add unit tests for this case

CHANGELOG.md

@@ -7,7 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## Unreleased
 
-// Add your changes here and then delete this line
+### Added
+
+ - Support to search multiple markets at once.
+ - Support to search all available Spotify markets.
 
 ## [2.13.0] - 2020-06-25
 

spotipy/client.py

@@ -34,6 +34,67 @@ class Spotify(object):
     """
     max_retries = 3
     default_retry_codes = (429, 500, 502, 503, 504)
+    country_codes = [
+        "AD",
+        "AR",
+        "AU",
+        "AT",
+        "BE",
+        "BO",
+        "BR",
+        "BG",
+        "CA",
+        "CL",
+        "CO",
+        "CR",
+        "CY",
+        "CZ",
+        "DK",
+        "DO",
+        "EC",
+        "SV",
+        "EE",
+        "FI",
+        "FR",
+        "DE",
+        "GR",
+        "GT",
+        "HN",
+        "HK",
+        "HU",
+        "IS",
+        "ID",
+        "IE",
+        "IT",
+        "JP",
+        "LV",
+        "LI",
+        "LT",
+        "LU",
+        "MY",
+        "MT",
+        "MX",
+        "MC",
+        "NL",
+        "NZ",
+        "NI",
+        "NO",
+        "PA",
+        "PY",
+        "PE",
+        "PH",
+        "PL",
+        "PT",
+        "SG",
+        "ES",
+        "SK",
+        "SE",
+        "CH",
+        "TW",
+        "TR",
+        "GB",
+        "US",
+        "UY"]
 
     def __init__(
         self,
@@ -447,19 +508,37 @@ class Spotify(object):
         tlist = [self._get_id("episode", e) for e in episodes]
         return self._get("episodes/?ids=" + ",".join(tlist), market=market)
 
-    def search(self, q, limit=10, offset=0, type="track", market=None):
+    def search(self, q, limit=10, offset=0, type="track", market=None, total=None):
         """ searches for an item
 
             Parameters:
                 - q - the search query (see how to write a query in the
                       official documentation https://developer.spotify.com/documentation/web-api/reference/search/search/)  # noqa
-                - limit  - the number of items to return (min = 1, default = 10, max = 50)
+                - limit  - the number of items to return (min = 1, default = 10, max = 50). If a search is to be done on multiple
+                            markets, then this limit is applied to each market. (e.g. search US, CA, MX each with a limit of 10).
                 - offset - the index of the first item to return
                 - type - the type of item to return. One of 'artist', 'album',
                          'track', 'playlist', 'show', or 'episode'
                 - market - An ISO 3166-1 alpha-2 country code or the string
-                           from_token.
+                           from_token. Can supply list of markets. Pass "ALL" to search all country codes.
+                - total - the total number of results to return if multiple markets are supplied in the search.
         """
+
+        if (isinstance(market, str) and market.upper() == "ALL"):
+            warnings.warn(
+                "Searching all markets is poorly performing.",
+                UserWarning,
+            )
+            return self._search_multiple_markets(q, limit, offset, type, self.country_codes, total)
+
+        elif isinstance(market, list) or isinstance(market, tuple):
+            warnings.warn(
+                "Searching multiple markets is poorly performing.",
+                UserWarning,
+            )
+            return self._search_multiple_markets(q, limit, offset, type, market, total)
+
+        else:
             return self._get(
                 "search", q=q, limit=limit, offset=offset, type=type, market=market
             )
@@ -1481,3 +1560,28 @@ class Spotify(object):
 
     def _get_uri(self, type, id):
         return "spotify:" + type + ":" + self._get_id(type, id)
+
+    def _search_multiple_markets(self, q, limit, offset, type, markets, total):
+        results = {
+                type + 's': {
+                    'href': [],
+                    'items': [],
+                    'limit': limit,
+                    'next': None,
+                    'offset': 0,
+                    'previous': None,
+                    'total': 0
+                }
+            }
+        for country in markets:
+            result = self._get(
+                "search", q=q, limit=limit, offset=offset, type=type, market=country
+            )
+            results[type + 's']['href'].append(result[type + 's']['href'])
+            results[type + 's']['items'] += result[type + 's']['items']
+            results[type + 's']['total'] += result[type + 's']['total']
+            if total and len(results[type + 's']['items']) >= total:
+                # splice 'items' to only include number of results requested
+                results[type + 's']['items'] = results[type + 's']['items'][:total]
+                return results
+        return results

spotipy/oauth2.py

@@ -32,6 +32,7 @@ logger = logging.getLogger(__name__)
 
 class SpotifyOauthError(Exception):
     """ Error during Auth Code or Implicit Grant flow """
+
     def __init__(self, message, error=None, error_description=None, *args, **kwargs):
         self.error = error
         self.error_description = error_description
@@ -41,6 +42,7 @@ class SpotifyOauthError(Exception):
 
 class SpotifyStateError(SpotifyOauthError):
     """ The state sent and state recieved were different """
+
     def __init__(self, local_state=None, remote_state=None, message=None,
                  error=None, error_description=None, *args, **kwargs):
         if not message:

tests/integration/test_non_user_endpoints.py

@@ -175,6 +175,40 @@ class AuthTestSpotipy(unittest.TestCase):
         self.assertTrue(len(results['artists']['items']) > 0)
         self.assertTrue(results['artists']['items'][0]['name'] == 'Weezer')
 
+    def test_artist_search_with_multiple_markets(self):
+        TOTAL = 3
+        results_single = self.spotify.search(q='weezer', type='artist', market='US')
+        results_multiple = self.spotify.search(q='weezer', type='artist',
+                                               market=['GB', 'US', 'AU'])
+        results_all = self.spotify.search(q='weezer', type='artist', market="ALL")
+        results_limited = self.spotify.search(q='weezer', type='artist',
+                                                market=['GB', 'US', 'AU'], total=TOTAL)
+
+        results_tuple = self.spotify.search(q='weezer', type='artist',
+                                            market=('GB', 'US', 'AU'), total=TOTAL)
+
+        self.assertTrue('artists' in results_multiple)
+        self.assertTrue('artists' in results_all)
+        self.assertTrue('artists' in results_limited)
+        self.assertTrue('artists' in results_tuple)
+
+        self.assertTrue(len(results_multiple['artists']['items']) > 0)
+        self.assertTrue(len(results_all['artists']['items']) > 0)
+        self.assertTrue(len(results_limited['artists']['items']) > 0)
+        self.assertTrue(len(results_tuple['artists']['items']) > 0)
+
+        self.assertTrue(len(results_all['artists']['items']) >
+                        len(results_multiple['artists']['items']))
+        self.assertTrue(len(results_multiple['artists']['items'])
+                        > len(results_single['artists']['items']))
+
+        self.assertTrue(results_multiple['artists']['items'][0]['name'] == 'Weezer')
+        self.assertTrue(results_all['artists']['items'][0]['name'] == 'Weezer')
+        self.assertTrue(results_limited['artists']['items'][0]['name'] == 'Weezer')
+        self.assertTrue(results_tuple['artists']['items'][0]['name'] == 'Weezer')
+
+        self.assertTrue(len(results_limited['artists']['items']) <= TOTAL)
+
     def test_artist_albums(self):
         results = self.spotify.artist_albums(self.weezer_urn)
         self.assertTrue('items' in results)