<commit_before>import copy
import re
import rospy
import roslib.message
import roslib.msgs
from functools import reduce

from rqt_py_common.topic_helpers import get_field_type

def get_field_type_capable_with_index(field_string):
    m = re.search('(.+)(\[[0-9]+\])', field_string)
    if m:
        return get_field_type(m.group(1))
    else:
        return get_field_type(field_string)

def make_topic_strings(msg_instance, string=''):
    if isinstance(msg_instance, list):
        array_instance = get_field_type_capable_with_index(string)[0]()
        return make_topic_strings(array_instance, string + '[0]')
    try:
        return [make_topic_strings(msg_instance.__getattribute__(slot),
                                   string + '/' + slot)
                for slot in msg_instance.__slots__]
    except AttributeError:
        return string

def set_msg_attribute_value(msg_instance, topic_name, type, attributes,
                            array_index, value):
    message_target = msg_instance
    if len(attributes) >= 2:
        full_string = topic_name
        for attr in attributes[:-1]:
            full_string += '/' + attr
            m = re.search('(\w+)\[([0-9]+)\]', attr)
            if m:
                index = int(m.group(2))
                attr = m.group(1)
                array_type = get_field_type_capable_with_index(full_string)[0]
                while len(message_target.__getattribute__(attr)) <= index:
                    message_target.__getattribute__(attr).append(array_type())
                message_target = message_target.__getattribute__(attr)[index]
            elif get_field_type_capable_with_index(full_string)[1]:  # this is impossible
                array_type = get_field_type_capable_with_index(full_string)[0]
                if len(message_target.__getattribute__(attr)) == 0:
                    message_target.__getattribute__(attr).append(array_type())
                message_target = message_target.__getattribute__(attr)[0]
            else:
                message_target = message_target.__getattribute__(attr)
    if array_index is not None:
        array = message_target.__getattribute__(attributes[-1])
        while len(array) <= array_index:
            array.append(type())
        array[array_index] = value
        message_target.__setattr__(attributes[-1], array)
    else:
        message_target.__setattr__(attributes[-1], value)
    message_target = value

def flatten(complicated_list):
    if isinstance(complicated_list, list):
        return reduce(lambda a, b: a + flatten(b), complicated_list, [])
    else:
        return [complicated_list]

def find_topic_name(full_text, topic_dict):
    if full_text[0] != '/':
        full_text = '/' + full_text
    # This is topic
    if full_text in topic_dict:
        return (full_text, None, None)
    splited_text = full_text.split('/')[1:]
    topic_name = ''
    while topic_name not in topic_dict and splited_text:
        topic_name += '/' + splited_text[0]
        splited_text = splited_text[1:]
    if splited_text:
        m = re.search('(\w+)\[([0-9]+)\]', splited_text[-1])
        if m:
            splited_text[-1] = m.group(1)
            return (topic_name, splited_text, int(m.group(2)))
        else:
            return (topic_name, splited_text, None)
    else:
        return (None, None, None)

def get_value_type(topic_type_str, attributes):
    # for Header -> std_msgs/Header
    topic_type_str = roslib.msgs.resolve_type(topic_type_str, '')
    if not attributes:
        return (None, False)
    try:
        _, spec = roslib.msgs.load_by_type(topic_type_str)
    except roslib.msgs.MsgSpecException as e:
        return (None, False)
    try:
        head_attribute = attributes[0].split('[')[0]
        index = spec.names.index(head_attribute)
        field = spec.parsed_fields()[index]
        attr_type = field.base_type
        if field.is_builtin:
            if attr_type in ['int8', 'int16', 'int32', 'int64']:
                return (int, field.is_array)
            if attr_type in ['byte', 'uint8', 'uint16', 'uint32', 'uint64']:
                return ('uint', field.is_array)
            elif attr_type in ['float32', 'float64']:
                return (float, field.is_array)
            elif attr_type == 'string':
                return (str, field.is_array)
            elif attr_type == 'bool':
                return (bool, field.is_array)
            else:
                print 'not support %s' % attr_type
                return (None, False)
        else:
            return get_value_type(field.base_type, attributes[1:])
    except ValueError as e:
        return (None, False)
    return (None, False)

def make_text(topic_name, attributes, array_index):
    text = topic_name + '/' + '/'.join(attributes)
    if array_index is not None:
        text += '[%d]' % array_index
    return text

class TopicPublisher(object):

    def __init__(self, topic_name, message_class):
        self._name = topic_name
        self._publisher = rospy.Publisher(
            topic_name, message_class, queue_size=100)
        self._message = message_class()

    def get_topic_name(self):
        return self._name

    def publish(self):
        self._publisher.publish(self._message)

    def get_message(self):
        return self._message

class EasyPublisherModel(object):

    def __init__(self, publisher_class=TopicPublisher):
        self._publishers = {}
        self._publisher_class = publisher_class

    def publish_topic(self, topic_name):
        if topic_name in self._publishers:
            self._publishers[topic_name].publish()

    def get_publisher(self, topic_name):
        if topic_name in self._publishers:
            return self._publishers[topic_name]
        else:
            return None

    def _add_publisher_if_not_exists(self, topic_name, message_class):
        if topic_name not in self._publishers:
            self._publishers[topic_name] = self._publisher_class(
                topic_name, message_class)

    def get_topic_names(self):
        _, _, topic_types = rospy.get_master().getTopicTypes()
        return sorted([x[0] for x in topic_types])

    def expand_attribute(self, input_text, array_index=None):
        text = copy.copy(input_text)
        try:
            msg_type, is_array = get_field_type_capable_with_index(text)
            if is_array:
                if array_index is None:
                    text += '[0]'
                else:
                    array_string = '[%d]' % array_index
                    if not text.endswith(array_string):
                        text += array_string
            if msg_type == int: # for time ? not support
                return []
            elif msg_type:
                return flatten(make_topic_strings(msg_type(), text))
            else:
                return []
        except AttributeError as e:
            return []

    def resister_topic_by_text(self, text):
        _, _, topic_types = rospy.get_master().getTopicTypes()
        topic_dict = dict(topic_types)
        topic_name, attributes, array_index = find_topic_name(text, topic_dict)
        if not topic_name:
            rospy.logerr('%s not found' % text)
            return None
        topic_type_str = topic_dict[topic_name]
        message_class = roslib.message.get_message_class(topic_type_str)
        self._add_publisher_if_not_exists(topic_name, message_class)
        builtin_type, is_array = get_value_type(topic_type_str, attributes)
        return (topic_name, attributes, builtin_type, is_array, array_index)

    def shutdown(self):
        for pub in list(self._publishers.values()):
            pub.unregister()
<commit_msg>Fix NoneType error for make_topic_string<commit_after>import copy
import re
import rospy
import roslib.message
import roslib.msgs
from functools import reduce

from rqt_py_common.topic_helpers import get_field_type

def get_field_type_capable_with_index(field_string):
    m = re.search('(.+)(\[[0-9]+\])', field_string)
    if m:
        return get_field_type(m.group(1))
    else:
        return get_field_type(field_string)

def make_topic_strings(msg_instance, string=''):
    if msg_instance is None:
        return string
    if isinstance(msg_instance, list):
        array_instance = get_field_type_capable_with_index(string)[0]()
        return make_topic_strings(array_instance, string + '[0]')
    try:
        return [make_topic_strings(msg_instance.__getattribute__(slot),
                                   string + '/' + slot)
                for slot in msg_instance.__slots__]
    except AttributeError:
        return string

def set_msg_attribute_value(msg_instance, topic_name, type, attributes,
                            array_index, value):
    message_target = msg_instance
    if len(attributes) >= 2:
        full_string = topic_name
        for attr in attributes[:-1]:
            full_string += '/' + attr
            m = re.search('(\w+)\[([0-9]+)\]', attr)
            if m:
                index = int(m.group(2))
                attr = m.group(1)
                array_type = get_field_type_capable_with_index(full_string)[0]
                while len(message_target.__getattribute__(attr)) <= index:
                    message_target.__getattribute__(attr).append(array_type())
                message_target = message_target.__getattribute__(attr)[index]
            elif get_field_type_capable_with_index(full_string)[1]:  # this is impossible
                array_type = get_field_type_capable_with_index(full_string)[0]
                if len(message_target.__getattribute__(attr)) == 0:
                    message_target.__getattribute__(attr).append(array_type())
                message_target = message_target.__getattribute__(attr)[0]
            else:
                message_target = message_target.__getattribute__(attr)
    if array_index is not None:
        array = message_target.__getattribute__(attributes[-1])
        while len(array) <= array_index:
            array.append(type())
        array[array_index] = value
        message_target.__setattr__(attributes[-1], array)
    else:
        message_target.__setattr__(attributes[-1], value)
    message_target = value

def flatten(complicated_list):
    if isinstance(complicated_list, list):
        return reduce(lambda a, b: a + flatten(b), complicated_list, [])
    else:
        return [complicated_list]

def find_topic_name(full_text, topic_dict):
    if full_text[0] != '/':
        full_text = '/' + full_text
    # This is topic
    if full_text in topic_dict:
        return (full_text, None, None)
    splited_text = full_text.split('/')[1:]
    topic_name = ''
    while topic_name not in topic_dict and splited_text:
        topic_name += '/' + splited_text[0]
        splited_text = splited_text[1:]
    if splited_text:
        m = re.search('(\w+)\[([0-9]+)\]', splited_text[-1])
        if m:
            splited_text[-1] = m.group(1)
            return (topic_name, splited_text, int(m.group(2)))
        else:
            return (topic_name, splited_text, None)
    else:
        return (None, None, None)

