#!/usr/bin/env python
'''Registry of stages for streamcorpus_pipeline.

.. This software is released under an MIT/X11 open source license.
   Copyright 2012-2015 Diffeo, Inc.

Stages are implemented as callable Python objects.  In almost all
cases the name of the stage in the configuration file is the same as
its class name.

Readers
=======

.. autoclass:: streamcorpus_pipeline._convert_kba_json.convert_kba_json
.. autoclass:: streamcorpus_pipeline._local_storage.from_local_chunks
.. autoclass:: streamcorpus_pipeline._local_storage.from_local_files
.. autoclass:: streamcorpus_pipeline._kvlayer.from_kvlayer
.. autoclass:: streamcorpus_pipeline._s3_storage.from_s3_chunks
.. autoclass:: streamcorpus_pipeline._john_smith.john_smith
.. autoclass:: streamcorpus_pipeline._yaml_files_list.yaml_files_list
.. autoclass:: streamcorpus_pipeline._spinn3r_feed_storage.from_spinn3r_feed
.. autoclass:: streamcorpus_pipeline._serifxml.from_serifxml

Incremental transforms
======================

.. autoclass:: streamcorpus_pipeline._clean_html.clean_html
.. autoclass:: streamcorpus_pipeline.force_clean_html.force_clean_html
.. autoclass:: streamcorpus_pipeline._clean_visible.clean_visible
.. autoclass:: streamcorpus_pipeline.offsets.xpath_offsets
.. autoclass:: streamcorpus_pipeline._pdf_to_text.pdf_to_text
.. autoclass:: streamcorpus_pipeline._docx_to_text.docx_to_text
.. autoclass:: streamcorpus_pipeline._title.title
.. autoclass:: streamcorpus_pipeline._filters.debug_filter
.. autoclass:: streamcorpus_pipeline._filters.filter_domains
.. autoclass:: streamcorpus_pipeline._filters.filter_domains_substrings
.. autoclass:: streamcorpus_pipeline._fix_text.fix_text
.. autoclass:: streamcorpus_pipeline._dedup.dedup
.. autoclass:: streamcorpus_pipeline._dump_label_stats.dump_label_stats
.. autoclass:: streamcorpus_pipeline._filters.id_filter
.. autofunction:: streamcorpus_pipeline._guess_media_type.file_type_stats
.. autoclass:: streamcorpus_pipeline._filters.filter_languages
.. autoclass:: streamcorpus_pipeline._filters.filter_tagger_ids
.. autoclass:: streamcorpus_pipeline._find.find
.. autoclass:: streamcorpus_pipeline._find.find_doc_ids
.. autoclass:: streamcorpus_pipeline._guess_media_type.guess_media_type
.. autoclass:: streamcorpus_pipeline._handle_unconvertible_spinn3r.handle_unconvertible_spinn3r
.. autoclass:: streamcorpus_pipeline._hyperlink_labels.hyperlink_labels
.. autoclass:: streamcorpus_pipeline._upgrade_streamcorpus.keep_annotated
.. autoclass:: streamcorpus_pipeline._language.language
.. autoclass:: streamcorpus_pipeline._nilsimsa.nilsimsa
.. autoclass:: streamcorpus_pipeline._filters.remove_raw
.. autoclass:: streamcorpus_pipeline._filters.replace_raw
.. autoclass:: streamcorpus_pipeline._filters.dump_stream_id_abs_url
.. autoclass:: streamcorpus_pipeline._upgrade_streamcorpus.upgrade_streamcorpus
.. autoclass:: streamcorpus_pipeline._upgrade_streamcorpus_v0_3_0.upgrade_streamcorpus_v0_3_0
.. autoclass:: streamcorpus_pipeline._tokenizer.nltk_tokenizer

Batch transforms
================

.. autoclass:: streamcorpus_pipeline._taggers.byte_offset_align_labels
.. autoclass:: streamcorpus_pipeline._taggers.line_offset_align_labels
.. autoclass:: streamcorpus_pipeline._taggers.name_align_labels
.. autoclass:: streamcorpus_pipeline._lingpipe.lingpipe
.. autoclass:: streamcorpus_pipeline._serif.serif

Writers
=======

.. autoclass:: streamcorpus_pipeline._local_storage.to_local_chunks
.. autoclass:: streamcorpus_pipeline._local_storage.to_local_tarballs
.. autoclass:: streamcorpus_pipeline._kvlayer.to_kvlayer
.. autoclass:: streamcorpus_pipeline._s3_storage.to_s3_chunks
.. autoclass:: streamcorpus_pipeline._s3_storage.to_s3_tarballs

API
===

All stage objects are constructed with a single parameter, the
dictionary of configuration specific to the stage.  The stage objects
can be passed directly to the ``streamcorpus_pipeline.Pipeline``
constructor.

.. currentmodule:: streamcorpus_pipeline.stages

.. autoclass:: StageRegistry
   :members:
   :show-inheritance:

.. autoclass:: PipelineStages
   :members:
   :show-inheritance:

.. autoclass:: Configured
   :members:
   :show-inheritance:

.. autoclass:: BatchTransform
   :members:
   :show-inheritance:

.. autoclass:: IncrementalTransform
   :members:
   :show-inheritance:

'''
from __future__ import absolute_import
from abc import ABCMeta, abstractmethod
from collections import Callable, MutableMapping
import imp
import logging
import time

import pkg_resources

logger = logging.getLogger(__name__)

class StageRegistry(MutableMapping):
    '''List of known stages for a pipeline.

    This is a dictionary, and so ``registry[stagename]`` will retrieve
    the class object or constructor function for the stage.  Each
    stage in the dictionary is a callable of a single parameter, the
    configuration dictionary, which returns a callable appropriate
    for its stage type.  Typical patterns are actual classes and
    wrapper functions::

        class a_stage(streamcorpus_pipeline.stages.Configurable):
            config_name = 'a_stage'
            def __init__(self, config, *args, **kwargs):
                super(AStage, self).__init__(config, *args, **kwargs)
                self.param = self.config.get('param')
            def __call__(self, si, context):
                logger.debug('param is %s', self.param)
                return si

        def b_stage(config):
            param = config.get('param')
            def do_stage(si, context):
               logger.debug('param is %s', param)
               return si
            return do_stage

        registry = StageRegistry()
        registry['a_stage'] = a_stage
        registry['b_stage'] = b_stage

    '''
    def __init__(self, *args, **kwargs):
        super(StageRegistry, self).__init__(*args, **kwargs)
        self.stages = {}

    # MutableMapping special methods

    def __getitem__(self, k):
        if k not in self.stages:
            raise KeyError('%r not in %r' % (k, self.stages.keys()))
        return self.stages[k]

    def __setitem__(self, k, v):
        self.stages[k] = v

    def __delitem__(self, k):
        del self.stages[k]

    def __iter__(self):
        return self.stages.iterkeys()

    def __len__(self):
        return len(self.stages)

    # StageRegistry-specific methods
    def tryload_stage(self, moduleName, functionName, name=None):
        '''Try to load a stage into self, ignoring errors.

        If loading a module fails because of some subordinate load
        failure, just give a warning and move on.  On success the
        stage is added to the stage dictionary.

        :param str moduleName: name of the Python module
        :param str functionName: name of the stage constructor
        :param str name: name of the stage, defaults to `functionName`

        '''
        if name is None:
            name = functionName
        try:
            mod = __import__(moduleName, globals(), locals(), [functionName])
        except ImportError, exc:
            logger.warn('cannot load stage {0}: cannot load module {1}'
                        .format(name, moduleName), exc_info=exc)
            return

        if not hasattr(mod, functionName):
            logger.warn('cannot load stage {0}: module {1} missing {2}'
                        .format(name, moduleName, functionName))
            return

        self[name] = getattr(mod, functionName)

    def load_external_stages(self, path):
        '''Add external stages from the Python module in `path`.

        `path` must be a path to a Python module source that contains
        a `Stages` dictionary, which is a map from stage name to callable.

        :param str path: path to the module file
        '''
        mod = imp.load_source('', path)
        self.update(mod.Stages)

    def load_module_stages(self, mod):
        '''Add external stages from the Python module `mod`.

        If `mod` is a string, then it will be interpreted as the name
        of a module; otherwise it is an actual module object.  The
        module should exist somewhere in :data:`sys.path`.  The module
        must contain a `Stages` dictionary, which is a map from stage
        name to callable.

        :param mod: name of the module or the module itself
        :raise exceptions.ImportError: if `mod` cannot be loaded or does
          not contain ``Stages``

        '''
        if isinstance(mod, basestring):
            mod = __import__(mod, globals=globals(), locals=locals(),
                             fromlist=['Stages'], level=0)
        if not hasattr(mod, 'Stages'):
            raise ImportError(mod)
        self.update(mod.Stages)

    def init_stage(self, name, config):

        '''Construct and configure a stage from known stages.

        `name` must be the name of one of the stages in this.  `config`
        is the configuration dictionary of the containing object, and its `name`
        member will be passed into the stage constructor.

        :param str name: name of the stage
        :param dict config: parent object configuration
        :return: callable stage
        :raise exceptions.KeyError: if `name` is not a known stage

        '''
        subconfig = config.get(name, {})
        ctor = self[name]
        return ctor(subconfig)

class Configured(object):
    '''Any object containing a configuration.

    The configuration is passed to the constructor.

    .. attribute:: config

        The configuration dictionary for this object.

    '''
    def __init__(self, config, *args, **kwargs):
        super(Configured, self).__init__(*args, **kwargs)
        self.config = config

class BatchTransform(Configured):
    '''Transform that acts on an entire :class:`streamcorpus.Chunk` file.'''
    __metaclass__ = ABCMeta

    @abstractmethod
    def process_path(self, chunk_path):
        '''Process a chunk file at a specific path.

        Implementations should work in place, writing results out to the
        same `chunk_path`.  The pipeline guarantees that this stage
        will have a key ``tmp_dir_path`` in its configuration, and this
        can be used to store an intermediate file which is then renamed
        over `chunk_path`.
        '''
        raise NotImplementedError('BatchTransform.process_path not implemented')

    @abstractmethod
    def shutdown(self):
        '''Do an orderly shutdown of this stage.

        If the stage spawns a child process, kill it, and do any other
        required cleanup.  :meth:`streamcorpus_pipeline.Pipeline._cleanup`
        will call this method on every batch transform stage, regardless
        of whether this is the currently running stage or not.
        '''
        raise NotImplementedError('shutdown is a required method for all BatchTransforms')

class IncrementalTransform(Configured, Callable):
    '''Transform that acts on individual streamcorpus StreamItem objects.

    The pipeline can run any stage with an appropriate :meth:`__call__`
    method.  The :meth:`__call__` or :meth:`process_item` methods are
    called once per stream item, reusing the same transform object.

    '''
    __metaclass__ = ABCMeta
    @abstractmethod
    def process_item(self, stream_item, context):
        '''Process a single :class:`streamcorpus.StreamItem` object.

        `context` is shared across all stages.  It is guaranteed to
        contain ``i_str`` nominally containing the original input file
        name and ``data`` nominally containing the auxiliary data from
        the work item.

        This function may modify `stream_item` in place, construct
        a new item and return it, or return :const:`None` to drop
        the item.

        :param stream_item: stream item to process
        :paramtype stream_item: :class:`streamcorpus.StreamItem`
        :param dict context: additional shared context data
        :return: same or different stream item or :const:`None`

        '''
        raise NotImplementedError('BatchTransform.process_path not implemented')

    def __call__(self, stream_item, context):
        '''Entry point from the pipeline.

        This implementation calls :meth:`process_item`.

        '''
        return self.process_item(stream_item, context)

    def shutdown(self):
        '''Do an orderly shutdown of this stage.

        Incremental transforms are generally simple Python code and
        a complicated shutdown is not required.  The pipeline does not
        call this method.
        '''
        pass

