import xml.sax
from Queue import Queue
from disco.core import Job, result_iterator
from disco.worker.classic.func import chain_reader

"""
For using this example, you should obtain an sml corpus and do the following:

1. Add the current directory to the python path
    $ export PYTHONPATH=$PYTHONPATH:.

2. For the xml file, we will use a very small portion of wikipedia dump:
    $ wget --no-check-certificate https://raw.githubusercontent.com/pooya/discostuff/master/sample.xml

(The actual wikipedia dump can be found here (very large file):
http://dumps.wikimedia.org/enwiki/latest/enwiki-latest-pages-articles.xml.bz2)

3. Use the xml_reader function in this file to extract the desried tags:
    $ ddfs chunk data:xml:read ./sample.xml --reader xml_reader.xml_reader

   for the wikipedia corpus, the desired tag is "text"

4. Then you can run any job that uses the content of this ddfs tag.
    $ python xml_reader.py
"""

XML_TAG = "text"
DDFS_TAG = "data:xml:read"

class ABContentHandler(xml.sax.ContentHandler):
    def __init__(self, q):
        xml.sax.ContentHandler.__init__(self)
        self.q = q

    def startElement(self, name, attrs):
        self.tag = name
        self.content = ""

    def endElement(self, name):
        if self.tag == XML_TAG:
            self.q.put(self.content)
        self.tag = ""

    def characters(self, content):
        if self.tag == XML_TAG:
            self.content += content

class ReadProxy(object):
    def __init__(self, stream, q):
        self.stream = stream
        self.buffer = ""
        self.q = q

    def read(self, size):
        if self.buffer == "":
            try:
                self.buffer = self.stream.next()
            except:
                self.q.put(0)
                return ""
        if size < len(self.buffer):
            buffer = self.buffer[:size]
            self.buffer = self.buffer[size:]
            return buffer
        else:
            return self._read()

    def _read(self):
        buffer = self.buffer
        self.buffer = ""
        return buffer

def xml_reader(stream, size, url, params):
    q = Queue()
    xml_reader.rproxy = ReadProxy(stream, q)
    s = xml.sax.parse(xml_reader.rproxy, ABContentHandler(q))
    counter = 0
    while True:
        item = q.get()
        if item == 0:
            return
        yield counter, item
        counter += 1
        q.task_done()

def map(line, params):
    yield line, 1

if __name__ == '__main__':
    job = Job().run(input=["tag://" + DDFS_TAG],
                    map=map,
                    map_reader = chain_reader)

    for line, count in result_iterator(job.wait(show=True)):
        print(line, count)

@@ -2224,24 +2224,8 @@
 q))%0A
-    counter = 0%0A
     
@@ -2318,106 +2318,420 @@
 eld 
-counter, item%0A        counter += 1%0A        q.task_done()%0A%0Adef map(line, params):%0A    yield line, 1
+item%0A        q.task_done()%0A%0Adef map(line, params):%0A    import __builtin__%0A    unwanted = u%22,!.#()%5D%5B%7B%7D-%3E%3C=%7C/%5C%22'*:?%22%0A    words = line.translate(__builtin__.dict.fromkeys(%5Bord(x) for x in%0A        unwanted%5D, u%22 %22)).lower()%0A    for word in words.split():%0A        yield word, 1%0A%0Adef reduce(iter, params):%0A    from disco.util import kvgroup%0A    for word, counts in kvgroup(sorted(iter)):%0A        yield word, sum(counts)
 %0A%0Aif
@@ -2833,16 +2833,51 @@
 ap=map,%0A
+                    reduce=reduce,%0A
         

Change the xml_reader to yet another count_words example.