def get_value_type(topic_type_str, attributes):
    # for Header -> std_msgs/Header
    topic_type_str = roslib.msgs.resolve_type(topic_type_str, '')
    if not attributes:
        return (None, False)
    try:
        _, spec = roslib.msgs.load_by_type(topic_type_str)
    except roslib.msgs.MsgSpecException as e:
        return (None, False)
    try:
        head_attribute = attributes[0].split('[')[0]
        index = spec.names.index(head_attribute)
        field = spec.parsed_fields()[index]
        attr_type = field.base_type
        if field.is_builtin:
            if attr_type in ['int8', 'int16', 'int32', 'int64']:
                return (int, field.is_array)
            if attr_type in ['byte', 'uint8', 'uint16', 'uint32', 'uint64']:
                return ('uint', field.is_array)
            elif attr_type in ['float32', 'float64']:
                return (float, field.is_array)
            elif attr_type == 'string':
                return (str, field.is_array)
            elif attr_type == 'bool':
                return (bool, field.is_array)
            else:
                print 'not support %s' % attr_type
                return (None, False)
        else:
            return get_value_type(field.base_type, attributes[1:])
    except ValueError as e:
        return (None, False)
    return (None, False)

def make_text(topic_name, attributes, array_index):
    text = topic_name + '/' + '/'.join(attributes)
    if array_index is not None:
        text += '[%d]' % array_index
    return text

class TopicPublisher(object):

    def __init__(self, topic_name, message_class):
        self._name = topic_name
        self._publisher = rospy.Publisher(
            topic_name, message_class, queue_size=100)
        self._message = message_class()

    def get_topic_name(self):
        return self._name

    def publish(self):
        self._publisher.publish(self._message)

    def get_message(self):
        return self._message

class EasyPublisherModel(object):

    def __init__(self, publisher_class=TopicPublisher):
        self._publishers = {}
        self._publisher_class = publisher_class

    def publish_topic(self, topic_name):
        if topic_name in self._publishers:
            self._publishers[topic_name].publish()

    def get_publisher(self, topic_name):
        if topic_name in self._publishers:
            return self._publishers[topic_name]
        else:
            return None

    def _add_publisher_if_not_exists(self, topic_name, message_class):
        if topic_name not in self._publishers:
            self._publishers[topic_name] = self._publisher_class(
                topic_name, message_class)

    def get_topic_names(self):
        _, _, topic_types = rospy.get_master().getTopicTypes()
        return sorted([x[0] for x in topic_types])

    def expand_attribute(self, input_text, array_index=None):
        text = copy.copy(input_text)
        try:
            msg_type, is_array = get_field_type_capable_with_index(text)
            if is_array:
                if array_index is None:
                    text += '[0]'
                else:
                    array_string = '[%d]' % array_index
                    if not text.endswith(array_string):
                        text += array_string
            if msg_type == int: # for time ? not support
                return []
            elif msg_type:
                return flatten(make_topic_strings(msg_type(), text))
            else:
                return []
        except AttributeError as e:
            return []

    def resister_topic_by_text(self, text):
        _, _, topic_types = rospy.get_master().getTopicTypes()
        topic_dict = dict(topic_types)
        topic_name, attributes, array_index = find_topic_name(text, topic_dict)
        if not topic_name:
            rospy.logerr('%s not found' % text)
            return None
        topic_type_str = topic_dict[topic_name]
        message_class = roslib.message.get_message_class(topic_type_str)
        self._add_publisher_if_not_exists(topic_name, message_class)
        builtin_type, is_array = get_value_type(topic_type_str, attributes)
        return (topic_name, attributes, builtin_type, is_array, array_index)

    def shutdown(self):
        for pub in list(self._publishers.values()):
            pub.unregister()
<|endoftext|>


<commit_before>from .utils import thumbmd5
from . import tasks
from .log import Log
from . import work
from flask import (
    current_app,
    copy_current_request_context,
)
from .bed.baidupan import BaiduPan
from redis import Redis
from contextlib import contextmanager
from functools import wraps

log = Log(__name__)

class CoreError(Exception):
    pass

def baidupan_redis_key(md5, width):
    return ':'.join(['baidupan', thumbmd5(md5, width)])

def sessioned(f):
    @wraps(f)
    def g(self, *args, **kargs):
        if 'session' not in kargs:
            with self.make_session() as session:
                kargs['session'] = session
                kargs['commit'] = True
                return f(self, *args, **kargs)
        else:
            if 'commit' not in kargs:
                kargs['commit'] = False
            return f(self, *args, **kargs)
    return g

class Core(object):

    def __init__(
        self,
        db,
        baidupan_cookie=None,
        baidupan_timeout=30,
    ):
        self.db = db
        self.redis = Redis()
        if baidupan_cookie:
            self.baidupan = BaiduPan(baidupan_cookie)
        self.baidupan_timeout = baidupan_timeout

    @contextmanager
    def make_session(self):
        s = self.db.create_scoped_session()
        try:
            yield s
        finally:
            s.remove()

    def thumbnail_mtime_bi_md5(self, md5):
        return self.db.thumbnail_mtime_bi_md5(md5)

    def thumbnail_bi_md5(self, md5, width):
        return self.db.thumbnail_bi_md5(md5, width)

    def thumbnail_cache_timeout_bi_md5(self, md5):
        return self.db.thumbnail_cache_timeout_bi_md5(md5)

    def imgur_bi_md5(self, md5):
        return self.db.imgur_bi_md5(md5)

    def baidupan_thumbnail_linkout(self, md5, width):
        uri = self.redis.get(baidupan_redis_key(md5, width))
        if uri:
            uri = uri.decode('ascii')
            return uri

        if not hasattr(self, 'baidupan'):
            return None

        detail = self.baidupan.uri_detail(
            thumbmd5(md5, width),
            life=self.baidupan_timeout,
        )
        if detail:
            life = int(max(self.baidupan_timeout, detail.life - 60))
            log.debug('{} life: {}', detail.uri, life)
            self.redis.setex(
                baidupan_redis_key(md5, width),
                detail.uri.encode('ascii'),
                life,
            )
            return detail.uri

        work.nonblock_call(
            copy_current_request_context(tasks.persist_thumbnail_to_baidupan),
            kargs=dict(
                makeapp=lambda: current_app,
                md5=md5,
                width=width,
            ),
            bound='io',
        )

    def imgur_thumbnail_linkout(self, md5, width):
        bim = self.db.imgur_bi_md5(thumbmd5(md5, width))
        if bim:
            return bim.link.replace('http://', 'https://')
        work.nonblock_call(
            copy_current_request_context(tasks.persist_thumbnail_to_imgur),
            kargs=dict(
                makeapp=lambda: current_app,
                md5=md5,
                width=width,
            ),
            bound='io',
        )

    def thumbnail_linkout(self, md5, width):
        for method in (
            self.baidupan_thumbnail_linkout,
            self.imgur_thumbnail_linkout,
        ):
            uri = method(md5, width)
            if uri:
                return uri

    def art_bi_md5(self, md5):
        return self.db.art_bi_md5(md5)

    @sessioned
    def art_detail_bi_md5(self, md5, session, commit):
        art = (
            session.query(self.db.Entry)
            .filter_by(md5=md5)
            .options(self.db.joinedload(self.db.Entry.posts, inner=True))
            .options(self.db.subqueryload(self.db.Entry.tags))
            .first()
        )
        if art:
            # touch it, to prevent session expire
            art.plus_count
        return art

    @contextmanager
    def scoped_all_session(self):
        try:
            yield
        finally:
            from sqlalchemy.orm.session import Session
            Session.close_all()

    def user_bi_id(self, id):
        return self.db.get_user_bi_id(id)

    def user_bi_openid(self, openid):
        return self.db.get_user_bi_openid(openid)

    @sessioned
    def has_plused(self, user, art, session, commit):
        '''http://stackoverflow.com/a/15314973/238472'''
        return session.scalar(
            self.db.select([
                self.db.exists('1')
                .select_from(self.db.table(self.db.Plus))
                .where(self.db.and_(
                    self.db.Plus.user_id == user.id,
                    self.db.Plus.entry_id == art.id,
                ))
            ])
        )

    @sessioned
    def plus_count(self, art_id, session, commit):
        return session.scalar(
            self.db.select([self.db.func.count('*')])
            .select_from(self.db.table(self.db.Plus))
            .where(self.db.Plus.entry_id == art_id)
        )

    @sessioned
    def plus(self, user_id, art_id, session, commit):
        session.execute(
            self.db.table(self.db.Plus).insert(dict(
                user_id=user_id,
                entry_id=art_id,
            ))
        )
        if commit:
            from sqlalchemy.exc import IntegrityError
            try:
                session.commit()
            except IntegrityError as e:
                code, _ = e.orig
                # http://stackoverflow.com/q/8072537/238472
                if code == 1062:
                    log.warning('duplicated plus ({}, {})', user_id, art_id)
                    pass
                else:
                    session.rollback()

    @sessioned
    def minus(self, user_id, art_id, session, commit):
        session.execute(
            self.db.table(self.db.Plus).delete(
                self.db.and_(
                    self.db.Plus.user_id == user_id,
                    self.db.Plus.entry_id == art_id,
                )
            )
        )
        if commit:
            try:
                session.commit()
            except:
                session.rollback()
                log.exception('minus ({}, {}) failed', user_id, art_id)
                raise CoreError(
                    'minus ({}, {}) failed due to database error',
                    user_id,
                    art_id
                )

    @sessioned
    def set_baidupcs_access_token(self, value, session, commit):
        session.merge(self.db.Meta(
            id='baidupcs_access_token',
            value=value.encode('ascii')
        ))
        if commit:
            session.commit()

    @sessioned
    def set_baidupcs_refresh_token(self, value, session, commit):
        session.merge(self.db.Meta(
            id='baidupcs_refresh_token',
            value=value.encode('ascii')
        ))
        if commit:
            session.commit()

    @sessioned
    def baidupcs_access_token(self, session, commit):
        meta = session.query(self.db.Meta).get('baidupcs_access_token')
        if meta is not None:
            return meta.value.decode('ascii')