class PipelineStages(StageRegistry):
    '''The list of stages, preloaded with streamcorpus_pipeline stages.'''
    def __init__(self, *args, **kwargs):
        start = time.time()
        super(PipelineStages, self).__init__(*args, **kwargs)

        # data source readers (read data from somewhere into pipeline)
        self.tryload_stage('_convert_kba_json', 'convert_kba_json')
        self.tryload_stage('_local_storage', 'from_local_chunks')
        self.tryload_stage('_local_storage', 'from_local_files')
        self.tryload_stage('_kvlayer', 'from_kvlayer')
        self.tryload_stage('_s3_storage', 'from_s3_chunks')
        self.tryload_stage('_john_smith', 'john_smith')
        self.tryload_stage('_yaml_files_list', 'yaml_files_list')
        self.tryload_stage('_spinn3r_feed_storage', 'from_spinn3r_feed')
        self.tryload_stage('_serifxml', 'from_serifxml')

        # StreamItem stages
        # (alphabetical by stage name)
        self.tryload_stage('_clean_html', 'clean_html')
        self.tryload_stage('force_clean_html', 'force_clean_html')
        self.tryload_stage('_clean_visible', 'clean_visible')
        self.tryload_stage('offsets', 'xpath_offsets')
        self.tryload_stage('_pdf_to_text', 'pdf_to_text')
        self.tryload_stage('_docx_to_text', 'docx_to_text')
        self.tryload_stage('_title', 'title')
        self.tryload_stage('_filters', 'debug_filter')
        self.tryload_stage('_filters', 'filter_domains')
        self.tryload_stage('_filters', 'filter_domains_substrings')
        self.tryload_stage('_fix_text', 'fix_text')
        self.tryload_stage('_dedup', 'dedup')
        self.tryload_stage('_dump_label_stats', 'dump_label_stats')
        self.tryload_stage('_filters', 'id_filter')
        self.tryload_stage('_guess_media_type', 'file_type_stats')
        self.tryload_stage('_filters', 'filter_languages')
        self.tryload_stage('_filters', 'filter_tagger_ids')
        self.tryload_stage('_find', 'find')
        self.tryload_stage('_find', 'find_doc_ids')
        self.tryload_stage('_guess_media_type', 'guess_media_type')
        self.tryload_stage('_handle_unconvertible_spinn3r', 'handle_unconvertible_spinn3r')
        self.tryload_stage('_hyperlink_labels', 'hyperlink_labels')
        self.tryload_stage('_upgrade_streamcorpus', 'keep_annotated')
        self.tryload_stage('_upgrade_streamcorpus', 'keep_annotated', 'keep_annotatoted')
        self.tryload_stage('_language', 'language')
        self.tryload_stage('_nilsimsa', 'nilsimsa')
        self.tryload_stage('_filters', 'remove_raw')
        self.tryload_stage('_filters', 'replace_raw')
        self.tryload_stage('_filters', 'dump_stream_id_abs_url')
        self.tryload_stage('_set_source', 'set_source')
        self.tryload_stage('_upgrade_streamcorpus', 'upgrade_streamcorpus')
        self.tryload_stage('_upgrade_streamcorpus_v0_3_0', 'upgrade_streamcorpus_v0_3_0')
        self.tryload_stage('_tokenizer', 'nltk_tokenizer')

        # BatchTransform
        self.tryload_stage('_taggers', 'byte_offset_align_labels')
        self.tryload_stage('_taggers', 'char_offset_align_labels')
        #self.tryload_stage('_taggers', 'line_offset_align_labels')
        self.tryload_stage('_taggers', 'name_align_labels')
        self.tryload_stage('_taggers', 'multi_token_match_align_labels')
        self.tryload_stage('_lingpipe', 'lingpipe')
        self.tryload_stage('_serif', 'serif')

        # 'writers' move data out of the pipeline
        self.tryload_stage('_local_storage', 'to_local_chunks')
        self.tryload_stage('_local_storage', 'to_local_tarballs')
        self.tryload_stage('_kvlayer', 'to_kvlayer')
        self.tryload_stage('_s3_storage', 'to_s3_chunks')
        self.tryload_stage('_s3_storage', 'to_s3_tarballs')

        # load from setuptools 'entry_points' of other installed packages
        for entry_point in pkg_resources.iter_entry_points('streamcorpus_pipeline.stages'):
            try:
                name = entry_point.name
                stage_constructor = entry_point.load()
                self[name] = stage_constructor
            except Exception, exc:
                import traceback
                print(traceback.format_exc(exc))
                logger.error('failure loading plugin entry point: %r', entry_point and entry_point.name, exc_info=True)
                logger.error('IGNORING plug loading failure and continuing on...')

        logger.debug('streamcorpus_pipeline.PipelineStages init in %s sec', time.time() - start)


#!/usr/bin/env python
'''
Pipeline stage for extracting hyperlinks to particular domains as
labels generated by the author.

This software is released under an MIT/X11 open source license.

Copyright 2012-2014 Diffeo, Inc.
'''
from __future__ import absolute_import
import re
import sys
import logging

from streamcorpus import add_annotation, Offset, OffsetType, Annotator, Target, Label
from streamcorpus_pipeline._clean_visible import make_clean_visible
from streamcorpus_pipeline.stages import Configured

logger = logging.getLogger(__name__)

anchors_re = re.compile('''(?P<before>(.|\n)*?)''' + \
                        '''(?P<ahref>\<a\s+(.|\n)*?href''' + \
                        '''(?P<preequals>(\s|\n)*)=(?P<postequals>(\s|\n)*)''' + \
                        '''(?P<quote>("|')?)(?P<href>[^"]*)(?P=quote)''' + \
                        '''(?P<posthref>(.|\n)*?)\>)''' + \
                        '''(?P<anchor>(.|\n)*)''', re.I)

def read_to( idx_bytes, stop_bytes=None, run_bytes=None ):
    '''
    iterates through idx_bytes until a byte in stop_bytes or a byte
    not in run_bytes.

    :rtype (int, string): idx of last byte and all of bytes including
    the terminal byte from stop_bytes or not in run_bytes
    '''
    idx = None
    vals = []
    next_b = None
    while 1:
        try:
            idx, next_b = idx_bytes.next()
        except StopIteration:
            ## maybe something going wrong?
            idx = None
            next_b = None
            break
        ## stop when we see any byte in stop_bytes
        if stop_bytes and next_b in stop_bytes:
            break
        ## stop when we see any byte not in run_bytes
        if run_bytes and next_b not in run_bytes:
            break
        ## assemble the ret_val
        vals.append( next_b )

    ## return whatever we have assembled
    return idx, b''.join(vals), next_b

