Use wptools.core continuation support in wptools.category #119

siznax · May 8, 2018 · 2852bf6 · 2852bf6
1 parent 82d9454
commit 2852bf6
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 20 deletions.
diff --git a/tests/test_basic.py b/tests/test_basic.py
@@ -105,7 +105,7 @@ def test_category_get_members_continue(self):
         cat = wptools.category('TEST')
         cat.cache['category'] = category_cmcontinue.cache
         cat._set_data('category')
-        self.assertTrue('cmcontinue' in cat.data)
+        self.assertTrue('continue' in cat.data)
         self.assertEqual(len(cat.data['members']), 1)
 
         qry = cat._query('category', wptools.query.WPToolsQuery())

diff --git a/wptools/category.py b/wptools/category.py
@@ -94,25 +94,15 @@ def _query(self, action, qobj):
         if action == 'random':
             return qobj.random(namespace=14)
         elif action == 'category':
-            qry = qobj.category(title=title, pageid=pageid)
-            if self.data.get('cmcontinue'):
-                qry += "&cmcontinue=%s" % self.data['cmcontinue']
-            return qry
+            return qobj.category(title, pageid, self._continue_params())
 
     def _set_data(self, action):
         """
         Set category member data from API response
         """
         data = self._load_response(action)
 
-        try:
-            cmcontinue = data.get('continue').get('cmcontinue')
-            if cmcontinue:
-                self.data['cmcontinue'] = cmcontinue
-                del self.cache['category']
-        except AttributeError:
-            if 'cmcontinue' in self.data:
-                del self.data['cmcontinue']
+        self._handle_continuations(data, 'category')
 
         if action == 'category':
             members = data.get('query').get('categorymembers')
@@ -151,9 +141,8 @@ def get_members(self, show=True, proxy=None, timeout=0):
 
         self._get('category', show, proxy, timeout)
 
-        if self.data.get('cmcontinue'):
-            while self.data.get('cmcontinue'):
-                self._get('category', show, proxy, timeout)
+        while self.data.get('continue'):
+            self._get('category', show, proxy, timeout)
 
         return self
 

diff --git a/wptools/query.py b/wptools/query.py
@@ -121,14 +121,14 @@ def __init__(self, lang='en', variant=None, wiki=None):
         self.domain = domain_name(self.wiki)
         self.uri = self.wiki_uri(self.wiki)
 
-    def category(self, title, pageid=None, limit=500, namespace=None):
+    def category(self, title, pageid=None, cparams=None, namespace=None):
         """
         Returns category query string
         """
         query = self.LIST.substitute(WIKI=self.uri, LIST='categorymembers')
+        status = pageid or title
 
-        if limit:
-            query += "&cmlimit=%d" % limit
+        query += "&cmlimit=500"
 
         if namespace is not None:
             query += "&cmnamespace=%d" % namespace
@@ -142,7 +142,11 @@ def category(self, title, pageid=None, limit=500, namespace=None):
         if pageid:
             query += "&cmpageid=%d" % pageid
 
-        self.set_status('categorymembers', pageid or title)
+        if cparams:
+            query += cparams
+            status += ' (%s)' % cparams
+
+        self.set_status('categorymembers', status)
 
         return query