<commit_before>import os
import sys
import glob
import json
import tempfile
from io import BytesIO

if sys.version_info[0] < 3:
    import robotparser
else:
    from urllib import robotparser

import mock
from nose.tools import assert_equal, assert_raises
import requests
from .. import (Scraper, HTTPError, HTTPMethodUnavailableError,
                RobotExclusionError, urllib_URLError)
from .. import _user_agent as default_user_agent
from ..cache import FileCache

HTTPBIN = 'http://httpbin.org/'

class FakeResponse(object):
    def __init__(self, url, code, content, encoding='utf-8', headers=None):
        self.url = url
        self.status_code = code
        self.content = content
        self.text = str(content)
        self.encoding = encoding
        self.headers = headers or {}

def request_200(method, url, *args, **kwargs):
    return FakeResponse(url, 200, b'ok')
mock_200 = mock.Mock(wraps=request_200)

def test_constructor():
    # timeout=0 means None
    s = Scraper(timeout=0)
    assert s.timeout is None

def test_get():
    s = Scraper(requests_per_minute=0, follow_robots=False)
    resp = s.urlopen(HTTPBIN + 'get?woo=woo')
    assert_equal(resp.response.code, 200)
    assert_equal(json.loads(resp)['args']['woo'], 'woo')

def test_post():
    s = Scraper(requests_per_minute=0, follow_robots=False)
    resp = s.urlopen(HTTPBIN + 'post', 'POST', b'woo=woo')
    assert_equal(resp.response.code, 200)
    resp_json = json.loads(resp)
    assert_equal(resp_json['form']['woo'], 'woo')
    assert_equal(resp_json['headers']['Content-Type'],
                     'application/x-www-form-urlencoded')

def test_request_throttling():
    s = Scraper(requests_per_minute=30, follow_robots=False,
                accept_cookies=False)
    assert_equal(s.requests_per_minute, 30)

    mock_sleep = mock.Mock()
    mock_200.reset_mock()

    # check that sleep is called
    with mock.patch('time.sleep', mock_sleep):
        with mock.patch.object(s._session, 'request', mock_200):
            s.urlopen('http://dummy/')
            s.urlopen('http://dummy/')
            s.urlopen('http://dummy/')
            assert_equal(mock_sleep.call_count, 2)
            # should have slept for ~2 seconds
            assert 1.8 <= mock_sleep.call_args[0][0] <= 2.2
            assert_equal(mock_200.call_count, 3)

    # unthrottled, should be able to fit in lots of calls
    s.requests_per_minute = 0
    mock_200.reset_mock()
    mock_sleep.reset_mock()

    with mock.patch('time.sleep', mock_sleep):
        with mock.patch.object(s._session, 'request', mock_200):
            s.urlopen('http://dummy/')
            s.urlopen('http://dummy/')
            s.urlopen('http://dummy/')
            assert_equal(mock_sleep.call_count, 0)
            assert_equal(mock_200.call_count, 3)

def test_user_agent():
    s = Scraper(requests_per_minute=0, follow_robots=False)
    resp = s.urlopen(HTTPBIN + 'user-agent')
    ua = json.loads(resp)['user-agent']
    assert_equal(ua, default_user_agent)

    s.user_agent = 'a different agent'
    resp = s.urlopen(HTTPBIN + 'user-agent')
    ua = json.loads(resp)['user-agent']
    assert_equal(ua, 'a different agent')

def test_default_to_http():
    s = Scraper(requests_per_minute=0, follow_robots=False)
    with mock.patch.object(s._session, 'request', mock_200):
        assert_equal('http://dummy/', s.urlopen("dummy/").response.url)

def test_follow_robots():
    s = Scraper(requests_per_minute=0, follow_robots=True)

    with mock.patch.object(s._session, 'request', mock_200):
        # check that a robots.txt is created
        s.urlopen(HTTPBIN)
        assert HTTPBIN + 'robots.txt' in s._robot_parsers

        # set a fake robots.txt for http://dummy
        parser = robotparser.RobotFileParser()
        parser.parse(['User-agent: *', 'Disallow: /private/', 'Allow: /'])
        s._robot_parsers['http://dummy/robots.txt'] = parser

        # anything behind private fails
        with assert_raises(RobotExclusionError):
            s.urlopen("http://dummy/private/secret.html")
        # but others work
        assert_equal(200, s.urlopen("http://dummy/").response.code)

        # turn off follow_robots, everything works
        s.follow_robots = False
        assert_equal(200,
            s.urlopen("http://dummy/private/secret.html").response.code)

def test_error_context():
    # create error dir
    error_dir = tempfile.mkdtemp()
    s = Scraper(requests_per_minute=0, follow_robots=False,
                error_dir=error_dir)

    with mock.patch.object(s._session, 'request', mock_200):
        # try twice to test file name collision
        try:
            with s.urlopen("http://dummy/"):
                raise Exception('test')
        except:
            pass
        try:
            with s.urlopen("http://dummy/"):
                raise Exception('test')
        except:
            pass

    assert os.path.isfile(os.path.join(error_dir, "http:,,dummy,"))
    assert os.path.isfile(os.path.join(error_dir, "http:,,dummy,-1"))

    # tear down error dir
    for path in glob.iglob(os.path.join(error_dir, "*")):
        os.remove(path)
    os.rmdir(error_dir)

def test_404():
    s = Scraper(requests_per_minute=0, follow_robots=False)
    with assert_raises(HTTPError):
        s.urlopen(HTTPBIN + 'status/404')

    s.raise_errors = False
    resp = s.urlopen(HTTPBIN + 'status/404')
    assert_equal(404, resp.response.code)

def test_500():
    s = Scraper(requests_per_minute=0, follow_robots=False)

    with assert_raises(HTTPError):
        s.urlopen(HTTPBIN + 'status/500')

    s.raise_errors = False
    resp = s.urlopen(HTTPBIN + 'status/500')
    assert_equal(500, resp.response.code)

def test_follow_redirect():
    s = Scraper(requests_per_minute=0, follow_robots=False)

    redirect_url = HTTPBIN + 'redirect/1'
    final_url = HTTPBIN + 'get'

    resp = s.urlopen(redirect_url)
    assert_equal(final_url, resp.response.url)
    assert_equal(redirect_url, resp.response.requested_url)
    assert_equal(200, resp.response.code)

    s.follow_redirects = False
    resp = s.urlopen(redirect_url)
    assert_equal(redirect_url, resp.response.url)
    assert_equal(redirect_url, resp.response.requested_url)
    assert_equal(302, resp.response.code)

def test_caching():
    cache_dir = tempfile.mkdtemp()
    s = Scraper(requests_per_minute=0, follow_robots=False,
                cache_obj=FileCache(cache_dir))

    resp = s.urlopen(HTTPBIN + 'status/200')
    assert not resp.response.fromcache
    resp = s.urlopen(HTTPBIN + 'status/200')
    assert resp.response.fromcache

    for path in glob.iglob(os.path.join(cache_dir, "*")):
        os.remove(path)
    os.rmdir(cache_dir)

def test_urlretrieve():
    s = Scraper(requests_per_minute=0, follow_robots=False)

    with mock.patch.object(s._session, 'request', mock_200):
        fname, resp = s.urlretrieve("http://dummy/")
        with open(fname) as f:
            assert_equal(f.read(), 'ok')
            assert_equal(200, resp.code)
        os.remove(fname)

        (fh, set_fname) = tempfile.mkstemp()
        fname, resp = s.urlretrieve("http://dummy/", set_fname)
        assert_equal(fname, set_fname)
        with open(set_fname) as f:
            assert_equal(f.read(), 'ok')
            assert_equal(200, resp.code)
        os.remove(set_fname)

## TODO: on these retry tests it'd be nice to ensure that it tries
## 3 times for 500 and once for 404

def test_retry():
    s = Scraper(retry_attempts=3, retry_wait_seconds=0.001,
                follow_robots=False, raise_errors=False)

    # On the first call return a 500, then a 200
    mock_request = mock.Mock(side_effect=[
        FakeResponse('http://dummy/', 500, 'failure!'),
        FakeResponse('http://dummy/', 200, 'success!')
    ])

    with mock.patch.object(s._session, 'request', mock_request):
        resp = s.urlopen('http://dummy/')
    assert_equal(mock_request.call_count, 2)

    # 500 always
    mock_request = mock.Mock(return_value=FakeResponse('http://dummy/', 500,
                                                       'failure!'))

    with mock.patch.object(s._session, 'request', mock_request):
        resp = s.urlopen('http://dummy/')
    assert_equal(resp.response.code, 500)
    assert_equal(mock_request.call_count, 4)

def test_retry_404():
    s = Scraper(retry_attempts=3, retry_wait_seconds=0.001,
                follow_robots=False, raise_errors=False)

    # On the first call return a 404, then a 200
    mock_request = mock.Mock(side_effect=[
        FakeResponse('http://dummy/', 404, 'failure!'),
        FakeResponse('http://dummy/', 200, 'success!')
    ])

    with mock.patch.object(s._session, 'request', mock_request):
        resp = s.urlopen('http://dummy/', retry_on_404=True)
    assert_equal(mock_request.call_count, 2)

    # 404 always
    mock_request = mock.Mock(return_value=FakeResponse('http://dummy/', 404,
                                                       'failure!'))

    # retry on 404 true, 4 tries
    with mock.patch.object(s._session, 'request', mock_request):
        resp = s.urlopen('http://dummy/', retry_on_404=True)
    assert_equal(resp.response.code, 404)
    assert_equal(mock_request.call_count, 4)

    # retry on 404 false, just one more try
    with mock.patch.object(s._session, 'request', mock_request):
        resp = s.urlopen('http://dummy/', retry_on_404=False)
    assert_equal(resp.response.code, 404)
    assert_equal(mock_request.call_count, 5)

def test_timeout_retry():
    count = []

    # On the first call raise timeout
    def side_effect(*args, **kwargs):
        if count:
            return FakeResponse('http://dummy/', 200, 'success!')
        count.append(1)
        raise requests.Timeout('timed out :(')

    mock_request = mock.Mock(side_effect=side_effect)

    s = Scraper(retry_attempts=0, retry_wait_seconds=0.001,
                follow_robots=False)

    with mock.patch.object(s._session, 'request', mock_request):
        # first, try without retries
        # try only once, get the error
        with assert_raises(requests.Timeout):
            s.urlopen("http://dummy/")
        assert_equal(mock_request.call_count, 1)

    # reset and try again with retries
    mock_request.reset_mock()
    count = []
    s = Scraper(retry_attempts=2, retry_wait_seconds=0.001,
                follow_robots=False)
    with mock.patch.object(s._session, 'request', mock_request):
        resp = s.urlopen("http://dummy/")
        # get the result, take two tries
        assert_equal(resp, "success!")
        assert_equal(mock_request.call_count, 2)