def iter_attrs( idx_bytes ):
    '''
    called when idx_chars is just past "'])
        attr_vals = []

        ## stop if we hit the end of the tag, or end of idx_bytes
        if next_b is None or next_b == '>':
            return

        idx, space, quote = read_to(idx_bytes, run_bytes = [' ', '\t', '\n', '\r'])
        if quote not in ['"', "'"]:
            ## caught start of the property value
            attr_vals = [quote]
            quote = ' '

        idx, attr_val, next_b = read_to(idx_bytes, [quote, '>'])
        ## next_b had better either balance the start quote, end the
        ## tag, or end idx_bytes
        assert next_b in [quote, '>', None], attr_val
        attr_vals.append( attr_val )

        yield idx, attr_name.strip(), b''.join(attr_vals).strip()

class hyperlink_labels(Configured):
    '''Creates document labels from hyperlinks in ``body.clean_html``.

    The labels are attached to the stream item body with an annotator
    ID of ``author``.  Any HTML ```` matching the selection
    criteria will be turned into a label.

    You generally must set either ``all_domains`` or
    ``domain_substrings``.  A typical configuration will look like:

    .. code-block:: yaml

        streamcorpus_pipeline:
          incremental_transforms: [ ..., hyperlink_labels, ... ]
          hyperlink_labels:
            require_abs_url: true
            domain_substrings: [ "trec-kba.org" ]

    Configuration options:

    .. code-block:: yaml

        all_domains: false
        domain_substrings: [ 'trec-kba.org' ]

    A label will only be produced if ``all_domains`` is true, or if
    the hostname part of the URL has one of ``domain_substrings`` as a
    substring.  Note that the default configuration is ``all_domains``
    false and an empty ``domain_substrings`` list, so no labels will
    be produced without additional configuration.

    .. code-block:: yaml

        require_abs_url: true

    Only produce labels for fully-qualified ``http://...`` URLs.  False
    by default.

    .. code-block:: yaml

        offset_types: [BYTES]

    A list containing at least one of ``BYTES`` and ``LINES``
    indicating what sort of document offset should be attached to the
    label.  Only the first value is used.  ``LINES`` is not
    recommended.

    .. code-block:: yaml

        require_clean_html: true

    Cannot be changed.

    '''
    config_name = 'hyperlink_labels'
    default_config = {
        'require_clean_html': True,
        'require_abs_url': False,
        'all_domains': False,
        'domain_substrings': [],
        'offset_types': ['BYTES'],
    }

    def __init__(self, *args, **kwargs):
        super(hyperlink_labels, self).__init__(*args, **kwargs)
        self.offset_type = getattr(OffsetType, self.config['offset_types'][0])
        if self.offset_type != OffsetType.BYTES:
            logger.warn('using offset_type other than BYTES: %r' % self.offset_type)

    def href_filter(self, href):
        '''
        Test whether an href string meets criteria specified by
        configuration parameters 'require_abs_url', which means "does
        it look like it is probably an absolute URL?" and
        'domain_substrings'.  It searches for each of the
        domain_substrings in the href individually, and if any match,
        then returns True.

        :param: href string
        :returns bool:
        '''
        if self.config['require_abs_url']:
            if not href.lower().startswith(('http://', 'https://')):
                return False
        if self.config['all_domains']:
            ## blanket accept all domains as labels
            return True

        if self.config['domain_substrings']:
            parts = href.split('/')
            if len(parts) < 3:
                return False
            domain = parts[2].lower()
            for substring in self.config['domain_substrings']:
                try:
                    if substring in domain:
                        return True
                except Exception, exc:
                    logger.warn('%r in %r raised', substring, domain, exc_info=True)
                    return False

    def line_href_anchors(self):
        '''
        simple, regex-based extractor of anchor tags, so we can
        compute LINE offsets for anchor texts and associate them with
        their href.

        Generates tuple(href_string, first_byte, byte_length, anchor_text)

        Also, this mangles the body.clean_html so that LINE offsets
        uniquely identify tokens in anchor tags -- quite a kludge.
        '''
        idx = 0
        new_clean_html = ''
        newlines_added = 0
        ## split doc up into pieces that end on an anchor tag
        parts = self.clean_html.split('')
        assert len(''.join(parts) ) == len(self.clean_html)
        for i in range(len(parts)):
            part = parts[i]

            ## try to get an A tag out:
            m = anchors_re.match(part)

            ## if not, then just keep going
            if not m:
                new_clean_html += part
                if i < len(parts) - 1:
                    new_clean_html += ''
                continue

            before = m.group('before')
            ahref = m.group('ahref')

            ## construct a text containing bytes up to the anchor text
            pre_anchor_increment = before + ahref

            ## increment the index to get line number for the anchor
            idx += len( pre_anchor_increment.splitlines() )
            first = idx

            ## usually this will be one, but it could be more than
            ## that when an anchor text contains newlines
            length = len(m.group('anchor').split('\n'))

            ## construct new clean_html with these newlines inserted
            new_clean_html += pre_anchor_increment + '\n' + m.group('anchor') + '\n'

            newlines_added += 2

            ## update the index for the next loop
            idx += length - 1

            yield m.group('href'), first, length, m.group('anchor')

        ## replace clean_html with our new one that has newlines inserted
        assert len(self.clean_html) == len(new_clean_html) - newlines_added
        self.clean_html = new_clean_html

    def char_href_anchors(self, chars=False):
        for tup in self.byte_href_anchors(chars=True):
            yield tup

    def byte_href_anchors(self, chars=False):
        '''
        simple, regex-based extractor of anchor tags, so we can
        compute BYTE offsets for anchor texts and associate them with
        their href.

        Generates tuple(href_string, first_byte, byte_length, anchor_text)
        '''
        input_buffer = self.clean_html
        if chars:
            input_buffer = input_buffer.decode('utf8')
        idx = 0
        ## split doc up into pieces that end on an anchor tag
        parts = input_buffer.split('')
        assert len(''.join(parts) ) == len(input_buffer)
        for part in parts:
            ## try to get an A tag out:
            m = anchors_re.match(part)

            if not m:
                idx += len(part) + 4
                continue

            before = m.group('before')
            ahref = m.group('ahref')

            ## increment the index to get line number for the anchor
            idx += len(before) + len(ahref)
            first = idx

            length = len(m.group('anchor'))

            ## update the index for the next loop
            # include anchor plus the
            idx += length + 4

            if chars:
                yield m.group('href').encode('utf8'), first, length, m.group('anchor').encode('utf8')
            else:
                yield m.group('href'), first, length, m.group('anchor')

        assert idx - 4 == len(input_buffer)

    def byte_href_anchors_state_machine(self):
        '''
        byte-based state machine extractor of anchor tags, so we can
        compute byte offsets for anchor texts and associate them with
        their href.

        Generates tuple(href_string, first_byte, byte_length, anchor_text)
        '''
        tag_depth = 0
        a_tag_depth = 0
        vals = []
        href = None
        idx_bytes = enumerate( self.clean_html )
        while 1:
            end_idx, val, next_b = read_to( idx_bytes, '<' )
            tag_depth += 1

            if href:
                ## must be inside an anchor tag, so accumulate the
                ## whole anchor
                assert a_tag_depth > 0, (href, self.clean_html)
                vals.append(val)

            ## figure out if start of an "A" anchor tag or close
            ## of a previous tag
            idx, next_b1 = idx_bytes.next()
            if next_b1.lower() == 'a':
                ## could be start of "A" tag
                idx, next_b2 = idx_bytes.next()
                if next_b2 == ' ':
                    a_tag_depth += 1

                    href = None
                    for idx, attr_name, attr_val in iter_attrs( idx_bytes ):
                        if attr_name.lower() == 'href':
                            href = attr_val
                        if idx is None:
                            ## doc ended mid tag, so invalid HTML--> just bail
                            return

                    first = idx + 1

                    ## if we got an href, then we want to keep the
                    ## first byte idx of the anchor:
                    if href:
                        ## Someone could nest an A tag inside another
                        ## A tag, which is invalid (even in HTML5), so
                        ## vals could be nonempty.  We only generate
                        ## the leaf-level A tags in these rare cases
                        ## of nested A tags, so reset it:
                        vals = []

            elif next_b1 == '/':
                idx, next_b1 = idx_bytes.next()
                if next_b1 == 'a':
                    ## could be end of "A" tag
                    idx, next_b2 = idx_bytes.next()
                    if next_b2 == '>':
                        a_tag_depth -= 1
                        if href:
                            ## join is much faster than using += above
                            anchor = b''.join(vals)
                            length = len(anchor)

                            ## yield the data
                            yield href, first, len(anchor), anchor

                            ## reset, no yield again in a nested A tag
                            href = None

            else:
                ## the next_b was not part of  or a nested <a tag,
                ## so keep it in the output
                vals.append(next_b)

    def make_labels(self, clean_html, clean_visible=None):
        '''
        Make a list of Labels for 'author' and the filtered hrefs &
        anchors
        '''
        if   self.offset_type == OffsetType.BYTES:
            parser = self.byte_href_anchors

        elif self.offset_type == OffsetType.CHARS:
            parser = self.char_href_anchors

        elif self.offset_type == OffsetType.LINES:
            parser = self.line_href_anchors

        labels = []
        ## make clean_html accessible as a class property so we can
        self.clean_html = clean_html
        for href, first, length, value in parser():
            if self.href_filter(href):
                '''
                if clean_visible:
                    _check_html = self.clean_html.splitlines()[first-10:10+first+length]
                    _check_visi =   clean_visible.splitlines()[first:first+length]
                    if not make_clean_visible(_check_html) == _check_visi:
                        print len(self.clean_html.splitlines())
                        print len(clean_visible.splitlines())

                        print href
                        print '\t html: %r' % _check_html
                        print '\t visi: %r' % _check_visi
                '''
                ## add a label for every href
                label = Label(
                    annotator = Annotator(annotator_id = 'author'),
                    target = Target(target_id = href),
                    )
                ## the offset type is specified by the config
                label.offsets[self.offset_type] = Offset(
                    first=first, length=length,
                    value=value,
                    ## the string name of the content field, not the
                    ## content itself :-)
                    content_form='clean_html')
                labels.append(label)

        return labels

    def __call__(self, stream_item, context):
        '''
        Act as an incremental transform in the kba.pipeline
        '''
        ## right now, we only do clean_html
        assert self.config.get('require_clean_html', True)

        if stream_item.body and stream_item.body.clean_html:
            labels = self.make_labels(stream_item.body.clean_html,
                                      stream_item.body.clean_visible)
            if labels:
                if self.offset_type == OffsetType.LINES:
                    ## for LINES-type labels, must replace clean_html
                    ## with a new one that has newlines inserted
                    stream_item.body.clean_html = self.clean_html

                ## Remove any previous author labels
                stream_item.body.labels['author'] = []

                ## also add the new labels
                add_annotation(stream_item.body, *labels)
        return stream_item

if __name__ == '__main__':
    clean_html = sys.stdin.read()
    for m in anchors_re.finditer(clean_html):
        print m.group('href'), m.group('anchor')


'''Example of how to transform a corpus into the streamcorpus format
and save it as a chunk file.

~/streamcorpus-pipeline$ python examples/john_smith_chunk_writer.py data/john-smith/original foo.sc
~/streamcorpus-pipeline$ streamcorpus_dump --count foo.sc
197     0       foo.sc

Copyright 2012-2014 Diffeo, Inc.
'''

import argparse
import logging
import os

## this assumes that streamcorpus has been installed; it is in pypi
import streamcorpus

logger = logging.getLogger(__name__)

def paths(input_dir):
    'yield all file paths under input_dir'
    for root, dirs, fnames in os.walk(input_dir):
        for i_fname in fnames:
            i_path = os.path.join(root, i_fname)
            yield i_path

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        'input_dir',
        help='path to a directory containing files to put in a streamcorpus.Chunk file')
    parser.add_argument(
        'output_path',
        help='file path to create the Chunk file')
    args = parser.parse_args()

    o_fh = open(args.output_path, 'wb')
    o_chunk = streamcorpus.Chunk(file_obj=o_fh, mode='wb')

    ## NB: can also pass a path into Chunk.  If it ends in .xz, then
    ## it will handle file compress for you.
    #o_chunk = streamcorpus.Chunk('output.sc.xz', mode='wb')

    for i_path in paths(args.input_dir):
        ### In the example code below, strings that start with
        ### 'headers.' should be replaced by code that pulls the
        ### indicated information from scrapy

        ## Every StreamItem has a stream_time property.  It usually
        ## comes from the document creation time.  This may be either
        ## a unix-time number or a string like
        creation_time = '2000-01-01T12:34:00.000123Z'  ## should come from headers.last_modified

        stream_item = streamcorpus.make_stream_item(
            creation_time,
            'headers.absolute URL')

        ## These docs came from the authors of the paper cited above.
        stream_item.source = 'scrapy'

        i_file = open(i_path)
        stream_item.body.raw = i_file.read()

        stream_item.body.media_type = 'headers.mime_type'
        stream_item.body.encoding = 'headers.encoding'

        o_chunk.add(stream_item)

    o_chunk.close()

if __name__ == '__main__':
    main()


'''
Provides a simple key=value functionality built on a cassandra.

This software is released under an MIT/X11 open source license.

