diff --git a/earwigbot/tasks/wikiproject_tagger.py b/earwigbot/tasks/wikiproject_tagger.py index 88867f1..eb16e05 100644 --- a/earwigbot/tasks/wikiproject_tagger.py +++ b/earwigbot/tasks/wikiproject_tagger.py @@ -22,14 +22,17 @@ import re +from earwigbot import exceptions from earwigbot.tasks import Task +from earwigbot.wiki import constants class WikiProjectTagger(Task): """A task to tag talk pages with WikiProject banners. Usage: :command:`earwigbot -t wikiproject_tagger PATH --banner BANNER (--category CAT | --file FILE) [--summary SUM] - [--append TEXT] [--autoassess] [--nocreate] [--recursive NUM]` + [--append TEXT] [--autoassess] [--nocreate] [--recursive NUM] + [--site SITE]` .. glossary:: @@ -57,6 +60,8 @@ class WikiProjectTagger(Task): ``--recursive NUM`` recursively go through subcategories up to a maximum depth of ``NUM``, or if ``NUM`` isn't provided, go infinitely (this can be dangerous) + ``--site SITE`` + the ID of the site to tag pages on, defaulting to the... default site """ name = "wikiproject_tagger" @@ -91,7 +96,58 @@ class WikiProjectTagger(Task): def run(self, **kwargs): """Main entry point for the bot task.""" + if "file" not in kwargs and "category" not in kwargs: + log = "No pages to tag; I need either a 'category' or a 'file' passed as kwargs" + self.logger.error(log) + return + if "banner" not in kwargs: + log = "Needs a banner to add passed as the 'banner' kwarg" + self.logger.error(log) + return + + banner = kwargs["banner"] + summary = kwargs.get("summary", "Adding $3 to article talk page.") + append = kwargs.get("append") + autoassess = kwargs.get("autoassess", False) + nocreate = kwargs.get("nocreate", False) + recursive = kwargs.get("recursive", 0) + site = self.bot.wiki.get_site(name=kwargs.get("site")) + if "category" in kwargs: - pass - elif "file" in kwargs: - pass + title = kwargs["category"] + prefix = name.split(":", 1)[0] + ns_cat = site.namespace_id_to_name(constants.NS_CATEGORY) + if prefix == title: + title = u":".join((ns_cat, title)) + else: + try: + site.namespace_name_to_id(prefix) + except exceptions.NamespaceNotFoundError: + title = u":".join((ns_cat, title)) + self.process_category(title, recursive) + + if "file" in kwargs: + with open(kwargs["file"], "r") as fileobj: + for line in fileobj: + if line.strip(): + line = line.decode("utf8") + if line.startswith("[[") and line.endswith("]]"): + line = line[2:-2] + page = site.get_page(line) + if page.namespace == constants.NS_CATEGORY: + self.process_category(page, recursive) + else: + self.process_page(page) + + def process_category(self, page, recursive): + for member in page.get_members(): + if member.namespace == constants.NS_CATEGORY: + if recursive is True: + self.process_category(member, True) + elif recursive: + self.process_category(member, recursive - 1) + else: + self.process_page(member) + + def process_page(self, page): + raise NotImplementedError(page)