Merge pull request #466 from pkkid/batchsearch

Make it easier to set a container size using search

Author: Hellowlol

plexapi/base.py

@@ -142,19 +142,21 @@ def fetchItem(self, ekey, cls=None, **kwargs):
         clsname = cls.__name__ if cls else 'None'
         raise NotFound('Unable to find elem: cls=%s, attrs=%s' % (clsname, kwargs))
 
-    def fetchItems(self, ekey, cls=None, **kwargs):
+    def fetchItems(self, ekey, cls=None, container_start=None, container_size=None, **kwargs):
         """ Load the specified key to find and build all items with the specified tag
             and attrs. See :func:`~plexapi.base.PlexObject.fetchItem` for more details
             on how this is used.
 
-            Use container_start and container_size for pagination.
+            Parameters:
+                container_start (None, int): offset to get a subset of the data
+                container_size (None, int): How many items in data
+
         """
         url_kw = {}
-        for key, value in dict(kwargs).items():
-            if key == "container_start":
-                url_kw["X-Plex-Container-Start"] = kwargs.pop(key)
-            if key == "container_size":
-                url_kw["X-Plex-Container-Size"] = kwargs.pop(key)
+        if container_start is not None:
+            url_kw["X-Plex-Container-Start"] = container_start
+        if container_size is not None:
+            url_kw["X-Plex-Container-Size"] = container_size
 
         if ekey is None:
             raise BadRequest('ekey was not provided')

plexapi/library.py

@@ -360,6 +360,40 @@ def _loadData(self, data):
         # Private attrs as we dont want a reload.
         self._total_size = None
 
+    def fetchItems(self, ekey, cls=None, container_start=None, container_size=None, **kwargs):
+        """ Load the specified key to find and build all items with the specified tag
+            and attrs. See :func:`~plexapi.base.PlexObject.fetchItem` for more details
+            on how this is used.
+
+            Parameters:
+                container_start (None, int): offset to get a subset of the data
+                container_size (None, int): How many items in data
+
+        """
+        url_kw = {}
+        if container_start is not None:
+            url_kw["X-Plex-Container-Start"] = container_start
+        if container_size is not None:
+            url_kw["X-Plex-Container-Size"] = container_size
+
+        if ekey is None:
+            raise BadRequest('ekey was not provided')
+        data = self._server.query(ekey, params=url_kw)
+
+        if '/all' in ekey:
+            # totalSize is only included in the xml response
+            # if container size is used.
+            total_size = data.attrib.get("totalSize") or data.attrib.get("size")
+            self._total_size = utils.cast(int, total_size)
+
+        items = self.findItems(data, cls, ekey, **kwargs)
+
+        librarySectionID = data.attrib.get('librarySectionID')
+        if librarySectionID:
+            for item in items:
+                item.librarySectionID = librarySectionID
+        return items
+
     @property
     def totalSize(self):
         if self._total_size is None:
@@ -505,9 +539,9 @@ def listChoices(self, category, libtype=None, **kwargs):
         key = '/library/sections/%s/%s%s' % (self.key, category, utils.joinArgs(args))
         return self.fetchItems(key, cls=FilterChoice)
 
-    def search(self, title=None, sort=None, maxresults=999999, libtype=None, **kwargs):
-        """ Search the library. If there are many results, they will be fetched from the server
-            in batches of X_PLEX_CONTAINER_SIZE amounts. If you're only looking for the first <num>
+    def search(self, title=None, sort=None, maxresults=None,
+               libtype=None, container_start=0, container_size=X_PLEX_CONTAINER_SIZE, **kwargs):
+        """ Search the library. The http requests will be batched in container_size. If you're only looking for the first <num>
             results, it would be wise to set the maxresults option to that amount so this functions
             doesn't iterate over all results on the server.
 
@@ -518,6 +552,8 @@ def search(self, title=None, sort=None, maxresults=999999, libtype=None, **kwarg
                 maxresults (int): Only return the specified number of results (optional).
                 libtype (str): Filter results to a spcifiec libtype (movie, show, episode, artist,
                     album, track; optional).
+                container_start (int): default 0
+                container_size (int): default X_PLEX_CONTAINER_SIZE in your config file.
                 **kwargs (dict): Any of the available filters for the current library section. Partial string
                         matches allowed. Multiple matches OR together. Negative filtering also possible, just add an
                         exclamation mark to the end of filter name, e.g. `resolution!=1x1`.
@@ -549,15 +585,37 @@ def search(self, title=None, sort=None, maxresults=999999, libtype=None, **kwarg
             args['sort'] = self._cleanSearchSort(sort)
         if libtype is not None:
             args['type'] = utils.searchType(libtype)
-        # iterate over the results
-        results, subresults = [], '_init'
-        args['X-Plex-Container-Start'] = 0
-        args['X-Plex-Container-Size'] = min(X_PLEX_CONTAINER_SIZE, maxresults)
-        while subresults and maxresults > len(results):
+
+        results = []
+        subresults = []
+        offset = container_start
+
+        if maxresults is not None:
+            container_size = min(container_size, maxresults)
+        while True:
             key = '/library/sections/%s/all%s' % (self.key, utils.joinArgs(args))
-            subresults = self.fetchItems(key)
-            results += subresults[:maxresults - len(results)]
-            args['X-Plex-Container-Start'] += args['X-Plex-Container-Size']
+            subresults = self.fetchItems(key, container_start=container_start,
+                                         container_size=container_size)
+            if not len(subresults):
+                if offset > self.totalSize:
+                    log.info("container_start is higher then the number of items in the library")
+                break
+
+            results.extend(subresults)
+
+            # self.totalSize is not used as a condition in the while loop as
+            # this require a additional http request.
+            # self.totalSize is updated from .fetchItems
+            wanted_number_of_items = self.totalSize - offset
+            if maxresults is not None:
+                wanted_number_of_items = min(maxresults, wanted_number_of_items)
+                container_size = min(container_size, maxresults - len(results))
+
+            if wanted_number_of_items <= len(results):
+                break
+
+            container_start += container_size
+
         return results
 
     def _cleanSearchFilter(self, category, value, libtype=None):

tests/test_library.py

@@ -260,3 +260,8 @@ def test_crazy_search(plex, movie):
     ), "Unable to search movie by year."
     assert movie not in movies.search(year=2007), "Unable to filter movie by year."
     assert movie in movies.search(actor=movie.actors[0].tag)
+    assert len(movies.search(container_start=2, maxresults=1)) == 1
+    assert len(movies.search(container_size=None)) == 4
+    assert len(movies.search(container_size=1)) == 4
+    assert len(movies.search(container_start=9999, container_size=1)) == 0
+    assert len(movies.search(container_start=2, container_size=1)) == 2