Skip to content

Extract page categories #190

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
zaemyung wants to merge 3 commits into
base: master
Choose a base branch
from
Open

Extract page categories #190

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
1ebc770
extract page categories
Oct 3, 2019
511e7f1
move category extraction, usr argparse, exclude sortkey
Oct 3, 2019
1305cf6
add --category_surface
Oct 3, 2019
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
27 changes: 25 additions & 2 deletions WikiExtractor.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -143,6 +143,11 @@ def __eq__ (self, other):
# Drop tables from the article
keep_tables = False,

##
# Extract and save page categories
extract_categories = False,
category_surface = 'Category',

##
# Whether to preserve links in output
keepLinks = False,
Expand Down Expand Up @@ -556,7 +561,7 @@ def __init__(self, id, revid, title, lines):
self.recursion_exceeded_3_errs = 0 # parameter recursion
self.template_title_errs = 0

def write_output(self, out, text):
def write_output(self, out, text, categories=[]):
"""
:param out: a memory file
:param text: the text of the page
Expand All @@ -569,6 +574,9 @@ def write_output(self, out, text):
'title': self.title,
'text': "\n".join(text)
}
if options.extract_categories:
json_data['categories'] = categories

if options.print_revision:
json_data['revid'] = self.revid
# We don't use json.dump(data, out) because we want to be
Expand All @@ -583,6 +591,10 @@ def write_output(self, out, text):
header = '<doc id="%s" revid="%s" url="%s" title="%s">\n' % (self.id, self.revid, url, self.title)
else:
header = '<doc id="%s" url="%s" title="%s">\n' % (self.id, url, self.title)
if options.extract_categories:
categories_line = ', '.join(categories)
header = header[:-2] + ' categories="%s">\n' % (categories_line)

footer = "\n</doc>\n"
if out == sys.stdout: # option -a or -o -
header = header.encode('utf-8')
Expand Down Expand Up @@ -644,6 +656,11 @@ def extract(self, out):
# $text = $frame->expand( $dom );
#
text = self.transform(text)
categories = []
if options.extract_categories:
# Extract categories but not sortkeys
rgx_category = r'\[\[%s:([^|\]]+)(?:|[^\]]+)?\]\]' % options.category_surface
categories = re.findall(rgx_category, text)
text = self.wiki2text(text)
text = compact(self.clean(text))
# from zwChan
Expand All @@ -652,7 +669,7 @@ def extract(self, out):
if sum(len(line) for line in text) < options.min_text_length:
return

self.write_output(out, text)
self.write_output(out, text, categories)

errs = (self.template_title_errs,
self.recursion_exceeded_1_errs,
Expand Down Expand Up @@ -3153,6 +3170,10 @@ def main():
help="comma separated list of elements that will be removed from the article text")
groupP.add_argument("--keep_tables", action="store_true", default=options.keep_tables,
help="Preserve tables in the output article text (default=%(default)s)")
groupP.add_argument("--extract_categories", action="store_true", default=options.extract_categories,
help="Whether to extract page categories, when used, --category_surface is also required (default=%(default)s)")
groupP.add_argument("--category_surface", type=str, default=options.category_surface,
help="How `Category` is written in the wiki namespace. E.g. for English: `Category`, for French: `Catégorie`")
default_process_count = max(1, cpu_count() - 1)
parser.add_argument("--processes", type=int, default=default_process_count,
help="Number of processes to use (default %(default)s)")
Expand Down Expand Up @@ -3187,6 +3208,8 @@ def main():
options.expand_templates = args.no_templates
options.filter_disambig_pages = args.filter_disambig_pages
options.keep_tables = args.keep_tables
options.extract_categories = args.extract_categories
options.category_surface = args.category_surface

try:
power = 'kmg'.find(args.bytes[-1].lower()) + 1
Expand Down