diff options
Diffstat (limited to 'backend/hist_data/enwiki/gen_img_data.py')
| -rwxr-xr-x | backend/hist_data/enwiki/gen_img_data.py | 51 |
1 files changed, 14 insertions, 37 deletions
diff --git a/backend/hist_data/enwiki/gen_img_data.py b/backend/hist_data/enwiki/gen_img_data.py index b4ade9f..9aa3863 100755 --- a/backend/hist_data/enwiki/gen_img_data.py +++ b/backend/hist_data/enwiki/gen_img_data.py @@ -8,17 +8,15 @@ The program can be re-run with an updated set of page IDs, and will skip already-processed page IDs. """ -import re -import os, bz2, html, urllib.parse +import os, re +import bz2, html, urllib.parse import sqlite3 DUMP_FILE = 'enwiki-20220501-pages-articles-multistream.xml.bz2' INDEX_DB = 'dump_index.db' -PAGEVIEW_DB = 'pageview_data.db' IMG_DB = 'img_data.db' # The database to create DB_FILE = os.path.join('..', 'data.db') -MAX_IMGS_PER_CTG = 20000 -# +# Regexps ID_LINE_REGEX = re.compile(r'<id>(.*)</id>') IMG_LINE_REGEX = re.compile(r'.*\| *image *= *([^|]*)') BRACKET_IMG_REGEX = re.compile(r'\[\[(File:[^|]*).*]]') @@ -35,7 +33,7 @@ def genData(pageIds: set[int], dumpFile: str, indexDb: str, imgDb: str) -> None: if imgDbCur.execute('SELECT name FROM sqlite_master WHERE type="table" AND name="page_imgs"').fetchone() is None: # Create tables if not present imgDbCur.execute('CREATE TABLE page_imgs (page_id INT PRIMARY KEY, title TEXT UNIQUE, img_name TEXT)') - # 'img_name' may be NULL + # 'img_name' values are set to NULL to indicate page IDs where no image was found imgDbCur.execute('CREATE INDEX page_imgs_idx ON page_imgs(img_name)') else: # Check for already-processed page IDs @@ -179,49 +177,28 @@ def getImageName(content: list[str]) -> str | None: return None return None -def getInputPageIdsFromDb(dbFile: str, pageviewDb: str, indexDb: str, maxImgsPerCtg: int) -> set[int]: +def getInputPageIdsFromDb(dbFile: str, indexDb: str) -> set[int]: print('Getting event data') - titleToCtg: dict[str, str] = {} + titles: set[str] = set() dbCon = sqlite3.connect(dbFile) - for title, ctg in dbCon.execute('SELECT title, ctg from events'): - titleToCtg[title] = ctg + for (title,) in dbCon.execute('SELECT title from events'): + titles.add(title) dbCon.close() - print('Getting top images for each event category') - ctgToTitles: dict[str, list[str]] = {} - dbCon = sqlite3.connect(pageviewDb) - for (title,) in dbCon.execute('SELECT title FROM views ORDER BY views DESC'): - if title not in titleToCtg: - continue - ctg = titleToCtg[title] - if ctg not in ctgToTitles: - ctgToTitles[ctg] = [] - elif len(ctgToTitles[ctg]) == maxImgsPerCtg: - continue - ctgToTitles[ctg].append(title) - del titleToCtg[title] - dbCon.close() - for title, ctg in titleToCtg.items(): # Account for titles without view counts - if ctg not in ctgToTitles: - ctgToTitles[ctg] = [] - elif len(ctgToTitles[ctg]) == maxImgsPerCtg: - continue - ctgToTitles[ctg].append(title) print('Getting page IDs') pageIds: set[int] = set() dbCon = sqlite3.connect(indexDb) dbCur = dbCon.cursor() - for ctg in ctgToTitles: - for title in ctgToTitles[ctg]: - row = dbCur.execute('SELECT id FROM offsets WHERE title = ?', (title,)).fetchone() - if row: - pageIds.add(row[0]) + for title in titles: + row = dbCur.execute('SELECT id FROM offsets WHERE title = ?', (title,)).fetchone() + if row: + pageIds.add(row[0]) dbCon.close() - print(f'Result: {len(pageIds)} out of {len(titleToCtg)}') + print(f'Result: {len(pageIds)} out of {len(titles)}') return pageIds if __name__ == '__main__': import argparse parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) parser.parse_args() # - pageIds = getInputPageIdsFromDb(DB_FILE, PAGEVIEW_DB, INDEX_DB, MAX_IMGS_PER_CTG) + pageIds = getInputPageIdsFromDb(DB_FILE, INDEX_DB) genData(pageIds, DUMP_FILE, INDEX_DB, IMG_DB) |