<commit_before>from math import copysign

from odoo import _, api, exceptions, fields, models
from odoo.tools.float_utils import float_compare

class AccountTax(models.Model):
    """Inherit to implement the tax using avatax API"""

    _inherit = "account.tax"

    is_avatax = fields.Boolean("Is Avatax")

    @api.model
    def _get_avalara_tax_domain(self, tax_rate, doc_type):
        return [
            ("amount", "=", tax_rate),
            ("is_avatax", "=", True),
            (
                "company_id",
                "=",
                self.env.company.id,
            ),
        ]

    @api.model
    def _get_avalara_tax_name(self, tax_rate, doc_type=None):
        return _("{}%*").format(str(tax_rate))

    @api.model
    def get_avalara_tax(self, tax_rate, doc_type):
        domain = self._get_avalara_tax_domain(tax_rate, doc_type)
        tax = self.with_context(active_test=False).search(domain, limit=1)
        if tax and not tax.active:
            tax.active = True
        if not tax:
            domain = self._get_avalara_tax_domain(0, doc_type)
            tax_template = self.search(domain, limit=1)
            if not tax_template:
                raise exceptions.UserError(
                    _("Please configure Avatax Tax for Company %s:")
                    % self.env.company.name
                )
            # If you get a unique constraint error here,
            # check the data for your existing Avatax taxes.
            vals = {
                "amount": tax_rate,
                "name": self._get_avalara_tax_name(tax_rate, doc_type),
            }
            tax = tax_template.sudo().copy(default=vals)
            # Odoo core does not use the name set in default dict
            tax.name = vals.get("name")
        return tax

    def compute_all(
        self,
        price_unit,
        currency=None,
        quantity=1.0,
        product=None,
        partner=None,
        is_refund=False,
        handle_price_include=True,
    ):
        """
        Adopted as the central point to inject custom tax computations.
        Avatax logic is triggered if the "avatax_invoice" is set in the context.
        To find the Avatax amount, we search an Invoice line with the same
        quantity, price and product.
        """
        res = super().compute_all(
            price_unit,
            currency,
            quantity,
            product,
            partner,
            is_refund,
            handle_price_include,
        )
        avatax_invoice = self.env.context.get("avatax_invoice")
        if avatax_invoice:
            # Find the Avatax amount in the invoice Lines
            # Looks up the line for the current product, price_unit, and quantity
            # Note that the price_unit used must consider discount
            base = res["total_excluded"]
            digits = 6
            avatax_amount = None
            for line in avatax_invoice.invoice_line_ids:
                price_unit = line.currency_id._convert(
                    price_unit,
                    avatax_invoice.company_id.currency_id,
                    avatax_invoice.company_id,
                    avatax_invoice.date,
                )
                if (
                    line.product_id == product
                    and float_compare(line.quantity, quantity, digits) == 0
                ):
                    line_price = line._get_avatax_amount(qty=1)
                    if float_compare(line_price, -price_unit, digits) == 0:
                        avatax_amount = copysign(line.avatax_amt_line, base)
                        break
            if avatax_amount is None:
                avatax_amount = 0.0
                raise exceptions.UserError(
                    _(
                        "Incorrect retrieval of Avatax amount for Invoice %s:"
                        " product %s, price_unit %f, quantity %f"
                    )
                    % (avatax_invoice, product.display_name, -price_unit, quantity)
                )
            for tax_item in res["taxes"]:
                if tax_item["amount"] != 0:
                    tax_item["amount"] = avatax_amount
            res["total_included"] = base + avatax_amount
        return res
<commit_msg>[FIX] Improved the account_avatax module code.<commit_after>from math import copysign

from odoo import _, api, exceptions, fields, models
from odoo.tools.float_utils import float_compare

class AccountTax(models.Model):
    """Inherit to implement the tax using avatax API"""

    _inherit = "account.tax"

    is_avatax = fields.Boolean("Is Avatax")

    @api.model
    def _get_avalara_tax_domain(self, tax_rate, doc_type):
        return [
            ("amount", "=", tax_rate),
            ("is_avatax", "=", True),
            (
                "company_id",
                "=",
                self.env.company.id,
            ),
        ]

    @api.model
    def _get_avalara_tax_name(self, tax_rate, doc_type=None):
        return _("{}%*").format(str(tax_rate))

    @api.model
    def get_avalara_tax(self, tax_rate, doc_type):
        domain = self._get_avalara_tax_domain(tax_rate, doc_type)
        tax = self.with_context(active_test=False).search(domain, limit=1)
        if tax and not tax.active:
            tax.active = True
        if not tax:
            domain = self._get_avalara_tax_domain(0, doc_type)
            tax_template = self.search(domain, limit=1)
            if not tax_template:
                raise exceptions.UserError(
                    _("Please configure Avatax Tax for Company %s:")
                    % self.env.company.name
                )
            # If you get a unique constraint error here,
            # check the data for your existing Avatax taxes.
            vals = {
                "amount": tax_rate,
                "name": self._get_avalara_tax_name(tax_rate, doc_type),
            }
            tax = tax_template.sudo().copy(default=vals)
            # Odoo core does not use the name set in default dict
            tax.name = vals.get("name")
        return tax

    def compute_all(
        self,
        price_unit,
        currency=None,
        quantity=1.0,
        product=None,
        partner=None,
        is_refund=False,
        handle_price_include=True,
    ):
        """
        Adopted as the central point to inject custom tax computations.
        Avatax logic is triggered if the "avatax_invoice" is set in the context.
        To find the Avatax amount, we search an Invoice line with the same
        quantity, price and product.
        """
        res = super().compute_all(
            price_unit,
            currency,
            quantity,
            product,
            partner,
            is_refund,
            handle_price_include,
        )
        avatax_invoice = self.env.context.get("avatax_invoice")
        if avatax_invoice:
            # Find the Avatax amount in the invoice Lines
            # Looks up the line for the current product, price_unit, and quantity
            # Note that the price_unit used must consider discount
            base = res["total_excluded"]
            digits = 6
            avatax_amount = None
            for line in avatax_invoice.invoice_line_ids:
                price_unit = line.currency_id._convert(
                    price_unit,
                    avatax_invoice.company_id.currency_id,
                    avatax_invoice.company_id,
                    avatax_invoice.date,
                )
                if (
                    line.product_id == product
                    and float_compare(line.quantity, quantity, digits) == 0
                ):
                    avatax_amount = copysign(line.avatax_amt_line, base)
                    break
            if avatax_amount is None:
                avatax_amount = 0.0
                raise exceptions.UserError(
                    _(
                        "Incorrect retrieval of Avatax amount for Invoice %s:"
                        " product %s, price_unit %f, quantity %f"
                    )
                    % (avatax_invoice, product.display_name, -price_unit, quantity)
                )
            for tax_item in res["taxes"]:
                if tax_item["amount"] != 0:
                    tax_item["amount"] = avatax_amount
            res["total_included"] = base + avatax_amount
        return res
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-

import time
import logging
import contextlib

LOGGER = logging.getLogger("pibooth")

@contextlib.contextmanager
def timeit(description):
    LOGGER.info(description)
    start = time.time()
    try:
        yield
    finally:
        LOGGER.debug("    -> took %0.3f seconds.", time.time() - start)

class PoolingTimer(object):

    """
    Timer to be used in a pooling loop to check if timeout has been exceed.
    """

    def __init__(self, timeout, start=True):
        if timeout < 0:
            raise ValueError("PoolingTimer timeout can not be lower than zero")
        else:
            self.timeout = timeout
            self.time = None
            if start:
                self.start()

    def __enter__(self):
        """Start timer if used as context manager.
        """
        self.start()
        return self

    def __exit__(self, *args):
        """Stop timer if used as context manager.
        """
        self.time = None

    def start(self):
        """Start the timer.
        """
        self.time = time.time()

    def remaining(self):
        """Return the remaining seconds.
        """
        if self.time is None:
            remain = float(self.timeout)
        else:
            remain = self.timeout - (time.time() - self.time)
            if remain < 0.0:
                remain = 0.0
        return remain

    def elapsed(self):
        """Return the elapsed seconds.
        """
        if self.time is None:
            return 0.0
        else:
            return time.time() - self.time

    def is_timeout(self):
        """Return True if the timer is in timeout.
        """
        if self.time is None:
            raise RuntimeError("PoolingTimer has never been started")
        else:
            return (time.time() - self.time) > self.timeout
<commit_msg>Improve readability of the timeit function logs<commit_after># -*- coding: utf-8 -*-

import sys
import time
import logging
import contextlib

LOGGER = logging.getLogger("pibooth")

