Add support for searching genre stations. (#45)

Add support for searching genre stations.

2 files changed, 82 insertions, 15 deletions

diff --git a/pandora/client.py b/pandora/client.py
index 7c9791e..48077cb 100644
--- a/pandora/client.py
+++ b/pandora/client.py
@@ -174,12 +174,18 @@ class APIClient(BaseAPIClient):
         return self("bookmark.deleteArtistBookmark",
                     bookmarkToken=bookmark_token)
 
-    def search(self, search_text):
+    def search(self, search_text,
+               include_near_matches=False,
+               include_genre_stations=False):
         from .models.pandora import SearchResult
 
-        return SearchResult.from_json(self,
-                                      self("music.search",
-                                           searchText=search_text))
+        return SearchResult.from_json(
+            self,
+            self("music.search",
+                 searchText=search_text,
+                 includeNearMatches=include_near_matches,
+                 includeGenreStations=include_genre_stations)
+        )
 
     def add_feedback(self, track_token, positive):
         return self("station.addFeedback",
diff --git a/pandora/models/pandora.py b/pandora/models/pandora.py
index f2fd36e..8fa8e21 100644
--- a/pandora/models/pandora.py
+++ b/pandora/models/pandora.py
@@ -284,37 +284,98 @@ class BookmarkList(PandoraModel):
 
 class SearchResultItem(PandoraModel):
 
-    artist = Field("artistName")
-    song_name = Field("songName")
     score = Field("score")
-    likely_match = Field("likelyMatch", default=False)
     token = Field("musicToken")
 
     @property
     def is_song(self):
-        return self.token.startswith('S')
+        return isinstance(self, SongSearchResultItem)
 
     @property
     def is_artist(self):
-        return self.token.startswith('R')
+        return isinstance(self, ArtistSearchResultItem) and \
+               self.token.startswith("R")
 
     @property
     def is_composer(self):
-        return self.token.startswith('C')
+        return isinstance(self, ArtistSearchResultItem) and \
+               self.token.startswith("C")
+
+    @property
+    def is_genre_station(self):
+        return isinstance(self, GenreStationSearchResultItem)
 
     def create_station(self):
-        if self.is_song:
-            self._api_client.create_station(track_token=self.token)
+        raise NotImplementedError
+
+    @classmethod
+    def from_json(cls, api_client, data):
+        if data["musicToken"].startswith("S"):
+            return SongSearchResultItem.from_json(api_client, data)
+
+        elif data["musicToken"].startswith(("R", "C")):
+            return ArtistSearchResultItem.from_json(api_client, data)
+
+        elif data["musicToken"].startswith("G"):
+            return GenreStationSearchResultItem.from_json(api_client, data)
         else:
-            self._api_client.create_station(artist_token=self.token)
+            raise NotImplementedError("Unknown result token type '{}'"
+                                      .format(data["musicToken"]))
+
+
+class ArtistSearchResultItem(SearchResultItem):
+
+    score = Field("score")
+    token = Field("musicToken")
+    artist = Field("artistName")
+    likely_match = Field("likelyMatch", default=False)
+
+    def create_station(self):
+        self._api_client.create_station(artist_token=self.token)
+
+    @classmethod
+    def from_json(cls, api_client, data):
+        return super(SearchResultItem, cls).from_json(api_client, data)
+
+
+class SongSearchResultItem(SearchResultItem):
+
+    score = Field("score")
+    token = Field("musicToken")
+    artist = Field("artistName")
+    song_name = Field("songName")
+
+    def create_station(self):
+        self._api_client.create_station(track_token=self.token)
+
+    @classmethod
+    def from_json(cls, api_client, data):
+        return super(SearchResultItem, cls).from_json(api_client, data)
+
+
+class GenreStationSearchResultItem(SearchResultItem):
+
+    score = Field("score")
+    token = Field("musicToken")
+    station_name = Field("stationName")
+
+    def create_station(self):
+        self._api_client.create_station(search_token=self.token)
+
+    @classmethod
+    def from_json(cls, api_client, data):
+        return super(SearchResultItem, cls).from_json(api_client, data)
 
 
 class SearchResult(PandoraModel):
 
     nearest_matches_available = Field("nearMatchesAvailable")
     explanation = Field("explanation")
-    songs = Field("songs", formatter=SearchResultItem.from_json_list)
-    artists = Field("artists", formatter=SearchResultItem.from_json_list)
+    songs = Field("songs", formatter=SongSearchResultItem.from_json_list)
+    artists = Field("artists", formatter=ArtistSearchResultItem.from_json_list)
+    genre_stations = Field(
+        "genreStations",
+        formatter=GenreStationSearchResultItem.from_json_list)
 
 
 class GenreStationList(PandoraDictListModel):