<commit_before># -*- coding: utf-8 -*-
import unittest
from unittest import mock
import os.path
import urllib.error
from urllib.parse import urlsplit, urlencode, urljoin

from openid import message, xrds
from openid.consumer import discover
from . import support

@mock.patch('urllib.request.urlopen', support.urlopen)
@support.gentests
class Failure(unittest.TestCase):
    data = [
        ('network_error', ('http://network.error/',)),
        ('not_found', ('/404',)),
        ('header_found', ('/?' + urlencode({'header': 'X-XRDS-Location: http://unittest/404'}),)),
        ('server_error', ('/?status=500',)),
    ]

    def _test(self, path):
        url = urljoin('http://unittest', path)
        self.assertRaises(urllib.error.URLError, discover.discover, url)

@mock.patch('urllib.request.urlopen', support.urlopen)
class Discovery(unittest.TestCase):
    def test_unicode(self):
        """
        Check page with unicode and HTML entities
        """
        services = discover.discoverall('http://unittest/unicode.html')
        self.assertEqual(len(services), 0)

    def test_unicode_undecodable_html2(self):
        """
        Check page with unicode and HTML entities that can not be decoded
        but xrds document is found before it matters
        """
        with open(os.path.join(support.DATAPATH, 'unicode3.html'), encoding='utf-8') as f:
            self.assertRaises(UnicodeDecodeError, f.read)
        services = discover.discoverall('http://unittest/unicode3.html')
        self.assertEqual(len(services), 1)

    def test_noOpenID(self):
        services = discover.discoverall('http://unittest/junk.txt')
        self.assertFalse(services)

    def test_yadisEmpty(self):
        services = discover.discoverall('http://unittest/yadis_0entries.xrds')
        self.assertFalse(services)

    def test_html_yadis_empty(self):
        # The HTML document contains OpenID links but also refers to an empty Yadis
        # document which we should prefer, by the Yadis spec.
        services = discover.discoverall('http://unittest/openid_and_yadis.html')
        self.assertFalse(services)

    def test_xrds_with_header(self):
        # Even if we got a valid XRDS document but with a link to another XRDS
        # location we should prefer that location, by the Yadis spec.
        location_url = 'http://unittest/openid2_xrds.xrds'
        params = {'header': 'X-XRDS-Location: %s' % location_url}
        url = 'http://unittest/openid_1_and_2_xrds.xrds?' + urlencode(params)
        services = discover.discoverall(url)
        self.assertEqual(support.urlopen.request.get_full_url(), location_url)

    def test_fragment(self):
        url = 'http://unittest/openid.html'
        service = discover.discover(url + '#fragment')
        self.assertEqual(service.claimed_id, url)

    def test_add_protocol(self):
        url = 'unittest:8000/'
        discover.discoverall(url)
        self.assertEqual(support.urlopen.request.get_full_url(), 'http://' + url)

    def test_wrong_protocol(self):
        url = 'ssh://unittest/'
        self.assertRaises(ValueError, discover.discoverall, url)

    def test_html1And2(self):
        url = 'http://unittest/openid_1_and_2.html'
        services = discover.discoverall(url)
        self.assertEqual(len(services), 2)
        for s in services:
            self.assertEqual(s.server_url, 'http://www.myopenid.com/server')
            self.assertEqual(s.local_id, 'http://smoker.myopenid.com/')
            self.assertEqual(s.claimed_id, url)

    def test_service_sort(self):
        services = discover.discoverall('http://unittest/multiple_services.xrds')
        self.assertEqual(len(services), 3)
        self.assertTrue(discover.OPENID_IDP_2_0_TYPE in services[0].types)
        self.assertTrue(discover.OPENID_2_0_TYPE in services[1].types)
        self.assertTrue(discover.OPENID_1_0_TYPE in services[2].types)

    def test_redirected_claimed_id(self):
        claimed_id = 'http://unittest/openid2_xrds.xrds'
        url = 'http://unittest/200.txt?' + urlencode({'redirect': claimed_id})
        service = discover.discover(url)
        self.assertEqual(claimed_id, service.claimed_id)

    def test_xri_idp(self):
        service = discover.discover('=iname.idp')
        self.assertEqual(service.server_url, 'http://www.livejournal.com/openid/server.bml')

    def test_two_services(self):
        services = discover.discoverall('=twoservices')
        self.assertEqual(len(services), 2)
        self.assertTrue(services[0].local_id, 'http://smoker.myopenid.com/')
        self.assertTrue(services[1].local_id, 'http://frank.livejournal.com/')

    def test_xriNoCanonicalID(self):
        self.assertRaises(xrds.XRDSError, discover.discoverall, '=iname*empty')

@mock.patch('urllib.request.urlopen', support.urlopen)
@support.gentests
class Services(unittest.TestCase):
    data = [
        # arguments: url, types, claimed_id, local_id, canonical_id
        # last four arguments can be None, in which case the url is used instead
        ('no_delegate', ('http://unittest/openid_no_delegate.html', ['1.1'], None, None)),
        ('html1', ('http://unittest/openid.html', ['1.1'], None, 'http://smoker.myopenid.com/')),
        ('html2', ('http://unittest/openid2.html', ['2.0'], None, 'http://smoker.myopenid.com/')),
        ('yadis1_no_delegate', ('http://unittest/yadis_no_delegate.xrds', ['1.0'], None, None)),
        ('yadis2_no_local_id', ('http://unittest/openid2_xrds_no_local_id.xrds', ['2.0'], None, None)),
        ('yadis2', ('http://unittest/openid2_xrds.xrds', ['2.0'], None, 'http://smoker.myopenid.com/')),
        ('yadis2_op', ('http://unittest/yadis_idp.xrds', ['2.0 OP'], False, False)),
        ('yadis2_op_delegate', ('http://unittest/yadis_idp_delegate.xrds', ['2.0 OP'], False, False)),
        ('yadis1_and_2', ('http://unittest/openid_1_and_2_xrds.xrds', ['2.0', '1.1'], None, 'http://smoker.myopenid.com/')),
        ('xri', ('=iname', ['1.0'], '=!1000', 'http://smoker.myopenid.com/')),
        ('xri_normalize', ('xri://=iname', ['1.0'], '=!1000', 'http://smoker.myopenid.com/')),
    ]

    def _test(self, url, types, claimed_id, local_id):
        service = discover.discover(url)
        if claimed_id is None:
            claimed_id = url
        if local_id is None:
            local_id = url
        self.assertEqual(service.server_url, 'http://www.myopenid.com/server')
        if types == ['2.0 OP']:
            self.assertFalse(claimed_id)
            self.assertFalse(local_id)
            self.assertFalse(service.claimed_id)
            self.assertFalse(service.local_id)
            self.assertFalse(service.identity())
            self.assertFalse(service.compat_mode())
            self.assertTrue(service.is_op_identifier())
            self.assertEqual(service.ns(), discover.OPENID2_NS)
        else:
            self.assertEqual(claimed_id, service.claimed_id)
            self.assertEqual(local_id, service.identity())

        openid_types = {
            '1.1': discover.OPENID_1_1_TYPE,
            '1.0': discover.OPENID_1_0_TYPE,
            '2.0': discover.OPENID_2_0_TYPE,
            '2.0 OP': discover.OPENID_IDP_2_0_TYPE,
            }

        types = [openid_types[t] for t in types]
        self.assertEqual(types, service.types)

class Endpoint(unittest.TestCase):
    def test_openid_2(self):
        endpoint = discover.Service()
        self.assertFalse(endpoint.compat_mode())
        endpoint = discover.Service([discover.OPENID_2_0_TYPE])
        self.assertFalse(endpoint.compat_mode())
        endpoint = discover.Service([discover.OPENID_IDP_2_0_TYPE])
        self.assertFalse(endpoint.compat_mode())
        self.assertEqual(endpoint.ns(), discover.OPENID2_NS)

    def test_openid_1(self):
        endpoint = discover.Service([discover.OPENID_1_1_TYPE])
        self.assertTrue(endpoint.compat_mode())
        self.assertEqual(endpoint.ns(), discover.OPENID1_NS)

    def test_is_op_identifier(self):
        endpoint = discover.Service([
            discover.OPENID_2_0_TYPE,
            discover.OPENID_1_0_TYPE,
            discover.OPENID_1_1_TYPE,
        ])
        self.assertFalse(endpoint.is_op_identifier())
        endpoint = discover.Service([discover.OPENID_IDP_2_0_TYPE])
        self.assertTrue(endpoint.is_op_identifier())

    def test_local_id(self):
        endpoint = discover.Service(claimed_id='claimed_id')
        self.assertEqual(endpoint.identity(), 'claimed_id')
        endpoint = discover.Service(claimed_id='claimed_id', local_id='local_id')
        self.assertEqual(endpoint.identity(), 'local_id')
        endpoint = discover.Service()
        self.assertEqual(endpoint.identity(), None)

@support.gentests
class TestDiscoverFunction(unittest.TestCase):
    data = [
        ('uri', ('http://unittest', 'uri')),
        ('bogus', ('not a URL or XRI', 'uri')),
        ('xri', ('xri://=something', 'xri')),
        ('xri_char', ('=something', 'xri')),
    ]

    @mock.patch('openid.consumer.discover.discoverURI')
    @mock.patch('openid.consumer.discover.discoverXRI')
    def _test(self, value, target, xri, uri):
        discover.discoverall(value)
        hit, miss = (uri, xri) if target == 'uri' else (xri, uri)
        hit.assert_called_with(value)
        self.assertEqual(miss.call_count, 0)

if __name__ == '__main__':
    unittest.main()
<commit_msg>Reformatted imports in test_discover<commit_after># -*- coding: utf-8 -*-
import unittest
from unittest import mock
import os.path
import urllib.error
from urllib.parse import urlencode, urljoin

from openid import message, xrds, discover
from . import support

@mock.patch('urllib.request.urlopen', support.urlopen)
@support.gentests
class Failure(unittest.TestCase):
    data = [
        ('network_error', ('http://network.error/',)),
        ('not_found', ('/404',)),
        ('header_found', ('/?' + urlencode({'header': 'X-XRDS-Location: http://unittest/404'}),)),
        ('server_error', ('/?status=500',)),
    ]

    def _test(self, path):
        url = urljoin('http://unittest', path)
        self.assertRaises(urllib.error.URLError, discover.discover, url)