class BlockConsoleHandler(logging.StreamHandler):

    pattern_indent = '+< '
    pattern_blocks = '|  '
    pattern_dedent = '+> '
    current_indent = ''

    def emit(self, record):
        cls = self.__class__
        record.msg = '{}{}'.format(cls.current_indent, record.msg)
        if cls.current_indent.endswith(cls.pattern_indent):
            cls.current_indent = (cls.current_indent[:-len(cls.pattern_indent)] + cls.pattern_blocks)
        elif cls.current_indent.endswith(cls.pattern_dedent):
            cls.current_indent = cls.current_indent[:-len(cls.pattern_dedent)]
        logging.StreamHandler.emit(self, record)

    @classmethod
    def indent(cls):
        """Begin a new log block (works only if level is DEBUG).
        """
        if logging.getLogger().getEffectiveLevel() < logging.INFO:
            cls.current_indent += cls.pattern_indent

    @classmethod
    def dedent(cls):
        """End the current log block (works only if level is DEBUG).
        """
        if logging.getLogger().getEffectiveLevel() < logging.INFO:
            cls.current_indent = (cls.current_indent[:-len(cls.pattern_blocks)] + cls.pattern_dedent)

def configure_logging(level=logging.INFO, msgfmt=logging.BASIC_FORMAT, datefmt=None):
    """Configure root logger for console printing.
    """
    root = logging.getLogger()
    if not root.handlers:
        hdlr = BlockConsoleHandler(sys.stdout)
        hdlr.setFormatter(logging.Formatter(msgfmt, datefmt))
        root.addHandler(hdlr)
        if level is not None:
            root.setLevel(level)

@contextlib.contextmanager
def timeit(description):
    """Measure time execution.
    """
    BlockConsoleHandler.indent()
    LOGGER.info(description)
    start = time.time()
    try:
        yield
    finally:
        BlockConsoleHandler.dedent()
        LOGGER.debug("took %0.3f seconds", time.time() - start)

class PoolingTimer(object):

    """
    Timer to be used in a pooling loop to check if timeout has been exceed.
    """

    def __init__(self, timeout, start=True):
        if timeout < 0:
            raise ValueError("PoolingTimer timeout can not be lower than zero")
        else:
            self.timeout = timeout
            self.time = None
            if start:
                self.start()

    def __enter__(self):
        """Start timer if used as context manager.
        """
        self.start()
        return self

    def __exit__(self, *args):
        """Stop timer if used as context manager.
        """
        self.time = None

    def start(self):
        """Start the timer.
        """
        self.time = time.time()

    def remaining(self):
        """Return the remaining seconds.
        """
        if self.time is None:
            remain = float(self.timeout)
        else:
            remain = self.timeout - (time.time() - self.time)
            if remain < 0.0:
                remain = 0.0
        return remain

    def elapsed(self):
        """Return the elapsed seconds.
        """
        if self.time is None:
            return 0.0
        else:
            return time.time() - self.time

    def is_timeout(self):
        """Return True if the timer is in timeout.
        """
        if self.time is None:
            raise RuntimeError("PoolingTimer has never been started")
        else:
            return (time.time() - self.time) > self.timeout
<|endoftext|>


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

import _env
import random
import time
import traceback
import requests
from pprint import pformat
from functools import wraps
from requests.exceptions import TooManyRedirects, Timeout

from config.config import CONFIG
from extract import extract_all
from lib._db import get_db
from utils import UrlManager, IncrId
from web_util import (
    change_ip, logged, cookie_dict_from_cookie_str, CurlStrParser,
)

def retry(retries=CONFIG.CRAWLER.RETRY or 5, sleep=CONFIG.CRAWLER.SLEEP,
          changeip=False):
    """一个失败请求重试，或者使用下边这个功能强大的retrying
    pip install retrying
    https://github.com/rholder/retrying
    文章：常见的爬虫策略
    http://mp.weixin.qq.com/s?__biz=MzAwMDU1MTE1OQ==&mid=2653547274&idx=1&sn=52e5037b163146c1656eedce2da1ecd8&scene=1&srcid=0527MEXhNRZATtlTPhinD5Re#rd

    :param retries: number int of retry times.
    301 Moved Temporarily
    401 Unauthorized
    403 Forbidden
    404 Not Found
    408 Request Timeout
    429 Too Many Requests
    503 Service Unavailable
    """
    def _retry(func):
        @wraps(func)
        def _wrapper(*args, **kwargs):
            index = 0
            while index < retries:
                index += 1
                try:
                    response = func(*args, **kwargs)
                    if response and (
                        LagouCrawler.is_block_html(response.text) or
                        LagouCrawler.is_check_html(response.text)
                    ):
                        sleep_time = (sleep ** index + random.randint(1, 10))
                        if sleep_time > 300:   # 5 mins
                            change_ip()
                            continue
                        else:
                            print('sleep for %ds' % sleep_time)
                            time.sleep(sleep_time)
                            continue

                    if response.status_code in (301, 302, 404, 500):
                        print('status_code', response.status_code)
                        break
                    elif response.status_code != 200:
                        print(response.status_code)
                        if changeip:
                            change_ip()
                        continue
                    else:
                        break
                except Exception as e:
                    traceback.print_exc()
                    response = None
                    if isinstance(e, Timeout):
                        if sleep is not None:
                            time.sleep(sleep + random.randint(10, 15))
                        continue
                    elif isinstance(e, TooManyRedirects):
                        break

            return response
        return _wrapper
    return _retry

_get = requests.get    # 防止循环引用

@retry(5)
def get(*args, **kwds):
    kwds.setdefault('timeout', 10)
    headers = {
        'User-Agent': 'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_9_5) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/48.0.2564.116 Safari/537.36',
    }
    kwds.setdefault('headers', headers)
    return _get(*args, **kwds)

@logged
class LagouCrawler(object):
    curl_str = """
    curl 'http://www.lagou.com/activityapi/icon/showIcon.json?callback=jQuery11130673730597542487_1469756732278&type=POSITION&ids=2034591%2C2147192%2C1899225%2C2112714%2C1993280%2C2107221%2C1980427%2C959204%2C1570458%2C1382996%2C2164841%2C1535725%2C2015991%2C1909703%2C1924731%2C1924585%2C1917417%2C1961327%2C1949207%2C1949217%2C1961114%2C1962767%2C1915882%2C1958811%2C1929575%2C1929708%2C1926524%2C1914752&_=1469756732279' -H 'Cookie: ctk=1469756728; JSESSIONID=006FA63ABE28DD910325F0A2B21D80DD; LGMOID=20160729094529-D8AB7E5EBC00B32D65F29DC499FDEEE0; Hm_lvt_4233e74dff0ae5bd0a3d81c6ccf756e6=1469756733; Hm_lpvt_4233e74dff0ae5bd0a3d81c6ccf756e6=1469756733' -H 'X-Anit-Forge-Code: 0' -H 'Accept-Encoding: gzip, deflate, sdch' -H 'Accept-Language: zh-CN,zh;q=0.8,en-US;q=0.6,en;q=0.4' -H 'User-Agent: Mozilla/5.0 (Windows NT 6.1; WOW64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/50.0.2661.75 Safari/537.36' -H 'Accept: text/javascript, application/javascript, application/ecmascript, application/x-ecmascript, */*; q=0.01' -H 'Referer: http://www.lagou.com/' -H 'X-Requested-With: XMLHttpRequest' -H 'Connection: keep-alive' -H 'X-Anit-Forge-Token: None' --compressed
    """
    base_url = CurlStrParser(curl_str).get_url()
    headers = CurlStrParser(curl_str).get_headers_dict()
    db = get_db('htmldb')
    col = getattr(db, 'lagou_html')    # collection
    sleep = 10

    def __init__(self, domain):
        self.domain = domain
        self.url_manager = UrlManager(domain)
        self.incr_id = IncrId(self.__class__.__name__)

    def add_url(self, url):
        self.url_manager.add_url(url)

    def delay_url(self, url, nums=10):
        self.logger.info('delay url: %s', url)
        self.url_manager.delay_url(url, nums)

    def add_url_list(self):
        for i in range(1, 532):
            url = 'http://www.lagou.com/upload/sitemap/xml/lagou_sitemap_%d.xml'%i
            self.logger.info('sitemap url: %s', url)
            html = self.get_response(url).text
            all_loc_url = extract_all('<loc>', '</loc>', html)
            self.logger.info('%s', pformat(all_loc_url))
            self.add_url(all_loc_url)

    def update_headers(self, changeip=True):
        if changeip:
            change_ip()
        r = get(self.base_url)
        h = cookie_dict_from_cookie_str(r.headers.get('Set-Cookie'))
        cookies_dict = cookie_dict_from_cookie_str(self.headers['Cookie'])
        cookies_dict.update(h)
        self.headers['Cookie'] = cookies_dict
        self.logger.info('headers: %s', pformat(self.headers))

    def get_response(self, url, **kwargs):
        if CONFIG.CRAWLER.USE_PROXY:
            kwargs.setdefault('proxies', CONFIG.CRAWLER.PROXIES)
        self.logger.info('now crawler: %s', url)
        return get(url, headers=self.headers, **kwargs)

    def url_nums(self):
        return self.url_manager.url_nums()

    def next_url(self, inorder=True):
        if inorder:
            return self.url_manager.first_url()
        else:
            return self.url_manager.last_url()

    def remove_url(self, url):
        self.logger.info('remove url: %s', url)
        return self.url_manager.remove_url(url)

    @staticmethod
    def is_block_html(html):
        return 'blocked_404' in html

    @staticmethod
    def is_check_html(html):
        return '访问验证-拉勾网' in html

    def save_html(self, url, html):
        self.logger.info('save html of url: %s', url)
        if html:
            self.col.update(
                {
                    '_id': self.incr_id.get(),
                    'url': url,
                },
                {
                    '$set': {'html': html}
                },
                upsert=True
            )

    def run(self):
        if not self.url_nums():
            self.add_url_list()

        self.update_headers()

        while self.url_nums() > 0:
            if self.sleep:
                time.sleep(self.sleep + random.randint(1, 5))

            url = self.next_url()
            if url is not None:
                r = self.get_response(url)
                if not r:
                    self.delay_url(url)
                    self.update_headers()
                    continue

                html = r.text
                if self.is_block_html(html) or self.is_check_html(html):
                    print('验证码页面')
                    break
                else:
                    self.save_html(url, html)
                    self.remove_url(url)

