<commit_before>import time
from enum import IntEnum

from PIL import Image
from pylab import *

num_of_colors = 2
colormap = plt.get_cmap('inferno')
colors = [colormap(k) for k in np.linspace(0.0, 0.8, num_of_colors)]
sec_colors = [colormap(k + 0.1) for k in np.linspace(0.0, 0.8, num_of_colors)]

class Cell(IntEnum):
    EMPTY = 0
    RIGHT = 1
    DOWN = 2
    BOTH = 3

class BML:
    def __init__(self, width, height, p, model, animation_ratio=10):
        self.width = width
        self.height = height
        self.cells = np.zeros((height, width))
        self.model = model  # multidimentional, probabilistic, klein bottle, chain
        self.step = 0
        self.velocity = []
        self.animation_ration = animation_ratio

        num_of_cars = int(width * height * (p / 2))
        # Get a list of indices
        inds = list(np.ndindex(self.cells.shape))
        # Shuffle the indices in-place
        np.random.shuffle(inds)
        # Access array elements using the indices to do cool stuff
        for i in inds[:num_of_cars]:
            self.cells[i] = Cell.RIGHT

        for i in inds[num_of_cars:2 * num_of_cars]:
            self.cells[i] = Cell.DOWN

        self.total_number = 2 * num_of_cars

        self.animation = [self.get_image()]

    def make_step(self):
        if self.model == 1:
            vel = self.step_right()
            vel += self.step_down()
            self.velocity.append(vel / self.total_number)
        elif self.model == 2:
            self.velocity.append(self.step_all() / self.total_number)
        elif self.model == 3:
            self.velocity.append(self.step_all_together() / self.total_number)  # TODO
        self.step += 1
        if self.step % self.animation_ration == 0:
            self.animation.append(self.get_image())

    def step_down(self):
        d = np.diff(np.r_[self.cells, [self.cells[0, :]]], axis=0)
        down_indexes = np.where(np.logical_and(self.cells == Cell.DOWN, d == Cell.EMPTY - Cell.DOWN))
        empty_indexes = self.down_set(down_indexes)
        self.cells[down_indexes] = Cell.EMPTY
        self.cells[empty_indexes] = Cell.DOWN
        return len(down_indexes[0])

    def step_right(self):
        d = np.diff(np.c_[self.cells, self.cells[:, 0]])
        right_indexes = np.where(np.logical_and(self.cells == Cell.RIGHT, d == Cell.EMPTY - Cell.RIGHT))
        empty_indexes = self.right_set(right_indexes)
        self.cells[right_indexes] = Cell.EMPTY
        self.cells[empty_indexes] = Cell.RIGHT
        return len(right_indexes[0])

    def step_all(self):
        dr = np.c_[self.cells[:, -1], self.cells[:, :-1]] - self.cells
        dd = np.r_[[self.cells[-1, :]], self.cells[:-1, :]] - self.cells
        empty_indexes_both = np.transpose(np.where(np.logical_and(self.cells == Cell.EMPTY,
                                                                  np.logical_and(dd == Cell.DOWN - Cell.EMPTY,
                                                                                dr == Cell.RIGHT - Cell.EMPTY))))

        if empty_indexes_both.shape[0] > 0:
            np.random.shuffle(empty_indexes_both)
            right_nums = int(empty_indexes_both.shape[0] / 2)
            # @formatter:off
            empty_indexes_right = tuple(np.transpose(np.r_[
                                                      np.transpose(np.where(
                                                                   np.logical_and(self.cells == Cell.EMPTY,
                                                                   np.logical_and(dr == Cell.RIGHT - Cell.EMPTY,
                                                                                dd != Cell.DOWN - Cell.EMPTY)))),
                                                      empty_indexes_both[:right_nums]]))
            empty_indexes_down = tuple(np.transpose(np.r_[
                                                      np.transpose(np.where(
                                                              np.logical_and(self.cells == Cell.EMPTY,
                                                              np.logical_and(dd == Cell.DOWN - Cell.EMPTY,
                                                                             dr != Cell.RIGHT - Cell.EMPTY)))),
                                                      empty_indexes_both[right_nums:]]))
            # @formatter:on
        else:
            empty_indexes_right = np.where(np.logical_and(self.cells == Cell.EMPTY, dr == Cell.RIGHT - Cell.EMPTY))
            empty_indexes_down = np.where(np.logical_and(self.cells == Cell.EMPTY, dd == Cell.DOWN - Cell.EMPTY))

        right_indexes = self.left_set(empty_indexes_right)
        down_indexes = self.up_set(empty_indexes_down)

        self.cells[right_indexes] = Cell.EMPTY
        self.cells[down_indexes] = Cell.EMPTY
        self.cells[empty_indexes_right] = Cell.RIGHT
        self.cells[empty_indexes_down] = Cell.DOWN

        # if self.step % 100 == 1:
        #    self.save()
        return len(right_indexes[0]) + len(down_indexes[0])

    def step_all_together(self):
        dr = np.c_[self.cells[:, -1], self.cells[:, :-1]] - self.cells
        dd = np.r_[[self.cells[-1, :]], self.cells[:-1, :]] - self.cells

        down_cond = np.logical_or(dd == Cell.DOWN - Cell.EMPTY, dd == Cell.BOTH - Cell.EMPTY)
        right_cond = np.logical_or(dr == Cell.RIGHT - Cell.EMPTY, dr == Cell.BOTH - Cell.EMPTY)
        empty = self.cells == Cell.EMPTY
        empty_indexes_both = np.where(np.logical_and(empty, np.logical_and(down_cond, right_cond)))

        empty_indexes_right = np.where(np.logical_and(empty, np.logical_and(right_cond, np.logical_not(down_cond))))
        empty_indexes_down = np.where(np.logical_and(empty, np.logical_and(down_cond, np.logical_not(right_cond))))

        right_indexes = self.left_set(empty_indexes_right)
        down_indexes = self.up_set(empty_indexes_down)

        self.cells[right_indexes] -= 1
        self.cells[down_indexes] -= 2
        self.cells[empty_indexes_right] = Cell.RIGHT
        self.cells[empty_indexes_down] = Cell.DOWN

        if len(empty_indexes_both[0]) > 0:
            both_right_indexes = self.left_set(empty_indexes_both)
            both_down_indexes = self.up_set(empty_indexes_both)
            self.cells[both_right_indexes] -= 1
            self.cells[both_down_indexes] -= 2
            self.cells[empty_indexes_both] = Cell.BOTH
        # if self.step % 100 == 1:
        #    self.save()
        return len(right_indexes[0]) + len(down_indexes[0]) + len(empty_indexes_both[0]) * 2

        return -1

    def upper(self, i, j):
        return self.height - 1 if i == 0 else i - 1, j

    def left(self, i, j):
        return i, self.width - 1 if j == 0 else j - 1

    def upper_t(self, ind):
        return self.height - 1 if ind[0] == 0 else ind[0] - 1, ind[1]

    def left_t(self, ind):
        return ind[0], self.width - 1 if ind[1] == 0 else ind[1] - 1

    def down_t(self, ind):
        return 0 if ind[0] == self.height - 1 else ind[0] + 1, ind[1]

    def right_t(self, ind):
        return ind[0], 0 if ind[1] == self.width - 1 else ind[1] + 1

    def down_set(self, inds):
        return tuple([[(x + 1) % self.height for x in inds[0]], inds[1]])

    def right_set(self, inds):
        return tuple([inds[0], [(x + 1) % self.width for x in inds[1]]])

    def up_set(self, inds):
        return tuple([[(x - 1) % self.height for x in inds[0]], inds[1]])

    def left_set(self, inds):
        return tuple([inds[0], [(x - 1) % self.width for x in inds[1]]])

    def run(self, steps):
        start = time.perf_counter()
        for i in range(steps):
            self.make_step()
        total_time = time.perf_counter() - start
        print('{} steps for {}x{} map run in {}s. Average of {}s per step'.format(steps, self.height, self.width,
                                                                                total_time, total_time / steps))

    def get_image(self, pixel_size=2):
        """
        Returns image
        """
        visual = np.zeros((self.height, self.width, 3)).astype('uint8')
        down_indexes = np.where(self.cells == Cell.DOWN)
        right_indexes = np.where(self.cells == Cell.RIGHT)
        both_indexes = np.where(self.cells == Cell.BOTH)
        visual[:, :, :] = 255
        visual[:, :, 0][down_indexes] = 0  # down is blue

        visual[:, :, 1][down_indexes] = 0
        visual[:, :, 1][right_indexes] = 0
        visual[:, :, 1][both_indexes] = 0  # both is magenta

        visual[:, :, 2][right_indexes] = 0  # right are red
        im = Image.fromarray(visual, mode='RGB')
        im = im.resize((self.height * pixel_size, self.width * pixel_size))
        return im

    def save(self, pixel_size=10):
        """
        Saves current state as png image
        """
        visual = np.zeros((self.height, self.width, 3)).astype('uint8')
        down_indexes = np.where(self.cells == Cell.DOWN)
        right_indexes = np.where(self.cells == Cell.RIGHT)
        both_indexes = np.where(self.cells == Cell.BOTH)
        visual[:, :, :] = 255
        visual[:, :, 0][down_indexes] = 0  # down is blue

        visual[:, :, 1][down_indexes] = 0
        visual[:, :, 1][right_indexes] = 0
        visual[:, :, 1][both_indexes] = 0  # both is magenta

        visual[:, :, 2][right_indexes] = 0  # right are red
        im = Image.fromarray(visual, mode='RGB')
        im = im.resize((self.height * pixel_size, self.width * pixel_size))
        im.save("visual" + str(self.step) + ".png")

    def save_animation(self):
        self.animation[0].save("progress.gif", append_images=self.animation[1:], save_all=True, optimize=True,
                               duration=10)

    def plot_velocity(self):
        """
        Builds plot of velocity as a function of step number until now
        """
        params = {'axes.labelsize': 8, 'font.size': 8, 'legend.fontsize': 10, 'xtick.labelsize': 10,
                  'ytick.labelsize': 10, 'text.usetex': False, 'figure.figsize': [4.5, 4.5]}
        rcParams.update(params)
        fig = figure()  # no frame
        ax = fig.add_subplot(111)
        ax.plot(self.velocity, linewidth=1, color=colors[1])
        ax.set_ylim(0, 1)
        ax.set_yticks(np.arange(0, 1, 0.2))

        ax.get_xaxis().tick_bottom()
        ax.get_yaxis().tick_left()

        ax.set_xlabel('Step')
        ax.set_ylabel('Velocity')

        ax.grid(axis='y', color="0.9", linestyle='-', linewidth=1)
        ax.set_axisbelow(True)

        plt.tight_layout()

        fig.savefig('velocity.png', dpi=600)
        # print(self.velocity)

automat = BML(256, 256, 0.98, 3, 10)
automat.save()
automat.run(15000)
automat.save()
automat.save_animation()
automat.plot_velocity()
<commit_msg>Remove unused code<commit_after>import time
from enum import IntEnum

from PIL import Image
from pylab import *

num_of_colors = 2
colormap = plt.get_cmap('inferno')
colors = [colormap(k) for k in np.linspace(0.0, 0.8, num_of_colors)]
sec_colors = [colormap(k + 0.1) for k in np.linspace(0.0, 0.8, num_of_colors)]

class Cell(IntEnum):
    EMPTY = 0
    RIGHT = 1
    DOWN = 2
    BOTH = 3

class BML:
    def __init__(self, width, height, p, model, move_probability=1, animation_ratio=10):
        self.width = width
        self.height = height
        self.cells = np.zeros((height, width))
        self.model = model  # multidimentional, probabilistic, klein bottle, chain
        self.step = 0
        self.velocity = []
        self.animation_ration = animation_ratio
        self.move_probability = move_probability

        num_of_cars = int(width * height * (p / 2))
        # Get a list of indices
        inds = list(np.ndindex(self.cells.shape))
        # Shuffle the indices in-place
        np.random.shuffle(inds)
        # Access array elements using the indices to do cool stuff
        for i in inds[:num_of_cars]:
            self.cells[i] = Cell.RIGHT

        for i in inds[num_of_cars:2 * num_of_cars]:
            self.cells[i] = Cell.DOWN

        self.total_number = 2 * num_of_cars

        self.animation = [self.get_image()]

    def make_step(self):
        if self.model == 1:
            vel = self.step_right()
            vel += self.step_down()
            self.velocity.append(vel / self.total_number)
        elif self.model == 2:
            self.velocity.append(self.step_all() / self.total_number)
        elif self.model == 3:
            self.velocity.append(self.step_all_together() / self.total_number)
        self.step += 1
        if self.step % self.animation_ration == 0:
            self.animation.append(self.get_image())

    def step_down(self):
        d = np.diff(np.r_[self.cells, [self.cells[0, :]]], axis=0)  # TODO: transpose in some way and do one function
        down_indexes = np.where(np.logical_and(self.cells == Cell.DOWN, d == Cell.EMPTY - Cell.DOWN))
        empty_indexes = self.down_set(down_indexes)
        self.cells[down_indexes] = Cell.EMPTY
        self.cells[empty_indexes] = Cell.DOWN
        return len(down_indexes[0])

    def step_right(self):
        d = np.diff(np.c_[self.cells, self.cells[:, 0]])
        right_indexes = np.where(np.logical_and(self.cells == Cell.RIGHT, d == Cell.EMPTY - Cell.RIGHT))
        empty_indexes = self.right_set(right_indexes)
        self.cells[right_indexes] = Cell.EMPTY
        self.cells[empty_indexes] = Cell.RIGHT
        return len(right_indexes[0])

    def step_all(self):
        dr = np.c_[self.cells[:, -1], self.cells[:, :-1]] - self.cells
        dd = np.r_[[self.cells[-1, :]], self.cells[:-1, :]] - self.cells

        empty_cells = self.cells == Cell.EMPTY

        empty_indexes_both = np.transpose(np.where(np.logical_and(empty_cells,
                                                                  np.logical_and(dd == Cell.DOWN - Cell.EMPTY,
                                                                                dr == Cell.RIGHT - Cell.EMPTY))))

        if empty_indexes_both.shape[0] > 0:
            np.random.shuffle(empty_indexes_both)
            right_nums = int(empty_indexes_both.shape[0] / 2)
            # @formatter:off
            empty_indexes_right = tuple(np.transpose(np.r_[
                                                      np.transpose(np.where(
                                                                   np.logical_and(empty_cells,
                                                                   np.logical_and(dr == Cell.RIGHT - Cell.EMPTY,
                                                                                dd != Cell.DOWN - Cell.EMPTY)))),
                                                      empty_indexes_both[:right_nums]]))
            empty_indexes_down = tuple(np.transpose(np.r_[
                                                      np.transpose(np.where(
                                                              np.logical_and(empty_cells,
                                                              np.logical_and(dd == Cell.DOWN - Cell.EMPTY,
                                                                             dr != Cell.RIGHT - Cell.EMPTY)))),
                                                      empty_indexes_both[right_nums:]]))
            # @formatter:on
        else:
            empty_indexes_right = np.where(np.logical_and(empty_cells, dr == Cell.RIGHT - Cell.EMPTY))
            empty_indexes_down = np.where(np.logical_and(empty_cells, dd == Cell.DOWN - Cell.EMPTY))

        right_indexes = self.left_set(empty_indexes_right)
        down_indexes = self.up_set(empty_indexes_down)

        self.cells[right_indexes] = Cell.EMPTY
        self.cells[down_indexes] = Cell.EMPTY
        self.cells[empty_indexes_right] = Cell.RIGHT
        self.cells[empty_indexes_down] = Cell.DOWN

        # if self.step % 100 == 1:
        #    self.save()
        return len(right_indexes[0]) + len(down_indexes[0])

    def step_all_together(self):
        dr = np.c_[self.cells[:, -1], self.cells[:, :-1]] - self.cells
        dd = np.r_[[self.cells[-1, :]], self.cells[:-1, :]] - self.cells

        down_cond = np.logical_or(dd == Cell.DOWN - Cell.EMPTY, dd == Cell.BOTH - Cell.EMPTY)
        right_cond = np.logical_or(dr == Cell.RIGHT - Cell.EMPTY, dr == Cell.BOTH - Cell.EMPTY)
        empty_cells = self.cells == Cell.EMPTY
        empty_indexes_both = np.where(np.logical_and(empty_cells, np.logical_and(down_cond, right_cond)))

        empty_indexes_right = np.where(
            np.logical_and(empty_cells, np.logical_and(right_cond, np.logical_not(down_cond))))
        empty_indexes_down = np.where(
            np.logical_and(empty_cells, np.logical_and(down_cond, np.logical_not(right_cond))))

        right_indexes = self.left_set(empty_indexes_right)
        down_indexes = self.up_set(empty_indexes_down)

        self.cells[right_indexes] -= 1
        self.cells[down_indexes] -= 2
        self.cells[empty_indexes_right] = Cell.RIGHT
        self.cells[empty_indexes_down] = Cell.DOWN

        if len(empty_indexes_both[0]) > 0:
            both_right_indexes = self.left_set(empty_indexes_both)
            both_down_indexes = self.up_set(empty_indexes_both)
            self.cells[both_right_indexes] -= 1
            self.cells[both_down_indexes] -= 2
            self.cells[empty_indexes_both] = Cell.BOTH

        return len(right_indexes[0]) + len(down_indexes[0]) + len(empty_indexes_both[0]) * 2

    def down_set(self, inds):
        return tuple([[(x + 1) % self.height for x in inds[0]], inds[1]])

    def right_set(self, inds):
        return tuple([inds[0], [(x + 1) % self.width for x in inds[1]]])

    def up_set(self, inds):
        return tuple([[(x - 1) % self.height for x in inds[0]], inds[1]])

    def left_set(self, inds):
        return tuple([inds[0], [(x - 1) % self.width for x in inds[1]]])

    def run(self, steps):
        start = time.perf_counter()
        for i in range(steps):
            self.make_step()
        total_time = time.perf_counter() - start
        print('{} steps for {}x{} map run in {}s. Average of {}s per step'.format(steps, self.height, self.width,
                                                                                total_time, total_time / steps))

    def get_image(self, pixel_size=2):
        """
        Returns image
        """
        visual = np.zeros((self.height, self.width, 3)).astype('uint8')
        down_indexes = np.where(self.cells == Cell.DOWN)
        right_indexes = np.where(self.cells == Cell.RIGHT)
        both_indexes = np.where(self.cells == Cell.BOTH)
        visual[:, :, :] = 255
        visual[:, :, 0][down_indexes] = 0  # down is blue

        visual[:, :, 1][down_indexes] = 0
        visual[:, :, 1][right_indexes] = 0
        visual[:, :, 1][both_indexes] = 0  # both is magenta

        visual[:, :, 2][right_indexes] = 0  # right are red
        im = Image.fromarray(visual, mode='RGB')
        im = im.resize((self.height * pixel_size, self.width * pixel_size))
        return im

    def save(self, pixel_size=10):
        """
        Saves current state as png image
        """
        visual = np.zeros((self.height, self.width, 3)).astype('uint8')
        down_indexes = np.where(self.cells == Cell.DOWN)
        right_indexes = np.where(self.cells == Cell.RIGHT)
        both_indexes = np.where(self.cells == Cell.BOTH)
        visual[:, :, :] = 255
        visual[:, :, 0][down_indexes] = 0  # down is blue

        visual[:, :, 1][down_indexes] = 0
        visual[:, :, 1][right_indexes] = 0
        visual[:, :, 1][both_indexes] = 0  # both is magenta

        visual[:, :, 2][right_indexes] = 0  # right are red
        im = Image.fromarray(visual, mode='RGB')
        im = im.resize((self.height * pixel_size, self.width * pixel_size))
        im.save("visual" + str(self.step) + ".png")

    def save_animation(self):
        self.animation[0].save("progress.gif", append_images=self.animation[1:], save_all=True, optimize=True,
                               duration=10)

    def plot_velocity(self):
        """
        Builds plot of velocity as a function of step number until now
        """
        params = {'axes.labelsize': 8, 'font.size': 8, 'legend.fontsize': 10, 'xtick.labelsize': 10,
                  'ytick.labelsize': 10, 'text.usetex': False, 'figure.figsize': [4.5, 4.5]}
        rcParams.update(params)
        fig = figure()  # no frame
        ax = fig.add_subplot(111)
        ax.plot(self.velocity, linewidth=1, color=colors[1])
        ax.set_ylim(0, 1)
        ax.set_yticks(np.arange(0, 1, 0.2))

        ax.get_xaxis().tick_bottom()
        ax.get_yaxis().tick_left()

        ax.set_xlabel('Step')
        ax.set_ylabel('Velocity')

        ax.grid(axis='y', color="0.9", linestyle='-', linewidth=1)
        ax.set_axisbelow(True)

        plt.tight_layout()

        fig.savefig('velocity.png', dpi=600)
        # print(self.velocity)

automat = BML(256, 256, 0.8, 2, 1, 10)
automat.save()
automat.run(150)
automat.save()
automat.save_animation()
automat.plot_velocity()
<|endoftext|>


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

'''
This module contains classes of objects serving as clients
for remote and local spam listing services
'''
from __future__ import unicode_literals

# pylint: disable=redefined-builtin
from builtins import zip, map, str, range, object
from dns import name
from dns.resolver import NXDOMAIN, query
from future.moves.urllib.parse import urlparse
from future.utils import raise_from
from requests import get, post
from requests.exceptions import HTTPError

from .exceptions import UnathorizedAPIKeyError, UnknownCodeError, \
    InvalidHostError
from .structures import AddressListItem, hostname_or_ip, non_ipv6_host
from .validation import accepts_valid_urls, accepts_valid_host

class HostList(object):
    ''' A class of clients for local or remote host list services '''
    def __init__(self, host_factory):
        ''' Constructor

        :param host_factory: a function responsible for
        creating valid host objects. It may raise InvalidHostError
        (or its subclasses) if a value passed to it is not
        a valid host of type accepted by the factory.
        '''
        self._host_factory = host_factory

    def _contains(self, host_value):
        ''' Check if the service lists an item
        matching given host value

        :param host_value: a host value
        :returns: True if the service lists a matching
        value
        '''
        raise NotImplementedError

    def _get_match_and_classification(self, host_value):
        ''' Get a listed value that matches
        given host value and its classification

        :param host_value: a host value
        :returns: a tuple containing listed item and its classification as
        a tuple containing all classification groups to which the item belongs
        '''
        raise NotImplementedError

    @accepts_valid_host
    def __contains__(self, host_value):
        ''' Check if given host value is listed by the service

        :param host_value: a string representing a valid host
        :returns: True if the host is listed
        :raises InvalidHostError: if the argument is not a valid host string
        '''
        try:
            host_object = self._host_factory(host_value)
        except InvalidHostError:
            return False
        return self._contains(host_object)

    @accepts_valid_host
    def lookup(self, host_value):
        ''' Get an object representing a host value
        matched by this host

        :param host_value: a value of the host of a type
        that can be listed by the service
        :returns: an instance of AddressListItem representing
        a matched value
        :raises InvalidHostError: if the argument is not a valid host string
        '''
        try:
            host_object = self._host_factory(host_value)
        except InvalidHostError:
            return None
        result = self._get_match_and_classification(
            host_object
        )
        host_item, classification = result
        if host_item is not None:
            return AddressListItem(
                host_item.to_unicode(),
                self,
                classification
            )
        return None

    @accepts_valid_urls
    def any_match(self, urls):
        ''' Check if any of given urls has a listed host

        :param urls: an iterable containing urls
        :returns: True if any host is listed
        :raises InvalidURLError: if there are any invalid urls in the sequence
        '''
        return any(urlparse(u).hostname in self for u in urls)

    @accepts_valid_urls
    def lookup_matching(self, urls):
        '''Get objects representing hosts in given urls
        that match listed hosts

        :param urls: an iterable containing urls
        :returns: items representing hosts matching the listed ones
        :raises InvalidURLError: if there are any invalid urls in the sequence
        '''
        hosts = (urlparse(u).hostname for u in urls)
        for val in hosts:
            item = self.lookup(val)
            if item is not None:
                yield item

    @accepts_valid_urls
    def filter_matching(self, urls):
        ''' Get urls with hosts matching listed ones

        :param urls: an iterable containing url addresses to filter
        :returns: a list containing matching urls
        :raises InvalidURLError: if there are any invalid urls in the sequence
        '''
        def is_match(url):
            return urlparse(url).hostname in self
        return (u for u in urls if is_match(u))

class DNSBL(HostList):
    ''' Represents a DNSBL service '''
    def __init__(
            self,
            identifier,
            query_suffix,
            classification_map,
            host_factory):
        ''' Create new DNSBL object

        :param identifier: a value designating DNSBL service provider:
        its name or url address.
        :param query_suffix: a suffix added to DNSBL query address
        :param classification_map: item classes associated with
        DNSBL query return codes
        :param host_factory: a callable object that returns an object
         representing host and providing method for getting a relative
         domain pertaining to it.
        '''
        self._identifier = identifier
        self._query_suffix = name.from_text(query_suffix)
        self._classification_map = classification_map
        self._host_factory = host_factory
        super(DNSBL, self).__init__(host_factory)

    def _query(self, host_object):
        ''' Query DNSBL service for given value

        :param host_object: an object representing host,
         created by _host_factory
        :returns: an instance of dns.resolver.Answer for
         given value, if it is listed. Otherwise, it returns None
        '''
        host_to_query = host_object.relative_domain
        query_name = host_to_query.derelativize(self._query_suffix)
        try:
            return query(query_name)
        except NXDOMAIN:
            return None

    def __str__(self):
        return str(self._identifier)

    def _contains(self, host_object):
        return bool(self._query(host_object))

    def _get_entry_classification(self, code):
        return [self._classification_map[code]]

    def _get_match_and_classification(self, host_object):
        answers = self._query(host_object)
        if answers is None:
            return None, None
        try:
            classification = set()
            for answer in answers:
                last_octet = answer.to_text().split('.')[-1]
                classes = self._get_entry_classification(int(last_octet))
                classification.update(classes)
            return host_object, classification
        except KeyError as ex:
            msg_tpl = "The code '{}' has no corresponding classification value"
            msg = msg_tpl.format(ex.args[0])
            raise_from(UnknownCodeError(msg), ex)

def get_powers_of_2(_sum):
    ''' Get powers of with a given sum

    This function transforms given integer to a binary string.
    A reversed value limited to digits of binary number is extracted
    from it, and each of its characters is enumerated.

    Each digit is tested for not being 0. If the test passes, the index
    associated with the digit is used as an exponent to get the next
    value in the sequence to be returned.

    :param _sum: a sum of all elements of the sequence to be returned
    :returns: a list of powers of two whose sum is given
    '''
    return [2**y for y, x in enumerate(bin(_sum)[:1:-1]) if int(x)]

class BitmaskingDNSBL(DNSBL):
    ''' A class representing DNSBL services mapping listed items to
    sums of classification codes

    Each classification code is a power of two.

    '''
    def _get_entry_classification(self, code):
        codes = get_powers_of_2(code)
        return [cl for c in codes for cl
                in DNSBL._get_entry_classification(self, c)]

class HpHosts(HostList):
    ''' hpHosts client '''
    identifier = ' http://www.hosts-file.net/'
    _NOT_LISTED = 'Not Listed'

    def __init__(self, client_name):
        '''Constructor

        :param client_name: name of client using the service
        '''
        self.app_id = client_name
        super(HpHosts, self).__init__(non_ipv6_host)

    def _query(self, host_object, classification=False):
        ''' Query the client for data of given host

        :param host_object: an object representing a host value
        :param classification: if True: hpHosts is queried also
         for classification for given host, if listed
        :returns: content of response to GET request to hpHosts
         for data on the given host
        '''
        template = 'http://verify.hosts-file.net/?v={}&s={}'
        url = template.format(self.app_id, host_object.to_unicode())
        url = url + '&class=true' if classification else url
        return get(url).text

    def _contains(self, host_object):
        return self._NOT_LISTED not in self._query(host_object)

    def _get_match_and_classification(self, host_object):
        data = self._query(host_object, True)
        if self._NOT_LISTED in data:
            return None, None
        elements = data.split(',')
        classification = set(elements[1:])
        return host_object, classification

class GoogleSafeBrowsing(object):
    ''' Google Safe Browsing lookup API client '''
    protocol_version = '3.1'
    max_urls_per_request = 500

    def __init__(self, client_name, app_version, api_key):
        ''' Create new instance

        :param client_name: name of application using the API
        :param app_version: version of the application
        :param api_key: API key given by Google:
        https://developers.google.com/safe-browsing/key_signup
        '''
        self.api_key = api_key
        self.client_name = client_name
        self.app_version = app_version
        self._request_address_val = ''

    @property
    def _request_address(self):
        ''' Get address of POST request to the service '''
        if not self._request_address_val:
            template = (
                'https://sb-ssl.google.com/safebrowsing/api/lookup'
                '?client={0}&key={1}&appver={2}&pver={3}'
            )
            self._request_address_val = template.format(
                self.client_name,
                self.api_key,
                self.app_version,
                self.protocol_version
            )
        return self._request_address_val

    def _query_once(self, urls):
        ''' Perform a single POST request using lookup API

        :param urls: a sequence of urls to put in request body
        :returns: a response object
        :raises UnathorizedAPIKeyError: when the API key for this instance
        is not valid
        :raises HTTPError: if the HTTPError was raised for a HTTP code
        other than 401, the exception is reraised
        '''
        request_body = '{}\n{}'.format(len(urls), '\n'.join(urls))
        response = post(self._request_address, request_body)
        try:
            response.raise_for_status()
        except HTTPError as error:
            if response.status_code == 401:
                msg = 'The API key is not authorized'
                raise_from(UnathorizedAPIKeyError(msg), error)
            else:
                raise
        return response

    def _query(self, urls):
        ''' Test urls for being listed by the service

        :param urls: a sequence of urls  to be tested
        :returns: a tuple containing chunk of urls and a response pertaining
          to them if the code of response was 200, which means at least one
          of the queried URLs is matched in either the phishing, malware,
           or unwanted software lists.
        '''
        urls = list(set(urls))
        for i in range(0, len(urls), self.max_urls_per_request):
            chunk = urls[i:i+self.max_urls_per_request]
            response = self._query_once(chunk)
            if response.status_code == 200:
                yield chunk, response

    @accepts_valid_urls
    def any_match(self, urls):
        ''' Check if the service recognizes any of given urls as spam

        :param urls: a sequence of urls to be tested
        :returns: True if any of the urls was recognized as spam
        :raises InvalidURLError: if there are any invalid urls in the sequence
        '''
        return any(self._query(urls))

    def _get_match_and_classification(self, urls):
        ''' Get classification for all matching urls

        :param urls: a sequence of urls to test
        :return: a tuple containing matching url and classification
        string pertaining to it
        '''
        for url_list, response in self._query(urls):
            classification_set = response.text.splitlines()
            for url, _class in zip(url_list, classification_set):
                if _class != 'ok':
                    yield url, _class

    @accepts_valid_urls
    def lookup_matching(self, urls):
        ''' Get items for all listed urls

        :param urls: a sequence of urls to be tested
        :returns: objects representing listed urls
        :raises InvalidURLError: if there are any invalid urls in the sequence
        '''
        for url, _class in self._get_match_and_classification(urls):
            classification = set(_class.split(','))
            yield AddressListItem(url, self, classification)

    @accepts_valid_urls
    def filter_matching(self, urls):
        ''' Get all listed urls

        :param urls: a sequence of urls to be tested
        :returns: spam urls
        :raises InvalidURLError: if there are any invalid urls in the sequence
        '''
        for url, _ in self._get_match_and_classification(urls):
            yield url

class HostCollection(HostList):
    ''' Provides a container for ip addresses and domain names.

    May be used as a local whitelist or blacklist.
    '''
    def __init__(self, identifier, classification, hosts=()):
        ''' Create new instance

        :param identifier: an identifier of this instance of host collection
        :param classification: a list or tuple containing strings representing
        types of items, assigned to each element of the collection
        :param hosts: a sequence of ip adresses and hostnames
        '''
        self.identifier = identifier
        self.classification = set(classification)
        self.hosts = set()
        for host_value in hosts:
            self.add(host_value)
        super(HostCollection, self).__init__(hostname_or_ip)

    @property
    def _host_objects(self):
        for h in self.hosts:
            yield self._host_factory(h)

    def _contains(self, host_object):
        def test(host):
            return host_object.is_match(host)
        return any(map(test, self._host_objects))

    def _get_match_and_classification(self, host_object):
        for val in self._host_objects:
            if host_object.is_match(val):
                return val, self.classification
        return None, None

    def add(self, host_value):
        ''' Add the given value to collection

        :param host: an ip address or a hostname
        :raises InvalidHostError: raised when the given value
        is not a valid ip address nor a hostname
        '''
        host_obj = self._host_factory(host_value)
        hosts = set(self.hosts)
        for listed_obj, listed in ((self._host_factory(h), h) for h in hosts):
            if host_obj.is_match(listed_obj):
                return
            elif listed_obj.is_subdomain(host_obj):
                self.hosts.remove(listed)
        self.hosts.add(host_value)
<commit_msg>Change the way HostCollection.hosts property is initialized<commit_after># -*- coding: utf-8 -*-

'''
This module contains classes of objects serving as clients
for remote and local spam listing services
'''
from __future__ import unicode_literals

# pylint: disable=redefined-builtin
from builtins import zip, map, str, range, object
from dns import name
from dns.resolver import NXDOMAIN, query
from future.moves.urllib.parse import urlparse
from future.utils import raise_from
from requests import get, post
from requests.exceptions import HTTPError

from .exceptions import UnathorizedAPIKeyError, UnknownCodeError, \
    InvalidHostError
from .structures import AddressListItem, hostname_or_ip, non_ipv6_host
from .validation import accepts_valid_urls, accepts_valid_host

class HostList(object):
    ''' A class of clients for local or remote host list services '''
    def __init__(self, host_factory):
        ''' Constructor

        :param host_factory: a function responsible for
        creating valid host objects. It may raise InvalidHostError
        (or its subclasses) if a value passed to it is not
        a valid host of type accepted by the factory.
        '''
        self._host_factory = host_factory

    def _contains(self, host_value):
        ''' Check if the service lists an item
        matching given host value

        :param host_value: a host value
        :returns: True if the service lists a matching
        value
        '''
        raise NotImplementedError

    def _get_match_and_classification(self, host_value):
        ''' Get a listed value that matches
        given host value and its classification

        :param host_value: a host value
        :returns: a tuple containing listed item and its classification as
        a tuple containing all classification groups to which the item belongs
        '''
        raise NotImplementedError

    @accepts_valid_host
    def __contains__(self, host_value):
        ''' Check if given host value is listed by the service

        :param host_value: a string representing a valid host
        :returns: True if the host is listed
        :raises InvalidHostError: if the argument is not a valid host string
        '''
        try:
            host_object = self._host_factory(host_value)
        except InvalidHostError:
            return False
        return self._contains(host_object)

    @accepts_valid_host
    def lookup(self, host_value):
        ''' Get an object representing a host value
        matched by this host

        :param host_value: a value of the host of a type
        that can be listed by the service
        :returns: an instance of AddressListItem representing
        a matched value
        :raises InvalidHostError: if the argument is not a valid host string
        '''
        try:
            host_object = self._host_factory(host_value)
        except InvalidHostError:
            return None
        result = self._get_match_and_classification(
            host_object
        )
        host_item, classification = result
        if host_item is not None:
            return AddressListItem(
                host_item.to_unicode(),
                self,
                classification
            )
        return None

    @accepts_valid_urls
    def any_match(self, urls):
        ''' Check if any of given urls has a listed host

        :param urls: an iterable containing urls
        :returns: True if any host is listed
        :raises InvalidURLError: if there are any invalid urls in the sequence
        '''
        return any(urlparse(u).hostname in self for u in urls)

    @accepts_valid_urls
    def lookup_matching(self, urls):
        '''Get objects representing hosts in given urls
        that match listed hosts

        :param urls: an iterable containing urls
        :returns: items representing hosts matching the listed ones
        :raises InvalidURLError: if there are any invalid urls in the sequence
        '''
        hosts = (urlparse(u).hostname for u in urls)
        for val in hosts:
            item = self.lookup(val)
            if item is not None:
                yield item

    @accepts_valid_urls
    def filter_matching(self, urls):
        ''' Get urls with hosts matching listed ones

        :param urls: an iterable containing url addresses to filter
        :returns: a list containing matching urls
        :raises InvalidURLError: if there are any invalid urls in the sequence
        '''
        def is_match(url):
            return urlparse(url).hostname in self
        return (u for u in urls if is_match(u))

class DNSBL(HostList):
    ''' Represents a DNSBL service '''
    def __init__(
            self,
            identifier,
            query_suffix,
            classification_map,
            host_factory):
        ''' Create new DNSBL object

        :param identifier: a value designating DNSBL service provider:
        its name or url address.
        :param query_suffix: a suffix added to DNSBL query address
        :param classification_map: item classes associated with
        DNSBL query return codes
        :param host_factory: a callable object that returns an object
         representing host and providing method for getting a relative
         domain pertaining to it.
        '''
        self._identifier = identifier
        self._query_suffix = name.from_text(query_suffix)
        self._classification_map = classification_map
        self._host_factory = host_factory
        super(DNSBL, self).__init__(host_factory)

    def _query(self, host_object):
        ''' Query DNSBL service for given value

        :param host_object: an object representing host,
         created by _host_factory
        :returns: an instance of dns.resolver.Answer for
         given value, if it is listed. Otherwise, it returns None
        '''
        host_to_query = host_object.relative_domain
        query_name = host_to_query.derelativize(self._query_suffix)
        try:
            return query(query_name)
        except NXDOMAIN:
            return None

    def __str__(self):
        return str(self._identifier)

    def _contains(self, host_object):
        return bool(self._query(host_object))

    def _get_entry_classification(self, code):
        return [self._classification_map[code]]

    def _get_match_and_classification(self, host_object):
        answers = self._query(host_object)
        if answers is None:
            return None, None
        try:
            classification = set()
            for answer in answers:
                last_octet = answer.to_text().split('.')[-1]
                classes = self._get_entry_classification(int(last_octet))
                classification.update(classes)
            return host_object, classification
        except KeyError as ex:
            msg_tpl = "The code '{}' has no corresponding classification value"
            msg = msg_tpl.format(ex.args[0])
            raise_from(UnknownCodeError(msg), ex)

def get_powers_of_2(_sum):
    ''' Get powers of with a given sum

    This function transforms given integer to a binary string.
    A reversed value limited to digits of binary number is extracted
    from it, and each of its characters is enumerated.

    Each digit is tested for not being 0. If the test passes, the index
    associated with the digit is used as an exponent to get the next
    value in the sequence to be returned.

    :param _sum: a sum of all elements of the sequence to be returned
    :returns: a list of powers of two whose sum is given
    '''
    return [2**y for y, x in enumerate(bin(_sum)[:1:-1]) if int(x)]

class BitmaskingDNSBL(DNSBL):
    ''' A class representing DNSBL services mapping listed items to
    sums of classification codes

    Each classification code is a power of two.

    '''
    def _get_entry_classification(self, code):
        codes = get_powers_of_2(code)
        return [cl for c in codes for cl
                in DNSBL._get_entry_classification(self, c)]

class HpHosts(HostList):
    ''' hpHosts client '''
    identifier = ' http://www.hosts-file.net/'
    _NOT_LISTED = 'Not Listed'

    def __init__(self, client_name):
        '''Constructor

        :param client_name: name of client using the service
        '''
        self.app_id = client_name
        super(HpHosts, self).__init__(non_ipv6_host)

    def _query(self, host_object, classification=False):
        ''' Query the client for data of given host

        :param host_object: an object representing a host value
        :param classification: if True: hpHosts is queried also
         for classification for given host, if listed
        :returns: content of response to GET request to hpHosts
         for data on the given host
        '''
        template = 'http://verify.hosts-file.net/?v={}&s={}'
        url = template.format(self.app_id, host_object.to_unicode())
        url = url + '&class=true' if classification else url
        return get(url).text

    def _contains(self, host_object):
        return self._NOT_LISTED not in self._query(host_object)

    def _get_match_and_classification(self, host_object):
        data = self._query(host_object, True)
        if self._NOT_LISTED in data:
            return None, None
        elements = data.split(',')
        classification = set(elements[1:])
        return host_object, classification

class GoogleSafeBrowsing(object):
    ''' Google Safe Browsing lookup API client '''
    protocol_version = '3.1'
    max_urls_per_request = 500

    def __init__(self, client_name, app_version, api_key):
        ''' Create new instance

        :param client_name: name of application using the API
        :param app_version: version of the application
        :param api_key: API key given by Google:
        https://developers.google.com/safe-browsing/key_signup
        '''
        self.api_key = api_key
        self.client_name = client_name
        self.app_version = app_version
        self._request_address_val = ''

    @property
    def _request_address(self):
        ''' Get address of POST request to the service '''
        if not self._request_address_val:
            template = (
                'https://sb-ssl.google.com/safebrowsing/api/lookup'
                '?client={0}&key={1}&appver={2}&pver={3}'
            )
            self._request_address_val = template.format(
                self.client_name,
                self.api_key,
                self.app_version,
                self.protocol_version
            )
        return self._request_address_val

    def _query_once(self, urls):
        ''' Perform a single POST request using lookup API

        :param urls: a sequence of urls to put in request body
        :returns: a response object
        :raises UnathorizedAPIKeyError: when the API key for this instance
        is not valid
        :raises HTTPError: if the HTTPError was raised for a HTTP code
        other than 401, the exception is reraised
        '''
        request_body = '{}\n{}'.format(len(urls), '\n'.join(urls))
        response = post(self._request_address, request_body)
        try:
            response.raise_for_status()
        except HTTPError as error:
            if response.status_code == 401:
                msg = 'The API key is not authorized'
                raise_from(UnathorizedAPIKeyError(msg), error)
            else:
                raise
        return response

    def _query(self, urls):
        ''' Test urls for being listed by the service

        :param urls: a sequence of urls  to be tested
        :returns: a tuple containing chunk of urls and a response pertaining
          to them if the code of response was 200, which means at least one
          of the queried URLs is matched in either the phishing, malware,
           or unwanted software lists.
        '''
        urls = list(set(urls))
        for i in range(0, len(urls), self.max_urls_per_request):
            chunk = urls[i:i+self.max_urls_per_request]
            response = self._query_once(chunk)
            if response.status_code == 200:
                yield chunk, response

    @accepts_valid_urls
    def any_match(self, urls):
        ''' Check if the service recognizes any of given urls as spam

        :param urls: a sequence of urls to be tested
        :returns: True if any of the urls was recognized as spam
        :raises InvalidURLError: if there are any invalid urls in the sequence
        '''
        return any(self._query(urls))

    def _get_match_and_classification(self, urls):
        ''' Get classification for all matching urls

        :param urls: a sequence of urls to test
        :return: a tuple containing matching url and classification
        string pertaining to it
        '''
        for url_list, response in self._query(urls):
            classification_set = response.text.splitlines()
            for url, _class in zip(url_list, classification_set):
                if _class != 'ok':
                    yield url, _class

    @accepts_valid_urls
    def lookup_matching(self, urls):
        ''' Get items for all listed urls

        :param urls: a sequence of urls to be tested
        :returns: objects representing listed urls
        :raises InvalidURLError: if there are any invalid urls in the sequence
        '''
        for url, _class in self._get_match_and_classification(urls):
            classification = set(_class.split(','))
            yield AddressListItem(url, self, classification)

    @accepts_valid_urls
    def filter_matching(self, urls):
        ''' Get all listed urls

        :param urls: a sequence of urls to be tested
        :returns: spam urls
        :raises InvalidURLError: if there are any invalid urls in the sequence
        '''
        for url, _ in self._get_match_and_classification(urls):
            yield url

class HostCollection(HostList):
    ''' Provides a container for ip addresses and domain names.

    May be used as a local whitelist or blacklist.
    '''
    def __init__(self, identifier, classification, hosts=None):
        ''' Create new instance

        :param identifier: an identifier of this instance of host collection
        :param classification: a list or tuple containing strings representing
        types of items, assigned to each element of the collection
        :param hosts: an object storing ip adresses and hostnames. It
        must be iterable and have .add and .remove methods.
        '''
        self.identifier = identifier
        self.classification = set(classification)
        self.hosts = hosts if hosts is not None else set()
        super(HostCollection, self).__init__(hostname_or_ip)

    @property
    def _host_objects(self):
        for h in self.hosts:
            yield self._host_factory(h)

    def _contains(self, host_object):
        def test(host):
            return host_object.is_match(host)
        return any(map(test, self._host_objects))

    def _get_match_and_classification(self, host_object):
        for val in self._host_objects:
            if host_object.is_match(val):
                return val, self.classification
        return None, None

    def add(self, host_value):
        ''' Add the given value to collection

        :param host: an ip address or a hostname
        :raises InvalidHostError: raised when the given value
        is not a valid ip address nor a hostname
        '''
        host_obj = self._host_factory(host_value)
        hosts = set(self.hosts)
        for listed_obj, listed in ((self._host_factory(h), h) for h in hosts):
            if host_obj.is_match(listed_obj):
                return
            elif listed_obj.is_subdomain(host_obj):
                self.hosts.remove(listed)
        self.hosts.add(host_value)
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
from __future__ import absolute_import
from datetime import datetime
from contextlib import contextmanager

from PyQt4.QtNetwork import (
    QNetworkAccessManager, QNetworkProxyQuery,
    QNetworkReply, QNetworkRequest
)
from PyQt4.QtWebKit import QWebFrame
from twisted.python import log

from splash.qtutils import qurl2ascii, OPERATION_NAMES
from splash import har
from splash.request_middleware import (
    AdblockMiddleware,
    AllowedDomainsMiddleware,
    AllowedSchemesMiddleware,
    RequestLoggingMiddleware,
    AdblockRulesRegistry,
)

# See: http://pyqt.sourceforge.net/Docs/PyQt4/qnetworkreply.html#NetworkError-enum
REQUEST_ERRORS = {
    QNetworkReply.NoError : 'no error condition. Note: When the HTTP protocol returns a redirect no error will be reported. You can check if there is a redirect with the QNetworkRequest::RedirectionTargetAttribute attribute.',
    QNetworkReply.ConnectionRefusedError : 'the remote server refused the connection (the server is not accepting requests)',
    QNetworkReply.RemoteHostClosedError : 'the remote server closed the connection prematurely, before the entire reply was received and processed',
    QNetworkReply.HostNotFoundError : 'the remote host name was not found (invalid hostname)',
    QNetworkReply.TimeoutError : 'the connection to the remote server timed out',
    QNetworkReply.OperationCanceledError : 'the operation was canceled via calls to abort() or close() before it was finished.',
    QNetworkReply.SslHandshakeFailedError : 'the SSL/TLS handshake failed and the encrypted channel could not be established. The sslErrors() signal should have been emitted.',
    QNetworkReply.TemporaryNetworkFailureError : 'the connection was broken due to disconnection from the network, however the system has initiated roaming to another access point. The request should be resubmitted and will be processed as soon as the connection is re-established.',
    QNetworkReply.ProxyConnectionRefusedError : 'the connection to the proxy server was refused (the proxy server is not accepting requests)',
    QNetworkReply.ProxyConnectionClosedError : 'the proxy server closed the connection prematurely, before the entire reply was received and processed',
    QNetworkReply.ProxyNotFoundError : 'the proxy host name was not found (invalid proxy hostname)',
    QNetworkReply.ProxyTimeoutError : 'the connection to the proxy timed out or the proxy did not reply in time to the request sent',
    QNetworkReply.ProxyAuthenticationRequiredError : 'the proxy requires authentication in order to honour the request but did not accept any credentials offered (if any)',
    QNetworkReply.ContentAccessDenied : 'the access to the remote content was denied (similar to HTTP error 401)',
    QNetworkReply.ContentOperationNotPermittedError : 'the operation requested on the remote content is not permitted',
    QNetworkReply.ContentNotFoundError : 'the remote content was not found at the server (similar to HTTP error 404)',
    QNetworkReply.AuthenticationRequiredError : 'the remote server requires authentication to serve the content but the credentials provided were not accepted (if any)',
    QNetworkReply.ContentReSendError : 'the request needed to be sent again, but this failed for example because the upload data could not be read a second time.',
    QNetworkReply.ProtocolUnknownError : 'the Network Access API cannot honor the request because the protocol is not known',
    QNetworkReply.ProtocolInvalidOperationError : 'the requested operation is invalid for this protocol',
    QNetworkReply.UnknownNetworkError : 'an unknown network-related error was detected',
    QNetworkReply.UnknownProxyError : 'an unknown proxy-related error was detected',
    QNetworkReply.UnknownContentError : 'an unknown error related to the remote content was detected',
    QNetworkReply.ProtocolFailure : 'a breakdown in protocol was detected (parsing error, invalid or unexpected responses, etc.)',
}

class ProxiedQNetworkAccessManager(QNetworkAccessManager):
    """
    This QNetworkAccessManager subclass enables "splash proxy factories"
    support. Qt provides similar functionality via setProxyFactory method,
    but standard QNetworkProxyFactory is not flexible enough.

    It also sets up some extra logging, provides a way to get
    the "source" request (that was made to Splash itself) and tracks
    information about requests/responses.
    """

    _REQUEST_ID = QNetworkRequest.User + 1

    REQUEST_CREATED = "created"
    REQUEST_FINISHED = "finished"
    REQUEST_HEADERS_RECEIVED = "headers"

    def __init__(self, verbosity):
        super(ProxiedQNetworkAccessManager, self).__init__()
        self.sslErrors.connect(self._sslErrors)
        self.finished.connect(self._finished)
        self.verbosity = verbosity
        self._next_id = 0

        assert self.proxyFactory() is None, "Standard QNetworkProxyFactory is not supported"

    def _sslErrors(self, reply, errors):
        reply.ignoreSslErrors()

    def _finished(self, reply):
        reply.deleteLater()

    def createRequest(self, operation, request, outgoingData=None):
        """
        This method is called when a new request is sent;
        it must return a reply object to work with.
        """
        start_time = datetime.utcnow()

        request = self._wrapRequest(request)
        har_entry = self._harEntry(request, create=True)
        if har_entry is not None:
            page_id = self._getWebPageAttribute(request, "page_id")
            if outgoingData is None:
                bodySize = -1
            else:
                bodySize = outgoingData.size()
            har_entry.update({
                '_tmp': {
                    'start_time': start_time,
                    'request_start_sending_time': start_time,
                    'request_sent_time': start_time,

                    # 'outgoingData': outgoingData,
                    'state': self.REQUEST_CREATED,
                },
                "pageref": page_id,
                "startedDateTime": har.format_datetime(start_time),
                "request": {
                    "method": OPERATION_NAMES.get(operation, '?'),
                    "url": unicode(request.url().toString()),
                    "httpVersion": "HTTP/1.1",
                    "cookies": har.request_cookies2har(request),
                    "queryString": har.querystring2har(request.url()),
                    "headers": har.headers2har(request),

                    "headersSize" : har.headers_size(request),
                    "bodySize": bodySize,
                },
                "response": {},
                "cache": {},
                "timings": {
                    "send": 0,
                    "wait": 0,
                    "receive": 0,
                },
            })

        with self._proxyApplied(request):
            reply = super(ProxiedQNetworkAccessManager, self).createRequest(
                operation, request, outgoingData
            )
            reply.error.connect(self._handleError)
            reply.finished.connect(self._handleFinished)
            reply.metaDataChanged.connect(self._handleMetaData)
            reply.downloadProgress.connect(self._handleDownloadProgress)

        return reply

    @contextmanager
    def _proxyApplied(self, request):
        """
        This context manager temporary sets a proxy based on request options.
        """
        old_proxy = self.proxy()
        splash_proxy_factory = self._getSplashProxyFactory(request)
        if splash_proxy_factory:
            proxy_query = QNetworkProxyQuery(request.url())
            proxy = splash_proxy_factory.queryProxy(proxy_query)[0]
            self.setProxy(proxy)
        try:
            yield
        finally:
            self.setProxy(old_proxy)

    def _wrapRequest(self, request):
        request = QNetworkRequest(request)
        request.setAttribute(self._REQUEST_ID, self._next_id)
        self._next_id += 1
        return request

    def _getRequestId(self, request=None):
        if request is None:
            request = self.sender().request()
        return request.attribute(self._REQUEST_ID).toPyObject()

    def _harEntry(self, request=None, create=False):
        """
        Return a mutable dictionary for request/response
        information storage.
        """
        if request is None:
            request = self.sender().request()

        entries = self._getWebPageAttribute(request, 'network_entries_map')
        if entries is None:
            return

        req_id = self._getRequestId(request)
        if create:
            assert req_id not in entries
            entries[req_id] = {
                "_idx": req_id,
            }
        return entries[req_id]

    def _getSplashProxyFactory(self, request):
        return self._getWebPageAttribute(request, 'splash_proxy_factory')

    def _getWebPageAttribute(self, request, attribute):
        web_frame = request.originatingObject()
        if isinstance(web_frame, QWebFrame):
            return getattr(web_frame.page(), attribute, None)

    def _handleError(self, error_id):
        error_msg = REQUEST_ERRORS.get(error_id, 'unknown error')
        self.log("Download error %d: %s ({url})" % (error_id, error_msg), self.sender(), min_level=1)

    def _handleFinished(self):
        reply = self.sender()
        har_entry = self._harEntry()
        if har_entry is not None:
            har_entry["_tmp"]["state"] = self.REQUEST_FINISHED

            now = datetime.utcnow()
            start_time = har_entry['_tmp']['start_time']
            response_start_time = har_entry['_tmp']['response_start_time']

            receive_time = har.get_duration(response_start_time, now)
            total_time = har.get_duration(start_time, now)

            har_entry["timings"]["receive"] = receive_time
            har_entry["time"] = total_time

            if not har_entry["timings"]["send"]:
                wait_time = har_entry["timings"]["wait"]
                har_entry["timings"]["send"] = total_time - receive_time - wait_time

            har_entry["response"].update(har.reply2har(reply))
            if 'bodySize' not in har_entry["response"]:
                har_entry["response"]["bodySize"] = -1

        self.log("Finished downloading {url}", reply)

    def _handleMetaData(self):
        reply = self.sender()
        har_entry = self._harEntry()
        if har_entry is not None:
            har_entry["_tmp"]["state"] = self.REQUEST_HEADERS_RECEIVED
            har_entry["response"].update(har.reply2har(reply))

            now = datetime.utcnow()
            request_sent = har_entry["_tmp"]["request_sent_time"]
            har_entry["_tmp"]["response_start_time"] = now
            har_entry["timings"]["wait"] = har.get_duration(request_sent, now)

        self.log("Headers received for {url}", reply, min_level=3)

    def _handleDownloadProgress(self, received, total):
        har_entry = self._harEntry()
        har_entry["response"]["bodySize"] = int(received)

        if total == -1:
            total = '?'
        self.log("Downloaded %d/%s of {url}" % (received, total), self.sender(), min_level=3)

    def _handleUploadProgress(self, sent, total):
        har_entry = self._harEntry()
        har_entry["request"]["bodySize"] = int(sent)

        now = datetime.utcnow()
        if sent == 0:
            # it is a moment the sending is started
            start_time = har_entry["_tmp"]["request_start_time"]
            har_entry["_tmp"]["request_start_sending_time"] = now
            har_entry["timings"]["blocked"] = har.get_duration(start_time, now)

        har_entry["_tmp"]["request_sent_time"] = now

        if sent == total:
            start_sending_time = har_entry["_tmp"]["request_start_sending_time"]
            har_entry["timings"]["send"] = har.get_duration(start_sending_time, now)

        if total == -1:
            total = '?'
        self.log("Uploaded %d/%s of {url}" % (sent, total), self.sender(), min_level=3)

    def _getSplashRequest(self, request):
        return self._getWebPageAttribute(request, 'splash_request')

    def log(self, msg, reply=None, min_level=2):
        if self.verbosity < min_level:
            return

        if not reply:
            url = ''
        else:
            url = qurl2ascii(reply.url())
            if not url:
                return

        msg = msg.format(url=url)
        log.msg(msg, system='network2')

class SplashQNetworkAccessManager(ProxiedQNetworkAccessManager):
    """
    This QNetworkAccessManager provides:

    * proxy support;
    * request middleware support;
    * additional logging.

    """
    def __init__(self, filters_path, allowed_schemes, verbosity):
        super(SplashQNetworkAccessManager, self).__init__(verbosity=verbosity)

        self.request_middlewares = []
        if self.verbosity >= 2:
            self.request_middlewares.append(RequestLoggingMiddleware())

        if allowed_schemes:
            self.request_middlewares.append(
                AllowedSchemesMiddleware(allowed_schemes, verbosity=verbosity)
            )

        self.request_middlewares.append(AllowedDomainsMiddleware(verbosity=verbosity))

        if filters_path is not None:
            self.adblock_rules = AdblockRulesRegistry(filters_path, verbosity=verbosity)
            self.request_middlewares.append(
                AdblockMiddleware(self.adblock_rules, verbosity=verbosity)
            )
        else:
            self.adblock_rules = None

    def createRequest(self, operation, request, outgoingData=None):
        splash_request = self._getSplashRequest(request)
        if splash_request:
            for filter in self.request_middlewares:
                request = filter.process(request, splash_request, operation, outgoingData)
        return super(SplashQNetworkAccessManager, self).createRequest(operation, request, outgoingData)

    def unknownFilters(self, filter_names):
        names = [f for f in filter_names.split(',') if f]
        if self.adblock_rules is None:
            return names
        return [name for name in names
                if not (self.adblock_rules.filter_is_known(name) or name=='none')]
<commit_msg>show send time as 0 if it is insignificant<commit_after># -*- coding: utf-8 -*-
from __future__ import absolute_import
from datetime import datetime
from contextlib import contextmanager

from PyQt4.QtNetwork import (
    QNetworkAccessManager, QNetworkProxyQuery,
    QNetworkReply, QNetworkRequest
)
from PyQt4.QtWebKit import QWebFrame
from twisted.python import log

from splash.qtutils import qurl2ascii, OPERATION_NAMES
from splash import har
from splash.request_middleware import (
    AdblockMiddleware,
    AllowedDomainsMiddleware,
    AllowedSchemesMiddleware,
    RequestLoggingMiddleware,
    AdblockRulesRegistry,
)

# See: http://pyqt.sourceforge.net/Docs/PyQt4/qnetworkreply.html#NetworkError-enum
REQUEST_ERRORS = {
    QNetworkReply.NoError : 'no error condition. Note: When the HTTP protocol returns a redirect no error will be reported. You can check if there is a redirect with the QNetworkRequest::RedirectionTargetAttribute attribute.',
    QNetworkReply.ConnectionRefusedError : 'the remote server refused the connection (the server is not accepting requests)',
    QNetworkReply.RemoteHostClosedError : 'the remote server closed the connection prematurely, before the entire reply was received and processed',
    QNetworkReply.HostNotFoundError : 'the remote host name was not found (invalid hostname)',
    QNetworkReply.TimeoutError : 'the connection to the remote server timed out',
    QNetworkReply.OperationCanceledError : 'the operation was canceled via calls to abort() or close() before it was finished.',
    QNetworkReply.SslHandshakeFailedError : 'the SSL/TLS handshake failed and the encrypted channel could not be established. The sslErrors() signal should have been emitted.',
    QNetworkReply.TemporaryNetworkFailureError : 'the connection was broken due to disconnection from the network, however the system has initiated roaming to another access point. The request should be resubmitted and will be processed as soon as the connection is re-established.',
    QNetworkReply.ProxyConnectionRefusedError : 'the connection to the proxy server was refused (the proxy server is not accepting requests)',
    QNetworkReply.ProxyConnectionClosedError : 'the proxy server closed the connection prematurely, before the entire reply was received and processed',
    QNetworkReply.ProxyNotFoundError : 'the proxy host name was not found (invalid proxy hostname)',
    QNetworkReply.ProxyTimeoutError : 'the connection to the proxy timed out or the proxy did not reply in time to the request sent',
    QNetworkReply.ProxyAuthenticationRequiredError : 'the proxy requires authentication in order to honour the request but did not accept any credentials offered (if any)',
    QNetworkReply.ContentAccessDenied : 'the access to the remote content was denied (similar to HTTP error 401)',
    QNetworkReply.ContentOperationNotPermittedError : 'the operation requested on the remote content is not permitted',
    QNetworkReply.ContentNotFoundError : 'the remote content was not found at the server (similar to HTTP error 404)',
    QNetworkReply.AuthenticationRequiredError : 'the remote server requires authentication to serve the content but the credentials provided were not accepted (if any)',
    QNetworkReply.ContentReSendError : 'the request needed to be sent again, but this failed for example because the upload data could not be read a second time.',
    QNetworkReply.ProtocolUnknownError : 'the Network Access API cannot honor the request because the protocol is not known',
    QNetworkReply.ProtocolInvalidOperationError : 'the requested operation is invalid for this protocol',
    QNetworkReply.UnknownNetworkError : 'an unknown network-related error was detected',
    QNetworkReply.UnknownProxyError : 'an unknown proxy-related error was detected',
    QNetworkReply.UnknownContentError : 'an unknown error related to the remote content was detected',
    QNetworkReply.ProtocolFailure : 'a breakdown in protocol was detected (parsing error, invalid or unexpected responses, etc.)',
}

class ProxiedQNetworkAccessManager(QNetworkAccessManager):
    """
    This QNetworkAccessManager subclass enables "splash proxy factories"
    support. Qt provides similar functionality via setProxyFactory method,
    but standard QNetworkProxyFactory is not flexible enough.

    It also sets up some extra logging, provides a way to get
    the "source" request (that was made to Splash itself) and tracks
    information about requests/responses.
    """

    _REQUEST_ID = QNetworkRequest.User + 1

    REQUEST_CREATED = "created"
    REQUEST_FINISHED = "finished"
    REQUEST_HEADERS_RECEIVED = "headers"

    def __init__(self, verbosity):
        super(ProxiedQNetworkAccessManager, self).__init__()
        self.sslErrors.connect(self._sslErrors)
        self.finished.connect(self._finished)
        self.verbosity = verbosity
        self._next_id = 0

        assert self.proxyFactory() is None, "Standard QNetworkProxyFactory is not supported"

    def _sslErrors(self, reply, errors):
        reply.ignoreSslErrors()

    def _finished(self, reply):
        reply.deleteLater()

    def createRequest(self, operation, request, outgoingData=None):
        """
        This method is called when a new request is sent;
        it must return a reply object to work with.
        """
        start_time = datetime.utcnow()

        request = self._wrapRequest(request)
        har_entry = self._harEntry(request, create=True)
        if har_entry is not None:
            page_id = self._getWebPageAttribute(request, "page_id")
            if outgoingData is None:
                bodySize = -1
            else:
                bodySize = outgoingData.size()
            har_entry.update({
                '_tmp': {
                    'start_time': start_time,
                    'request_start_sending_time': start_time,
                    'request_sent_time': start_time,

                    # 'outgoingData': outgoingData,
                    'state': self.REQUEST_CREATED,
                },
                "pageref": page_id,
                "startedDateTime": har.format_datetime(start_time),
                "request": {
                    "method": OPERATION_NAMES.get(operation, '?'),
                    "url": unicode(request.url().toString()),
                    "httpVersion": "HTTP/1.1",
                    "cookies": har.request_cookies2har(request),
                    "queryString": har.querystring2har(request.url()),
                    "headers": har.headers2har(request),

                    "headersSize" : har.headers_size(request),
                    "bodySize": bodySize,
                },
                "response": {},
                "cache": {},
                "timings": {
                    "send": 0,
                    "wait": 0,
                    "receive": 0,
                },
            })

        with self._proxyApplied(request):
            reply = super(ProxiedQNetworkAccessManager, self).createRequest(
                operation, request, outgoingData
            )
            reply.error.connect(self._handleError)
            reply.finished.connect(self._handleFinished)
            reply.metaDataChanged.connect(self._handleMetaData)
            reply.downloadProgress.connect(self._handleDownloadProgress)

        return reply

    @contextmanager
    def _proxyApplied(self, request):
        """
        This context manager temporary sets a proxy based on request options.
        """
        old_proxy = self.proxy()
        splash_proxy_factory = self._getSplashProxyFactory(request)
        if splash_proxy_factory:
            proxy_query = QNetworkProxyQuery(request.url())
            proxy = splash_proxy_factory.queryProxy(proxy_query)[0]
            self.setProxy(proxy)
        try:
            yield
        finally:
            self.setProxy(old_proxy)

    def _wrapRequest(self, request):
        request = QNetworkRequest(request)
        request.setAttribute(self._REQUEST_ID, self._next_id)
        self._next_id += 1
        return request

    def _getRequestId(self, request=None):
        if request is None:
            request = self.sender().request()
        return request.attribute(self._REQUEST_ID).toPyObject()

    def _harEntry(self, request=None, create=False):
        """
        Return a mutable dictionary for request/response
        information storage.
        """
        if request is None:
            request = self.sender().request()

        entries = self._getWebPageAttribute(request, 'network_entries_map')
        if entries is None:
            return

        req_id = self._getRequestId(request)
        if create:
            assert req_id not in entries
            entries[req_id] = {
                "_idx": req_id,
            }
        return entries[req_id]

    def _getSplashProxyFactory(self, request):
        return self._getWebPageAttribute(request, 'splash_proxy_factory')

    def _getWebPageAttribute(self, request, attribute):
        web_frame = request.originatingObject()
        if isinstance(web_frame, QWebFrame):
            return getattr(web_frame.page(), attribute, None)

    def _handleError(self, error_id):
        error_msg = REQUEST_ERRORS.get(error_id, 'unknown error')
        self.log("Download error %d: %s ({url})" % (error_id, error_msg), self.sender(), min_level=1)

    def _handleFinished(self):
        reply = self.sender()
        har_entry = self._harEntry()
        if har_entry is not None:
            har_entry["_tmp"]["state"] = self.REQUEST_FINISHED

            now = datetime.utcnow()
            start_time = har_entry['_tmp']['start_time']
            response_start_time = har_entry['_tmp']['response_start_time']

            receive_time = har.get_duration(response_start_time, now)
            total_time = har.get_duration(start_time, now)

            har_entry["timings"]["receive"] = receive_time
            har_entry["time"] = total_time

            if not har_entry["timings"]["send"]:
                wait_time = har_entry["timings"]["wait"]
                har_entry["timings"]["send"] = total_time - receive_time - wait_time
                if har_entry["timings"]["send"] < 1e-6:
                    har_entry["timings"]["send"] = 0

            har_entry["response"].update(har.reply2har(reply))
            if 'bodySize' not in har_entry["response"]:
                har_entry["response"]["bodySize"] = -1

        self.log("Finished downloading {url}", reply)

    def _handleMetaData(self):
        reply = self.sender()
        har_entry = self._harEntry()
        if har_entry is not None:
            har_entry["_tmp"]["state"] = self.REQUEST_HEADERS_RECEIVED
            har_entry["response"].update(har.reply2har(reply))

            now = datetime.utcnow()
            request_sent = har_entry["_tmp"]["request_sent_time"]
            har_entry["_tmp"]["response_start_time"] = now
            har_entry["timings"]["wait"] = har.get_duration(request_sent, now)

        self.log("Headers received for {url}", reply, min_level=3)

    def _handleDownloadProgress(self, received, total):
        har_entry = self._harEntry()
        har_entry["response"]["bodySize"] = int(received)

        if total == -1:
            total = '?'
        self.log("Downloaded %d/%s of {url}" % (received, total), self.sender(), min_level=3)

    def _handleUploadProgress(self, sent, total):
        har_entry = self._harEntry()
        har_entry["request"]["bodySize"] = int(sent)

        now = datetime.utcnow()
        if sent == 0:
            # it is a moment the sending is started
            start_time = har_entry["_tmp"]["request_start_time"]
            har_entry["_tmp"]["request_start_sending_time"] = now
            har_entry["timings"]["blocked"] = har.get_duration(start_time, now)

        har_entry["_tmp"]["request_sent_time"] = now

        if sent == total:
            start_sending_time = har_entry["_tmp"]["request_start_sending_time"]
            har_entry["timings"]["send"] = har.get_duration(start_sending_time, now)

        if total == -1:
            total = '?'
        self.log("Uploaded %d/%s of {url}" % (sent, total), self.sender(), min_level=3)

    def _getSplashRequest(self, request):
        return self._getWebPageAttribute(request, 'splash_request')

    def log(self, msg, reply=None, min_level=2):
        if self.verbosity < min_level:
            return

        if not reply:
            url = ''
        else:
            url = qurl2ascii(reply.url())
            if not url:
                return

        msg = msg.format(url=url)
        log.msg(msg, system='network2')

class SplashQNetworkAccessManager(ProxiedQNetworkAccessManager):
    """
    This QNetworkAccessManager provides:

    * proxy support;
    * request middleware support;
    * additional logging.

    """
    def __init__(self, filters_path, allowed_schemes, verbosity):
        super(SplashQNetworkAccessManager, self).__init__(verbosity=verbosity)

        self.request_middlewares = []
        if self.verbosity >= 2:
            self.request_middlewares.append(RequestLoggingMiddleware())

        if allowed_schemes:
            self.request_middlewares.append(
                AllowedSchemesMiddleware(allowed_schemes, verbosity=verbosity)
            )

        self.request_middlewares.append(AllowedDomainsMiddleware(verbosity=verbosity))

        if filters_path is not None:
            self.adblock_rules = AdblockRulesRegistry(filters_path, verbosity=verbosity)
            self.request_middlewares.append(
                AdblockMiddleware(self.adblock_rules, verbosity=verbosity)
            )
        else:
            self.adblock_rules = None

    def createRequest(self, operation, request, outgoingData=None):
        splash_request = self._getSplashRequest(request)
        if splash_request:
            for filter in self.request_middlewares:
                request = filter.process(request, splash_request, operation, outgoingData)
        return super(SplashQNetworkAccessManager, self).createRequest(operation, request, outgoingData)

    def unknownFilters(self, filter_names):
        names = [f for f in filter_names.split(',') if f]
        if self.adblock_rules is None:
            return names
        return [name for name in names
                if not (self.adblock_rules.filter_is_known(name) or name=='none')]
<|endoftext|>


<commit_before>from .core import HttpLocust, Locust, TaskSet, task
from .event import Events
from .exception import InterruptTaskSet, ResponseError, RescheduleTaskImmediately
from .sequential_taskset import SequentialTaskSet
from .wait_time import between, constant, constant_pacing

events = Events()

__version__ = "0.14.5"
<commit_msg>Bump locust version to 1.0, in preparation of the release.<commit_after>from .core import HttpLocust, Locust, TaskSet, task
from .event import Events
from .exception import InterruptTaskSet, ResponseError, RescheduleTaskImmediately
from .sequential_taskset import SequentialTaskSet
from .wait_time import between, constant, constant_pacing

events = Events()

__version__ = "1.0.0"
<|endoftext|>


<commit_before>from collections import namedtuple

W = 800
H = 600
RED = 200, 0, 0
BLUE = 0, 0, 200

ball = Rect((W/2, H/2), (30, 30))
Direction = namedtuple('Direction', 'x y')
ball_dir = Direction(5, 5)

bat = Rect((W/2, 0.96 * H), (150, 15))

def draw():
    screen.clear()
    screen.draw.filled_rect(ball, RED)
    screen.draw.filled_rect(bat, RED)

def on_key_down():
    import sys
    sys.exit()

def on_mouse_move(pos):
    x, y = pos
    bat.center = (x, bat.center[1])

def move(ball):
    """returns a new ball at a new position
    """
    global ball_dir
    ball.move_ip(ball_dir)

def boubce():
    if ball.x > W or ball.x <= 0:
        ball_dir = Direction(-1 * ball_dir.x, ball_dir.y)

    if ball.y > H or ball.y <= 0:
        ball_dir = Direction(ball_dir.x, ball_dir.y * -1)

def update():
    move(ball)
<commit_msg>bat collide fix<commit_after>from collections import namedtuple
import sys

W = 800
H = 600
RED = 200, 0, 0
BLUE = 0, 0, 200

ball = Rect((W/2, H/2), (30, 30))
Direction = namedtuple('Direction', 'x y')
ball_dir = Direction(5, -5)

bat = Rect((W/2, 0.96 * H), (150, 15))

def draw():
    screen.clear()
    screen.draw.filled_rect(ball, RED)
    screen.draw.filled_rect(bat, RED)

def on_key_down():
    sys.exit()

def on_mouse_move(pos):
    x, y = pos
    bat.center = (x, bat.center[1])

def move(ball):
    """returns a new ball at a new position
    """
    global ball_dir
    ball.move_ip(ball_dir)

    if ball.x > W or ball.x <= 0:
        ball_dir = Direction(-1 * ball_dir.x, ball_dir.y)

    if ball.y <= 0:
        ball_dir = Direction(ball_dir.x, ball_dir.y * -1)

    if ball.colliderect(bat):
        ball_dir = Direction(ball_dir.x, - abs(ball_dir.y))

    if ball.y > H:
        sys.exit()

def update():
    move(ball)
<|endoftext|>


<commit_before>import imp
import subprocess
import os
import shutil

import requests
from jinja2 import Environment, DebugUndefined, meta

ALLOWED_TYPES = [
    type(True), type(0), type(0.0), type(0l), type(0+0j), type(''), type(u''),
    type([]), type((1,)), type(set()), type({})]

def __walk(root, ignore):
    """Replicates :func:`os.walk` but filters out any files or folders whose name is found in *ignore*.
    """
    for dirpath, dirnames, filenames in os.walk(root):
        dirnames[:] = filter(lambda dn: dn not in ignore, dirnames)
        filenames = filter(lambda fn: fn not in ignore, filenames)
        yield dirpath, dirnames, filenames

def __filter_copy(path):
    """Jinja2 filter that reads the contents of the file at *path* and returns it.
    """
    with open(path, 'rb') as fd:
        return fd.read()

def __filter_download(url):
    """Jinja2 filter that downloads the file at *url* and returns its contents.
    """
    return requests.get(url).content

def __make_env():
    """Creates a :class:`jinja2.Environment` and prepares it to be used by jingerly.
    We use a :class:`jinja2.DebugUndefined` undefined object so that when template variables are not provided they are left alone. This means we can use jingerly to template projects that use Jinja2
    """
    env = Environment(undefined=DebugUndefined)
    env.filters['copy'] = __filter_copy
    env.filters['download'] = __filter_download
    return env

def __make_variables(template_dir, output_dir, kwargs):
    """Returns a dictionary of the variables that will be passed into the template.
    Combines *template_dir* and *output_dir* with *kwargs* as 'IN' and 'OUT' repectively.
    Then processes *output_dir/jingerly.env* (if it exists).
    """
    variables = kwargs.copy()
    variables['IN'] = template_dir
    variables['OUT'] = output_dir
    env_path = os.path.join(template_dir, 'jingerly.env')
    if os.path.isfile(env_path):
        env = imp.load_source('module.name', env_path)
        for name in dir(env):
            if name.startswith('__'):
                continue
            attr = getattr(env, name)
            if type(attr) in ALLOWED_TYPES:
                variables[name] = attr
    return variables

def __make_renderer(env, variables):
    """Returns a function that takes some text and creates and renders a template from it using *env* and *variables*
    """
    def renderer(text):
        try:
            template = env.from_string(text)
        except UnicodeDecodeError as ude:
            print text
            raise ude
        return template.render(**variables)
    return renderer

def __process_files(root, files, renderer):
    """Pass the contents of each file in *files* to *renderer*, changes the file name if needed.
    """
    for f in files:
        file_path = os.path.join(root, f)
        file_name = renderer(f)
        with open(file_path, 'rb') as fd:
            file_contents = renderer(fd.read())
        if file_name != f:
            os.remove(file_path)
            file_path = os.path.join(root, file_name)
        with open(file_path, 'wb') as fd:
            fd.write(file_contents)

def __process_dirs(root, dirs, renderer):
    """Changes any directory name in *dirs* that needs changing (rendering).
    """
    dir_names = []
    for d in dirs:
        dir_path = os.path.join(root, d)
        dir_name = renderer(d)
        dir_names.append(dir_name)
        if dir_name != d:
            shutil.move(dir_path, os.path.join(root, dir_name))
    return dir_names

def __run_script(script_path, renderer):
    """Renders the script at *script_path* then calls it"""
    if os.path.isfile(script_path):
        with open(script_path, 'rb') as fd:
            file_contents = renderer(fd.read())
        with open(script_path, 'wb') as fd:
            fd.write(file_contents)
        subprocess.call(script_path)

def __run_pre(output_dir, renderer):
    """Runs the `jingerly.pre` script"""
    script_path = os.path.join(output_dir, 'jingerly.pre')
    __run_script(script_path, renderer)

def __run_post(output_dir, renderer):
    """Runs the `jingerly.post` script"""
    script_path = os.path.join(output_dir, 'jingerly.post')
    __run_script(script_path, renderer)

def __clean(path):
    """Deletes the `jingerly` specific files, if they exist.
    """
    for name in ['jingerly.pre', 'jingerly.post', 'jingerly.env']:
        try:
            os.remove(os.path.join(path, name))
        except OSError:
            pass

def render(template_dir, output_dir, _ignore=None, **kwargs):
    """:func:`jingerly.__walk` through every file and directory in *template_dir*, render it using variable values in *kwargs*, and save the output in *output_dir*.
    The order things happen in is:

    1. Copy entire template to output directory
    2. Read variables from `jingerly.env`
    3. Run `jingerly.pre`
    4. Walk through output_dir processings files and directories
    5. Run `jingerly.post`
    6. Clean up any jingerly specific files
    """
    shutil.copytree(template_dir, output_dir)
    if _ignore is None:
        _ignore = ['.DS_Store', '.git']
    env = __make_env()
    variables = __make_variables(template_dir, output_dir, kwargs)
    renderer = __make_renderer(env, variables)
    __run_pre(output_dir, renderer)
    for root, dirs, files in __walk(output_dir, _ignore):
        __process_files(
            root, files, renderer)
        dirs[:] = __process_dirs(
            root, dirs, renderer)
    __run_post(output_dir, renderer)
    __clean(output_dir)

def find_variables(template_dir, _ignore=None):
    """:func:`jingerly.__walk` through every file and directory in *template_dir* returning a list of all of the variables.
    This is used for the interactive feature of the CLI.
    """
    if _ignore is None:
        _ignore = ['.DS_Store', '.git', 'jingerly.env']
    env = __make_env()
    known = __make_variables(template_dir, template_dir, {})
    unknown = set()
    for root, dirs, files in __walk(template_dir, _ignore):
        for f in files:
            ast = env.parse(f)
            unknown |= meta.find_undeclared_variables(ast)
            with open(os.path.join(root, f), 'rb') as fd:
                ast = env.parse(fd.read())
                unknown |= meta.find_undeclared_variables(ast)
        for d in dirs:
            ast = env.parse(d)
            unknown |= meta.find_undeclared_variables(ast)
    return filter(lambda n: n not in known, unknown)
<commit_msg>print file name too<commit_after>import imp
import subprocess
import os
import shutil

import requests
from jinja2 import Environment, DebugUndefined, meta

ALLOWED_TYPES = [
    type(True), type(0), type(0.0), type(0l), type(0+0j), type(''), type(u''),
    type([]), type((1,)), type(set()), type({})]

def __walk(root, ignore):
    """Replicates :func:`os.walk` but filters out any files or folders whose name is found in *ignore*.
    """
    for dirpath, dirnames, filenames in os.walk(root):
        dirnames[:] = filter(lambda dn: dn not in ignore, dirnames)
        filenames = filter(lambda fn: fn not in ignore, filenames)
        yield dirpath, dirnames, filenames

def __filter_copy(path):
    """Jinja2 filter that reads the contents of the file at *path* and returns it.
    """
    with open(path, 'rb') as fd:
        return fd.read()

def __filter_download(url):
    """Jinja2 filter that downloads the file at *url* and returns its contents.
    """
    return requests.get(url).content

def __make_env():
    """Creates a :class:`jinja2.Environment` and prepares it to be used by jingerly.
    We use a :class:`jinja2.DebugUndefined` undefined object so that when template variables are not provided they are left alone. This means we can use jingerly to template projects that use Jinja2
    """
    env = Environment(undefined=DebugUndefined)
    env.filters['copy'] = __filter_copy
    env.filters['download'] = __filter_download
    return env

def __make_variables(template_dir, output_dir, kwargs):
    """Returns a dictionary of the variables that will be passed into the template.
    Combines *template_dir* and *output_dir* with *kwargs* as 'IN' and 'OUT' repectively.
    Then processes *output_dir/jingerly.env* (if it exists).
    """
    variables = kwargs.copy()
    variables['IN'] = template_dir
    variables['OUT'] = output_dir
    env_path = os.path.join(template_dir, 'jingerly.env')
    if os.path.isfile(env_path):
        env = imp.load_source('module.name', env_path)
        for name in dir(env):
            if name.startswith('__'):
                continue
            attr = getattr(env, name)
            if type(attr) in ALLOWED_TYPES:
                variables[name] = attr
    return variables

def __make_renderer(env, variables):
    """Returns a function that takes some text and creates and renders a template from it using *env* and *variables*
    """
    def renderer(text):
        try:
            template = env.from_string(text)
        except UnicodeDecodeError as ude:
            print text
            raise ude
        return template.render(**variables)
    return renderer

def __process_files(root, files, renderer):
    """Pass the contents of each file in *files* to *renderer*, changes the file name if needed.
    """
    for f in files:
        file_path = os.path.join(root, f)
        file_name = renderer(f)
        with open(file_path, 'rb') as fd:
            try:
                file_contents = renderer(fd.read())
            except UnicodeDecodeError as ude:
                print file_path
                raise ude
        if file_name != f:
            os.remove(file_path)
            file_path = os.path.join(root, file_name)
        with open(file_path, 'wb') as fd:
            fd.write(file_contents)

def __process_dirs(root, dirs, renderer):
    """Changes any directory name in *dirs* that needs changing (rendering).
    """
    dir_names = []
    for d in dirs:
        dir_path = os.path.join(root, d)
        dir_name = renderer(d)
        dir_names.append(dir_name)
        if dir_name != d:
            shutil.move(dir_path, os.path.join(root, dir_name))
    return dir_names

def __run_script(script_path, renderer):
    """Renders the script at *script_path* then calls it"""
    if os.path.isfile(script_path):
        with open(script_path, 'rb') as fd:
            file_contents = renderer(fd.read())
        with open(script_path, 'wb') as fd:
            fd.write(file_contents)
        subprocess.call(script_path)

def __run_pre(output_dir, renderer):
    """Runs the `jingerly.pre` script"""
    script_path = os.path.join(output_dir, 'jingerly.pre')
    __run_script(script_path, renderer)

def __run_post(output_dir, renderer):
    """Runs the `jingerly.post` script"""
    script_path = os.path.join(output_dir, 'jingerly.post')
    __run_script(script_path, renderer)

def __clean(path):
    """Deletes the `jingerly` specific files, if they exist.
    """
    for name in ['jingerly.pre', 'jingerly.post', 'jingerly.env']:
        try:
            os.remove(os.path.join(path, name))
        except OSError:
            pass

def render(template_dir, output_dir, _ignore=None, **kwargs):
    """:func:`jingerly.__walk` through every file and directory in *template_dir*, render it using variable values in *kwargs*, and save the output in *output_dir*.
    The order things happen in is:

    1. Copy entire template to output directory
    2. Read variables from `jingerly.env`
    3. Run `jingerly.pre`
    4. Walk through output_dir processings files and directories
    5. Run `jingerly.post`
    6. Clean up any jingerly specific files
    """
    shutil.copytree(template_dir, output_dir)
    if _ignore is None:
        _ignore = ['.DS_Store', '.git']
    env = __make_env()
    variables = __make_variables(template_dir, output_dir, kwargs)
    renderer = __make_renderer(env, variables)
    __run_pre(output_dir, renderer)
    for root, dirs, files in __walk(output_dir, _ignore):
        __process_files(
            root, files, renderer)
        dirs[:] = __process_dirs(
            root, dirs, renderer)
    __run_post(output_dir, renderer)
    __clean(output_dir)

def find_variables(template_dir, _ignore=None):
    """:func:`jingerly.__walk` through every file and directory in *template_dir* returning a list of all of the variables.
    This is used for the interactive feature of the CLI.
    """
    if _ignore is None:
        _ignore = ['.DS_Store', '.git', 'jingerly.env']
    env = __make_env()
    known = __make_variables(template_dir, template_dir, {})
    unknown = set()
    for root, dirs, files in __walk(template_dir, _ignore):
        for f in files:
            ast = env.parse(f)
            unknown |= meta.find_undeclared_variables(ast)
            with open(os.path.join(root, f), 'rb') as fd:
                ast = env.parse(fd.read())
                unknown |= meta.find_undeclared_variables(ast)
        for d in dirs:
            ast = env.parse(d)
            unknown |= meta.find_undeclared_variables(ast)
    return filter(lambda n: n not in known, unknown)
<|endoftext|>