Copyright 2012-2013 Diffeo, Inc.
'''
import json
import random
import hashlib
import logging
logger = logging.getLogger(__name__)

## get the Cassandra client library
import pycassa
from pycassa.pool import ConnectionPool
from pycassa.system_manager import SystemManager, SIMPLE_STRATEGY, \
    TIME_UUID_TYPE, TIME_UUID_TYPE, ASCII_TYPE, BYTES_TYPE, \
    COUNTER_COLUMN_TYPE
from pycassa.types import CounterColumnType, UTF8Type

def _delete_namespace(config):
    sm = SystemManager(config['storage_addresses'][0])
    sm.drop_keyspace(config['namespace'])
    sm.close()

class Cassa(object):
    '''
    Provides a simple key=value functionality built on a cassandra
    table with a key and a single column.  Used in ZookeeperTaskQueue
    to replace the two tables that grow in size with the number of
    tasks rather than the number of workers.
    '''

    def __init__(self, namespace, server_list=['localhost:9160']):
        # save cassandra server
        self.server_list = server_list
        self.namespace = namespace
        self._closed = False

        #setup_logging(self)

        # Connect to the server creating the namespace if it doesn't
        # already exist
        try:
            self.pool = ConnectionPool(namespace, self.server_list, max_retries=500, pool_timeout=600, timeout=10)
        except pycassa.InvalidRequestException:
            self._create_namespace(namespace)
            self.pool = ConnectionPool(namespace, self.server_list, max_retries=500, pool_timeout=600, timeout=10)

        try:
            self._tasks = pycassa.ColumnFamily(self.pool, 'tasks')
        except pycassa.NotFoundException:
            self._create_column_family('tasks',
                                       key_validation_class=ASCII_TYPE,
                                       bytes_columns=['task_data'])
            self._tasks = pycassa.ColumnFamily(self.pool, 'tasks')

        try:
            self._available = pycassa.ColumnFamily(self.pool, 'available')
        except pycassa.NotFoundException:
            self._create_column_family('available',
                                        key_validation_class=ASCII_TYPE,
                                        bytes_columns=['available'])
            self._available = pycassa.ColumnFamily(self.pool, 'available')

        try:
            self._task_count = pycassa.ColumnFamily(self.pool, 'task_count')
        except pycassa.NotFoundException:
            self._create_counter_column_family('task_count',
                                       key_validation_class=ASCII_TYPE,
                                       counter_columns=['task_count'])
            self._task_count = pycassa.ColumnFamily(self.pool, 'task_count')
            self._task_count.insert('RowKey', {'task_count': 0})

        try:
            self._available_count = pycassa.ColumnFamily(self.pool, 'available_count')
        except pycassa.NotFoundException:
            self._create_counter_column_family('available_count',
                                       key_validation_class=ASCII_TYPE,
                                       counter_columns=['available_count'])
            self._available_count = pycassa.ColumnFamily(self.pool, 'available_count')
            self._available_count.insert('RowKey', {'available_count': 0})

    def delete_namespace(self):
        sm = SystemManager(random.choice(self.server_list))
        sm.drop_keyspace(self.namespace)
        sm.close()

    def _create_namespace(self, namespace):
        sm = SystemManager(random.choice(self.server_list))
        sm.create_keyspace(namespace, SIMPLE_STRATEGY, {'replication_factor': '1'})
        sm.close()

    def _create_column_family(self, family, bytes_columns=[],
                              key_validation_class=TIME_UUID_TYPE):
        '''
        Creates a column family of the name 'family' and sets any of
        the names in the bytes_column list to have the BYTES_TYPE.

        key_validation_class defaults to TIME_UUID_TYPE and could also
        be ASCII_TYPE for md5 hash keys, like we use for 'inbound'
        '''
        sm = SystemManager(random.choice(self.server_list))
        # sys.create_column_family(self.namespace, family, super=False)
        sm.create_column_family(self.namespace, family, super=False,
                key_validation_class = key_validation_class,
                default_validation_class  = TIME_UUID_TYPE,
                column_name_class = ASCII_TYPE)
        for column in bytes_columns:
            sm.alter_column(self.namespace, family, column, BYTES_TYPE)
        sm.close()

    def _create_counter_column_family(self, family, counter_columns=[],
                              key_validation_class=UTF8Type):
        '''
        Creates a column family of the name 'family' and sets any of
        the names in the bytes_column list to have the BYTES_TYPE.

        key_validation_class defaults to TIME_UUID_TYPE and could also
        be ASCII_TYPE for md5 hash keys, like we use for 'inbound'
        '''
        sm = SystemManager(random.choice(self.server_list))
        # sys.create_column_family(self.namespace, family, super=False)
        sm.create_column_family(self.namespace, family, super=False,
                key_validation_class = key_validation_class,
                default_validation_class="CounterColumnType",
                column_name_class = ASCII_TYPE)
        for column in counter_columns:
            sm.alter_column(self.namespace, family, column, COUNTER_COLUMN_TYPE)
        sm.close()

    def tasks(self, key_prefix=''):
        '''
        generate the data objects for every task
        '''
        for row in self._tasks.get_range():
            logger.debug(row)
            if not row[0].startswith(key_prefix):
                continue
            data = json.loads(row[1]['task_data'])
            data['task_key'] = row[0]
            yield data

    def put_task(self, key, task_data):
        try:
            found = self._tasks.get(key, column_count=1)
            exists = True
        except pycassa.cassandra.ttypes.NotFoundException:
            exists = False

        self._tasks.insert(key, {'task_data': json.dumps(task_data)})
        if not exists:
            self._task_count.insert('RowKey', {'task_count': 1})
        return exists

    def get_task(self, key):
        data = self._tasks.get(key)
        return json.loads(data['task_data'])

    def pop_task(self, key):
        self._tasks.remove(key)
        self._task_count.insert('RowKey', {'task_count': -1})
        return key

    @property
    def task_keys(self):
        c = 0
        for key, _ in self._tasks.get_range(column_count=0, filter_empty=False):
            c += 1
            yield key

    def num_tasks(self):
        data = self._task_count.get('RowKey')
        return data['task_count']

    def num_available(self):
        data = self._available_count.get('RowKey')
        return data['available_count']

    def put_available(self, key):
        ## closest thing to storing only the key
        try:
            found = self._available.get(key, column_count=1)
            exists = True
        except pycassa.cassandra.ttypes.NotFoundException:
            exists = False

        if not exists:
            self._available.insert(key, {'available': ''})
            self._available_count.insert('RowKey', {'available_count': 1})

    #def push_batch(self, row_iter):
    #    '''
    #    Push opaque vertex data objects into the inbound queue
    #    '''
    #    return self._tasks.batch_insert({k: json.dumps(v) for k, v in row_iter})

    def get_random_available(self, max_iter=10000):
        '''
        get a random key out of the first max_iter rows
        '''
        c = 1
        keeper = None
        ## note the ConsistencyLevel here.  If we do not do this, and
        ## get all slick with things like column_count=0 and filter
        ## empty False, then we can get keys that were recently
        ## deleted... EVEN if the default consistency would seem to
        ## rule that out!

        ## note the random start key, so that we do not always hit the
        ## same place in the key range with all workers
        #random_key = hashlib.md5(str(random.random())).hexdigest()
        #random_key = '0' * 32
        #logger.debug('available.get_range(%r)' % random_key)
        ## scratch that idea: turns out that using a random start key
        ## OR using row_count=1 can cause get_range to hang for hours

        ## why we need ConsistencyLevel.ALL on a single node is not
        ## clear, but experience indicates it is needed.

        ## note that putting a finite row_count is problematic in two
        ## ways:
        # 1) if there are more workers than max_iter, some will not
        # get tasks
        #
        # 2) if there are more than max_iter records, then all workers
        # have to wade through all of these just to get a task!  What
        # we really want is a "pick random row" function, and that is
        # probably best implemented using CQL3 token function via the
        # cql python module instead of pycassa...
        for row in self._available.get_range(row_count=max_iter, read_consistency_level=pycassa.ConsistencyLevel.ALL):
        #for row in self._available.get_range(row_count=100):
            logger.debug('considering %r' % (row,))
            if random.random() < 1 / c:
                keeper = row[0]
            if c == max_iter:
                break
            c += 1
        return keeper

    def in_available(self, key):
        try:
            row = self._available.get(key)
            return True
        except pycassa.NotFoundException:
            return False

    def pop_available(self, key):
        self._available.remove(key, write_consistency_level=pycassa.ConsistencyLevel.ALL)
        self._available_count.insert('RowKey', {'available_count': -1})
        assert not self.in_available(key)
        return key

    def close(self):
        self._closed = True
        if hasattr(self, 'pool'):
            self.pool.dispose()


#!/usr/bin/env python
'''
streamcorpus_pipeline.TaggerBatchTransform for LingPipe

This software is released under an MIT/X11 open source license.

Copyright 2012-2014 Diffeo, Inc.
'''
from __future__ import absolute_import
import hashlib
import itertools
import logging
import os
import re
import sys
import time
import traceback
import uuid

from nltk.tokenize import WhitespaceTokenizer

from streamcorpus import Token, Sentence, EntityType, Chunk, Offset, \
    OffsetType, Gender, MentionType, Attribute, AttributeType
from streamcorpus_pipeline.stages import Configured
from streamcorpus_pipeline._taggers import TaggerBatchTransform

logger = logging.getLogger(__name__)

## map LingPipe's specific strings to streamcorpus.EntityType
_ENTITY_TYPES = {
    'PERSON': EntityType.PER,
    'FEMALE_PRONOUN': EntityType.PER,
    'MALE_PRONOUN': EntityType.PER,
    'ORGANIZATION': EntityType.ORG,
    'LOCATION': EntityType.LOC,
    }

_PRONOUNS = {
    'FEMALE_PRONOUN': Gender.FEMALE,
    'MALE_PRONOUN': Gender.MALE,
    }

filename_re = re.compile('''.*?<FILENAME docid="(?P<stream_id>.*?)">(?P<tagged_doc>(.|\n)*?)</FILENAME>''')
sentence_re = re.compile('''(?P<before>(.|\n)*?)<s i="(?P<sentence_id>.*?)">(?P<tagged_sentence>(.|\n)*?)</s>''')
ner_re = re.compile('''(?P<before>(.|\n)*?)(<ENAMEX ID="(?P<chain_id>.*?)" TYPE="(?P<entity_type>.*?)">(?P<entity_string>(.|\n)*?)</ENAMEX>)?''')

## detect any kind of whitespace, including unicode whitespace, which
## should include \u200b, even though python2.7 apparently lost this
## knowledge:  http://bugs.python.org/issue10567
only_whitespace = re.compile(u'''^(\s|\n|\u200b)*$''', re.UNICODE)

def files(text):
    '''
    Iterate over <FILENAME> XML-like tags and tokenize with nltk
    '''
    for f_match in filename_re.finditer(text):
        yield f_match.group('stream_id'), f_match.group('tagged_doc')

