# descriptors.massproduced
#
# Generates Descriptor classes whose functionality can be
# implemented with a single function.

from __future__ import print_function, unicode_literals, division

import os

from descriptors import Descriptor
from descriptors.builtin_types import builtins_str, builtins_camel
from descriptors.docstrings import docstrings

range_funcs = [
    #(Descriptor name, function to satisfy, __init__ arg names)
    ("Positive", lambda x: x > 0, []),
    ("SemiPositive", lambda x: x >= 0, []),
    ("Negative", lambda x: x < 0, []),
    ("SemiNegative", lambda x: x <= 0, []),
    ("NotZero", lambda x: x != 0, []),
    ("GreaterThan", lambda x, a: x > a, ["threshold"]),
    ("GreaterThanOrEqual", lambda x, a: x >= a, ["threshold"]),
    ("LessThan", lambda x, a: x < a, ["threshold"]),
    ("LessThanOrEqual", lambda x, a: x <= a, ["threshold"]),
    ("InRange", lambda x, a, b: a <= x <= b, ["lower_bound", "upper_bound"])]

misc_funcs = [
    ("NotNone", lambda x: x is not None, []),
    ("Callable", lambda x: callable(x), []),
    ("HasAttr", lambda x, a: hasattr(x, a), ["attribute"]),
    ("Satisfies", lambda x, a: a(x), ["function"]),
    ("NotSatisfies", lambda x, a: not a(x), ["function"]),
    ("Length", lambda x, a: len(x) == a, ["length"]),
    ("MinLength", lambda x, a: len(x) >= a, ["min_length"]),
    ("MaxLength", lambda x, a: len(x) <= a, ["max_length"]),
    ("ExistingPath", lambda x: os.path.exists(x), [])]

# Turn the builtins tuple into a list like the other funcs lists
builtin_func_els = ", ".join([
    '("{name}", lambda obj: isinstance(obj, {ty}), [])'.format(
        name=b_camel,
        ty=b_str)
    for b_str, b_camel in zip(builtins_str, builtins_camel)])
list_code = "builtin_funcs = [{}]".format(builtin_func_els)
exec(list_code, locals())

funcs = builtin_funcs + range_funcs + misc_funcs

def create_init(attrs):
    """Create an __init__ method that sets all the attributes
    necessary for the function the Descriptor invokes to check the
    value.

    """
    args = ", ".join(attrs)
    vals = ", ".join(['getattr(self, "{}")'.format(attr) for attr in attrs])
    attr_lines = "\n    ".join(
        ["self.{attr} = {attr}".format(attr=attr) for attr in attrs])
    init_code = """def _init(self, {args}):
    super(self.__class__, self).__init__()
    {attr_lines}
    self.field_type += "({{}})".format(
        ", ".join([str(val) for val in [{vals}]]))
    """.format(args=args, attr_lines=attr_lines, vals=vals)
    exec(init_code, globals())
    return _init

def create_setter(func, attrs):
    """Create the __set__ method for the descriptor."""
    def _set(self, instance, value, name=None):
        args = [getattr(self, attr) for attr in attrs]
        if not func(value, *args):
            raise ValueError(self.err_msg(instance, value))
    return _set

def make_class(clsname, func, attrs):
    """Turn a funcs list element into a class object."""
    clsdict = {"__set__": create_setter(func, attrs)}
    if len(attrs) > 0:
        clsdict["__init__"] = create_init(attrs)
    clsobj = type(str(clsname), (Descriptor, ), clsdict)
    clsobj.__doc__ = docstrings.get(clsname)
    return clsobj

# create all descriptors classes from funcs and put them in this module
desc_dict = {
    name: make_class(name, func, attrs) for name, func, attrs in funcs}
globals().update(desc_dict)


import logging

class Dashboard:
    def __init__(self, display=None, data_feeds=None, notifications=None):
        self._display = display
        self._data_feeds = data_feeds or []
        self._notifications = notifications or []

    @property
    def display(self):
        return self._display

    @display.setter
    def display(self, display):
        self._display = display

    @property
    def data_feeds(self):
        return self._data_feeds

    def add_data_feeds(self, data_feeds):
        self._data_feeds += data_feeds

    @property
    def notifications(self):
        return self._notifications

    def add_notifications(self, notifications):
        self._notifications += notifications

class DashboardRunner:

    def __init__(self, dashboard):
        self._logger = logging.getLogger(__name__)
        self._dashboard = dashboard

    def cycle(self):
        """
        Cycles through notifications with latest results from data feeds.
        """
        messages = self.poll_datafeeds()
        notifications = self.process_notifications(messages)

        self.draw_notifications(notifications)

    def poll_datafeeds(self):
        messages = []
        for feed in self._dashboard.data_feeds:
            messages += feed.get_latest_messages()

        return messages

    def process_notifications(self, messages):
        for notification in self._dashboard.notifications:
            yield notification.create(messages)

    def draw_notifications(self, notifications):
        display = self._dashboard.display

        for notification in notifications:
            display.draw(notification)

class DashboardValidator:

    def validate(self, dashboard):
        self._check_display_supports_notification(dashboard)

    @staticmethod
    def _check_display_supports_notification(dashboard: Dashboard):
        supported_notifications = dashboard.display.get_supported_notifications()

        for notification in dashboard.notifications:
            output_types = notification.get_output_types()
            for output_type in output_types:
                if output_type not in supported_notifications:
                    raise DisplayDoesNotSupportNotification(dashboard.display, notification, output_type)

class ValidationException(Exception):
    def __init__(self, value):
        self.value = value

    def __str__(self):
        return repr(self.value)

class DisplayDoesNotSupportNotification(ValidationException):
    def __init__(self, display, notification, output_type_not_supported):
        super().__init__("Display '%s' does not support output type '%s'" % (display.name, output_type_not_supported))
        self.display = display
        self.notification = notification
        self.output_type_not_supported = output_type_not_supported


# descriptors.handmade
#
# Implements Descriptor classes that cannot easily be implemented
# with only a single function.

from __future__ import print_function, unicode_literals, division

import os
import re
from numbers import Number

from descriptors import Descriptor
from descriptors.massproduced import create_init

class In(Descriptor):
    """A descriptor that only allows assigning a value if that value is
    a member of a set of given elements.

    Example:
        class A(Validated):
            mood = In(set(["bad", "go away", "Why me!?", ":-("]))

        a = A()
        a.mood = "awesome"
        -> ValueError

    """
    def __init__(self, valid_values):
        create_init(["valid_values"])(self, valid_values)
        if not hasattr(valid_values, "__contains__"):
            raise TypeError(
                "Attempted to create an In-Descriptor with an argument that "
                "doesn't support membership testing.")
        self.valid_values = valid_values

    def __set__(self, instance, value, name=None):
        try:
            if not value in self.valid_values:
                raise ValueError(self.err_msg(instance, value))
        except TypeError:  # raised e.g. when checking if non-str in str
            raise ValueError(self.err_msg(instance, value))

class RegexMatch(Descriptor):
    """A descriptor that ensures the described attribute is only set
    to a string that matches the supplied regular expression.

    """
    def __init__(self, regex):
        create_init(["regex"])(self, regex)
        self.pattern = re.compile(regex)

    def __set__(self, instance, value, name=None):
        if not self.pattern.search(value):
            raise ValueError(self.err_msg(instance, value))

class NotRegexMatch(Descriptor):
    """A descriptor that ensures the described attribute is only set
    to a string that does not match the supplied regular expression.

    """
    def __init__(self, regex):
        create_init(["regex"])(self, regex)
        self.pattern = re.compile(regex)

    def __set__(self, instance, value, name=None):
        if self.pattern.search(value):
            raise ValueError(self.err_msg(instance, value))

class Apply(Descriptor):
    """A descriptor that sets the described attribute to the result
    obtained from applying the supplied function to the value being
    set.

    Example:
        class A(Validated):
            f = Apply(str.lower)

        a = A()
        a.f = "SoMe stRInG"
        print(a.f)
        -> some string

    """
    no_autoset = True

    def __init__(self, func):
        create_init(["func"])(self, func)
        if not callable(func):
            raise TypeError(
                "Tried to create an Apply-Descriptor with a argument "
                "that is not callable.")
        self.func = func

    def __set__(self, instance, value, name=None):
        value = self.func(value)
        super(self.__class__, self).__set__(instance, value, name)