def test_disable_compression():
    s = Scraper(disable_compression=True)

    headers = s._make_headers("http://dummy/")
    assert_equal(headers['accept-encoding'], 'text/*')

    # A supplied Accept-Encoding headers overrides the
    # disable_compression option
    s.headers['accept-encoding'] = '*'
    headers = s._make_headers('http://dummy/')
    assert_equal(headers['accept-encoding'], '*')

def test_callable_headers():
    s = Scraper(headers=lambda url: {'url': url})

    headers = s._make_headers('http://google.com')
    assert_equal(headers['url'], 'http://google.com')

    # Make sure it gets called freshly each time
    headers = s._make_headers('http://example.com')
    assert_equal(headers['url'], 'http://example.com')

def test_ftp_uses_urllib2():
    s = Scraper(requests_per_minute=0, follow_robots=False)
    urlopen = mock.Mock(return_value=BytesIO(b"ftp success!"))

    with mock.patch('scrapelib.urllib_urlopen', urlopen):
        r = s.urlopen('ftp://dummy/')
        assert r.response.code == 200
        assert r == "ftp success!"

def test_ftp_retries():
    count = []

    # On the first call raise URLError, then work
    def side_effect(*args, **kwargs):
        if count:
            return BytesIO(b"ftp success!")
        count.append(1)
        raise urllib_URLError('550: ftp failure!')

    mock_urlopen = mock.Mock(side_effect=side_effect)

    # retry on, retry_on_404 on (will retry due to 550 code)
    with mock.patch('scrapelib.urllib_urlopen', mock_urlopen):
        s = Scraper(retry_attempts=2, retry_wait_seconds=0.001,
                    follow_robots=False)
        r = s.urlopen('ftp://dummy/', retry_on_404=True)
        assert r == "ftp success!"
    assert_equal(mock_urlopen.call_count, 2)

    # retry off, retry_on_404 on (shouldn't matter)
    count = []
    mock_urlopen.reset_mock()
    with mock.patch('scrapelib.urllib_urlopen', mock_urlopen):
        s = Scraper(retry_attempts=0, retry_wait_seconds=0.001,
                    follow_robots=False)
        with assert_raises(urllib_URLError):
            s.urlopen('ftp://dummy/', retry_on_404=True)
    assert_equal(mock_urlopen.call_count, 1)

    # retry on, retry_on_404 off
    count = []
    mock_urlopen.reset_mock()
    with mock.patch('scrapelib.urllib_urlopen', mock_urlopen):
        s = Scraper(retry_attempts=2, retry_wait_seconds=0.001,
                    follow_robots=False)
        with assert_raises(urllib_URLError):
            s.urlopen('ftp://dummy/')
    assert_equal(mock_urlopen.call_count, 1)

def test_ftp_method_restrictions():
    s = Scraper(requests_per_minute=0, follow_robots=False)

    # only http(s) supports non-'GET' requests
    with assert_raises(HTTPMethodUnavailableError):
        s.urlopen("ftp://dummy/", method='POST')
<commit_msg>assert_raises for python 2.6<commit_after>import os
import sys
import glob
import json
import tempfile
from io import BytesIO

if sys.version_info[0] < 3:
    import robotparser
else:
    from urllib import robotparser

import mock
from nose.tools import assert_equal, assert_raises
import requests
from .. import (Scraper, HTTPError, HTTPMethodUnavailableError,
                RobotExclusionError, urllib_URLError)
from .. import _user_agent as default_user_agent
from ..cache import FileCache

HTTPBIN = 'http://httpbin.org/'

class FakeResponse(object):
    def __init__(self, url, code, content, encoding='utf-8', headers=None):
        self.url = url
        self.status_code = code
        self.content = content
        self.text = str(content)
        self.encoding = encoding
        self.headers = headers or {}

def request_200(method, url, *args, **kwargs):
    return FakeResponse(url, 200, b'ok')
mock_200 = mock.Mock(wraps=request_200)

def test_constructor():
    # timeout=0 means None
    s = Scraper(timeout=0)
    assert s.timeout is None

def test_get():
    s = Scraper(requests_per_minute=0, follow_robots=False)
    resp = s.urlopen(HTTPBIN + 'get?woo=woo')
    assert_equal(resp.response.code, 200)
    assert_equal(json.loads(resp)['args']['woo'], 'woo')

def test_post():
    s = Scraper(requests_per_minute=0, follow_robots=False)
    resp = s.urlopen(HTTPBIN + 'post', 'POST', b'woo=woo')
    assert_equal(resp.response.code, 200)
    resp_json = json.loads(resp)
    assert_equal(resp_json['form']['woo'], 'woo')
    assert_equal(resp_json['headers']['Content-Type'],
                     'application/x-www-form-urlencoded')

def test_request_throttling():
    s = Scraper(requests_per_minute=30, follow_robots=False,
                accept_cookies=False)
    assert_equal(s.requests_per_minute, 30)

    mock_sleep = mock.Mock()
    mock_200.reset_mock()

    # check that sleep is called
    with mock.patch('time.sleep', mock_sleep):
        with mock.patch.object(s._session, 'request', mock_200):
            s.urlopen('http://dummy/')
            s.urlopen('http://dummy/')
            s.urlopen('http://dummy/')
            assert_equal(mock_sleep.call_count, 2)
            # should have slept for ~2 seconds
            assert 1.8 <= mock_sleep.call_args[0][0] <= 2.2
            assert_equal(mock_200.call_count, 3)

    # unthrottled, should be able to fit in lots of calls
    s.requests_per_minute = 0
    mock_200.reset_mock()
    mock_sleep.reset_mock()

    with mock.patch('time.sleep', mock_sleep):
        with mock.patch.object(s._session, 'request', mock_200):
            s.urlopen('http://dummy/')
            s.urlopen('http://dummy/')
            s.urlopen('http://dummy/')
            assert_equal(mock_sleep.call_count, 0)
            assert_equal(mock_200.call_count, 3)

def test_user_agent():
    s = Scraper(requests_per_minute=0, follow_robots=False)
    resp = s.urlopen(HTTPBIN + 'user-agent')
    ua = json.loads(resp)['user-agent']
    assert_equal(ua, default_user_agent)

    s.user_agent = 'a different agent'
    resp = s.urlopen(HTTPBIN + 'user-agent')
    ua = json.loads(resp)['user-agent']
    assert_equal(ua, 'a different agent')

def test_default_to_http():
    s = Scraper(requests_per_minute=0, follow_robots=False)
    with mock.patch.object(s._session, 'request', mock_200):
        assert_equal('http://dummy/', s.urlopen("dummy/").response.url)

def test_follow_robots():
    s = Scraper(requests_per_minute=0, follow_robots=True)

    with mock.patch.object(s._session, 'request', mock_200):
        # check that a robots.txt is created
        s.urlopen(HTTPBIN)
        assert HTTPBIN + 'robots.txt' in s._robot_parsers

        # set a fake robots.txt for http://dummy
        parser = robotparser.RobotFileParser()
        parser.parse(['User-agent: *', 'Disallow: /private/', 'Allow: /'])
        s._robot_parsers['http://dummy/robots.txt'] = parser

        # anything behind private fails
        assert_raises(RobotExclusionError, s.urlopen,
                      "http://dummy/private/secret.html")
        # but others work
        assert_equal(200, s.urlopen("http://dummy/").response.code)

        # turn off follow_robots, everything works
        s.follow_robots = False
        assert_equal(200,
            s.urlopen("http://dummy/private/secret.html").response.code)

def test_error_context():
    # create error dir
    error_dir = tempfile.mkdtemp()
    s = Scraper(requests_per_minute=0, follow_robots=False,
                error_dir=error_dir)

    with mock.patch.object(s._session, 'request', mock_200):
        # try twice to test file name collision
        try:
            with s.urlopen("http://dummy/"):
                raise Exception('test')
        except:
            pass
        try:
            with s.urlopen("http://dummy/"):
                raise Exception('test')
        except:
            pass

    assert os.path.isfile(os.path.join(error_dir, "http:,,dummy,"))
    assert os.path.isfile(os.path.join(error_dir, "http:,,dummy,-1"))

    # tear down error dir
    for path in glob.iglob(os.path.join(error_dir, "*")):
        os.remove(path)
    os.rmdir(error_dir)

def test_404():
    s = Scraper(requests_per_minute=0, follow_robots=False)
    assert_raises(HTTPError, s.urlopen, HTTPBIN + 'status/404')

    s.raise_errors = False
    resp = s.urlopen(HTTPBIN + 'status/404')
    assert_equal(404, resp.response.code)

def test_500():
    s = Scraper(requests_per_minute=0, follow_robots=False)

    assert_raises(HTTPError, s.urlopen, HTTPBIN + 'status/500')

    s.raise_errors = False
    resp = s.urlopen(HTTPBIN + 'status/500')
    assert_equal(500, resp.response.code)

def test_follow_redirect():
    s = Scraper(requests_per_minute=0, follow_robots=False)

    redirect_url = HTTPBIN + 'redirect/1'
    final_url = HTTPBIN + 'get'

    resp = s.urlopen(redirect_url)
    assert_equal(final_url, resp.response.url)
    assert_equal(redirect_url, resp.response.requested_url)
    assert_equal(200, resp.response.code)

    s.follow_redirects = False
    resp = s.urlopen(redirect_url)
    assert_equal(redirect_url, resp.response.url)
    assert_equal(redirect_url, resp.response.requested_url)
    assert_equal(302, resp.response.code)

def test_caching():
    cache_dir = tempfile.mkdtemp()
    s = Scraper(requests_per_minute=0, follow_robots=False,
                cache_obj=FileCache(cache_dir))

    resp = s.urlopen(HTTPBIN + 'status/200')
    assert not resp.response.fromcache
    resp = s.urlopen(HTTPBIN + 'status/200')
    assert resp.response.fromcache

    for path in glob.iglob(os.path.join(cache_dir, "*")):
        os.remove(path)
    os.rmdir(cache_dir)

def test_urlretrieve():
    s = Scraper(requests_per_minute=0, follow_robots=False)

    with mock.patch.object(s._session, 'request', mock_200):
        fname, resp = s.urlretrieve("http://dummy/")
        with open(fname) as f:
            assert_equal(f.read(), 'ok')
            assert_equal(200, resp.code)
        os.remove(fname)

        (fh, set_fname) = tempfile.mkstemp()
        fname, resp = s.urlretrieve("http://dummy/", set_fname)
        assert_equal(fname, set_fname)
        with open(set_fname) as f:
            assert_equal(f.read(), 'ok')
            assert_equal(200, resp.code)
        os.remove(set_fname)

## TODO: on these retry tests it'd be nice to ensure that it tries
## 3 times for 500 and once for 404

def test_retry():
    s = Scraper(retry_attempts=3, retry_wait_seconds=0.001,
                follow_robots=False, raise_errors=False)

    # On the first call return a 500, then a 200
    mock_request = mock.Mock(side_effect=[
        FakeResponse('http://dummy/', 500, 'failure!'),
        FakeResponse('http://dummy/', 200, 'success!')
    ])

    with mock.patch.object(s._session, 'request', mock_request):
        resp = s.urlopen('http://dummy/')
    assert_equal(mock_request.call_count, 2)

    # 500 always
    mock_request = mock.Mock(return_value=FakeResponse('http://dummy/', 500,
                                                       'failure!'))

    with mock.patch.object(s._session, 'request', mock_request):
        resp = s.urlopen('http://dummy/')
    assert_equal(resp.response.code, 500)
    assert_equal(mock_request.call_count, 4)

def test_retry_404():
    s = Scraper(retry_attempts=3, retry_wait_seconds=0.001,
                follow_robots=False, raise_errors=False)

    # On the first call return a 404, then a 200
    mock_request = mock.Mock(side_effect=[
        FakeResponse('http://dummy/', 404, 'failure!'),
        FakeResponse('http://dummy/', 200, 'success!')
    ])

    with mock.patch.object(s._session, 'request', mock_request):
        resp = s.urlopen('http://dummy/', retry_on_404=True)
    assert_equal(mock_request.call_count, 2)

    # 404 always
    mock_request = mock.Mock(return_value=FakeResponse('http://dummy/', 404,
                                                       'failure!'))

    # retry on 404 true, 4 tries
    with mock.patch.object(s._session, 'request', mock_request):
        resp = s.urlopen('http://dummy/', retry_on_404=True)
    assert_equal(resp.response.code, 404)
    assert_equal(mock_request.call_count, 4)

    # retry on 404 false, just one more try
    with mock.patch.object(s._session, 'request', mock_request):
        resp = s.urlopen('http://dummy/', retry_on_404=False)
    assert_equal(resp.response.code, 404)
    assert_equal(mock_request.call_count, 5)

def test_timeout_retry():
    count = []

    # On the first call raise timeout
    def side_effect(*args, **kwargs):
        if count:
            return FakeResponse('http://dummy/', 200, 'success!')
        count.append(1)
        raise requests.Timeout('timed out :(')

    mock_request = mock.Mock(side_effect=side_effect)

    s = Scraper(retry_attempts=0, retry_wait_seconds=0.001,
                follow_robots=False)

    with mock.patch.object(s._session, 'request', mock_request):
        # first, try without retries
        # try only once, get the error
        assert_raises(requests.Timeout, s.urlopen, "http://dummy/")
        assert_equal(mock_request.call_count, 1)

    # reset and try again with retries
    mock_request.reset_mock()
    count = []
    s = Scraper(retry_attempts=2, retry_wait_seconds=0.001,
                follow_robots=False)
    with mock.patch.object(s._session, 'request', mock_request):
        resp = s.urlopen("http://dummy/")
        # get the result, take two tries
        assert_equal(resp, "success!")
        assert_equal(mock_request.call_count, 2)

def test_disable_compression():
    s = Scraper(disable_compression=True)

    headers = s._make_headers("http://dummy/")
    assert_equal(headers['accept-encoding'], 'text/*')

    # A supplied Accept-Encoding headers overrides the
    # disable_compression option
    s.headers['accept-encoding'] = '*'
    headers = s._make_headers('http://dummy/')
    assert_equal(headers['accept-encoding'], '*')

def test_callable_headers():
    s = Scraper(headers=lambda url: {'url': url})

    headers = s._make_headers('http://google.com')
    assert_equal(headers['url'], 'http://google.com')

    # Make sure it gets called freshly each time
    headers = s._make_headers('http://example.com')
    assert_equal(headers['url'], 'http://example.com')

def test_ftp_uses_urllib2():
    s = Scraper(requests_per_minute=0, follow_robots=False)
    urlopen = mock.Mock(return_value=BytesIO(b"ftp success!"))

    with mock.patch('scrapelib.urllib_urlopen', urlopen):
        r = s.urlopen('ftp://dummy/')
        assert r.response.code == 200
        assert r == "ftp success!"

def test_ftp_retries():
    count = []

    # On the first call raise URLError, then work
    def side_effect(*args, **kwargs):
        if count:
            return BytesIO(b"ftp success!")
        count.append(1)
        raise urllib_URLError('550: ftp failure!')

    mock_urlopen = mock.Mock(side_effect=side_effect)

    # retry on, retry_on_404 on (will retry due to 550 code)
    with mock.patch('scrapelib.urllib_urlopen', mock_urlopen):
        s = Scraper(retry_attempts=2, retry_wait_seconds=0.001,
                    follow_robots=False)
        r = s.urlopen('ftp://dummy/', retry_on_404=True)
        assert r == "ftp success!"
    assert_equal(mock_urlopen.call_count, 2)

    # retry off, retry_on_404 on (shouldn't matter)
    count = []
    mock_urlopen.reset_mock()
    with mock.patch('scrapelib.urllib_urlopen', mock_urlopen):
        s = Scraper(retry_attempts=0, retry_wait_seconds=0.001,
                    follow_robots=False)
        assert_raises(urllib_URLError, s.urlopen, 'ftp://dummy/',
                      retry_on_404=True)
    assert_equal(mock_urlopen.call_count, 1)

    # retry on, retry_on_404 off
    count = []
    mock_urlopen.reset_mock()
    with mock.patch('scrapelib.urllib_urlopen', mock_urlopen):
        s = Scraper(retry_attempts=2, retry_wait_seconds=0.001,
                    follow_robots=False)
        assert_raises(urllib_URLError, s.urlopen, 'ftp://dummy/')
    assert_equal(mock_urlopen.call_count, 1)

def test_ftp_method_restrictions():
    s = Scraper(requests_per_minute=0, follow_robots=False)

    # only http(s) supports non-'GET' requests
    assert_raises(HTTPMethodUnavailableError, s.urlopen, "ftp://dummy/",
                  method='POST')
<|endoftext|>


<commit_before>#!/usr/bin/python

# For inserting well-formed TZ-aware post dates into the database.
from datetime import datetime, timedelta
from pytz import utc

# So we can pause between API calls, per Tumblr's robots.txt.
from time import sleep

# For proper log message handling.
import logging

# Library for API calls, and secondary file with authentication keys.
import pytumblr
from secrets import consumer_key, secret_key, oauth_token, oauth_secret

# All of this is just so we can get and use the Django models for DB objects.
import django
from django.conf import settings
from django.db import connection
import sys
import os

sys.path.insert(0, "../../website")
os.environ.setdefault("DJANGO_SETTINGS_MODULE", "config.settings")
django.setup()

from noveltorpedo import models

class TumblrNotFound(ValueError):
    pass

# This is defined by Tumblr's API.
MAX_POSTS = 20

# Grab the Host object from the database and a client for connecting to the
# Tumblr API, since we'll need these in a couple of places.
host_attrs = {"url": "tumblr.com",
              "spider": "tumblr/fetch_tumblr.py",
              "wait": 1}
tumblr = models.Host.objects.get_or_create(**host_attrs)[0]
client = pytumblr.TumblrRestClient(consumer_key, secret_key,
                                   oauth_token, oauth_secret)

# Initialize the logger with basic stream output.
logger = logging.getLogger(__name__)
logger.addHandler(logging.StreamHandler())
logger.setLevel(logging.DEBUG)

def get_or_create_storyhost(blog):
    """Finds the StoryHost for a Tumblr blog, creating it if necessary.

    Either fetches the StoryHost corresponding to the given Tumblr username
    or creates it, initializing appropriate author and story rows at the same
    time. Doesn't currently handle the case where the blog has been added and
    then deleted, so could conceivably return a database entry for a blog
    that no longer exists.

    Args:
        blog: The username of a Tumblr blog (e.g. "litrocket-test").

    Returns:
        A StoryHost object corresponding to the given blog.

    Raises:
        TumblrNotFound: If there is no such blog and no existing StoryHost.
    """
    try:
        # Not using get_or_create() here, because if it doesn't exist, we
        # have more work to do than just creating a single object.
        storyhost = models.StoryHost.objects.get(url=blog+".tumblr.com")
    except models.StoryHost.DoesNotExist:
        try:
            blog_info = client.blog_info(blog)["blog"]
        except KeyError:
            raise TumblrNotFound(blog)
        story = models.Story()
        story.title = blog_info["title"]
        story.save()

        author = models.Author()
        author.name = blog
        author.save()
        story.authors.add(author)

        storyhost = models.StoryHost()
        storyhost.host = tumblr
        storyhost.url = blog + ".tumblr.com"
        storyhost.story = story
        # We initialize the last scraped time to the earliest time a datetime
        # can store. This is easier to compare later than a null value.
        storyhost.last_scraped = datetime.min.replace(tzinfo=utc)
        storyhost.save()
    return storyhost

def update_continuously(idle_time=10, minimum_delay=10):
    """Loops indefinitely, updating the database with new posts to Tumblr blogs.

    Specifically, finds the Tumblr StoryHost which is least recently updated,
    and if that update is longer ago than the minimum delay, checks it for new
    posts. Then repeats.

    Args:
        idle_time: Time, in seconds, to sleep when there's nothing to do.
        minimum_delay: Minimum time, in seconds, to wait before requesting
            updates to the same blog again.
    """
    logger.info("Entering continuous update loop.")
    while True:
        try:
            sh_manager = models.StoryHost.objects
            storyhost = sh_manager.filter(host=tumblr).earliest("last_scraped")
        except models.StoryHost.DoesNotExist:
            logger.debug("No storyhosts found. Idling for %s seconds.",
                         idle_time)
            sleep(idle_time)
            continue
        if (datetime.now(utc) - storyhost.last_scraped
                < timedelta(seconds=minimum_delay)):
            logger.debug("Least recent storyhost update is less than %s "
                         "seconds old. Idling for %s seconds.",
                         minimum_delay, idle_time)
            sleep(idle_time)
            continue
        logger.debug("Checking for updates to %s.", storyhost.story)
        update_story(storyhost)

def update_story(storyhost):
    """Fetches updates to the story corresponding to the given StoryHost.

    Makes one or more calls to the Tumblr API, adding segment entries for any
    posts which were published more recently than the last_updated value for
    the StoryHost. Will not currently notice posts that are backdated,
    deleted, or edited. Returns nothing.

    Args:
        storyhost: an existing StoryHost object.
    """
    # Start at the most recent post.
    offset = 0
    posts = get_posts(storyhost.url)
    try:
        post = posts.pop(0)
    except IndexError:
        # Blog has no text posts at all.
        # Set this in a way that will skip the update loop.
        post_date = storyhost.last_scraped
    else:
        post_date = datetime.fromtimestamp(post["timestamp"], utc)
    while post_date > storyhost.last_scraped:
        # Loop through posts, one at a time, as long as the current post we're
        # examining is more recent than the last time we updated the story.
        segment = models.StorySegment()
        segment.story = storyhost.story
        segment.title = post["title"]
        segment.contents = post["body"]
        segment.published = post_date
        segment.save()
        logger.debug("Adding segment: " + str(segment))
        try:
            post = posts.pop(0)
        except IndexError:
            # No more posts remain from the last API call.
            # Wait the amount of time that robots.txt requests, then
            # make another call to get more posts.
            sleep(tumblr.wait)
            offset = offset + MAX_POSTS
            posts = get_posts(storyhost.url, offset)
            try:
                post = posts.pop(0)
            except IndexError:
                # We got no posts from the new API call.
                # That means we've run out of posts overall.
                break
        # If we get back out here, we have a new post to loop on.
        post_date = datetime.fromtimestamp(post["timestamp"], utc)
    storyhost.last_scraped = datetime.now(utc)
    storyhost.save()

def get_posts(blog, offset=0, limit=MAX_POSTS):
    """Retrieves some text posts by the given account from the Tumblr API.

    Makes one call to the Tumblr API, requesting posts with the given
    parameters. Only asks for text posts, in the format that the user entered
    them (which may include HTML or Markdown, but needn't).

    Args:
        blog: The username of a Tumblr blog (e.g. "litrocket-test").
        offset: How many posts back from the most recent one to begin.
        limit: How many total posts to retrieve (bounded by the maximum
            defined by the API).

    Returns:
        A list of dictionaries, each one corresponding to a post, in
        descending order by publication timestamp. Relevant keys include
        "timestamp" (publication time in seconds since the epoch in UTC),
        "title", and "body".
    """
    return client.posts(blog, type="text", filter="text",
                        offset=offset, limit=limit)["posts"]

if __name__ == "__main__":
    """Add a Tumblr blog to the database, or start the update loop.

    If a username is given, attempts to add appropriate entries to the
    database (story, storyhost, and author) for that Tumblr account. Will
    not add it if those entries already exist, or if there is no such blog,
    but also won't notice if those entries already exist but the blog does
    not any more.

    With no arguments, will just start an infinite loop of checking for
    updates to the stories already recorded in the database. (Adding new
    stories while it's going, by running another instance with an argument
    or calling the get_or_create_storyhost() function from elsewhere, works
    and is the intended use case.)

    Command Line Args:
        (optional) The username of a Tumblr account (e.g. "litrocket-test").
    """
    try:
        get_or_create_storyhost(sys.argv[1])
    except TumblrNotFound as e:
        logger.error("No such tumblr: " + str(e))
    except IndexError:
        update_continuously()
<commit_msg>Capture and store NSFW flag on individual blogs.<commit_after>#!/usr/bin/python

# For inserting well-formed TZ-aware post dates into the database.
from datetime import datetime, timedelta
from pytz import utc

# So we can pause between API calls, per Tumblr's robots.txt.
from time import sleep

# For proper log message handling.
import logging

# Library for API calls, and secondary file with authentication keys.
import pytumblr
from secrets import consumer_key, secret_key, oauth_token, oauth_secret

# All of this is just so we can get and use the Django models for DB objects.
import django
from django.conf import settings
from django.db import connection
import sys
import os

sys.path.insert(0, "../../website")
os.environ.setdefault("DJANGO_SETTINGS_MODULE", "config.settings")
django.setup()

from noveltorpedo import models

class TumblrNotFound(ValueError):
    pass

# This is defined by Tumblr's API.
MAX_POSTS = 20

# Grab the Host object from the database and a client for connecting to the
# Tumblr API, since we'll need these in a couple of places.
host_attrs = {"url": "tumblr.com",
              "spider": "tumblr/fetch_tumblr.py",
              "wait": 1}
tumblr = models.Host.objects.get_or_create(**host_attrs)[0]
client = pytumblr.TumblrRestClient(consumer_key, secret_key,
                                   oauth_token, oauth_secret)

# Initialize the logger with basic stream output.
logger = logging.getLogger(__name__)
logger.addHandler(logging.StreamHandler())
logger.setLevel(logging.DEBUG)

def get_or_create_storyhost(blog):
    """Finds the StoryHost for a Tumblr blog, creating it if necessary.

    Either fetches the StoryHost corresponding to the given Tumblr username
    or creates it, initializing appropriate author and story rows at the same
    time. Doesn't currently handle the case where the blog has been added and
    then deleted, so could conceivably return a database entry for a blog
    that no longer exists.

    Args:
        blog: The username of a Tumblr blog (e.g. "litrocket-test").

    Returns:
        A StoryHost object corresponding to the given blog.

    Raises:
        TumblrNotFound: If there is no such blog and no existing StoryHost.
    """
    try:
        # Not using get_or_create() here, because if it doesn't exist, we
        # have more work to do than just creating a single object.
        storyhost = models.StoryHost.objects.get(url=blog+".tumblr.com")
    except models.StoryHost.DoesNotExist:
        try:
            blog_info = client.blog_info(blog)["blog"]
        except KeyError:
            raise TumblrNotFound(blog)
        story = models.Story()
        story.title = blog_info["title"]
        story.save()

        if blog_info["is_nsfw"]:
            nsfw = models.StoryAttribute()
            nsfw.story = story
            nsfw.key = "nsfw"
            nsfw.value = "TRUE"
            nsfw.save()

        author = models.Author()
        author.name = blog
        author.save()
        story.authors.add(author)

        storyhost = models.StoryHost()
        storyhost.host = tumblr
        storyhost.url = blog + ".tumblr.com"
        storyhost.story = story
        # We initialize the last scraped time to the earliest time a datetime
        # can store. This is easier to compare later than a null value.
        storyhost.last_scraped = datetime.min.replace(tzinfo=utc)
        storyhost.save()
    return storyhost

def update_continuously(idle_time=10, minimum_delay=10):
    """Loops indefinitely, updating the database with new posts to Tumblr blogs.

    Specifically, finds the Tumblr StoryHost which is least recently updated,
    and if that update is longer ago than the minimum delay, checks it for new
    posts. Then repeats.

    Args:
        idle_time: Time, in seconds, to sleep when there's nothing to do.
        minimum_delay: Minimum time, in seconds, to wait before requesting
            updates to the same blog again.
    """
    logger.info("Entering continuous update loop.")
    while True:
        try:
            sh_manager = models.StoryHost.objects
            storyhost = sh_manager.filter(host=tumblr).earliest("last_scraped")
        except models.StoryHost.DoesNotExist:
            logger.debug("No storyhosts found. Idling for %s seconds.",
                         idle_time)
            sleep(idle_time)
            continue
        if (datetime.now(utc) - storyhost.last_scraped
                < timedelta(seconds=minimum_delay)):
            logger.debug("Least recent storyhost update is less than %s "
                         "seconds old. Idling for %s seconds.",
                         minimum_delay, idle_time)
            sleep(idle_time)
            continue
        logger.debug("Checking for updates to %s.", storyhost.story)
        update_story(storyhost)

def update_story(storyhost):
    """Fetches updates to the story corresponding to the given StoryHost.

    Makes one or more calls to the Tumblr API, adding segment entries for any
    posts which were published more recently than the last_updated value for
    the StoryHost. Will not currently notice posts that are backdated,
    deleted, or edited. Returns nothing.

    Args:
        storyhost: an existing StoryHost object.
    """
    # Start at the most recent post.
    offset = 0
    posts = get_posts(storyhost.url)
    try:
        post = posts.pop(0)
    except IndexError:
        # Blog has no text posts at all.
        # Set this in a way that will skip the update loop.
        post_date = storyhost.last_scraped
    else:
        post_date = datetime.fromtimestamp(post["timestamp"], utc)
    while post_date > storyhost.last_scraped:
        # Loop through posts, one at a time, as long as the current post we're
        # examining is more recent than the last time we updated the story.
        segment = models.StorySegment()
        segment.story = storyhost.story
        segment.title = post["title"]
        segment.contents = post["body"]
        segment.published = post_date
        segment.save()
        logger.debug("Adding segment: " + str(segment))
        try:
            post = posts.pop(0)
        except IndexError:
            # No more posts remain from the last API call.
            # Wait the amount of time that robots.txt requests, then
            # make another call to get more posts.
            sleep(tumblr.wait)
            offset = offset + MAX_POSTS
            posts = get_posts(storyhost.url, offset)
            try:
                post = posts.pop(0)
            except IndexError:
                # We got no posts from the new API call.
                # That means we've run out of posts overall.
                break
        # If we get back out here, we have a new post to loop on.
        post_date = datetime.fromtimestamp(post["timestamp"], utc)
    storyhost.last_scraped = datetime.now(utc)
    storyhost.save()

def get_posts(blog, offset=0, limit=MAX_POSTS):
    """Retrieves some text posts by the given account from the Tumblr API.

    Makes one call to the Tumblr API, requesting posts with the given
    parameters. Only asks for text posts, in the format that the user entered
    them (which may include HTML or Markdown, but needn't).

    Args:
        blog: The username of a Tumblr blog (e.g. "litrocket-test").
        offset: How many posts back from the most recent one to begin.
        limit: How many total posts to retrieve (bounded by the maximum
            defined by the API).

    Returns:
        A list of dictionaries, each one corresponding to a post, in
        descending order by publication timestamp. Relevant keys include
        "timestamp" (publication time in seconds since the epoch in UTC),
        "title", and "body".
    """
    return client.posts(blog, type="text", filter="text",
                        offset=offset, limit=limit)["posts"]

if __name__ == "__main__":
    """Add a Tumblr blog to the database, or start the update loop.

    If a username is given, attempts to add appropriate entries to the
    database (story, storyhost, and author) for that Tumblr account. Will
    not add it if those entries already exist, or if there is no such blog,
    but also won't notice if those entries already exist but the blog does
    not any more.

    With no arguments, will just start an infinite loop of checking for
    updates to the stories already recorded in the database. (Adding new
    stories while it's going, by running another instance with an argument
    or calling the get_or_create_storyhost() function from elsewhere, works
    and is the intended use case.)

    Command Line Args:
        (optional) The username of a Tumblr account (e.g. "litrocket-test").
    """
    try:
        get_or_create_storyhost(sys.argv[1])
    except TumblrNotFound as e:
        logger.error("No such tumblr: " + str(e))
    except IndexError:
        update_continuously()
<|endoftext|>


<commit_before>from scrapy.command import ScrapyCommand, cmdline

class Command(ScrapyCommand):

    requires_project = False

    def syntax(self):
        return "<command>"

    def short_desc(self):
        return "Provides extended help for the given command"

    def run(self, args, opts):
        if not args:
            return False

        commands = cmdline.builtin_commands_dict()
        commands.update(cmdline.custom_commands_dict())

        cmdname = args[0]
        if cmdname in commands:
            cmd = commands[cmdname]
            help = getattr(cmd, 'help', None) or getattr(cmd, 'long_desc', None)
            print "%s: %s" % (cmdname, cmd.short_desc())
            print "usage: %s %s" % (cmdname, cmd.syntax())
            print
            print help()
        else:
            print "Unknown command: %s" % cmdname
<commit_msg>fixed bug in help command (thanks slav0nic for reporting)<commit_after>import textwrap

from scrapy.command import ScrapyCommand, cmdline

class Command(ScrapyCommand):

    requires_project = False

    def syntax(self):
        return "<command>"

    def short_desc(self):
        return "Provides extended help for the given command"

    def run(self, args, opts):
        if not args:
            return False

        commands = cmdline.get_commands_dict()

        cmdname = args[0]
        if cmdname in commands:
            cmd = commands[cmdname]
            help = getattr(cmd, 'help', None) or getattr(cmd, 'long_desc', None)
            title = "%s command" % cmdname
            print title
            print "-" * len(title)
            print
            print cmd.short_desc()
            print
            print "usage: %s %s" % (cmdname, cmd.syntax())
            print
            print "\n".join(textwrap.wrap(help()))
            print
            print "For a list of supported arguments use:"
            print
            print "   scrapy-ctl.py %s -h" % cmdname
            print
        else:
            print "Unknown command: %s" % cmdname
<|endoftext|>


<commit_before>from mailu import app, scheduler, dockercli

import subprocess
import os

def certbot_command(subcommand, *args):
    """ Run a certbot command while specifying the standard parameters.
    """
    command = [
        "certbot", subcommand,
        "-n",
        "--work-dir", "/tmp",
        "--logs-dir", "/tmp",
        "--config-dir", app.config["CERTS_PATH"],
        *args
    ]
    result = subprocess.run(command, stdout=subprocess.PIPE,
        stderr=subprocess.PIPE)
    return result

def certbot_install(domain):
    """ Install certificates for the given domain. Return True if a reload
    is required.
    """
    must_reload = False
    path = app.config["CERTS_PATH"]
    cert = os.path.join(path, "cert.pem")
    key = os.path.join(path, "key.pem")
    live_cert = os.path.join("live", domain, "fullchain.pem")
    live_key = os.path.join("live", domain, "privkey.pem")
    if not os.path.islink(cert) or os.readlink(cert) != live_cert:
        must_reload = True
        if os.path.exists(cert):
            os.unlink(cert)
        os.symlink(live_cert, cert)
    if not os.path.islink(key) or os.readlink(key) != live_key:
        must_reload = True
        if os.path.exists(key):
            os.unlink(key)
        os.symlink(live_key, key)
    return must_reload

@scheduler.scheduled_job('date')
@scheduler.scheduled_job('cron', hour=96, minute=0)
def generate_cert():
    print("Generating TLS certificates using Certbot")
    hostname = app.config["HOSTNAME"]
    email = "{}@{}".format(app.config["POSTMASTER"], app.config["DOMAIN"])
    result = certbot_command(
        "certonly",
        "--standalone",
        "--agree-tos",
        "--preferred-challenges", "http",
        "--email", email,
        "-d", hostname,
        # The port is hardcoded in the nginx image as well, we should find
        # a more suitable way to go but this will do until we have a proper
        # daemon handling certbot stuff
        "--http-01-port", "8081"
    )
    if result.returncode:
        print("Error while generating certificates:\n{}".format(
            result.stdout.decode("utf8") + result.stdout.decode("utf8")))
    else:
        print("Successfully generated or renewed TLS certificates")
        if certbot_install(hostname):
            print("Reloading TLS-dependant services")
            dockercli.reload("http", "smtp", "imap")
<commit_msg>Print stderr when certbot fails<commit_after>from mailu import app, scheduler, dockercli

import subprocess
import os

def certbot_command(subcommand, *args):
    """ Run a certbot command while specifying the standard parameters.
    """
    command = [
        "certbot", subcommand,
        "-n",
        "--work-dir", "/tmp",
        "--logs-dir", "/tmp",
        "--config-dir", app.config["CERTS_PATH"],
        *args
    ]
    result = subprocess.run(command, stdout=subprocess.PIPE,
        stderr=subprocess.PIPE)
    return result

def certbot_install(domain):
    """ Install certificates for the given domain. Return True if a reload
    is required.
    """
    must_reload = False
    path = app.config["CERTS_PATH"]
    cert = os.path.join(path, "cert.pem")
    key = os.path.join(path, "key.pem")
    live_cert = os.path.join("live", domain, "fullchain.pem")
    live_key = os.path.join("live", domain, "privkey.pem")
    if not os.path.islink(cert) or os.readlink(cert) != live_cert:
        must_reload = True
        if os.path.exists(cert):
            os.unlink(cert)
        os.symlink(live_cert, cert)
    if not os.path.islink(key) or os.readlink(key) != live_key:
        must_reload = True
        if os.path.exists(key):
            os.unlink(key)
        os.symlink(live_key, key)
    return must_reload

@scheduler.scheduled_job('date')
@scheduler.scheduled_job('cron', hour=96, minute=0)
def generate_cert():
    print("Generating TLS certificates using Certbot")
    hostname = app.config["HOSTNAME"]
    email = "{}@{}".format(app.config["POSTMASTER"], app.config["DOMAIN"])
    result = certbot_command(
        "certonly",
        "--standalone",
        "--agree-tos",
        "--preferred-challenges", "http",
        "--email", email,
        "-d", hostname,
        # The port is hardcoded in the nginx image as well, we should find
        # a more suitable way to go but this will do until we have a proper
        # daemon handling certbot stuff
        "--http-01-port", "8081"
    )
    if result.returncode:
        print("Error while generating certificates:\n{}".format(
            result.stdout.decode("utf8") + result.stderr.decode("utf8")))
    else:
        print("Successfully generated or renewed TLS certificates")
        if certbot_install(hostname):
            print("Reloading TLS-dependant services")
            dockercli.reload("http", "smtp", "imap")
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
from __future__ import unicode_literals

from collections import OrderedDict
from distutils.version import StrictVersion
from django import get_version
from django import forms
from django.conf import settings
from django.core.paginator import InvalidPage
from django.core.exceptions import PermissionDenied
from django.views.generic.edit import FormMixin
from django.views.generic import TemplateView
from django.http import HttpResponse
from django.utils.html import format_html
from django.shortcuts import render
from django.contrib.admin.templatetags.admin_static import static
from django.contrib.admin.options import IncorrectLookupParameters

ALL_VAR = 'all'
ORDER_VAR = 'o'
PAGE_VAR = 'p'
EXPORT_VAR = 'e'
CONTROL_VARS = [ALL_VAR, ORDER_VAR, PAGE_VAR, EXPORT_VAR]

class ReportList(object):

    def __init__(self, request, report):
        self.request = request
        self.report = report
        self.ordering_field_columns = self._get_ordering_field_columns()
        self.report.set_sort_params(*self._get_ordering())
        self.multi_page = False
        self.can_show_all = True
        self.paginator = None #self.report.get_paginator()
        try:
            self.page_num = int(self.request.GET.get(PAGE_VAR, 0))
        except ValueError:
            self.page_num = 0
        self.show_all = ALL_VAR in self.request.GET

    def get_query_string(self, new_params=None, remove=None):
        if new_params is None:
            new_params = {}
        if remove is None:
            remove = []
        params = self.request.GET.copy()
        for r in remove:
            for k in params.iterkeys():
                if k.startswith(r):
                    del params[k]

        for k, v in new_params.items():
            if v is None:
                if k in params:
                    del params[k]
            else:
                params[k] = v
        return '?%s' % params.urlencode()

    def _get_ordering(self):
        ordering = []
        order_params = self.request.GET.get(ORDER_VAR)
        if order_params:
            sort_values = order_params.split('.')
            fields = self.report.get_fields()
            for o in sort_values:
                if o.startswith('-'):
                    field = '-%s' % fields[int(o.replace('-', ''))][0]
                else:
                    field = fields[int(o)][0]
                ordering.append(field)
        return ordering

    def _get_ordering_field_columns(self):
        """
        Returns an OrderedDict of ordering field column numbers and asc/desc
        """
        # We must cope with more than one column having the same underlying sort
        # field, so we base things on column numbers.
        ordering = []
        ordering_fields = OrderedDict()
        if ORDER_VAR not in self.request.GET:
            # for ordering specified on ModelAdmin or model Meta, we don't know
            # the right column numbers absolutely, because there might be more
            # than one column associated with that ordering, so we guess.
            for field in ordering:
                if field.startswith('-'):
                    field = field[1:]
                    order_type = 'desc'
                else:
                    order_type = 'asc'
                for index, attr in enumerate(self.list_display):
                    if self.get_ordering_field(attr) == field:
                        ordering_fields[index] = order_type
                        break
        else:
            for p in self.request.GET[ORDER_VAR].split('.'):
                _, pfx, idx = p.rpartition('-')
                try:
                    idx = int(idx)
                except ValueError:
                    continue  # skip it
                ordering_fields[idx] = 'desc' if pfx == '-' else 'asc'
        return ordering_fields

    def headers(self):
        fields = self.report.get_fields()
        for i, field in enumerate(fields):
            name = field[0]
            label = field[1]
            if callable(getattr(self.report, name, name)):
                yield {
                    'label': label,
                    'class_attrib': format_html(' class="column-{0}"', name),
                    'sortable': False,
                }
                continue
            th_classes = ['sortable', 'column-{0}'.format(name)]
            order_type = ''
            new_order_type = 'asc'
            sort_priority = 0
            sorted_ = False
            # Is it currently being sorted on?
            if i in self.ordering_field_columns:
                sorted_ = True
                order_type = self.ordering_field_columns.get(i).lower()
                sort_priority = list(self.ordering_field_columns).index(i) + 1
                th_classes.append('sorted %sending' % order_type)
                new_order_type = {'asc': 'desc', 'desc': 'asc'}[order_type]
            # build new ordering param
            o_list_primary = []  # URL for making this field the primary sort
            o_list_remove = []  # URL for removing this field from sort
            o_list_toggle = []  # URL for toggling order type for this field
            make_qs_param = lambda t, n: ('-' if t == 'desc' else '') + str(n)
            for j, ot in self.ordering_field_columns.items():
                if j == i:  # Same column
                    param = make_qs_param(new_order_type, j)
                    # We want clicking on this header to bring the ordering to the
                    # front
                    o_list_primary.insert(0, param)
                    o_list_toggle.append(param)
                    # o_list_remove - omit
                else:
                    param = make_qs_param(ot, j)
                    o_list_primary.append(param)
                    o_list_toggle.append(param)
                    o_list_remove.append(param)
            if i not in self.ordering_field_columns:
                o_list_primary.insert(0, make_qs_param(new_order_type, i))
            yield {
                "label": label,
                "sortable": True,
                "sorted": sorted_,
                "ascending": order_type == "asc",
                "sort_priority": sort_priority,
                "url_primary": self.get_query_string({ORDER_VAR: '.'.join(o_list_primary)}),
                "url_remove": self.get_query_string({ORDER_VAR: '.'.join(o_list_remove)}),
                "url_toggle": self.get_query_string({ORDER_VAR: '.'.join(o_list_toggle)}),
                "class_attrib": format_html(' class="{0}"', ' '.join(th_classes)) if th_classes else '',
            }

    @property
    def totals(self):
        fields = self.report.get_fields()
        for idx, value in enumerate(self.report.iter_totals()):
            yield (self.report.get_alignment(fields[idx][0]), value)

    @property
    def results(self):
        fields = self.report.get_fields()
        for record in self.paginate():
            yield [(self.report.get_alignment(fields[idx][0]), value) for idx, value in enumerate(record)]

    def get_result_count(self):
        return len(self.report)

    def paginate(self):
        records = self.report.results
        self.paginator = self.report.get_paginator()
        result_count = self.paginator.count
        self.multi_page = result_count > self.report.get_list_per_page()
        self.can_show_all = result_count <= self.report.get_list_max_show_all()
        if not (self.show_all and self.can_show_all) and self.multi_page:
            try:
                records = self.paginator.page(self.page_num + 1).object_list
            except InvalidPage:
                raise IncorrectLookupParameters
        return records

class ReportView(TemplateView, FormMixin):

    report_class = None

    def dispatch(self, request, *args, **kwargs):
        return super(ReportView, self).dispatch(request, *args, **kwargs)

    def __init__(self, report_class, *args, **kwargs):
        super(ReportView, self).__init__(*args, **kwargs)
        self.report_class = report_class
        self.report = None

    def get_initial(self):
        initial = super(ReportView, self).get_initial()
        initial.update(self.report.get_initial())
        return initial

    @property
    def media(self):
        # taken from django.contrib.admin.options ModelAdmin
        extra = '' if settings.DEBUG else '.min'
        # if VERSION <= (1, 8):
        if StrictVersion(get_version()) < StrictVersion('1.9'):
            js = [
                'core.js',
                'admin/RelatedObjectLookups.js',
                'jquery%s.js' % extra,
                'jquery.init.js',
            ]
        else:
            js = [
                'core.js',
                'vendor/jquery/jquery%s.js' % extra,
                'jquery.init.js',
                'admin/RelatedObjectLookups.js',
                'actions%s.js' % extra,
                'urlify.js',
                'prepopulate%s.js' % extra,
                'vendor/xregexp/xregexp%s.js' % extra,
            ]
        return forms.Media(js=[static('admin/js/%s' % url) for url in js])

    def _export(self, form=None):
        if form is None:
            form = self.get_export_form()
        ctx = {
            'form': form,
            'back': '?%s' % '&'.join(['%s=%s' % param for param in self.request.GET.items()
                                      if param[0] != EXPORT_VAR]),
        }
        return render(self.request, 'admin/export.html', ctx)

    def post(self, request, *args, **kwargs):
        # TODO: FIXME: you shouldn't just pass request args to Report __init__() this is straight up dangerous.
        self.report = self.report_class(*args, **kwargs)
        if not self.report.has_permission(self.request):
            raise PermissionDenied()
        form = self.get_export_form(data=self.request.POST)
        if form.is_valid():
            context = self.get_context_data(**kwargs)
            filename = context['title'].lower().replace(' ', '_')
            response =  HttpResponse(content_type='text/csv')
            response['Content-Disposition'] = 'attachment;filename="%s.csv"' % filename
            self.report.to_csv(response, **form.cleaned_data)
            return response
        return self._export(form=form)

    def get(self, request, *args, **kwargs):
        self.report = self.report_class(*args, **kwargs)
        if not self.report.has_permission(request):
            raise PermissionDenied()
        if EXPORT_VAR in request.GET:
            return self._export()
        return super(ReportView, self).get(request, *args, **kwargs)

    def get_form_kwargs(self):
        kwargs = super(ReportView, self).get_form_kwargs()
        if self.request.method in ('GET', 'POST'):
            form_data = self.request.GET.copy()
            for key in CONTROL_VARS:
                if key in form_data:
                    del form_data[key]
            if form_data:
                kwargs.update({
                    'data': form_data,
                })
            else:
                kwargs.update({
                    'data': kwargs['initial']
                })
        return kwargs

    def get_form(self, form_class=None):
        if form_class is None:
            # If there's no form... there's no form.
            return None
        return super(ReportView, self).get_form(form_class)

    def get_form_class(self):
        return self.report_class.form_class

    def get_context_data(self, **kwargs):
        kwargs = super(ReportView, self).get_context_data(**kwargs)
        kwargs['media'] = self.media
        form = self.get_form(self.get_form_class())
        if form is not None:
            kwargs['form'] = form
            if form.is_valid():
                self.report.set_params(**form.cleaned_data)
        rl = ReportList(self.request, self.report)
        kwargs.update({
            'rl': rl,
            'title': self.report.get_title(),
            'has_filters': self.get_form_class() is not None,
            'help_text': self.report.get_help_text(),
            'description': self.report.get_description(),
            'export_path': rl.get_query_string({EXPORT_VAR: ''}),
            'totals': self.report.get_has_totals(),
            'totals_on_top': self.report.totals_on_top,
            'suit': ('suit' in settings.INSTALLED_APPS) or ('bootstrap_admin' in settings.INSTALLED_APPS),
        })
        return kwargs

    def get_template_names(self):
        return self.report.template_name

    def get_export_form(self, form_class=None, **kwargs):
        if form_class is None:
            form_class = self.report.get_export_form_class()
        return form_class(**kwargs)
<commit_msg>workaround missing opts in report's context<commit_after># -*- coding: utf-8 -*-
from __future__ import unicode_literals

import logging
from collections import OrderedDict
from distutils.version import StrictVersion

from django import get_version
from django import forms
from django.apps import apps
from django.conf import settings
from django.core.paginator import InvalidPage
from django.core.exceptions import PermissionDenied
from django.views.generic.edit import FormMixin
from django.views.generic import TemplateView
from django.http import HttpResponse
from django.utils.html import format_html
from django.shortcuts import render
from django.contrib.admin.templatetags.admin_static import static
from django.contrib.admin.options import IncorrectLookupParameters

logger = logging.getLogger(__name__)

ALL_VAR = 'all'
ORDER_VAR = 'o'
PAGE_VAR = 'p'
EXPORT_VAR = 'e'
CONTROL_VARS = [ALL_VAR, ORDER_VAR, PAGE_VAR, EXPORT_VAR]

class ReportList(object):

    def __init__(self, request, report):
        self.request = request
        self.report = report
        self.ordering_field_columns = self._get_ordering_field_columns()
        self.report.set_sort_params(*self._get_ordering())
        self.multi_page = False
        self.can_show_all = True
        self.paginator = None #self.report.get_paginator()
        try:
            self.page_num = int(self.request.GET.get(PAGE_VAR, 0))
        except ValueError:
            self.page_num = 0
        self.show_all = ALL_VAR in self.request.GET

    def get_query_string(self, new_params=None, remove=None):
        if new_params is None:
            new_params = {}
        if remove is None:
            remove = []
        params = self.request.GET.copy()
        for r in remove:
            for k in params.iterkeys():
                if k.startswith(r):
                    del params[k]

        for k, v in new_params.items():
            if v is None:
                if k in params:
                    del params[k]
            else:
                params[k] = v
        return '?%s' % params.urlencode()

    def _get_ordering(self):
        ordering = []
        order_params = self.request.GET.get(ORDER_VAR)
        if order_params:
            sort_values = order_params.split('.')
            fields = self.report.get_fields()
            for o in sort_values:
                if o.startswith('-'):
                    field = '-%s' % fields[int(o.replace('-', ''))][0]
                else:
                    field = fields[int(o)][0]
                ordering.append(field)
        return ordering

    def _get_ordering_field_columns(self):
        """
        Returns an OrderedDict of ordering field column numbers and asc/desc
        """
        # We must cope with more than one column having the same underlying sort
        # field, so we base things on column numbers.
        ordering = []
        ordering_fields = OrderedDict()
        if ORDER_VAR not in self.request.GET:
            # for ordering specified on ModelAdmin or model Meta, we don't know
            # the right column numbers absolutely, because there might be more
            # than one column associated with that ordering, so we guess.
            for field in ordering:
                if field.startswith('-'):
                    field = field[1:]
                    order_type = 'desc'
                else:
                    order_type = 'asc'
                for index, attr in enumerate(self.list_display):
                    if self.get_ordering_field(attr) == field:
                        ordering_fields[index] = order_type
                        break
        else:
            for p in self.request.GET[ORDER_VAR].split('.'):
                _, pfx, idx = p.rpartition('-')
                try:
                    idx = int(idx)
                except ValueError:
                    continue  # skip it
                ordering_fields[idx] = 'desc' if pfx == '-' else 'asc'
        return ordering_fields

    def headers(self):
        fields = self.report.get_fields()
        for i, field in enumerate(fields):
            name = field[0]
            label = field[1]
            if callable(getattr(self.report, name, name)):
                yield {
                    'label': label,
                    'class_attrib': format_html(' class="column-{0}"', name),
                    'sortable': False,
                }
                continue
            th_classes = ['sortable', 'column-{0}'.format(name)]
            order_type = ''
            new_order_type = 'asc'
            sort_priority = 0
            sorted_ = False
            # Is it currently being sorted on?
            if i in self.ordering_field_columns:
                sorted_ = True
                order_type = self.ordering_field_columns.get(i).lower()
                sort_priority = list(self.ordering_field_columns).index(i) + 1
                th_classes.append('sorted %sending' % order_type)
                new_order_type = {'asc': 'desc', 'desc': 'asc'}[order_type]
            # build new ordering param
            o_list_primary = []  # URL for making this field the primary sort
            o_list_remove = []  # URL for removing this field from sort
            o_list_toggle = []  # URL for toggling order type for this field
            make_qs_param = lambda t, n: ('-' if t == 'desc' else '') + str(n)
            for j, ot in self.ordering_field_columns.items():
                if j == i:  # Same column
                    param = make_qs_param(new_order_type, j)
                    # We want clicking on this header to bring the ordering to the
                    # front
                    o_list_primary.insert(0, param)
                    o_list_toggle.append(param)
                    # o_list_remove - omit
                else:
                    param = make_qs_param(ot, j)
                    o_list_primary.append(param)
                    o_list_toggle.append(param)
                    o_list_remove.append(param)
            if i not in self.ordering_field_columns:
                o_list_primary.insert(0, make_qs_param(new_order_type, i))
            yield {
                "label": label,
                "sortable": True,
                "sorted": sorted_,
                "ascending": order_type == "asc",
                "sort_priority": sort_priority,
                "url_primary": self.get_query_string({ORDER_VAR: '.'.join(o_list_primary)}),
                "url_remove": self.get_query_string({ORDER_VAR: '.'.join(o_list_remove)}),
                "url_toggle": self.get_query_string({ORDER_VAR: '.'.join(o_list_toggle)}),
                "class_attrib": format_html(' class="{0}"', ' '.join(th_classes)) if th_classes else '',
            }

    @property
    def totals(self):
        fields = self.report.get_fields()
        for idx, value in enumerate(self.report.iter_totals()):
            yield (self.report.get_alignment(fields[idx][0]), value)

    @property
    def results(self):
        fields = self.report.get_fields()
        for record in self.paginate():
            yield [(self.report.get_alignment(fields[idx][0]), value) for idx, value in enumerate(record)]

    def get_result_count(self):
        return len(self.report)

    def paginate(self):
        records = self.report.results
        self.paginator = self.report.get_paginator()
        result_count = self.paginator.count
        self.multi_page = result_count > self.report.get_list_per_page()
        self.can_show_all = result_count <= self.report.get_list_max_show_all()
        if not (self.show_all and self.can_show_all) and self.multi_page:
            try:
                records = self.paginator.page(self.page_num + 1).object_list
            except InvalidPage:
                raise IncorrectLookupParameters
        return records

class Opts(object):

    def __init__(self, report):
        self._report = report
        module = self._report.__class__.__module__
        app_config = apps.get_containing_app_config(module)
        self._app_label = app_config.label
        self._object_name = self._report.__class__.__name__

    def get_app_label(self):
        return self._app_label
    app_label = property(get_app_label)

    def get_object_name(self):
        return self._object_name
    object_name = property(get_object_name)

class ReportView(TemplateView, FormMixin):

    report_class = None

    def __init__(self, report_class, *args, **kwargs):
        super(ReportView, self).__init__(*args, **kwargs)
        self.report_class = report_class
        self.report = None

    def get_initial(self):
        initial = super(ReportView, self).get_initial()
        initial.update(self.report.get_initial())
        return initial

    @property
    def media(self):
        # taken from django.contrib.admin.options ModelAdmin
        extra = '' if settings.DEBUG else '.min'
        # if VERSION <= (1, 8):
        if StrictVersion(get_version()) < StrictVersion('1.9'):
            js = [
                'core.js',
                'admin/RelatedObjectLookups.js',
                'jquery%s.js' % extra,
                'jquery.init.js',
            ]
        else:
            js = [
                'core.js',
                'vendor/jquery/jquery%s.js' % extra,
                'jquery.init.js',
                'admin/RelatedObjectLookups.js',
                'actions%s.js' % extra,
                'urlify.js',
                'prepopulate%s.js' % extra,
                'vendor/xregexp/xregexp%s.js' % extra,
            ]
        return forms.Media(js=[static('admin/js/%s' % url) for url in js])

    def _export(self, form=None):
        if form is None:
            form = self.get_export_form()
        ctx = {
            'form': form,
            'back': '?%s' % '&'.join(['%s=%s' % param for param in self.request.GET.items()
                                      if param[0] != EXPORT_VAR]),
        }
        return render(self.request, 'admin/export.html', ctx)

    def post(self, request, *args, **kwargs):
        # TODO: FIXME: you shouldn't just pass request args to Report __init__() this is straight up dangerous.
        self.report = self.report_class(*args, **kwargs)
        if not self.report.has_permission(self.request):
            raise PermissionDenied()
        form = self.get_export_form(data=self.request.POST)
        if form.is_valid():
            context = self.get_context_data(**kwargs)
            filename = context['title'].lower().replace(' ', '_')
            response =  HttpResponse(content_type='text/csv')
            response['Content-Disposition'] = 'attachment;filename="%s.csv"' % filename
            self.report.to_csv(response, **form.cleaned_data)
            return response
        return self._export(form=form)

    def get(self, request, *args, **kwargs):
        self.report = self.report_class(*args, **kwargs)
        if not self.report.has_permission(request):
            raise PermissionDenied()
        if EXPORT_VAR in request.GET:
            return self._export()
        return super(ReportView, self).get(request, *args, **kwargs)

    def get_form_kwargs(self):
        kwargs = super(ReportView, self).get_form_kwargs()
        if self.request.method in ('GET', 'POST'):
            form_data = self.request.GET.copy()
            for key in CONTROL_VARS:
                if key in form_data:
                    del form_data[key]
            if form_data:
                kwargs.update({
                    'data': form_data,
                })
            else:
                kwargs.update({
                    'data': kwargs['initial']
                })
        return kwargs

    def get_form(self, form_class=None):
        if form_class is None:
            # If there's no form... there's no form.
            return None
        return super(ReportView, self).get_form(form_class)

    def get_form_class(self):
        return self.report_class.form_class

    def get_context_data(self, **kwargs):
        kwargs = super(ReportView, self).get_context_data(**kwargs)
        kwargs['media'] = self.media
        form = self.get_form(self.get_form_class())
        if form is not None:
            kwargs['form'] = form
            if form.is_valid():
                self.report.set_params(**form.cleaned_data)
        rl = ReportList(self.request, self.report)
        kwargs.update({
            'rl': rl,
            'opts': Opts(self.report),
            'title': self.report.get_title(),
            'has_filters': self.get_form_class() is not None,
            'help_text': self.report.get_help_text(),
            'description': self.report.get_description(),
            'export_path': rl.get_query_string({EXPORT_VAR: ''}),
            'totals': self.report.get_has_totals(),
            'totals_on_top': self.report.totals_on_top,
            'suit': ('suit' in settings.INSTALLED_APPS) or ('bootstrap_admin' in settings.INSTALLED_APPS),
        })
        return kwargs

    def get_template_names(self):
        return self.report.template_name

    def get_export_form(self, form_class=None, **kwargs):
        if form_class is None:
            form_class = self.report.get_export_form_class()
        return form_class(**kwargs)
<|endoftext|>


<commit_before>import numpy

class KalmanFilter:
    def __init__(self,A,C,Q,R,initial_x,initial_P):
        self.A = A # process update model
        self.C = C # observation model
        self.Q = Q # process covariance matrix
        self.R = R # measurement covariance matrix
        self.xhat_k1 = initial_x # a posteri state estimate from step (k-1)
        self.P_k1 = initial_P    # a posteri error estimate from step (k-1)

        self.ss = self.A.shape[0] # ndim in state space
        self.os = self.C.shape[0] # ndim in observation space
        self.AT = self.A.T
        self.CT = self.C.T

        self.n_skipped = 0

        if len(initial_x)!=self.ss:
            raise ValueError( 'initial_x must be a vector with ss components' )

    def step(self,y=None,isinitial=False):
        xhatminus, Pminus = self.step1__calculate_a_priori(isinitial=isinitial)
        return self.step2__calculate_a_posteri(xhatminus, Pminus, y=y)

    def step1__calculate_a_priori(self,isinitial=False):
        dot = numpy.dot # shorthand
        ############################################
        #          update state-space

        # compute a priori estimate of statespace
        if not isinitial:
            xhatminus = dot(self.A,self.xhat_k1)
            # compute a priori estimate of errors
            Pminus = dot(dot(self.A,self.P_k1),self.AT)+self.Q
        else:
            xhatminus = self.xhat_k1
            Pminus = self.P_k1

        return xhatminus, Pminus

    def step2__calculate_a_posteri(self,xhatminus,Pminus,y=None):
        """
        y represents the observation for this time-step
        """
        dot = numpy.dot # shorthand
        inv = numpy.linalg.inv

        ############################################
        #          incorporate observation

        if y is None:
            self.n_skipped += 1
        else:
            self.n_skipped = 0

        # With each skipped data point, measurement uncertainty doubles,
        # which means variance goes up 4x.

        factor = 2.0**self.n_skipped**2.0
        this_R = factor*self.R

        # calculate Kalman gain
        Knumerator = dot(Pminus,self.CT)
        Kdenominator = dot(dot(self.C,Pminus),self.CT)+this_R
        K = dot(Knumerator,inv(Kdenominator))

        # calculate a posteri state estimate
        if y is not None:
            residuals = y-dot(self.C,xhatminus) # error/innovation
            xhat = xhatminus+dot(K, residuals)
        else:
            xhat = xhatminus

        one_minus_KC = numpy.eye(self.ss)-dot(K,self.C)

        # compute a posteri estimate of errors
        P = dot(one_minus_KC,Pminus)

        # this step (k) becomes next step's prior (k-1)
        self.xhat_k1 = xhat
        self.P_k1 = P

        return xhat, P

def kalman_smoother(y,A,C,Q,R,init_x,init_V,valid_data_idx=None):
    """

    arguments
    ---------
    y - observations
    A - process update matrix
    C - state-to-observation matrix
    Q - process covariance matrix
    R - observation covariance matrix
    init_x - initial state
    init_V - initial error estimate
    valid_data_idx - (optional) indices to rows of y that are valid (None if all data valid)

    returns
    -------
    xsmooth - smoothed state estimates
    Vsmooth - smoothed error estimates

    Kalman smoother based on Kevin Murphy's Kalman toolbox for
    MATLAB(tm).

    N.B. Axes are swapped relative to Kevin Murphy's example, because
    in all my data, time is the first dimension."""

    def smooth_update(xsmooth_future,Vsmooth_future,xfilt,Vfilt,Vfilt_future,A,Q):
        dot = numpy.dot
        inv = numpy.linalg.inv

        xpred = dot(A,xfilt)
        Vpred = dot(A,numpy.dot(Vfilt,A.T)) + Q
        J = dot(Vfilt,numpy.dot(A.T,inv(Vpred))) # smoother gain matrix
        xsmooth = xfilt + dot(J, xsmooth_future-xpred)
        Vsmooth = Vfilt + dot(J,dot(Vsmooth_future-Vpred,J.T))
        return xsmooth, Vsmooth

    T, os = y.shape
    ss = len(A)

    kfilt = KalmanFilter(A,C,Q,R,init_x,init_V)
    # Forward pass
    xfilt = numpy.zeros((T,ss))
    Vfilt = numpy.zeros((T,ss,ss))

    for i in range(T):
        isinitial = i==0
        if i in valid_data_idx:
            y_i = y[i]
        else:
            y_i = None
        xfilt_i, Vfilt_i = kfilt.step(y=y_i,isinitial=isinitial)
        xfilt[i] = xfilt_i
        Vfilt[i] = Vfilt_i

    xsmooth = numpy.array(xfilt,copy=True)
    Vsmooth = numpy.array(Vfilt,copy=True)

    for t in range(T-2,-1,-1):
        xsmooth_t, Vsmooth_t = smooth_update(xsmooth[t+1,:],
                                             Vsmooth[t+1,:,:],
                                             xfilt[t,:],
                                             Vfilt[t,:,:],
                                             Vfilt[t,:,:],
                                             A,Q)
        xsmooth[t,:] = xsmooth_t
        Vsmooth[t,:,:] = Vsmooth_t

    return xsmooth, Vsmooth
<commit_msg>added comment<commit_after>import numpy

