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

Benjamin Mako Hill || Want to submit a patch?