class ForceNumeric(Descriptor):
    """A descriptor that only allows numeric values, with "numeric"
    meaning that the value is
        a) an instance of a subclass of numbers.Number, or
        b) a string that can be converted to a float or int.
    In case a) the value will be passed through as is, in case b)
    the string will be converted to an int or a float, depending
    on the numeric value represented by the string.

    Example:
        class A(Validated):
            f = Numeric()

        a = A()
        a.f = 7.0
        -> ok

        a.f = "7"
        print(a.f)
        -> 7

        a.f = "7.0"
        print(a.f)
        -> 7.0

        a.f = (7, 0)
        -> ValueError

    """
    no_autoset = True

    def __set__(self, instance, value, name=None):
        if not isinstance(value, Number):
            try:
                value = self.try_convert(value)
            except ValueError:
                raise ValueError(self.err_msg(instance, value))
        super(self.__class__, self).__set__(instance, value, name)

    @staticmethod
    def try_convert(value):
        """Convert value to a numeric value or raise a ValueError
        if that isn't possible.

        """
        convertible = ForceNumeric.is_convertible(value)
        if not convertible or isinstance(value, bool):
            raise ValueError
        if isinstance(str(value), str):
            return ForceNumeric.str_to_num(value)
        return float(value)

    @staticmethod
    def is_convertible(value):
        """Return True if value can be converted to a float."""
        try:
            float(value)
            return True
        except ValueError:
            return False
        except TypeError:
            return False

    @staticmethod
    def str_to_num(str_value):
        """Convert str_value to an int or a float, depending on the
        numeric value represented by str_value.

        """
        str_value = str(str_value)
        try:
            return int(str_value)
        except ValueError:
            return float(str_value)

class MadePath(Descriptor):
    """A descriptor that creates the path represented by the passed
    string if that path doesn't exist already.

    """
    def __set__(self, instance, value, name=None):
        if not os.path.exists(value):
            try:
                os.makedirs(value)
            except OSError as e:
                if os.path.exists(value):
                    # path was created in the time since we checked
                    pass
                else:
                    raise e

class EitherOr(Descriptor):
    """A descriptor asserting that either the described attribute OR
    other_attr are True-y, but not both or neither. The descriptor
    is only active if both attributes are not None to allow
    re-setting values.
    If other_attr does not exist, it is assumed to be None.

    Example:
        class A(Validated):
            f1 = EitherOr("f2")
            f2 = EitherOr("f1")

        a = A()
        a.f1 = True
        a.f2 = True
        -> ValueError

    """
    def __init__(self, other_attr):
        create_init(["other_attr"])(self, other_attr)

    def __set__(self, instance, value, name=None):
        if self.other_attr in instance.__dict__:
            other_value = getattr(instance, self.other_attr)
        else:
            other_value = None
        if value is not None and other_value is not None:
            if bool(value) == bool(other_value):
                raise ValueError(self.err_msg(instance, value))


from random import choice
from django import template

register = template.Library()

class GraphRenderer(template.Node):
    """
    This template tag is mean to render the graph object on the template
    """
    def __init__(self, graph, attr_string, dom_id):
        self.graph = template.Variable(graph)
        self.attr_string = attr_string
        self.dom_id = dom_id

    def render(self, context):
        graph = self.graph.resolve(context)
        ctx = {
            'dom_id': self.dom_id,
            'attr_string': self.attr_string,
            'graph_data': graph.json_data,
            'graph_options': graph.options
        }
        return """
        <div %(attr_string)s></div>
        <script type='text/javascript'>
            $.plot($("#%(dom_id)s"), %(graph_data)s, %(graph_options)s);
        </script>
        """ % ctx

@register.tag
def plot(parser, token):
    """
    Tag to plot graphs into the template
    """

    tokens = token.split_contents()
    tokens.pop(0)
    graph = tokens.pop(0)

    attrs = dict([token.split("=") for token in tokens])

    if 'id' not in attrs.keys():
        attrs['id'] = ''.join([chr(choice(range(65, 90))) for i in range(0, 5)])
    else:
        attrs['id'] = attrs['id'][1:len(attrs['id'])-1]

    attr_string = ''.join([" %s=%s" % (k, v) for k, v in attrs.iteritems()])
    return GraphRenderer(graph, attr_string, attrs['id'])


#!/usr/bin/env python
'''Convert raw text into cleaned HTML.

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

.. autoclass:: clean_html

'''
from __future__ import absolute_import
import logging
import re
import string

import html5lib
import lxml.etree
import lxml.html
import lxml.html.clean
import lxml.html.soupparser
from BeautifulSoup import UnicodeDammit

from streamcorpus_pipeline.emails import fix_emails
from streamcorpus_pipeline.stages import Configured

logger = logging.getLogger(__name__)

encoding_re = re.compile(
    '''(?P<start_xml>([^<]|\n)*?\<\?xml[^>]*)'''
    '''(?P<encoding>encoding=.*?)(?P<remainder>(\s|\?\>)(.|\n)*)''',
    re.I)

def drop_invalid_and_upper_utf8_chars(possibly_invalid_string):
    '''Clean unexpected Unicode characters, including non-BMP.

    Returns a copy of `possibly_invalid_string` with the following
    replaced by space (U+0020): ASCII control characters other than
    tab, carriage return, and newline; reserved code points for
    surrogate pairs; invalid characters U+FFFE and U+FFFF; and
    supplementary characters U+10000 and higher.

    :param unicode possibly_invalid_string: string to clean
    :return: cleaned string
    :returntype: :class:`unicode`

    '''
    return re.sub(ur'[^\t\r\n\u0020-\ud7ff\ue000-\ufffd]', u' ',
                  possibly_invalid_string)

def force_unicode(raw):
    '''Try really really hard to get a Unicode copy of a string.

    First try :class:`BeautifulSoup.UnicodeDammit` to try to force
    to Unicode; if that fails, assume UTF-8 encoding, and ignore
    all errors.

    :param str raw: string to coerce
    :return: Unicode approximation of `raw`
    :returntype: :class:`unicode`

    '''
    converted = UnicodeDammit(raw, isHTML=True)
    if not converted.unicode:
        converted.unicode = unicode(raw, 'utf8', errors='ignore')

    encoding_m = encoding_re.match(converted.unicode)
    if encoding_m:
        converted.unicode = \
            encoding_m.group('start_xml') + \
            encoding_m.group('remainder')

    return converted.unicode

def nice_decode(raw, stream_item=None, encoding=None):
    raw_decoded = None
    if encoding is not None:
        try:
            raw_decoded = raw.decode(encoding)
        except:
            logger.warn('encoding=%r provided, but failed',
                        encoding, exc_info=True)
            raw_decoded = None

    if raw_decoded is None and stream_item is not None:
        if stream_item.body and stream_item.body.encoding:
            encoding = stream_item.body.encoding
            try:
                raw_decoded = raw.decode(encoding)
            except:
                logger.warn('stream_item.body.encoding=%r provided, but failed',
                            encoding, exc_info=True)
                raw_decoded = None

        if raw_decoded and stream_item.body and stream_item.body.media_type:
            #text/html; charset=windows-1251
            parts = stream_item.body.media_type.split('=')
            if len(parts) >= 2:
                encoding = parts[1].strip()
                try:
                    raw_decoded = raw.decode(encoding)
                except:
                    logger.warn('stream_item.body.content_type=%r provided, but failed',
                                stream_item.body.content_type, exc_info=True)
                    raw_decoded = None

    return raw_decoded

def make_clean_html(raw, stream_item=None, encoding=None):
    '''Get a clean text representation of presumed HTML.

    Treat `raw` as though it is HTML, even if we have no idea what it
    really is, and attempt to get a properly formatted HTML document
    with all HTML-escaped characters converted to their unicode.

    This is called below by the `clean_html` transform stage, which
    interprets MIME-type.  If `character_encoding` is not provided,
    and `stream_item` is provided, then this falles back to
    :attr:`streamcorpus.StreamItem.body.encoding`.

    :param str raw: raw text to clean up
    :param stream_item: optional stream item with encoding metadata
    :type stream_item: :class:`streamcorpus.StreamItem`
    :returns: UTF-8-encoded byte string of cleaned HTML text
    :returntype: :class:`str`

    '''
    # Fix emails by protecting the <,> from HTML
    raw = fix_emails(raw)
    raw_decoded = nice_decode(raw, stream_item=stream_item, encoding=encoding)
    if raw_decoded is None:
        # give up on decoding it... maybe this should use force_unicode
        raw_decoded = raw

    # default attempt uses vanilla lxml.html
    try:
        root = lxml.html.document_fromstring(raw_decoded)
    except ValueError, exc:
        if 'with encoding declaration' in str(exc):
            root = lxml.html.document_fromstring(raw)
        else:
            raise

    # While we have the document parsed as a DOM, let's strip attributes.
    # (The HTML cleaner seems to only support whitelisting attributes.
    # As of now, we just want to blacklist a few.)
    lxml.etree.strip_attributes(root, 'class', 'id')

    # if that worked, then we will be able to generate a
    # valid HTML string
    fixed_html = lxml.html.tostring(root, encoding=unicode)

    # remove any ^M characters
    fixed_html = string.replace(fixed_html, '\r', ' ')

    # We drop utf8 characters that are above 0xFFFF as
    # Lingpipe seems to be doing the wrong thing with them.
    fixed_html = drop_invalid_and_upper_utf8_chars(fixed_html)

    # construct a Cleaner that removes any ``<script>`` tags,
    # Javascript, like an ``onclick`` attribute, comments, style
    # tags or attributes, ``<link>`` tags
    cleaner = lxml.html.clean.Cleaner(
        scripts=True, javascript=True,
        comments=True,
        # do not remove <html> <head> <title> etc
        page_structure=False,
        remove_tags=['base'],
        style=True, links=True)

    # now get the really sanitized HTML
    _clean_html = cleaner.clean_html(fixed_html)

    # generate pretty HTML in utf-8
    _clean_html = lxml.html.tostring(
        lxml.html.document_fromstring(_clean_html),
        method='html', encoding='utf-8',
        pretty_print=True,
        # include_meta_content_type=True
        )

    return uniform_html(_clean_html)