class KalmanFilter:
    def __init__(self,A,C,Q,R,initial_x,initial_P):
        self.A = A # process update model
        self.C = C # observation model
        self.Q = Q # process covariance matrix
        self.R = R # measurement covariance matrix
        self.xhat_k1 = initial_x # a posteri state estimate from step (k-1)
        self.P_k1 = initial_P    # a posteri error estimate from step (k-1)

        self.ss = self.A.shape[0] # ndim in state space
        self.os = self.C.shape[0] # ndim in observation space
        self.AT = self.A.T
        self.CT = self.C.T

        self.n_skipped = 0

        if len(initial_x)!=self.ss:
            raise ValueError( 'initial_x must be a vector with ss components' )

    def step(self,y=None,isinitial=False):
        xhatminus, Pminus = self.step1__calculate_a_priori(isinitial=isinitial)
        return self.step2__calculate_a_posteri(xhatminus, Pminus, y=y)

    def step1__calculate_a_priori(self,isinitial=False):
        dot = numpy.dot # shorthand
        ############################################
        #          update state-space

        # compute a priori estimate of statespace
        if not isinitial:
            xhatminus = dot(self.A,self.xhat_k1)
            # compute a priori estimate of errors
            Pminus = dot(dot(self.A,self.P_k1),self.AT)+self.Q
        else:
            xhatminus = self.xhat_k1
            Pminus = self.P_k1

        return xhatminus, Pminus

    def step2__calculate_a_posteri(self,xhatminus,Pminus,y=None):
        """
        y represents the observation for this time-step
        """
        dot = numpy.dot # shorthand
        inv = numpy.linalg.inv

        ############################################
        #          incorporate observation

        if y is None:
            self.n_skipped += 1
        else:
            self.n_skipped = 0

        # With each skipped data point, measurement uncertainty doubles,
        # which means variance goes up 4x.

        # XXX should re-visit this when frames are dropped!

        factor = 2.0**self.n_skipped**2.0
        this_R = factor*self.R

        # calculate Kalman gain
        Knumerator = dot(Pminus,self.CT)
        Kdenominator = dot(dot(self.C,Pminus),self.CT)+this_R
        K = dot(Knumerator,inv(Kdenominator))

        # calculate a posteri state estimate
        if y is not None:
            residuals = y-dot(self.C,xhatminus) # error/innovation
            xhat = xhatminus+dot(K, residuals)
        else:
            xhat = xhatminus

        one_minus_KC = numpy.eye(self.ss)-dot(K,self.C)

        # compute a posteri estimate of errors
        P = dot(one_minus_KC,Pminus)

        # this step (k) becomes next step's prior (k-1)
        self.xhat_k1 = xhat
        self.P_k1 = P

        return xhat, P

def kalman_smoother(y,A,C,Q,R,init_x,init_V,valid_data_idx=None):
    """

    arguments
    ---------
    y - observations
    A - process update matrix
    C - state-to-observation matrix
    Q - process covariance matrix
    R - observation covariance matrix
    init_x - initial state
    init_V - initial error estimate
    valid_data_idx - (optional) indices to rows of y that are valid (None if all data valid)

    returns
    -------
    xsmooth - smoothed state estimates
    Vsmooth - smoothed error estimates

    Kalman smoother based on Kevin Murphy's Kalman toolbox for
    MATLAB(tm).

    N.B. Axes are swapped relative to Kevin Murphy's example, because
    in all my data, time is the first dimension."""

    def smooth_update(xsmooth_future,Vsmooth_future,xfilt,Vfilt,Vfilt_future,A,Q):
        dot = numpy.dot
        inv = numpy.linalg.inv

        xpred = dot(A,xfilt)
        Vpred = dot(A,numpy.dot(Vfilt,A.T)) + Q
        J = dot(Vfilt,numpy.dot(A.T,inv(Vpred))) # smoother gain matrix
        xsmooth = xfilt + dot(J, xsmooth_future-xpred)
        Vsmooth = Vfilt + dot(J,dot(Vsmooth_future-Vpred,J.T))
        return xsmooth, Vsmooth

    T, os = y.shape
    ss = len(A)

    kfilt = KalmanFilter(A,C,Q,R,init_x,init_V)
    # Forward pass
    xfilt = numpy.zeros((T,ss))
    Vfilt = numpy.zeros((T,ss,ss))

    for i in range(T):
        isinitial = i==0
        if i in valid_data_idx:
            y_i = y[i]
        else:
            y_i = None
        xfilt_i, Vfilt_i = kfilt.step(y=y_i,isinitial=isinitial)
        xfilt[i] = xfilt_i
        Vfilt[i] = Vfilt_i

    xsmooth = numpy.array(xfilt,copy=True)
    Vsmooth = numpy.array(Vfilt,copy=True)

    for t in range(T-2,-1,-1):
        xsmooth_t, Vsmooth_t = smooth_update(xsmooth[t+1,:],
                                             Vsmooth[t+1,:,:],
                                             xfilt[t,:],
                                             Vfilt[t,:,:],
                                             Vfilt[t,:,:],
                                             A,Q)
        xsmooth[t,:] = xsmooth_t
        Vsmooth[t,:,:] = Vsmooth_t

    return xsmooth, Vsmooth
<|endoftext|>


<commit_before>def distance(n):
    if n == 1:
        return 0

    c = 1
    layer = 0

    while c**2 <= n:
        c += 2
        layer += 1

    p_start = (c-2)**2
    p_ur = p_start + c - 1
    p_ul = p_ur + c - 1
    p_ll = p_ul + c - 1
    p_lr = p_ll + c - 1

    offset = layer

    if n < p_ur:
        offset = n - (p_ur - layer)
        # print(n, p_ur, layer, (p_ur - layer))
    elif n < p_ul:
        offset = n - (p_ul - layer)
    elif n < p_ll:
        offset = n - (p_ll - layer)
    elif n < p_lr:
        offset = n - (p_lr - layer)

    offset = abs(offset)
    # print(offset, layer, p_start, p_ur, p_ul, p_ll, p_lr, n, c, c**2)
    return layer + offset

def test_distance():
    assert 0 == distance(1)
    assert 3 == distance(12)
    assert 2 == distance(23)
    assert 31 == distance(1024)

if __name__ == "__main__":
    print(distance(289326))
<commit_msg>Whitespace<commit_after>def distance(n):
    if n == 1:
        return 0

    c = 1
    layer = 0

    while c**2 <= n:
        c += 2
        layer += 1

    p_start = (c-2)**2
    p_ur = p_start + c - 1
    p_ul = p_ur + c - 1
    p_ll = p_ul + c - 1
    p_lr = p_ll + c - 1

    offset = layer

    if n < p_ur:
        offset = n - (p_ur - layer)
        # print(n, p_ur, layer, (p_ur - layer))
    elif n < p_ul:
        offset = n - (p_ul - layer)
    elif n < p_ll:
        offset = n - (p_ll - layer)
    elif n < p_lr:
        offset = n - (p_lr - layer)

    offset = abs(offset)
    # print(offset, layer, p_start, p_ur, p_ul, p_ll, p_lr, n, c, c**2)
    return layer + offset

def test_distance():
    assert 0 == distance(1)
    assert 3 == distance(12)
    assert 2 == distance(23)
    assert 31 == distance(1024)

if __name__ == "__main__":
    print(distance(289326))
<|endoftext|>


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

from computations import ZScores

def patch():
    """
    Patch the features of this library onto agate's core :class:`.Table` and :class:`.TableSet`.
    """
    import agate
    from agatestats.table import TableStats

    agate.Table.monkeypatch(TableStats)
<commit_msg>Fix computations import.<commit_after>#!/usr/bin/env python

from agatestats.computations import ZScores

def patch():
    """
    Patch the features of this library onto agate's core :class:`.Table` and :class:`.TableSet`.
    """
    import agate
    from agatestats.table import TableStats

    agate.Table.monkeypatch(TableStats)
<|endoftext|>


<commit_before>import django_filters
from crispy_forms.helper import FormHelper
from crispy_forms.layout import Layout, Row, Submit, Div
from django import forms

from agir.events.models import EventSubtype, Event

class EventFilterForm(forms.Form):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

        self.helper = FormHelper()
        self.helper.form_method = "GET"
        self.helper.layout = Layout(
            Row(Div("q", css_class="col-md-8"), Div("type", css_class="col-md-4")),
            "past",
        )
        self.helper.add_input(Submit("chercher", "Chercher"))

class EventFilter(django_filters.FilterSet):
    q = django_filters.CharFilter(method="filter_search", label="Chercher")

    type = django_filters.ChoiceFilter(
        field_name="subtype__type",
        lookup_expr="exact",
        label="Type d'événement",
        choices=EventSubtype.TYPE_CHOICES,
    )

    past = django_filters.BooleanFilter(
        label="Inclure les événements passés",
        method="filter_include_past",
        widget=forms.CheckboxInput(),
    )

    def __init__(self, data=None, *args, **kwargs):
        if data is not None:
            data = data.copy()
            if data.get("include_past") is None:
                data["include_past"] = False
            if data.get("include_hidden") is None:
                data["include_hidden"] = False

        super().__init__(data, *args, **kwargs)

    def filter_include_past(self, qs, name, value):
        if not value:
            return qs.upcoming(published_only=False)
        else:
            return qs

    def filter_search(self, qs, name, terms):
        if terms:
            return qs.search(terms)
        return qs

    class Meta:
        model = Event
        fields = []
        form = EventFilterForm

class EventAPIFilter(EventFilter, django_filters.rest_framework.FilterSet):
    pass
<commit_msg>fix: missing subtype filter on event map<commit_after>import django_filters
from crispy_forms.helper import FormHelper
from crispy_forms.layout import Layout, Row, Submit, Div
from django import forms

from agir.events.models import EventSubtype, Event
from agir.lib.filters import FixedModelMultipleChoiceFilter

class EventFilterForm(forms.Form):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

        self.helper = FormHelper()
        self.helper.form_method = "GET"
        self.helper.layout = Layout(
            Row(Div("q", css_class="col-md-8"), Div("type", css_class="col-md-4")),
            "past",
        )
        self.helper.add_input(Submit("chercher", "Chercher"))

class EventFilter(django_filters.rest_framework.FilterSet):
    q = django_filters.CharFilter(method="filter_search", label="Chercher")

    type = django_filters.ChoiceFilter(
        field_name="subtype__type",
        lookup_expr="exact",
        label="Type d'événement",
        choices=EventSubtype.TYPE_CHOICES,
    )

    subtype = FixedModelMultipleChoiceFilter(
        field_name="subtype",
        to_field_name="label",
        queryset=EventSubtype.objects.all(),
    )

    past = django_filters.BooleanFilter(
        label="Inclure les événements passés",
        method="filter_include_past",
        widget=forms.CheckboxInput(),
    )

    def __init__(self, data=None, *args, **kwargs):
        if data is not None:
            data = data.copy()
            if data.get("include_past") is None:
                data["include_past"] = False
            if data.get("include_hidden") is None:
                data["include_hidden"] = False

        super().__init__(data, *args, **kwargs)

    def filter_include_past(self, qs, name, value):
        if not value:
            return qs.upcoming(published_only=False)
        else:
            return qs

    def filter_search(self, qs, name, terms):
        if terms:
            return qs.search(terms)
        return qs

    class Meta:
        model = Event
        fields = []
        form = EventFilterForm

class EventAPIFilter(EventFilter, django_filters.rest_framework.FilterSet):
    pass
<|endoftext|>