<commit_msg>load tags of art detail in seperate query<commit_after>from .utils import thumbmd5
from . import tasks
from .log import Log
from . import work
from flask import (
    current_app,
    copy_current_request_context,
)
from .bed.baidupan import BaiduPan
from redis import Redis
from contextlib import contextmanager
from functools import wraps

log = Log(__name__)

class CoreError(Exception):
    pass

def baidupan_redis_key(md5, width):
    return ':'.join(['baidupan', thumbmd5(md5, width)])

def sessioned(f):
    @wraps(f)
    def g(self, *args, **kargs):
        if 'session' not in kargs:
            with self.make_session() as session:
                kargs['session'] = session
                kargs['commit'] = True
                return f(self, *args, **kargs)
        else:
            if 'commit' not in kargs:
                kargs['commit'] = False
            return f(self, *args, **kargs)
    return g

class Core(object):

    def __init__(
        self,
        db,
        baidupan_cookie=None,
        baidupan_timeout=30,
    ):
        self.db = db
        self.redis = Redis()
        if baidupan_cookie:
            self.baidupan = BaiduPan(baidupan_cookie)
        self.baidupan_timeout = baidupan_timeout

    @contextmanager
    def make_session(self):
        s = self.db.create_scoped_session()
        try:
            yield s
        finally:
            s.remove()

    def thumbnail_mtime_bi_md5(self, md5):
        return self.db.thumbnail_mtime_bi_md5(md5)

    def thumbnail_bi_md5(self, md5, width):
        return self.db.thumbnail_bi_md5(md5, width)

    def thumbnail_cache_timeout_bi_md5(self, md5):
        return self.db.thumbnail_cache_timeout_bi_md5(md5)

    def imgur_bi_md5(self, md5):
        return self.db.imgur_bi_md5(md5)

    def baidupan_thumbnail_linkout(self, md5, width):
        uri = self.redis.get(baidupan_redis_key(md5, width))
        if uri:
            uri = uri.decode('ascii')
            return uri

        if not hasattr(self, 'baidupan'):
            return None

        detail = self.baidupan.uri_detail(
            thumbmd5(md5, width),
            life=self.baidupan_timeout,
        )
        if detail:
            life = int(max(self.baidupan_timeout, detail.life - 60))
            log.debug('{} life: {}', detail.uri, life)
            self.redis.setex(
                baidupan_redis_key(md5, width),
                detail.uri.encode('ascii'),
                life,
            )
            return detail.uri

        work.nonblock_call(
            copy_current_request_context(tasks.persist_thumbnail_to_baidupan),
            kargs=dict(
                makeapp=lambda: current_app,
                md5=md5,
                width=width,
            ),
            bound='io',
        )

    def imgur_thumbnail_linkout(self, md5, width):
        bim = self.db.imgur_bi_md5(thumbmd5(md5, width))
        if bim:
            return bim.link.replace('http://', 'https://')
        work.nonblock_call(
            copy_current_request_context(tasks.persist_thumbnail_to_imgur),
            kargs=dict(
                makeapp=lambda: current_app,
                md5=md5,
                width=width,
            ),
            bound='io',
        )

    def thumbnail_linkout(self, md5, width):
        for method in (
            self.baidupan_thumbnail_linkout,
            self.imgur_thumbnail_linkout,
        ):
            uri = method(md5, width)
            if uri:
                return uri

    def art_bi_md5(self, md5):
        return self.db.art_bi_md5(md5)

    @sessioned
    def art_detail_bi_md5(self, md5, session, commit):
        art = (
            session.query(self.db.Entry)
            .filter_by(md5=md5)
            .options(self.db.joinedload(self.db.Entry.posts, inner=True))
            .first()
        )
        if art:
            # touch it, to prevent session expire
            art.plus_count
        return art

    @contextmanager
    def scoped_all_session(self):
        try:
            yield
        finally:
            from sqlalchemy.orm.session import Session
            Session.close_all()

    def user_bi_id(self, id):
        return self.db.get_user_bi_id(id)

    def user_bi_openid(self, openid):
        return self.db.get_user_bi_openid(openid)

    @sessioned
    def has_plused(self, user, art, session, commit):
        '''http://stackoverflow.com/a/15314973/238472'''
        return session.scalar(
            self.db.select([
                self.db.exists('1')
                .select_from(self.db.table(self.db.Plus))
                .where(self.db.and_(
                    self.db.Plus.user_id == user.id,
                    self.db.Plus.entry_id == art.id,
                ))
            ])
        )

    @sessioned
    def plus_count(self, art_id, session, commit):
        return session.scalar(
            self.db.select([self.db.func.count('*')])
            .select_from(self.db.table(self.db.Plus))
            .where(self.db.Plus.entry_id == art_id)
        )

    @sessioned
    def plus(self, user_id, art_id, session, commit):
        session.execute(
            self.db.table(self.db.Plus).insert(dict(
                user_id=user_id,
                entry_id=art_id,
            ))
        )
        if commit:
            from sqlalchemy.exc import IntegrityError
            try:
                session.commit()
            except IntegrityError as e:
                code, _ = e.orig
                # http://stackoverflow.com/q/8072537/238472
                if code == 1062:
                    log.warning('duplicated plus ({}, {})', user_id, art_id)
                    pass
                else:
                    session.rollback()

    @sessioned
    def minus(self, user_id, art_id, session, commit):
        session.execute(
            self.db.table(self.db.Plus).delete(
                self.db.and_(
                    self.db.Plus.user_id == user_id,
                    self.db.Plus.entry_id == art_id,
                )
            )
        )
        if commit:
            try:
                session.commit()
            except:
                session.rollback()
                log.exception('minus ({}, {}) failed', user_id, art_id)
                raise CoreError(
                    'minus ({}, {}) failed due to database error',
                    user_id,
                    art_id
                )

    @sessioned
    def set_baidupcs_access_token(self, value, session, commit):
        session.merge(self.db.Meta(
            id='baidupcs_access_token',
            value=value.encode('ascii')
        ))
        if commit:
            session.commit()

    @sessioned
    def set_baidupcs_refresh_token(self, value, session, commit):
        session.merge(self.db.Meta(
            id='baidupcs_refresh_token',
            value=value.encode('ascii')
        ))
        if commit:
            session.commit()

    @sessioned
    def baidupcs_access_token(self, session, commit):
        meta = session.query(self.db.Meta).get('baidupcs_access_token')
        if meta is not None:
            return meta.value.decode('ascii')
<|endoftext|>


<commit_before>#!/usr/bin/env python

try:
    from pyspark import SparkContext, SparkFiles
except:
    print "### NO PYSPARK"
import sys
import os
import platform
import socket
from hybridJaccard import HybridJaccard
import argparse
import json
import cgi
from htmltoken import tokenize
import crf_features
from base64 import b64encode, b64decode

# import snakebite for doing hdfs file manipulations
from snakebite.client import Client
from snakebite.errors import FileNotFoundException

def extract_body(main_json):
    try:
        return main_json["hasBodyPart"]["text"]
    except:
        pass

def extract_title(main_json):
    try:
        return main_json["hasTitlePart"]["text"]
    except:
        pass

def textTokens(texts):
    # Turn None into empty text
    texts = texts or ""
    # Allow for multiple texts
    texts = texts if isinstance(texts, list) else [texts]
    v = []
    for text in texts:
        try:
            for tok in genescaped(text):
                v.append([tok])
        except TypeError as e:
            print >> sys.stderr, "Error %s" % e
            print >> sys.stderr, type(text)
            print >> sys.stderr, "Computing textTokens of %s: %s" % (text, e)
        v.append("")
    return v

def genescaped(text, maxTokenLength=40):
    """All tokens in TEXT with any odd characters (such as <>&) encoded using HTML escaping"""
    for tok in tokenize(text, interpret=cgi.escape, keepTags=False):
        # Some ads have odd tokens like 1000 As in a row
        if len(tok) <= maxTokenLength:
            # yield tok
            yield tok.replace('\t', ' ')

### The URI + index mechanism isn't good enough to recover data when there are multiple sentences per URI
### which can occur from (a) title + body as separate documents (b) sentence breaking of body
### The intra-sentence space rows won't be maintained by reconstructTuple + groupBy
### Proposed future workarounds:
### (a) sentence-specific temporary URIs ending /processed/title or /processed/1
### (b) and/or sentence-specific indexing 1.1, 1.2, ... 1.N, 2.1, etc. could work

def reconstructTuple(tabsep):
    """Works only for old multi-line invocation of crf_test, not crf_test_b64"""
    fields = tabsep.split('\t')
    try:
        uri = fields[-3]
        idx = fields[-2]
    except:
        uri = "endOfDocument"
        idx = "0"
    return (uri, [idx] + fields)