def main():
    lagou_crawler = LagouCrawler('lagou.com')
    lagou_crawler.run()

if __name__ == '__main__':
    main()
<commit_msg>continue after block<commit_after>#!/usr/bin/env python
# -*- coding:utf-8 -*-

import _env
import random
import time
import traceback
import requests
from pprint import pformat
from functools import wraps
from requests.exceptions import TooManyRedirects, Timeout

from config.config import CONFIG
from extract import extract_all
from lib._db import get_db
from utils import UrlManager, IncrId
from web_util import (
    change_ip, logged, cookie_dict_from_cookie_str, CurlStrParser,
)

def retry(retries=CONFIG.CRAWLER.RETRY or 5, sleep=CONFIG.CRAWLER.SLEEP,
          changeip=False):
    """一个失败请求重试，或者使用下边这个功能强大的retrying
    pip install retrying
    https://github.com/rholder/retrying
    文章：常见的爬虫策略
    http://mp.weixin.qq.com/s?__biz=MzAwMDU1MTE1OQ==&mid=2653547274&idx=1&sn=52e5037b163146c1656eedce2da1ecd8&scene=1&srcid=0527MEXhNRZATtlTPhinD5Re#rd

    :param retries: number int of retry times.
    301 Moved Temporarily
    401 Unauthorized
    403 Forbidden
    404 Not Found
    408 Request Timeout
    429 Too Many Requests
    503 Service Unavailable
    """
    def _retry(func):
        @wraps(func)
        def _wrapper(*args, **kwargs):
            index = 0
            while index < retries:
                index += 1
                try:
                    response = func(*args, **kwargs)
                    if response and (
                        LagouCrawler.is_block_html(response.text) or
                        LagouCrawler.is_check_html(response.text)
                    ):
                        sleep_time = (sleep ** index + random.randint(1, 10))
                        if sleep_time > 300:   # 5 mins
                            change_ip()
                            continue
                        else:
                            print('sleep for %ds' % sleep_time)
                            time.sleep(sleep_time)
                            continue

                    if response.status_code in (301, 302, 404, 500):
                        print('status_code', response.status_code)
                        break
                    elif response.status_code != 200:
                        print(response.status_code)
                        if changeip:
                            change_ip()
                        continue
                    else:
                        break
                except Exception as e:
                    traceback.print_exc()
                    response = None
                    if isinstance(e, Timeout):
                        if sleep is not None:
                            time.sleep(sleep + random.randint(10, 15))
                        continue
                    elif isinstance(e, TooManyRedirects):
                        break

            return response
        return _wrapper
    return _retry

_get = requests.get    # 防止循环引用

@retry(5)
def get(*args, **kwds):
    kwds.setdefault('timeout', 10)
    headers = {
        'User-Agent': 'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_9_5) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/48.0.2564.116 Safari/537.36',
    }
    kwds.setdefault('headers', headers)
    return _get(*args, **kwds)

@logged
class LagouCrawler(object):
    curl_str = """
    curl 'http://www.lagou.com/activityapi/icon/showIcon.json?callback=jQuery11130673730597542487_1469756732278&type=POSITION&ids=2034591%2C2147192%2C1899225%2C2112714%2C1993280%2C2107221%2C1980427%2C959204%2C1570458%2C1382996%2C2164841%2C1535725%2C2015991%2C1909703%2C1924731%2C1924585%2C1917417%2C1961327%2C1949207%2C1949217%2C1961114%2C1962767%2C1915882%2C1958811%2C1929575%2C1929708%2C1926524%2C1914752&_=1469756732279' -H 'Cookie: ctk=1469756728; JSESSIONID=006FA63ABE28DD910325F0A2B21D80DD; LGMOID=20160729094529-D8AB7E5EBC00B32D65F29DC499FDEEE0; Hm_lvt_4233e74dff0ae5bd0a3d81c6ccf756e6=1469756733; Hm_lpvt_4233e74dff0ae5bd0a3d81c6ccf756e6=1469756733' -H 'X-Anit-Forge-Code: 0' -H 'Accept-Encoding: gzip, deflate, sdch' -H 'Accept-Language: zh-CN,zh;q=0.8,en-US;q=0.6,en;q=0.4' -H 'User-Agent: Mozilla/5.0 (Windows NT 6.1; WOW64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/50.0.2661.75 Safari/537.36' -H 'Accept: text/javascript, application/javascript, application/ecmascript, application/x-ecmascript, */*; q=0.01' -H 'Referer: http://www.lagou.com/' -H 'X-Requested-With: XMLHttpRequest' -H 'Connection: keep-alive' -H 'X-Anit-Forge-Token: None' --compressed
    """
    base_url = CurlStrParser(curl_str).get_url()
    headers = CurlStrParser(curl_str).get_headers_dict()
    db = get_db('htmldb')
    col = getattr(db, 'lagou_html')    # collection
    sleep = 10

    def __init__(self, domain):
        self.domain = domain
        self.url_manager = UrlManager(domain)
        self.incr_id = IncrId(self.__class__.__name__)

    def add_url(self, url):
        self.url_manager.add_url(url)

    def delay_url(self, url, nums=10):
        self.logger.info('delay url: %s', url)
        self.url_manager.delay_url(url, nums)

    def add_url_list(self):
        for i in range(1, 532):
            url = 'http://www.lagou.com/upload/sitemap/xml/lagou_sitemap_%d.xml'%i
            self.logger.info('sitemap url: %s', url)
            html = self.get_response(url).text
            all_loc_url = extract_all('<loc>', '</loc>', html)
            self.logger.info('%s', pformat(all_loc_url))
            self.add_url(all_loc_url)

    def update_headers(self, changeip=True):
        if changeip:
            change_ip()
        r = get(self.base_url)
        h = cookie_dict_from_cookie_str(r.headers.get('Set-Cookie'))
        cookies_dict = cookie_dict_from_cookie_str(self.headers['Cookie'])
        cookies_dict.update(h)
        self.headers['Cookie'] = cookies_dict
        self.logger.info('headers: %s', pformat(self.headers))

    def get_response(self, url, **kwargs):
        if CONFIG.CRAWLER.USE_PROXY:
            kwargs.setdefault('proxies', CONFIG.CRAWLER.PROXIES)
        self.logger.info('now crawler: %s', url)
        return get(url, headers=self.headers, **kwargs)

    def url_nums(self):
        return self.url_manager.url_nums()

    def next_url(self, inorder=True):
        if inorder:
            return self.url_manager.first_url()
        else:
            return self.url_manager.last_url()

    def remove_url(self, url):
        self.logger.info('remove url: %s', url)
        return self.url_manager.remove_url(url)

    @staticmethod
    def is_block_html(html):
        return 'blocked_404' in html

    @staticmethod
    def is_check_html(html):
        return '访问验证-拉勾网' in html

    def save_html(self, url, html):
        self.logger.info('save html of url: %s', url)
        if html:
            self.col.update(
                {
                    '_id': self.incr_id.get(),
                    'url': url,
                },
                {
                    '$set': {'html': html}
                },
                upsert=True
            )

    def run(self):
        if not self.url_nums():
            self.add_url_list()

        self.update_headers()

        while self.url_nums() > 0:
            if self.sleep:
                time.sleep(self.sleep + random.randint(1, 5))

            url = self.next_url()
            if url is not None:
                r = self.get_response(url)
                if not r:
                    self.delay_url(url)
                    self.update_headers()
                    continue

                html = r.text
                if self.is_block_html(html) or self.is_check_html(html):
                    print('验证码页面')
                    self.delay_url(url)
                    self.update_headers()
                    continue
                else:
                    self.save_html(url, html)
                    self.remove_url(url)

def main():
    lagou_crawler = LagouCrawler('lagou.com')
    lagou_crawler.run()

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


<commit_before>
import queue
import urllib.parse
import copy