class LingPipeParser(object):
    def __init__(self, config):
        self.config = config
        self.clear()

    def clear(self):
        self.tok_num = 0
        self.byte_idx = 0
        self.line_idx = 0
        self.word_tokenizer = WhitespaceTokenizer()

    def set(self, ner_dom):
        self.clear()
        ## nltk wants a unicode string, so decode, it and then we will
        ## re-encode it to carefully recover the byte offsets.  We
        ## must take care not to use any nltk components that insert
        ## new whitespace, such
        ## nltk.tokenize.treebank.TreebankTokenizer
        self.ner_dom = ner_dom
        self.attributes = []
        self.relations = []

    def sentences(self):
        '''
        Iterate over <s> XML-like tags and tokenize with nltk
        '''
        for sentence_id, node in enumerate(self.ner_dom.childNodes):
            ## increment the char index with any text before the <s>
            ## tag.  Crucial assumption here is that the LingPipe XML
            ## tags are inserted into the original byte array without
            ## modifying the portions that are not inside the
            ## LingPipe-added tags themselves.
            if node.nodeType == node.TEXT_NODE:
                ## we expect to only see TEXT_NODE instances with whitespace
                assert only_whitespace.match(node.data), repr(node.data)

                ## must convert back to utf-8 to have expected byte offsets
                self.byte_idx += len(node.data.encode('utf-8'))

                ## count full lines, i.e. only those that end with a \n
                # 'True' here means keep the trailing newlines
                for line in node.data.splitlines(True):
                    if line.endswith('\n'):
                        self.line_idx += 1
            else:
                logger.debug('getting tokens for sentence_id=%d' % sentence_id)
                more_sentence_remains = True
                while more_sentence_remains:
                    ## always a sentence
                    sent = Sentence()

                    ## this "node" came from for loop above, and it's
                    ## childNodes list might have been popped by a
                    ## previous pass through this while loop
                    tokens = iter( self.tokens( node ) )

                    while 1:
                        try:
                            tok = tokens.next()
                            sent.tokens.append(tok)
                            #logger.debug('got token: %r  %d %d' % (tok.token, tok.mention_id, tok.sentence_pos))

                        except StopIteration:
                            yield sent
                            more_sentence_remains = False
                            break

    def _make_token(self, start, end):
        '''
        Instantiates a Token from self._input_string[start:end]
        '''
        ## all thfift strings must be encoded first
        tok_string = self._input_string[start:end].encode('utf-8')
        if only_whitespace.match(tok_string):
            ## drop any tokens with only whitespace
            return None
        tok = Token()
        tok.token = tok_string
        tok.token_num = self.tok_num
        if 'BYTES' in self.config['offset_types']:
            tok.offsets[OffsetType.BYTES] = Offset(
                type =  OffsetType.BYTES,
                first=self.byte_idx + len(self._input_string[:start].encode('utf-8')),
                length=len(tok_string),
                value=self.config['offset_debugging'] and tok_string or None,
                )
        if 'LINES' in self.config['offset_types']:
            tok.offsets[OffsetType.LINES] = Offset(
                type =  OffsetType.LINES,
                first=self.line_idx,
                length=1,
                value=self.config['offset_debugging'] and tok_string or None,
                )
        self.tok_num += 1
        ## keep track of position within a sentence
        tok.sentence_pos = self.sent_pos
        self.sent_pos += 1
        return tok

    def tokens(self, sentence_dom):
        '''
        Tokenize all the words and preserve NER labels from ENAMEX tags
        '''
        ## keep track of sentence position, which is reset for each
        ## sentence, and used above in _make_token
        self.sent_pos = 0

        ## keep track of mention_id, so we can distinguish adjacent
        ## multi-token mentions within the same coref chain
        mention_id = 0

        while len(sentence_dom.childNodes) > 0:
            ## shrink the sentence_dom's child nodes.  In v0_2_0 this
            ## was required to cope with HitMaxi16.  Now it is just to
            ## save memory.
            node = sentence_dom.childNodes.pop(0)

            if node.nodeType == node.TEXT_NODE:
                ## process portion before an ENAMEX tag
                for line in node.data.splitlines(True):
                    self._input_string = line
                    for start, end in self.word_tokenizer.span_tokenize(line):
                        tok = self._make_token(start, end)
                        if tok:
                            yield tok

                    if line.endswith('\n'):
                        ## maintain the index to the current line
                        self.line_idx += 1

                    ## increment index pasat the 'before' portion
                    self.byte_idx += len(line.encode('utf-8'))

            else:
                ## process text inside an ENAMEX tag
                assert node.nodeName == 'ENAMEX', node.nodeName
                chain_id = node.attributes.get('ID').value
                entity_type = node.attributes.get('TYPE').value
                for node in node.childNodes:
                    assert node.nodeType == node.TEXT_NODE, node.nodeType
                    for line in node.data.splitlines(True):
                        self._input_string = line
                        for start, end in self.word_tokenizer.span_tokenize(line):
                            tok = self._make_token(start, end)
                            if tok:
                                if entity_type in _PRONOUNS:
                                    tok.mention_type = MentionType.PRO
                                    tok.entity_type = _ENTITY_TYPES[entity_type]

                                    ## create an attribute
                                    attr = Attribute(
                                        attribute_type=AttributeType.PER_GENDER,
                                        value=str(_PRONOUNS[entity_type])
                                        )
                                    self.attributes.append(attr)

                                else:
                                    ## regular entity_type
                                    tok.mention_type = MentionType.NAME
                                    tok.entity_type = _ENTITY_TYPES[entity_type]

                                tok.equiv_id = int(chain_id)
                                tok.mention_id = mention_id
                                yield tok

                        if line.endswith('\n'):
                            ## maintain the index to the current line
                            self.line_idx += 1

                        ## increment index pasat the 'before' portion
                        self.byte_idx += len(line.encode('utf-8'))

                ## increment mention_id within this sentence
                mention_id += 1

class lingpipe(TaggerBatchTransform):
    '''
    a streamcorpus_pipeline batch transform that converts a chunk into a new
    chunk with data generated by LingPipe
    '''
    config_name = 'lingpipe'
    tagger_id = 'lingpipe'

    template = \
        '''cd %(tagger_root_path)s/demos/generic/bin && ''' + \
        '''cat %(clean_visible_path)s | ./cmd_coref_en_news_muc6.sh ''' + \
        ''' "-contentType=text/html" "-includeElts=FILENAME" ''' + \
        ''' %(java_heap_size)s 1> %(ner_xml_path)s'''

    def get_sentences(self, ner_dom):
        '''parse the sentences and tokens out of the XML'''
        lp_parser = LingPipeParser(self.config)
        lp_parser.set(ner_dom)
        sentences = list( lp_parser.sentences() )
        return sentences, lp_parser.relations, lp_parser.attributes

if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('action',
                        metavar='postproc|align',
                        help='postproc')
    parser.add_argument('input_file', help='XML file from LingPipe')
    parser.add_argument('output_file', help='XML file to generate with OWPL data')
    parser.add_argument('--source_chunk', help='source chunk file that was input to the pipeline data')
    args = parser.parse_args()

    if args.action == 'postproc':
        text = open(args.input_file).read()
        print 'read %d bytes from %s' % (len(text), args.input_file)

        raise NotImplementedError('need to instantiate a LingPipeParser object here')

        for stream_id, tagged_doc in files(text):
            for sent in sentences(tagged_doc):  # pylint: disable=E0602
                for tok in sent.tokens:
                    if tok.entity_type is not None:
                        print tok, EntityType._VALUES_TO_NAMES[tok.entity_type]

    elif args.action == 'align':
        i_chunk = Chunk(path=args.source_chunk, mode='rb')
        o_chunk = Chunk(path=args.output_file,  mode='wb')
        align_chunk_with_ner(args.input_file, i_chunk, o_chunk)  # pylint: disable=E0602


#!/usr/bin/env python

import os
import sys
import fnmatch
import subprocess

## prepare to run PyTest as a command
from distutils.core import Command

from setuptools import setup, find_packages

from version import get_git_version
VERSION, SOURCE_LABEL = get_git_version()
PROJECT = 'streamcorpus_pipeline'
AUTHOR = 'Diffeo, Inc.'
AUTHOR_EMAIL = 'support@diffeo.com'
URL = 'http://github.com/trec-kba/streamcorpus-pipeline'
DESC = 'Tools for building streamcorpus objects, such as those used in TREC.'

def read_file(file_name):
    file_path = os.path.join(
        os.path.dirname(__file__),
        file_name
        )
    return open(file_path).read()

def recursive_glob(treeroot, pattern):
    results = []
    for base, dirs, files in os.walk(treeroot):
        goodfiles = fnmatch.filter(files, pattern)
        results.extend(os.path.join(base, f) for f in goodfiles)
    return results

def recursive_glob_with_tree(new_base, old_base, treeroot, pattern):
    '''generate a list of tuples(new_base, list(paths to put there)
    where the files are found inside of old_base/treeroot.
    '''
    results = []
    old_cwd = os.getcwd()
    os.chdir(old_base)
    for rel_base, dirs, files in os.walk(treeroot):
        goodfiles = fnmatch.filter(files, pattern)
        one_dir_results = []
        for f in goodfiles:
            one_dir_results.append(os.path.join(old_base, rel_base, f))
        results.append((os.path.join(new_base, rel_base), one_dir_results))
    os.chdir(old_cwd)
    return results

def _myinstall(pkgspec):
    subprocess.check_call(['pip', 'install', pkgspec])

class PyTest(Command):
    '''run py.test'''

    description = 'runs py.test to execute all tests'

    user_options = []

    def initialize_options(self):
        pass

    def finalize_options(self):
        pass

    def run(self):
        if self.distribution.install_requires:
            for ir in self.distribution.install_requires:
                _myinstall(ir)
        if self.distribution.tests_require:
            for ir in self.distribution.tests_require:
                _myinstall(ir)

        # reload sys.path for any new libraries installed
        import site
        site.main()
        print sys.path
        # use pytest to run tests
        pytest = __import__('pytest')
        if pytest.main(['-n', '8', '-vvs', 'streamcorpus_pipeline']):
            sys.exit(1)

setup(
    name=PROJECT,
    version=VERSION,
    description=DESC,
    license='MIT/X11 license http://opensource.org/licenses/MIT',
    long_description=read_file('README.md'),
    #source_label=SOURCE_LABEL,
    author=AUTHOR,
    author_email=AUTHOR_EMAIL,
    url=URL,
    packages = find_packages(),
    cmdclass={'test': PyTest,
    },
    classifiers=[
        'Development Status :: 3 - Alpha',
        'Topic :: Utilities',
        'License :: OSI Approved :: MIT License',  ## MIT/X11 license http://opensource.org/licenses/MIT
    ],
    tests_require=[
        'pytest',
        'ipdb',
        'pytest-cov',
        'pytest-xdist',
        'pytest-timeout',
        'pytest-capturelog',

        ## this does not work in python2.6
        'pytest-incremental',

        'epydoc',
    ],
    install_requires=[
        ## these are used by streamcorpus_pipeline_test
        'pytest-diffeo >= 0.1.7',

        'dblogger >= 0.4.0',
        'yakonfig >= 0.4.2',
        'importlib',
        'python-magic',
        'python-docx',
        'ftfy < 5',  # 5 and higher only support Python 3
        'kvlayer >= 0.4.10',
        'coordinate',
        'protobuf < 3',
        'requests',
        'streamcorpus>=0.3.48',
        'pyyaml',
        'nltk',
        'lxml',
        'BeautifulSoup',
        'html5lib>=0.999999999',
        'boto',
        'jellyfish',
        'nilsimsa >= 0.3',
        'regex != 2014.08.28',
        'sortedcollection',
        'python-docx',
        'pdfminer',
        'backports.lzma!=0.0.4',
        'backport_collections', # for python2.6
    ],
    extras_require = {
        'keyword_indexing': [
            'many_stop_words >= 0.2.0',
            'mmh3',
        ],
    },
    entry_points={
        'console_scripts': [
            'streamcorpus_pipeline = streamcorpus_pipeline.run:main',
            'streamcorpus_directory = streamcorpus_pipeline._directory:main',
            'streamcorpus_pipeline_test = streamcorpus_pipeline.tests.run:main',
            'streamcorpus_pipeline_keywords = streamcorpus_pipeline.keywords:main',
            'scp-xpaths = streamcorpus_pipeline.run_xpaths:main',
        ]
    },
    data_files = [
        ('docs/examples/streamcorpus-pipeline/', recursive_glob('examples', '*.py')),
        ('docs/examples/streamcorpus-pipeline/', recursive_glob('examples', '*.yaml')),
        ('data/streamcorpus-pipeline/john-smith/', [
                'data/john-smith/john-smith-tagged-by-lingpipe-serif-0-197.sc.xz',
                'data/john-smith/ground-truth.yaml',
                ]),
    ] + recursive_glob_with_tree('data/streamcorpus-pipeline', 'data', 'john-smith/original', '*') \
      + recursive_glob_with_tree('data/streamcorpus-pipeline', 'data', 'test', '*'),
    include_package_data = True,
    zip_safe = False
)


#!/usr/bin/env python
'''
Provides classes for loading chunk files from local storage and
putting them out into local storage.

.. This software is released under an MIT/X11 open source license.
   Copyright 2012-2016 Diffeo, Inc.
'''
from __future__ import absolute_import, division, print_function
from cStringIO import StringIO
from functools import partial
import gzip
import hashlib
import logging
import os
import re
import sys
import time

import requests

try:
    from dossier.fc import FeatureCollectionChunk as FCChunk
except ImportError:
    FCChunk = None

import streamcorpus
from streamcorpus import decrypt_and_uncompress, \
    parse_file_extensions, known_compression_schemes, \
    compress_and_encrypt_path, Chunk
from streamcorpus_pipeline._exceptions import FailedExtraction, \
    FailedVerification, ConfigurationError