def fakeFindBestMatch(words):
    if 'blue' in words:
        return 'blue'
    else:
        return 'NONE'

def alignToControlledVocab(harvested, vocabs):
    """harvested is a dict matching category 'eyeColor', 'hairType' to a function"""
    try:
        category = harvested['category']
        # f = vocabs[category]
        f = fakeFindBestMatch
        words = harvested['words']
        try:
            harvested['bestMatch'] = f(words)
        except Exception as e:
            return "sm.findBestMatch error:" + str(e) + ":" + str(words)
        return harvested
    except Exception as e:
        return str(e)
    return None

def vectorToUTF8(v, debug=False):
    "unicode only"

    def rowToUnicode(r):
        try:
            if isinstance(r, list):
                return u"\t".join([unicode(x) for x in r])
            else:
                return unicode(r)
        except:
            print >> sys.stderr, "error in rowToUnicode"
            return u""

    rows = []
    if v[-1] == u"":
        pass
    else:
        # print "appending1"
        v.append(u"")

    for r in v:
        rows.append(rowToUnicode(r))

    result = u"\n".join(rows)
    # result now a unicode object
    # here is the only place where we convert to UTF8
    return result.encode('utf-8')

### HISTORICAL
def computeSpans(v, verbose=False, indexed=False):
    # extract the data and write the result as vector
    currentLabel = None
    currentTokens = []
    spans = []
    def addSpan(u, l, words):
        spans.append( {"uri": u, "category": l, "words": " ".join(words) } )
        if verbose:
            print >> sys.stderr, "  Added %s" % (spans[-1],)

    uri = 'bogus'
    currentUri = None
    for row in v:
        if (len(row) <= 1):
            # blank/empty line: expecting "" but might be []
            if currentLabel and currentTokens:
                addSpan(currentUri, currentLabel, currentTokens)
            currentUri = None
            continue
        if (len(row) >= 4):
            # a typical row
            token = row[0]
            uri = row[-3] if indexed else row[-2]
            label = row[-1]
            if verbose:
                print >> sys.stderr, "Typical row: token %r uri %r: crflabel %r" % (token, uri, label)
            # now process this row
            if label == "O":
                # unlabeled row
                if currentLabel:
                    # so this concludes span in progress
                    addSpan(uri, currentLabel, currentTokens)
                    currentLabel = None
                    currentTokens = []
                else:
                    pass
            else:
                # Labeled row
                if label == currentLabel:
                    # continue span in progress
                    currentTokens.append(token)
                elif currentLabel and label != currentLabel:
                    # span/span boundary
                    # first conclude old one
                    addSpan(uri, currentLabel, currentTokens)
                    currentLabel = None
                    currentTokens = []
                    # then begin new one
                    currentLabel = label
                    currentTokens = [token]
                elif not currentLabel:
                    # begin novel span
                    currentLabel = label
                    currentTokens = [token]
                else:
                    raise Exception("Unexpected file structure")
        currentUri = uri

    if currentLabel and currentTokens:
        # Input ended without blank line after last marked span, so hallucinate one
        addSpan(uri, currentLabel, currentTokens)

    # Publish results
    return spans

def crfsmall(sc, input, output,
            limit=None, location='hdfs', outputFormat="text", numPartitions=None):

    crfConfigDir = os.path.join(os.path.dirname(__file__), "data/config")
    featureListFilename = os.path.join(crfConfigDir, "features.hair-eye")
    crfExecutable = os.path.join(os.path.dirname(__file__), "bin/crf_test_filter.sh")
    crfModelFilename = os.path.join(crfConfigDir, "dig-hair-eye-train.model")
    sc.addFile(crfExecutable)
    sc.addFile(crfModelFilename)

    # crfConfigDir = os.path.join(os.path.dirname(__file__), "data/config")
    # def cpath(n):
    #     return os.path.join(crfConfigDir, n)

    # smEyeColor = HybridJaccard(ref_path=cpath("eyeColor_reference_wiki.txt"),
    #                            config_path=cpath("eyeColor_config.txt"))
    # smHairColor = HybridJaccard(ref_path=cpath("hairColor_reference_wiki.txt"),
    #                             config_path=cpath("hairColor_config.txt"))
    # print smEyeColor, smHairColor

    # hypothesis1: data fetched this way prompts the lzo compression error
    # hypothesis2: but it doesn't matter, error is just a warning

    rdd_crfl = sc.sequenceFile(input)
    rdd_crfl.setName('rdd_crfl')

    if limit:
        rdd_crfl = sc.parallelize(rdd_crfl.take(limit))
    if numPartitions:
        rdd_crfl = rdd_crfl.repartition(numPartitions)

    rdd_json = rdd_crfl.mapValues(lambda x: json.loads(x))
    rdd_json.setName('rdd_json')
    # rdd_json.persist()

    rdd_texts = rdd_json.mapValues(lambda x: (textTokens(extract_body(x)), textTokens(extract_title(x))))
    rdd_texts.setName('rdd_texts')

    c = crf_features.CrfFeatures(featureListFilename)
    SEPARATOR = '&nbsp;',

    def makeMatrix(c, uri, bodyTokens, titleTokens):
        b = c.computeFeatMatrix(bodyTokens, False, addLabels=False, addIndex=False)
        s = c.computeFeatMatrix([SEPARATOR, ""], False, addLabels=False, addIndex=False)
        t = c.computeFeatMatrix(titleTokens, False, addLabels=False, addIndex=False)
        idx = 1
        for row in b:
            if row == u"":
                pass
            else:
                label = uri + "/%05d/%05d" % (0, idx)
                row.append(label)
                idx += 1
        idx = 1
        for row in s:
            if row == u"":
                pass
            else:
                label = uri + "/%05d/%05d" % (1, idx)
                row.append(label)
                idx += 1
        idx = 1
        for row in t:
            if row == u"":
                pass
            else:
                label = uri + "/%05d/%05d" % (2, idx)
                row.append(label)
                idx += 1
        # might be b[0:-1] + s[0:-1] + t?
        return b[0:-1] + s[0:-1] + t

    rdd_features = rdd_texts.map(lambda x: makeMatrix(c, x[0], x[1][0], x[1][1]))
    rdd_features.setName('rdd_features')
    # rdd_features.persist()

    # unicode/string representation of the feature matrix
    rdd_vector = rdd_features.map(lambda x: vectorToUTF8(x))
    rdd_vector.setName('rdd_vector')

    # all strings concatenated together, then base64 encoded into one input for crf_test
    rdd_pipeinput = sc.parallelize([b64encode(rdd_vector.reduce(lambda a,b: a+b))])
    rdd_pipeinput.setName('rdd_pipeinput')

    executable = SparkFiles.get(os.path.basename(crfExecutable))
    model = SparkFiles.get(os.path.basename(crfModelFilename))

    cmd = "%s %s" % (executable, model)

    # this result is base64 encoded
    rdd_crfoutput = rdd_pipeinput.pipe(cmd)
    rdd_crfoutput.setName('rdd_crfoutput')

    rdd_base64decode = rdd_crfoutput.map(lambda x: b64decode(x))
    # 1. break into physical lines
    # 2. turn each line into its own spark row
    # 3. drop any inter-document empty string markers
    rdd_lines = rdd_base64decode.map(lambda x: x.split("\n")).flatMap(lambda l: l).filter(lambda x: len(x)>1)

    def processOneLine(l):
        print l
        return l.split("\t")

    rdd_triples = rdd_lines.map(lambda l: processOneLine(l))
    rdd_triples.saveAsTextFile('out_rdd_triples')

    def organizeByOrigDoc(triple):
        try:
            (word, uri, label) = triple
            (parentUri, docId, wordId) = uri.rsplit('/', 2)
            return ( (parentUri, docId), (wordId, word, label) )
        except Exception as e:
            print "Can't destructure %r" % [triple]
            return ()

    rdd_reorg = rdd_triples.map(lambda l: organizeByOrigDoc(l))
    rdd_reorg.saveAsTextFile('out_rdd_reorg')

    rdd_sorted = rdd_reorg.sortByKey()
    rdd_sorted.saveAsTextFile('out_rdd_sorted')

    # each (parentUri, docId) has a sequence of (wordId, word, label)
    # we want to consider them in order (by wordId)

    rdd_grouped = rdd_sorted.groupByKey()

    def harvest(seq):
        allSpans = []
        lastIndex = -2
        lastLabel = None
        currentSpan = []
        for (wordId, word, label) in seq:
            currentIndex = int(wordId)
            if lastIndex+1 == currentIndex and lastLabel == label:
                # continuing current span
                currentSpan.append( (currentIndex, word, label) )
            else:
                # end current span
                if currentSpan:
                    allSpans.append(currentSpan)
                # begin new span
                currentSpan = [ (currentIndex, word, label) ]
                lastLabel = label
            lastIndex = currentIndex

        # end current span
        if currentSpan:
            allSpans.append(currentSpan)

        result = []
        for span in allSpans:
            words = []
            spanLabel = None
            for (wordIdx, word, label) in span:
                spanLabel = label
                words.append(word)
            result.append( (' '.join(words), spanLabel) )
        return result

    rdd_harvest = rdd_grouped.mapValues(lambda s: harvest(s))
    rdd_harvest.saveAsTextFile('out_rdd_harvest')
    exit(1)

    rdd_final = rdd_harvest

    if outputFormat == "sequence":
        rdd_final.saveAsSequenceFile(output)
    elif outputFormat == "text":
        rdd_final.saveAsTextFile(output)
    else:
        raise RuntimeError("Unrecognized output format: %s" % outputFormat)

