forked from volpino/wiki-network
/
countries_stats.py
executable file
·178 lines (149 loc) · 5.7 KB
/
countries_stats.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
#!/usr/bin/env python
##########################################################################
# #
# This program is free software; you can redistribute it and/or modify #
# it under the terms of the GNU General Public License as published by #
# the Free Software Foundation; version 2 of the License. #
# #
# This program is distributed in the hope that it will be useful, #
# but WITHOUT ANY WARRANTY; without even the implied warranty of #
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the #
# GNU General Public License for more details. #
# #
##########################################################################
from sonet.mediawiki import HistoryPageProcessor, \
get_translations, get_tags, \
explode_dump_filename, \
ts2dt
from sonet import lib
import csv
import sys
import logging
from sonet.timr import Timr
import pygeoip
from collections import Counter, OrderedDict
import datetime
from dateutil.rrule import rrule, MONTHLY
class CountriesPageProcessor(HistoryPageProcessor):
output = None
data = None
gi = None
countries = set()
csv_writer = None
_skip = None
_country = None
def __init__(self, **kwargs):
super(CountriesPageProcessor, self).__init__(**kwargs)
self.gi = pygeoip.GeoIP(kwargs["geoip"])
self.data = OrderedDict()
def flush(self):
"""
Flushes queue in the CSV output
"""
f = open(self.output, "w")
self.csv_writer = csv.DictWriter(f, ["date"] + list(self.countries))
self.csv_writer.writeheader()
for date in self.data:
to_write = Counter(date=date)
to_write.update(dict([(x, 0) for x in self.countries]))
to_write.update(self.data[date])
self.csv_writer.writerow(to_write)
f.close()
def process_timestamp(self, elem):
if self._skip:
return
self._date = ts2dt(elem.text)
def process_ip(self, elem):
if self._skip:
return
try:
self._country = self.gi.country_name_by_addr(elem.text)
except (pygeoip.GeoIPError, IndexError, KeyError):
logging.warn("Skipping IP %s", elem.text)
else:
if self._country:
self.countries.add(self._country)
else:
logging.warn("Can't find country for IP %s", elem.text)
self._country = "Unknown"
self.countries.add(self._country)
def process_revision(self, _):
if not self._country:
return
current_date = self._date.strftime("%Y/%m")
first_date = None # 2001 date mismatch
mismatch = False
for date in self.data:
first_date = date
break
if first_date and first_date > current_date:
mismatch = True
logging.warn("Date mismatch! Fixing... - %s %s", first_date,
current_date)
if not self.data or mismatch: # populate dict with all dates
start = self._date.date()
end = datetime.date.today()
for dt in rrule(MONTHLY, dtstart=start, until=end):
dt = dt.strftime("%Y/%m")
if not dt in self.data:
self.data[dt] = Counter()
self.data[current_date][self._country] += 1
self._date = None
self._country = None
def process_page(self, _):
pass
def process_redirect(self, _):
pass
def process_title(self, _):
pass
def dumps_checker(dump_name):
"""
Checks if wikimedia dump is *-meta-history in order to extract
revisions
>>> dumps_checker("wikimedia-pages-current")
Traceback (most recent call last):
AssertionError: Wrong dump file, required: *-meta-history
>>> dumps_checker("wikimedia-pages-meta-history")
"""
import re
assert re.search('.-(meta-history)', dump_name), \
"Wrong dump file, required: *-meta-history"
def main():
import optparse
p = optparse.OptionParser(
usage="usage: %prog [options] input_file geoip_db output_file")
p.add_option('-v', action="store_true", dest="verbose", default=False,
help="Verbose output (like timings)")
opts, files = p.parse_args()
if len(files) != 3:
p.error("Wrong parameters")
if opts.verbose:
logging.basicConfig(stream=sys.stderr,
level=logging.DEBUG,
format='%(asctime)s %(levelname)s %(message)s',
datefmt='%Y-%m-%d %H:%M:%S')
xml = files[0]
geoip_db = files[1]
output = files[2]
dumps_checker(xml)
lang, _, _ = explode_dump_filename(xml)
deflate, _lineno = lib.find_open_for_this_file(xml)
if _lineno:
src = deflate(xml, 51)
else:
src = deflate(xml)
translation = get_translations(src)
tag = get_tags(src,
tags='page,redirect,timestamp,ip,revision,title')
src.close()
src = deflate(xml)
processor = CountriesPageProcessor(tag=tag, lang=lang,
output=output,
userns=translation['User'],
geoip=geoip_db
)
with Timr('Processing'):
processor.start(src) ## PROCESSING
processor.flush()
if __name__ == "__main__":
main()