from streamcorpus_pipeline._get_name_info import get_name_info
from streamcorpus_pipeline._spinn3r_feed_storage import _generate_stream_items
from streamcorpus_pipeline.stages import Configured
from streamcorpus_pipeline._tarball_export import tarball_export

logger = logging.getLogger(__name__)

_message_versions = {
    'v0_1_0': streamcorpus.StreamItem_v0_1_0,
    'v0_2_0': streamcorpus.StreamItem_v0_2_0,
    'v0_3_0': streamcorpus.StreamItem_v0_3_0,
}

import boto
from boto.s3.key import Key
from boto.s3.connection import S3Connection
## stop Boto's built-in retries, so we can do our own
if not boto.config.has_section('Boto'):
    boto.config.add_section('Boto')
boto.config.set('Boto', 'num_retries', '0')

def _retry(func):
    '''
    Decorator for methods that need many retries, because of
    intermittent failures, such as AWS calls via boto, which has a
    non-back-off retry.
    '''
    def retry_func(self, *args, **kwargs):
        tries = 1
        while True:
            # If a handler allows execution to continue, then
            # fall through and do a back-off retry.
            try:
                return func(self, *args, **kwargs)
                break
            except OSError as exc:
                ## OSError: [Errno 24] Too many open files
                logger.error('assuming OSError unrecoverable')
                raise
            except FailedExtraction as exc:
                ## pass through exc to caller
                logger.error('FAIL(%d)', tries, exc_info=True)
                raise
            except FailedVerification as exc:
                logger.warn('FAIL(%d)', tries, exc_info=True)
                if tries >= self.config['tries']:
                    if self.config.get('suppress_failures'):
                        logger.warn('suppressing failure and breaking out of this loop; data may be corrupt, downstream will have to cope')
                        break
                    else:
                        raise
            except Exception as exc:
                logger.warn('FAIL(%d): having I/O trouble with S3', tries, exc_info=True)
                if tries >= self.config['tries']:
                    raise

            logger.warn('RETRYING (%d left)', self.config['tries'] - tries)
            time.sleep(3 * tries)
            tries += 1

    return retry_func

def timedop(what, datalen, fun):
    start_time = time.time()
    try:
        retval = fun()
    except:
        elapsed = time.time() - start_time
        if elapsed  > 0:
            logger.debug('%s %.1f bytes/second ERROR', what, float(datalen) / float(elapsed), exc_info=True)
        raise

    elapsed = time.time() - start_time
    if elapsed  > 0:
        logger.debug('%s %.1f bytes/second', what, float(datalen) / float(elapsed))
    return retval

def verify_md5(md5_expected, data, other_errors=None):
    "return True if okay, raise Exception if not"  # O_o ?
    md5_recv = hashlib.md5(data).hexdigest()
    if md5_expected != md5_recv:
        if other_errors is not None:
            logger.critical('\n'.join(other_errors))
        raise FailedVerification('original md5 = %r != %r = received md5' \
                                 % (md5_expected, md5_recv))
    return True

def get_bucket(config, bucket_name=None):
    '''This function is mostly about managing configuration, and then
    finally returns a boto.Bucket object.

    AWS credentials come first from config keys
    aws_access_key_id_path, aws_secret_access_key_path (paths to one
    line files); secondly from environment variables
    AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY; also from $HOME/.aws/credentials
    or the magic Amazon http://169.254.169.254/ service.  If credentials
    are not set in the config then behavior is the same as other AWS-based
    command-line tools.

    '''
    if not bucket_name:
        if 'bucket' not in config:
            raise ConfigurationError(
                'The "bucket" parameter is required for the s3 stages.')
        bucket_name = config['bucket']

    # get AWS credentials. first, from config; else, from env vars.
    # (boto will read environment variables and other normal places.)
    aws_access_key_id_path = config.get('aws_access_key_id_path')
    aws_secret_access_key_path = config.get('aws_secret_access_key_path')

    params = ()
    if aws_access_key_id_path and aws_secret_access_key_path:
        try:
            access = open(aws_access_key_id_path).read().strip()
            secret = open(aws_secret_access_key_path).read().strip()
            params = (access, secret)
        except:
            logger.error('failed reading aws credentials from configured file', exc_info=True)
            raise

    conn = S3Connection(*params)
    bucket = conn.get_bucket(bucket_name)
    return bucket

class from_s3_chunks(Configured):
    '''
    Reads data from Amazon S3 one key at a time. The type of data read
    can be specified with the ``input_format`` config option. The
    following values are legal: ``streamitem``, ``featurecollection``
    or ``spinn3r``.

    When the input format is ``streamitem`` or ``spinn3r``, then this
    reader produces a generator of :class:`streamcorpus.StreamItem`
    instances.

    When the input format is ``featurecollection``, then this reader
    produces a generator of ``dossier.fc.FeatureCollection`` instances.

    ``bucket`` is the s3 bucket to use if input paths are not full
    s3://{bucket}{path} URIs. ``aws_access_key_id_path`` and
    ``aws_secret_access_key_path`` should point to files containing
    your s3 credentials.  Alternately credentials can be in environment
    variables ``AWS_ACCESS_KEY_ID`` and ``AWS_SECRET_ACCESS_KEY`` or
    provided by other standard means.

    The rest of the configuration options are optional and are described
    in the following example:

    .. code-block:: yaml

        from_s3_chunks:
          # How to connect to S3
          # bucket can also come from input path if full uri s3://{bucket}{path}
          bucket: aws-publicdatasets
          # can also come from environment variable AWS_ACCESS_KEY_ID
          aws_access_key_id_path: keys/aws_access_key_id
          # can also come from environment variable AWS_SECRET_ACCESS_KEY
          aws_secret_access_key_path: keys/aws_secret_access_key

          # Optional parameters.

          # The number of times to try reading from s3. A value of
          # `1` means the download is tried exactly once.
          # The default value is `10`.
          tries: 1

          # This is prepended to every key given. Namely, all s3 URLs
          # are formed by `s3://{bucket}/{s3_path_prefix}/{input-key-name}`.
          # By default, it is empty.
          s3_path_prefix: some/s3/prefix

          # A path to a private GPG decryption key file.
          gpg_decryption_key_path: keys/gpg-key

          # The type of data to read from s3. Valid values are
          # "StreamItem", "spinn3r" or "FeatureCollection".
          # The default is "StreamItem".
          input_format: StreamItem

          # When the input format is "StreamItem", this indicates the
          # Thrift version to use. Defaults to "v0_3_0".
          streamcorpus_version: v0_3_0

          # When set, an md5 is extracted from the s3 key and is used
          # to verify the decrypted and decompressed content downloaded.
          # This is disabled by default.
          compare_md5_in_file_name: true

          # A temporary directory where intermediate files may reside.
          # Uses your system's default tmp directory (usually `/tmp`)
          # by default.
          tmp_dir_path: /tmp
    '''
    config_name = 'from_s3_chunks'
    default_config = {
        'compare_md5_in_file_name': False,
        's3_path_prefix': '',
        'tmp_dir_path': None,
        'tries': 10,
        'gpg_decryption_key_path': None,
        'input_format': 'StreamItem',
        'streamcorpus_version': 'v0_3_0',
    }

    def __init__(self, config):
        super(from_s3_chunks, self).__init__(config)
        self.gpg_decryption_key_path = config.get('gpg_decryption_key_path')

    def __call__(self, i_str):
        '''
        Takes a path suffix as a string over stdin and generates chunks from
        s3://<bucket><s3_prefix_path>/{i_str}.
        '''
        if i_str.startswith('s3://'):
            # full path including bucket name
            bucket_name, kpath = i_str[5:].split('/', 1)
            assert bucket_name and kpath, 'bad s3 url {0!r}'.format(i_str)
        else:
            # get bucket_name from config
            bucket_name = None
            kpath = os.path.join(self.config['s3_path_prefix'].strip(),
                                 i_str.strip())
        logger.info('from_s3_chunks: %s / %r', bucket_name, kpath)
        return self.get_chunk(bucket_name, kpath)

    def _decode(self, data):
        '''
        Given the raw data from s3, return a generator for the items
        contained in that data. A generator is necessary to support
        chunk files, but non-chunk files can be provided by a generator
        that yields exactly one item.

        Decoding works by case analysis on the config option
        ``input_format``. If an invalid ``input_format`` is given, then
        a ``ConfigurationError`` is raised.
        '''
        informat = self.config['input_format'].lower()
        if informat == 'spinn3r':
            return _generate_stream_items(data)
        elif informat == 'streamitem':
            ver = self.config['streamcorpus_version']
            if ver not in _message_versions:
                raise ConfigurationError(
                    'Not a valid streamcorpus version: %s '
                    '(choose from: %s)'
                    % (ver, ', '.join(_message_versions.keys())))

            message = _message_versions[ver]
            return streamcorpus.Chunk(data=data, message=message)
        elif informat == 'featurecollection' and FCChunk is not None:
            return FCChunk(data=data)
        else:
            raise ConfigurationError(
                'from_s3_chunks unknown input_format = %r'
                % informat)

    @_retry
    def get_chunk(self, bucket_name, key_path):
        '''return Chunk object full of records
        bucket_name may be None'''
        bucket = get_bucket(self.config, bucket_name=bucket_name)
        key = bucket.get_key(key_path)
        if key is None:
            raise FailedExtraction('Key "%s" does not exist.' % key_path)

        fh = StringIO()
        key.get_contents_to_file(fh)
        data = fh.getvalue()
        if not data:
            raise FailedExtraction('%s: no data (does the key exist?)'
                                   % key.key)

        chunk_type, compression, encryption = parse_file_extensions(key_path)
        if encryption == 'gpg':
            if not self.gpg_decryption_key_path:
                raise FailedExtraction('%s ends with ".gpg" but gpg_decryption_key_path=%s'
                                       % (key.key, self.gpg_decryption_key_path))

        _errors = []
        if compression or encryption:
            _errors, data = decrypt_and_uncompress(
                data,
                self.gpg_decryption_key_path,
                tmp_dir=self.config.get('tmp_dir_path'),
                compression=compression,
                )
            if not data:
                msg = 'decrypt_and_uncompress got no data for {0!r}, from {1} bytes' \
                    + ' downloaded, errors: {2}' \
                        .format(key_path, len(data), '\n'.join(_errors))
                logger.error(msg)
                raise FailedExtraction(msg)
            logger.info( '\n'.join(_errors) )

        if not self.config['compare_md5_in_file_name']:
            logger.warn('not checking md5 in file name, consider setting '
                        'from_s3_chunks:compare_md5_in_file_name')
        else:
            logger.info('Verifying md5 for "%s"...' % key.key)

            # The regex hammer.
            m = re.search('([a-z0-9]{32})(?:\.|$)', key.key)
            if m is None:
                raise FailedExtraction(
                    'Could not extract md5 from key "%s". '
                    'Perhaps you should disable compare_md5_in_file_name?'
                    % key.key)

            i_content_md5 = m.group(1)
            #import pdb; pdb.set_trace()
            verify_md5(i_content_md5, data, other_errors=_errors)
        return self._decode(data)