# Utility functions
def cached_property(f):
    def get(self):
        try:
            return self._property_cache[f]
        except AttributeError:
            self._property_cache = {}
            x = self._property_cache[f] = f(self)
            return x
        except KeyError:
            x = self._property_cache[f] = f(self)
            return x
    return property(get)

def void_fn(*args, **kw):
    pass

def id_fn(arg):
    return arg

def constant_fn(const):
    def inner(*args, **kw):
        return const
    return inner

class PageIterator:
    def __init__(self, source):
        self.source = source
        self.queue = queue.Queue()
        self.is_source_empty = False

    def __next__(self):
        self.prefetch(1)
        try:
            return self.queue.get_nowait()
        except queue.Empty:
            raise StopIteration()

    def __iter__(self):
        return self

    def fetch(self):
        try:
            next_page = next(self.source)
            for item in next_page:
                self.queue.put(item)
        except StopIteration:
            self.is_source_empty = True

    def prefetch(self, buffer_to):
        if self.buffered_size < buffer_to and not self.is_source_empty:
            self.fetch()
            self.prefetch(buffer_to)

    @property
    def page_size(self):
        return 50

    @property
    def buffered_size(self):
        return self.queue.qsize()  # TODO: check validity.
                                   # qsize method basically returns approximate
                                   # size

def with_prefix(prefix, value):
    if value.startswith(prefix):
        return value
    else:
        return prefix + value

class URLBuilder:
    DEFAULT_PARTS = {
        'scheme': 'http',
        'host': '',
        'path': '',
        'params': {},
        'fragment': ''
    }
    _parts = None

    def __init__(self, **parts_kw):
        object.__setattr__(self, '_frozen', False)
        for k, v in parts_kw.items():
            if k not in self.DEFAULT_PARTS:
                raise ValueError('{} is not allowed'.format(k))
        parts = copy.deepcopy(self.DEFAULT_PARTS)
        parts.update(parts_kw)
        parts['path'] = with_prefix('/', parts['path'])
        parts['fragment'] = with_prefix('#', parts['fragment'])
        self._parts = parts
        self._frozen = True

    def update_with(self, **parts_kw):
        for k, v in parts_kw.items():
            if k not in self.DEFAULT_PARTS:
                raise ValueError('{} is not allowed'.format(k))
        newparts = copy.deepcopy(self._parts)
        newparts.update(parts_kw)
        return type(self)(**newparts)

    def __getattr__(self, key):
        if key not in self.DEFAULT_PARTS:
            raise ValueError('{} is not allowed'.format(key))
        return self._parts[key]

    def __hasattr__(self, key):
        return key in self.DEFAULT_PARTS

    @property
    def query(self):
        def mk_pair(kv):
            return '{}={}'.format(kv[0],
                                  urllib.parse.quote(str(kv[1]), safe=''))
        pairs = map(mk_pair, self.params.items())
        return '?' + '&'.join(pairs)

    def __str__(self):
        return '{}://{}{}{}{}'.format(self.scheme,
                                      self.host,
                                      self.path,
                                      self.query,
                                      self.fragment)

    def copy(self):
        return copy.deepcopy(self)

    def __eq__(self, other):
        return str(self) == str(other)

    def __setattr__(self, name, value):
        if not self._frozen:
            object.__setattr__(self, name, value)
        else:
            raise TypeError('Cannot set name %r on object of type %s' % (
                            name, self.__class__.__name__))

    def to_request(self):
        pass
<commit_msg>Add update_params function to utile.URLBuilder<commit_after>
import queue
import urllib.parse
import copy

# Utility functions
def cached_property(f):
    def get(self):
        try:
            return self._property_cache[f]
        except AttributeError:
            self._property_cache = {}
            x = self._property_cache[f] = f(self)
            return x
        except KeyError:
            x = self._property_cache[f] = f(self)
            return x
    return property(get)

def void_fn(*args, **kw):
    pass

def id_fn(arg):
    return arg

def constant_fn(const):
    def inner(*args, **kw):
        return const
    return inner

class PageIterator:
    def __init__(self, source):
        self.source = source
        self.queue = queue.Queue()
        self.is_source_empty = False

    def __next__(self):
        self.prefetch(1)
        try:
            return self.queue.get_nowait()
        except queue.Empty:
            raise StopIteration()

    def __iter__(self):
        return self

    def fetch(self):
        try:
            next_page = next(self.source)
            for item in next_page:
                self.queue.put(item)
        except StopIteration:
            self.is_source_empty = True

    def prefetch(self, buffer_to):
        if self.buffered_size < buffer_to and not self.is_source_empty:
            self.fetch()
            self.prefetch(buffer_to)

    @property
    def page_size(self):
        return 50

    @property
    def buffered_size(self):
        return self.queue.qsize()  # TODO: check validity.
                                   # qsize method basically returns approximate
                                   # size

def with_prefix(prefix, value):
    if value.startswith(prefix):
        return value
    else:
        return prefix + value

class URLBuilder:
    DEFAULT_PARTS = {
        'scheme': 'http',
        'host': '',
        'path': '',
        'params': {},
        'fragment': ''
    }
    _parts = None

    def __init__(self, **parts_kw):
        object.__setattr__(self, '_frozen', False)
        for k, v in parts_kw.items():
            if k not in self.DEFAULT_PARTS:
                raise ValueError('{} is not allowed'.format(k))
        parts = copy.deepcopy(self.DEFAULT_PARTS)
        parts.update(parts_kw)
        parts['path'] = with_prefix('/', parts['path'])
        parts['fragment'] = with_prefix('#', parts['fragment'])
        self._parts = parts
        self._frozen = True

    def update_params(self, **kw):
        newparams = copy.deepcopy(self.params)
        newparams.update(kw)
        return self.update_with(params=newparams)

    def update_with(self, **parts_kw):
        for k, v in parts_kw.items():
            if k not in self.DEFAULT_PARTS:
                raise ValueError('{} is not allowed'.format(k))
        newparts = copy.deepcopy(self._parts)
        newparts.update(parts_kw)
        return type(self)(**newparts)

    def __getattr__(self, key):
        if key not in self.DEFAULT_PARTS:
            raise ValueError('{} is not allowed'.format(key))
        return self._parts[key]

    def __hasattr__(self, key):
        return key in self.DEFAULT_PARTS

    @property
    def query(self):
        def mk_pair(kv):
            return '{}={}'.format(kv[0],
                                  urllib.parse.quote(str(kv[1]), safe=''))
        pairs = map(mk_pair, self.params.items())
        return '?' + '&'.join(pairs)

    def __str__(self):
        return '{}://{}{}{}{}'.format(self.scheme,
                                      self.host,
                                      self.path,
                                      self.query,
                                      self.fragment)

    def copy(self):
        return copy.deepcopy(self)

    def __eq__(self, other):
        return str(self) == str(other)

    def __setattr__(self, name, value):
        if not self._frozen:
            object.__setattr__(self, name, value)
        else:
            raise TypeError('Cannot set name %r on object of type %s' % (
                            name, self.__class__.__name__))

    def to_request(self):
        pass
<|endoftext|>


<commit_before># ActivitySim
# See full license in LICENSE.txt.
import logging
import warnings
import os
import pandas as pd

from activitysim.core import tracing
from activitysim.core import config
from activitysim.core import inject
from activitysim.core import pipeline
from activitysim.core import expressions
from activitysim.core import chunk

from activitysim.core.steps.output import write_data_dictionary
from activitysim.core.steps.output import write_tables
from activitysim.core.steps.output import track_skim_usage

from activitysim.abm.tables import shadow_pricing

# We are using the naming conventions in the mtc_asim.h5 example
# file for our default list. This provides backwards compatibility
# with previous versions of ActivitySim in which only 'input_store'
# is given in the settings file.
DEFAULT_TABLE_LIST = [
    {'tablename': 'households',
     'h5_tablename': 'households',
     'index_col': 'household_id'},
    {'tablename': 'persons',
     'h5_tablename': 'persons',
     'index_col': 'person_id'},
    {'tablename': 'land_use',
     'h5_tablename': 'land_use_taz',
     'index_col': 'TAZ'}
]

logger = logging.getLogger(__name__)

def annotate_tables(model_settings, trace_label):

    trace_label = tracing.extend_trace_label(trace_label, 'annotate_tables')

    chunk.log_rss(trace_label)

    annotate_tables = model_settings.get('annotate_tables', [])

    if not annotate_tables:
        logger.warning(f"{trace_label} - annotate_tables setting is empty - nothing to do!")

    assert isinstance(annotate_tables, list), \
        f"annotate_tables settings should be a list but is {type(annotate_tables)}"

    t0 = tracing.print_elapsed_time()

    for table_info in annotate_tables:

        tablename = table_info['tablename']

        chunk.log_rss(f"{trace_label}.pre-get_table.{tablename}")

        df = inject.get_table(tablename).to_frame()
        chunk.log_df(trace_label, tablename, df)

        # - rename columns
        column_map = table_info.get('column_map', None)
        if column_map:

            warnings.warn(f"Setting 'column_map' has been changed to 'rename_columns'. "
                          f"Support for 'column_map' in annotate_tables  will be removed in future versions.",
                          FutureWarning)

            logger.info(f"{trace_label} - renaming {tablename} columns {column_map}")
            df.rename(columns=column_map, inplace=True)

        # - annotate
        annotate = table_info.get('annotate', None)
        if annotate:
            logger.info(f"{trace_label} - annotating {tablename} SPEC {annotate['SPEC']}")
            expressions.assign_columns(
                df=df,
                model_settings=annotate,
                trace_label=trace_label)

        chunk.log_df(trace_label, tablename, df)

        # - write table to pipeline
        pipeline.replace_table(tablename, df)

        del df
        chunk.log_df(trace_label, tablename, None)

@inject.step()
def initialize_landuse():

    trace_label = 'initialize_landuse'

    with chunk.chunk_log(trace_label, base=True):

        model_settings = config.read_model_settings('initialize_landuse.yaml', mandatory=True)

        annotate_tables(model_settings, trace_label)

        # instantiate accessibility (must be checkpointed to be be used to slice accessibility)
        accessibility = pipeline.get_table('accessibility')
        chunk.log_df(trace_label, "accessibility", accessibility)

@inject.step()
def initialize_households():

    trace_label = 'initialize_households'

    with chunk.chunk_log(trace_label, base=True):

        chunk.log_rss(f"{trace_label}.inside-yield")

        households = inject.get_table('households').to_frame()
        assert not households._is_view
        chunk.log_df(trace_label, "households", households)
        del households
        chunk.log_df(trace_label, "households", None)

        persons = inject.get_table('persons').to_frame()
        assert not persons._is_view
        chunk.log_df(trace_label, "persons", persons)
        del persons
        chunk.log_df(trace_label, "persons", None)

        model_settings = config.read_model_settings('initialize_households.yaml', mandatory=True)
        annotate_tables(model_settings, trace_label)

        # - initialize shadow_pricing size tables after annotating household and person tables
        # since these are scaled to model size, they have to be created while single-process
        # this can now be called as a stand alone model step instead, add_size_tables
        add_size_tables = model_settings.get('add_size_tables', True)
        if add_size_tables:
            warnings.warn(f"Calling add_size_tables from initialize will be removed in future versions.", FutureWarning)
            shadow_pricing.add_size_tables()
            mem.trace_memory_info(f"initialize_households after shadow_pricing.add_size_tables")

        # - preload person_windows
        person_windows = inject.get_table('person_windows').to_frame()
        chunk.log_df(trace_label, "person_windows", person_windows)

@inject.injectable(cache=True)
def preload_injectables():
    """
    preload bulky injectables up front - stuff that isn't inserted into the pipeline
    """

    logger.info("preload_injectables")

    inject.add_step('track_skim_usage', track_skim_usage)
    inject.add_step('write_data_dictionary', write_data_dictionary)
    inject.add_step('write_tables', write_tables)

    table_list = config.setting('input_table_list')

    # default ActivitySim table names and indices
    if table_list is None:
        logger.warning(
            "No 'input_table_list' found in settings. This will be a "
            "required setting in upcoming versions of ActivitySim.")

        new_settings = inject.get_injectable('settings')
        new_settings['input_table_list'] = DEFAULT_TABLE_LIST
        inject.add_injectable('settings', new_settings)

    # FIXME undocumented feature
    if config.setting('write_raw_tables'):

        # write raw input tables as csv (before annotation)
        csv_dir = config.output_file_path('raw_tables')
        if not os.path.exists(csv_dir):
            os.makedirs(csv_dir)  # make directory if needed

        table_names = [t['tablename'] for t in table_list]
        for t in table_names:
            df = inject.get_table(t).to_frame()
            df.to_csv(os.path.join(csv_dir, '%s.csv' % t), index=True)

    t0 = tracing.print_elapsed_time()

    # FIXME - still want to do this?
    # if inject.get_injectable('skim_dict', None) is not None:
    #     t0 = tracing.print_elapsed_time("preload skim_dict", t0, debug=True)
    #
    # if inject.get_injectable('skim_stack', None) is not None:
    #     t0 = tracing.print_elapsed_time("preload skim_stack", t0, debug=True)

    return True
<commit_msg>comment out future warning for now<commit_after># ActivitySim
# See full license in LICENSE.txt.
import logging
import warnings
import os
import pandas as pd

from activitysim.core import tracing
from activitysim.core import config
from activitysim.core import inject
from activitysim.core import pipeline
from activitysim.core import expressions
from activitysim.core import chunk

from activitysim.core.steps.output import write_data_dictionary
from activitysim.core.steps.output import write_tables
from activitysim.core.steps.output import track_skim_usage

from activitysim.abm.tables import shadow_pricing

# We are using the naming conventions in the mtc_asim.h5 example
# file for our default list. This provides backwards compatibility
# with previous versions of ActivitySim in which only 'input_store'
# is given in the settings file.
DEFAULT_TABLE_LIST = [
    {'tablename': 'households',
     'h5_tablename': 'households',
     'index_col': 'household_id'},
    {'tablename': 'persons',
     'h5_tablename': 'persons',
     'index_col': 'person_id'},
    {'tablename': 'land_use',
     'h5_tablename': 'land_use_taz',
     'index_col': 'TAZ'}
]

logger = logging.getLogger(__name__)

def annotate_tables(model_settings, trace_label):

    trace_label = tracing.extend_trace_label(trace_label, 'annotate_tables')

    chunk.log_rss(trace_label)

    annotate_tables = model_settings.get('annotate_tables', [])

    if not annotate_tables:
        logger.warning(f"{trace_label} - annotate_tables setting is empty - nothing to do!")

    assert isinstance(annotate_tables, list), \
        f"annotate_tables settings should be a list but is {type(annotate_tables)}"

    t0 = tracing.print_elapsed_time()

    for table_info in annotate_tables:

        tablename = table_info['tablename']

        chunk.log_rss(f"{trace_label}.pre-get_table.{tablename}")

        df = inject.get_table(tablename).to_frame()
        chunk.log_df(trace_label, tablename, df)

        # - rename columns
        column_map = table_info.get('column_map', None)
        if column_map:

            warnings.warn(f"Setting 'column_map' has been changed to 'rename_columns'. "
                          f"Support for 'column_map' in annotate_tables  will be removed in future versions.",
                          FutureWarning)

            logger.info(f"{trace_label} - renaming {tablename} columns {column_map}")
            df.rename(columns=column_map, inplace=True)

        # - annotate
        annotate = table_info.get('annotate', None)
        if annotate:
            logger.info(f"{trace_label} - annotating {tablename} SPEC {annotate['SPEC']}")
            expressions.assign_columns(
                df=df,
                model_settings=annotate,
                trace_label=trace_label)

        chunk.log_df(trace_label, tablename, df)

        # - write table to pipeline
        pipeline.replace_table(tablename, df)

        del df
        chunk.log_df(trace_label, tablename, None)

@inject.step()
def initialize_landuse():

    trace_label = 'initialize_landuse'

    with chunk.chunk_log(trace_label, base=True):

        model_settings = config.read_model_settings('initialize_landuse.yaml', mandatory=True)

        annotate_tables(model_settings, trace_label)

        # instantiate accessibility (must be checkpointed to be be used to slice accessibility)
        accessibility = pipeline.get_table('accessibility')
        chunk.log_df(trace_label, "accessibility", accessibility)

@inject.step()
def initialize_households():

    trace_label = 'initialize_households'

    with chunk.chunk_log(trace_label, base=True):

        chunk.log_rss(f"{trace_label}.inside-yield")

        households = inject.get_table('households').to_frame()
        assert not households._is_view
        chunk.log_df(trace_label, "households", households)
        del households
        chunk.log_df(trace_label, "households", None)

        persons = inject.get_table('persons').to_frame()
        assert not persons._is_view
        chunk.log_df(trace_label, "persons", persons)
        del persons
        chunk.log_df(trace_label, "persons", None)

        model_settings = config.read_model_settings('initialize_households.yaml', mandatory=True)
        annotate_tables(model_settings, trace_label)

        # - initialize shadow_pricing size tables after annotating household and person tables
        # since these are scaled to model size, they have to be created while single-process
        # this can now be called as a stand alone model step instead, add_size_tables
        add_size_tables = model_settings.get('add_size_tables', True)
        if add_size_tables:
            # warnings.warn(f"Calling add_size_tables from initialize will be removed in future versions.", FutureWarning)
            shadow_pricing.add_size_tables()

        # - preload person_windows
        person_windows = inject.get_table('person_windows').to_frame()
        chunk.log_df(trace_label, "person_windows", person_windows)