def uniform_html(html):
    '''Takes a utf-8-encoded string of HTML as input and returns a new
    HTML string with fixed quoting and close tags, which generally
    should not break any of the offsets and makes it easier for
    functions like
    :func:`streamcorpus_pipeline.offsets.char_offsets_to_xpaths` to
    operate without failures.

    '''
    doc = html5lib.parse(html.decode('utf-8'))
    config = {
        'omit_optional_tags': False,
        'encoding': 'utf-8',
        'quote_attr_values': 'always',
    }
    return html5lib.serializer.serialize(doc, **config)

class clean_html(Configured):
    '''Create `body.clean_html` from `body.raw`.

    Configuration options:

    .. code-block:: yaml

        require_language_code: en

    If set, only work on stream items whose language code is the
    specified value; pass others on unchanged.  Defaults to unset.

    .. code-block:: yaml

        include_language_codes: [en, ""]

    If set to a non-empty list, only work on stream items whose
    language code is one of the specified values; pass others on
    unchanged.  Defaults to empty list (process all languages).

    .. code-block:: yaml

        include_mime_types: ['text/html']

    If set to a non-empty list, only work on stream items whose
    mime type matches one of the specified values. If set to an
    empty list, then all stream items (assuming they match the
    above language criteria) will have ``clean_html`` populated
    from ``body.raw``. Default to list of just ``text/html``.

    '''
    config_name = 'clean_html'
    default_config = {
        'include_language_codes': [],
        'include_mime_types': ['text/html'],
    }

    def __init__(self, *args, **kwargs):
        super(clean_html, self).__init__(*args, **kwargs)
        self.require_code = self.config.get('require_language_code')
        self.codes = self.config.get('include_language_codes', [])
        self.include_mime_types = \
            [s.lower()
             for s in self.config.get('include_mime_types', ['text/html'])]

    def is_matching_mime_type(self, mime_type):
        '''This implements the MIME-type matching logic for deciding whether
        to run `make_clean_html`

        '''
        if len(self.include_mime_types) == 0:
            return True
        if mime_type is None:
            return False
        mime_type = mime_type.lower()
        # NB: startswith is necessary here, because encodings are
        # often appended to HTTP header Content-Type
        return any(mime_type.startswith(mt) for mt in self.include_mime_types)

    def __call__(self, stream_item, context):
        if not stream_item.body or not stream_item.body.raw:
            logger.warn('dropping stream_item %s with missing body',
                        stream_item.stream_id)
            return None

        lang = stream_item.body.language
        if not lang and (self.require_code or self.codes):
            logger.warn('skipping stream_item %s with missing language '
                        '(did you run the "language" transform?)',
                        stream_item.stream_id)
            return stream_item

        if self.require_code and self.require_code != lang.code:
            logger.info('skipping stream_item %s with language %r, '
                        'not required language %r',
                        stream_item.stream_id,
                        stream_item.body.language.code,
                        self.require_code)
            return stream_item

        if self.codes and lang.code not in self.codes:
            logger.info('skipping stream_item %s with language %r, '
                        'not one of included codes %r',
                        stream_item.stream_id, lang.code, self.codes)

        if self.is_matching_mime_type(stream_item.body.media_type):
            stream_item.body.clean_html = make_clean_html(
                stream_item.body.raw,
                stream_item=stream_item)
        else:
            logger.info('skipping stream_item %s with unrecognized '
                        'media_type %r (allowed mime types: %r)',
                        stream_item.stream_id, stream_item.body.media_type,
                        self.include_mime_types or 'EVERYTHING')
        return stream_item


'''
Incremental transforms that only pass specific StreamItems.

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

'''
from __future__ import absolute_import
from itertools import imap, chain
import logging
import re
from urlparse import urlparse

from backports import lzma

from streamcorpus_pipeline.stages import Configured

logger = logging.getLogger(__name__)

class dump_stream_id_abs_url(Configured):
    '''print :attr:`~streamcorpus.StreamItem.stream_id`,
    :attr:`~streamcorpus.StreamItem.abs_url` strings to stdout and do
    not pass any :class:`~streamcorpus.StreamItems` down the pipeline

    '''
    config_name = 'dump_abs_url'
    default_config = {}
    def __call__(self, si, context):
        print si.stream_id, si.abs_url
        return None

class debug_filter(Configured):
    '''Remove all stream items except specified ones.

    This is only needed to debug troublesome configurations if there
    are problems with specific stream items.  It has one
    configuration item:

    .. code-block:: yaml

        accept_stream_ids:
          - 1360783522-ebf059defa41b2812c1f969f28cdb45e

    Only stream IDs in the list are passed on, all others are dropped.
    Defaults to an empty list (drop all stream items).

    '''
    config_name = 'debug_filter'
    default_config = { 'accept_stream_ids': [] }
    def __call__(self, si, context):
        if si.stream_id in self.config['accept_stream_ids']:
            return si
        else:
            return None

class id_filter(Configured):
    '''Include or exclude specific stream items by doc_id or stream_id

    If specific stream items have been identified, e.g. as causing
    problems with a tagger or other stages, then this can remove them
    before they run.

    .. code-block:: yaml

      id_filter:
        excluded_stream_ids_path: file-of-one-stream-id-per-line.txt.xz
        excluded_stream_ids:
          - 1360783522-ebf059defa41b2812c1f969f28cdb45e
        included_stream_ids_path: file-of-one-stream-id-per-line.txt.xz
        included_stream_ids:
          - 1360783522-ebf059defa41b2812c1f969f28cdb45e
        included_doc_ids_path: file-of-one-doc-id-per-line.txt.xz
        included_doc_ids:
          - ebf059defa41b2812c1f969f28cdb45e
        excluded_doc_ids_path: file-of-one-doc-id-per-line.txt.xz
        excluded_doc_ids:
          - ebf059defa41b2812c1f969f28cdb45e

    Stream IDs in the `excluded_stream_ids` list are dropped, all
    others are passed on unmodified.  Defaults to an empty list (don't
    drop any stream items).

    In addition to listing Stream IDs in the yaml itself, you can pass
    `excluded_stream_ids_path` with one Stream ID per line.  This file
    can be optionally be xz-compressed.

    '''
    config_name = 'id_filter'
    default_config = {
        'excluded_stream_ids': [],
        'included_stream_ids': [],
        'excluded_doc_ids': [],
        'included_doc_ids': [],
        }

    def __init__(self, *args, **kwargs):
        super(id_filter, self).__init__(*args, **kwargs)
        self.included_doc_ids = set()
        self.excluded_doc_ids = set()
        self.included_stream_ids = set()
        self.excluded_stream_ids = set()
        for mode in ['excluded', 'included']:
            for id_type in ['stream_ids', 'doc_ids']:
                self.configure(mode, id_type)

    def configure(self, mode, id_type):
        mode_type = '%s_%s' % (mode, id_type)
        data = set(self.config.get(mode_type, []))
        path = self.config.get(mode_type + '_path')
        if path:
            if path.endswith('.xz'):
                fh = lzma.open(path)
            else:
                fh = open(path)
            map(data.add, fh.read().splitlines())
        setattr(self, mode_type, data)
        logger.info('finished loading %d %s to %s', len(data), id_type, mode)

    def __call__(self, si, context):
        if si.stream_id in self.excluded_stream_ids:
            return None
        if si.doc_id in self.excluded_stream_ids:
            return None

        if self.included_stream_ids or self.included_doc_ids:
            if si.stream_id in self.included_stream_ids:
                return si
            if si.doc_id in self.included_doc_ids:
                return si
        else:
            ## not filtering by inclusion
            return si

class filter_languages(Configured):
    '''Remove stream items that aren't in specific languages.

    This looks at the :attr:`~streamcorpus.StreamItem.body`
    :attr:`~streamcorpus.ContentItem.language` field.  It has
    two configuration options:

    .. code-block:: yaml

        filter_languages:
          allow_null_language: True
          included_language_codes: [en]

    If the language field is empty, the stream item is dropped unless
    `allow_null_language` is true (default value).  Otherwise, the
    stream item is dropped unless its language code is one of the
    `included_language_codes` values (defaults to empty list).

    '''
    config_name = 'filter_languages'
    default_config = { 'allow_null_language': True,
                       'included_language_codes': [] }
    def __call__(self, si, context):
        if not si.body.language and self.config['allow_null_language']:
            return si
        if si.body.language.code in self.config['included_language_codes']:
            return si
        ## otherwise return None, which excludes the language
        return None

class filter_tagger_ids(Configured):
    '''Remove stream items that lack a StreamItem.body.sentences entry
    from specified tagger_id(s)

    This looks at the :attr:`~streamcorpus.StreamItem.body.sentences`
    field.  It has two configuration options:

    .. code-block:: yaml

        filter_tagger_ids:
          tagger_ids_to_keep:
            - serif

    If the none of the tagger_ids_to_keep are in
    :attr:`~streamcorpus.StreamItem.body.sentences`, the stream item
    is dropped.  If one of the required tagger_ids is present, but
    points to an empty list, then it is also dropped.

    '''
    config_name = 'filter_tagger_ids'
    default_config = { 'tagger_ids_to_keep': [] }
    def __call__(self, si, context):
        if not si.body or not si.body.sentences:
            return None
        for tagger_id in self.config['tagger_ids_to_keep']:
            sentences = si.body.sentences.get(tagger_id)
            if sentences:
                return si
        return None

class remove_raw(Configured):
    '''Remove the raw form of the content.

    While the :attr:`~streamcorpus.StreamItem.body`
    :attr:`~streamcorpus.ContentItem.raw` form is important as inputs
    to create the "clean" forms, it may be bulky and not useful later.
    This strips the raw form, replacing it with an empty string.  It
    has one configuration option:

    .. code-block:: yaml

        remove_raw:
          if_clean_visible_remove_raw: True

    If `if_clean_visible_remove_raw` is set (defaults to false),
    the raw form is only removed if the clean\_visible form is
    available; otherwise it is always removed.

    '''
    config_name = 'remove_raw'
    default_config = { 'if_clean_visible_remove_raw': False }
    def __call__(self, si, context):
        if self.config['if_clean_visible_remove_raw']:
            if si.body.clean_visible:
                si.body.raw = ''
            ## otherwise leave body.raw unchanged
        else:
            si.body.raw = ''

        return si

class replace_raw(Configured):
    '''Replace the `raw` form of the content with `clean_html`, or delete
    the StreamItem if it lacks `clean_html`.

    While the :attr:`~streamcorpus.StreamItem.body`
    :attr:`~streamcorpus.ContentItem.raw` form is important as inputs
    to create the "clean" forms, it may be bulky and not useful later.
    This strips the raw form, replacing it with clean_html.

    '''
    config_name = 'replace_raw'
    default_config = {}
    def __call__(self, si, context=None):
        if si.body.clean_html:
            si.body.raw = si.body.clean_html
            return si
        else:
            return None # delete

def domain_name_cleanse(raw_string):
    '''extract a lower-case, no-slashes domain name from a raw string
    that might be a URL
    '''
    try:
        parts = urlparse(raw_string)
        domain = parts.netloc.split(':')[0]
    except:
        domain = ''
    if not domain:
        domain = raw_string
    if not domain:
        return ''
    domain = re.sub('\/', '', domain.strip().lower())
    return domain

def domain_name_left_cuts(domain):
    '''returns a list of strings created by splitting the domain on
    '.' and successively cutting off the left most portion
    '''
    cuts = []
    if domain:
        parts = domain.split('.')
        for i in range(len(parts)):
            cuts.append( '.'.join(parts[i:]))
    return cuts

class filter_domains(Configured):
    '''Remove stream items that are not from a specific domain by
    inspecting first :attr:`~StreamItem.abs_url` and then
    :attr:`~StreamItem.schost`.  domain name strings are cleansed.

    .. code-block:: yaml

        filter_domains:
          include_domains: [example.com]
          include_domains_path:
          - path-to-file-with-one-domain-per-line.txt
          - path-to-file-with-one-domain-per-line2.txt

    '''
    config_name = 'filter_domains'
    default_config = { 'include_domains': [] }
    def __init__(self, config, *args, **kwargs):
        super(filter_domains, self).__init__(config, *args, **kwargs)

        ## cleanse the input domains lists
        self.domains = set()

        include_domains = self.config.get('include_domains', [])
        map(self.domains.add, imap(domain_name_cleanse, include_domains))

        include_domains_path = self.config.get('include_domains_path', [])
        if not isinstance(include_domains_path, list):
            include_domains_path = [include_domains_path]
        map(self.domains.add, imap(
                domain_name_cleanse,
                chain(*imap(open, include_domains_path))))

        logger.info('filter_domains configured with %d domain names',
                    len(self.domains))
        logger.info('filter_domains.domains = %r', self.domains)

    def __call__(self, si, context=None):
        url_things = [si.abs_url, si.schost]
        for thing in url_things:
            if not thing:
                continue
            for cut in domain_name_left_cuts(domain_name_cleanse(thing)):
                if cut in self.domains:
                    logger.info('found: %r', cut)
                    return si

        ## otherwise return None, which excludes the stream item
        logger.debug('rejecting: %r %r', si.schost, si.abs_url)
        return None

class filter_domains_substrings(Configured):
    '''Remove stream items that are not from a domain whose name contains
    one of the specified strings by inspecting first
    :attr:`~StreamItem.abs_url` and then :attr:`~StreamItem.schost`.
    Substrings are used exactly as provided without cleansing.

    .. code-block:: yaml

        filter_domains_matching:
          include_substrings: [example]
          include_substrings_path:
          - path-to-file-with-one-substring-per-line.txt
          - path-to-file-with-one-substring-per-line2.txt

    '''
    config_name = 'filter_domains_substrings'
    default_config = { 'include_substrings': [] }
    def __init__(self, config, *args, **kwargs):
        super(filter_domains_substrings, self).__init__(config, *args, **kwargs)

        ## cleanse the input substrings lists
        self.substrings = set()

        include_substrings = self.config.get('include_substrings', [])
        map(self.substrings.add, include_substrings)

        include_substrings_path = self.config.get('include_substrings_path', [])
        if not isinstance(include_substrings_path, list):
            include_substrings_path = [include_substrings_path]
        map(self.substrings.add,
            imap(lambda s: s.strip(),
                 chain(*imap(open, include_substrings_path))))

        logger.info('filter_substrings configured with %d domain names',
                    len(self.substrings))
        logger.info('filter_domans_substrings.substrings = %r', self.substrings)

    def __call__(self, si, context=None):
        url_things = [si.abs_url, si.schost]
        for thing in url_things:
            if not thing:
                continue
            for substr in self.substrings:
                if substr in thing:
                    logger.info('found: %r in %r', substr, thing)
                    return si

        ## otherwise return None, which excludes the stream item
        logger.debug('rejecting: %r %r', si.schost, si.abs_url)
        return None


'''rudimentary keyword search indexing build on top of `kvlayer`

.. This software is released under an MIT/X11 open source license.
   Copyright 2012-2015 Diffeo, Inc.
'''
from __future__ import absolute_import
from collections import Counter, defaultdict
import logging

import mmh3

from many_stop_words import get_stop_words
from streamcorpus_pipeline._clean_visible import cleanse
from streamcorpus_pipeline._kvlayer_table_names import \
    HASH_TF_INDEX_TABLE, HASH_FREQUENCY_TABLE, HASH_KEYWORD_INDEX_TABLE, \
    kvlayer_key_to_stream_id, key_for_stream_item

logger = logging.getLogger(__name__)

DOCUMENT_HASH_KEY = 0
DOCUMENT_HASH_KEY_REPLACEMENT = 1

class keyword_indexer(object):
    '''Do simple token-based search on documents.

    This is a utility class used by the
    :class:`~streamcorpus_pipeline._kvlayer.to_kvlayer` writer to
    implement simple token-based search.

    .. automethod:: __init__
    .. automethod:: index

    '''

    def __init__(self, kvl, hash_docs=True, hash_frequencies=True,
                 hash_keywords=True, keyword_tagger_ids=None,
                 keyword_size_limit=128):
        '''Create a new keyword indexer.

        `kvl` should be preconfigured to handle the
        :mod:`streamcorpus_pipeline` tables; this constructor will not
        call
        :meth:`~kvlayer._abstract_storage.AbstractStorage.setup_namespace`.
        The various `hash` parameters control what :meth:`index`
        writes out, and if only reading the default values will work.

        :param kvl: :mod:`kvlayer` client object
        :param hash_docs: if true, record hash to document ID index
        :param hash_frequencies: if true, record hash document frequencies
        :param hash_keywords: if true, record hash to keyword reverse index
        :param keyword_tagger_ids: if not :const:`None`, only consider
          these tagger IDs
        :param int keyword_size_limit: largest token to consider

        '''
        self.client = kvl
        self.hash_docs = hash_docs
        self.hash_frequencies = hash_frequencies
        self.hash_keywords = hash_keywords
        self.keyword_tagger_ids = keyword_tagger_ids
        self.keyword_size_limit = keyword_size_limit

    # Class static, initialized on first use
    _stop_words = None

    @property
    def stop_words(self):
        if self._stop_words is None:
            self.__class__._stop_words = get_stop_words()
        return self._stop_words

    def make_hash(self, tok):
        '''Get a Murmur hash for a token.

        `tok` may be a :class:`unicode` string or a UTF-8-encoded
        byte string.  :data:`DOCUMENT_HASH_KEY`, hash value 0, is
        reserved for the document count, and this function remaps
        that value.

        '''
        (_, h) = self.make_hash_kw(tok)
        return h

    def make_hash_kw(self, tok):
        '''Get a Murmur hash and a normalized token.

        `tok` may be a :class:`unicode` string or a UTF-8-encoded
        byte string.  :data:`DOCUMENT_HASH_KEY`, hash value 0, is
        reserved for the document count, and this function remaps
        that value.

        :param tok: token to hash
        :return: pair of normalized `tok` and its hash

        '''
        if isinstance(tok, unicode):
            tok = tok.encode('utf-8')
        h = mmh3.hash(tok)
        if h == DOCUMENT_HASH_KEY:
            h = DOCUMENT_HASH_KEY_REPLACEMENT
        return (tok, h)

    def collect_words(self, si):
        '''Collect all of the words to be indexed from a stream item.

        This scans `si` for all of the configured tagger IDs.  It
        collects all of the token values (the
        :attr:`streamcorpus.Token.token`) and returns a
        :class:`collections.Counter` of them.

        :param si: stream item to scan
        :type si: :class:`streamcorpus.StreamItem`
        :return: counter of :class:`unicode` words to index
        :returntype: :class:`collections.Counter`

        '''
        counter = Counter()
        for tagger_id, sentences in si.body.sentences.iteritems():
            if ((self.keyword_tagger_ids is not None
                 and tagger_id not in self.keyword_tagger_ids)):
                continue
            for sentence in sentences:
                for token in sentence.tokens:
                    term = token.token  # always a UTF-8 byte string
                    term = term.decode('utf-8')
                    term = cleanse(term)
                    if ((self.keyword_size_limit is not None and
                         len(term) > self.keyword_size_limit)):
                        continue
                    if term not in self.stop_words:
                        counter[term] += 1
        return counter

    def index(self, si):
        '''Record index records for a single document.

        Which indexes this creates depends on the parameters to the
        constructor.  This records all of the requested indexes for
        a single document.

        '''
        if not si.body.clean_visible:
            logger.warn('stream item %s has no clean_visible part, '
                        'skipping keyword indexing', si.stream_id)
            return

        # Count tokens in si.clean_visible
        # We will recycle hash==0 for "# of documents"
        hash_counts = defaultdict(int)
        hash_counts[DOCUMENT_HASH_KEY] = 1
        hash_kw = defaultdict(int)
        words = self.collect_words(si)
        for tok, count in words.iteritems():
            (tok, tok_hash) = self.make_hash_kw(tok)
            hash_counts[tok_hash] += count
            hash_kw[tok] = tok_hash

        # Convert this and write it out
        if self.hash_docs:
            (k1, k2) = key_for_stream_item(si)
            kvps = [((h, k1, k2), n) for (h, n) in hash_counts.iteritems()
                    if h != DOCUMENT_HASH_KEY]
            self.client.put(HASH_TF_INDEX_TABLE, *kvps)

        if self.hash_frequencies:
            kvps = [((h,), 1) for h in hash_counts.iterkeys()]
            self.client.increment(HASH_FREQUENCY_TABLE, *kvps)

        if self.hash_keywords:
            kvps = [((h, t), 1) for (t, h) in hash_kw.iteritems()]
            self.client.increment(HASH_KEYWORD_INDEX_TABLE, *kvps)

    def invert_hash(self, tok_hash):
        '''Get strings that correspond to some hash.

        No string will correspond to :data:`DOCUMENT_HASH_KEY`; use
        :data:`DOCUMENT_HASH_KEY_REPLACEMENT` instead.

        :param int tok_hash: Murmur hash to query
        :return: list of :class:`unicode` strings

        '''
        return [tok_encoded.decode('utf8')
                for (_, tok_encoded) in
                self.client.scan_keys(HASH_KEYWORD_INDEX_TABLE,
                                      ((tok_hash,), (tok_hash,)))]

    def document_frequencies(self, hashes):
        '''Get document frequencies for a list of hashes.

        This will return all zeros unless the index was written with
        `hash_frequencies` set.  If :data:`DOCUMENT_HASH_KEY` is
        included in `hashes`, that value will be returned with the
        total number of documents indexed.  If you are looking for
        documents with that hash, pass
        :data:`DOCUMENT_HASH_KEY_REPLACEMENT` instead.

        :param hashes: hashes to query
        :paramtype hashes: list of :class:`int`
        :return: map from hash to document frequency

        '''
        result = {}
        for (k, v) in self.client.get(HASH_FREQUENCY_TABLE,
                                      *[(h,) for h in hashes]):
            if v is None:
                v = 0
            result[k[0]] = v
        return result

    def lookup(self, h):
        '''Get stream IDs for a single hash.

        This yields strings that can be retrieved using
        :func:`streamcorpus_pipeline._kvlayer.get_kvlayer_stream_item`,
        or fed back into :mod:`coordinate` or other job queue systems.

        Note that for common terms this can return a large number of
        stream IDs!  This is a scan over a dense region of a
        :mod:`kvlayer` table so it should be reasonably efficient,
        but be prepared for it to return many documents in a large
        corpus.  Blindly storing the results in a :class:`list`
        may be inadvisable.

        This will return nothing unless the index was written with
        :attr:`hash_docs` set.  No document will correspond to
        :data:`DOCUMENT_HASH_KEY`; use
        :data:`DOCUMENT_HASH_KEY_REPLACEMENT` instead.

        :param int h: Murmur hash to look up

        '''
        for (_, k1, k2) in self.client.scan_keys(HASH_TF_INDEX_TABLE,
                                                 ((h,), (h,))):
            yield kvlayer_key_to_stream_id((k1, k2))

    def lookup_tf(self, h):
        '''Get stream IDs and term frequencies for a single hash.

        This yields pairs of strings that can be retrieved using
        :func:`streamcorpus_pipeline._kvlayer.get_kvlayer_stream_item`
        and the corresponding term frequency.

        ..see:: :meth:`lookup`

        '''
        for ((_, k1, k2), v) in self.client.scan(HASH_TF_INDEX_TABLE,
                                                 ((h,), (h,))):
            yield (kvlayer_key_to_stream_id((k1, k2)), v)


"""Read spinn3r protostreams directly into the pipeline.

Purpose
=======

Read a spinn3r.com protostream file, the spinn3r.com Web API, or
continuously fetch spinn3r.com streams as a reader.

In your streamcorpus-pipeline YAML configuration, specify::

    reader: from_spinn3r_feed

Give a local file name containing a fetched spinn3r protostream file
as the input to the pipeline.

Implementation details
======================

The feed file itself consists of a sequence of length-prefixed chunks,
with alternating type and data chunks.  The file continues until a
"type" chunk specifies "END".  Each of the data chunks should produce
one StreamItem into the rest of the pipeline.

`from_spinn3r_feed` is the main reader.  It is exported to the
pipeline configuration in `streamcorpus_pipeline.stages`.

The spinn3r data can be injected externally.  Create a dictionary
mapping keys to the protostream data from spinn3r, cause it to be
passed to the stage constructor, set a configuration parameter
``use_prefetched: true``, and invoke the pipeline with the provided
key as the input filename.  The key can be any valid string.

`ProtoStreamReader` will read chunks of the spinn3r document, typically
in the spinn3rApi_pb2.Entry Google protobuf format.

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

"""

from __future__ import absolute_import

import itertools
import json
import logging
from StringIO import StringIO
import zlib

from google.protobuf.internal.decoder import _DecodeVarint

import streamcorpus
from streamcorpus_pipeline._exceptions import ConfigurationError
from streamcorpus_pipeline._spinn3r.protoStream_pb2 \
    import ProtoStreamDelimiter, ProtoStreamHeader
from streamcorpus_pipeline._spinn3r.spinn3rApi_pb2 import Entry
from streamcorpus_pipeline.stages import Configured

logger = logging.getLogger(__name__)