class to_s3_chunks(Configured):
    '''Copies chunk files on disk to Amazon S3. The type of data written
    can be specified with the ``output_format`` config option. The
    following values are legal: ``streamitem`` and ``featurecollection``.

    N.B. The format is only necessary for construction the
    ``output_name``. The format is also used to pick between an ``fc``
    (for feature collections) extension and a ``sc`` (for stream items)
    extension.

    The following configuration options are mandatory:
    ``bucket`` is the s3 bucket to use. ``aws_access_key_id_path``
    and ``aws_secret_access_key_path`` should point to files
    containing your s3 credentials.

    ``output_name`` is also required and is expanded in the same way as
    the :class:`~streamcorpus_pipeline._local_storage.to_local_chunks`
    writer.  The filename always has ``.{sc,fc}.{xz,gz,sz}`` appended to it
    (depending on the output format specified), and correspondingly,
    the output file is always compressed.  If GPG keys are provided,
    then ``.gpg`` is appended and the file is encrypted.

    The rest of the configuration options are optional and are
    described in the following example:

    .. code-block:: yaml

        to_s3_chunks:
          # Mandatory
          bucket: aws-publicdatasets
          aws_access_key_id_path: keys/aws_access_key_id
          aws_secret_access_key_path: keys/aws_secret_access_key
          output_name: "%(date_hour)s/%(source)s-%(num)d-%(input_fname)s-%(md5)s"

          # Optional parameters.

          # The number of times to try writing to s3. A value of
          # `1` means the upload is tried exactly once.
          # The default value is `10`.
          # (This also applies to the verification step.)
          tries: 1

          # When set, the file uploaded will be private.
          # Default: false
          is_private: false

          # When set, the file will be re-downloaded from Amazon, decrypted,
          # decompressed and checksummed against the data sent to Amazon.
          # (This used to be "verify_via_http", but this more broadly named
          # option applies even when "is_private" is true.)
          #
          # Default: true
          verify: true

          # If verification fails `tries` times, then the default
          # behavior is to exit, which can cause a coordinate
          # fork_worker parent to retry the whole job.
          #
          # Default: false
          suppress_failures: false

          # This is prepended to every key given. Namely, all s3 URLs
          # are formed by `s3://{bucket}/{s3_path_prefix}/{input-key-name}`.
          # By default, it is empty.
          s3_path_prefix: some/s3/prefix

          # Paths to GPG keys. Note that if you're using verification,
          # then a decryption key must be given.
          # Default values: None
          gpg_encryption_key_path: keys/gpg-key.pub
          gpg_decryption_key_path: keys/gpg-key.private

          # GPG recipient.
          # Default value: trec-kba
          gpg_recipient: trec-kba

          # Removes the intermediate chunk file from disk.
          # Default: true.
          cleanup_tmp_files: true

          # A temporary directory where intermediate files may reside.
          # Uses your system's default tmp directory (usually `/tmp`)
          # by default.
          tmp_dir_path: /tmp

          # Change compression scheme; default is .xz for greatest
          # compression of archival content (S3 charges by the byte).
          # xz is also the slowest, so other options can make more
          # sense in some applications.  Snappy (.sz) is the fastest
          # and still much better than no compression at all.  Choices
          # are "xz", "sz", "gz", ""
          compression: xz

    '''
    config_name = 'to_s3_chunks'
    default_config = {
        's3_path_prefix': '',
        'tries': 10,
        'gpg_encryption_key_path': None,
        'gpg_decryption_key_path': None,
        'gpg_recipient': 'trec-kba',
        'verify': True,
        'suppress_failures': False,
        'is_private': False,
        'output_format': 'StreamItem',
        'tmp_dir_path': '/tmp',
        'cleanup_tmp_files': True,
        'compression': 'xz',
        # require: bucket, output_name, aws_access_key_id_path,
        #          aws_secret_access_key_path
    }

    def __init__(self, config):
        super(to_s3_chunks, self).__init__(config)

        ## TODO: use something like verify_config
        if self.compression not in known_compression_schemes:
            raise Exception('to_s3_chunks "compression: %r" not in %r' % (
                self.compression, known_compression_schemes))

        logger.critical('compression: %s' % self.config.get('compression'))

        # Backwards compatibility.
        if 'verify_via_http' in self.config:
            logger.warning('Update your config! Use "verify" instead of '
                           '"verify_via_http". The latter is deprecated.')
            self.config['verify'] = self.config.pop('verify_via_http')
            if self.config['verify'] and self.config['is_private']:
                logger.warning('Nonsensical config "verify_via_http=true" and '
                               '"is_private=true". Will verify with boto.')
        self.bucket = get_bucket(self.config)

    def __call__(self, t_path, name_info, i_str):
        '''
        Load chunk from t_path and put it into the right place in s3
        using the output_name template from the config
        '''
        logger.info('to_s3_chunks.call(t_path=%r, name_info=%r, i_str=%r)',
                    t_path, name_info, i_str)
        # Getting name info actually assembles an entire chunk in memory
        # from `t_path`, so we now need to tell it which chunk type to use.
        try:
            more_name_info = get_name_info(t_path, i_str=i_str,
                                           chunk_type=self.chunk_type)
        except Exception, exc:
            logger.critical('failed get_name_info(%r, %r', t_path, i_str, exc_info=True)
            raise

        self.name_info = dict(name_info, **more_name_info)

        if self.name_info['num'] == 0:
            return None

        o_path = self.s3key_name
        logger.info('%s: \n\t%r\n\tfrom: %r\n\tby way of %r',
                    self.__class__.__name__, o_path, i_str, t_path)

        t_path2 = self.prepare_on_disk(t_path)
        data_len = os.path.getsize(t_path2)
        if data_len == 0:
            logger.critical('data is now zero bytes!')
        logger.debug('prepared %s bytes of %r', data_len, t_path2)
        self.put_data(o_path, t_path2, self.name_info['md5'], data_len)

        self.cleanup(t_path, t_path2)
        logger.info('%s finished:\n\t input: %s\n\toutput: %s',
                    self.__class__.__name__, i_str, o_path)
        return [o_path]

    @property
    def outfmt(self):
        return self.config['output_format'].lower()

    @property
    def compression(self):
        return self.config['compression'].lower()

    @property
    def chunk_type(self):
        if self.outfmt == 'featurecollection' and FCChunk is not None:
            if self.config.get('verify'):
                return partial(FCChunk, inline_md5=True)
            return FCChunk
        elif self.outfmt == 'streamitem':
            return streamcorpus.Chunk
        else:
            raise ConfigurationError(
                'Invalid output format: "%s". Choose one of StreamItem '
                'or FeatureCollection.' % self.config['output_format'])

    @property
    def s3key_name(self):
        o_fname = self.config['output_name'] % self.name_info
        ext = 'fc' if self.outfmt == 'featurecollection' else 'sc'
        o_path = os.path.join(self.config['s3_path_prefix'],
                              '%s.%s' % (o_fname, ext))
        if self.compression:
            o_path += '.' + self.compression
        if self.config.get('gpg_encryption_key_path') is not None:
            o_path += '.gpg'
        self.name_info['s3_output_path'] = o_path
        return o_path

    def prepare_on_disk(self, t_path):
        logger.debug('gpg_encryption_key_path: %r', self.config.get('gpg_encryption_key_path'))
        _errors, t_path2 = compress_and_encrypt_path(
            t_path,
            self.config.get('gpg_encryption_key_path'),
            gpg_recipient=self.config['gpg_recipient'],
            tmp_dir=self.config.get('tmp_dir_path'),
            compression=self.config.get('compression'),
        )
        if len(_errors) > 0:
            logger.error('compress and encrypt errors: %r', _errors)
        return t_path2

    def cleanup(self, *files):
        if not self.config['cleanup_tmp_files']:
            return
        for f in files:
            try:
                os.remove(f)
            except Exception as exc:
                logger.info('%s --> failed to remove %s' % (exc, f))

    @_retry
    def put_data(self, key_path, t_path, md5, data_len=None):
        timedop('s3 put', data_len, lambda: self.put(key_path, t_path))
        if self.config['verify']:
            ok = timedop('s3 verify', data_len, lambda: self.verify(key_path, md5))
            if not ok:
                raise Exception('verify failed putting {0!r}'.format(key_path))

    def put(self, o_path, t_path):
        key = Key(self.bucket, o_path)
        key.set_contents_from_filename(t_path)

        if not self.config.get('is_private', False):
            # Makes the file have a public URL.
            key.set_acl('public-read')

    @_retry
    def verify(self, o_path, md5):
        chunk_format, compression, encryption = parse_file_extensions(o_path)
        logger.info('verifying %s --> %r, %r, %r', o_path, chunk_format, compression, encryption)
        if self.config.get('gpg_encryption_key_path') \
                and not self.config.get('gpg_decryption_key_path'):
            raise ConfigurationError(
                'When "verify=true" and "gpg_encryption_key_path" is set, '
                '"gpg_decryption_key_path" must also be set.')

        if self.config.get('is_private', False):
            rawdata = self.private_data(o_path)
        else:
            rawdata = self.public_data(o_path)

        if not rawdata:
            logger.error('got no data out of reading the data')

        errors, data = decrypt_and_uncompress(
            rawdata,
            self.config.get('gpg_decryption_key_path'),
            tmp_dir=self.config.get('tmp_dir_path'),
            compression=compression,
        )
        if not data:
            logger.error('got no data back from decrypt_and_uncompress %r, (size=%r), errors: %r', o_path, len(rawdata), errors)
            return False

        ### Let's not use both belt and suspenders.  md5 is enough.
        ### Of note, if this gets killed by something, e.g. coordinate
        ### fork_worker parent, then since this is one of the slowest
        ### parts, you might see cbor getting killed, which can look
        ### like this:
        ##    >>> c = list(f)
        ##    ^CTraceback (most recent call last):
        ##    File "<stdin>", line 1, in <module>
        ##    File "/ebs/home/jrf/streamcorpus/py/src/streamcorpus/_chunk.py", line 381, in __iter__
        ##    for msg in self.read_msg_impl():
        ##    File "/ebs/home/jrf/streamcorpus/py/src/streamcorpus/_cbor_chunk.py", line 111, in read_msg_impl
        ##    ob = cbor.load(self._i_chunk_fh)
        ##    RuntimeError: unknown error decoding TAG...
        #
        #try:
        #    count = len(list(self.chunk_type(data=data)))
        #except Exception, exc:
        #    count = None
        #    logger.critical('\n\n********\n\nfailure on %r\n\n********\n\n', o_path, exc_info=True)
        #logger.info('attempting verify of %r %r in %r', count, chunk_format, o_path)
        logger.info('attempting verify of %r in %r', chunk_format, o_path)
        return verify_md5(md5, data, other_errors=errors)

    def public_data(self, o_path):
        url = 'http://s3.amazonaws.com/%(bucket)s/%(o_path)s' % {
            'bucket': self.config['bucket'],
            'o_path': o_path,
        }
        logger.info('public fetching %r', url)
        return requests.get(url).content

    def private_data(self, o_path):
        return self.bucket.get_key(o_path).get_contents_as_string()

