]> projects.mako.cc - wikiq/blob - wikiq.cpp
dfa7251d5a9e0ea943be334480cc8539fb19094a
[wikiq] / wikiq.cpp
1 /* 
2  * An XML parser for Wikipedia Data dumps.
3  * Converts XML files to tab-separated values files readable by spreadsheets
4  * and statistical packages.
5  */
6
7 #include <stdio.h>
8 #include <iostream>
9 #include <string.h>
10 #include <ctype.h>
11 #include <stdlib.h>
12 #include "expat.h"
13 #include <getopt.h>
14 #include "disorder.h"
15 #include "md5.h"
16 #include "dtl/dtl.hpp"
17 #include <vector>
18 #include <map>
19 #include <pcrecpp.h>
20
21
22 using namespace std;
23
24 // timestamp of the form 2003-11-07T00:43:23Z
25 #define DATE_LENGTH 10
26 #define TIME_LENGTH 8
27 #define TIMESTAMP_LENGTH 20
28
29 #define MEGABYTE 1048576
30 #define FIELD_BUFFER_SIZE 1024
31
32 // this can be changed at runtime if we encounter an article larger than 10mb
33 size_t text_buffer_size = 10 * MEGABYTE;
34
35 enum elements { 
36     TITLE, ARTICLEID, REVISION, REVID, TIMESTAMP, CONTRIBUTOR, 
37     EDITOR, EDITORID, MINOR, COMMENT, UNUSED, TEXT
38 }; 
39
40 enum block { TITLE_BLOCK, REVISION_BLOCK, CONTRIBUTOR_BLOCK, SKIP };
41
42 enum outtype { FULL, SIMPLE };
43
44 typedef struct {
45
46     // pointers to once-allocated buffers
47     char *title;
48     char *articleid;
49     char *revid;
50     char *date;
51     char *time;
52     char *timestamp;
53     char *anon;
54     char *editor;
55     char *editorid;
56     char *comment;
57     char *text;
58     vector<string> last_text_tokens;
59     vector<pcrecpp::RE> title_regexes;
60     vector<string> diff_regex_names;
61     vector<pcrecpp::RE> diff_regexes;
62     map<string, string> revision_md5; // used for detecting reversions
63
64     // track string size of the elements, to prevent O(N^2) processing in charhndl
65     // when we have to take strlen for every character which we append to the buffer
66     size_t title_size;
67     size_t articleid_size;
68     size_t revid_size;
69     size_t date_size;
70     size_t time_size;
71     size_t timestamp_size;
72     size_t anon_size;
73     size_t editor_size;
74     size_t editorid_size;
75     size_t comment_size;
76     size_t text_size;
77
78     bool minor;
79     
80     enum elements element;
81     enum block position;
82     enum outtype output_type;
83     
84 } revisionData;
85
86
87 /* free_data and clean_data
88  * Takes a pointer to the data struct and an integer {0,1} indicating if the 
89  * title data needs to be cleared as well.
90  * Also, frees memory dynamically allocated to store data.
91  */ 
92 static void
93 clean_data(revisionData *data, int title)
94 {
95     // reset title (if we are switching articles)
96     if (title) {
97         data->title[0] = '\0';
98         data->articleid[0] = '\0';
99         data->title_size = 0;
100         data->articleid_size = 0;
101     }
102
103     // reset text fields
104     data->revid[0] = '\0';
105     data->date[0] = '\0';
106     data->time[0] = '\0';
107     data->timestamp[0] = '\0';
108     data->anon[0] = '\0';
109     data->editor[0] = '\0';
110     data->editorid[0] = '\0';
111     data->comment[0] = '\0';
112     data->text[0] = '\0';
113
114     // reset length tracking
115     data->revid_size = 0;
116     data->date_size = 0;
117     data->time_size = 0;
118     data->timestamp_size = 0;
119     data->anon_size = 0;
120     data->editor_size = 0;
121     data->editorid_size = 0;
122     data->comment_size = 0;
123     data->text_size = 0;
124
125     // reset flags and element type info
126     data->minor = false;
127     data->element = UNUSED;
128
129 }
130
131 // presently unused
132 static void
133 free_data(revisionData *data, int title)
134 {
135     if (title) {
136         //printf("freeing article\n");
137         free(data->title);
138         free(data->articleid);
139     }
140     free(data->revid);
141     free(data->date);
142     free(data->time);
143     free(data->timestamp);
144     free(data->anon);
145     free(data->editor);
146     free(data->editorid);
147     free(data->comment);
148     free(data->text);
149     data->last_text_tokens.clear();
150 }
151
152 void cleanup_revision(revisionData *data) {
153     clean_data(data, 0);
154 }
155
156 void cleanup_article(revisionData *data) {
157     clean_data(data, 1);
158     data->last_text_tokens.clear();
159     data->revision_md5.clear();
160 }
161
162
163 static void 
164 init_data(revisionData *data, outtype output_type)
165 {
166     data->text = (char*) malloc(text_buffer_size);
167     data->comment = (char*) malloc(FIELD_BUFFER_SIZE);
168     data->title = (char*) malloc(FIELD_BUFFER_SIZE);
169     data->articleid = (char*) malloc(FIELD_BUFFER_SIZE);
170     data->revid = (char*) malloc(FIELD_BUFFER_SIZE);
171     data->date = (char*) malloc(FIELD_BUFFER_SIZE);
172     data->time = (char*) malloc(FIELD_BUFFER_SIZE);
173     data->timestamp = (char*) malloc(FIELD_BUFFER_SIZE);
174     data->anon = (char*) malloc(FIELD_BUFFER_SIZE);
175     data->editor = (char*) malloc(FIELD_BUFFER_SIZE);
176     data->editorid = (char*) malloc(FIELD_BUFFER_SIZE);
177     data->minor = false;
178
179     // resets the data fields, null terminates strings, sets lengths
180     clean_data(data, 1);
181
182     data->output_type = output_type;
183 }
184
185 /* for debugging only, prints out the state of the data struct
186  */
187 static void
188 print_state(revisionData *data) 
189 {
190     printf("element = %i\n", data->element);
191     printf("output_type = %i\n", data->output_type);
192     printf("title = %s\n", data->title);
193     printf("articleid = %s\n", data->articleid);
194     printf("revid = %s\n", data->revid);
195     printf("date = %s\n", data->date);
196     printf("time = %s\n", data->time);
197     printf("anon = %s\n", data->anon);
198     printf("editor = %s\n", data->editor);
199     printf("editorid = %s\n", data->editorid);
200     printf("minor = %s\n", (data->minor ? "1" : "0"));
201     printf("comment = %s\n", data->comment); 
202     printf("text = %s\n", data->text);
203     printf("\n");
204
205 }
206
207
208 /* 
209  * write a line of comma-separated value formatted data to standard out
210  * follows the form:
211  * title,articleid,revid,date,time,anon,editor,editorid,minor,comment
212  * (str)  (int)    (int) (str)(str)(bin)(str)   (int)   (bin) (str)
213  *
214  * it is called right before cleanup_revision() and cleanup_article()
215  */
216 static void
217 write_row(revisionData *data)
218 {
219
220     // get md5sum
221     md5_state_t state;
222     md5_byte_t digest[16];
223     char md5_hex_output[2 * 16 + 1];
224     md5_init(&state);
225     md5_append(&state, (const md5_byte_t *)data->text, data->text_size);
226     md5_finish(&state, digest);
227     int di;
228     for (di = 0; di < 16; ++di) {
229         sprintf(md5_hex_output + di * 2, "%02x", digest[di]);
230     }
231
232     string reverted_to;
233     map<string, string>::iterator prev_revision = data->revision_md5.find(md5_hex_output);
234     if (prev_revision != data->revision_md5.end()) {
235         reverted_to = prev_revision->second; // id of previous revision
236     }
237     data->revision_md5[md5_hex_output] = data->revid;
238
239     string text = string(data->text, data->text_size);
240     vector<string> text_tokens;
241     size_t pos = 0;
242     size_t start = 0;
243     while ((pos = text.find_first_of(" \n\t\r", pos)) != string::npos) {
244         //cout << "\"\"\"" << text.substr(start, pos - start) << "\"\"\"" << endl;
245         text_tokens.push_back(text.substr(start, pos - start));
246         start = pos;
247         ++pos;
248     }
249
250     // look to see if (a) we've passed in a list of /any/ title_regexes
251     // and (b) if all of the title_regex_matches match
252     // if (a) is true and (b) is not, we return
253     bool any_title_regex_match = false;
254     if (!data->title_regexes.empty()) {
255         for (vector<pcrecpp::RE>::iterator r = data->title_regexes.begin(); r != data->title_regexes.end(); ++r) {
256             pcrecpp::RE& title_regex = *r;
257             if (title_regex.PartialMatch(data->title)) {
258                 any_title_regex_match = true;
259                 break;
260             }
261         }
262         if (!any_title_regex_match) {
263             return;
264         }
265     }
266
267     //vector<string> additions;
268     //vector<string> deletions;
269     string additions;
270     string deletions;
271
272     vector<bool> diff_regex_matches_adds;
273     vector<bool> diff_regex_matches_dels;
274
275     if (data->last_text_tokens.empty()) {
276         additions = data->text;
277     } else {
278         // do the diff
279         
280         dtl::Diff< string, vector<string> > d(data->last_text_tokens, text_tokens);
281         //d.onOnlyEditDistance();
282         d.compose();
283
284         vector<pair<string, dtl::elemInfo> > ses_v = d.getSes().getSequence();
285         for (vector<pair<string, dtl::elemInfo> >::iterator sit=ses_v.begin(); sit!=ses_v.end(); ++sit) {
286             switch (sit->second.type) {
287             case dtl::SES_ADD:
288                 //cout << "ADD: \"" << sit->first << "\"" << endl;
289                 additions += sit->first;
290                 break;
291             case dtl::SES_DELETE:
292                 //cout << "DEL: \"" << sit->first << "\"" << endl;
293                 deletions += sit->first;
294                 break;
295             }
296         }
297     }
298     
299     if (!additions.empty()) {
300         //cout << "ADD: " << additions << endl;
301         for (vector<pcrecpp::RE>::iterator r = data->diff_regexes.begin(); r != data->diff_regexes.end(); ++r) {
302             pcrecpp::RE& diff_regex = *r;
303             diff_regex_matches_adds.push_back(diff_regex.PartialMatch(additions));
304         }
305     }
306
307     if (!deletions.empty()) {
308         //cout << "DEL: " << deletions << endl;
309         for (vector<pcrecpp::RE>::iterator r = data->diff_regexes.begin(); r != data->diff_regexes.end(); ++r) {
310             pcrecpp::RE& diff_regex = *r;
311             diff_regex_matches_dels.push_back(diff_regex.PartialMatch(deletions));
312         }
313     }
314
315     data->last_text_tokens = text_tokens;
316
317
318     // print line of tsv output
319     cout
320         << data->title << "\t"
321         << data->articleid << "\t"
322         << data->revid << "\t"
323         << data->date << " "
324         << data->time << "\t"
325         << ((data->editor[0] != '\0') ? "FALSE" : "TRUE") << "\t"
326         << data->editor << "\t"
327         << data->editorid << "\t"
328         << ((data->minor) ? "TRUE" : "FALSE") << "\t"
329         << (unsigned int) data->text_size << "\t"
330         << shannon_H(data->text, data->text_size) << "\t"
331         << md5_hex_output << "\t"
332         << reverted_to << "\t"
333         << (int) additions.size() << "\t"
334         << (int) deletions.size();
335
336     for (int n = 0; n < data->diff_regex_names.size(); ++n) {
337         cout << "\t" << ((!diff_regex_matches_adds.empty() && diff_regex_matches_adds.at(n)) ? "TRUE" : "FALSE")
338              << "\t" << ((!diff_regex_matches_dels.empty() && diff_regex_matches_dels.at(n)) ? "TRUE" : "FALSE");
339     }
340     cout << endl;
341
342     // 
343     if (data->output_type == FULL) {
344         cout << "comment:" << data->comment << endl
345              << "text:" << endl << data->text << endl;
346     }
347
348 }
349
350 void
351 split_timestamp(revisionData *data) 
352 {
353     char *t = data->timestamp;
354     strncpy(data->date, data->timestamp, DATE_LENGTH);
355     char *timeinstamp = &data->timestamp[DATE_LENGTH+1];
356     strncpy(data->time, timeinstamp, TIME_LENGTH);
357 }
358
359 // like strncat but with previously known length
360 char*
361 strlcatn(char *dest, const char *src, size_t dest_len, size_t n)
362 {
363    size_t i;
364
365    for (i = 0 ; i < n && src[i] != '\0' ; i++)
366        dest[dest_len + i] = src[i];
367    dest[dest_len + i] = '\0';
368
369    return dest;
370 }
371
372 static void
373 charhndl(void* vdata, const XML_Char* s, int len)
374
375     revisionData* data = (revisionData*) vdata;
376     size_t bufsz;
377     if (data->element != UNUSED && data->position != SKIP) {
378         switch (data->element) {
379             case TEXT:
380                     // check if we'd overflow our buffer
381                     bufsz = data->text_size + len;
382                     if (bufsz + 1 > text_buffer_size) {
383                         data->text = (char*) realloc(data->text, bufsz + 1);
384                         text_buffer_size = bufsz + 1;
385                     }
386                     strlcatn(data->text, s, data->text_size, len);
387                     data->text_size = bufsz;
388                     break;
389             case COMMENT:
390                     strlcatn(data->comment, s, data->comment_size, len);
391                     data->comment_size += len;
392                     break;
393             case TITLE:
394                     strlcatn(data->title, s, data->title_size, len);
395                     data->title_size += len;
396                     break;
397             case ARTICLEID:
398                    // printf("articleid = %s\n", t);
399                     strlcatn(data->articleid, s, data->articleid_size, len);
400                     data->articleid_size += len;
401                     break;
402             case REVID:
403                    // printf("revid = %s\n", t);
404                     strlcatn(data->revid, s, data->revid_size, len);
405                     data->revid_size += len;
406                     break;
407             case TIMESTAMP: 
408                     strlcatn(data->timestamp, s, data->timestamp_size, len);
409                     data->timestamp_size += len;
410                     if (strlen(data->timestamp) == TIMESTAMP_LENGTH)
411                         split_timestamp(data);
412                     break;
413             case EDITOR:
414                     strlcatn(data->editor, s, data->editor_size, len);
415                     data->editor_size += len;
416                     break;
417             case EDITORID: 
418                     //printf("editorid = %s\n", t);
419                     strlcatn(data->editorid, s, data->editorid_size, len);
420                     data->editorid_size += len;
421                     break;
422             /* the following are implied or skipped:
423             case MINOR: 
424                     printf("found minor element\n");  doesn't work
425                     break;                   minor tag is just a tag
426             case UNUSED: 
427             */
428             default: break;
429         }
430     }
431 }
432
433 static void
434 start(void* vdata, const XML_Char* name, const XML_Char** attr)
435 {
436     revisionData* data = (revisionData*) vdata;
437     
438     if (strcmp(name,"title") == 0) {
439         cleanup_article(data); // cleans up data from last article
440         data->element = TITLE;
441         data->position = TITLE_BLOCK;
442     } else if (data->position != SKIP) {
443         if (strcmp(name,"revision") == 0) {
444             data->element = REVISION;
445             data->position = REVISION_BLOCK;
446         } else if (strcmp(name, "contributor") == 0) {
447             data->element = CONTRIBUTOR;
448             data->position = CONTRIBUTOR_BLOCK;
449         } else if (strcmp(name,"id") == 0)
450             switch (data->position) {
451                 case TITLE_BLOCK:
452                     data->element = ARTICLEID;
453                     break;
454                 case REVISION_BLOCK: 
455                     data->element = REVID;
456                     break;
457                 case CONTRIBUTOR_BLOCK:
458                     data->element = EDITORID;
459                     break;
460             }
461     
462         // minor tag has no character data, so we parse here
463         else if (strcmp(name,"minor") == 0) {
464             data->element = MINOR;
465             data->minor = true; 
466         }
467         else if (strcmp(name,"timestamp") == 0)
468             data->element = TIMESTAMP;
469
470         else if (strcmp(name, "username") == 0)
471             data->element = EDITOR;
472
473         else if (strcmp(name,"ip") == 0) 
474             data->element = EDITORID;
475
476         else if (strcmp(name,"comment") == 0)
477             data->element = COMMENT;
478
479         else if (strcmp(name,"text") == 0)
480             data->element = TEXT;
481
482         else if (strcmp(name,"page") == 0 
483                 || strcmp(name,"mediawiki") == 0
484                 || strcmp(name,"restrictions") == 0
485                 || strcmp(name,"siteinfo") == 0)
486             data->element = UNUSED;
487     }
488
489 }
490
491
492 static void
493 end(void* vdata, const XML_Char* name)
494 {
495     revisionData* data = (revisionData*) vdata;
496     if (strcmp(name, "revision") == 0 && data->position != SKIP) {
497         write_row(data); // crucial... :)
498         cleanup_revision(data);  // also crucial
499     } else {
500         data->element = UNUSED; // sets our state to "not-in-useful"
501     }                           // thus avoiding unpleasant character data 
502                                 // b/w tags (newlines etc.)
503 }
504
505 void print_usage(char* argv[]) {
506     cerr << "usage: <wikimedia dump xml> | " << argv[0] << "[options]" << endl
507          << endl
508          << "options:" << endl
509          << "  -v   verbose mode prints text and comments after each line of tab separated data" << endl
510          << "  -n   name of the following regex (e.g. -n name -r \"...\")" << endl
511          << "  -r   regex to check against additions and deletions" << endl
512          << "  -t   parse revisions only from pages whose titles match regex(es)" << endl
513          << endl
514          << "Takes a wikimedia data dump XML stream on standard in, and produces" << endl
515          << "a tab-separated stream of revisions on standard out:" << endl
516          << endl
517          << "title, articleid, revid, timestamp, anon, editor, editorid, minor," << endl
518          << "text_length, text_entropy, text_md5, reversion, additions_size, deletions_size" << endl
519          << ".... and additional fields for each regex executed against add/delete diffs" << endl
520          << endl
521          << "Boolean fields are TRUE/FALSE except in the case of reversion, which is blank" << endl
522          << "unless the article is a revert to a previous revision, in which case, it" << endl
523          << "contains the revision ID of the revision which was reverted to." << endl
524          << endl
525          << "author: Erik Garrison <erik@hypervolu.me>" << endl;
526 }
527
528
529 int
530 main(int argc, char *argv[])
531 {
532     
533     enum outtype output_type;
534     int dry_run = 0;
535     // in "simple" output, we don't print text and comments
536     output_type = SIMPLE;
537     char c;
538     string diff_regex_name;
539
540     // the user data struct which is passed to callback functions
541     revisionData data;
542
543     while ((c = getopt(argc, argv, "hvn:r:t:")) != -1)
544         switch (c)
545         {
546             case 'd':
547                 dry_run = 1;
548                 break;
549             case 'v':
550                 output_type = FULL;
551                 break;
552             case 'n':
553                 diff_regex_name = optarg;
554                 break;
555             case 'r':
556                 data.diff_regexes.push_back(pcrecpp::RE(optarg, pcrecpp::UTF8()));
557                 data.diff_regex_names.push_back(diff_regex_name);
558                 if (!diff_regex_name.empty()) {
559                     diff_regex_name.clear();
560                 }
561                 break;
562             case 'h':
563                 print_usage(argv);
564                 exit(0);
565                 break;
566             case 't':
567                 data.title_regexes.push_back(pcrecpp::RE(optarg, pcrecpp::UTF8()));
568                 break;
569         }
570
571     if (dry_run) { // lets us print initialization options
572         printf("simple_output = %i\n", output_type);
573         exit(1);
574     }
575
576     // create a new instance of the expat parser
577     XML_Parser parser = XML_ParserCreate("UTF-8");
578
579     // initialize the elements of the struct to default values
580     init_data(&data, output_type);
581
582
583     // makes the parser pass "data" as the first argument to every callback 
584     XML_SetUserData(parser, &data);
585     void (*startFnPtr)(void*, const XML_Char*, const XML_Char**) = start;
586     void (*endFnPtr)(void*, const XML_Char*) = end;
587     void (*charHandlerFnPtr)(void*, const XML_Char*, int) = charhndl;
588
589     // sets start and end to be the element start and end handlers
590     XML_SetElementHandler(parser, startFnPtr, endFnPtr);
591     // sets charhndl to be the callback for character data
592     XML_SetCharacterDataHandler(parser, charHandlerFnPtr);
593
594     bool done;
595     char buf[BUFSIZ];
596
597     // write header
598
599     cout << "title" << "\t"
600         << "articleid" << "\t"
601         << "revid" << "\t"
602         << "date" << " "
603         << "time" << "\t"
604         << "anon" << "\t"
605         << "editor" << "\t"
606         << "editor_id" << "\t"
607         << "minor" << "\t"
608         << "text_size" << "\t"
609         << "text_entropy" << "\t"
610         << "text_md5" << "\t"
611         << "reversion" << "\t"
612         << "additions_size" << "\t"
613         << "deletions_size";
614
615     int n = 0;
616     if (!data.diff_regexes.empty()) {
617         for (vector<pcrecpp::RE>::iterator r = data.diff_regexes.begin(); r != data.diff_regexes.end(); ++r, ++n) {
618             if (data.diff_regex_names.at(n).empty()) {
619                 cout << "\t" << "regex_" << n << "_add"
620                      << "\t" << "regex_" << n << "_del";
621             } else {
622                 cout << "\t" << data.diff_regex_names.at(n) << "_add"
623                      << "\t" << data.diff_regex_names.at(n) << "_del";
624             }
625         }
626     }
627     cout << endl;
628     
629     // shovel data into the parser
630     do {
631         
632         // read into buf a bufferfull of data from standard input
633         size_t len = fread(buf, 1, BUFSIZ, stdin);
634         done = len < BUFSIZ; // checks if we've got the last bufferfull
635         
636         // passes the buffer of data to the parser and checks for error
637         //   (this is where the callbacks are invoked)
638         if (XML_Parse(parser, buf, len, done) == XML_STATUS_ERROR) {
639             cerr << "XML ERROR: " << XML_ErrorString(XML_GetErrorCode(parser)) << " at line "
640                  << (int) XML_GetCurrentLineNumber(parser) << endl;
641             return 1;
642         }
643     } while (!done);
644    
645
646     XML_ParserFree(parser);
647
648     return 0;
649 }

Benjamin Mako Hill || Want to submit a patch?