Improve Playlist.__iter__

pafy/playlist.py

@@ -137,8 +137,6 @@ def parseISO8591(duration):
 
 
 class Playlist(object):
-    _items = None
-
     def __init__(self, playlist_url, basic, gdata, size, callback):
         playlist_id = extract_playlist_id(playlist_url)
 
@@ -158,6 +156,7 @@ def __init__(self, playlist_url, basic, gdata, size, callback):
         self._callback = callback
         self._pageToken = None
         self._have_basic = False
+        self._items = []
 
     @classmethod
     def from_dict(cls, pl, basic, gdata, size, callback):
@@ -211,22 +210,16 @@ def __len__(self):
         return self._len
 
     def __iter__(self):
-        if self._items is not None:
-            for i in self._items:
-                yield i
-            if self._pageToken == -1:
-                return
-
-        items = self._items
-        if items is None:
-            items = []
+        for i in self._items:
+            yield i
 
         # playlist items specific metadata
         query = {'part': 'snippet',
                  'maxResults': 50,
                  'playlistId': self.plid}
 
-        while True:
+        # Use -1 to represent having reached the last page
+        while self._pageToken != -1:
             if self._pageToken:
                 query['pageToken'] = self._pageToken
             playlistitems = call_gdata('playlistItems', query)
@@ -237,12 +230,13 @@ def __iter__(self):
                                      for i in playlistitems['items'])}
             wdata = call_gdata('videos', query2)
 
+            index = len(self._items)
             for v in wdata['items']:
                 vid_data = dict_for_playlist(v)
 
                 try:
                     pafy_obj = new(v['id'],
-                                   basic=self._basic, gdata=self._gdata,
+                                   basic=False, gdata=False,
                                    size=self._size, callback=self._callback)
 
                 except IOError as e:
@@ -251,23 +245,25 @@ def __iter__(self):
                     continue
 
                 pafy_obj.populate_from_playlist(vid_data)
-                items.append(pafy_obj)
+                self._items.append(pafy_obj)
                 if self._callback:
                     self._callback("Added video: %s" % vid_data['title'])
-                yield pafy_obj
 
-            if not playlistitems.get('nextPageToken'):
-                self._pageToken = -1
-                self._len = len(items)
-                break
+            self._pageToken = playlistitems.get('nextPageToken', -1)
+            if self._pageToken == -1:
+                self._len = len(self._items)
 
-            query['pageToken'] = playlistitems['nextPageToken']
-            self._pageToken = playlistitems['nextPageToken']
+            # Do not yield until self._items and self._pageToken are set
+            for i in self._items[index:]:
+                if self._basic:
+                    i._fetch_basic()
+                if self._gdata:
+                    i._fetch_gdata()
 
-        self._items = items
+                yield i
 
     def __getitem__(self, index):
-        if self._items is not None and index < len(self._items):
+        if index < len(self._items):
             return self._items[index]
 
         try: