<commit_before># game_of_life.py
# Conway's game of life
import curses
import time
import random
import signal
import sys

def signal_handler(signal, frame):
    """ Make Ctrl-c exit close curses window """
    #logging.info("Caught SIGINT, closing")
    curses.endwin()
    sys.exit(0)

class World:
    """ World class """
    def __init__(self, size_x=10, size_y=10):
        #logging.info("Instantianting World")
        self.size_y = size_y
        self.size_x = size_x
        self.make_neigh_cache()
        self.zero()
        self.random(int(size_x * size_y * 1 / 3))

    def zero(self):
        """ Set the world to all zeros """
        world = {}
        self.world = world

    def random(self, num):
        #logging.info(
        #    "In World.random: num={0}, size_x={1}, size_y={2}".format(
        #        num, self.size_x, self.size_y))
        self.zero()
        for i in range(num):
            x = random.randint(0, self.size_x)
            y = random.randint(0, self.size_y)
            self.world[(x, y)] = 1

    def print_screen(self, world=False, start_pos=None, end_pos=None):
        """Print world to screen"""
        if start_pos and end_pos:
            # Print world from start_pos to end_pos
            start_x, start_y = start_pos
            end_x, end_y = end_pos
        else:
            # No positions specified. Print the whole world
            start_x, start_y = self.min_pos()
            end_x, end_y = self.max_pos()
        for y in range(start_y, end_y + 1):
            for x in range(start_x, end_x + 1):
                if self.world[(x, y)]:
                    print("#", end="")
                else:
                    print("_", end="")
            print()

    def print_curses(self, screen, start_pos, end_pos):
        """Print the world from start_pos to end_pos on ncurses screen"""
        start_x, start_y = start_pos
        end_x, end_y = end_pos
        offset_x = start_x - 1
        offset_y = start_y - 1
        line_of_lines = [''.join(['#' if self.world[(x,y)] else '-'
                                  for x in range(start_x, end_x + 1)
                                 ]) for y in range(start_y, end_y + 1)]
        # TODO Add colors?

        #for y in range(start_y, end_y + 1):
        #    for x in range(start_x, end_x + 1):
        #        if self.world.get((x, y), False):
        #            screen.addch(y - offset_y, x - offset_x, "#",
        #                         curses.color_pair(1))
        #        else:
        #            screen.addch(y - offset_y, x - offset_x, ".")
        for line_no, line in enumerate(line_of_lines):
            screen.addstr(line_no + 1, 1, line)
        screen.refresh()

    def min_pos(self):
        """ Return the top, left position with alive cell"""
        cells = list(self.world.keys())
        if not cells:
            # World empty
            return (0, 0)
        else:
            leftmost = cells[0][0]
            upmost = cells[0][1]
            for cell in cells:
                x, y = cell
                if x < leftmost:
                    leftmost = x
                if y < upmost:
                    upmost = y
        return (leftmost, upmost)

    def max_pos(self):
        """ Return the bottom, right position with alive cell"""
        cells = list(self.world.keys())
        if not cells:
            # World empty
            return (10, 10)
        else:
            rightmost = cells[0][0]
            downmost = cells[0][1]
            for cell in cells:
                x, y = cell
                if x > rightmost:
                    rightmost = x
                if y > downmost:
                    downmost = y
        return (rightmost, downmost)

    def _new_cell(self, cell, neighbours):
        """Returns the new cell based on how many alive neighbours there is"""
        if cell == 1:
            if neighbours < 2:
                return 0
            elif neighbours < 4:
                return 1
            else:
                return 0
        if cell == 0:
            if neighbours == 3:
                return 1
            else:
                return 0
    def make_neigh_cache(self):
        self._neigh_cache = [(x, y) for x in range(-1,2) for y in range(-1,2) if not
                            (x == 0 and y == 0)]
        self._makro_cell_cache = [(x, y) for x in range(-1,2) for y in
                                 range(-1,2)]

    def calculate_neighbours(self, pos):
        """ calculate the number of neighbours of cell pos """
        neighbours = 0
        pos_x, pos_y = pos
        for x, y in self._neigh_cache:
            neighbours += self.world[(pos_x + x, pos_y + y)]
        return neighbours

    def update_cell(self, new_world, pos):
        """ Update the cell at position pos """
        #logging.info("In World.update_cell: updating cell ({0}, {1})".format(
        #    pos[0], pos[1]
        #))
        neighbours = self.calculate_neighbours(pos)
        new_cell = self._new_cell(self.world[(pos[0], pos[1])], neighbours)
        if new_cell == 1:
            new_world[(pos)] = 1
        return new_world

    def update_neighbours(self, new_world, pos, updated_cells):
        """Update the cells in new_world around cell in position pos"""
        for x, y in self._makro_cell_cache:
            if (pos[0] + x, pos[1] + y) in updated_cells:
                # Cell already updated, continue
                continue
            else:
                # Cell not updated yet, updating
                new_world = self.update_cell(new_world, (pos[0] + x, pos[1] + y))
                updated_cells.add((pos[0] + x, pos[1] + y))
        return new_world, updated_cells

    def update(self):
        """ Update the current world one step """
        new_world = {}
        updated_cells = set()
        for pos, cell in self.world.items():
            new_world, updated_cells = self.update_neighbours(new_world, pos,
                                                              updated_cells)
        self.world = new_world

    def __str__(self):
        string = ""
        for y in range(self.size_y):
            for x in range(self.size_x):
                if self.world[(x, y)]:
                    string += "#"
                else:
                    string += "-"
            string += "\n"
        return string

class Game:
    def __init__(self, start=False, mode="screen", size_x=20, size_y=20,
                 max_size=False):

        self.size_x = size_x
        self.size_y = size_y
        self.mode = mode

        if self.mode == "curses":
            self.init_curses(max_size)

        self.world = World(self.size_x, self.size_y)
        self.top_corner = (0, 0)
        self.bottom_corner = (self.size_x, self.size_y)

    def init_curses(self, max_size):
        """ Initilize the curses screen """
        self.screen = curses.initscr()
        curses.start_color()
        if not curses.has_colors():
            self.exit("Error, no color support", 1)
        else:
            # We have color support
            # Pair 1, green on black background
            curses.init_pair(1, curses.COLOR_GREEN, curses.COLOR_BLACK)
        if max_size:
            # Get maximum size of screen
            size_y_max, size_x_max = self.screen.getmaxyx()
            # Need space for border
            self.size_y, self.size_x = size_y_max - 3 , size_x_max - 3

        self.screen.border(0)
        curses.noecho()
        curses.cbreak()
        self.screen.keypad(1)
        curses.curs_set(0)

    def ask_user(self, question):
        """ Ask the users a question, return answer as string """
        q_size = len(question)
        self.screen.addstr(self.size_y // 2, self.size_x // 2 - q_size // 2, question + ' ')
        answer = self.screen.getstr().decode("UTF-8")
        return answer

    def prompt(self):
        """ Wait for user input and update world:
            space: continue one generation
            w or up: move screen up
            a or left: move screen left
            s or down: move screen down
            f or right: move screen right
            r: ask user how many generations to run and then run them
        """
        answer = self.screen.getch()
        #logging.debug("Got char %d, symbol %c" % (answer, chr(answer)))
        if answer == ord('w') or answer == curses.KEY_UP:
            # Move world down
            self.top_corner = (self.top_corner[0], self.top_corner[1] - 1)
            self.bottom_corner = (self.bottom_corner[0],
                                  self.bottom_corner[1] - 1)
            self.print_world()
        elif answer == ord('s') or answer == curses.KEY_DOWN:
            # Move world up
            self.top_corner = (self.top_corner[0], self.top_corner[1] + 1)
            self.bottom_corner = (self.bottom_corner[0],
                                  self.bottom_corner[1] + 1)
            self.print_world()
        elif answer == ord('a') or answer == curses.KEY_LEFT:
            # Move world right
            self.top_corner = (self.top_corner[0] - 1, self.top_corner[1])
            self.bottom_corner = (self.bottom_corner[0] - 1,
                                  self.bottom_corner[1])
            self.print_world()
        elif answer == ord('d') or answer == curses.KEY_RIGHT:
            # Move world left
            self.top_corner = (self.top_corner[0] + 1, self.top_corner[1])
            self.bottom_corner = (self.bottom_corner[0] + 1,
                                  self.bottom_corner[1])
            self.print_world()
        elif answer == ord("r"):
            # Ask user how many generations to animate
            try:
                num = int(self.ask_user("How many generations"))
            except ValueError:
                self.ask_user("Enter integer please")
                return
            self.animate(num)

        elif answer == ord(" "):
            # Update world
            #logging.info("In Game.prompt: Updating world one generation")
            self.world.update()
            self.print_world()

        elif answer == ord("q"):
            # quit game
            #logging.info(
            #    "In Game.prompt: got 'q': game is quitting. World = {0}".format(
            #        self.world.world))
            self.exit("Quitting", 0)

        else:
            self.prompt()

    def print_world(self):
        if self.mode == "curses":
            self.world.print_curses(self.screen, self.top_corner,
                                    self.bottom_corner)
        elif self.mode == "screen":
            self.world.print_screen()
        else:
            pass

    def kill_screen(self):
        """ Kill the screen and cleanup """
        if self.mode == "curses":
            curses.endwin()

    def exit(self, msg=None, status=0):
        self.kill_screen()
        if msg:
            print(msg)
        if type(status) != type(int):
            status = 0
        sys.exit(status)

    def animate(self, steps, dt=0.2):
        """ Update and print the screen 'step' times"""
        #logging.info("Started animation")
        for i in range(steps):
            self.world.update()
            self.print_world()
            time.sleep(dt)
        #logging.info("Animation finished")

signal.signal(signal.SIGINT, signal_handler)
<commit_msg>Reverted successfully to dict<commit_after># game_of_life.py
# Conway's game of life
import curses
import time
import random
import signal
import sys

def signal_handler(signal, frame):
    """ Make Ctrl-c exit close curses window """
    curses.endwin()
    sys.exit(0)

class World:
    """ World class """
    def __init__(self, size_x=10, size_y=10):
        self.size_y = size_y
        self.size_x = size_x
        self.make_neigh_cache()
        self.zero()
        self.random(int(size_x * size_y * 1 / 3))

    def zero(self):
        """ Set the world to all zeros """
        world = {}
        self.world = world

    def random(self, num):
        self.zero()
        for i in range(num):
            x = random.randint(0, self.size_x)
            y = random.randint(0, self.size_y)
            self.world[(x, y)] = 1

    def print_screen(self, world=False, start_pos=None, end_pos=None):
        """Print world to screen"""
        if start_pos and end_pos:
            # Print world from start_pos to end_pos
            start_x, start_y = start_pos
            end_x, end_y = end_pos
        else:
            # No positions specified. Print the whole world
            start_x, start_y = self.min_pos()
            end_x, end_y = self.max_pos()
        for y in range(start_y, end_y + 1):
            for x in range(start_x, end_x + 1):
                if self.world.get((x, y), False):
                    print("#", end="")
                else:
                    print("_", end="")
            print()

    def print_curses(self, screen, start_pos, end_pos):
        """Print the world from start_pos to end_pos on ncurses screen"""
        start_x, start_y = start_pos
        end_x, end_y = end_pos
        offset_x = start_x - 1
        offset_y = start_y - 1
        line_of_lines = [''.join(['#' if self.world.get((x,y), False) else '-'
                                  for x in range(start_x, end_x + 1)
                                 ]) for y in range(start_y, end_y + 1)]
        # TODO Add colors?
        for line_no, line in enumerate(line_of_lines):
            screen.addstr(line_no + 1, 1, line)
        screen.refresh()

    def min_pos(self):
        """ Return the top, left position with alive cell"""
        cells = list(self.world.keys())
        if not cells:
            # World empty
            return (0, 0)
        else:
            leftmost = cells[0][0]
            upmost = cells[0][1]
            for cell in cells:
                x, y = cell
                if x < leftmost:
                    leftmost = x
                if y < upmost:
                    upmost = y
        return (leftmost, upmost)

    def max_pos(self):
        """ Return the bottom, right position with alive cell"""
        cells = list(self.world.keys())
        if not cells:
            # World empty
            return (10, 10)
        else:
            rightmost = cells[0][0]
            downmost = cells[0][1]
            for cell in cells:
                x, y = cell
                if x > rightmost:
                    rightmost = x
                if y > downmost:
                    downmost = y
        return (rightmost, downmost)

    def _new_cell(self, cell, neighbours):
        """Returns the new cell based on how many alive neighbours there is"""
        if cell == 1:
            if neighbours < 2:
                return 0
            elif neighbours < 4:
                return 1
            else:
                return 0
        if cell == 0:
            if neighbours == 3:
                return 1
            else:
                return 0
    def make_neigh_cache(self):
        self._neigh_cache = [(x, y) for x in range(-1,2) for y in range(-1,2) if not
                            (x == 0 and y == 0)]
        self._makro_cell_cache = [(x, y) for x in range(-1,2) for y in
                                 range(-1,2)]

    def calculate_neighbours(self, pos):
        """ calculate the number of neighbours of cell pos """
        neighbours = 0
        pos_x, pos_y = pos
        for x, y in self._neigh_cache:
            neighbours += self.world.get((pos_x + x, pos_y + y), 0)
        return neighbours

    def update_cell(self, new_world, pos):
        """ Update the cell at position pos """
        neighbours = self.calculate_neighbours(pos)
        new_cell = self._new_cell(self.world.get((pos[0], pos[1]), 0), neighbours)
        if new_cell == 1:
            new_world[(pos)] = 1
        return new_world

    def update_neighbours(self, new_world, pos, updated_cells):
        """Update the cells in new_world around cell in position pos"""
        for x, y in self._makro_cell_cache:
            if (pos[0] + x, pos[1] + y) in updated_cells:
                # Cell already updated, continue
                continue
            else:
                # Cell not updated yet, updating
                new_world = self.update_cell(new_world, (pos[0] + x, pos[1] + y))
                updated_cells.add((pos[0] + x, pos[1] + y))
        return new_world, updated_cells

    def update(self):
        """ Update the current world one step """
        new_world = {}
        updated_cells = set()
        for pos, cell in self.world.items():
            new_world, updated_cells = self.update_neighbours(new_world, pos,
                                                              updated_cells)
        self.world = new_world

    def __str__(self):
        string = ""
        for y in range(self.size_y):
            for x in range(self.size_x):
                if self.world.get((x, y), False):
                    string += "#"
                else:
                    string += "-"
            string += "\n"
        return string

class Game:
    def __init__(self, start=False, mode="screen", size_x=20, size_y=20,
                 max_size=False):

        self.size_x = size_x
        self.size_y = size_y
        self.mode = mode

        if self.mode == "curses":
            self.init_curses(max_size)

        self.world = World(self.size_x, self.size_y)
        self.top_corner = (0, 0)
        self.bottom_corner = (self.size_x, self.size_y)

    def init_curses(self, max_size):
        """ Initilize the curses screen """
        self.screen = curses.initscr()
        curses.start_color()
        if not curses.has_colors():
            self.exit("Error, no color support", 1)
        else:
            # We have color support
            # Pair 1, green on black background
            curses.init_pair(1, curses.COLOR_GREEN, curses.COLOR_BLACK)
        if max_size:
            # Get maximum size of screen
            size_y_max, size_x_max = self.screen.getmaxyx()
            # Need space for border
            self.size_y, self.size_x = size_y_max - 3 , size_x_max - 3

        self.screen.border(0)
        curses.noecho()
        curses.cbreak()
        self.screen.keypad(1)
        curses.curs_set(0)

    def ask_user(self, question):
        """ Ask the users a question, return answer as string """
        q_size = len(question)
        self.screen.addstr(self.size_y // 2, self.size_x // 2 - q_size // 2, question + ' ')
        answer = self.screen.getstr().decode("UTF-8")
        return answer

    def prompt(self):
        """ Wait for user input and update world:
            space: continue one generation
            w or up: move screen up
            a or left: move screen left
            s or down: move screen down
            f or right: move screen right
            r: ask user how many generations to run and then run them
        """
        answer = self.screen.getch()
        if answer == ord('w') or answer == curses.KEY_UP:
            # Move world down
            self.top_corner = (self.top_corner[0], self.top_corner[1] - 1)
            self.bottom_corner = (self.bottom_corner[0],
                                  self.bottom_corner[1] - 1)
            self.print_world()
        elif answer == ord('s') or answer == curses.KEY_DOWN:
            # Move world up
            self.top_corner = (self.top_corner[0], self.top_corner[1] + 1)
            self.bottom_corner = (self.bottom_corner[0],
                                  self.bottom_corner[1] + 1)
            self.print_world()
        elif answer == ord('a') or answer == curses.KEY_LEFT:
            # Move world right
            self.top_corner = (self.top_corner[0] - 1, self.top_corner[1])
            self.bottom_corner = (self.bottom_corner[0] - 1,
                                  self.bottom_corner[1])
            self.print_world()
        elif answer == ord('d') or answer == curses.KEY_RIGHT:
            # Move world left
            self.top_corner = (self.top_corner[0] + 1, self.top_corner[1])
            self.bottom_corner = (self.bottom_corner[0] + 1,
                                  self.bottom_corner[1])
            self.print_world()
        elif answer == ord("r"):
            # Ask user how many generations to animate
            try:
                num = int(self.ask_user("How many generations"))
            except ValueError:
                self.ask_user("Enter integer please")
                return
            self.animate(num)

        elif answer == ord(" "):
            # Update world
            self.world.update()
            self.print_world()

        elif answer == ord("q"):
            # quit game
            self.exit("Quitting", 0)

        else:
            self.prompt()

    def print_world(self):
        if self.mode == "curses":
            self.world.print_curses(self.screen, self.top_corner,
                                    self.bottom_corner)
        elif self.mode == "screen":
            self.world.print_screen()
        else:
            pass

    def kill_screen(self):
        """ Kill the screen and cleanup """
        if self.mode == "curses":
            curses.endwin()

    def exit(self, msg=None, status=0):
        self.kill_screen()
        if msg:
            print(msg)
        if type(status) != type(int):
            status = 0
        sys.exit(status)

    def animate(self, steps, dt=0.2):
        """ Update and print the screen 'step' times"""
        for i in range(steps):
            self.world.update()
            self.print_world()
            time.sleep(dt)

signal.signal(signal.SIGINT, signal_handler)
<|endoftext|>


<commit_before>import numpy as np

def regular_grid(ar_shape, n_points):
    """Find `n_points` regularly spaced along `ar_shape`.

    The returned points (as slices) should be as close to cubically-spaced as
    possible.

    Parameters
    ----------
    ar_shape : array-like of ints
        The shape of the space embedding the grid. `len(ar_shape)` is the
        number of dimensions.
    n_points : int
        The (approximate) number of points to embed in the space.

    Returns
    -------
    slices : list of slice objects
        A slice along each dimension of `ar_shape`, such that the intersection
        of all the slices give the coordinates of regularly spaced points.
    """
    ar_shape = np.asanyarray(ar_shape)
    ndim = len(ar_shape)
    unsort_dim_idxs = np.argsort(np.argsort(ar_shape))
    sorted_dims = np.sort(ar_shape)
    space_size = float(np.prod(ar_shape))
    if space_size <= n_points:
        return [slice(None)] * ndim
    stepsizes = (space_size / n_points) ** (1.0 / ndim) * np.ones(ndim)
    if (sorted_dims < stepsizes).any():
        for dim in range(ndim):
            stepsizes[dim] = sorted_dims[dim]
            space_size = float(np.prod(sorted_dims[dim+1:]))
            stepsizes[dim+1:] = ((space_size / n_points) **
                                            (1.0 / (ndim - dim - 1)))
            if (sorted_dims >= stepsizes).all():
                break
    starts = np.floor(stepsizes/2)
    stepsizes = np.round(stepsizes)
    slices = [slice(start, None, step) for
                                        start, step in zip(starts, stepsizes)]
    slices = [slices[i] for i in unsort_dim_idxs]
    return slices
<commit_msg>Improve PEP8 and Python 3 compliance<commit_after>import numpy as np

def regular_grid(ar_shape, n_points):
    """Find `n_points` regularly spaced along `ar_shape`.

    The returned points (as slices) should be as close to cubically-spaced as
    possible.

    Parameters
    ----------
    ar_shape : array-like of ints
        The shape of the space embedding the grid. `len(ar_shape)` is the
        number of dimensions.
    n_points : int
        The (approximate) number of points to embed in the space.

    Returns
    -------
    slices : list of slice objects
        A slice along each dimension of `ar_shape`, such that the intersection
        of all the slices give the coordinates of regularly spaced points.
    """
    ar_shape = np.asanyarray(ar_shape)
    ndim = len(ar_shape)
    unsort_dim_idxs = np.argsort(np.argsort(ar_shape))
    sorted_dims = np.sort(ar_shape)
    space_size = float(np.prod(ar_shape))
    if space_size <= n_points:
        return [slice(None)] * ndim
    stepsizes = (space_size / n_points) ** (1.0 / ndim) * np.ones(ndim)
    if (sorted_dims < stepsizes).any():
        for dim in range(ndim):
            stepsizes[dim] = sorted_dims[dim]
            space_size = float(np.prod(sorted_dims[dim+1:]))
            stepsizes[dim+1:] = ((space_size / n_points) **
                                            (1.0 / (ndim - dim - 1)))
            if (sorted_dims >= stepsizes).all():
                break
    starts = stepsizes // 2
    stepsizes = np.round(stepsizes)
    slices = [slice(start, None, step) for
                                        start, step in zip(starts, stepsizes)]
    slices = [slices[i] for i in unsort_dim_idxs]
    return slices
<|endoftext|>


<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-
##
## Author: Adriano Monteiro Marques <adriano@umitproject.org>
##
## Copyright (C) 2011 S2S Network Consultoria e Tecnologia da Informacao LTDA
##
## This program is free software: you can redistribute it and/or modify
## it under the terms of the GNU Affero General Public License as
## published by the Free Software Foundation, either version 3 of the
## License, or (at your option) any later version.
##
## This program is distributed in the hope that it will be useful,
## but WITHOUT ANY WARRANTY; without even the implied warranty of
## MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
## GNU Affero General Public License for more details.
##
## You should have received a copy of the GNU Affero General Public License
## along with this program.  If not, see <http://www.gnu.org/licenses/>.
##

import decimal

from django.db import models
from django.core.cache import cache
from django.conf import settings

from dbextra.fields import ListField
from geoip.ip import convert_ip, convert_int_ip

CACHE_EXPIRATION = 60*60 # 1 hour, since this doesn't change any often
LOCATION_CACHE_KEY = "location_%s"
PUBLIC_KEY_AGENT_CACHE_KEY = "pkey_agent_%s"
AGENT_CACHE_KEY = "agent_%s"
IP_RANGE_CACHE_KEY = "ip_range_%s"
IP_RANGES_CACHE_KEY = "ip_ranges_%s"
NETWORK_LIST_AGENTS_CACHE_KEY = "network_logged_agents_%s"
PREFIX_KEY = "region_prefix_%s"
CLOSEST_LOCATION_KEY='closest_location_%s_%s'
CLOSEST_LOCATIONS_KEY='closes_locationss_%s_%s'

BAN_FLAGS = dict(
    abuse=1,
    robot=2,
    other=4
)

class LocationAggregation(models.Model):
    lat = models.DecimalField(decimal_places=20, max_digits=23) # Base Latitude
    lon = models.DecimalField(decimal_places=20, max_digits=23) # Base Longitude
    locations = ListField(py_type=int)

    def add(self, location):
        if location.id in self.locations:
            return

        self.locations.append(location.id)
        self.save()

    @staticmethod
    def add_location(location):
        """Is going to add the location to the aggregation ranges so we can link
        locations that are close to each other.
        """
        lat = int(location.lat)
        lon = int(location.lon)

        aggs = LocationAggregation.objects.filter(lat=lat, lon=lon)

        if not aggs:
            # There are no aggregations for such region. Let's create them now then
            agg = LocationAggregation()
            agg.lat = int(location.lat)
            agg.lon = int(location.lon)
            agg.save()

            agg.add(location)
            location.add_aggregation(agg)
        else:
            agg = aggs[0]
            agg.add(location)
            location.add_aggregation(agg)

class LocationNamesAggregation(models.Model):
    prefix = models.CharField(max_length=200)
    names = ListField()
    locations = ListField(py_type=int)

    @staticmethod
    def add_location(location):
        for i in xrange(1, len(location.fullname)):
            prefix = location.fullname[:i].lower()

            agg = LocationNamesAggregation.objects.filter(prefix=prefix)

            if agg:
                agg = agg[0]
                if location.id in agg.locations:
                    continue

                agg.names.append(location.fullname)
                agg.locations.append(location.id)
            else:
                agg = LocationNamesAggregation()
                agg.prefix = prefix
                agg.names.append(location.fullname)
                agg.locations.append(location.id)

            agg.save()

    @staticmethod
    def get_names(prefix):
        prefix = prefix.lower()
        names = cache.get(PREFIX_KEY % prefix, [])
        if names:
            return names

        agg = LocationNamesAggregation.objects.filter(prefix=prefix)
        if agg:
            names = agg[0].names

        cache.set(PREFIX_KEY, names)
        return names

    def __unicode__(self):
        return self.prefix

class Location(models.Model):
    ip_range_ids = ListField(py_type=int)
    fullname = models.CharField(max_length=300, blank=True, null=True)
    country_name = models.CharField(max_length=100, blank=True, null=True)
    country_code = models.CharField(max_length=2)
    state_region = models.CharField(max_length=2)
    city = models.CharField(max_length=255)
    zipcode = models.CharField(max_length=6)
    lat = models.DecimalField(decimal_places=20, max_digits=23)
    lon = models.DecimalField(decimal_places=20, max_digits=23)
    aggregations = ListField(py_type=int)
    nodes_count = models.IntegerField(default=0)

    def add_aggregation(self, aggregation):
        if aggregation.id in self.aggregations:
            return

        self.aggregations.append(aggregation.id)
        self.save()

    def __unicode__(self):
        return "%s, %s" % (self.city, self.country_name) \
                                if self.city != '' else self.country_code

    @staticmethod
    def retrieve_location(name):
        """If location is empty we consider this suggestion to be world wide.
        If location doesn't exist in our datastore, we try to get its
        coordinates. If successful, we create a new entry. Otherwise, we
        consider the suggestion to be world wide.
        """
        location = Location.objects.filter(fullname__startswith=name)
        if not location:
            # TODO
            location = Location.objects.filter(fullname__startswith=name)

        if location:
            return location[0]

    @property
    def ip_ranges(self):
        key = IP_RANGES_CACHE_KEY % self.id
        ranges = cache.get(key, False)
        if not ranges:
            ranges = IPRange.objects.filter(location_id=self.id)
            cache.set(key, ranges, CACHE_EXPIRATION)
        return ranges

    @staticmethod
    def add_ip_range(ip_range):
        location = ip_range.location
        if location:
            if ip_range.id in location.ip_range_ids:
                return location

            location.ip_range_ids.append(ip_range.id)
            location.save()
            return location

        #location = Location()
        #location.ip_range_ids = [ip_range.id]
        #location.country_code = ip_range.country_code
        #location.state_region = ip_range.state_region
        #location.city = ip_range.city
        #location.zipcode = ip_range.zipcode
        #location.lat = ip_range.lat
        #location.lon = ip_range.lon
        #location.save()

        return location

    def save(self, *args, **kwargs):
        super(Location, self).save(*args, **kwargs)

    def generateAggregations(self):
        # Create proper aggregation for the current lat/lon
        LocationAggregation.add_location(self)
        LocationNamesAggregation.add_location(self)

    @staticmethod
    def closest_location(lat, lon):
        """Not necessarily the closest location, but one in the closest range.
        TODO: make sure we save the locations in order of distance inside the
        aggregations so that we can retrieve the closest location rather than a
        random one in the closest range.
        """
        location = cache.get(CLOSEST_LOCATION_KEY % (lat, lon), None)
        if location is None:
            aggs = LocationAggregation.objects.filter(lat=int(lat), lon=int(lon))[:1]
            region = Location.objects.get(pk=aggs.locations[0])
            cache.set(CLOSEST_LOCATION_KEY % (lat, lon), region)
        return region

    @staticmethod
    def closest_locations(lat, lon):
        locations = cache.get(CLOSEST_LOCATIONS_KEY % (lat, lon), None)
        if locations is None:
            aggs = LocationAggregation.objects.filter(lat=int(lat), lon=int(lon))
            locations = []
            for agg in aggs:
                locations += [Location.objects.get(pk=id) for id in agg.locations]

            cache.set(CLOSEST_LOCATIONS_KEY % (lat, lon), locations)

        return locations

try:
    UNKNOWN_LOCATION = Location.objects.get_or_create(fullname='Unknown',
                                                      country_name='Unknown',
                                                      country_code='UN',
                                                      state_region='UN',
                                                      city='Unknown',
                                                      zipcode='',
                                                      lat=decimal.Decimal('0.0'),
                                                      lon=decimal.Decimal('0.0'))[0]
except:
    UNKNOWN_LOCATION = None

class BannedNetworks(models.Model):
    """This is an aggregation of Banned Networks that is built from
    IPRange model. This is in order for making retrieval of ban list faster
    and cheaper.
    """
    location_id = models.IntegerField()
    iprange_id = models.IntegerField()
    start_number = models.IntegerField()
    end_number = models.IntegerField()
    nodes_count = models.IntegerField()
    flags = models.IntegerField()

    @property
    def location(self):
        key = LOCATION_CACHE_KEY % self.location_id
        location = cache.get(key, False)
        if not location:
            location = Location.objects.get(id=self.location_id)
            cache.set(key, location, CACHE_EXPIRATION)
        return location

    @property
    def iprange(self):
        key = IP_RANGE_CACHE_KEY % self.iprange_id
        iprange = cache.get(key, False)
        if not iprange:
            iprange = IPRange.objects.get(id=self.iprange_id)
            cache.set(key, iprange, CACHE_EXPIRATION)
        return iprange

class IPRange(models.Model):
    location_id = models.IntegerField()
    start_number = models.PositiveIntegerField(db_index=True) #need to index for MySQL
    end_number = models.PositiveIntegerField()
    nodes_count = models.IntegerField(default=0)
    banned = models.BooleanField(default=False)
    ban_flags = models.IntegerField(default=0)

    def __unicode__(self):
        return "%s - %s" % (convert_int_ip(self.start_number),
                            convert_int_ip(self.end_number))

    def ban(self, flags):
        self.banned = True
        import pdb; pdb.set_trace()
        if self.ban_flags is None or self.ban_flags == 0:
            self.ban_flags = flags
        else:
            self.ban_flags |= flags
        self.save()

        banet = BannedNetworks.objects.filter(iprange_id=self.id)
        if banet:
            banet.flags |= flags
        else:
            banet = BannedNetworks()
            banet.location_id = self.location_id
            banet.iprange_id = self.id
            banet.start_number = self.start_number
            banet.end_number = self.end_number
            banet.nodes_count = self.nodes_count # the amount of nodes at the moment of the ban
            banet.flags = flags

        banet.save()

    def unban(self):
        self.banned = False
        self.ban_flags = 0
        self.save()

        banet = BannedNetworks.objects.filter(iprange_id=self.id)
        if banet:
            banet.delete()

    @staticmethod
    def ip_location(ip):
        if type(ip) != type(0):
            ip = convert_ip(ip)

        try:
            return IPRange.objects.filter(
                start_number__lte=ip).order_by('-start_number')[0]

        except IndexError:
            return IPRange.objects.get_or_create(location_id=UNKNOWN_LOCATION.id,
                                                 start_number=ip,
                                                 end_number=ip)[0]

    @property
    def location(self):
        key = LOCATION_CACHE_KEY % self.location_id
        location = cache.get(key, False)
        if not location:
            location = Location.objects.get(id=self.location_id)
            cache.set(key, location, CACHE_EXPIRATION)
        return location

    @property
    def logged_agents(self):
        key = NETWORK_LIST_AGENTS_CACHE_KEY % self.id
        logged_agents = cache.get(key, False)
        if not logged_agents:
            count = settings.MAX_AGENTSLIST_RESPONSE
            from agents.models import LoggedAgent

            # TODO: Randomize results
            logged_agents = LoggedAgent.objects.filter(location_id=self.location_id)[:count]
            cache.set(key, logged_agents, CACHE_EXPIRATION)
        return logged_agents

    def save(self, *args, **kwargs):
        new = self.id is None

        super(IPRange, self).save(*args, **kwargs)

        #if new:
        Location.add_ip_range(self)

    def dump(self):
        return dict(
                    start_ip=convert_int_ip(self.start_number),
                    end_ip=convert_int_ip(self.end_number)
                    )

<commit_msg>Use PositiveIntegerField where possible.<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-
##
## Author: Adriano Monteiro Marques <adriano@umitproject.org>
##
## Copyright (C) 2011 S2S Network Consultoria e Tecnologia da Informacao LTDA
##
## This program is free software: you can redistribute it and/or modify
## it under the terms of the GNU Affero General Public License as
## published by the Free Software Foundation, either version 3 of the
## License, or (at your option) any later version.
##
## This program is distributed in the hope that it will be useful,
## but WITHOUT ANY WARRANTY; without even the implied warranty of
## MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
## GNU Affero General Public License for more details.
##
## You should have received a copy of the GNU Affero General Public License
## along with this program.  If not, see <http://www.gnu.org/licenses/>.
##

import decimal

from django.db import models
from django.core.cache import cache
from django.conf import settings

from dbextra.fields import ListField
from geoip.ip import convert_ip, convert_int_ip

CACHE_EXPIRATION = 60*60 # 1 hour, since this doesn't change any often
LOCATION_CACHE_KEY = "location_%s"
PUBLIC_KEY_AGENT_CACHE_KEY = "pkey_agent_%s"
AGENT_CACHE_KEY = "agent_%s"
IP_RANGE_CACHE_KEY = "ip_range_%s"
IP_RANGES_CACHE_KEY = "ip_ranges_%s"
NETWORK_LIST_AGENTS_CACHE_KEY = "network_logged_agents_%s"
PREFIX_KEY = "region_prefix_%s"
CLOSEST_LOCATION_KEY='closest_location_%s_%s'
CLOSEST_LOCATIONS_KEY='closes_locationss_%s_%s'

BAN_FLAGS = dict(
    abuse=1,
    robot=2,
    other=4
)

class LocationAggregation(models.Model):
    lat = models.DecimalField(decimal_places=20, max_digits=23) # Base Latitude
    lon = models.DecimalField(decimal_places=20, max_digits=23) # Base Longitude
    locations = ListField(py_type=int)

    def add(self, location):
        if location.id in self.locations:
            return

        self.locations.append(location.id)
        self.save()

    @staticmethod
    def add_location(location):
        """Is going to add the location to the aggregation ranges so we can link
        locations that are close to each other.
        """
        lat = int(location.lat)
        lon = int(location.lon)

        aggs = LocationAggregation.objects.filter(lat=lat, lon=lon)

        if not aggs:
            # There are no aggregations for such region. Let's create them now then
            agg = LocationAggregation()
            agg.lat = int(location.lat)
            agg.lon = int(location.lon)
            agg.save()

            agg.add(location)
            location.add_aggregation(agg)
        else:
            agg = aggs[0]
            agg.add(location)
            location.add_aggregation(agg)

class LocationNamesAggregation(models.Model):
    prefix = models.CharField(max_length=200)
    names = ListField()
    locations = ListField(py_type=int)

    @staticmethod
    def add_location(location):
        for i in xrange(1, len(location.fullname)):
            prefix = location.fullname[:i].lower()

            agg = LocationNamesAggregation.objects.filter(prefix=prefix)

            if agg:
                agg = agg[0]
                if location.id in agg.locations:
                    continue

                agg.names.append(location.fullname)
                agg.locations.append(location.id)
            else:
                agg = LocationNamesAggregation()
                agg.prefix = prefix
                agg.names.append(location.fullname)
                agg.locations.append(location.id)

            agg.save()

    @staticmethod
    def get_names(prefix):
        prefix = prefix.lower()
        names = cache.get(PREFIX_KEY % prefix, [])
        if names:
            return names

        agg = LocationNamesAggregation.objects.filter(prefix=prefix)
        if agg:
            names = agg[0].names

        cache.set(PREFIX_KEY, names)
        return names

    def __unicode__(self):
        return self.prefix

class Location(models.Model):
    ip_range_ids = ListField(py_type=int)
    fullname = models.CharField(max_length=300, blank=True, null=True)
    country_name = models.CharField(max_length=100, blank=True, null=True)
    country_code = models.CharField(max_length=2)
    state_region = models.CharField(max_length=2)
    city = models.CharField(max_length=255)
    zipcode = models.CharField(max_length=6)
    lat = models.DecimalField(decimal_places=20, max_digits=23)
    lon = models.DecimalField(decimal_places=20, max_digits=23)
    aggregations = ListField(py_type=int)
    nodes_count = models.PositiveIntegerField(default=0)

    def add_aggregation(self, aggregation):
        if aggregation.id in self.aggregations:
            return

        self.aggregations.append(aggregation.id)
        self.save()

    def __unicode__(self):
        return "%s, %s" % (self.city, self.country_name) \
                                if self.city != '' else self.country_code

    @staticmethod
    def retrieve_location(name):
        """If location is empty we consider this suggestion to be world wide.
        If location doesn't exist in our datastore, we try to get its
        coordinates. If successful, we create a new entry. Otherwise, we
        consider the suggestion to be world wide.
        """
        location = Location.objects.filter(fullname__startswith=name)
        if not location:
            # TODO
            location = Location.objects.filter(fullname__startswith=name)

        if location:
            return location[0]

    @property
    def ip_ranges(self):
        key = IP_RANGES_CACHE_KEY % self.id
        ranges = cache.get(key, False)
        if not ranges:
            ranges = IPRange.objects.filter(location_id=self.id)
            cache.set(key, ranges, CACHE_EXPIRATION)
        return ranges

    @staticmethod
    def add_ip_range(ip_range):
        location = ip_range.location
        if location:
            if ip_range.id in location.ip_range_ids:
                return location

            location.ip_range_ids.append(ip_range.id)
            location.save()
            return location

        #location = Location()
        #location.ip_range_ids = [ip_range.id]
        #location.country_code = ip_range.country_code
        #location.state_region = ip_range.state_region
        #location.city = ip_range.city
        #location.zipcode = ip_range.zipcode
        #location.lat = ip_range.lat
        #location.lon = ip_range.lon
        #location.save()

        return location

    def save(self, *args, **kwargs):
        super(Location, self).save(*args, **kwargs)

    def generateAggregations(self):
        # Create proper aggregation for the current lat/lon
        LocationAggregation.add_location(self)
        LocationNamesAggregation.add_location(self)

    @staticmethod
    def closest_location(lat, lon):
        """Not necessarily the closest location, but one in the closest range.
        TODO: make sure we save the locations in order of distance inside the
        aggregations so that we can retrieve the closest location rather than a
        random one in the closest range.
        """
        location = cache.get(CLOSEST_LOCATION_KEY % (lat, lon), None)
        if location is None:
            aggs = LocationAggregation.objects.filter(lat=int(lat), lon=int(lon))[:1]
            region = Location.objects.get(pk=aggs.locations[0])
            cache.set(CLOSEST_LOCATION_KEY % (lat, lon), region)
        return region

    @staticmethod
    def closest_locations(lat, lon):
        locations = cache.get(CLOSEST_LOCATIONS_KEY % (lat, lon), None)
        if locations is None:
            aggs = LocationAggregation.objects.filter(lat=int(lat), lon=int(lon))
            locations = []
            for agg in aggs:
                locations += [Location.objects.get(pk=id) for id in agg.locations]

            cache.set(CLOSEST_LOCATIONS_KEY % (lat, lon), locations)

        return locations

try:
    UNKNOWN_LOCATION = Location.objects.get_or_create(fullname='Unknown',
                                                      country_name='Unknown',
                                                      country_code='UN',
                                                      state_region='UN',
                                                      city='Unknown',
                                                      zipcode='',
                                                      lat=decimal.Decimal('0.0'),
                                                      lon=decimal.Decimal('0.0'))[0]
except:
    UNKNOWN_LOCATION = None

class BannedNetworks(models.Model):
    """This is an aggregation of Banned Networks that is built from
    IPRange model. This is in order for making retrieval of ban list faster
    and cheaper.
    """
    location_id = models.PositiveIntegerField()
    iprange_id = models.PositiveIntegerField()
    start_number = models.PositiveIntegerField()
    end_number = models.PositiveIntegerField()
    nodes_count = models.PositiveIntegerField()
    flags = models.IntegerField()

    @property
    def location(self):
        key = LOCATION_CACHE_KEY % self.location_id
        location = cache.get(key, False)
        if not location:
            location = Location.objects.get(id=self.location_id)
            cache.set(key, location, CACHE_EXPIRATION)
        return location

    @property
    def iprange(self):
        key = IP_RANGE_CACHE_KEY % self.iprange_id
        iprange = cache.get(key, False)
        if not iprange:
            iprange = IPRange.objects.get(id=self.iprange_id)
            cache.set(key, iprange, CACHE_EXPIRATION)
        return iprange

class IPRange(models.Model):
    location_id = models.PositiveIntegerField()
    start_number = models.PositiveIntegerField(db_index=True) #need to index for MySQL
    end_number = models.PositiveIntegerField()
    nodes_count = models.PositiveIntegerField(default=0)
    banned = models.BooleanField(default=False)
    ban_flags = models.IntegerField(default=0)

    def __unicode__(self):
        return "%s - %s" % (convert_int_ip(self.start_number),
                            convert_int_ip(self.end_number))

    def ban(self, flags):
        self.banned = True
        import pdb; pdb.set_trace()
        if self.ban_flags is None or self.ban_flags == 0:
            self.ban_flags = flags
        else:
            self.ban_flags |= flags
        self.save()

        banet = BannedNetworks.objects.filter(iprange_id=self.id)
        if banet:
            banet.flags |= flags
        else:
            banet = BannedNetworks()
            banet.location_id = self.location_id
            banet.iprange_id = self.id
            banet.start_number = self.start_number
            banet.end_number = self.end_number
            banet.nodes_count = self.nodes_count # the amount of nodes at the moment of the ban
            banet.flags = flags

        banet.save()

    def unban(self):
        self.banned = False
        self.ban_flags = 0
        self.save()

        banet = BannedNetworks.objects.filter(iprange_id=self.id)
        if banet:
            banet.delete()

    @staticmethod
    def ip_location(ip):
        if type(ip) != type(0):
            ip = convert_ip(ip)

        try:
            return IPRange.objects.filter(
                start_number__lte=ip).order_by('-start_number')[0]

        except IndexError:
            return IPRange.objects.get_or_create(location_id=UNKNOWN_LOCATION.id,
                                                 start_number=ip,
                                                 end_number=ip)[0]

    @property
    def location(self):
        key = LOCATION_CACHE_KEY % self.location_id
        location = cache.get(key, False)
        if not location:
            location = Location.objects.get(id=self.location_id)
            cache.set(key, location, CACHE_EXPIRATION)
        return location

    @property
    def logged_agents(self):
        key = NETWORK_LIST_AGENTS_CACHE_KEY % self.id
        logged_agents = cache.get(key, False)
        if not logged_agents:
            count = settings.MAX_AGENTSLIST_RESPONSE
            from agents.models import LoggedAgent

            # TODO: Randomize results
            logged_agents = LoggedAgent.objects.filter(location_id=self.location_id)[:count]
            cache.set(key, logged_agents, CACHE_EXPIRATION)
        return logged_agents

    def save(self, *args, **kwargs):
        new = self.id is None

        super(IPRange, self).save(*args, **kwargs)

        #if new:
        Location.add_ip_range(self)

    def dump(self):
        return dict(
                    start_ip=convert_int_ip(self.start_number),
                    end_ip=convert_int_ip(self.end_number)
                    )

<|endoftext|>


<commit_before>import codecs
import csv
import operator
import os.path
import re

import prettytable
import six
import sqlalchemy
import sqlparse

from .column_guesser import ColumnGuesserMixin

try:
    from pgspecial.main import PGSpecial
except ImportError:
    PGSpecial = None

def unduplicate_field_names(field_names):
    """Append a number to duplicate field names to make them unique. """
    res = []
    for k in field_names:
        if k in res:
            i = 1
            while k + '_' + str(i) in res:
                i += 1
            k += '_' + str(i)
        res.append(k)
    return res

class UnicodeWriter(object):
    """
    A CSV writer which will write rows to CSV file "f",
    which is encoded in the given encoding.
    """

    def __init__(self, f, dialect=csv.excel, encoding="utf-8", **kwds):
        # Redirect output to a queue
        self.queue = six.StringIO()
        self.writer = csv.writer(self.queue, dialect=dialect, **kwds)
        self.stream = f
        self.encoder = codecs.getincrementalencoder(encoding)()

    def writerow(self, row):
        if six.PY2:
            _row = [s.encode("utf-8") if hasattr(s, "encode") else s
                    for s in row]
        else:
            _row = row
        self.writer.writerow(_row)
        # Fetch UTF-8 output from the queue ...
        data = self.queue.getvalue()
        if six.PY2:
            data = data.decode("utf-8")
            # ... and reencode it into the target encoding
            data = self.encoder.encode(data)
        # write to the target stream
        self.stream.write(data)
        # empty queue
        self.queue.truncate(0)
        self.queue.seek(0)

    def writerows(self, rows):
        for row in rows:
            self.writerow(row)

class CsvResultDescriptor(object):
    """Provides IPython Notebook-friendly output for the feedback after a ``.csv`` called."""

    def __init__(self, file_path):
        self.file_path = file_path

    def __repr__(self):
        return 'CSV results at %s' % os.path.join(
            os.path.abspath('.'), self.file_path)

    def _repr_html_(self):
        return 'CSV results' % os.path.join('.', 'files',
                                                             self.file_path)

def _nonbreaking_spaces(match_obj):
    """
    Make spaces visible in HTML by replacing all `` `` with `` ``

    Call with a ``re`` match object.  Retain group 1, replace group 2
    with nonbreaking speaces.
    """
    spaces = ' ' * len(match_obj.group(2))
    return '%s%s' % (match_obj.group(1), spaces)

_cell_with_spaces_pattern = re.compile(r'()( {2,})')

class ResultSet(list, ColumnGuesserMixin):
    """
    Results of a SQL query.

    Can access rows listwise, or by string value of leftmost column.
    """

    def __init__(self, sqlaproxy, sql, config):
        self.keys = sqlaproxy.keys()
        self.sql = sql
        self.config = config
        self.limit = config.autolimit
        style_name = config.style
        self.style = prettytable.__dict__[style_name.upper()]
        if sqlaproxy.returns_rows:
            if self.limit:
                list.__init__(self, sqlaproxy.fetchmany(size=self.limit))
            else:
                list.__init__(self, sqlaproxy.fetchall())
            self.field_names = unduplicate_field_names(self.keys)
            self.pretty = PrettyTable(self.field_names, style=self.style)
            # self.pretty.set_style(self.style)
        else:
            list.__init__(self, [])
            self.pretty = None

    def _repr_html_(self):
        _cell_with_spaces_pattern = re.compile(r'()( {2,})')
        if self.pretty:
            self.pretty.add_rows(self)
            result = self.pretty.get_html_string()
            result = _cell_with_spaces_pattern.sub(_nonbreaking_spaces, result)
            if self.config.displaylimit and len(
                    self) > self.config.displaylimit:
                result = '%s\n%d rows, truncated to displaylimit of %d' % (
                    result, len(self), self.config.displaylimit)
            return result
        else:
            return None

    def __str__(self, *arg, **kwarg):
        self.pretty.add_rows(self)
        return str(self.pretty or '')

    def __getitem__(self, key):
        """
        Access by integer (row position within result set)
        or by string (value of leftmost column)
        """
        try:
            return list.__getitem__(self, key)
        except TypeError:
            result = [row for row in self if row[0] == key]
            if not result:
                raise KeyError(key)
            if len(result) > 1:
                raise KeyError('%d results for "%s"' % (len(result), key))
            return result[0]

    def dict(self):
        """Returns a single dict built from the result set

        Keys are column names; values are a tuple"""
        return dict(zip(self.keys, zip(*self)))

    def dicts(self):
        "Iterator yielding a dict for each row"
        for row in self:
            yield dict(zip(self.keys, row))

    def DataFrame(self):
        "Returns a Pandas DataFrame instance built from the result set."
        import pandas as pd
        frame = pd.DataFrame(self, columns=(self and self.keys) or [])
        return frame

    def pie(self, key_word_sep=" ", title=None, **kwargs):
        """Generates a pylab pie chart from the result set.

        ``matplotlib`` must be installed, and in an
        IPython Notebook, inlining must be on::

            %%matplotlib inline

        Values (pie slice sizes) are taken from the
        rightmost column (numerical values required).
        All other columns are used to label the pie slices.

        Parameters
        ----------
        key_word_sep: string used to separate column values
                      from each other in pie labels
        title: Plot title, defaults to name of value column

        Any additional keyword arguments will be passsed
        through to ``matplotlib.pylab.pie``.
        """
        self.guess_pie_columns(xlabel_sep=key_word_sep)
        import matplotlib.pylab as plt
        pie = plt.pie(self.ys[0], labels=self.xlabels, **kwargs)
        plt.title(title or self.ys[0].name)
        return pie

    def plot(self, title=None, **kwargs):
        """Generates a pylab plot from the result set.

        ``matplotlib`` must be installed, and in an
        IPython Notebook, inlining must be on::

            %%matplotlib inline

        The first and last columns are taken as the X and Y
        values.  Any columns between are ignored.

        Parameters
        ----------
        title: Plot title, defaults to names of Y value columns

        Any additional keyword arguments will be passsed
        through to ``matplotlib.pylab.plot``.
        """
        import matplotlib.pylab as plt
        self.guess_plot_columns()
        self.x = self.x or range(len(self.ys[0]))
        coords = reduce(operator.add, [(self.x, y) for y in self.ys])
        plot = plt.plot(*coords, **kwargs)
        if hasattr(self.x, 'name'):
            plt.xlabel(self.x.name)
        ylabel = ", ".join(y.name for y in self.ys)
        plt.title(title or ylabel)
        plt.ylabel(ylabel)
        return plot

    def bar(self, key_word_sep=" ", title=None, **kwargs):
        """Generates a pylab bar plot from the result set.

        ``matplotlib`` must be installed, and in an
        IPython Notebook, inlining must be on::

            %%matplotlib inline

        The last quantitative column is taken as the Y values;
        all other columns are combined to label the X axis.

        Parameters
        ----------
        title: Plot title, defaults to names of Y value columns
        key_word_sep: string used to separate column values
                      from each other in labels

        Any additional keyword arguments will be passsed
        through to ``matplotlib.pylab.bar``.
        """
        import matplotlib.pylab as plt
        self.guess_pie_columns(xlabel_sep=key_word_sep)
        plot = plt.bar(range(len(self.ys[0])), self.ys[0], **kwargs)
        if self.xlabels:
            plt.xticks(range(len(self.xlabels)), self.xlabels, rotation=45)
        plt.xlabel(self.xlabel)
        plt.ylabel(self.ys[0].name)
        return plot

    def csv(self, filename=None, **format_params):
        """Generate results in comma-separated form.  Write to ``filename`` if given.
           Any other parameters will be passed on to csv.writer."""
        if not self.pretty:
            return None  # no results
        self.pretty.add_rows(self)
        if filename:
            encoding = format_params.get('encoding', 'utf-8')
            if six.PY2:
                outfile = open(filename, 'wb')
            else:
                outfile = open(filename, 'w', newline='', encoding=encoding)
        else:
            outfile = six.StringIO()
        writer = UnicodeWriter(outfile, **format_params)
        writer.writerow(self.field_names)
        for row in self:
            writer.writerow(row)
        if filename:
            outfile.close()
            return CsvResultDescriptor(filename)
        else:
            return outfile.getvalue()

def interpret_rowcount(rowcount):
    if rowcount < 0:
        result = 'Done.'
    else:
        result = '%d rows affected.' % rowcount
    return result

class FakeResultProxy(object):
    """A fake class that pretends to behave like the ResultProxy from
    SqlAlchemy.
    """

    def __init__(self, cursor, headers):
        if cursor is None:
            cursor = []
            headers = []
        if isinstance(cursor, list):
            self.from_list(source_list=cursor)
        else:
            self.fetchall = cursor.fetchall
            self.fetchmany = cursor.fetchmany
            self.rowcount = cursor.rowcount
        self.keys = lambda: headers
        self.returns_rows = True

    def from_list(self, source_list):
        "Simulates SQLA ResultProxy from a list."

        self.fetchall = lambda: source_list
        self.rowcount = len(source_list)

        def fetchmany(size):
            pos = 0
            while pos < len(source_list):
                yield source_list[pos:pos+size]
                pos += size

        self.fetchmany = fetchmany

# some dialects have autocommit
# specific dialects break when commit is used:
_COMMIT_BLACKLIST_DIALECTS = ('mssql', 'clickhouse', 'teradata')

def _commit(conn, config):
    """Issues a commit, if appropriate for current config and dialect"""

    _should_commit = config.autocommit and all(
        dialect not in str(conn.dialect)
        for dialect in _COMMIT_BLACKLIST_DIALECTS)

    if _should_commit:
        try:
            conn.session.execute('commit')
        except sqlalchemy.exc.OperationalError:
            pass  # not all engines can commit

def run(conn, sql, config, user_namespace):
    if sql.strip():
        for statement in sqlparse.split(sql):
            first_word = sql.strip().split()[0].lower()
            if first_word == 'begin':
                raise Exception("ipython_sql does not support transactions")
            if first_word.startswith('\\') and 'postgres' in str(conn.dialect):
                if not PGSpecial:
                    raise ImportError('pgspecial not installed')
                pgspecial = PGSpecial()
                _, cur, headers, _ = pgspecial.execute(
                    conn.session.connection.cursor(), statement)[0]
                result = FakeResultProxy(cur, headers)
            else:
                txt = sqlalchemy.sql.text(statement)
                result = conn.session.execute(txt, user_namespace)
            _commit(conn=conn, config=config)
            if result and config.feedback:
                print(interpret_rowcount(result.rowcount))
        resultset = ResultSet(result, statement, config)
        if config.autopandas:
            return resultset.DataFrame()
        else:
            return resultset
        #returning only last result, intentionally
    else:
        return 'Connected: %s' % conn.name

class PrettyTable(prettytable.PrettyTable):
    def __init__(self, *args, **kwargs):
        self.row_count = 0
        self.displaylimit = None
        return super(PrettyTable, self).__init__(*args, **kwargs)

    def add_rows(self, data):
        if self.row_count and (data.config.displaylimit == self.displaylimit):
            return  # correct number of rows already present
        self.clear_rows()
        self.displaylimit = data.config.displaylimit
        if self.displaylimit == 0:
            self.displaylimit = None  # TODO: remove this to make 0 really 0
        if self.displaylimit in (None, 0):
            self.row_count = len(data)
        else:
            self.row_count = min(len(data), self.displaylimit)
        for row in data[:self.displaylimit]:
            self.add_row(row)
<commit_msg>Import reduce function from functools<commit_after>import codecs
import csv
import operator
import os.path
import re
from functools import reduce

import prettytable
import six
import sqlalchemy
import sqlparse

from .column_guesser import ColumnGuesserMixin

try:
    from pgspecial.main import PGSpecial
except ImportError:
    PGSpecial = None

def unduplicate_field_names(field_names):
    """Append a number to duplicate field names to make them unique. """
    res = []
    for k in field_names:
        if k in res:
            i = 1
            while k + '_' + str(i) in res:
                i += 1
            k += '_' + str(i)
        res.append(k)
    return res

class UnicodeWriter(object):
    """
    A CSV writer which will write rows to CSV file "f",
    which is encoded in the given encoding.
    """

    def __init__(self, f, dialect=csv.excel, encoding="utf-8", **kwds):
        # Redirect output to a queue
        self.queue = six.StringIO()
        self.writer = csv.writer(self.queue, dialect=dialect, **kwds)
        self.stream = f
        self.encoder = codecs.getincrementalencoder(encoding)()

    def writerow(self, row):
        if six.PY2:
            _row = [s.encode("utf-8") if hasattr(s, "encode") else s
                    for s in row]
        else:
            _row = row
        self.writer.writerow(_row)
        # Fetch UTF-8 output from the queue ...
        data = self.queue.getvalue()
        if six.PY2:
            data = data.decode("utf-8")
            # ... and reencode it into the target encoding
            data = self.encoder.encode(data)
        # write to the target stream
        self.stream.write(data)
        # empty queue
        self.queue.truncate(0)
        self.queue.seek(0)

    def writerows(self, rows):
        for row in rows:
            self.writerow(row)

class CsvResultDescriptor(object):
    """Provides IPython Notebook-friendly output for the feedback after a ``.csv`` called."""

    def __init__(self, file_path):
        self.file_path = file_path

    def __repr__(self):
        return 'CSV results at %s' % os.path.join(
            os.path.abspath('.'), self.file_path)

    def _repr_html_(self):
        return 'CSV results' % os.path.join('.', 'files',
                                                             self.file_path)

def _nonbreaking_spaces(match_obj):
    """
    Make spaces visible in HTML by replacing all `` `` with `` ``

    Call with a ``re`` match object.  Retain group 1, replace group 2
    with nonbreaking speaces.
    """
    spaces = ' ' * len(match_obj.group(2))
    return '%s%s' % (match_obj.group(1), spaces)

_cell_with_spaces_pattern = re.compile(r'()( {2,})')

class ResultSet(list, ColumnGuesserMixin):
    """
    Results of a SQL query.

    Can access rows listwise, or by string value of leftmost column.
    """

    def __init__(self, sqlaproxy, sql, config):
        self.keys = sqlaproxy.keys()
        self.sql = sql
        self.config = config
        self.limit = config.autolimit
        style_name = config.style
        self.style = prettytable.__dict__[style_name.upper()]
        if sqlaproxy.returns_rows:
            if self.limit:
                list.__init__(self, sqlaproxy.fetchmany(size=self.limit))
            else:
                list.__init__(self, sqlaproxy.fetchall())
            self.field_names = unduplicate_field_names(self.keys)
            self.pretty = PrettyTable(self.field_names, style=self.style)
            # self.pretty.set_style(self.style)
        else:
            list.__init__(self, [])
            self.pretty = None

    def _repr_html_(self):
        _cell_with_spaces_pattern = re.compile(r'()( {2,})')
        if self.pretty:
            self.pretty.add_rows(self)
            result = self.pretty.get_html_string()
            result = _cell_with_spaces_pattern.sub(_nonbreaking_spaces, result)
            if self.config.displaylimit and len(
                    self) > self.config.displaylimit:
                result = '%s\n%d rows, truncated to displaylimit of %d' % (
                    result, len(self), self.config.displaylimit)
            return result
        else:
            return None

    def __str__(self, *arg, **kwarg):
        self.pretty.add_rows(self)
        return str(self.pretty or '')

    def __getitem__(self, key):
        """
        Access by integer (row position within result set)
        or by string (value of leftmost column)
        """
        try:
            return list.__getitem__(self, key)
        except TypeError:
            result = [row for row in self if row[0] == key]
            if not result:
                raise KeyError(key)
            if len(result) > 1:
                raise KeyError('%d results for "%s"' % (len(result), key))
            return result[0]

    def dict(self):
        """Returns a single dict built from the result set

        Keys are column names; values are a tuple"""
        return dict(zip(self.keys, zip(*self)))

    def dicts(self):
        "Iterator yielding a dict for each row"
        for row in self:
            yield dict(zip(self.keys, row))

    def DataFrame(self):
        "Returns a Pandas DataFrame instance built from the result set."
        import pandas as pd
        frame = pd.DataFrame(self, columns=(self and self.keys) or [])
        return frame

    def pie(self, key_word_sep=" ", title=None, **kwargs):
        """Generates a pylab pie chart from the result set.

        ``matplotlib`` must be installed, and in an
        IPython Notebook, inlining must be on::

            %%matplotlib inline

        Values (pie slice sizes) are taken from the
        rightmost column (numerical values required).
        All other columns are used to label the pie slices.

        Parameters
        ----------
        key_word_sep: string used to separate column values
                      from each other in pie labels
        title: Plot title, defaults to name of value column

        Any additional keyword arguments will be passsed
        through to ``matplotlib.pylab.pie``.
        """
        self.guess_pie_columns(xlabel_sep=key_word_sep)
        import matplotlib.pylab as plt
        pie = plt.pie(self.ys[0], labels=self.xlabels, **kwargs)
        plt.title(title or self.ys[0].name)
        return pie

    def plot(self, title=None, **kwargs):
        """Generates a pylab plot from the result set.

        ``matplotlib`` must be installed, and in an
        IPython Notebook, inlining must be on::

            %%matplotlib inline

        The first and last columns are taken as the X and Y
        values.  Any columns between are ignored.

        Parameters
        ----------
        title: Plot title, defaults to names of Y value columns

        Any additional keyword arguments will be passsed
        through to ``matplotlib.pylab.plot``.
        """
        import matplotlib.pylab as plt
        self.guess_plot_columns()
        self.x = self.x or range(len(self.ys[0]))
        coords = reduce(operator.add, [(self.x, y) for y in self.ys])
        plot = plt.plot(*coords, **kwargs)
        if hasattr(self.x, 'name'):
            plt.xlabel(self.x.name)
        ylabel = ", ".join(y.name for y in self.ys)
        plt.title(title or ylabel)
        plt.ylabel(ylabel)
        return plot

    def bar(self, key_word_sep=" ", title=None, **kwargs):
        """Generates a pylab bar plot from the result set.

        ``matplotlib`` must be installed, and in an
        IPython Notebook, inlining must be on::

            %%matplotlib inline

        The last quantitative column is taken as the Y values;
        all other columns are combined to label the X axis.

        Parameters
        ----------
        title: Plot title, defaults to names of Y value columns
        key_word_sep: string used to separate column values
                      from each other in labels

        Any additional keyword arguments will be passsed
        through to ``matplotlib.pylab.bar``.
        """
        import matplotlib.pylab as plt
        self.guess_pie_columns(xlabel_sep=key_word_sep)
        plot = plt.bar(range(len(self.ys[0])), self.ys[0], **kwargs)
        if self.xlabels:
            plt.xticks(range(len(self.xlabels)), self.xlabels, rotation=45)
        plt.xlabel(self.xlabel)
        plt.ylabel(self.ys[0].name)
        return plot

    def csv(self, filename=None, **format_params):
        """Generate results in comma-separated form.  Write to ``filename`` if given.
           Any other parameters will be passed on to csv.writer."""
        if not self.pretty:
            return None  # no results
        self.pretty.add_rows(self)
        if filename:
            encoding = format_params.get('encoding', 'utf-8')
            if six.PY2:
                outfile = open(filename, 'wb')
            else:
                outfile = open(filename, 'w', newline='', encoding=encoding)
        else:
            outfile = six.StringIO()
        writer = UnicodeWriter(outfile, **format_params)
        writer.writerow(self.field_names)
        for row in self:
            writer.writerow(row)
        if filename:
            outfile.close()
            return CsvResultDescriptor(filename)
        else:
            return outfile.getvalue()

def interpret_rowcount(rowcount):
    if rowcount < 0:
        result = 'Done.'
    else:
        result = '%d rows affected.' % rowcount
    return result

class FakeResultProxy(object):
    """A fake class that pretends to behave like the ResultProxy from
    SqlAlchemy.
    """

    def __init__(self, cursor, headers):
        if cursor is None:
            cursor = []
            headers = []
        if isinstance(cursor, list):
            self.from_list(source_list=cursor)
        else:
            self.fetchall = cursor.fetchall
            self.fetchmany = cursor.fetchmany
            self.rowcount = cursor.rowcount
        self.keys = lambda: headers
        self.returns_rows = True

    def from_list(self, source_list):
        "Simulates SQLA ResultProxy from a list."

        self.fetchall = lambda: source_list
        self.rowcount = len(source_list)

        def fetchmany(size):
            pos = 0
            while pos < len(source_list):
                yield source_list[pos:pos+size]
                pos += size

        self.fetchmany = fetchmany

# some dialects have autocommit
# specific dialects break when commit is used:
_COMMIT_BLACKLIST_DIALECTS = ('mssql', 'clickhouse', 'teradata')

def _commit(conn, config):
    """Issues a commit, if appropriate for current config and dialect"""

    _should_commit = config.autocommit and all(
        dialect not in str(conn.dialect)
        for dialect in _COMMIT_BLACKLIST_DIALECTS)

    if _should_commit:
        try:
            conn.session.execute('commit')
        except sqlalchemy.exc.OperationalError:
            pass  # not all engines can commit

def run(conn, sql, config, user_namespace):
    if sql.strip():
        for statement in sqlparse.split(sql):
            first_word = sql.strip().split()[0].lower()
            if first_word == 'begin':
                raise Exception("ipython_sql does not support transactions")
            if first_word.startswith('\\') and 'postgres' in str(conn.dialect):
                if not PGSpecial:
                    raise ImportError('pgspecial not installed')
                pgspecial = PGSpecial()
                _, cur, headers, _ = pgspecial.execute(
                    conn.session.connection.cursor(), statement)[0]
                result = FakeResultProxy(cur, headers)
            else:
                txt = sqlalchemy.sql.text(statement)
                result = conn.session.execute(txt, user_namespace)
            _commit(conn=conn, config=config)
            if result and config.feedback:
                print(interpret_rowcount(result.rowcount))
        resultset = ResultSet(result, statement, config)
        if config.autopandas:
            return resultset.DataFrame()
        else:
            return resultset
        #returning only last result, intentionally
    else:
        return 'Connected: %s' % conn.name

class PrettyTable(prettytable.PrettyTable):
    def __init__(self, *args, **kwargs):
        self.row_count = 0
        self.displaylimit = None
        return super(PrettyTable, self).__init__(*args, **kwargs)

    def add_rows(self, data):
        if self.row_count and (data.config.displaylimit == self.displaylimit):
            return  # correct number of rows already present
        self.clear_rows()
        self.displaylimit = data.config.displaylimit
        if self.displaylimit == 0:
            self.displaylimit = None  # TODO: remove this to make 0 really 0
        if self.displaylimit in (None, 0):
            self.row_count = len(data)
        else:
            self.row_count = min(len(data), self.displaylimit)
        for row in data[:self.displaylimit]:
            self.add_row(row)
<|endoftext|>


<commit_before>import asynchat
import asyncore
import re
import socket
import time
import sys
import zlib
import ssl
from errno import ECONNRESET, ENOTCONN, ESHUTDOWN, ECONNABORTED, EWOULDBLOCK
import select
import threading
from . import adns
from warnings import warn

class SocketPanic (Exception): pass
class TimeOut (Exception): pass

class AsynConnect (asynchat.async_chat):
	ac_in_buffer_size = 4096
	ac_out_buffer_size = 4096
	zombie_timeout = 120
	keep_alive = 300

	def __init__ (self, address, lock = None, logger = None):
		self.address = address
		self.lock = lock
		self.logger = logger

		self.request_count = 0
		self.event_time = time.time ()
		self._cv = threading.Condition ()
		self.active = 0
		self.ready = None
		self.affluent = None
		self.initialize ()
		asynchat.async_chat.__init__ (self)

	def log (self, msg, logtype):
		if self.request is not None and hasattr (self.request, "log"):
			self.request.log (msg, logtype)
		elif self.logger:
			self.logger (msg, logtype)
		else:
			warn ("No logger")

	def trace (self):
		if self.request is not None and hasattr (self.request, "trace"):
			self.request.trace ()
		elif self.logger:
			self.logger.trace ()
		else:
			warn ("No logger for traceback")

	def duplicate (self):
		return self.__class__ (self.address, self.lock, self.logger)

	def initialize (self):
		self.request = None
		self.got_data = False
		self.close_it = True
		self.errcode = 0
		self.errmsg = ""
		self.request_buffer = b""

	def clean_shutdown_control (self, phase, time_in_this_phase):
		if phase == 2:
			self.error (12, "Server Entered Shutdown Process")
			self.handle_close ()
		return 0

	def readable (self):
		if self.affluent is not None:
			return asynchat.async_chat.readable (self)	and self.affluent ()
		return asynchat.async_chat.readable (self)

	def writable (self):
		if self.ready is not None:
			return asynchat.async_chat.writable (self) and self.ready ()
		return asynchat.async_chat.writable (self)

	def handle_write (self):
		self.initiate_send()

	def maintern (self):
		if self.isactive () and time.time () - self.event_time > self.zombie_timeout:
			# do not user close_socket (), this func might be called in the thread, and currently in select.select()
			self.handle_close ()

	def is_deletable (self, timeout):
		if time.time () - self.event_time > timeout:
			if not self.isactive ():
				self.abort ()
				return True
		return False

	def is_in_channel (self, map = None):
		if map is None:
			map = self._map
		return self._fileno in map

	def del_channel (self, map=None):
		fd = self._fileno
		if map is None:
			map = self._map
		if fd in map:
			del map[fd]

	def close_socket (self):
		self.connected = False
		self.accepting = False
		self.del_channel ()
		self._fineno = None

		if self.socket:
			try:
				self.socket.close()
			except socket.error as why:
				if why.args[0] not in (ENOTCONN, EBADF):
					raise

	def close (self, force = False):
		if force:
			self.close_it = True

		if self.request:
			self.request.done (self.errcode, self.errmsg)

		if self.connected and self.close_it:
			self.close_socket ()
		else:
			self.del_channel ()

		self.request = None
		self.set_active (False)

	def abort (self):
		self.close_socket ()
		self.request = None
		self.set_active (False)

	def error (self, code, msg):
		self.close_it = True
		self.errcode = code
		self.errmsg = msg

	def set_active (self, flag, nolock = False):
		if flag:
			flag = time.time ()
		else:
			flag = 0

		if nolock or self.lock is None:
			self.active = flag
			return

		self.lock.acquire ()
		self.active = flag
		self.request_count += 1
		self.lock.release ()

	def get_active (self, nolock = False):
		if nolock or self.lock is None:
			return self.active
		self.lock.acquire ()
		active = self.active
		self.lock.release ()
		return active

	def isactive (self):
		return self.get_active () > 0

	def isconnected (self):
		return self.connected

	def get_request_count (self):
		return self.request_count

	def add_channel (self, map = None):
		return asynchat.async_chat.add_channel (self, map)

	def create_socket (self, family, type):
		self.family_and_type = family, type
		sock = socket.socket (family, type)
		sock.setblocking (0)
		self.set_socket (sock)

	def connect (self, force = 0):
		self.event_time = time.time ()
		self.create_socket (socket.AF_INET, socket.SOCK_STREAM)

		if adns.query:
			res = adns.get (self.address [0], "A")
			if res:
				ip = res [-1]["data"]
				if ip:
					asynchat.async_chat.connect (self, (ip, self.address [1]))
					return

		asynchat.async_chat.connect (self, self.address)

	def recv (self, buffer_size):
		self.event_time = time.time ()
		try:
			data = self.socket.recv(buffer_size)
			if not data:
				if not self.got_data: # disconnected by server
					self.log ("connection closed by remote server maybe caused by keep-alive timeout, retry connect...", "info")
					if self.reconnect ():
						return b''
				self.handle_close ()
				return b''
			else:
				return data

		except socket.error as why:
			if why.errno in asyncore._DISCONNECTED:
				if not self.got_data: # disconnected by server
					self.log ("_DISCONNECTED Error in recv (), retry connect...", "info")
					if self.reconnect ():
						return b''
				self.close_it = True
				self.handle_close ()
				return b''
			else:
				raise

	def send (self, data):
		self.event_time = time.time ()
		try:
			return self.socket.send(data)
		except socket.error as why:
			if why.errno == EWOULDBLOCK:
				return 0
			elif why.errno in asyncore._DISCONNECTED:
				self.close_it = True
				self.handle_close ()
				return 0
			else:
				raise

	def reconnect (self):
		self.close_socket ()
		return self.request.retry ()

	def close_if_over_keep_live (self):
		if time.time () - self.event_time > self.keep_alive:
			if self.connected:
				self.close_socket ()
			return True
		return False

	def start_request (self, request):
		self.initialize ()
		self.request = request

		if self.connected:
			self.close_if_over_keep_live () # check keep-alive

		try:
			if not self.connected:
				self.connect ()
			else:
				self.add_channel ()

		except:
			self.handle_error ()

	def initiate_send (self):
		if self.is_in_channel ():
			asynchat.async_chat.initiate_send (self)

	# proxy POST need no init_send
	def push (self, thing, init_send = True):
		if type (thing) is type (""):
			asynchat.async_chat.push (self, thing)
		else:
			self.push_with_producer (thing, init_send)

	def push_with_producer (self, producer, init_send = True):
		self.producer_fifo.append (producer)
		if init_send:
			self.initiate_send ()

	def collect_incoming_data (self, data):
		if not self.request:
			return # already closed
		self.got_data = True
		self.request.collect_incoming_data (data)

	def found_terminator (self):
		if not self.request:
			return # already closed
		self.request.found_terminator ()

	def set_timeout (self, timeout):
		self.zombie_timeout = timeout

	def handle_connect (self):
		pass

	def handle_timeout (self):
		self.log ("socket timeout", "fail")
		self.error (13, "Socket Timeout")
		self.handle_close ()

	def handle_expt (self):
		self.loggger ("socket panic", "fail")
		self.error (14, "Socket Panic")
		self.handle_close ()

	def handle_error (self):
		self.trace ()
		self.error (11, "Exception")
		self.handle_close()

	def handle_expt_event(self):
		err = self.socket.getsockopt(socket.SOL_SOCKET, socket.SO_ERROR)
		if err != 0:
			self.log ("SO_ERROR %d Occurred" % err, "warn")
			self.error (10, "Socket %d Error" % err)
			self.handle_close ()
		else:
			self.handle_expt ()

class AsynSSLConnect (AsynConnect):
	ac_in_buffer_size = 65535 # generally safe setting 65535 for SSL
	ac_out_buffer_size = 65535

	def connect (self, force = 0):
		self.handshaking = False
		AsynConnect.connect (self, force)

	def handle_connect_event(self):
		if not self.handshaking:
			err = self.socket.getsockopt(socket.SOL_SOCKET, socket.SO_ERROR)
			if err != 0:
				raise socket.error(err, _strerror(err))

			self.socket = ssl.wrap_socket (self.socket, do_handshake_on_connect = False)
			self.handshaking = True

		try:
			self.socket.do_handshake ()
		except ssl.SSLError as why:
			if why.args[0] in (ssl.SSL_ERROR_WANT_READ, ssl.SSL_ERROR_WANT_WRITE):
				return # retry handshake
			raise ssl.SSLError(why)

		# handshaking done
		self.handle_connect()
		self.connected = True

	def recv (self, buffer_size):
		self.event_time = time.time ()
		try:
			data = self.socket.recv (buffer_size)
			if not data:
				if not self.got_data: # disconnected by server
					self.log ("SSL connection closed by remote server maybe caused by keep-alive timeout, retry connect...", "info")
					if self.reconnect ():
						return b''
				self.handle_close ()
				return b''
			else:
				return data

		except ssl.SSLError as why:
			if why.errno == ssl.SSL_ERROR_WANT_READ:
				return b'' # retry
			# closed connection
			elif why.errno == ssl.SSL_ERROR_ZERO_RETURN:
				if not self.got_data: # disconnected by server
					self.log ("SSL_ERROR_ZERO_RETURN Error Occurred in recv (), retry connect...", "info")
					if self.reconnect ():
						return b''
				self.close_it = True
				self.handle_close ()
				return b''

			# eof error
			elif why.errno == ssl.SSL_ERROR_EOF:
				self.log ("SSL_ERROR_EOF Error Occurred in recv ()", "warn")
				self.close_it = True
				self.handle_close ()
				return b''
			else:
				raise

	def send (self, data):
		self.event_time = time.time ()
		try:
			return self.socket.send(data)

		except ssl.SSLError as why:
			if why.errno == ssl.SSL_ERROR_WANT_WRITE:
				return 0
			else:
				raise

<commit_msg>fix asyncon initiate_send, asyndns results #5<commit_after>import asynchat
import asyncore
import re
import socket
import time
import sys
import zlib
import ssl
from errno import ECONNRESET, ENOTCONN, ESHUTDOWN, ECONNABORTED, EWOULDBLOCK
import select
import threading
from . import adns
from warnings import warn

class SocketPanic (Exception): pass
class TimeOut (Exception): pass

class AsynConnect (asynchat.async_chat):
	ac_in_buffer_size = 4096
	ac_out_buffer_size = 4096
	zombie_timeout = 120
	keep_alive = 300

	def __init__ (self, address, lock = None, logger = None):
		self.address = address
		self.lock = lock
		self.logger = logger

		self.request_count = 0
		self.event_time = time.time ()
		self._cv = threading.Condition ()
		self.active = 0
		self.ready = None
		self.affluent = None
		self.initialize ()
		asynchat.async_chat.__init__ (self)

	def log (self, msg, logtype):
		if self.request is not None and hasattr (self.request, "log"):
			self.request.log (msg, logtype)
		elif self.logger:
			self.logger (msg, logtype)
		else:
			warn ("No logger")

	def trace (self):
		if self.request is not None and hasattr (self.request, "trace"):
			self.request.trace ()
		elif self.logger:
			self.logger.trace ()
		else:
			warn ("No logger for traceback")

	def duplicate (self):
		return self.__class__ (self.address, self.lock, self.logger)

	def initialize (self):
		self.request = None
		self.got_data = False
		self.close_it = True
		self.errcode = 0
		self.errmsg = ""
		self.request_buffer = b""

	def clean_shutdown_control (self, phase, time_in_this_phase):
		if phase == 2:
			self.error (12, "Server Entered Shutdown Process")
			self.handle_close ()
		return 0

	def readable (self):
		if self.affluent is not None:
			return asynchat.async_chat.readable (self)	and self.affluent ()
		return asynchat.async_chat.readable (self)

	def writable (self):
		if self.ready is not None:
			return asynchat.async_chat.writable (self) and self.ready ()
		return asynchat.async_chat.writable (self)

	def handle_write (self):
		self.initiate_send()

	def maintern (self):
		if self.isactive () and time.time () - self.event_time > self.zombie_timeout:
			# do not user close_socket (), this func might be called in the thread, and currently in select.select()
			self.handle_close ()

	def is_deletable (self, timeout):
		if time.time () - self.event_time > timeout:
			if not self.isactive ():
				self.abort ()
				return True
		return False

	def is_channel_in_map (self, map = None):
		if map is None:
			map = self._map
		return self._fileno in map

	def del_channel (self, map=None):
		fd = self._fileno
		if map is None:
			map = self._map
		if fd in map:
			del map[fd]

	def close_socket (self):
		self.connected = False
		self.accepting = False
		self.del_channel ()
		self._fineno = None

		if self.socket:
			try:
				self.socket.close()
			except socket.error as why:
				if why.args[0] not in (ENOTCONN, EBADF):
					raise

	def close (self, force = False):
		if force:
			self.close_it = True

		if self.request:
			self.request.done (self.errcode, self.errmsg)

		if self.connected and self.close_it:
			self.close_socket ()
		else:
			self.del_channel ()

		self.request = None
		self.set_active (False)

	def abort (self):
		self.close_socket ()
		self.request = None
		self.set_active (False)

	def error (self, code, msg):
		self.close_it = True
		self.errcode = code
		self.errmsg = msg

	def set_active (self, flag, nolock = False):
		if flag:
			flag = time.time ()
		else:
			flag = 0

		if nolock or self.lock is None:
			self.active = flag
			return

		self.lock.acquire ()
		self.active = flag
		self.request_count += 1
		self.lock.release ()

	def get_active (self, nolock = False):
		if nolock or self.lock is None:
			return self.active
		self.lock.acquire ()
		active = self.active
		self.lock.release ()
		return active

	def isactive (self):
		return self.get_active () > 0

	def isconnected (self):
		return self.connected

	def get_request_count (self):
		return self.request_count

	def add_channel (self, map = None):
		return asynchat.async_chat.add_channel (self, map)

	def create_socket (self, family, type):
		self.family_and_type = family, type
		sock = socket.socket (family, type)
		sock.setblocking (0)
		self.set_socket (sock)

	def connect (self, force = 0):
		self.event_time = time.time ()
		self.create_socket (socket.AF_INET, socket.SOCK_STREAM)

		if adns.query:
			res = adns.get (self.address [0], "A")
			if res:
				ip = res [-1]["data"]
				if ip:
					asynchat.async_chat.connect (self, (ip, self.address [1]))
					return

		asynchat.async_chat.connect (self, self.address)

	def recv (self, buffer_size):
		self.event_time = time.time ()
		try:
			data = self.socket.recv(buffer_size)
			if not data:
				if not self.got_data: # disconnected by server
					self.log ("connection closed by remote server maybe caused by keep-alive timeout, retry connect...", "info")
					if self.reconnect ():
						return b''
				self.handle_close ()
				return b''
			else:
				return data

		except socket.error as why:
			if why.errno in asyncore._DISCONNECTED:
				if not self.got_data: # disconnected by server
					self.log ("_DISCONNECTED Error in recv (), retry connect...", "info")
					if self.reconnect ():
						return b''
				self.close_it = True
				self.handle_close ()
				return b''
			else:
				raise

	def send (self, data):
		self.event_time = time.time ()
		try:
			return self.socket.send(data)
		except socket.error as why:
			if why.errno == EWOULDBLOCK:
				return 0
			elif why.errno in asyncore._DISCONNECTED:
				self.close_it = True
				self.handle_close ()
				return 0
			else:
				raise

	def reconnect (self):
		self.close_socket ()
		return self.request.retry ()

	def close_if_over_keep_live (self):
		if time.time () - self.event_time > self.keep_alive:
			if self.connected:
				self.close_socket ()
			return True
		return False

	def start_request (self, request):
		self.initialize ()
		self.request = request

		if self.connected:
			self.close_if_over_keep_live () # check keep-alive

		try:
			if not self.connected:
				self.connect ()
			else:
				self.add_channel ()

		except:
			self.handle_error ()

	def initiate_send (self):
		if self.is_channel_in_map ():
			asynchat.async_chat.initiate_send (self)

	# proxy POST need no init_send
	def push (self, thing, init_send = True):
		if type (thing) is type (""):
			asynchat.async_chat.push (self, thing)
		else:
			self.push_with_producer (thing, init_send)

	def push_with_producer (self, producer, init_send = True):
		self.producer_fifo.append (producer)
		if init_send:
			self.initiate_send ()

	def collect_incoming_data (self, data):
		if not self.request:
			return # already closed
		self.got_data = True
		self.request.collect_incoming_data (data)

	def found_terminator (self):
		if not self.request:
			return # already closed
		self.request.found_terminator ()

	def set_timeout (self, timeout):
		self.zombie_timeout = timeout

	def handle_connect (self):
		pass

	def handle_timeout (self):
		self.log ("socket timeout", "fail")
		self.error (13, "Socket Timeout")
		self.handle_close ()

	def handle_expt (self):
		self.loggger ("socket panic", "fail")
		self.error (14, "Socket Panic")
		self.handle_close ()

	def handle_error (self):
		self.trace ()
		self.error (11, "Exception")
		self.handle_close()

	def handle_expt_event(self):
		err = self.socket.getsockopt(socket.SOL_SOCKET, socket.SO_ERROR)
		if err != 0:
			self.log ("SO_ERROR %d Occurred" % err, "warn")
			self.error (10, "Socket %d Error" % err)
			self.handle_close ()
		else:
			self.handle_expt ()

class AsynSSLConnect (AsynConnect):
	ac_in_buffer_size = 65535 # generally safe setting 65535 for SSL
	ac_out_buffer_size = 65535

	def connect (self, force = 0):
		self.handshaking = False
		AsynConnect.connect (self, force)

	def handle_connect_event(self):
		if not self.handshaking:
			err = self.socket.getsockopt(socket.SOL_SOCKET, socket.SO_ERROR)
			if err != 0:
				raise socket.error(err, _strerror(err))

			self.socket = ssl.wrap_socket (self.socket, do_handshake_on_connect = False)
			self.handshaking = True

		try:
			self.socket.do_handshake ()
		except ssl.SSLError as why:
			if why.args[0] in (ssl.SSL_ERROR_WANT_READ, ssl.SSL_ERROR_WANT_WRITE):
				return # retry handshake
			raise ssl.SSLError(why)

		# handshaking done
		self.handle_connect()
		self.connected = True

	def recv (self, buffer_size):
		self.event_time = time.time ()
		try:
			data = self.socket.recv (buffer_size)
			if not data:
				if not self.got_data: # disconnected by server
					self.log ("SSL connection closed by remote server maybe caused by keep-alive timeout, retry connect...", "info")
					if self.reconnect ():
						return b''
				self.handle_close ()
				return b''
			else:
				return data

		except ssl.SSLError as why:
			if why.errno == ssl.SSL_ERROR_WANT_READ:
				return b'' # retry
			# closed connection
			elif why.errno == ssl.SSL_ERROR_ZERO_RETURN:
				if not self.got_data: # disconnected by server
					self.log ("SSL_ERROR_ZERO_RETURN Error Occurred in recv (), retry connect...", "info")
					if self.reconnect ():
						return b''
				self.close_it = True
				self.handle_close ()
				return b''

			# eof error
			elif why.errno == ssl.SSL_ERROR_EOF:
				self.log ("SSL_ERROR_EOF Error Occurred in recv ()", "warn")
				self.close_it = True
				self.handle_close ()
				return b''
			else:
				raise

	def send (self, data):
		self.event_time = time.time ()
		try:
			return self.socket.send(data)

		except ssl.SSLError as why:
			if why.errno == ssl.SSL_ERROR_WANT_WRITE:
				return 0
			else:
				raise

<|endoftext|>