Merge branch 'master' of github.com:makoshark/COVID-19_Digital_Observatory
[covid19.git] / wikipedia / scripts / fetch_enwiki_revisions.py
1 #!/usr/bin/env python3
2
3 ###############################################################################
4 #
5 # This script assumes the presence of the COVID-19 repo.
6
7 # It (1) reads in the article list and then (2) calls the Wikimedia API to 
8 # fetch view information for each article. Output is to (3) JSON and TSV.
9 #
10 ###############################################################################
11
12 import argparse
13 import logging
14 import os.path
15 import json
16 import subprocess
17 import datetime
18
19 from requests import Request
20 from csv import DictWriter
21 from mw import api
22
23
24 def parse_args():
25
26     parser = argparse.ArgumentParser(description='Call the views API to collect Wikipedia revision data.')
27     parser.add_argument('-o', '--output_folder', help='Where to save output', default="wikipedia/data", type=str)
28     parser.add_argument('-i', '--article_file', help='File listing article names', default="wikipedia/resources/enwp_wikiproject_covid19_articles.txt", type=str)
29     parser.add_argument('-L', '--logging_level', help='Logging level. Options are debug, info, warning, error, critical. Default: info.', default='info', type=str), 
30     parser.add_argument('-W', '--logging_destination', help='Logging destination file. (default: standard error)', type=str), 
31     args = parser.parse_args()
32     return(args)
33
34 def main():
35     args = parse_args()
36
37     output_path = args.output_folder
38     article_filename = args.article_file
39
40     #handle -L
41     loglevel_mapping = { 'debug' : logging.DEBUG,
42                          'info' : logging.INFO,
43                          'warning' : logging.WARNING,
44                          'error' : logging.ERROR,
45                          'critical' : logging.CRITICAL }
46
47     if args.logging_level in loglevel_mapping:
48         loglevel = loglevel_mapping[args.logging_level]
49     else:
50         print("Choose a valid log level: debug, info, warning, error, or critical") 
51         exit
52
53     #handle -W
54     if args.logging_destination:
55         logging.basicConfig(filename=args.logging_destination, filemode='a', level=loglevel)
56     else:
57         logging.basicConfig(level=loglevel)
58
59     export_git_hash = subprocess.check_output(['git', 'rev-parse', 'HEAD']).decode().strip()
60     export_git_short_hash = subprocess.check_output(['git', 'rev-parse', '--short', 'HEAD']).decode().strip()
61     export_time = str(datetime.datetime.now())
62     export_date = datetime.datetime.today().strftime("%Y%m%d")
63
64     logging.info(f"Starting run at {export_time}")
65     logging.info(f"Last commit: {export_git_hash}")
66
67     json_output_filename = os.path.join(output_path, f"digobs_covid19-wikipedia-enwiki_revisions-{export_date}.json")
68     tsv_output_filename =  os.path.join(output_path, f"digobs_covid19-wikipedia-enwiki_revisions-{export_date}.tsv")
69     
70     api_session = api.Session("https://en.wikipedia.org/w/api.php")
71
72     # list of properties from the API we want to gather (basically all of
73     # them supported by mediawik-utilities)
74
75     rv_props =  {'revid' : 'ids',
76                  'timestamp' : 'timestamp',
77                  'user' : 'user',
78                  'userid' : 'userid',
79                  'size' : 'size',
80                  'sha1' : 'sha1',
81                  'contentmodel' : 'contentmodel',
82                  'tags' : 'tags',
83                  'comment' : 'comment',
84                  'content' : 'content' }
85
86     exclude_from_tsv = ['tags', 'comment', 'content']
87
88     # load the list of articles
89     with open(article_filename, 'r') as infile:
90         article_list = [art.strip() for art in list(infile)]
91
92     def get_revisions_for_page(title):
93         return api_session.revisions.query(properties=rv_props.values(),
94                                            titles={title},
95                                            direction="newer")
96
97     tsv_fields = ['title', 'pageid', 'namespace']
98     tsv_fields = tsv_fields + list(rv_props.keys())
99
100     # drop fields that we identified for exclusion
101     tsv_fields = [e for e in tsv_fields if e not in exclude_from_tsv]
102
103     # add special export fields
104     tsv_fields = tsv_fields + ['url', 'export_timestamp', 'export_commit']
105
106     export_info = { 'git_commit' : export_git_hash,
107                     'timestamp' : export_time }
108
109     with open(json_output_filename, 'w') as json_output, \
110          open(tsv_output_filename, 'w') as tsv_output:
111
112         tsv_writer = DictWriter(tsv_output, fieldnames=tsv_fields, delimiter="\t")
113         tsv_writer.writeheader()
114
115         for article in article_list:
116             logging.info(f"pulling revisions for: {article}")
117             for rev in get_revisions_for_page(article):
118                 logging.debug(f"processing raw revision: {rev}")
119
120                 # add export metadata
121                 rev['exported'] = export_info
122
123                 # save the json version of the code
124                 print(json.dumps(rev), file=json_output)
125
126                 # handle missing data
127                 if "sha1" not in rev:
128                     rev["sha1"] = ""
129
130                 # add page title information
131                 rev['title'] = rev['page']['title']
132                 rev['pageid'] = rev['page']['pageid']
133                 rev['namespace'] = rev['page']['ns']
134
135                 # construct a URL
136                 rev['url'] = Request('GET', 'https://en.wikipedia.org/w/index.php',
137                                      params={'title' : rev['title'].replace(" ", "_"),
138                                             'oldid' : rev['revid']}).prepare().url
139
140                 rev['export_timestamp'] = export_time
141                 rev['export_commit'] = export_git_short_hash
142
143                 tsv_writer.writerow({k: rev[k] for k in tsv_fields})
144             break
145
146 if __name__ == "__main__":
147
148     main()

Community Data Science Collective || Want to submit a patch?