class to_s3_tarballs(to_s3_chunks):
    '''
    The same as :class:`streamcorpus_pipeline._s3_storage.to_s3_chunks`,
    except it puts stream items into a gzipped tarball instead of
    chunks. This writer does not do any encryption.

    In addition to the required parameters for ``to_s3_chunks``, this
    writer has three more required parameters: ``tarinfo_name`` (which
    supports ``output_name`` substitution semantics), ``tarinfo_uname``
    and ``tarinfo_gname``.
    '''
    config_name = 'to_s3_tarballs'

    def __init__(self, config):
        super(to_s3_tarballs, self).__init__(config)
        if self.config['output_format'].lower() != 'streamitem':
            raise ConfigurationError(
                'to_s3_tarballs only supports "output_format=streamitem", '
                'but given "output_format=%s"' % self.config['output_format'])

    @property
    def s3key_name(self):
        o_fname = self.config['output_name'] % self.name_info
        return os.path.join(self.config['s3_path_prefix'], o_fname + '.tar.gz')

    def prepare_on_disk(self, t_path):
        t_path2 = tarball_export(self.config, t_path, self.name_info)
        # Cheat a bit here. We are checking the md5 of the full compressed
        # archive instead of the decompressed/decrypted chunk (because this is
        # a tarball, not a chunk).
        # TODO: if md5 is actually needed, read the file and calculate it
        #self.name_info['md5'] = hashlib.md5(data).hexdigest()
        return t_path2

    @_retry
    def redownload_verify(self, o_path, md5):
        key = Key(get_bucket(self.config), o_path)
        data = key.get_contents_as_string()
        logger.info( 'got back SIs: %d' % len( list( Chunk(data=data) ) ))
        return verify_md5(md5, data)

    @_retry
    def verify(self, o_path, md5):
        if self.config.get('is_private', False):
            rawdata = self.private_data(o_path)
        else:
            rawdata = self.public_data(o_path)
        return verify_md5(md5, rawdata)


'''table names used by `{from,to}_kvlayer` for `streamcorpus_pipeline`

.. This software is released under an MIT/X11 open source license.
   Copyright 2012-2015 Diffeo, Inc.

This module defines :mod:`kvlayer` table names.  Other modules that
need to access these tables should call:

.. code-block:: python

   import kvlayer
   from streamcorpus_pipeline._kvlayer_table_names import \
     STREAM_ITEM_TABLE_DEFS, STREAM_ITEM_VALUE_DEFS

   client = kvlayer.client()
   client.setup_namespace(STREAM_ITEM_TABLE_DEFS, STREAM_ITEM_VALUE_DEFS)

.. autofunction:: stream_id_to_kvlayer_key
.. autofunction:: kvlayer_key_to_stream_id
.. autofunction:: key_for_stream_item

.. autodata:: STREAM_ITEM_TABLE_DEFS
.. autodata:: STREAM_ITEM_VALUE_DEFS
.. autodata:: STREAM_ITEMS_TABLE
.. autodata:: STREAM_ITEMS_SOURCE_INDEX
.. autodata:: STREAM_ITEMS_TIME_INDEX
.. autodata:: HASH_TF_INDEX_TABLE
.. autodata:: HASH_FREQUENCY_TABLE
.. autodata:: HASH_KEYWORD_INDEX_TABLE
.. autodata:: INDEX_TABLE_NAMES
.. autodata:: WITH_SOURCE
.. autodata:: BY_TIME
.. autodata:: KEYWORDS
.. autodata:: HASH_TF_SID
.. autodata:: HASH_FREQUENCY
.. autodata:: HASH_KEYWORD

'''
from __future__ import absolute_import
import base64
import hashlib
import string

from kvlayer import COUNTER

# table names:

#: Name of the table holding stream items.  Keys are pairs of a
#: 16-byte binary string holding the document ID and an :class:`int`
#: holding the stream time.  Values are compressed serialized stream
#: items.
STREAM_ITEMS_TABLE = 'si'

#: Name of the table holding stream item source strings.  Keys are the
#: same as :data:`STREAM_ITEMS_TABLE`.  Values are simple strings
#: holding the stream item "source" string.  If the record is absent
#: for a given stream item, its source is :const:`None`.
STREAM_ITEMS_SOURCE_INDEX = 'sisi'

#: Name of the table holding stream item keys ordered by time.
#: Keys are the reverse of :data:`STREAM_ITEMS_TABLE`, pairs of
#: an :class:`int` holding the stream time and a 16-byte binary
#: string holding the document ID.  Values are empty.
STREAM_ITEMS_TIME_INDEX = 'siti'

#: Name of the table holding keyword index data.  Keys are tuples of
#: the Murmur hash of the keyword, the document ID, and the stream time;
#: the last two parts make up the stream ID.  Values are ints with the
#: number of times this keyword appears in the document.
HASH_TF_INDEX_TABLE = 'sitf'

#: Name of the table holding document frequency data for hash tokens.
#: Keys are the Murmur hash of the keyword.  Values are ints with the
#: number of documents this hash appears in.
HASH_FREQUENCY_TABLE = 'sihf'

#: Name of the table holding keyword index terms.  Keys are pairs of an
#: integer hash and the corresponding string.  Values are counters
#: holding the number of documents the terms appear in.
HASH_KEYWORD_INDEX_TABLE = 'sihk'

#: Dictionary mapping :mod:`kvlayer` table name to key tuple type,
#: suitable to be passed to
#: :meth:`kvlayer._abstract_storage.AbstractStorage.setup_namespace`
STREAM_ITEM_TABLE_DEFS = {
    # StreamItem primary key is:
    # (binary 16 byte md5 of abs_url, timestamp)
    STREAM_ITEMS_TABLE: (str, int),
    STREAM_ITEMS_SOURCE_INDEX: (str, int),
    STREAM_ITEMS_TIME_INDEX: (int, str),

    # keyword index tables:
    HASH_TF_INDEX_TABLE: (int, str, int),
    HASH_FREQUENCY_TABLE: (int,),
    HASH_KEYWORD_INDEX_TABLE: (int, str),

    # TODO: maybe reintroduce index-by-time
}

#: Dictionary mapping :mod:`kvlayer` table name to value type, suitable
#: to be passed to
#: :meth:`kvlayer._abstract_storage.AbstractStorage.setup_namespace`
STREAM_ITEM_VALUE_DEFS = {
    # Unspecified tables default to str value
    HASH_TF_INDEX_TABLE: int,
    HASH_FREQUENCY_TABLE: COUNTER,
    HASH_KEYWORD_INDEX_TABLE: COUNTER,
}

# index names, used for configuration, not kvlayer table names

#: :class:`~streamcorpus_pipeline._kvlayer.to_kvlayer` index
#: configuration to generate the all of the keyword index tables.
KEYWORDS = 'keywords'

#: :class:`~streamcorpus_pipeline._kvlayer.to_kvlayer` index
#: configuration to generate the document keyword index.
HASH_TF_SID = 'hash_tf_sid'

#: :class:`~streamcorpus_pipeline._kvlayer.to_kvlayer` index
#: configuration to record hash frequencies.
HASH_FREQUENCY = 'hash_frequency'

#: :class:`~streamcorpus_pipeline._kvlayer.to_kvlayer` index
#: configuration to generate the hash-to-keyword index.
HASH_KEYWORD = 'hash_keyword'

#: :class:`~streamcorpus_pipeline._kvlayer.to_kvlayer` index
#: configuration to generate the source-type index.
WITH_SOURCE = 'with_source'

#: :class:`~streamcorpus_pipeline._kvlayer.to_kvlayer` index
#: configuration to generate the timestamp index.
BY_TIME = 'by_time'

#: Dictionary mapping a configuration name from the
#: :class:`~streamcorpus_pipeline._kvlayer.to_kvlayer` writer to
#: its underlying table name.
INDEX_TABLE_NAMES = {
    WITH_SOURCE: STREAM_ITEMS_SOURCE_INDEX,
    BY_TIME: STREAM_ITEMS_TIME_INDEX,
    HASH_FREQUENCY: HASH_FREQUENCY_TABLE,
    HASH_KEYWORD: HASH_KEYWORD_INDEX_TABLE,
    HASH_TF_SID: HASH_TF_INDEX_TABLE,
}

def stream_id_to_kvlayer_key(stream_id):
    '''Convert a text stream ID to a kvlayer key.

    The return tuple can be used directly as a key in the
    :data:`STREAM_ITEMS_TABLE` table.

    :param str stream_id: stream ID to convert
    :return: :mod:`kvlayer` key tuple
    :raise exceptions.KeyError: if `stream_id` is malformed

    '''
    # Reminder: stream_id is 1234567890-123456789abcdef...0
    # where the first part is the (decimal) epoch_ticks and the second
    # part is the (hex) doc_id
    parts = stream_id.split('-')
    if len(parts) != 2:
        raise KeyError('invalid stream_id ' + stream_id)
    epoch_ticks_s = parts[0]
    doc_id_s = parts[1]
    if not epoch_ticks_s.isdigit():
        raise KeyError('invalid stream_id ' + stream_id)
    if doc_id_s.lstrip(string.hexdigits) != '':
        raise KeyError('invalid stream_id ' + stream_id)
    return (base64.b16decode(doc_id_s.upper()), int(epoch_ticks_s))

def kvlayer_key_to_stream_id(k):
    '''Convert a kvlayer key to a text stream ID.

    `k` should be of the same form produced by
    :func:`stream_id_to_kvlayer_key`.

    :param k: :mod:`kvlayer` key tuple
    :return: converted stream ID
    :returntype str:

    '''
    abs_url_hash, epoch_ticks = k
    return '{0}-{1}'.format(epoch_ticks,
                            base64.b16encode(abs_url_hash).lower())

def key_for_stream_item(si):
    '''Get a kvlayer key from a stream item.

    The return tuple can be used directly as a key in the
    :data:`STREAM_ITEMS_TABLE` table.  Note that this recalculates the
    stream ID, and if the internal data on the stream item is inconsistent
    then this could return a different result from
    :func:`stream_id_to_kvlayer_key`.

    :param si: stream item to get key for
    :return: :mod:`kvlayer` key tuple

    '''
    # get binary 16 byte digest
    urlhash = hashlib.md5(si.abs_url).digest()
    return (urlhash, int(si.stream_time.epoch_ticks))


# Copyright (c) 2017, Matt Layman and contributors

import argparse
import logging
import sys

from httpie.plugins import plugin_manager
from werkzeug.serving import run_simple

from httpony.application import make_app

def main(argv=sys.argv):
    args = parse(argv)
    """Serve up some ponies."""
    hostname = args.listen
    port = args.port
    print(
        "Making all your dreams for a pony come true on http://{0}:{1}.\n"
        "Press Ctrl+C to quit.\n".format(hostname, port))

    # Hush, werkzeug.
    logging.getLogger('werkzeug').setLevel(logging.CRITICAL)

    plugin_manager.load_installed_plugins()
    app = make_app()
    run_simple(hostname, port, app)

def parse(argv):
    """Parse the user arguments."""
    parser = build_parser()

    # argparse expects the executable to be removed from argv.
    args = parser.parse_args(argv[1:])

    return args

def build_parser():
    """Build the parser that will have all available commands and options."""
    description = (
        'HTTPony (pronounced aych-tee-tee-pony) is a simple HTTP '
        'server that pretty prints HTTP requests to a terminal. It '
        'is a useful aide for developing clients that send HTTP '
        'requests. HTTPony acts as a sink for a client so that a '
        'developer can understand what the client is sending.')
    parser = argparse.ArgumentParser(description=description)
    parser.add_argument(
        '-l', '--listen', help='set the IP address or hostname',
        default='localhost')
    parser.add_argument(
        '-p', '--port', help='set the port', default=8000, type=int)

    return parser

if __name__ == '__main__':
    main()