class ProtoStreamReader(object):
    """Reader object for spinn3r "protostream" files.

    These files are produced by spinn3r's external API.  The actual
    files consist of length-prefixed chunks, where the length is in
    protobuf varint format.  The actual file consists of a header,
    then alternating "delimiter" and data chunks.  The file ends when
    a delimiter chunk has an "end" flag.

    Create this with a file-like object, and then iterate through
    the produced chunks.  By default it will produce spinn3rApi_pb2.Entry
    objects out.  For example:

    >>> with open(filename, 'rb') as f:
    ...   reader = ProtoStreamReader(f)
    ...   for entry in reader:
    ...     ...

    The `header` property will contain the
    protoStream_pb2.ProtoStreamHeader object from the file header.

    """

    def __init__(self, f, entry_type=Entry):
        """Create a new ProtoStreamReader.

        ``f``
          file-like object to read
        ``entry_type``
          type of objects in this feed

        """
        super(ProtoStreamReader, self).__init__()
        self.f = f
        self._prefix = ''
        self._header = None
        self._entry_type = entry_type

    def _unread(self, b):
        """Push byte array 'b' on to the front of the prefix list; this
        undoes part of a _read().  This bytes will be the next to read."""
        self._prefix = b + self._prefix

    def _read(self, n):
        """Read (up to) 'n' bytes from the underlying file.  If any bytes
        have been pushed in with _unread() those are returned first."""
        if n <= len(self._prefix):
            # the read can be fulfilled entirely from the prefix
            result = self._prefix[:n]
            self._prefix = self._prefix[n:]
            return result
        # otherwise we need to read some
        n -= len(self._prefix)
        result = self._prefix + self.f.read(n)
        self._prefix = ""
        return result

    def _read_varint(self):
        """Read exactly a varint out of the underlying file."""
        buf = self._read(8)
        (n, l) = _DecodeVarint(buf, 0)
        self._unread(buf[l:])
        return n

    def _read_block(self):
        """Read a length from the file, then read that many bytes.
        Return the read block."""
        l = self._read_varint()
        return self._read(l)

    def _read_a(self, cls):
        """Read some protobuf-encoded object stored in a single block
        out of the file."""
        o = cls()
        o.ParseFromString(self._read_block())
        return o

    def _bootstrap(self):
        """Read the header block out of the underlying file.
        After this (and forevermore afterwards) the file will point
        at the start of a ProtoStreamDelimiter block."""
        self._header = self._read_a(ProtoStreamHeader)

    @property
    def header(self):
        if self._header is None:
            self._bootstrap()
        return self._header

    def __iter__(self):
        """Iterate through the objects in the file."""
        if self._header is None:
            self._bootstrap()
        while True:
            delim = self._read_a(ProtoStreamDelimiter)
            if delim.delimiter_type == ProtoStreamDelimiter.END:
                return
            assert delim.delimiter_type == ProtoStreamDelimiter.ENTRY
            yield self._read_a(self._entry_type)

class from_spinn3r_feed(Configured):
    """streamcorpus-pipeline reader for spinn3r.com feeds.

    Accepted configuration items include:

    ``use_prefetched``
      If set to a boolean value, always/never use prefetched data
      (default: unset; use prefetched data if present)
    ``publisher_type``
      If set, only return stream items whose publisher type in the
      spinn3r feed exactly matches this string

    A dictionary from URL to prefetched data can be passed as a
    parameter to the stage constructor.  If `use_prefetched` is
    :const:`True` then all input strings must be present in the
    prefetch dictionary, and this stage never makes an outgoing
    network connection.  If `use_prefetched` is :const:`False` then
    the prefetch dictionary is ignored.  Otherwise if an input string
    is present in the prefetch dictionary, then the prefetched data is
    used, and if not, it is fetched from the network.

    """
    config_name = 'from_spinn3r_feed'

    def __init__(self, config={}, prefetched=None):
        super(from_spinn3r_feed, self).__init__(config)
        if prefetched is None:
            prefetched = {}
        self.prefetched = prefetched

    def __call__(self, i_str):
        # Do we have prefetched data for i_str?  Can/should/must we
        # use it?
        if ((self.config.get('use_prefetched', False) and
             i_str not in self.prefetched)):
            raise ConfigurationError('from_spinn3r_feed "use_prefetched" '
                                     'is set, but prefetched content is '
                                     'missing')
        if ((self.config.get('use_prefetched', True) and
             i_str in self.prefetched)):
            logger.debug('using prefetched content for {0}'.format(i_str))
            stream = StringIO(self.prefetched[i_str])
        else:
            logger.debug('getting local content from {0}'.format(i_str))
            stream = open(i_str, 'rb')

        try:
            count = 0
            # spinn3r publisher type == streamitem source
            source = self.config.get('publisher_type')
            for si in _make_stream_items(stream):
                if source is not None and si.source != source:
                    continue
                count += 1
                yield si
            logger.info('produced {0} stream items'.format(count))
        finally:
            stream.close()

def _generate_stream_items(data):
    return _make_stream_items(StringIO(data))

def _make_stream_items(f):
    """Given a spinn3r feed, produce a sequence of valid StreamItems.

    Because of goopy Python interactions, you probably need to call
    this and re-yield its results, as

    >>> with open(filename, 'rb') as f:
    ...   for si in _make_stream_items(f):
    ...     yield si

    """
    reader = ProtoStreamReader(f)
    return itertools.ifilter(
        lambda x: x is not None,
        itertools.imap(_make_stream_item, reader))

def _make_stream_item(entry):
    """Given a single spinn3r feed entry, produce a single StreamItem.

    Returns 'None' if a complete item can't be constructed.

    """
    # get standard metadata, assuming it's present...
    if not hasattr(entry, 'permalink_entry'):
        return None
    pe = entry.permalink_entry

    # ...and create a streamitem...
    si = streamcorpus.make_stream_item(
        pe.date_found[:-1] + '.0Z',
        pe.canonical_link.href.encode('utf8'))
    if not si.stream_time:
        logger.debug('failed to generate stream_time from {0!r}'
                     .format(pe.date_found))
        return None
    if not si.abs_url:
        logger.debug('failed to generate abs_url from {0!r}'
                     .format(pe.canonical_link.href))
        return None

    # ...filling in the actual data
    si.body = _make_content_item(
        pe.content,
        alternate_data=entry.feed_entry.content.data)
    if not si.body:
        return None
    if not si.body.raw:
        return None

    if pe.content_extract.data:
        si.other_content['extract'] = _make_content_item(pe.content_extract)
    si.other_content['title'] = streamcorpus.ContentItem(
        raw=pe.title.encode('utf8'),
        media_type=pe.content_extract.mime_type,
        encoding='UTF-8')
    si.other_content['feed_entry_title'] = streamcorpus.ContentItem(
        raw=entry.feed_entry.title.encode('utf8'),
        media_type=entry.feed_entry.content.mime_type,
        encoding='UTF-8')
    if entry.feed_entry.content.data:
        si.other_content['feed_entry'] = _make_content_item(
            entry.feed_entry.content)
    si.source_metadata['lang'] = pe.lang[0].code
    si.source_metadata['author'] = json.dumps(
        dict(
            name=pe.author[0].name,
            email=pe.author[0].email,
            link=pe.author[0].link[0].href,
        )
    )
    si.source = entry.source.publisher_type
    return si

def _make_content_item(node, mime_type=None, alternate_data=None):
    """Create a ContentItem from a node in the spinn3r data tree.

    The ContentItem is created with raw data set to ``node.data``,
    decompressed if the node's encoding is 'zlib', and UTF-8
    normalized, with a MIME type from ``node.mime_type``.

    ``node``
      the actual node from the spinn3r protobuf data
    ``mime_type``
      string MIME type to use (defaults to ``node.mime_type``)
    ``alternate_data``
      alternate (compressed) data to use, if ``node.data`` is missing
      or can't be decompressed

    """
    raw = node.data
    if getattr(node, 'encoding', None) == 'zlib':
        try:
            raw = zlib.decompress(node.data)
        except Exception, exc:
            if alternate_data is not None:
                try:
                    raw = zlib.decompress(alternate_data)
                except Exception:
                    raise exc  # the original exception
                else:
                    raise
    if mime_type is None:
        mime_type = node.mime_type
    raw = raw.decode('utf8').encode('utf8')
    return streamcorpus.ContentItem(raw=raw, media_type=mime_type)


'''Read and write stream items to/from :mod:`kvlayer`.

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

.. autoclass:: from_kvlayer
.. autoclass:: to_kvlayer

'''
from __future__ import absolute_import
import base64
import logging
import re
import struct

import kvlayer
import streamcorpus
from streamcorpus_pipeline.stages import Configured
from streamcorpus_pipeline._kvlayer_table_names import \
    BY_TIME, WITH_SOURCE, KEYWORDS, HASH_TF_SID, HASH_FREQUENCY, \
    HASH_KEYWORD, \
    stream_id_to_kvlayer_key, kvlayer_key_to_stream_id, key_for_stream_item, \
    STREAM_ITEM_TABLE_DEFS, STREAM_ITEM_VALUE_DEFS, INDEX_TABLE_NAMES, \
    STREAM_ITEMS_TABLE, STREAM_ITEMS_SOURCE_INDEX, STREAM_ITEMS_TIME_INDEX
import yakonfig

try:
    from ._kvlayer_keyword_search import keyword_indexer
except ImportError, exc:
    keyword_indexer = str(exc)

logger = logging.getLogger(__name__)

class WrongNumberException(Exception):
    pass

def iter_one_or_ex(gen):
    # for when you want exactly one result out of an iterable.
    # raise WrongNumberException otherwise.
    val = None
    any = False
    for tv in gen:
        if not any:
            any = True
            val = tv
        else:
            raise WrongNumberException("only wanted one, but got another")
    if not any:
        raise WrongNumberException("didn't get any")
    return val