@inject.injectable(cache=True)
def preload_injectables():
    """
    preload bulky injectables up front - stuff that isn't inserted into the pipeline
    """

    logger.info("preload_injectables")

    inject.add_step('track_skim_usage', track_skim_usage)
    inject.add_step('write_data_dictionary', write_data_dictionary)
    inject.add_step('write_tables', write_tables)

    table_list = config.setting('input_table_list')

    # default ActivitySim table names and indices
    if table_list is None:
        logger.warning(
            "No 'input_table_list' found in settings. This will be a "
            "required setting in upcoming versions of ActivitySim.")

        new_settings = inject.get_injectable('settings')
        new_settings['input_table_list'] = DEFAULT_TABLE_LIST
        inject.add_injectable('settings', new_settings)

    # FIXME undocumented feature
    if config.setting('write_raw_tables'):

        # write raw input tables as csv (before annotation)
        csv_dir = config.output_file_path('raw_tables')
        if not os.path.exists(csv_dir):
            os.makedirs(csv_dir)  # make directory if needed

        table_names = [t['tablename'] for t in table_list]
        for t in table_names:
            df = inject.get_table(t).to_frame()
            df.to_csv(os.path.join(csv_dir, '%s.csv' % t), index=True)

    t0 = tracing.print_elapsed_time()

    # FIXME - still want to do this?
    # if inject.get_injectable('skim_dict', None) is not None:
    #     t0 = tracing.print_elapsed_time("preload skim_dict", t0, debug=True)
    #
    # if inject.get_injectable('skim_stack', None) is not None:
    #     t0 = tracing.print_elapsed_time("preload skim_stack", t0, debug=True)

    return True
<|endoftext|>


<commit_before>#-*-coding:utf-8-*-

import curses

__all__ = ['Picker', 'pick']

KEYS_ENTER = (curses.KEY_ENTER, ord('\n'), ord('\r'))
KEYS_UP = (curses.KEY_UP, ord('k'))
KEYS_DOWN = (curses.KEY_DOWN, ord('j'))

class Picker(object):
    """The :class:`Picker <Picker>` object

    :param options: a list of options to choose from
    :param title: (optional) a title above options list
    :param indicator: (optional) custom the selection indicator
    :param default_index: (optional) set this if the default selected option is not the first one
    """

    def __init__(self, options, title=None, indicator='*', default_index=0):

        if len(options) == 0:
            raise ValueError('options should not be an empty list')

        self.options = options
        self.title = title
        self.indicator = indicator

        if default_index >= len(options):
            raise ValueError('default_index should be less than the length of options')

        self.index = default_index
        self.custom_handlers = {}

    def register_custom_handler(self, key, func):
        self.custom_handlers[key] = func

    def move_up(self):
        self.index -= 1
        if self.index < 0:
            self.index = len(self.options) - 1

    def move_down(self):
        self.index += 1
        if self.index >= len(self.options):
            self.index = 0

    def get_selected(self):
        """return the current selected option as a tuple: (option, index)
        """
        return self.options[self.index], self.index

    def get_title_lines(self):
        if self.title:
            return self.title.split('\n') + ['']
        return []

    def get_option_lines(self):
        lines = []

        for index, option in enumerate(self.options):
            if index == self.index:
                prefix = self.indicator
            else:
                prefix = len(self.indicator) * ' '
            line = '{0} {1}'.format(prefix, option)
            lines.append(line)

        return lines

    def get_lines(self):
        title_lines = self.get_title_lines()
        option_lines = self.get_option_lines()
        lines = title_lines + option_lines
        current_line = self.index + len(title_lines) + 1
        return lines, current_line

    def draw(self):
        """draw the curses ui on the screen, handle scroll if needed"""
        self.screen.clear()

        x, y = 1, 1  # start point
        max_y, max_x = self.screen.getmaxyx()
        max_rows = max_y - y  # the max rows we can draw

        lines, current_line = self.get_lines()

        # calculate how many lines we should scroll, relative to the top
        scroll_top = getattr(self, 'scroll_top', 0)
        if current_line <= scroll_top:
            scroll_top = 0
        elif current_line - scroll_top > max_rows:
            scroll_top = current_line - max_rows
        self.scroll_top = scroll_top

        lines_to_draw = lines[scroll_top:scroll_top+max_rows]

        for line in lines_to_draw:
            self.screen.addstr(y, x, line)
            y += 1

        self.screen.refresh()

    def run_loop(self):
        while True:
            self.draw()
            c = self.screen.getch()
            if c in KEYS_UP:
                self.move_up()
            elif c in KEYS_DOWN:
                self.move_down()
            elif c in KEYS_ENTER:
                return self.get_selected()
            elif c in self.custom_handlers:
                ret = self.custom_handlers[c](self)
                if ret:
                    return ret

    def config_curses(self):
        # use the default colors of the terminal
        curses.use_default_colors()
        # hide the cursor
        curses.curs_set(0)

    def _start(self, screen):
        self.screen = screen
        self.config_curses()
        return self.run_loop()

    def start(self):
        return curses.wrapper(self._start)

def pick(options, title=None, indicator='*', default_index=0):
    """Construct and start a :class:`Picker <Picker>`.

    Usage::

      >>> from pick import pick
      >>> title = 'Please choose an option: '
      >>> options = ['option1', 'option2', 'option3']
      >>> option, index = pick(options, title)
    """
    picker = Picker(options, title, indicator, default_index)
    return picker.start()
<commit_msg>truncate lines if too long<commit_after>#-*-coding:utf-8-*-

import curses

__all__ = ['Picker', 'pick']

KEYS_ENTER = (curses.KEY_ENTER, ord('\n'), ord('\r'))
KEYS_UP = (curses.KEY_UP, ord('k'))
KEYS_DOWN = (curses.KEY_DOWN, ord('j'))

class Picker(object):
    """The :class:`Picker <Picker>` object

    :param options: a list of options to choose from
    :param title: (optional) a title above options list
    :param indicator: (optional) custom the selection indicator
    :param default_index: (optional) set this if the default selected option is not the first one
    """

    def __init__(self, options, title=None, indicator='*', default_index=0):

        if len(options) == 0:
            raise ValueError('options should not be an empty list')

        self.options = options
        self.title = title
        self.indicator = indicator

        if default_index >= len(options):
            raise ValueError('default_index should be less than the length of options')

        self.index = default_index
        self.custom_handlers = {}

    def register_custom_handler(self, key, func):
        self.custom_handlers[key] = func

    def move_up(self):
        self.index -= 1
        if self.index < 0:
            self.index = len(self.options) - 1

    def move_down(self):
        self.index += 1
        if self.index >= len(self.options):
            self.index = 0

    def get_selected(self):
        """return the current selected option as a tuple: (option, index)
        """
        return self.options[self.index], self.index

    def get_title_lines(self):
        if self.title:
            return self.title.split('\n') + ['']
        return []

    def get_option_lines(self):
        lines = []

        for index, option in enumerate(self.options):
            if index == self.index:
                prefix = self.indicator
            else:
                prefix = len(self.indicator) * ' '
            line = '{0} {1}'.format(prefix, option)
            lines.append(line)

        return lines

    def get_lines(self):
        title_lines = self.get_title_lines()
        option_lines = self.get_option_lines()
        lines = title_lines + option_lines
        current_line = self.index + len(title_lines) + 1
        return lines, current_line

    def draw(self):
        """draw the curses ui on the screen, handle scroll if needed"""
        self.screen.clear()

        x, y = 1, 1  # start point
        max_y, max_x = self.screen.getmaxyx()
        max_rows = max_y - y  # the max rows we can draw

        lines, current_line = self.get_lines()

        # calculate how many lines we should scroll, relative to the top
        scroll_top = getattr(self, 'scroll_top', 0)
        if current_line <= scroll_top:
            scroll_top = 0
        elif current_line - scroll_top > max_rows:
            scroll_top = current_line - max_rows
        self.scroll_top = scroll_top

        lines_to_draw = lines[scroll_top:scroll_top+max_rows]

        for line in lines_to_draw:
            self.screen.addnstr(y, x, line, max_x-2)
            y += 1

        self.screen.refresh()

    def run_loop(self):
        while True:
            self.draw()
            c = self.screen.getch()
            if c in KEYS_UP:
                self.move_up()
            elif c in KEYS_DOWN:
                self.move_down()
            elif c in KEYS_ENTER:
                return self.get_selected()
            elif c in self.custom_handlers:
                ret = self.custom_handlers[c](self)
                if ret:
                    return ret

    def config_curses(self):
        # use the default colors of the terminal
        curses.use_default_colors()
        # hide the cursor
        curses.curs_set(0)

    def _start(self, screen):
        self.screen = screen
        self.config_curses()
        return self.run_loop()

    def start(self):
        return curses.wrapper(self._start)

def pick(options, title=None, indicator='*', default_index=0):
    """Construct and start a :class:`Picker <Picker>`.

    Usage::

      >>> from pick import pick
      >>> title = 'Please choose an option: '
      >>> options = ['option1', 'option2', 'option3']
      >>> option, index = pick(options, title)
    """
    picker = Picker(options, title, indicator, default_index)
    return picker.start()
<|endoftext|>


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

# I'm tired of having "du -hs / | sort -h" choke on /proc, /home
# etc.

import os
import sys
import subprocess
import re

MY_NAME = "pickle_sickle"
INSPECTEE = "/" # an absolute path to inspect WITH trailing slash
EXCLUDE_US = ("proc", "home", "run", "lost+found")
EXCLUDED = tuple(INSPECTEE + ent for ent in EXCLUDE_US)
MEGABYTES = re.compile(r"^b'([0-9]+).*$")

def do_du(dir):
  """calls "du -ms" on dir if appropriate. returns the size
     in Mb, else None. dir is an absolute path!"""
  dnll = open("/dev/null", "w")
  if dir not in EXCLUDED:
    try:
      du = str(subprocess.check_output(("du", "-ms", dir),
      stderr = dnll))
    except subprocess.CalledProcessError as err:
      annoy("du didn't work on " + dir + ": " + str(err.output))
      du = str(err.output)
    search = MEGABYTES.search(du)
    if not search:
      annoy("du didn't work on " + dir)
    else:
      return int(search.group(1))
  return None

def print_sizes(list_of_sizes):
  """prints a human-readable list of dirs and sizes. expects
     a list of tuples (dir, size_in_mb), and prints these out."""
  print_magic = 1
  for tup in list_of_sizes:
    if len(tup[0]) > print_magic:
      print_magic = len(tup[0])
  print_fmt = "%" + str(print_magic) + "s%26s"
  for tup in list_of_sizes:
    name = tup[0]
    size = tup[1]
    if size > 1024:
      size /= 1024
      size = str("%s GB" % size)
    else:
      size = str("%s MB" % size)
    printable = str(print_fmt % (name, size))
    print(printable)
  return

def annoy(msg):
  """generic shortened form for printing errors."""
  sys.stderr.write(MY_NAME + ": ")
  sys.stderr.write(msg)
  sys.stderr.write("\n")
  sys.stderr.flush()
  return

def do_work(pr=False):
  """main entry point. returns a list of tuples, each one being
     (dir, size_in_mb)."""
  retv = []
  for ent in os.listdir(INSPECTEE):
    dirent_str = INSPECTEE + ent
    dirent_siz = do_du(dirent_str)
    if dirent_siz:
      retv.append((dirent_str, do_du(dirent_str)))
  retv = sorted(retv, key=(lambda x: x[1]))
  if pr:
    print_sizes(retv)
  return retv

if __name__ == "__main__":
  do_work(pr=True)
<commit_msg>do things in KB. more general. still extensible.<commit_after>#! /usr/bin/env python3

# I'm tired of having "du -hs / | sort -h" choke on /proc, /home
# etc.

import os
import sys
import subprocess
import re

MY_NAME = "pickle_sickle"
INSPECTEE = "/" # an absolute path to inspect WITH trailing slash
EXCLUDE_US = ("proc", "home", "run", "lost+found")
EXCLUDED = tuple(INSPECTEE + ent for ent in EXCLUDE_US)
MEGABYTES = re.compile(r"^b'([0-9]+).*$")

DEBUGGING = False

def do_du(dir):
  """calls "du -ks" on dir if appropriate. returns the size
     in Kb, else None. dir is an absolute path!"""
  dnll = open("/dev/null", "w")
  if dir not in EXCLUDED:
    try:
      du = str(subprocess.check_output(("du", "-ks", dir),
      stderr = dnll))
    except subprocess.CalledProcessError as err:
      annoy("du didn't work on " + dir + ": " + str(err.output))
      du = str(err.output)
    search = MEGABYTES.search(du)
    if not search:
      annoy("du didn't work on " + dir)
    else:
      return int(search.group(1))
  return None

def print_sizes(list_of_sizes):
  """prints a human-readable list of dirs and sizes. expects
     a list of tuples (dir, size_in_kb), and prints these out."""
  print_magic = 1
  for tup in list_of_sizes:
    if len(tup[0]) > print_magic:
      print_magic = len(tup[0])
  print_fmt = "%" + str(print_magic) + "s%26s"
  for tup in list_of_sizes:
    name = tup[0]
    size = tup[1]
    if size > 1024*1024:
      size /= (1024*1024)
      size = str("%.2f GB" % size)
    elif size > 1024:
      size /= 1024
      size = str("%.2f MB" % size)
    else:
      size = str("%.2f KB" % size)
    printable = str(print_fmt % (name, size))
    print(printable)
  return

def annoy(msg):
  """generic shortened form for printing errors."""
  if not DEBUGGING:
    return
  sys.stderr.write(MY_NAME + ": ")
  sys.stderr.write(msg)
  sys.stderr.write("\n")
  sys.stderr.flush()
  return

def check_args():
  """ if given an argument, take it! """
  argc = len(sys.argv) - 1
  if argc:
    if os.path.exists(sys.argv[1]):
      global INSPECTEE
      INSPECTEE = str("%s/" % sys.argv[1])
    else:
      annoy(str("%s doesn't exist!" % sys.argv[1]))
  return

def do_work(pr=False):
  """main entry point. returns a list of tuples, each one being
     (dir, size_in_mb)."""
  check_args()
  retv = []
  for ent in os.listdir(INSPECTEE):
    dirent_str = INSPECTEE + ent
    dirent_siz = do_du(dirent_str)
    if dirent_siz:
      retv.append((dirent_str, do_du(dirent_str)))
  retv = sorted(retv, key=(lambda x: x[1]))
  if pr:
    print_sizes(retv)
  return retv

if __name__ == "__main__":
  do_work(pr=True)
<|endoftext|>


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

import os
import json
import hashlib
import time
import glob

import click
import requests
import redis

_base = 'http://localhost:3000'
_here = os.path.dirname(__file__)

def get_blogposts():
    items = json.load(open(os.path.join(_here, 'blogposts.json')))['items']
    for i, item in enumerate(items):
        url = 'http://www.peterbe.com/plog/%s' % item['slug']
        yield (item['title'], url, len(items) - i, None)

def get_events():
    items = json.load(open(os.path.join(_here, 'airmoevents.json')))['items']
    for i, item in enumerate(items):
        group = item['group']
        group = group != 'public' and group or ''
        yield (item['title'], item['url'], item['popularity'], group)

def get_items(jsonfile):
    items = json.load(open(jsonfile))['items']
    for i, item in enumerate(items):
        if 'group' in item:
            group = item['group']
            group = group != 'public' and group or ''
        else:
            group = None
        yield (item['title'], item['url'], item.get('popularity'), group)

def populate(database, destination, domain, jsonfile, flush=False, bulk=False):
    c = redis.StrictRedis(host='localhost', port=6379, db=database)
    if flush:
        c.flushdb()
    key = hashlib.md5(open(__file__).read()).hexdigest()

    print "KEY", key
    print "DOMAIN", domain
    c.hset('$domainkeys', key, domain)
    c.sadd('$userdomains$peterbe', key)

    items = get_items(jsonfile)
    # items = get_blogposts()
    #items = get_events()
    t0 = time.time()
    if bulk:
        _in_bulk(destination, items, key)
    else:
        _one_at_a_time(destination, items, key)
    t1 = time.time()
    print "TOOK", t1 - t0

def _one_at_a_time(destination, items, key):
    for title, url, popularity, group in items:
        _url = destination + '/v1'
        data = {
            'title': title,
            'url': url,
            'popularity': popularity,
            'group': group,
        }
        #print (url, title, popularity, group)
        r = requests.post(
            _url,
            data=data,
            headers={'Auth-Key': key}
        )
        #print r.status_code
        assert r.status_code == 201, r.status_code

def _in_bulk(destination, items, key):
    data = {
        'documents': [
            dict(
                title=t,
                url=u,
                popularity=p,
                group=g
            )
            for t, u, p, g in items
        ]
    }
    _url = destination + '/v1/bulk'
    r = requests.post(
        _url,
        data=json.dumps(data),
        headers={'Auth-Key': key}
    )
    assert r.status_code == 201, r.status_code

_json_files = glob.glob(os.path.join(_here, '*.json'))

@click.command()
@click.option('--database', '-d', default=8)
@click.option('--destination', default='https://autocompeter.com')
@click.option('--domain', default='autocompeter.com')
@click.option('--dataset', default='blogposts.json')
@click.option('--flush', default=False, is_flag=True)
@click.option('--bulk', default=True, is_flag=True)
def run(database, destination, domain, dataset, flush=False, bulk=True):
    # print (database, domain, flush)
    for filename in _json_files:
        if os.path.basename(filename) == dataset:
            jsonfile = filename
            break
    else:
        raise ValueError("dataset %r not recognized" % dataset)
    populate(
        database,
        destination,
        domain,
        jsonfile,
        flush=flush,
        bulk=bulk,
    )

if __name__ == '__main__':
    run()
<commit_msg>always bulk load with the populate tool<commit_after>#!/usr/bin/env python

