plum erg

scripts/categories.py

@@ -0,0 +1,91 @@
+from mwclient import Site, APIError
+from mwclient.page import Page
+import argparse, json, sys
+
+import unidecode
+# unaccented_string = unidecode.unidecode(accented_string)
+
+ap = argparse.ArgumentParser("")
+ap.add_argument("--wikiprotocol", default="http")
+ap.add_argument("--wikihost", default="localhost")
+ap.add_argument("--wikipath", default="/mw/")
+ap.add_argument("--limit", default=None)
+ap.add_argument("--output", default=sys.stdout, type=argparse.FileType('w'))
+# ap.add_argument("--user", default=None)
+# ap.add_argument("--password", default=None)
+args = ap.parse_args()
+
+site = Site((args.wikiprotocol, args.wikihost), path=args.wikipath)
+ 
+def category_subcats (site, cattitle, objects=True):
+    cmcontinue = None
+    ret = []
+    while True:
+        if cmcontinue == None:
+            resp = site.api("query", list="categorymembers", cmtitle=cattitle, cmtype="subcat", cmlimit=50)
+        else:
+            resp = site.api("query", list="categorymembers", cmtitle=cattitle, cmtype="subcat", cmlimit=50, cmcontinue=cmcontinue)
+        ret.extend([x['title'] for x in resp['query']['categorymembers']])
+        if 'continue' in resp:
+            cmcontinue = resp['continue']['cmcontinue']
+        else:
+            break
+    if objects:
+        # print "converting to page objects ({0})".format(len(ret))
+        ret = [site.pages[x] for x in ret]
+    return ret
+
+def strip_namespace (x):
+    if ":" in x:
+       return x.split(":", 1)[1]
+    return x
+
+cats = list(site.allcategories())
+cats.sort(key=lambda x: unidecode.unidecode(x.name))
+# cats_by_name = {}
+cats = [{ 'title': cat.page_title, 'name': cat.name } for cat in cats]
+index = {}
+for c in cats:
+    index[c['title']] = c
+
+for c in cats:
+    sc = category_subcats(site, c['name'], objects=False)
+    sc = [strip_namespace(x) for x in sc]
+    if sc:
+        for subcat in sc:
+            subcat = index[subcat]
+            if 'parent' not in subcat:
+                subcat['parent'] = c
+
+def make_hierarchy (items):
+    root = {'children': []}
+    for item in items:
+        if 'parent' in item:
+            if 'children' not in item['parent']:
+                item['parent']['children'] = []
+            item['parent']['children'].append(item)
+            item['parent'] = item['parent']['title']
+            # del item['parent']
+        else:
+            root['children'].append(item)
+    return root
+
+def flatten (root, depth=0):
+    for x in root['children']:
+        r = {'title': x['title'], 'name': x['name'], 'depth': depth}
+        if 'parent' in x:
+            r['parent'] = x['parent']
+        yield(r)
+        if 'children' in x:
+            for item in flatten(x, depth+1):
+                yield item
+
+import json
+root = make_hierarchy(cats)
+#print (json.dumps(root, indent=2), file=args.output)
+
+flat = list(flatten(root))
+print (json.dumps(flat, indent=2), file=args.output)
+
+
+