class from_kvlayer(Configured):
    '''Read documents from a :mod:`kvlayer` supported database.

    .. currentmodule:: streamcorpus_pipeline._kvlayer

    :mod:`kvlayer` must be included in the top-level configuration.
    Typically this reads stream items previously written by
    :class:`to_kvlayer`.

    The input string is a sequence of stream IDs; see :meth:`__call__`
    for details.  If no values are specified, scans the entire
    :data:`~streamcorpus_pipeline._kvlayer_table_names.STREAM_ITEMS_TABLE`
    table.

    .. automethod:: __call__

    '''
    config_name = 'from_kvlayer'

    @staticmethod
    def check_config(config, name):
        yakonfig.check_toplevel_config(kvlayer, name)

    def __init__(self, *args, **kwargs):
        super(from_kvlayer, self).__init__(*args, **kwargs)
        self.client = kvlayer.client()
        self.client.setup_namespace(STREAM_ITEM_TABLE_DEFS,
                                    STREAM_ITEM_VALUE_DEFS)

    def _loadkey(self, k):
        key, data = iter_one_or_ex(self.client.get(STREAM_ITEMS_TABLE, k))
        errors, data = streamcorpus.decrypt_and_uncompress(data)
        if errors:
            if not isinstance(errors, Exception):
                errors = Exception(repr(errors))
            raise errors
        yield streamcorpus.deserialize(data)

    def _loadrange(self, keya, keyb):
        for key, data in self.client.scan(STREAM_ITEMS_TABLE, (keya, keyb)):
            errors, data = streamcorpus.decrypt_and_uncompress(data)
            if errors:
                logger.error('could not decrypet_and_uncompress %s: %s',
                             key, errors)
                continue
            yield streamcorpus.deserialize(data)

    def __call__(self, i_str):
        '''Actually scan the table, yielding stream items.

        `i_str` consists of a sequence of 20-byte encoded stream IDs.
        These are the 16 bytes of the document ID, plus the 4-byte
        big-endian stream time.  :func:`serialize_si` can produce
        these.  `i_str` can then consist of a single encoded stream
        ID; an encoded stream ID, a literal ``<``, and a second
        encoded stream ID to specify a range of documents; or a list
        of either of the preceding separated by literal ``;``.

        .. todo:: make this support keyword index strings

        '''
        if not i_str:
            for si in self._loadrange(None, None):
                yield si
            return
        for si in parse_keys_and_ranges(i_str, self._loadkey,
                                        self._loadrange):
            yield si

_STREAM_ID_RE = re.compile(r'[0-9]+-[0-9a-fA-F]{32}')

def parse_keys_and_ranges(i_str, keyfunc, rangefunc):
    '''Parse the :class:`from_kvlayer` input string.

    This accepts two formats.  In the textual format, it accepts any
    number of stream IDs in timestamp-docid format, separated by ``,``
    or ``;``, and processes those as individual stream IDs.  In the
    binary format, it accepts 20-byte key blobs (16 bytes md5 hash, 4
    bytes timestamp) split by ``;`` or ``<``; e.g., ``a<f;x`` loads
    scans keys `a` through `f` and loads singly key `x`.

    `keyfunc` and `rangefunc` are run as generators and their yields
    are yielded from this function.

    '''
    while i_str:
        m = _STREAM_ID_RE.match(i_str)
        if m:
            # old style text stream_id
            for retval in keyfunc(stream_id_to_kvlayer_key(m.group())):
                yield retval
            i_str = i_str[m.end():]
            while i_str and ((i_str[0] == ',') or (i_str[0] == ';')):
                i_str = i_str[1:]
            continue

        if len(i_str) == SI_KEY_LENGTH:
            # one key, get it.
            key = parse_si_key(i_str)
            for retval in keyfunc(key):
                yield retval
            return

        keya = i_str[:SI_KEY_LENGTH]
        splitc = i_str[SI_KEY_LENGTH]
        if splitc == '<':
            # range
            keyb = i_str[SI_KEY_LENGTH+1:SI_KEY_LENGTH+1+SI_KEY_LENGTH]
            i_str = i_str[SI_KEY_LENGTH+1+SI_KEY_LENGTH:]
            keya = parse_si_key(keya)
            keyb = parse_si_key(keyb)
            for retval in rangefunc(keya, keyb):
                yield retval
        elif splitc == ';':
            # keya is single key to load
            keya = parse_si_key(keya)
            for retval in keyfunc(keya):
                yield retval
            i_str = i_str[SI_KEY_LENGTH+1+1:]
        else:
            logger.error('bogus key splitter %s, %r', splitc, i_str)
            return
        # continue parsing i_str

def get_kvlayer_stream_item(client, stream_id):
    '''Retrieve a :class:`streamcorpus.StreamItem` from :mod:`kvlayer`.

    This function requires that `client` already be set up properly::

        client = kvlayer.client()
        client.setup_namespace(STREAM_ITEM_TABLE_DEFS,
                               STREAM_ITEM_VALUE_DEFS)
        si = get_kvlayer_stream_item(client, stream_id)

    `stream_id` is in the form of
    :data:`streamcorpus.StreamItem.stream_id` and contains the
    ``epoch_ticks``, a hyphen, and the ``doc_id``.

    :param client: kvlayer client object
    :type client: :class:`kvlayer.AbstractStorage`
    :param str stream_id: stream Id to retrieve
    :return: corresponding :class:`streamcorpus.StreamItem`
    :raise exceptions.KeyError: if `stream_id` is malformed or does
      not correspond to anything in the database

    '''
    if client is None:
        client = kvlayer.client()
        client.setup_namespace(STREAM_ITEM_TABLE_DEFS,
                               STREAM_ITEM_VALUE_DEFS)
    key = stream_id_to_kvlayer_key(stream_id)
    for k, v in client.get(STREAM_ITEMS_TABLE, key):
        if v is not None:
            errors, bytestr = streamcorpus.decrypt_and_uncompress(v)
            return streamcorpus.deserialize(bytestr)
    raise KeyError(stream_id)

def make_doc_id_range(doc_id):
    '''Construct a tuple(begin, end) of one-tuple kvlayer keys from a
    hexdigest doc_id.

    '''
    assert len(doc_id) == 32, 'expecting 32 hex string, not: %r' % doc_id
    bin_docid = base64.b16decode(doc_id.upper())
    doc_id_range = ((bin_docid,), (bin_docid,))
    return doc_id_range

def get_kvlayer_stream_item_by_doc_id(client, doc_id):
    '''Retrieve :class:`streamcorpus.StreamItem`s from :mod:`kvlayer`.

    Namely, it returns an iterator over all documents with the given
    docid. The docid should be an md5 hash of the document's abs_url.

    :param client: kvlayer client object
    :type client: :class:`kvlayer.AbstractStorage`
    :param str doc_id: doc id of documents to retrieve
    :return: generator of :class:`streamcorpus.StreamItem`
    '''
    if client is None:
        client = kvlayer.client()
        client.setup_namespace(STREAM_ITEM_TABLE_DEFS,
                               STREAM_ITEM_VALUE_DEFS)
    doc_id_range = make_doc_id_range(doc_id)
    for k, v in client.scan(STREAM_ITEMS_TABLE, doc_id_range):
        if v is not None:
            errors, bytestr = streamcorpus.decrypt_and_uncompress(v)
            yield streamcorpus.deserialize(bytestr)

def get_kvlayer_stream_ids_by_doc_id(client, doc_id):
    '''Retrieve stream ids from :mod:`kvlayer`.

    Namely, it returns an iterator over all stream ids with the given
    docid. The docid should be an md5 hash of the document's abs_url.

    :param client: kvlayer client object
    :type client: :class:`kvlayer.AbstractStorage`
    :param str doc_id: doc id of documents to retrieve
    :return: generator of str
    '''
    if client is None:
        client = kvlayer.client()
        client.setup_namespace(STREAM_ITEM_TABLE_DEFS,
                               STREAM_ITEM_VALUE_DEFS)
    doc_id_range = make_doc_id_range(doc_id)
    for k in client.scan_keys(STREAM_ITEMS_TABLE, doc_id_range):
        yield kvlayer_key_to_stream_id(k)

def delete_kvlayer_stream_item(client, stream_id):
    if client is None:
        client = kvlayer.client()
        client.setup_namespace(STREAM_ITEM_TABLE_DEFS,
                               STREAM_ITEM_VALUE_DEFS)
    key = stream_id_to_kvlayer_key(stream_id)
    client.delete(STREAM_ITEMS_TABLE, key)

