Merge branch 'continue-categories'

Author: siznax

tests/category_cmcontinue.py

@@ -0,0 +1,24 @@
+# -*- coding:utf-8 -*-
+
+query = 'https://en.wikipedia.org/w/api.php?action=query&formatversion=2&list=categorymembers&cmlimit=500&cmtitle=Category%3ABAFTA%20winners%20%28people%29'
+
+response = r"""{
+    "batchcomplete": true,
+    "continue": {
+        "cmcontinue": "page|412745372f4907044b352f2f41273d03424b352f2f412704412745372f490121018f7f8f7f8f808f09|42525291",
+        "continue": "-||"
+    },
+    "query": {
+        "categorymembers": [
+            {
+                "pageid": 22167530,
+                "ns": 0,
+                "title": "Allison Abbate"
+            }
+        ]
+    }
+}"""
+
+cache = {'query': query, 
+         'response': response, 
+         'info': {'content': 'TEST', 'status': 200}}

tests/stress.py

@@ -11,6 +11,7 @@
 
 import argparse
 import random
+import sys
 import time
 
 import wptools
@@ -77,12 +78,17 @@ def print_header(delay, lang, pages):
     msg.append("delay: %d lang: %s pages: %s" % (delay, langstr, pagestr))
     msgstr = " ".join(msg)
 
+    header = [msgstr]
+    header.append("=" * len(msgstr))
+    header.append("Python " + sys.version)
+    header.append('-' * len(msgstr))
+
     if len(pages) > 1:
         print("Getting top %s.wikipedia.org pages" % lang)
         for i, title in enumerate(pages[:10]):
             print(" %d. %s" % (i + 1, title))
 
-    print("%s\n%s" % (msgstr, "=" * len(msgstr)))
+    print("\n".join(header))
 
 
 def main(args):

tests/test_basic.py

@@ -9,6 +9,7 @@
 import wptools
 
 from . import category
+from . import category_cmcontinue
 from . import disambiguation
 from . import imageinfo
 from . import labels_1
@@ -84,16 +85,33 @@ def test_category_get_members(self):
         cat = wptools.category('TEST')
         cat.cache['category'] = category.cache
         cat._set_data('category')
-        self.assertTrue(len(cat.data['members']), 92)
+        self.assertEqual(len(cat.data['members']), 68)
+        self.assertEqual(len(cat.data['subcategories']), 24)
         self.assertTrue('requests' not in cat.data)
 
     def test_category_get_members_namespace(self):
         cat = wptools.category('TEST', namespace=0)
         cat.cache['category'] = category.cache
         cat._set_data('category')
-        self.assertTrue(len(cat.data['members']), 92)
+        self.assertEqual(len(cat.data['members']), 68)
         self.assertTrue('requests' not in cat.data)
 
+    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.assertEqual(len(cat.data['members']), 1)
+
+        qry = cat._query('category', wptools.query.WPToolsQuery())
+        self.assertTrue('&cmcontinue=page|' in qry)
+        self.assertTrue(qry.endswith('|42525291'))
+
+        cat.cache['category'] = category.cache
+        cat._set_data('category')
+        self.assertTrue('cmcontinue' not in cat.data)
+        self.assertEqual(len(cat.data['members']), 69)
+
     def test_category_query(self):
         cat = wptools.category('TEST')
         qobj = wptools.query.WPToolsQuery()

wptools/category.py

@@ -66,6 +66,24 @@ def __init__(self, *args, **kwargs):
         if not pageid and not title:
             self.get_random()
 
+    def _add_members(self, catmembers):
+        """
+        Adds category members and subcategories to data
+        """
+        members = [x for x in catmembers if x['ns'] == 0]
+        subcats = [x for x in catmembers if x['ns'] == 14]
+
+        if 'members' in self.data:
+            self.data['members'].extend(members)
+        else:
+            self.data.update({'members': members})
+
+        if subcats:
+            if 'subcategories' in self.data:
+                self.data['subcategories'].extend(subcats)
+            else:
+                self.data.update({'subcategories': subcats})
+
     def _query(self, action, qobj):
         """
         Form query to enumerate category
@@ -76,17 +94,30 @@ def _query(self, action, qobj):
         if action == 'random':
             return qobj.random(namespace=14)
         elif action == 'category':
-            return qobj.category(title=title, pageid=pageid)
+            qry = qobj.category(title=title, pageid=pageid)
+            if self.data.get('cmcontinue'):
+                qry += "&cmcontinue=%s" % self.data['cmcontinue']
+            return qry
 
     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']
+
         if action == 'category':
             members = data.get('query').get('categorymembers')
-            self.data.update({'members': members})
+            if members:
+                self._add_members(members)
 
         if action == 'random':
             rand = data['query']['random'][0]
@@ -120,6 +151,10 @@ 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)
+
         return self
 
     def get_random(self, show=True, proxy=None, timeout=0):