@mock.patch('urllib.request.urlopen', support.urlopen)
class Discovery(unittest.TestCase):
    def test_unicode(self):
        """
        Check page with unicode and HTML entities
        """
        services = discover.discoverall('http://unittest/unicode.html')
        self.assertEqual(len(services), 0)

    def test_unicode_undecodable_html2(self):
        """
        Check page with unicode and HTML entities that can not be decoded
        but xrds document is found before it matters
        """
        with open(os.path.join(support.DATAPATH, 'unicode3.html'), encoding='utf-8') as f:
            self.assertRaises(UnicodeDecodeError, f.read)
        services = discover.discoverall('http://unittest/unicode3.html')
        self.assertEqual(len(services), 1)

    def test_noOpenID(self):
        services = discover.discoverall('http://unittest/junk.txt')
        self.assertFalse(services)

    def test_yadisEmpty(self):
        services = discover.discoverall('http://unittest/yadis_0entries.xrds')
        self.assertFalse(services)

    def test_html_yadis_empty(self):
        # The HTML document contains OpenID links but also refers to an empty Yadis
        # document which we should prefer, by the Yadis spec.
        services = discover.discoverall('http://unittest/openid_and_yadis.html')
        self.assertFalse(services)

    def test_xrds_with_header(self):
        # Even if we got a valid XRDS document but with a link to another XRDS
        # location we should prefer that location, by the Yadis spec.
        location_url = 'http://unittest/openid2_xrds.xrds'
        params = {'header': 'X-XRDS-Location: %s' % location_url}
        url = 'http://unittest/openid_1_and_2_xrds.xrds?' + urlencode(params)
        services = discover.discoverall(url)
        self.assertEqual(support.urlopen.request.get_full_url(), location_url)

    def test_fragment(self):
        url = 'http://unittest/openid.html'
        service = discover.discover(url + '#fragment')
        self.assertEqual(service.claimed_id, url)

    def test_add_protocol(self):
        url = 'unittest:8000/'
        discover.discoverall(url)
        self.assertEqual(support.urlopen.request.get_full_url(), 'http://' + url)

    def test_wrong_protocol(self):
        url = 'ssh://unittest/'
        self.assertRaises(ValueError, discover.discoverall, url)

    def test_html1And2(self):
        url = 'http://unittest/openid_1_and_2.html'
        services = discover.discoverall(url)
        self.assertEqual(len(services), 2)
        for s in services:
            self.assertEqual(s.server_url, 'http://www.myopenid.com/server')
            self.assertEqual(s.local_id, 'http://smoker.myopenid.com/')
            self.assertEqual(s.claimed_id, url)

    def test_service_sort(self):
        services = discover.discoverall('http://unittest/multiple_services.xrds')
        self.assertEqual(len(services), 3)
        self.assertTrue(discover.OPENID_IDP_2_0_TYPE in services[0].types)
        self.assertTrue(discover.OPENID_2_0_TYPE in services[1].types)
        self.assertTrue(discover.OPENID_1_0_TYPE in services[2].types)

    def test_redirected_claimed_id(self):
        claimed_id = 'http://unittest/openid2_xrds.xrds'
        url = 'http://unittest/200.txt?' + urlencode({'redirect': claimed_id})
        service = discover.discover(url)
        self.assertEqual(claimed_id, service.claimed_id)

    def test_xri_idp(self):
        service = discover.discover('=iname.idp')
        self.assertEqual(service.server_url, 'http://www.livejournal.com/openid/server.bml')

    def test_two_services(self):
        services = discover.discoverall('=twoservices')
        self.assertEqual(len(services), 2)
        self.assertTrue(services[0].local_id, 'http://smoker.myopenid.com/')
        self.assertTrue(services[1].local_id, 'http://frank.livejournal.com/')

    def test_xriNoCanonicalID(self):
        self.assertRaises(xrds.XRDSError, discover.discoverall, '=iname*empty')

@mock.patch('urllib.request.urlopen', support.urlopen)
@support.gentests
class Services(unittest.TestCase):
    data = [
        # arguments: url, types, claimed_id, local_id, canonical_id
        # last four arguments can be None, in which case the url is used instead
        ('no_delegate', ('http://unittest/openid_no_delegate.html', ['1.1'], None, None)),
        ('html1', ('http://unittest/openid.html', ['1.1'], None, 'http://smoker.myopenid.com/')),
        ('html2', ('http://unittest/openid2.html', ['2.0'], None, 'http://smoker.myopenid.com/')),
        ('yadis1_no_delegate', ('http://unittest/yadis_no_delegate.xrds', ['1.0'], None, None)),
        ('yadis2_no_local_id', ('http://unittest/openid2_xrds_no_local_id.xrds', ['2.0'], None, None)),
        ('yadis2', ('http://unittest/openid2_xrds.xrds', ['2.0'], None, 'http://smoker.myopenid.com/')),
        ('yadis2_op', ('http://unittest/yadis_idp.xrds', ['2.0 OP'], False, False)),
        ('yadis2_op_delegate', ('http://unittest/yadis_idp_delegate.xrds', ['2.0 OP'], False, False)),
        ('yadis1_and_2', ('http://unittest/openid_1_and_2_xrds.xrds', ['2.0', '1.1'], None, 'http://smoker.myopenid.com/')),
        ('xri', ('=iname', ['1.0'], '=!1000', 'http://smoker.myopenid.com/')),
        ('xri_normalize', ('xri://=iname', ['1.0'], '=!1000', 'http://smoker.myopenid.com/')),
    ]

    def _test(self, url, types, claimed_id, local_id):
        service = discover.discover(url)
        if claimed_id is None:
            claimed_id = url
        if local_id is None:
            local_id = url
        self.assertEqual(service.server_url, 'http://www.myopenid.com/server')
        if types == ['2.0 OP']:
            self.assertFalse(claimed_id)
            self.assertFalse(local_id)
            self.assertFalse(service.claimed_id)
            self.assertFalse(service.local_id)
            self.assertFalse(service.identity())
            self.assertFalse(service.compat_mode())
            self.assertTrue(service.is_op_identifier())
            self.assertEqual(service.ns(), discover.OPENID2_NS)
        else:
            self.assertEqual(claimed_id, service.claimed_id)
            self.assertEqual(local_id, service.identity())

        openid_types = {
            '1.1': discover.OPENID_1_1_TYPE,
            '1.0': discover.OPENID_1_0_TYPE,
            '2.0': discover.OPENID_2_0_TYPE,
            '2.0 OP': discover.OPENID_IDP_2_0_TYPE,
            }

        types = [openid_types[t] for t in types]
        self.assertEqual(types, service.types)

class Endpoint(unittest.TestCase):
    def test_openid_2(self):
        endpoint = discover.Service()
        self.assertFalse(endpoint.compat_mode())
        endpoint = discover.Service([discover.OPENID_2_0_TYPE])
        self.assertFalse(endpoint.compat_mode())
        endpoint = discover.Service([discover.OPENID_IDP_2_0_TYPE])
        self.assertFalse(endpoint.compat_mode())
        self.assertEqual(endpoint.ns(), discover.OPENID2_NS)

    def test_openid_1(self):
        endpoint = discover.Service([discover.OPENID_1_1_TYPE])
        self.assertTrue(endpoint.compat_mode())
        self.assertEqual(endpoint.ns(), discover.OPENID1_NS)

    def test_is_op_identifier(self):
        endpoint = discover.Service([
            discover.OPENID_2_0_TYPE,
            discover.OPENID_1_0_TYPE,
            discover.OPENID_1_1_TYPE,
        ])
        self.assertFalse(endpoint.is_op_identifier())
        endpoint = discover.Service([discover.OPENID_IDP_2_0_TYPE])
        self.assertTrue(endpoint.is_op_identifier())

    def test_local_id(self):
        endpoint = discover.Service(claimed_id='claimed_id')
        self.assertEqual(endpoint.identity(), 'claimed_id')
        endpoint = discover.Service(claimed_id='claimed_id', local_id='local_id')
        self.assertEqual(endpoint.identity(), 'local_id')
        endpoint = discover.Service()
        self.assertEqual(endpoint.identity(), None)

@support.gentests
class TestDiscoverFunction(unittest.TestCase):
    data = [
        ('uri', ('http://unittest', 'uri')),
        ('bogus', ('not a URL or XRI', 'uri')),
        ('xri', ('xri://=something', 'xri')),
        ('xri_char', ('=something', 'xri')),
    ]

    @mock.patch('openid.consumer.discover.discoverURI')
    @mock.patch('openid.consumer.discover.discoverXRI')
    def _test(self, value, target, xri, uri):
        discover.discoverall(value)
        hit, miss = (uri, xri) if target == 'uri' else (xri, uri)
        hit.assert_called_with(value)
        self.assertEqual(miss.call_count, 0)

if __name__ == '__main__':
    unittest.main()
<|endoftext|>


<commit_before>"""Functions for finding the foreground in an image with a clean
background (i.e. mostly white or black)"""

from __future__ import division
from collections import namedtuple
import numpy as np
from crash_kiss.config import BLACK, WHITE, PURPLE, FULL_DEPTH, config
import crash_kiss.util as util

def find_foreground(img, background, threshold):
    """Find the foreground of the image by subracting each RGB element
    in the image by the background. If the background has been reduced
    to a simple int or float, we'll try to avoid calling `np.abs`
    by checking to see if the background value is near 0 or 255."""
    if len(img.shape) == 2:
        return _compare_pixels(img, background, threshold)
    rgb_views = [img[:, :, idx] for idx in range(img.shape[-1])]

    # the foreground is a 2D array where 1==foreground 0==background
    fg = _compare_pixels(rgb_views[0], background, threshold)
    for view in rgb_views[1:]:
        bg = fg == 0
        new_data = _compare_pixels(view[bg], background, threshold)
        fg[bg] = new_data
    return fg

def _compare_pixels(img, background, threshold):
    """Compare a 2-D or 3-D image array
    to a background value given a certain threshold"""
    is_num = isinstance(background, int)
    if background - BLACK <= 5:
        diff = img - background > threshold
    elif WHITE - background <= 5:
        diff = background - img > threshold
    else:
        diff = np.abs(img - background) > threshold
    if len(diff.shape) == 3:
        diff = np.any(diff, axis=2)  # we're using a 3D array
    return diff

def simplify_background(background, config):
    """See if the background's RGB elements are similar.
    If each element in the background is similar enough, we can do a simple
    array - int operation instead of the more expensive array - [R, G, B]
    or the even pricier array - <array of shape (NROWS, 1, 3)> operation."""
    while isinstance(background, np.ndarray):
        bg_change_tolerance = config["bg_change_tolerance"]
        bmax, bmin = background.max(axis=0), background.min(axis=0)
        diff = (bmax - bmin) < bg_change_tolerance
        if len(background.shape) >= 2:
            diff = np.all(diff)
        if diff:
            background = np.median(background, axis=0)
        else:
            break  # we've simplified the background as much as we can
    return background

_MID_FG = 0xFFFF

@profile
def center_smash(img, fg, bounds):
    """Move the rows of each subject together until they touch.
    Write over the vacated space with whatever the row's negative space
    is (probably white or transparent pixels)."""
    start, stop, fg_mid = bounds
    rlen = img.shape[1] - stop
    max_depth = fg.shape[1] // 4
    center = start + 2*max_depth
    side_len = fg.shape[1] // 2
    fg_l = fg[:, :bounds.fg_mid]
    fg_l = util.invert_horizontal(fg_l)
    fg_r = fg[:, bounds.fg_mid:]

    lstart = np.argmax(fg_l, axis=1)
    lstart[fg_l[:, 0] == 1] = _MID_FG
    rstart = np.argmax(fg_r, axis=1)
    rstart[fg_r[:, 0] == 1] = _MID_FG

    for irow, ls, rs, frow in zip(img, lstart, rstart, fg):
        lmov = rmov = max_depth
        if not ls and not rs:
            # no contact can be made
            irow[lmov: start+lmov] = irow[:start]
            irow[stop-rmov: -rmov] = irow[stop:]
        elif not np.any(frow[max_depth: -max_depth]):
            # no contact can be made but
            # we'll select & shift the foreground differently
            irow[lmov: start+lmov*2] = irow[:start+max_depth]
            irow[center: -rmov] = irow[stop-rmov:]
        elif not ls:
            # no contact can be made, but the order we do things matters
            # because the background of the left side could cover up the
            # foreground of the right side
            subj = irow[stop-(side_len-rs): stop].copy()
            irow[lmov: start+lmov] = irow[:start]
            irow[stop-rmov: -rmov] = irow[stop:]
            irow[stop-rmov-len(subj): stop-rmov] = subj
        elif not rs:
            # no contact can be made, but the order we do things matters
            # because the background of the right side could cover up the
            # foreground of the left side
            subj = irow[start: start + (side_len-ls)].copy()
            irow[stop-rmov: -rmov] = irow[stop:]
            irow[lmov: start+lmov] = irow[:start]
            irow[start+lmov: start+lmov+len(subj)] = subj
        elif rs and ls and np.any(frow[max_depth: -max_depth]):
            # contact will be made (this is the "crash" or "smash")
            subj = irow[start: stop][frow]
            subjl = len(subj)
            offs = rs - ls
            fstart = start + fg_mid + offs - (subjl // 2)
            lmov = fstart - start
            rmov = stop - (fstart + subjl)
            irow[fstart: fstart + subjl] = subj
            irow[lmov: start+lmov] = irow[:start]
            r1 = irow[stop-rmov:-rmov]
            s2 = rlen - len(r1)
            r2 = irow[stop+s2:]
            r1[:] = r2
        else:
            # contact won't be made, but white space may cover
            # either side of the subject if we're not careful
            lsubj = irow[start: start+rs].copy()
            rsubj = irow[stop-side_len+rs:stop].copy()
            irow[stop-rmov: -rmov] = irow[stop:]
            irow[lmov: start+lmov] = irow[:start]
            irow[start+lmov:start+lmov+len(lsubj)] = lsubj
        irow[:lmov] = WHITE
        irow[-rmov:] = WHITE

def get_foreground_area(img, max_depth):
    bounds = _get_fg_bounds(img.shape, max_depth)
    return img[:, bounds.start:bounds.stop], bounds

_fg_bounds = namedtuple("fg_bounds", "start stop fg_mid")

def _get_fg_bounds(img_shape, max_depth):
    """Returns start, stop idx of the 'smashable foreground'
    area in the middle of an image.

    Indexing the image with `img[:, start:stop]` will successfully
    select the foreground columns."""
    width = img_shape[1]
    half = width // 2
    if max_depth >= half or max_depth == FULL_DEPTH:
        max_depth = width // 4
    start = (width // 2) - (max_depth * 2)
    stop = start + max_depth * 4
    assert stop - start == max_depth * 4
    fg_mid = (stop - start) // 2
    return _fg_bounds(start, stop, fg_mid)

def reveal_foreground(img, foreground, bounds):
    max_depth = foreground.shape[1] // 4
    start, stop, fg_mid = bounds
    critical_fg = foreground[:, max_depth: -max_depth]
    img[:, start: stop][foreground] = PURPLE
    img[:, start+max_depth: stop-max_depth][critical_fg] = BLACK
    img[:, start-1:start+1] = PURPLE
    img[:, stop-1:stop+1] = PURPLE
    img[:, start+max_depth*2] = BLACK
    img[:, start+max_depth] = BLACK
    img[:, stop-max_depth] = BLACK

def reveal_background(img, foreground, bounds):
    img[:, bounds.start: bounds.stop][foreground] = WHITE

<commit_msg>Added case in center_smash<commit_after>"""Functions for finding the foreground in an image with a clean
background (i.e. mostly white or black)"""

from __future__ import division
from collections import namedtuple
import numpy as np
from crash_kiss.config import BLACK, WHITE, PURPLE, FULL_DEPTH, config
import crash_kiss.util as util

def find_foreground(img, background, threshold):
    """Find the foreground of the image by subracting each RGB element
    in the image by the background. If the background has been reduced
    to a simple int or float, we'll try to avoid calling `np.abs`
    by checking to see if the background value is near 0 or 255."""
    if len(img.shape) == 2:
        return _compare_pixels(img, background, threshold)
    rgb_views = [img[:, :, idx] for idx in range(img.shape[-1])]

    # the foreground is a 2D array where 1==foreground 0==background
    fg = _compare_pixels(rgb_views[0], background, threshold)
    for view in rgb_views[1:]:
        bg = fg == 0
        new_data = _compare_pixels(view[bg], background, threshold)
        fg[bg] = new_data
    return fg

def _compare_pixels(img, background, threshold):
    """Compare a 2-D or 3-D image array
    to a background value given a certain threshold"""
    is_num = isinstance(background, int)
    if background - BLACK <= 5:
        diff = img - background > threshold
    elif WHITE - background <= 5:
        diff = background - img > threshold
    else:
        diff = np.abs(img - background) > threshold
    if len(diff.shape) == 3:
        diff = np.any(diff, axis=2)  # we're using a 3D array
    return diff

def simplify_background(background, config):
    """See if the background's RGB elements are similar.
    If each element in the background is similar enough, we can do a simple
    array - int operation instead of the more expensive array - [R, G, B]
    or the even pricier array - <array of shape (NROWS, 1, 3)> operation."""
    while isinstance(background, np.ndarray):
        bg_change_tolerance = config["bg_change_tolerance"]
        bmax, bmin = background.max(axis=0), background.min(axis=0)
        diff = (bmax - bmin) < bg_change_tolerance
        if len(background.shape) >= 2:
            diff = np.all(diff)
        if diff:
            background = np.median(background, axis=0)
        else:
            break  # we've simplified the background as much as we can
    return background

_MID_FG = 0xFFFF

@profile
def center_smash(img, fg, bounds):
    """Move the rows of each subject together until they touch.
    Write over the vacated space with whatever the row's negative space
    is (probably white or transparent pixels)."""
    start, stop, fg_mid = bounds
    rlen = img.shape[1] - stop
    max_depth = fg.shape[1] // 4
    center = start + 2*max_depth
    side_len = fg.shape[1] // 2
    fg_l = fg[:, :bounds.fg_mid]
    fg_l = util.invert_horizontal(fg_l)
    fg_r = fg[:, bounds.fg_mid:]

    lstart = np.argmax(fg_l, axis=1)
    lstart[fg_l[:, 0] == 1] = _MID_FG
    rstart = np.argmax(fg_r, axis=1)
    rstart[fg_r[:, 0] == 1] = _MID_FG

    for irow, ls, rs, frow in zip(img, lstart, rstart, fg):
        lmov = rmov = max_depth
        if rs == _MID_FG and ls == _MID_FG:
            # NOTE: smash behavior not fully defined yet
            # (if the subject overlaps the middle)
            lsubj = irow[start+max_depth: center][frow[max_depth:fg_mid]].copy()
            rsubj = irow[center: center+max_depth][frow[fg_mid:fg_mid+max_depth]].copy()
            llen = len(lsubj)
            rlen = len(rsubj)
            lmov = max_depth - llen
            rmov = max_depth - rlen
            irow[lmov: center] = irow[:center - lmov]
            irow[center - llen: center] = lsubj
            irow[center: -rmov] = irow[center + rmov:]
            irow[center: center + rlen] = rsubj
        elif not ls and not rs:
            irow[center + rmov: -rmov] = irow[stop:]
            irow[lmov: start+lmov] = irow[:start]
        elif not ls or not rs and (ls + rs > max_depth):
            # no contact can be made
            irow[center: -rmov] = irow[center + rmov:]
            irow[lmov: center] = irow[:start + lmov]
        elif ls and not rs:
            # no contact is possible, but there's a left-over-center overshoot
            irow[center + rmov: -rmov] = irow[stop:]
            overshoot = ls - max_depth
            assert overshoot > 0
            print rs, ls, overshoot
            irow[lmov: center+overshoot] = irow[:center-ls]
        elif rs:
            # no contact is possible
            # a right-over-center overshoot is possible
            irow[:] = [200, 200, 200]
        elif (ls >= max_depth) and (rs >= max_depth):
            # no contact can be made but
            # we'll select & shift the foreground differently
            irow[lmov: start+lmov*2] = irow[:start+max_depth]
            irow[center: -rmov] = irow[stop-rmov:]
        elif rs + ls <= side_len:
            # contact will be made (this is the "crash" or "smash")
            irow[:] = [100, 0, 0]  # dark red
            continue
            subj = irow[start: stop][frow]
            subjl = len(subj)
            offs = rs - ls
            fstart = start + fg_mid + offs - (subjl // 2)
            lmov = fstart - start
            rmov = stop - (fstart + subjl)
            irow[fstart: fstart + subjl] = subj
            irow[lmov: start+lmov] = irow[:start]
            r1 = irow[stop-rmov:-rmov]
            s2 = rlen - len(r1)
            r2 = irow[stop+s2:]
            r1[:] = r2
        else:
            # contact won't be made, but white space may cover
            # either side of the subject if we're not careful
            irow[:] = [128, 255, 0]  # bright green
            continue
            lsubj = irow[start: start+rs].copy()
            rsubj = irow[stop-side_len+rs:stop].copy()
            irow[center: -rs] = irow[center+rs:]
            irow[lmov: start+lmov] = irow[:start]
            irow[start+lmov:start+lmov+len(lsubj)] = lsubj
            irow[center+rs-rmov:center+rs-rmov+len(rsubj)] = rsubj
        irow[:lmov] = WHITE
        irow[-rmov:] = WHITE

def get_foreground_area(img, max_depth):
    bounds = _get_fg_bounds(img.shape, max_depth)
    return img[:, bounds.start:bounds.stop], bounds

_fg_bounds = namedtuple("fg_bounds", "start stop fg_mid")

def _get_fg_bounds(img_shape, max_depth):
    """Returns start, stop idx of the 'smashable foreground'
    area in the middle of an image.

    Indexing the image with `img[:, start:stop]` will successfully
    select the foreground columns."""
    width = img_shape[1]
    half = width // 2
    if max_depth >= half or max_depth == FULL_DEPTH:
        max_depth = width // 4
    start = (width // 2) - (max_depth * 2)
    stop = start + max_depth * 4
    assert stop - start == max_depth * 4
    fg_mid = (stop - start) // 2
    return _fg_bounds(start, stop, fg_mid)

def reveal_foreground(img, foreground, bounds):
    max_depth = foreground.shape[1] // 4
    start, stop, fg_mid = bounds
    critical_fg = foreground[:, max_depth: -max_depth]
    img[:, start: stop][foreground] = PURPLE
    img[:, start+max_depth: stop-max_depth][critical_fg] = BLACK
    img[:, start-1:start+1] = PURPLE
    img[:, stop-1:stop+1] = PURPLE
    img[:, start+max_depth*2] = BLACK
    img[:, start+max_depth] = BLACK
    img[:, stop-max_depth] = BLACK

def reveal_background(img, foreground, bounds):
    img[:, bounds.start: bounds.stop][foreground] = WHITE

<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
# vim: tabstop=4 shiftwidth=4 softtabstop=4
#
# Copyright (C) 2014-2018 GEM Foundation
#
# OpenQuake is free software: you can redistribute it and/or modify it
# under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# OpenQuake is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with OpenQuake. If not, see <http://www.gnu.org/licenses/>.
import os
import sys
import getpass
import logging
from openquake.baselib import sap, config, datastore
from openquake.baselib.general import safeprint
from openquake.hazardlib import valid
from openquake.commonlib import logs, readinput
from openquake.engine import engine as eng
from openquake.engine.export import core
from openquake.engine.utils import confirm
from openquake.engine.tools.make_html_report import make_report
from openquake.server import dbserver
from openquake.calculators.export import export
from openquake.commands.abort import abort

HAZARD_CALCULATION_ARG = "--hazard-calculation-id"
MISSING_HAZARD_MSG = "Please specify '%s=<id>'" % HAZARD_CALCULATION_ARG

def read(calc_id, username=None):
    """
    :param calc_id: a calculation ID
    :param username: if given, restrict the search to the user's calculations
    :returns: the associated DataStore instance
    """
    if isinstance(calc_id, str) or calc_id < 0 and not username:
        # get the last calculation in the datastore of the current user
        return datastore.read(calc_id)
    job = logs.dbcmd('get_job', calc_id, username)
    if job:
        return datastore.read(job.ds_calc_dir + '.hdf5')
    else:
        # calc_id can be present in the datastore and not in the database:
        # this happens if the calculation was run with `oq run`
        return datastore.read(calc_id)

def get_job_id(job_id, username=None):
    job = logs.dbcmd('get_job', job_id, username)
    if not job:
        sys.exit('Job %s not found' % job_id)
    return job.id

def run_job(job_ini, log_level='info', log_file=None, exports='',
            username=getpass.getuser(), **kw):
    """
    Run a job using the specified config file and other options.

    :param str job_ini:
        Path to calculation config (INI-style) files.
    :param str log_level:
        'debug', 'info', 'warn', 'error', or 'critical'
    :param str log_file:
        Path to log file.
    :param exports:
        A comma-separated string of export types requested by the user.
    :param username:
        Name of the user running the job
    :param kw:
        Extra parameters like hazard_calculation_id and calculation_mode
    """
    job_id = logs.init('job', getattr(logging, log_level.upper()))
    with logs.handle(job_id, log_level, log_file):
        job_ini = os.path.abspath(job_ini)
        oqparam = eng.job_from_file(job_ini, job_id, username, **kw)
        kw['username'] = username
        eng.run_calc(job_id, oqparam, exports, **kw)
        for line in logs.dbcmd('list_outputs', job_id, False):
            safeprint(line)
    return job_id

def run_tile(job_ini, sites_slice):
    """
    Used in tiling calculations
    """
    return run_job(job_ini, sites_slice=(sites_slice.start, sites_slice.stop))

def del_calculation(job_id, confirmed=False):
    """
    Delete a calculation and all associated outputs.
    """
    if logs.dbcmd('get_job', job_id) is None:
        print('There is no job %d' % job_id)
        return

    if confirmed or confirm(
            'Are you sure you want to (abort and) delete this calculation and '
            'all associated outputs?\nThis action cannot be undone. (y/n): '):
        try:
            abort.func(job_id)
            resp = logs.dbcmd('del_calc', job_id, getpass.getuser())
        except RuntimeError as err:
            safeprint(err)
        else:
            if 'success' in resp:
                print('Removed %d' % job_id)
            else:
                print(resp['error'])

def smart_run(job_ini, oqparam, log_level, log_file, exports, reuse_hazard):
    """
    Run calculations by storing their hazard checksum and reusing previous
    calculations if requested.
    """
    haz_checksum = readinput.get_checksum32(oqparam, hazard=True)
    # retrieve an old calculation with the right checksum, if any
    job = logs.dbcmd('get_job_from_checksum', haz_checksum)
    reuse = reuse_hazard and job and os.path.exists(job.ds_calc_dir + '.hdf5')
    # recompute the hazard and store the checksum
    ebr = (oqparam.calculation_mode in ('event_based_risk', 'ebrisk') and
           'gmfs' not in oqparam.inputs)
    if ebr:
        kw = dict(calculation_mode='event_based')
        if (oqparam.sites or 'sites' in oqparam.inputs or
                'site_model' in oqparam.inputs):
            # remove exposure from the hazard
            kw['exposure_file'] = ''
    else:
        kw = {}
    if not reuse:
        hc_id = run_job(job_ini, log_level, log_file, exports, **kw)
        if job is None:
            logs.dbcmd('add_checksum', hc_id, haz_checksum)
        elif not reuse_hazard or not os.path.exists(job.ds_calc_dir + '.hdf5'):
            logs.dbcmd('update_job_checksum', hc_id, haz_checksum)
        if ebr:
            job_id = run_job(job_ini, log_level, log_file,
                             exports, hazard_calculation_id=hc_id)
    else:
        hc_id = job.id
        logging.info('Reusing job #%d', job.id)
        job_id = run_job(job_ini, log_level, log_file,
                         exports, hazard_calculation_id=hc_id)
    if oqparam.aggregate_by:
        logging.info('Exporting aggregated data')
        dstore = datastore.read(job_id)
        aggby = 'aggregate_by/%s/' % ','.join(oqparam.aggregate_by)
        fnames = []
        fnames.extend(export((aggby + 'avg_losses', 'csv'), dstore))
        fnames.extend(export((aggby + 'curves', 'csv'), dstore))
        for fname in fnames:
            logging.info('Exported %s', fname)

@sap.Script
def engine(log_file, no_distribute, yes, config_file, make_html_report,
           upgrade_db, db_version, what_if_I_upgrade, run,
           list_hazard_calculations, list_risk_calculations,
           delete_calculation, delete_uncompleted_calculations,
           hazard_calculation_id, list_outputs, show_log,
           export_output, export_outputs, exports='',
           log_level='info', reuse_hazard=False):
    """
    Run a calculation using the traditional command line API
    """
    if not run:
        # configure a basic logging
        logs.init()

    if config_file:
        config.read(os.path.abspath(os.path.expanduser(config_file)),
                    soft_mem_limit=int, hard_mem_limit=int, port=int,
                    multi_user=valid.boolean, multi_node=valid.boolean)

    if no_distribute:
        os.environ['OQ_DISTRIBUTE'] = 'no'

    # check if the datadir exists
    datadir = datastore.get_datadir()
    if not os.path.exists(datadir):
        os.makedirs(datadir)

    dbserver.ensure_on()
    # check if we are talking to the right server
    err = dbserver.check_foreign()
    if err:
        sys.exit(err)

    if upgrade_db:
        msg = logs.dbcmd('what_if_I_upgrade', 'read_scripts')
        if msg.startswith('Your database is already updated'):
            pass
        elif yes or confirm('Proceed? (y/n) '):
            logs.dbcmd('upgrade_db')
        sys.exit(0)

    if db_version:
        safeprint(logs.dbcmd('db_version'))
        sys.exit(0)

    if what_if_I_upgrade:
        safeprint(logs.dbcmd('what_if_I_upgrade', 'extract_upgrade_scripts'))
        sys.exit(0)

    # check if the db is outdated
    outdated = logs.dbcmd('check_outdated')
    if outdated:
        sys.exit(outdated)

    # hazard or hazard+risk
    if hazard_calculation_id == -1:
        # get the latest calculation of the current user
        hc_id = get_job_id(hazard_calculation_id, getpass.getuser())
    elif hazard_calculation_id:
        # make it possible to use calculations made by another user
        hc_id = get_job_id(hazard_calculation_id)
    else:
        hc_id = None
    if run:
        log_file = os.path.expanduser(log_file) \
            if log_file is not None else None
        job_inis = [os.path.expanduser(f) for f in run]
        if len(job_inis) == 1 and not hc_id:
            # init logs before calling get_oqparam
            logs.init('nojob', getattr(logging, log_level.upper()))
            # not using logs.handle that logs on the db
            oq = readinput.get_oqparam(job_inis[0])
            smart_run(job_inis[0], oq, log_level, log_file,
                      exports, reuse_hazard)
            return
        for i, job_ini in enumerate(job_inis):
            open(job_ini, 'rb').read()  # IOError if the file does not exist
            job_id = run_job(job_ini, log_level, log_file,
                             exports, hazard_calculation_id=hc_id)
            if not hc_id:  # use the first calculation as base for the others
                hc_id = job_id
    # hazard
    elif list_hazard_calculations:
        for line in logs.dbcmd(
                'list_calculations', 'hazard', getpass.getuser()):
            safeprint(line)
    elif delete_calculation is not None:
        del_calculation(delete_calculation, yes)
    # risk
    elif list_risk_calculations:
        for line in logs.dbcmd('list_calculations', 'risk', getpass.getuser()):
            safeprint(line)

    # export
    elif make_html_report:
        safeprint('Written %s' % make_report(make_html_report))
        sys.exit(0)

    elif list_outputs is not None:
        hc_id = get_job_id(list_outputs)
        for line in logs.dbcmd('list_outputs', hc_id):
            safeprint(line)
    elif show_log is not None:
        hc_id = get_job_id(show_log)
        for line in logs.dbcmd('get_log', hc_id):
            safeprint(line)

    elif export_output is not None:
        output_id, target_dir = export_output
        dskey, calc_id, datadir = logs.dbcmd('get_output', int(output_id))
        for line in core.export_output(
                dskey, calc_id, datadir, os.path.expanduser(target_dir),
                exports or 'csv,xml'):
            safeprint(line)

    elif export_outputs is not None:
        job_id, target_dir = export_outputs
        hc_id = get_job_id(job_id)
        for line in core.export_outputs(
                hc_id, os.path.expanduser(target_dir), exports or 'csv,xml'):
            safeprint(line)

    elif delete_uncompleted_calculations:
        logs.dbcmd('delete_uncompleted_calculations', getpass.getuser())

    else:
        engine.parentparser.prog = 'oq engine'
        engine.parentparser.print_usage()

engine._add('log_file', '--log-file', '-L', help='''\
Location where to store log messages; if not specified, log messages
will be printed to the console (to stderr)''')
engine._add('no_distribute', '--no-distribute', '--nd', help='''\
Disable calculation task distribution and run the
computation in a single process. This is intended for
use in debugging and profiling.''', action='store_true')
engine.flg('yes', 'Automatically answer "yes" when asked to confirm an action')
engine.opt('config_file', 'Custom openquake.cfg file, to override default '
           'configurations')
engine._add('make_html_report', '--make-html-report', '--r',
            help='Build an HTML report of the computation at the given date',
            metavar='YYYY-MM-DD|today')
engine.flg('upgrade_db', 'Upgrade the openquake database')
engine.flg('db_version', 'Show the current version of the openquake database')
engine.flg('what_if_I_upgrade', 'Show what will happen to the openquake '
           'database if you upgrade')
engine._add('run', '--run', help='Run a job with the specified config file',
            metavar='JOB_INI', nargs='+')
engine._add('list_hazard_calculations', '--list-hazard-calculations', '--lhc',
            help='List hazard calculation information', action='store_true')
engine._add('list_risk_calculations', '--list-risk-calculations', '--lrc',
            help='List risk calculation information', action='store_true')
engine._add('delete_calculation', '--delete-calculation', '--dc',
            help='Delete a calculation and all associated outputs',
            metavar='CALCULATION_ID', type=int)
engine._add('delete_uncompleted_calculations',
            '--delete-uncompleted-calculations', '--duc',
            help='Delete all the uncompleted calculations',
            action='store_true')
engine._add('hazard_calculation_id', '--hazard-calculation-id', '--hc',
            help='Use the given job as input for the next job')
engine._add('list_outputs', '--list-outputs', '--lo',
            help='List outputs for the specified calculation',
            metavar='CALCULATION_ID')
engine._add('show_log', '--show-log', '--sl',
            help='Show the log of the specified calculation',
            metavar='CALCULATION_ID')
engine._add('export_output', '--export-output', '--eo',
            nargs=2, metavar=('OUTPUT_ID', 'TARGET_DIR'),
            help='Export the desired output to the specified directory')
engine._add('export_outputs', '--export-outputs', '--eos',
            nargs=2, metavar=('CALCULATION_ID', 'TARGET_DIR'),
            help='Export all of the calculation outputs to the '
            'specified directory')
engine.opt('exports', 'Comma-separated string specifing the export formats, '
           'in order of priority')
engine.opt('log_level', 'Defaults to "info"',
           choices=['debug', 'info', 'warn', 'error', 'critical'])
engine.flg('reuse_hazard', 'Reuse the event based hazard if available')
<commit_msg>[skip hazardlib]<commit_after># -*- coding: utf-8 -*-
# vim: tabstop=4 shiftwidth=4 softtabstop=4
#
# Copyright (C) 2014-2018 GEM Foundation
#
# OpenQuake is free software: you can redistribute it and/or modify it
# under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# OpenQuake is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with OpenQuake. If not, see <http://www.gnu.org/licenses/>.
import os
import sys
import getpass
import logging
from openquake.baselib import sap, config, datastore
from openquake.baselib.general import safeprint
from openquake.hazardlib import valid
from openquake.commonlib import logs, readinput
from openquake.engine import engine as eng
from openquake.engine.export import core
from openquake.engine.utils import confirm
from openquake.engine.tools.make_html_report import make_report
from openquake.server import dbserver
from openquake.calculators.export import export
from openquake.commands.abort import abort

HAZARD_CALCULATION_ARG = "--hazard-calculation-id"
MISSING_HAZARD_MSG = "Please specify '%s=<id>'" % HAZARD_CALCULATION_ARG

def read(calc_id, username=None):
    """
    :param calc_id: a calculation ID
    :param username: if given, restrict the search to the user's calculations
    :returns: the associated DataStore instance
    """
    if isinstance(calc_id, str) or calc_id < 0 and not username:
        # get the last calculation in the datastore of the current user
        return datastore.read(calc_id)
    job = logs.dbcmd('get_job', calc_id, username)
    if job:
        return datastore.read(job.ds_calc_dir + '.hdf5')
    else:
        # calc_id can be present in the datastore and not in the database:
        # this happens if the calculation was run with `oq run`
        return datastore.read(calc_id)

def get_job_id(job_id, username=None):
    job = logs.dbcmd('get_job', job_id, username)
    if not job:
        sys.exit('Job %s not found' % job_id)
    return job.id

def run_job(job_ini, log_level='info', log_file=None, exports='',
            username=getpass.getuser(), **kw):
    """
    Run a job using the specified config file and other options.

    :param str job_ini:
        Path to calculation config (INI-style) files.
    :param str log_level:
        'debug', 'info', 'warn', 'error', or 'critical'
    :param str log_file:
        Path to log file.
    :param exports:
        A comma-separated string of export types requested by the user.
    :param username:
        Name of the user running the job
    :param kw:
        Extra parameters like hazard_calculation_id and calculation_mode
    """
    job_id = logs.init('job', getattr(logging, log_level.upper()))
    with logs.handle(job_id, log_level, log_file):
        job_ini = os.path.abspath(job_ini)
        oqparam = eng.job_from_file(job_ini, job_id, username, **kw)
        kw['username'] = username
        eng.run_calc(job_id, oqparam, exports, **kw)
        for line in logs.dbcmd('list_outputs', job_id, False):
            safeprint(line)
    return job_id

def run_tile(job_ini, sites_slice):
    """
    Used in tiling calculations
    """
    return run_job(job_ini, sites_slice=(sites_slice.start, sites_slice.stop))

def del_calculation(job_id, confirmed=False):
    """
    Delete a calculation and all associated outputs.
    """
    if logs.dbcmd('get_job', job_id) is None:
        print('There is no job %d' % job_id)
        return

    if confirmed or confirm(
            'Are you sure you want to (abort and) delete this calculation and '
            'all associated outputs?\nThis action cannot be undone. (y/n): '):
        try:
            abort.func(job_id)
            resp = logs.dbcmd('del_calc', job_id, getpass.getuser())
        except RuntimeError as err:
            safeprint(err)
        else:
            if 'success' in resp:
                print('Removed %d' % job_id)
            else:
                print(resp['error'])

def smart_run(job_ini, oqparam, log_level, log_file, exports, reuse_hazard):
    """
    Run calculations by storing their hazard checksum and reusing previous
    calculations if requested.
    """
    haz_checksum = readinput.get_checksum32(oqparam, hazard=True)
    # retrieve an old calculation with the right checksum, if any
    job = logs.dbcmd('get_job_from_checksum', haz_checksum)
    reuse = reuse_hazard and job and os.path.exists(job.ds_calc_dir + '.hdf5')
    # recompute the hazard and store the checksum
    ebr = (oqparam.calculation_mode in ('event_based_risk', 'ebrisk') and
           'gmfs' not in oqparam.inputs)
    if ebr:
        kw = dict(calculation_mode='event_based')
        if (oqparam.sites or 'sites' in oqparam.inputs or
                'site_model' in oqparam.inputs):
            # remove exposure from the hazard
            kw['exposure_file'] = ''
    else:
        kw = {}
    if not reuse:
        hc_id = run_job(job_ini, log_level, log_file, exports, **kw)
        if job is None:
            logs.dbcmd('add_checksum', hc_id, haz_checksum)
        elif not reuse_hazard or not os.path.exists(job.ds_calc_dir + '.hdf5'):
            logs.dbcmd('update_job_checksum', hc_id, haz_checksum)
        if ebr:
            job_id = run_job(job_ini, log_level, log_file,
                             exports, hazard_calculation_id=hc_id)
        else:
            job_id = hc_id
    else:
        hc_id = job.id
        logging.info('Reusing job #%d', job.id)
        job_id = run_job(job_ini, log_level, log_file,
                         exports, hazard_calculation_id=hc_id)
    if oqparam.aggregate_by:
        logging.info('Exporting aggregated data')
        dstore = datastore.read(job_id)
        aggby = 'aggregate_by/%s/' % ','.join(oqparam.aggregate_by)
        fnames = []
        fnames.extend(export((aggby + 'avg_losses', 'csv'), dstore))
        fnames.extend(export((aggby + 'curves', 'csv'), dstore))
        for fname in fnames:
            logging.info('Exported %s', fname)

@sap.Script
def engine(log_file, no_distribute, yes, config_file, make_html_report,
           upgrade_db, db_version, what_if_I_upgrade, run,
           list_hazard_calculations, list_risk_calculations,
           delete_calculation, delete_uncompleted_calculations,
           hazard_calculation_id, list_outputs, show_log,
           export_output, export_outputs, exports='',
           log_level='info', reuse_hazard=False):
    """
    Run a calculation using the traditional command line API
    """
    if not run:
        # configure a basic logging
        logs.init()

    if config_file:
        config.read(os.path.abspath(os.path.expanduser(config_file)),
                    soft_mem_limit=int, hard_mem_limit=int, port=int,
                    multi_user=valid.boolean, multi_node=valid.boolean)

    if no_distribute:
        os.environ['OQ_DISTRIBUTE'] = 'no'

    # check if the datadir exists
    datadir = datastore.get_datadir()
    if not os.path.exists(datadir):
        os.makedirs(datadir)

    dbserver.ensure_on()
    # check if we are talking to the right server
    err = dbserver.check_foreign()
    if err:
        sys.exit(err)

    if upgrade_db:
        msg = logs.dbcmd('what_if_I_upgrade', 'read_scripts')
        if msg.startswith('Your database is already updated'):
            pass
        elif yes or confirm('Proceed? (y/n) '):
            logs.dbcmd('upgrade_db')
        sys.exit(0)

    if db_version:
        safeprint(logs.dbcmd('db_version'))
        sys.exit(0)

    if what_if_I_upgrade:
        safeprint(logs.dbcmd('what_if_I_upgrade', 'extract_upgrade_scripts'))
        sys.exit(0)

    # check if the db is outdated
    outdated = logs.dbcmd('check_outdated')
    if outdated:
        sys.exit(outdated)

    # hazard or hazard+risk
    if hazard_calculation_id == -1:
        # get the latest calculation of the current user
        hc_id = get_job_id(hazard_calculation_id, getpass.getuser())
    elif hazard_calculation_id:
        # make it possible to use calculations made by another user
        hc_id = get_job_id(hazard_calculation_id)
    else:
        hc_id = None
    if run:
        log_file = os.path.expanduser(log_file) \
            if log_file is not None else None
        job_inis = [os.path.expanduser(f) for f in run]
        if len(job_inis) == 1 and not hc_id:
            # init logs before calling get_oqparam
            logs.init('nojob', getattr(logging, log_level.upper()))
            # not using logs.handle that logs on the db
            oq = readinput.get_oqparam(job_inis[0])
            smart_run(job_inis[0], oq, log_level, log_file,
                      exports, reuse_hazard)
            return
        for i, job_ini in enumerate(job_inis):
            open(job_ini, 'rb').read()  # IOError if the file does not exist
            job_id = run_job(job_ini, log_level, log_file,
                             exports, hazard_calculation_id=hc_id)
            if not hc_id:  # use the first calculation as base for the others
                hc_id = job_id
    # hazard
    elif list_hazard_calculations:
        for line in logs.dbcmd(
                'list_calculations', 'hazard', getpass.getuser()):
            safeprint(line)
    elif delete_calculation is not None:
        del_calculation(delete_calculation, yes)
    # risk
    elif list_risk_calculations:
        for line in logs.dbcmd('list_calculations', 'risk', getpass.getuser()):
            safeprint(line)

    # export
    elif make_html_report:
        safeprint('Written %s' % make_report(make_html_report))
        sys.exit(0)

    elif list_outputs is not None:
        hc_id = get_job_id(list_outputs)
        for line in logs.dbcmd('list_outputs', hc_id):
            safeprint(line)
    elif show_log is not None:
        hc_id = get_job_id(show_log)
        for line in logs.dbcmd('get_log', hc_id):
            safeprint(line)

    elif export_output is not None:
        output_id, target_dir = export_output
        dskey, calc_id, datadir = logs.dbcmd('get_output', int(output_id))
        for line in core.export_output(
                dskey, calc_id, datadir, os.path.expanduser(target_dir),
                exports or 'csv,xml'):
            safeprint(line)

    elif export_outputs is not None:
        job_id, target_dir = export_outputs
        hc_id = get_job_id(job_id)
        for line in core.export_outputs(
                hc_id, os.path.expanduser(target_dir), exports or 'csv,xml'):
            safeprint(line)

    elif delete_uncompleted_calculations:
        logs.dbcmd('delete_uncompleted_calculations', getpass.getuser())

    else:
        engine.parentparser.prog = 'oq engine'
        engine.parentparser.print_usage()

engine._add('log_file', '--log-file', '-L', help='''\
Location where to store log messages; if not specified, log messages
will be printed to the console (to stderr)''')
engine._add('no_distribute', '--no-distribute', '--nd', help='''\
Disable calculation task distribution and run the
computation in a single process. This is intended for
use in debugging and profiling.''', action='store_true')
engine.flg('yes', 'Automatically answer "yes" when asked to confirm an action')
engine.opt('config_file', 'Custom openquake.cfg file, to override default '
           'configurations')
engine._add('make_html_report', '--make-html-report', '--r',
            help='Build an HTML report of the computation at the given date',
            metavar='YYYY-MM-DD|today')
engine.flg('upgrade_db', 'Upgrade the openquake database')
engine.flg('db_version', 'Show the current version of the openquake database')
engine.flg('what_if_I_upgrade', 'Show what will happen to the openquake '
           'database if you upgrade')
engine._add('run', '--run', help='Run a job with the specified config file',
            metavar='JOB_INI', nargs='+')
engine._add('list_hazard_calculations', '--list-hazard-calculations', '--lhc',
            help='List hazard calculation information', action='store_true')
engine._add('list_risk_calculations', '--list-risk-calculations', '--lrc',
            help='List risk calculation information', action='store_true')
engine._add('delete_calculation', '--delete-calculation', '--dc',
            help='Delete a calculation and all associated outputs',
            metavar='CALCULATION_ID', type=int)
engine._add('delete_uncompleted_calculations',
            '--delete-uncompleted-calculations', '--duc',
            help='Delete all the uncompleted calculations',
            action='store_true')
engine._add('hazard_calculation_id', '--hazard-calculation-id', '--hc',
            help='Use the given job as input for the next job')
engine._add('list_outputs', '--list-outputs', '--lo',
            help='List outputs for the specified calculation',
            metavar='CALCULATION_ID')
engine._add('show_log', '--show-log', '--sl',
            help='Show the log of the specified calculation',
            metavar='CALCULATION_ID')
engine._add('export_output', '--export-output', '--eo',
            nargs=2, metavar=('OUTPUT_ID', 'TARGET_DIR'),
            help='Export the desired output to the specified directory')
engine._add('export_outputs', '--export-outputs', '--eos',
            nargs=2, metavar=('CALCULATION_ID', 'TARGET_DIR'),
            help='Export all of the calculation outputs to the '
            'specified directory')
engine.opt('exports', 'Comma-separated string specifing the export formats, '
           'in order of priority')
engine.opt('log_level', 'Defaults to "info"',
           choices=['debug', 'info', 'warn', 'error', 'critical'])
engine.flg('reuse_hazard', 'Reuse the event based hazard if available')
<|endoftext|>


<commit_before>"""
sentry.options.store
~~~~~~~~~~~~~~~~~~~~

:copyright: (c) 2010-2015 by the Sentry Team, see AUTHORS for more details.
:license: BSD, see LICENSE for more details.
"""
from __future__ import absolute_import, print_function

import logging
from collections import namedtuple
from time import time

from django.utils import timezone
from sentry.db.models.query import create_or_update
from sentry.models import Option
from sentry.utils.hashlib import md5

Key = namedtuple('Key', ('name', 'default', 'type', 'flags', 'ttl', 'grace', 'cache_key'))

CACHE_FETCH_ERR = 'Unable to fetch option cache for %s'
CACHE_UPDATE_ERR = 'Unable to update option cache for %s'

logger = logging.getLogger('sentry')

def _make_cache_key(key):
    return 'o:%s' % md5(key).hexdigest()

def _make_cache_value(key, value):
    now = int(time())
    return (
        value,
        now + key.ttl,
        now + key.ttl + key.grace,
    )

class OptionsStore(object):
    """
    Abstraction for the Option storage logic that should be driven
    by the OptionsManager.

    OptionsStore is gooey and raw. It provides no protection over
    what goes into the store. It only knows that it's reading/writing
    to the right place. If using the OptionsStore directly, it's your
    job to do validation of the data. You should probably go through
    OptionsManager instead, unless you need raw access to something.
    """

    def __init__(self, cache=None, ttl=None):
        if cache is None:
            from sentry.cache import default_cache
            cache = default_cache
        self.cache = cache
        self.ttl = ttl
        self.flush_local_cache()

    def make_key(self, name, default, type, flags, ttl, grace):
        return Key(name, default, type, flags, int(ttl), int(grace), _make_cache_key(name))

    def get(self, key):
        """
        Fetches a value from the options store.
        """
        result = self.get_cache(key)
        if result is not None:
            return result

        result = self.get_store(key)
        if result is not None:
            return result

        # As a last ditch effort, let's hope we have a key
        # in local cache that's possibly stale
        return self.get_local_cache(key, grace=True)

    def get_cache(self, key):
        """
        First check agaist our local in-process cache, falling
        back to the network cache.
        """
        value = self.get_local_cache(key)
        if value is not None:
            return value

        cache_key = key.cache_key
        try:
            value = self.cache.get(cache_key)
        except Exception:
            logger.warn(CACHE_FETCH_ERR, key.name, exc_info=True)
            value = None
        else:
            if key.ttl > 0:
                self._local_cache[cache_key] = _make_cache_value(key, value)
        return value

    def get_local_cache(self, key, grace=False):
        """
        Attempt to fetch a key out of the local cache.

        If the key exists, but is beyond expiration, we only
        return it if grace=True. This forces the key to be returned
        in a disaster scenario as long as we're still holding onto it.
        This allows the OptionStore to pave over potential network hiccups
        by returning a stale value.
        """
        try:
            value, expires, grace = self._local_cache[key.cache_key]
        except KeyError:
            return None

        now = int(time())

        # Key is within normal expiry window, so just return it
        if now < expires:
            return value

        # If we're able to accept within grace window, return it
        if grace and now < grace:
            return value

        # Let's clean up values if we're beyond grace.
        if now > grace:
            del self._local_cache[key.cache_key]

        # If we're outside the grace window, even if we ask for it
        # in grace, too bad. The value is considered bad.
        return None

    def get_store(self, key):
        """
        Attempt to fetch value from the database. If successful,
        also set it back in the cache.
        """
        try:
            value = Option.objects.get(key=key.name).value
        except Option.DoesNotExist:
            value = None
        except Exception as e:
            logger.exception(unicode(e))
            value = None
        else:
            # we only attempt to populate the cache if we were previously
            # able to successfully talk to the backend
            # NOTE: There is definitely a race condition here between updating
            # the store and the cache
            try:
                self.set_cache(key, value)
            except Exception:
                logger.warn(CACHE_UPDATE_ERR, key.name, exc_info=True)
        return value

    def set(self, key, value):
        """
        Store a value in the option store. Value must get persisted to database first,
        then attempt caches. If it fails datastore, the entire operation blows up.
        If cache fails, we ignore silently since it'll get repaired later by sync_options.
        A boolean is returned to indicate if the network set succeeds.
        """
        self.set_store(key, value)
        return self.set_cache(key, value)

    def set_store(self, key, value):
        create_or_update(
            model=Option,
            key=key.name,
            values={
                'value': value,
                'last_updated': timezone.now(),
            }
        )

    def set_cache(self, key, value):
        cache_key = key.cache_key
        if key.ttl > 0:
            self._local_cache[cache_key] = _make_cache_value(key, value)
        try:
            self.cache.set(cache_key, value, self.ttl)
            return True
        except Exception:
            logger.warn(CACHE_UPDATE_ERR, key.name, exc_info=True)
            return False

    def delete(self, key):
        """
        Remove key out of option stores. This operation must succeed on the
        database first. If database fails, an exception is raised.
        If database succeeds, caches are then allowed to fail silently.
        A boolean is returned to indicate if the network deletion succeeds.
        """
        self.delete_store(key)
        return self.delete_cache(key)

    def delete_store(self, key):
        Option.objects.filter(key=key.name).delete()

    def delete_cache(self, key):
        cache_key = key.cache_key
        try:
            del self._local_cache[cache_key]
        except KeyError:
            pass
        try:
            self.cache.delete(cache_key)
            return True
        except Exception:
            logger.warn(CACHE_UPDATE_ERR, key, exc_info=True)
            return False

    def expire_local_cache(self):
        """
        Iterate over our local cache items, and
        remove the keys that are beyond their grace time.
        """
        if not self._local_cache:
            return

        to_expire = []
        now = int(time())
        for k, (_, _, grace) in self._local_cache.iteritems():
            if now > grace:
                to_expire.append(k)

        for k in to_expire:
            del self._local_cache[k]

    def flush_local_cache(self):
        """
        Empty store's local in-process cache.
        """
        self._local_cache = {}
<commit_msg>Guard against KeyError for race conditions<commit_after>"""
sentry.options.store
~~~~~~~~~~~~~~~~~~~~

:copyright: (c) 2010-2015 by the Sentry Team, see AUTHORS for more details.
:license: BSD, see LICENSE for more details.
"""
from __future__ import absolute_import, print_function

import logging
from collections import namedtuple
from time import time

from django.utils import timezone
from sentry.db.models.query import create_or_update
from sentry.models import Option
from sentry.utils.hashlib import md5

Key = namedtuple('Key', ('name', 'default', 'type', 'flags', 'ttl', 'grace', 'cache_key'))

CACHE_FETCH_ERR = 'Unable to fetch option cache for %s'
CACHE_UPDATE_ERR = 'Unable to update option cache for %s'

logger = logging.getLogger('sentry')

def _make_cache_key(key):
    return 'o:%s' % md5(key).hexdigest()

def _make_cache_value(key, value):
    now = int(time())
    return (
        value,
        now + key.ttl,
        now + key.ttl + key.grace,
    )

class OptionsStore(object):
    """
    Abstraction for the Option storage logic that should be driven
    by the OptionsManager.

    OptionsStore is gooey and raw. It provides no protection over
    what goes into the store. It only knows that it's reading/writing
    to the right place. If using the OptionsStore directly, it's your
    job to do validation of the data. You should probably go through
    OptionsManager instead, unless you need raw access to something.
    """

    def __init__(self, cache=None, ttl=None):
        if cache is None:
            from sentry.cache import default_cache
            cache = default_cache
        self.cache = cache
        self.ttl = ttl
        self.flush_local_cache()

    def make_key(self, name, default, type, flags, ttl, grace):
        return Key(name, default, type, flags, int(ttl), int(grace), _make_cache_key(name))

    def get(self, key):
        """
        Fetches a value from the options store.
        """
        result = self.get_cache(key)
        if result is not None:
            return result

        result = self.get_store(key)
        if result is not None:
            return result

        # As a last ditch effort, let's hope we have a key
        # in local cache that's possibly stale
        return self.get_local_cache(key, grace=True)

    def get_cache(self, key):
        """
        First check agaist our local in-process cache, falling
        back to the network cache.
        """
        value = self.get_local_cache(key)
        if value is not None:
            return value

        cache_key = key.cache_key
        try:
            value = self.cache.get(cache_key)
        except Exception:
            logger.warn(CACHE_FETCH_ERR, key.name, exc_info=True)
            value = None
        else:
            if key.ttl > 0:
                self._local_cache[cache_key] = _make_cache_value(key, value)
        return value

    def get_local_cache(self, key, grace=False):
        """
        Attempt to fetch a key out of the local cache.

        If the key exists, but is beyond expiration, we only
        return it if grace=True. This forces the key to be returned
        in a disaster scenario as long as we're still holding onto it.
        This allows the OptionStore to pave over potential network hiccups
        by returning a stale value.
        """
        try:
            value, expires, grace = self._local_cache[key.cache_key]
        except KeyError:
            return None

        now = int(time())

        # Key is within normal expiry window, so just return it
        if now < expires:
            return value

        # If we're able to accept within grace window, return it
        if grace and now < grace:
            return value

        # Let's clean up values if we're beyond grace.
        if now > grace:
            try:
                del self._local_cache[key.cache_key]
            except KeyError:
                # This could only exist in a race condition
                # where another thread has already deleted this key,
                # but we'll guard ourselves against it Justin Case.
                pass

        # If we're outside the grace window, even if we ask for it
        # in grace, too bad. The value is considered bad.
        return None

    def get_store(self, key):
        """
        Attempt to fetch value from the database. If successful,
        also set it back in the cache.
        """
        try:
            value = Option.objects.get(key=key.name).value
        except Option.DoesNotExist:
            value = None
        except Exception as e:
            logger.exception(unicode(e))
            value = None
        else:
            # we only attempt to populate the cache if we were previously
            # able to successfully talk to the backend
            # NOTE: There is definitely a race condition here between updating
            # the store and the cache
            try:
                self.set_cache(key, value)
            except Exception:
                logger.warn(CACHE_UPDATE_ERR, key.name, exc_info=True)
        return value

    def set(self, key, value):
        """
        Store a value in the option store. Value must get persisted to database first,
        then attempt caches. If it fails datastore, the entire operation blows up.
        If cache fails, we ignore silently since it'll get repaired later by sync_options.
        A boolean is returned to indicate if the network set succeeds.
        """
        self.set_store(key, value)
        return self.set_cache(key, value)

    def set_store(self, key, value):
        create_or_update(
            model=Option,
            key=key.name,
            values={
                'value': value,
                'last_updated': timezone.now(),
            }
        )

    def set_cache(self, key, value):
        cache_key = key.cache_key
        if key.ttl > 0:
            self._local_cache[cache_key] = _make_cache_value(key, value)
        try:
            self.cache.set(cache_key, value, self.ttl)
            return True
        except Exception:
            logger.warn(CACHE_UPDATE_ERR, key.name, exc_info=True)
            return False

    def delete(self, key):
        """
        Remove key out of option stores. This operation must succeed on the
        database first. If database fails, an exception is raised.
        If database succeeds, caches are then allowed to fail silently.
        A boolean is returned to indicate if the network deletion succeeds.
        """
        self.delete_store(key)
        return self.delete_cache(key)

    def delete_store(self, key):
        Option.objects.filter(key=key.name).delete()

    def delete_cache(self, key):
        cache_key = key.cache_key
        try:
            del self._local_cache[cache_key]
        except KeyError:
            pass
        try:
            self.cache.delete(cache_key)
            return True
        except Exception:
            logger.warn(CACHE_UPDATE_ERR, key.name, exc_info=True)
            return False

    def expire_local_cache(self):
        """
        Iterate over our local cache items, and
        remove the keys that are beyond their grace time.
        """
        if not self._local_cache:
            return

        to_expire = []
        now = int(time())
        for k, (_, _, grace) in self._local_cache.iteritems():
            if now > grace:
                to_expire.append(k)

        for k in to_expire:
            try:
                del self._local_cache[k]
            except KeyError:
                # This could only exist in a race condition
                # where another thread has already deleted this key,
                # but we'll guard ourselves against it Justin Case.
                pass

    def flush_local_cache(self):
        """
        Empty store's local in-process cache.
        """
        self._local_cache = {}
<|endoftext|>


<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
    crawler/crawler.py
    ~~~~~~~~~~~~~~~~~~

    Yeah, I'm a crawler. I can crawl a songlist, or one page of songlists or
    crawl all the songlist. I'm at your service.
"""

from __future__ import absolute_import

import re

import requests
from lxml import html

from crawler import config
from crawler import logger

class Songlist(object):
    def __init__(self, url):
        self.url = url
        self.tree = html.fromstring(requests.get(url).text)

    def _get_num(self, css_expression):
        text_in_tag = self.tree.cssselect(css_expression)[0].text_content()
        matched_num = re.search(r'\d+', text_in_tag)
        if matched_num is None:
            return 0
        num = int(matched_num.group())
        return num

    @property
    def name(self):
        return self.tree.cssselect('h2')[0].text

    @property
    def songlist_id(self):
        return int(re.search(r'(?<=id=)\d+$', self.url).group())

    @property
    def plays(self):
        return int(self.tree.cssselect('#play-count')[0].text)

    @property
    def comments(self):
        return self._get_num('.u-btni-cmmt i')

    @property
    def shares(self):
        return self._get_num('.u-btni-share i')

    @property
    def favourites(self):
        return self._get_num('.u-btni-fav i')

    @property
    def tags(self):
        if self.tree.cssselect('.tags'):
            tags = [tag.text for tag in self.tree.cssselect('.u-tag i')]
            return ', '.join(tags)
        else:
            return ''

    @property
    def meta(self):
        songlist_meta = {
            'name': self.name,
            'id': self.songlist_id,
            'url': self.url,
            'plays': self.plays,
            'comments': self.comments,
            'shares': self.shares,
            'favourites': self.favourites,
            'tags': self.tags
        }
        return songlist_meta

class Crawler(object):
    def __init__(self):
        self.redis_server = config.redis_server
        self.base_url = 'http://music.163.com'
        self.logger = logger.create_logger('crawler')

    def crawl_one_songlist(self, songlist_url):
        songlist = Songlist(songlist_url)

        key = 'wangyi:songlist:{0}'.format(songlist.songlist_id)
        self.redis_server.hmset(key, songlist.meta)
        self.logger.info('Crawled songlist: {0}'.format(songlist_url))

        if key not in self.redis_server.lrange('wangyi:songlists', 0, -1):
            self.redis_server.lpush('wangyi:songlists', key)

    def crawl_one_page(self, page_url):
        self.logger.info('Start to crawl the page: {0}'.format(page_url))
        tree = html.fromstring(requests.get(page_url).text)
        songlists = tree.cssselect('.u-cover > a')
        for songlist in songlists:
            self.crawl_one_songlist(self.base_url + songlist.get('href'))
        self.logger.info('The page: {0} has been crawled'.format(page_url))

        next_page = tree.cssselect('.znxt')[0].get('href')
        if next_page == 'javascript:void(0)':
            return None
        else:
            return self.base_url + next_page

    def crawl_the_site(self, start_url):
        self.logger.info('Start to crawl the site from: {0}'.format(start_url))
        next_page = self.crawl_one_page(start_url)
        while next_page is not None:
            next_page = self.crawl_one_page(next_page)
        self.logger.info(
            'Finish crawling the site from: {0}'.format(start_url))
<commit_msg>Make sure get the full num when it's high, fix #17<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
    crawler/crawler.py
    ~~~~~~~~~~~~~~~~~~

    Yeah, I'm a crawler. I can crawl a songlist, or one page of songlists or
    crawl all the songlist. I'm at your service.
"""

from __future__ import absolute_import

import re

import requests
from lxml import html

from crawler import config
from crawler import logger

class Songlist(object):
    def __init__(self, url):
        self.url = url
        self.tree = html.fromstring(requests.get(url).text)

    def _get_num(self, css_expression):
        text_in_tag = self.tree.cssselect(css_expression)[0].text_content()
        matched_num = re.search(r'\d+', text_in_tag)
        if matched_num is None:
            return 0
        num = int(matched_num.group())
        return num

    @property
    def name(self):
        return self.tree.cssselect('h2')[0].text

    @property
    def songlist_id(self):
        return int(re.search(r'(?<=id=)\d+$', self.url).group())

    @property
    def plays(self):
        return int(self.tree.cssselect('#play-count')[0].text)

    @property
    def comments(self):
        text_in_tag = self.tree.cssselect('.u-btni-cmmt i')[0].text_content()
        matched_num = re.search(r'\d+', text_in_tag)
        if matched_num is None:
            return 0
        num = int(matched_num.group())
        return num

    @property
    def shares(self):
        return int(self.tree.cssselect('.u-btni-share')[0].get('data-count'))

    @property
    def favourites(self):
        return int(self.tree.cssselect('.u-btni-fav')[0].get('data-count'))

    @property
    def tags(self):
        if self.tree.cssselect('.tags'):
            tags = [tag.text for tag in self.tree.cssselect('.u-tag i')]
            return ', '.join(tags)
        else:
            return ''

    @property
    def meta(self):
        songlist_meta = {
            'name': self.name,
            'id': self.songlist_id,
            'url': self.url,
            'plays': self.plays,
            'comments': self.comments,
            'shares': self.shares,
            'favourites': self.favourites,
            'tags': self.tags
        }
        return songlist_meta

class Crawler(object):
    def __init__(self):
        self.redis_server = config.redis_server
        self.base_url = 'http://music.163.com'
        self.logger = logger.create_logger('crawler')

    def crawl_one_songlist(self, songlist_url):
        songlist = Songlist(songlist_url)

        key = 'wangyi:songlist:{0}'.format(songlist.songlist_id)
        self.redis_server.hmset(key, songlist.meta)
        self.logger.info('Crawled songlist: {0}'.format(songlist_url))

        if key not in self.redis_server.lrange('wangyi:songlists', 0, -1):
            self.redis_server.lpush('wangyi:songlists', key)

    def crawl_one_page(self, page_url):
        self.logger.info('Start to crawl the page: {0}'.format(page_url))
        tree = html.fromstring(requests.get(page_url).text)
        songlists = tree.cssselect('.u-cover > a')
        for songlist in songlists:
            self.crawl_one_songlist(self.base_url + songlist.get('href'))
        self.logger.info('The page: {0} has been crawled'.format(page_url))

        next_page = tree.cssselect('.znxt')[0].get('href')
        if next_page == 'javascript:void(0)':
            return None
        else:
            return self.base_url + next_page

    def crawl_the_site(self, start_url):
        self.logger.info('Start to crawl the site from: {0}'.format(start_url))
        next_page = self.crawl_one_page(start_url)
        while next_page is not None:
            next_page = self.crawl_one_page(next_page)
        self.logger.info(
            'Finish crawling the site from: {0}'.format(start_url))
<|endoftext|>


<commit_before>
from . import extract_metadata
from . import detect_features
from . import match_features
from . import create_tracks
from . import reconstruct
from . import bundle
from . import mesh
from . import undistort
from . import compute_depthmaps
from . import export_ply
from . import export_openmvs
from . import export_visualsfm
from . import export_geocoords
from . import create_submodels
from . import align_submodels

opensfm_commands = [
    extract_metadata,
    detect_features,
    match_features,
    create_tracks,
    reconstruct,
    bundle,
    mesh,
    undistort,
    compute_depthmaps,
    export_ply,
    export_openmvs,
    export_visualsfm,
    export_geocoords,
    create_submodels,
    align_submodels,
]
<commit_msg>fix: add missing commands to make test_commands pass again<commit_after>
from . import extract_metadata
from . import detect_features
from . import match_features
from . import create_tracks
from . import reconstruct
from . import bundle
from . import mesh
from . import undistort
from . import compute_depthmaps
from . import export_ply
from . import export_openmvs
from . import export_visualsfm
from . import export_pmvs
from . import export_bundler
from . import export_colmap
from . import export_geocoords
from . import create_submodels
from . import align_submodels

opensfm_commands = [
    extract_metadata,
    detect_features,
    match_features,
    create_tracks,
    reconstruct,
    bundle,
    mesh,
    undistort,
    compute_depthmaps,
    export_ply,
    export_openmvs,
    export_visualsfm,
    export_pmvs,
    export_bundler,
    export_colmap,
    export_geocoords,
    create_submodels,
    align_submodels,
]
<|endoftext|>


<commit_before>"""
sentry.options.types
~~~~~~~~~~~~~~~~~~~~

:copyright: (c) 2010-2015 by the Sentry Team, see AUTHORS for more details.
:license: BSD, see LICENSE for more details.
"""
from __future__ import absolute_import, print_function

from sentry.utils.yaml import safe_load
from yaml.parser import ParserError
from yaml.scanner import ScannerError

__all__ = (
    'InvalidTypeError',
    'Any', 'Bool', 'Int', 'Float', 'String', 'Dict', 'Sequence',
)

class InvalidTypeError(TypeError):
    pass

class OptionType(object):
    "Base OptionType that provides type coersion"
    name = ''
    # Default value to be returned when initializing
    default = None
    # Types that do not need to be coerced
    expected_types = (object,)
    # Types that are acceptable for coersion
    compatible_types = (basestring,)

    def __call__(self, value=None):
        if value is None:
            return self._default()
        if self.test(value):
            return value
        if isinstance(value, self.compatible_types):
            rv = self.convert(value)
            # Make sure convert was able to do the right thing
            # and give us the type we were expecting
            if self.test(rv):
                return rv
        raise InvalidTypeError('{!r} is not a valid {}'.format(value, repr(self)))

    def convert(self, value):
        return value

    def _default(self):
        return self.default

    def test(self, value):
        "Check if the value needs to be coerced or not"
        return isinstance(value, self.expected_types)

    def __repr__(self):
        return self.name

class AnyType(OptionType):
    "A type that accepts any value and does no coersion"
    name = 'any'

class BoolType(OptionType):
    "Coerce a boolean from a string"
    name = 'boolean'
    default = False
    expected_types = (bool,)

    def convert(self, value):
        if value.lower() in ('y', 'yes', 't', 'true', '1'):
            return True
        if value.lower() in ('n', 'no', 'f', 'false', '0'):
            return False

class IntType(OptionType):
    "Coerce an integer from a string"
    name = 'integer'
    default = 0
    expected_types = (int,)

    def convert(self, value):
        try:
            return int(value)
        except ValueError:
            return

class FloatType(OptionType):
    "Coerce a float from a string or integer"
    name = 'float'
    default = 0.0
    expected_types = (float,)
    compatible_types = (basestring, float, int)

    def convert(self, value):
        try:
            return float(value)
        except ValueError:
            return

class StringType(OptionType):
    "String type without any coersion, must be a string"
    name = 'string'
    default = u''
    expected_types = (basestring,)
    compatible_types = (basestring,)

class DictType(OptionType):
    "Coerce a dict out of a json/yaml string"
    name = 'dictionary'
    expected_types = (dict,)

    def _default(self):
        # make sure we create a fresh dict each time
        return {}

    def convert(self, value):
        try:
            return safe_load(value)
        except (AttributeError, ParserError, ScannerError):
            return

class SequenceType(OptionType):
    "Coerce a tuple out of a json/yaml string or a list"
    name = 'sequence'
    default = ()
    expected_types = (tuple,)
    compatible_types = (basestring, tuple, list)

    def convert(self, value):
        if isinstance(value, basestring):
            try:
                value = safe_load(value)
            except (AttributeError, ParserError, ScannerError):
                return
        if isinstance(value, list):
            value = tuple(value)
        return value

# Initialize singletons of each type for easy reuse
Any = AnyType()
Bool = BoolType()
Int = IntType()
Float = FloatType()
String = StringType()
Dict = DictType()
Sequence = SequenceType()

# Mapping for basic types into what their OptionType is
_type_mapping = {
    bool: Bool,
    int: Int,
    float: Float,
    str: String,
    unicode: String,
    dict: Dict,
    tuple: Sequence,
    list: Sequence,
}

def type_from_value(value):
    "Fetch OptionType based on a primitive value"
    return _type_mapping[type(value)]
<commit_msg>Use triple quotes for docstrings<commit_after>"""
sentry.options.types
~~~~~~~~~~~~~~~~~~~~

:copyright: (c) 2010-2015 by the Sentry Team, see AUTHORS for more details.
:license: BSD, see LICENSE for more details.
"""
from __future__ import absolute_import, print_function

from sentry.utils.yaml import safe_load
from yaml.parser import ParserError
from yaml.scanner import ScannerError

__all__ = (
    'InvalidTypeError',
    'Any', 'Bool', 'Int', 'Float', 'String', 'Dict', 'Sequence',
)

class InvalidTypeError(TypeError):
    pass

class OptionType(object):
    """Base OptionType that provides type coersion"""
    name = ''
    # Default value to be returned when initializing
    default = None
    # Types that do not need to be coerced
    expected_types = (object,)
    # Types that are acceptable for coersion
    compatible_types = (basestring,)

    def __call__(self, value=None):
        if value is None:
            return self._default()
        if self.test(value):
            return value
        if isinstance(value, self.compatible_types):
            rv = self.convert(value)
            # Make sure convert was able to do the right thing
            # and give us the type we were expecting
            if self.test(rv):
                return rv
        raise InvalidTypeError('{!r} is not a valid {}'.format(value, repr(self)))

    def convert(self, value):
        return value

    def _default(self):
        return self.default

    def test(self, value):
        """Check if the value needs to be coerced or not"""
        return isinstance(value, self.expected_types)

    def __repr__(self):
        return self.name

class AnyType(OptionType):
    """A type that accepts any value and does no coersion"""
    name = 'any'

class BoolType(OptionType):
    "Coerce a boolean from a string"
    name = 'boolean'
    default = False
    expected_types = (bool,)

    def convert(self, value):
        if value.lower() in ('y', 'yes', 't', 'true', '1'):
            return True
        if value.lower() in ('n', 'no', 'f', 'false', '0'):
            return False

class IntType(OptionType):
    """Coerce an integer from a string"""
    name = 'integer'
    default = 0
    expected_types = (int,)

    def convert(self, value):
        try:
            return int(value)
        except ValueError:
            return

class FloatType(OptionType):
    """Coerce a float from a string or integer"""
    name = 'float'
    default = 0.0
    expected_types = (float,)
    compatible_types = (basestring, float, int)

    def convert(self, value):
        try:
            return float(value)
        except ValueError:
            return

class StringType(OptionType):
    """String type without any coersion, must be a string"""
    name = 'string'
    default = u''
    expected_types = (basestring,)
    compatible_types = (basestring,)

class DictType(OptionType):
    """Coerce a dict out of a json/yaml string"""
    name = 'dictionary'
    expected_types = (dict,)

    def _default(self):
        # make sure we create a fresh dict each time
        return {}

    def convert(self, value):
        try:
            return safe_load(value)
        except (AttributeError, ParserError, ScannerError):
            return

class SequenceType(OptionType):
    """Coerce a tuple out of a json/yaml string or a list"""
    name = 'sequence'
    default = ()
    expected_types = (tuple,)
    compatible_types = (basestring, tuple, list)

    def convert(self, value):
        if isinstance(value, basestring):
            try:
                value = safe_load(value)
            except (AttributeError, ParserError, ScannerError):
                return
        if isinstance(value, list):
            value = tuple(value)
        return value

# Initialize singletons of each type for easy reuse
Any = AnyType()
Bool = BoolType()
Int = IntType()
Float = FloatType()
String = StringType()
Dict = DictType()
Sequence = SequenceType()

# Mapping for basic types into what their OptionType is
_type_mapping = {
    bool: Bool,
    int: Int,
    float: Float,
    str: String,
    unicode: String,
    dict: Dict,
    tuple: Sequence,
    list: Sequence,
}

def type_from_value(value):
    """Fetch OptionType based on a primitive value"""
    return _type_mapping[type(value)]
<|endoftext|>