class to_kvlayer(Configured):
    '''Writer that puts stream items in :mod:`kvlayer`.

    :mod:`kvlayer` must be included in the top-level configuration.
    Compressed stream items are written to the
    :data:`~streamcorpus_pipeline._kvlayer_table_names.STREAM_ITEMS_TABLE`
    table.

    This writer has one configuration parameter: `indexes` adds
    additional indexes on the data.  Supported index types include
    :data:`~streamcorpus_pipeline._kvlayer_table_names.WITH_SOURCE`,
    :data:`~streamcorpus_pipeline._kvlayer_table_names.BY_TIME`,
    :data:`~streamcorpus_pipeline._kvlayer_table_names.HASH_TF_SID`,
    :data:`~streamcorpus_pipeline._kvlayer_table_names.HASH_FREQUENCIES`,
    and
    :data:`~streamcorpus_pipeline._kvlayer_table_names.HASH_KEYWORD`.
    The special index type
    :data:`~streamcorpus_pipeline._kvlayer_table_names.KEYWORDS`
    expands to all of the keyword-index tables.

    If keyword indexing is enabled, it indexes the tokens generated
    by some or all of the taggers that have been run, less a set of
    stop words.  The configuration parameter `keyword_tagger_ids` may
    be set to a list of tagger IDs; if it is set, only the tokens
    from those tagger IDs will be indexed.  If it is set to ``null``
    or left unset, all taggers' tokens will be indexed as distinct
    words.  Any tokens whose size is greater than `keyword_size_limit`
    (default 128) will not be indexed.

    '''
    config_name = 'to_kvlayer'
    default_config = {'indexes': [],
                      'keyword_size_limit': 128}

    @staticmethod
    def check_config(config, name):
        yakonfig.check_toplevel_config(kvlayer, name)
        for ndx in config['indexes']:
            if ndx != KEYWORDS and ndx not in INDEX_TABLE_NAMES:
                raise yakonfig.ConfigurationError(
                    'invalid {0} indexes type {1!r}'
                    .format(name, ndx))
            if isinstance(keyword_indexer, basestring):
                if ndx in [KEYWORDS, HASH_TF_SID, HASH_FREQUENCY,
                           HASH_KEYWORD]:
                    raise yakonfig.ConfigurationError(
                        'cannot configure {0} index {1}: {2}'
                        .format(name, ndx, keyword_indexer))

    @staticmethod
    def normalize_config(config):
        if KEYWORDS in config['indexes']:
            config['indexes'].remove(KEYWORDS)
            config['indexes'].append(HASH_TF_SID)
            config['indexes'].append(HASH_FREQUENCY)
            config['indexes'].append(HASH_KEYWORD)

    def __init__(self, *args, **kwargs):
        super(to_kvlayer, self).__init__(*args, **kwargs)
        self.client = kvlayer.client()
        self.client.setup_namespace(STREAM_ITEM_TABLE_DEFS,
                                    STREAM_ITEM_VALUE_DEFS)

        self.keyword_indexer = None
        hash_docs = HASH_TF_SID in self.config['indexes']
        hash_frequencies = HASH_FREQUENCY in self.config['indexes']
        hash_keywords = HASH_KEYWORD in self.config['indexes']
        keyword_tagger_ids = self.config.get('keyword_tagger_ids', None)
        keyword_size_limit = self.config['keyword_size_limit']
        if hash_docs or hash_frequencies or hash_keywords:
            self.keyword_indexer = keyword_indexer(
                self.client, hash_docs=hash_docs,
                hash_frequencies=hash_frequencies,
                hash_keywords=hash_keywords,
                keyword_tagger_ids=keyword_tagger_ids,
                keyword_size_limit=keyword_size_limit)

    def __call__(self, t_path, name_info, i_str):
        si_keys = []
        sitable = TableBuffer(self.client, STREAM_ITEMS_TABLE)
        outputs = {}
        indexes = self.config.get('indexes', [])
        for index_name in indexes:
            itn = INDEX_TABLE_NAMES.get(index_name)
            if itn is not None:
                outputs[itn] = TableBuffer(self.client, itn)

        for si in streamcorpus.Chunk(t_path):
            si_key = key_for_stream_item(si)
            data = streamcorpus.serialize(si)
            errors, data = streamcorpus.compress_and_encrypt(data)
            assert not errors, errors
            sitable.put(si_key, data)
            si_keys.append(serialize_si_key(si_key))

            for index_name in indexes:
                index_func = INDEX_FUNCTIONS.get(index_name)
                if index_func is not None:
                    for tablename, kv in index_func(si):
                        outputs[tablename].put(*kv)

            if self.keyword_indexer:
                self.keyword_indexer.index(si)

        sitable.flush()
        for outbuf in outputs.itervalues():
            outbuf.flush()

        return si_keys

SI_KEY_LENGTH = 20

def serialize_si_key(si_key):
    '''
    Return packed bytes representation of StreamItem kvlayer key.
    The result is 20 bytes, 16 of md5 hash, 4 of int timestamp.
    '''
    if len(si_key[0]) != 16:
        raise ValueError('bad StreamItem key, expected 16 byte '
                         'md5 hash binary digest, got: {0!r}'.format(si_key))
    return struct.pack('>16si', si_key[0], si_key[1])

def parse_si_key(si_key_bytes):
    '''
    Parse result of serialize_si_key()
    Return kvlayer key tuple.
    '''
    return struct.unpack('>16si', si_key_bytes)

def streamitem_to_key_data(si):
    '''
    extract the parts of a StreamItem that go into a kvlayer key,
    convert StreamItem to blob for storage.

    return (kvlayer key tuple), data blob
    '''
    key = key_for_stream_item(si)
    data = streamcorpus.serialize(si)
    errors, data = streamcorpus.compress_and_encrypt(data)
    assert not errors, errors
    return key, data

def index_source(si):
    '''
    yield (table name, (key, value))
    '''
    if si.source:
        yield (STREAM_ITEMS_SOURCE_INDEX, (key_for_stream_item(si), si.source))

def index_time(si):
    key = key_for_stream_item(si)
    yield (STREAM_ITEMS_TIME_INDEX, ((key[1], key[0]), ''))

# functions take a StreamItem and yield (table name, (key, value))
INDEX_FUNCTIONS = {
    WITH_SOURCE: index_source,
    BY_TIME: index_time,
}

class TableBuffer(object):
    def __init__(self, kvl, table_name, buffer_count=100):
        self.kvl = kvl
        self.table_name = table_name
        self.buffer_count = buffer_count
        self.buffer = []

    def put(self, key, value):
        self.buffer.append((key, value))
        if len(self.buffer) >= self.buffer_count:
            self.kvl.put(self.table_name, *self.buffer)
            self.buffer = []

    def flush(self):
        self.kvl.put(self.table_name, *self.buffer)
        self.buffer = []


import os
from collections import deque
from contextlib import contextmanager
from functools import update_wrapper

from wimpy.compat import range
from wimpy.compat import zip_longest
from wimpy.exceptions import WimpyError

__all__ = [
    "cached_property",
    "working_directory",
    "strip_prefix",
    "strip_suffix",
    "ceiling_division",
    "grouper",
    "chunks",
    "is_subsequence",
]

class cached_property(object):
    """
    non-data descriptor: property is computed once per instance and then replaces itself
    with an ordinary attribute.  Deleting the attribute invalidates the cache.
    """

    def __init__(self, func):
        update_wrapper(self, func)
        self.func = func

    def __get__(self, obj, cls):
        if obj is None:
            return self
        value = obj.__dict__[self.func.__name__] = self.func(obj)
        return value

@contextmanager
def working_directory(path):
    """Change working directory and restore the previous on exit"""
    prev_dir = os.getcwd()
    os.chdir(str(path))
    try:
        yield
    finally:
        os.chdir(prev_dir)

def strip_prefix(s, prefix, strict=False):
    """Removes the prefix, if it's there, otherwise returns input string unchanged.
    If strict is True, also ensures the prefix was present"""
    if s.startswith(prefix):
        return s[len(prefix) :]
    elif strict:
        raise WimpyError("string doesn't start with prefix")
    return s

def strip_suffix(s, suffix, strict=False):
    """Removes the suffix, if it's there, otherwise returns input string unchanged.
    If strict is True, also ensures the suffix was present"""
    if s.endswith(suffix):
        return s[: len(s) - len(suffix)]
    elif strict:
        raise WimpyError("string doesn't end with suffix")
    return s

def ceiling_division(numerator, denominator):
    """Divide and round up"""
    # Implementation relies on Python's // operator using floor division (round down)
    # This might surprise C users who expect rounding towards zero
    return -(-numerator // denominator)

def grouper(iterable, n, fillvalue=None):
    """Yield successive non-overlapping chunks of size n from iterable"""
    args = [iter(iterable)] * n
    return zip_longest(*args, fillvalue=fillvalue)

def chunks(iterable, chunk_size=3, overlap=0):
    # we'll use a deque to hold the values because it automatically
    # discards any extraneous elements if it grows too large
    if chunk_size < 1:
        raise WimpyError("chunk size too small")
    if overlap >= chunk_size:
        raise WimpyError("overlap too large")
    queue = deque(maxlen=chunk_size)
    it = iter(iterable)
    i = 0
    try:
        # start by filling the queue with the first group
        for i in range(chunk_size):
            queue.append(next(it))
        while True:
            yield tuple(queue)
            # after yielding a chunk, get enough elements for the next chunk
            for i in range(chunk_size - overlap):
                queue.append(next(it))
    except StopIteration:
        # if the iterator is exhausted, yield any remaining elements
        i += overlap
        if i > 0:
            yield tuple(queue)[-i:]

def is_subsequence(needle, haystack):
    """Are all the elements of needle contained in haystack, and in the same order?
    There may be other elements interspersed throughout"""
    it = iter(haystack)
    for element in needle:
        if element not in it:
            return False
    return True