def main(argv=None):
    '''this is called if run from command line'''
    parser = argparse.ArgumentParser()
    parser.add_argument('-i','--input', required=True)
    parser.add_argument('-o','--output', required=True)
    parser.add_argument('-p','--numPartitions', required=False, default=None, type=int)
    parser.add_argument('-l','--limit', required=False, default=None, type=int)
    parser.add_argument('-v','--verbose', required=False, help='verbose', action='store_true')
    args=parser.parse_args()

    location = "hdfs"
    try:
        if "avatar" in platform.node():
            location = "local"
    except:
        pass
    try:
        if "avatar" in socket.gethostname():
            location = "local"
    except:
        pass
    print "### location %s" % location

    if not args.numPartitions:
        if location == "local":
            args.numPartitions = 2
        elif location == "hdfs":
            args.numPartitions = 50

    sc = SparkContext(appName="crfsmall")
    crfsmall(sc, args.input, args.output,
             limit=args.limit,
             location=location,
             outputFormat="text",
             numPartitions=args.numPartitions)

# call main() if this is run as standalone
if __name__ == "__main__":
    sys.exit(main())
<commit_msg>flattens and uniquifies<commit_after>#!/usr/bin/env python

try:
    from pyspark import SparkContext, SparkFiles
except:
    print "### NO PYSPARK"
import sys
import os
import platform
import socket
from hybridJaccard import HybridJaccard
import argparse
import json
import cgi
from htmltoken import tokenize
import crf_features
from base64 import b64encode, b64decode

# import snakebite for doing hdfs file manipulations
from snakebite.client import Client
from snakebite.errors import FileNotFoundException

def extract_body(main_json):
    try:
        return main_json["hasBodyPart"]["text"]
    except:
        pass

def extract_title(main_json):
    try:
        return main_json["hasTitlePart"]["text"]
    except:
        pass

def textTokens(texts):
    # Turn None into empty text
    texts = texts or ""
    # Allow for multiple texts
    texts = texts if isinstance(texts, list) else [texts]
    v = []
    for text in texts:
        try:
            for tok in genescaped(text):
                v.append([tok])
        except TypeError as e:
            print >> sys.stderr, "Error %s" % e
            print >> sys.stderr, type(text)
            print >> sys.stderr, "Computing textTokens of %s: %s" % (text, e)
        v.append("")
    return v

def genescaped(text, maxTokenLength=40):
    """All tokens in TEXT with any odd characters (such as <>&) encoded using HTML escaping"""
    for tok in tokenize(text, interpret=cgi.escape, keepTags=False):
        # Some ads have odd tokens like 1000 As in a row
        if len(tok) <= maxTokenLength:
            # yield tok
            yield tok.replace('\t', ' ')

### The URI + index mechanism isn't good enough to recover data when there are multiple sentences per URI
### which can occur from (a) title + body as separate documents (b) sentence breaking of body
### The intra-sentence space rows won't be maintained by reconstructTuple + groupBy
### Proposed future workarounds:
### (a) sentence-specific temporary URIs ending /processed/title or /processed/1
### (b) and/or sentence-specific indexing 1.1, 1.2, ... 1.N, 2.1, etc. could work

def reconstructTuple(tabsep):
    """Works only for old multi-line invocation of crf_test, not crf_test_b64"""
    fields = tabsep.split('\t')
    try:
        uri = fields[-3]
        idx = fields[-2]
    except:
        uri = "endOfDocument"
        idx = "0"
    return (uri, [idx] + fields)

def fakeFindBestMatch(words):
    if 'blue' in words:
        return 'blue'
    else:
        return 'NONE'

def alignToControlledVocab(harvested, vocabs):
    """harvested is a dict matching category 'eyeColor', 'hairType' to a function"""
    try:
        category = harvested['category']
        # f = vocabs[category]
        f = fakeFindBestMatch
        words = harvested['words']
        try:
            harvested['bestMatch'] = f(words)
        except Exception as e:
            return "sm.findBestMatch error:" + str(e) + ":" + str(words)
        return harvested
    except Exception as e:
        return str(e)
    return None

def vectorToUTF8(v, debug=False):
    "unicode only"

    def rowToUnicode(r):
        try:
            if isinstance(r, list):
                return u"\t".join([unicode(x) for x in r])
            else:
                return unicode(r)
        except:
            print >> sys.stderr, "error in rowToUnicode"
            return u""

    rows = []
    if v[-1] == u"":
        pass
    else:
        # print "appending1"
        v.append(u"")

    for r in v:
        rows.append(rowToUnicode(r))

    result = u"\n".join(rows)
    # result now a unicode object
    # here is the only place where we convert to UTF8
    return result.encode('utf-8')

### HISTORICAL
def computeSpans(v, verbose=False, indexed=False):
    # extract the data and write the result as vector
    currentLabel = None
    currentTokens = []
    spans = []
    def addSpan(u, l, words):
        spans.append( {"uri": u, "category": l, "words": " ".join(words) } )
        if verbose:
            print >> sys.stderr, "  Added %s" % (spans[-1],)

    uri = 'bogus'
    currentUri = None
    for row in v:
        if (len(row) <= 1):
            # blank/empty line: expecting "" but might be []
            if currentLabel and currentTokens:
                addSpan(currentUri, currentLabel, currentTokens)
            currentUri = None
            continue
        if (len(row) >= 4):
            # a typical row
            token = row[0]
            uri = row[-3] if indexed else row[-2]
            label = row[-1]
            if verbose:
                print >> sys.stderr, "Typical row: token %r uri %r: crflabel %r" % (token, uri, label)
            # now process this row
            if label == "O":
                # unlabeled row
                if currentLabel:
                    # so this concludes span in progress
                    addSpan(uri, currentLabel, currentTokens)
                    currentLabel = None
                    currentTokens = []
                else:
                    pass
            else:
                # Labeled row
                if label == currentLabel:
                    # continue span in progress
                    currentTokens.append(token)
                elif currentLabel and label != currentLabel:
                    # span/span boundary
                    # first conclude old one
                    addSpan(uri, currentLabel, currentTokens)
                    currentLabel = None
                    currentTokens = []
                    # then begin new one
                    currentLabel = label
                    currentTokens = [token]
                elif not currentLabel:
                    # begin novel span
                    currentLabel = label
                    currentTokens = [token]
                else:
                    raise Exception("Unexpected file structure")
        currentUri = uri

    if currentLabel and currentTokens:
        # Input ended without blank line after last marked span, so hallucinate one
        addSpan(uri, currentLabel, currentTokens)

    # Publish results
    return spans

def crfsmall(sc, input, output,
            limit=None, location='hdfs', outputFormat="text", numPartitions=None):

    crfConfigDir = os.path.join(os.path.dirname(__file__), "data/config")
    featureListFilename = os.path.join(crfConfigDir, "features.hair-eye")
    crfExecutable = os.path.join(os.path.dirname(__file__), "bin/crf_test_filter.sh")
    crfModelFilename = os.path.join(crfConfigDir, "dig-hair-eye-train.model")
    sc.addFile(crfExecutable)
    sc.addFile(crfModelFilename)

    # crfConfigDir = os.path.join(os.path.dirname(__file__), "data/config")
    # def cpath(n):
    #     return os.path.join(crfConfigDir, n)

    # smEyeColor = HybridJaccard(ref_path=cpath("eyeColor_reference_wiki.txt"),
    #                            config_path=cpath("eyeColor_config.txt"))
    # smHairColor = HybridJaccard(ref_path=cpath("hairColor_reference_wiki.txt"),
    #                             config_path=cpath("hairColor_config.txt"))
    # print smEyeColor, smHairColor

    # hypothesis1: data fetched this way prompts the lzo compression error
    # hypothesis2: but it doesn't matter, error is just a warning

    rdd_crfl = sc.sequenceFile(input)
    rdd_crfl.setName('rdd_crfl')

    if limit:
        rdd_crfl = sc.parallelize(rdd_crfl.take(limit))
    if numPartitions:
        rdd_crfl = rdd_crfl.repartition(numPartitions)

    rdd_json = rdd_crfl.mapValues(lambda x: json.loads(x))
    rdd_json.setName('rdd_json')
    # rdd_json.persist()

    rdd_texts = rdd_json.mapValues(lambda x: (textTokens(extract_body(x)), textTokens(extract_title(x))))
    rdd_texts.setName('rdd_texts')

    c = crf_features.CrfFeatures(featureListFilename)
    SEPARATOR = '&nbsp;',

    def makeMatrix(c, uri, bodyTokens, titleTokens):
        b = c.computeFeatMatrix(bodyTokens, False, addLabels=False, addIndex=False)
        s = c.computeFeatMatrix([SEPARATOR, ""], False, addLabels=False, addIndex=False)
        t = c.computeFeatMatrix(titleTokens, False, addLabels=False, addIndex=False)
        idx = 1
        for row in b:
            if row == u"":
                pass
            else:
                label = uri + "/%05d/%05d" % (0, idx)
                row.append(label)
                idx += 1
        idx = 1
        for row in s:
            if row == u"":
                pass
            else:
                label = uri + "/%05d/%05d" % (1, idx)
                row.append(label)
                idx += 1
        idx = 1
        for row in t:
            if row == u"":
                pass
            else:
                label = uri + "/%05d/%05d" % (2, idx)
                row.append(label)
                idx += 1
        # might be b[0:-1] + s[0:-1] + t?
        return b[0:-1] + s[0:-1] + t

    rdd_features = rdd_texts.map(lambda x: makeMatrix(c, x[0], x[1][0], x[1][1]))
    rdd_features.setName('rdd_features')
    # rdd_features.persist()

    # unicode/string representation of the feature matrix
    rdd_vector = rdd_features.map(lambda x: vectorToUTF8(x))
    rdd_vector.setName('rdd_vector')

    # all strings concatenated together, then base64 encoded into one input for crf_test
    rdd_pipeinput = sc.parallelize([b64encode(rdd_vector.reduce(lambda a,b: a+b))])
    rdd_pipeinput.setName('rdd_pipeinput')

    executable = SparkFiles.get(os.path.basename(crfExecutable))
    model = SparkFiles.get(os.path.basename(crfModelFilename))

    cmd = "%s %s" % (executable, model)

    # this result is base64 encoded
    rdd_crfoutput = rdd_pipeinput.pipe(cmd)
    rdd_crfoutput.setName('rdd_crfoutput')

    rdd_base64decode = rdd_crfoutput.map(lambda x: b64decode(x))
    ### There may be a need to utf8 decode this data ###
    ### There are values like \xed\xa0\xbd which might be a broken emoji
    # 1. break into physical lines
    # 2. turn each line into its own spark row
    # 3. drop any inter-document empty string markers
    rdd_lines = rdd_base64decode.map(lambda x: x.split("\n")).flatMap(lambda l: l).filter(lambda x: len(x)>1)

    def processOneLine(l):
        return l.split("\t")

    rdd_triples = rdd_lines.map(lambda l: processOneLine(l))
    rdd_triples.saveAsTextFile('out_rdd_triples')

    def organizeByOrigDoc(triple):
        try:
            (word, uri, label) = triple
            (parentUri, docId, wordId) = uri.rsplit('/', 2)
            return ( (parentUri, docId), (wordId, word, label) )
        except Exception as e:
            print "Can't destructure %r" % [triple]
            return ()

    rdd_reorg = rdd_triples.map(lambda l: organizeByOrigDoc(l))
    rdd_reorg.saveAsTextFile('out_rdd_reorg')

    rdd_sorted = rdd_reorg.sortByKey()
    rdd_sorted.saveAsTextFile('out_rdd_sorted')

    # each (parentUri, docId) has a sequence of (wordId, word, label)
    # we want to consider them in order (by wordId)

    rdd_grouped = rdd_sorted.groupByKey()

    def harvest(seq):
        allSpans = []
        lastIndex = -2
        lastLabel = None
        currentSpan = []
        for (wordId, word, label) in seq:
            currentIndex = int(wordId)
            if lastIndex+1 == currentIndex and lastLabel == label:
                # continuing current span
                currentSpan.append( (currentIndex, word, label) )
            else:
                # end current span
                if currentSpan:
                    allSpans.append(currentSpan)
                # begin new span
                currentSpan = [ (currentIndex, word, label) ]
                lastLabel = label
            lastIndex = currentIndex

        # end current span
        if currentSpan:
            allSpans.append(currentSpan)

        result = []
        for span in allSpans:
            words = []
            spanLabel = None
            for (wordIdx, word, label) in span:
                spanLabel = label
                words.append(word)
            result.append( (' '.join(words), spanLabel) )
        return result

    # ( (parentUri, docId), [ (words1, category1), (words2, category2), ... ]
    rdd_harvest = rdd_grouped.mapValues(lambda s: harvest(s))
    rdd_harvest.saveAsTextFile('out_rdd_harvest')

    # rdd_flat = rdd_harvest.flatMap(lambda r: [ (e[0][0], e[1]) for e in r ])
    # rdd_flat = rdd_harvest.map(lambda r: (r[0][0], r[1]))

    # parentUri -> (words, category)
    # we use .distinct() because (e.g.) both title and body might have the same feature
    rdd_flat = rdd_harvest.map(lambda r: (r[0][0], r[1])).flatMapValues(lambda x: x).distinct()
    rdd_flat.saveAsTextFile('out_rdd_flat')

    # rdd_aligned = rdd_flat.mapValues(lambda x: alignToControlledVocab(x, {"eyeColor": smEyeColor.findBestMatch,
    #                                                                       "hairType": smHairColor.findBestMatch}))
    exit(1)

    rdd_final = rdd_harvest

    if outputFormat == "sequence":
        rdd_final.saveAsSequenceFile(output)
    elif outputFormat == "text":
        rdd_final.saveAsTextFile(output)
    else:
        raise RuntimeError("Unrecognized output format: %s" % outputFormat)

def main(argv=None):
    '''this is called if run from command line'''
    parser = argparse.ArgumentParser()
    parser.add_argument('-i','--input', required=True)
    parser.add_argument('-o','--output', required=True)
    parser.add_argument('-p','--numPartitions', required=False, default=None, type=int)
    parser.add_argument('-l','--limit', required=False, default=None, type=int)
    parser.add_argument('-v','--verbose', required=False, help='verbose', action='store_true')
    args=parser.parse_args()

    location = "hdfs"
    try:
        if "avatar" in platform.node():
            location = "local"
    except:
        pass
    try:
        if "avatar" in socket.gethostname():
            location = "local"
    except:
        pass
    print "### location %s" % location

    if not args.numPartitions:
        if location == "local":
            args.numPartitions = 2
        elif location == "hdfs":
            args.numPartitions = 50

    sc = SparkContext(appName="crfsmall")
    crfsmall(sc, args.input, args.output,
             limit=args.limit,
             location=location,
             outputFormat="text",
             numPartitions=args.numPartitions)

# call main() if this is run as standalone
if __name__ == "__main__":
    sys.exit(main())
<|endoftext|>


<commit_before>#!/usr/bin/env python3
import argparse
import collections
import concurrent.futures
import errno
import logging
import json
import os
import random
import re
import shutil
import socket
import ssl
import time

import jinja2
import requests

from lxml import etree, html

PARALLELISM = 16
USER_AGENT = "Mozilla/5.0 Firefox/35.0 compatible HTTPSWatchAU Bot https://httpswatch.com.au"
ANALYTICS = """
"""
META_XPATH = etree.XPath(
    "//meta[not(ancestor::noscript) and re:test(@http-equiv, \"^refresh$\", \"i\")]",
    namespaces={"re": "http://exslt.org/regular-expressions"}
)
MIXED_CONTENT_XPATH = etree.XPath(
    """
    //link[@ref="stylesheet" and starts-with(@href, "http://")] |
    //script[starts-with(@src, "http://")] |
    //img[not(ancestor::noscript) and starts-with(@src, "http://")]
    """
)

log = logging.getLogger("check_https")

class SiteInfo:

    def __init__(self):
        self.domain = None
        self.ssllabs_grade = None
        self.secure_connection_works = None
        self.can_load_https_page = None
        self.mixed_content = None
        self.sts = None
        self.https_redirects_to_http = None
        self.http_redirects_to_https = None
        self.checks = []

    def new_check(self):
        c = Check()
        self.checks.append(c)
        return c

class Check:

    def __init__(self):
        pass

    def succeed(self, desc):
        self.failed = False
        self.icon = "good"
        self.description = desc

    def fail(self, desc):
        self.failed = True
        self.icon = "bad"
        self.description = desc

class Not200(Exception):

    def __init__(self, status):
        super(Exception, self).__init__()
        self.status = status

def fetch_through_redirects(url):
    tree = None
    while True:
        cont = False
        resp = requests.get(
            url,
            verify="moz-certs.pem",
            headers={"User-Agent": USER_AGENT},
            timeout=10,
            stream=True,
        )
        try:
            if resp.status_code != 200:
                raise Not200(resp.status_code)
            # Convince urllib3 to decode gzipped pages.
            resp.raw.decode_content = True
            tree = html.parse(resp.raw)
        finally:
            resp.close()
        # Check for sneaky <meta> redirects.
        for meta in META_XPATH(tree):
            m = re.match("0;\s*url=['\"](.+?)['\"]", meta.get("content"))
            if m is not None:
                url = m.groups()[0]
                cont = True
                break
        if not cont:
            break
    return resp, tree

def has_mixed_content(tree):
    s = MIXED_CONTENT_XPATH(tree)
    return len(s) >= 1

def check_secure_connection(info):
    # Guilty until proven innocent.
    info.secure_connection_works = False
    good_connection = info.new_check()
    try:
        addrs = socket.getaddrinfo(info.domain, 443, socket.AF_INET, proto=socket.IPPROTO_TCP)
    except socket.gaierror:
        good_connection.fail("DNS lookup failed.")
        return
    addr_info = random.choice(addrs)
    sock = socket.socket(addr_info[0], addr_info[1], addr_info[2])
    sock.settimeout(10)
    context = ssl.SSLContext(ssl.PROTOCOL_SSLv23)
    context.options |= ssl.OP_NO_SSLv2 | ssl.OP_NO_SSLv3
    # Some platforms (OS X) do not have OP_NO_COMPRESSION
    context.options |= getattr(ssl, "OP_NO_COMPRESSION", 0)
    context.verify_mode = ssl.CERT_REQUIRED
    context.check_hostname = True
    context.load_verify_locations("moz-certs.pem")
    secure_sock = context.wrap_socket(sock, server_hostname=info.domain)
    try:
        secure_sock.connect(addr_info[4])
    except ConnectionRefusedError:
        good_connection.fail("Nothing is listening on port 443.")
        return
    except socket.timeout:
        good_connection.fail("connect() to port 443 times out.")
        return
    except ssl.SSLError as e:
        if e.reason == "CERTIFICATE_VERIFY_FAILED":
            desc = "Certificate not trusted by Mozilla cert store."
        else:
            desc = "A TLS-related error ({}) occurs when trying to connect.".format(e.reason)
        good_connection.fail(desc)
        return
    except ssl.CertificateError:
        good_connection.fail("Certificate hostname verification fails.")
        return
    except OSError as e:
        error_name = errno.errorcode[e.errno]
        good_connection.fail("connect() returns with error {}.".format(error_name))
        return
    finally:
        secure_sock.close()
    msg = "A verified TLS connection can be established. "
    if info.ssllabs_grade is not None:
        msg += "SSL Labs grade is " + info.ssllabs_grade + "."
    else:
        msg += "(SSL Labs report)"
    info.secure_connection_works = True
    good_connection.succeed(msg.format(info.domain))

def check_https_page(info):
    info.can_load_https_page = False
    https_load = info.new_check()
    try:
        resp, tree = fetch_through_redirects("https://{}".format(info.domain))
        info.https_redirects_to_http = resp.url.startswith("http://")
        if info.https_redirects_to_http:
            https_load.fail("The HTTPS site redirects to HTTP.")
            return
        info.mixed_content = tree is not None and has_mixed_content(tree)
        if info.mixed_content:
            https_load.fail("The HTML page loaded over HTTPS has mixed content.")
            return
        good_sts = info.new_check()
        sts = resp.headers.get("Strict-Transport-Security")
        if sts is not None:
            m = re.search("max-age=(\d+)", sts)
            if m is not None:
                info.sts = int(m.group(1))
                if info.sts >= 2592000:
                    good_sts.succeed("Strict-Transport-Security header is set with a long max-age directive.")
                else:
                    good_sts.fail("Strict-Transport-Security header is set but the max-age is less than 30 days.")
            else:
                good_sts.fail("Strict-Transport-Security header doesn't contain a max-age directive.")
        else:
            good_sts.fail("Strict-Transport-Security header is not set.")
    except requests.Timeout:
        https_load.fail("Requesting HTTPS page times out.")
        return
    except Not200 as e:
        https_load.fail("The HTTPS site returns an error status ({}) on request.".format(e.status))
        return
    info.can_load_https_page = True
    https_load.succeed("A page can be successfully fetched over HTTPS.")

def check_http_page(info):
    http_redirect = info.new_check()
    try:
        resp, tree = fetch_through_redirects("http://{}".format(info.domain))
        info.http_redirects_to_https = resp.url.startswith("https://")
        if info.http_redirects_to_https:
            http_redirect.succeed("HTTP site redirects to HTTPS.")
        else:
            http_redirect.fail("HTTP site doesn't redirect to HTTPS.")
    except requests.Timeout:
        http_redirect.fail("The HTTP site times out.")
        return
    except Not200 as e:
        http_redirect.fail("The HTTP site returns an error status ({}) on request.".format(e.status))
        return

def check_site(site):
    log.info("Checking {}".format(site["domain"]))
    info = SiteInfo()
    info.domain = site["domain"]
    info.ssllabs_grade = site.get("ssllabs_grade")

    check_secure_connection(info)
    if not info.secure_connection_works:
        log.info("Couldn't connect securely to {}, so aborting further checks.".format(info.domain))
        return info

    check_https_page(info)
    if not info.can_load_https_page:
        log.info("Couldn't load HTTPS page for {}, so aborting further checks.".format(info.domain))
        return info

    check_http_page(info)

    return info

def set_site_template_data_from_info(site, info):
    site["checks"] = info.checks

    if (not info.secure_connection_works or
            info.https_redirects_to_http or
            info.mixed_content):
        status = "bad"
    elif (not info.can_load_https_page or
          not info.http_redirects_to_https or
          info.sts is None or info.sts < 2592000):
        status = "mediocre"
    else:
        status = "good"
    site["status"] = status

def regenerate_everything(ssllabs_grades_file):
    with open("config/meta.json", "r", encoding="utf-8") as fp:
        meta = json.load(fp)
    ssllabs_grades = {}
    if ssllabs_grades_file is not None:
        with open(ssllabs_grades_file, "r", encoding="utf-8") as fp:
            ssllabs_grades = json.load(fp)
    futures = []
    executor = concurrent.futures.ThreadPoolExecutor(max_workers=PARALLELISM)
    for listing in meta["listings"]:
        if "external" in listing:
            continue
        with open("config/{}.json".format(listing["shortname"]), encoding="utf-8") as fp:
            listing["data"] = json.load(fp)
        for cat in listing["data"]["categories"]:
            for site in cat["sites"]:
                domain = site["domain"]
                if domain in ssllabs_grades:
                    site["ssllabs_grade"] = ssllabs_grades[domain]
                futures.append(executor.submit(check_site, site))
    with executor:
        while True:
            done, not_done = concurrent.futures.wait(futures, timeout=1)
            print("{}/{}".format(len(done), len(done) + len(not_done)))
            if not not_done:
                break
        infos = {}
        for f in futures:
            info = f.result()
            infos[info.domain] = info

    for listing in meta["listings"]:
        if "external" in listing:
            continue
        for cat in listing["data"]["categories"]:
            cat_status = collections.Counter()
            for site in cat["sites"]:
                info = infos[site["domain"]]
                set_site_template_data_from_info(site, info)
                cat_status[site["status"]] += 1
            cat["counts"] = cat_status

    return meta

def encode_check(o):
    if not isinstance(o, Check):
        raise TypeError
    return o.__dict__

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--cached", action="store_true")
    parser.add_argument("--ssllabs")

    args = parser.parse_args()

    if args.cached:
        with open("cache.json", "r", encoding="utf-8") as fp:
            meta = json.load(fp)
    else:
        meta = regenerate_everything(args.ssllabs)
        with open("cache.json", "w", encoding="utf-8") as fp:
            json.dump(meta, fp, default=encode_check)

    # Write out results.
    env = jinja2.Environment(loader=jinja2.FileSystemLoader("templates"))
    env.globals["analytics"] = ANALYTICS
    env.globals["update_time"] = time.strftime("%Y-%m-%d %H:%MZ", time.gmtime())
    env.globals["meta"] = meta
    try:
        os.mkdir("out")
    except FileExistsError:
        pass
    with open("out/about.html", "w", encoding="utf-8") as fp:
        fp.write(env.get_template("about.html.jinja").render())
    listing_tmp = env.get_template("listing.html.jinja")
    for listing in meta["listings"]:
        if "external" in listing:
            continue
        out_fn = "out/{}.html".format(listing["shortname"])
        with open(out_fn, "w", encoding="utf-8") as fp:
            fp.write(listing_tmp.render(listing=listing))
        if meta["default_page"] == listing["shortname"]:
            shutil.copy(out_fn, "out/index.html")

if __name__ == "__main__":
    main()
<commit_msg>handle connection errors and untrusted certificates<commit_after>#!/usr/bin/env python3
import argparse
import collections
import concurrent.futures
import errno
import logging
import json
import os
import random
import re
import shutil
import socket
import ssl
import time

import jinja2
import requests

from lxml import etree, html

PARALLELISM = 16
USER_AGENT = "Mozilla/5.0 Firefox/35.0 compatible HTTPSWatchAU Bot https://httpswatch.com.au"
ANALYTICS = """
"""
META_XPATH = etree.XPath(
    "//meta[not(ancestor::noscript) and re:test(@http-equiv, \"^refresh$\", \"i\")]",
    namespaces={"re": "http://exslt.org/regular-expressions"}
)
MIXED_CONTENT_XPATH = etree.XPath(
    """
    //link[@ref="stylesheet" and starts-with(@href, "http://")] |
    //script[starts-with(@src, "http://")] |
    //img[not(ancestor::noscript) and starts-with(@src, "http://")]
    """
)

log = logging.getLogger("check_https")

class SiteInfo:

    def __init__(self):
        self.domain = None
        self.ssllabs_grade = None
        self.secure_connection_works = None
        self.can_load_https_page = None
        self.mixed_content = None
        self.sts = None
        self.https_redirects_to_http = None
        self.http_redirects_to_https = None
        self.checks = []

    def new_check(self):
        c = Check()
        self.checks.append(c)
        return c

class Check:

    def __init__(self):
        pass

    def succeed(self, desc):
        self.failed = False
        self.icon = "good"
        self.description = desc

    def fail(self, desc):
        self.failed = True
        self.icon = "bad"
        self.description = desc

class Not200(Exception):

    def __init__(self, status):
        super(Exception, self).__init__()
        self.status = status

class ConnectionError(Exception):

    def __init__(self, status):
        super(Exception, self).__init__()
        self.status = status

class InvalidCert(Exception):

    def __init__(self, status):
        super(Exception, self).__init__()
        self.status = status

def fetch_through_redirects(url):
    tree = None
    while True:
        cont = False
        try:
            resp = requests.get(
                url,
                verify="moz-certs.pem",
                headers={"User-Agent": USER_AGENT},
                timeout=10,
                stream=True,
            )
        except requests.exceptions.ConnectionError:
            raise ConnectionError('Connection reset by peer')
        except requests.exceptions.SSLError as e:
            # certificate presentied is not trusted?
            # we don't actually need to do anything different, but we
            # have to catch the exception
            raise InvalidCert('Untrusted certificate')
        try:
            if resp.status_code != 200:
                raise Not200(resp.status_code)
            # Convince urllib3 to decode gzipped pages.
            resp.raw.decode_content = True
            tree = html.parse(resp.raw)
        finally:
            resp.close()
        # Check for sneaky <meta> redirects.
        for meta in META_XPATH(tree):
            m = re.match("0;\s*url=['\"](.+?)['\"]", meta.get("content"))
            if m is not None:
                url = m.groups()[0]
                cont = True
                break
        if not cont:
            break
    return resp, tree

def has_mixed_content(tree):
    s = MIXED_CONTENT_XPATH(tree)
    return len(s) >= 1

def check_secure_connection(info):
    # Guilty until proven innocent.
    info.secure_connection_works = False
    good_connection = info.new_check()
    try:
        addrs = socket.getaddrinfo(info.domain, 443, socket.AF_INET, proto=socket.IPPROTO_TCP)
    except socket.gaierror:
        good_connection.fail("DNS lookup failed.")
        return
    addr_info = random.choice(addrs)
    sock = socket.socket(addr_info[0], addr_info[1], addr_info[2])
    sock.settimeout(10)
    context = ssl.SSLContext(ssl.PROTOCOL_SSLv23)
    context.options |= ssl.OP_NO_SSLv2 | ssl.OP_NO_SSLv3
    # Some platforms (OS X) do not have OP_NO_COMPRESSION
    context.options |= getattr(ssl, "OP_NO_COMPRESSION", 0)
    context.verify_mode = ssl.CERT_REQUIRED
    context.check_hostname = True
    context.load_verify_locations("moz-certs.pem")
    secure_sock = context.wrap_socket(sock, server_hostname=info.domain)
    try:
        secure_sock.connect(addr_info[4])
    except ConnectionRefusedError:
        good_connection.fail("Nothing is listening on port 443.")
        return
    except socket.timeout:
        good_connection.fail("connect() to port 443 times out.")
        return
    except ssl.SSLError as e:
        if e.reason == "CERTIFICATE_VERIFY_FAILED":
            desc = "Certificate not trusted by Mozilla cert store."
        else:
            desc = "A TLS-related error ({}) occurs when trying to connect.".format(e.reason)
        good_connection.fail(desc)
        return
    except ssl.CertificateError:
        good_connection.fail("Certificate hostname verification fails.")
        return
    except OSError as e:
        error_name = errno.errorcode[e.errno]
        good_connection.fail("connect() returns with error {}.".format(error_name))
        return
    finally:
        secure_sock.close()
    msg = "A verified TLS connection can be established. "
    if info.ssllabs_grade is not None:
        msg += "SSL Labs grade is " + info.ssllabs_grade + "."
    else:
        msg += "(SSL Labs report)"
    info.secure_connection_works = True
    good_connection.succeed(msg.format(info.domain))

def check_https_page(info):
    info.can_load_https_page = False
    https_load = info.new_check()
    try:
        resp, tree = fetch_through_redirects("https://{}".format(info.domain))
        info.https_redirects_to_http = resp.url.startswith("http://")
        if info.https_redirects_to_http:
            https_load.fail("The HTTPS site redirects to HTTP.")
            return
        info.mixed_content = tree is not None and has_mixed_content(tree)
        if info.mixed_content:
            https_load.fail("The HTML page loaded over HTTPS has mixed content.")
            return
        good_sts = info.new_check()
        sts = resp.headers.get("Strict-Transport-Security")
        if sts is not None:
            m = re.search("max-age=(\d+)", sts)
            if m is not None:
                info.sts = int(m.group(1))
                if info.sts >= 2592000:
                    good_sts.succeed("Strict-Transport-Security header is set with a long max-age directive.")
                else:
                    good_sts.fail("Strict-Transport-Security header is set but the max-age is less than 30 days.")
            else:
                good_sts.fail("Strict-Transport-Security header doesn't contain a max-age directive.")
        else:
            good_sts.fail("Strict-Transport-Security header is not set.")
    except requests.Timeout:
        https_load.fail("Requesting HTTPS page times out.")
        return
    except Not200 as e:
        https_load.fail("The HTTPS site returns an error status ({}) on request.".format(e.status))
        return
    except ConnectionError as e:
        https_load.fail("The HTTPS site returns a connection error: ({}).".format(e.status))
        return
    except InvalidCert as e:
        https_load.fail("The HTTPS site returns an invalid certificate: ({}).".format(e.status))
        return
    info.can_load_https_page = True
    https_load.succeed("A page can be successfully fetched over HTTPS.")

def check_http_page(info):
    http_redirect = info.new_check()
    try:
        resp, tree = fetch_through_redirects("http://{}".format(info.domain))
        info.http_redirects_to_https = resp.url.startswith("https://")
        if info.http_redirects_to_https:
            http_redirect.succeed("HTTP site redirects to HTTPS.")
        else:
            http_redirect.fail("HTTP site doesn't redirect to HTTPS.")
    except requests.Timeout:
        http_redirect.fail("The HTTP site times out.")
        return
    except Not200 as e:
        http_redirect.fail("The HTTP site returns an error status ({}) on request.".format(e.status))
        return
    except ConnectionError as e:
        https_load.fail("The HTTPS site returns a connection error: ({}).".format(e.status))
        return

def check_site(site):
    log.info("Checking {}".format(site["domain"]))
    info = SiteInfo()
    info.domain = site["domain"]
    info.ssllabs_grade = site.get("ssllabs_grade")

    check_secure_connection(info)
    if not info.secure_connection_works:
        log.info("Couldn't connect securely to {}, so aborting further checks.".format(info.domain))
        return info

    check_https_page(info)
    if not info.can_load_https_page:
        log.info("Couldn't load HTTPS page for {}, so aborting further checks.".format(info.domain))
        return info

    check_http_page(info)

    return info

def set_site_template_data_from_info(site, info):
    site["checks"] = info.checks

    if (not info.secure_connection_works or
            info.https_redirects_to_http or
            info.mixed_content):
        status = "bad"
    elif (not info.can_load_https_page or
          not info.http_redirects_to_https or
          info.sts is None or info.sts < 2592000):
        status = "mediocre"
    else:
        status = "good"
    site["status"] = status

def regenerate_everything(ssllabs_grades_file):
    with open("config/meta.json", "r", encoding="utf-8") as fp:
        meta = json.load(fp)
    ssllabs_grades = {}
    if ssllabs_grades_file is not None:
        with open(ssllabs_grades_file, "r", encoding="utf-8") as fp:
            ssllabs_grades = json.load(fp)
    futures = []
    executor = concurrent.futures.ThreadPoolExecutor(max_workers=PARALLELISM)
    for listing in meta["listings"]:
        if "external" in listing:
            continue
        with open("config/{}.json".format(listing["shortname"]), encoding="utf-8") as fp:
            listing["data"] = json.load(fp)
        for cat in listing["data"]["categories"]:
            for site in cat["sites"]:
                domain = site["domain"]
                if domain in ssllabs_grades:
                    site["ssllabs_grade"] = ssllabs_grades[domain]
                futures.append(executor.submit(check_site, site))
    with executor:
        while True:
            done, not_done = concurrent.futures.wait(futures, timeout=1)
            print("{}/{}".format(len(done), len(done) + len(not_done)))
            if not not_done:
                break
        infos = {}
        for f in futures:
            info = f.result()
            infos[info.domain] = info

    for listing in meta["listings"]:
        if "external" in listing:
            continue
        for cat in listing["data"]["categories"]:
            cat_status = collections.Counter()
            for site in cat["sites"]:
                info = infos[site["domain"]]
                set_site_template_data_from_info(site, info)
                cat_status[site["status"]] += 1
            cat["counts"] = cat_status

    return meta

def encode_check(o):
    if not isinstance(o, Check):
        raise TypeError
    return o.__dict__

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--cached", action="store_true")
    parser.add_argument("--ssllabs")

    args = parser.parse_args()

    if args.cached:
        with open("cache.json", "r", encoding="utf-8") as fp:
            meta = json.load(fp)
    else:
        meta = regenerate_everything(args.ssllabs)
        with open("cache.json", "w", encoding="utf-8") as fp:
            json.dump(meta, fp, default=encode_check)

    # Write out results.
    env = jinja2.Environment(loader=jinja2.FileSystemLoader("templates"))
    env.globals["analytics"] = ANALYTICS
    env.globals["update_time"] = time.strftime("%Y-%m-%d %H:%MZ", time.gmtime())
    env.globals["meta"] = meta
    try:
        os.mkdir("out")
    except FileExistsError:
        pass
    with open("out/about.html", "w", encoding="utf-8") as fp:
        fp.write(env.get_template("about.html.jinja").render())
    listing_tmp = env.get_template("listing.html.jinja")
    for listing in meta["listings"]:
        if "external" in listing:
            continue
        out_fn = "out/{}.html".format(listing["shortname"])
        with open(out_fn, "w", encoding="utf-8") as fp:
            fp.write(listing_tmp.render(listing=listing))
        if meta["default_page"] == listing["shortname"]:
            shutil.copy(out_fn, "out/index.html")

if __name__ == "__main__":
    main()
<|endoftext|>