import os
import json
import hashlib
import time
import glob

import click
import requests
import redis

_base = 'http://localhost:3000'
_here = os.path.dirname(__file__)

def get_blogposts():
    items = json.load(open(os.path.join(_here, 'blogposts.json')))['items']
    for i, item in enumerate(items):
        url = 'http://www.peterbe.com/plog/%s' % item['slug']
        yield (item['title'], url, len(items) - i, None)

def get_events():
    items = json.load(open(os.path.join(_here, 'airmoevents.json')))['items']
    for i, item in enumerate(items):
        group = item['group']
        group = group != 'public' and group or ''
        yield (item['title'], item['url'], item['popularity'], group)

def get_items(jsonfile):
    items = json.load(open(jsonfile))['items']
    for i, item in enumerate(items):
        if 'group' in item:
            group = item['group']
            group = group != 'public' and group or ''
        else:
            group = None
        yield (item['title'], item['url'], item.get('popularity'), group)

def populate(database, destination, domain, jsonfile, flush=False, bulk=False):
    c = redis.StrictRedis(host='localhost', port=6379, db=database)
    if flush:
        c.flushdb()
    key = hashlib.md5(open(__file__).read()).hexdigest()

    print "KEY", key
    print "DOMAIN", domain
    c.hset('$domainkeys', key, domain)
    c.sadd('$userdomains$peterbe', key)

    items = get_items(jsonfile)
    # items = get_blogposts()
    #items = get_events()
    t0 = time.time()
    if bulk:
        _in_bulk(destination, items, key)
    else:
        _one_at_a_time(destination, items, key)
    t1 = time.time()
    print "TOOK", t1 - t0

def _one_at_a_time(destination, items, key):
    for title, url, popularity, group in items:
        _url = destination + '/v1'
        data = {
            'title': title,
            'url': url,
            'popularity': popularity,
            'group': group,
        }
        #print (url, title, popularity, group)
        r = requests.post(
            _url,
            data=data,
            headers={'Auth-Key': key}
        )
        #print r.status_code
        assert r.status_code == 201, r.status_code

def _in_bulk(destination, items, key):
    data = {
        'documents': [
            dict(
                title=t,
                url=u,
                popularity=p,
                group=g
            )
            for t, u, p, g in items
        ]
    }
    _url = destination + '/v1/bulk'
    r = requests.post(
        _url,
        data=json.dumps(data),
        headers={'Auth-Key': key}
    )
    assert r.status_code == 201, r.status_code

_json_files = glob.glob(os.path.join(_here, '*.json'))

@click.command()
@click.option('--database', '-d', default=8)
@click.option('--destination', default='https://autocompeter.com')
@click.option('--domain', default='autocompeter.com')
@click.option('--dataset', default='blogposts.json')
@click.option('--flush', default=False, is_flag=True)
@click.option('--no-bulk', default=False, is_flag=True)
def run(database, destination, domain, dataset, flush=False, no_bulk=True):
    # print (database, domain, flush)
    bulk = not no_bulk
    for filename in _json_files:
        if os.path.basename(filename) == dataset:
            jsonfile = filename
            break
    else:
        raise ValueError("dataset %r not recognized" % dataset)
    populate(
        database,
        destination,
        domain,
        jsonfile,
        flush=flush,
        bulk=bulk,
    )

if __name__ == '__main__':
    run()
<|endoftext|>


<commit_before>"""Locations where we look for configs, install stuff, etc"""

import sys
import os
import tempfile
from pip.backwardcompat import get_python_lib

def running_under_virtualenv():
    """
    Return True if we're running inside a virtualenv, False otherwise.

    """
    return hasattr(sys, 'real_prefix')

if running_under_virtualenv():
    ## FIXME: is build/ a good name?
    build_prefix = os.path.join(sys.prefix, 'build')
    src_prefix = os.path.join(sys.prefix, 'src')
else:
    #Use tempfile to create a temporary folder
    build_prefix = tempfile.mkdtemp('-build', 'pip-')
    src_prefix = tempfile.mkdtemp('-src', 'pip-')
    #This is a terrible hack- since pip relies on this directory not being created yet
    #We will delete it now, and have pip recreate it later
    os.rmdir(build_prefix)
    os.rmdir(src_prefix)

# under Mac OS X + virtualenv sys.prefix is not properly resolved
# it is something like /path/to/python/bin/..
build_prefix = os.path.abspath(build_prefix)
src_prefix = os.path.abspath(src_prefix)

# FIXME doesn't account for venv linked to global site-packages

site_packages = get_python_lib()
user_dir = os.path.expanduser('~')
if sys.platform == 'win32':
    bin_py = os.path.join(sys.prefix, 'Scripts')
    # buildout uses 'bin' on Windows too?
    if not os.path.exists(bin_py):
        bin_py = os.path.join(sys.prefix, 'bin')
    user_dir = os.environ.get('APPDATA', user_dir) # Use %APPDATA% for roaming
    default_storage_dir = os.path.join(user_dir, 'pip')
    default_config_file = os.path.join(default_storage_dir, 'pip.ini')
    default_log_file = os.path.join(default_storage_dir, 'pip.log')
else:
    bin_py = os.path.join(sys.prefix, 'bin')

    #Use ~/.config/pip instead of ~/.pip- cleaner home folder
    #On some systems, we may have to create this, on others it probably exists
    if not os.path.exists(os.path.join(user_dir, '.config')):
        os.mkdir(os.path.join(user_dir, '.config'))
    default_storage_dir = os.path.join(user_dir, '.config', 'pip')

    default_config_file = os.path.join(default_storage_dir, 'pip.conf')
    default_log_file = os.path.join(default_storage_dir, 'pip.log')
    # Forcing to use /usr/local/bin for standard Mac OS X framework installs
    # Also log to ~/Library/Logs/ for use with the Console.app log viewer
    if sys.platform[:6] == 'darwin' and sys.prefix[:16] == '/System/Library/':
        bin_py = '/usr/local/bin'
        default_log_file = os.path.join(user_dir, 'Library/Logs/pip.log')
<commit_msg>Switched to use XDG_CONFIG_DIR environment variable if possible (default to ~/.config if not)<commit_after>"""Locations where we look for configs, install stuff, etc"""

import sys
import os
import tempfile
from pip.backwardcompat import get_python_lib

def running_under_virtualenv():
    """
    Return True if we're running inside a virtualenv, False otherwise.

    """
    return hasattr(sys, 'real_prefix')

if running_under_virtualenv():
    ## FIXME: is build/ a good name?
    build_prefix = os.path.join(sys.prefix, 'build')
    src_prefix = os.path.join(sys.prefix, 'src')
else:
    #Use tempfile to create a temporary folder
    build_prefix = tempfile.mkdtemp('-build', 'pip-')
    src_prefix = tempfile.mkdtemp('-src', 'pip-')
    # This is a terrible hack- since pip relies on this directory not being created yet
    # We will delete it now, and have pip recreate it later
    os.rmdir(build_prefix)
    os.rmdir(src_prefix)

# under Mac OS X + virtualenv sys.prefix is not properly resolved
# it is something like /path/to/python/bin/..
build_prefix = os.path.abspath(build_prefix)
src_prefix = os.path.abspath(src_prefix)

# FIXME doesn't account for venv linked to global site-packages

site_packages = get_python_lib()
user_dir = os.path.expanduser('~')
if sys.platform == 'win32':
    bin_py = os.path.join(sys.prefix, 'Scripts')
    # buildout uses 'bin' on Windows too?
    if not os.path.exists(bin_py):
        bin_py = os.path.join(sys.prefix, 'bin')
    user_dir = os.environ.get('APPDATA', user_dir) # Use %APPDATA% for roaming
    default_storage_dir = os.path.join(user_dir, 'pip')
    default_config_file = os.path.join(default_storage_dir, 'pip.ini')
    default_log_file = os.path.join(default_storage_dir, 'pip.log')
else:
    bin_py = os.path.join(sys.prefix, 'bin')

    # Use XDG_CONFIG_DIR instead of the ~/.pip
    # On some systems, we may have to create this, on others it probably exists
    xdg_dir = os.path.join(user_dir, '.config')
    xdg_dir = os.environ.get('XDG_CONFIG_DIR', xdg_dir)
    if not os.path.exists(xdg_dir):
        os.mkdir(xdg_dir)
    default_storage_dir = os.path.join(xdg_dir, 'pip')

    default_config_file = os.path.join(default_storage_dir, 'pip.conf')
    default_log_file = os.path.join(default_storage_dir, 'pip.log')
    # Forcing to use /usr/local/bin for standard Mac OS X framework installs
    # Also log to ~/Library/Logs/ for use with the Console.app log viewer
    if sys.platform[:6] == 'darwin' and sys.prefix[:16] == '/System/Library/':
        bin_py = '/usr/local/bin'
        default_log_file = os.path.join(user_dir, 'Library/Logs/pip.log')
<|endoftext|>