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

# Importing pygame modules
import random, sys, pygame
from pygame.locals import *
from explosion import Blowup

# Set variables, like screen width and height
# globals
FPS = 30
WINDOWWIDTH = 800
WINDOWHEIGHT = 600
TILESIZE = 40
BUTTONHEIGHT = 20
BUTTONWIDTH = 40
TEXT_HEIGHT = 25
TEXT_LEFT_POSN = 10
BOARDWIDTH = 10
BOARDHEIGHT = 10
DISPLAYWIDTH = 200
XMARGIN = int((WINDOWWIDTH - (BOARDWIDTH * TILESIZE) - (DISPLAYWIDTH + 20)) / 2)
YMARGIN = int((WINDOWHEIGHT - (BOARDHEIGHT * TILESIZE)) / 2)

BLACK   = (  0,   0,   0)
WHITE   = (255, 255, 255)
GREEN   = (  0, 204,   0)
GRAY    = ( 60,  60,  60)
BLUE    = (  0,  50, 255)
YELLOW  = (255, 255,   0)
DARKGRAY =( 40,  40,  40)

BGCOLOR = GRAY
BUTTONCOLOR = GREEN
TEXTCOLOR = WHITE
TILECOLOR = GREEN
BORDERCOLOR = BLUE
TEXTSHADOWCOLOR = BLUE
SHIPCOLOR = YELLOW
HIGHLIGHTCOLOR = BLUE

def main():
    global DISPLAYSURF, FPSCLOCK, BASICFONT, HELP_SURF, HELP_RECT, NEW_SURF, \
           NEW_RECT, SHOTS_SURF, SHOTS_RECT, BIGFONT, COUNTER_SURF, \
           COUNTER_RECT
    pygame.init()
    FPSCLOCK = pygame.time.Clock()
    DISPLAYSURF = pygame.display.set_mode((WINDOWWIDTH, WINDOWHEIGHT))
    BASICFONT = pygame.font.Font('freesansbold.ttf', 20)
    BIGFONT = pygame.font.Font('freesansbold.ttf', 100)

    # create buttons
    HELP_SURF = BASICFONT.render("HELP", True, WHITE)
    HELP_RECT = HELP_SURF.get_rect()
    HELP_RECT.topleft = (WINDOWWIDTH - 180, WINDOWHEIGHT - 350)
    NEW_SURF = BASICFONT.render("NEW GAME", True, WHITE)
    NEW_RECT = NEW_SURF.get_rect()
    NEW_RECT.topleft = (WINDOWWIDTH - 200, WINDOWHEIGHT - 200)

    # 'Shots:' label
    SHOTS_SURF = BASICFONT.render("Shots: ", True, WHITE)
    SHOTS_RECT = SHOTS_SURF.get_rect()
    SHOTS_RECT.topleft = (WINDOWWIDTH - 750, WINDOWHEIGHT - 570)

    pygame.display.set_caption('Battleship')

    while True:
        run_game()
        show_text_screen('Game Over')

def run_game():
    revealed_tiles = generate_default_tiles(False)
    main_board = generate_default_tiles((None, None))
    ship_objs = make_ships() # list of ships to be used, holds list of tuples
                             # of coords
    main_board = add_ships_to_board(main_board, ship_objs)
    mousex, mousey = 0, 0
    counter = [] # counter to track number of shots fired

    while True:
        # counter display (it needs to be here in order to refresh it)
        COUNTER_SURF = BASICFONT.render(str(len(counter)), True, WHITE)
        COUNTER_RECT = SHOTS_SURF.get_rect()
        COUNTER_RECT.topleft = (WINDOWWIDTH - 680, WINDOWHEIGHT - 570)
        # end of the counter
        DISPLAYSURF.fill(BGCOLOR)
        DISPLAYSURF.blit(HELP_SURF, HELP_RECT)
        DISPLAYSURF.blit(NEW_SURF, NEW_RECT)
        DISPLAYSURF.blit(SHOTS_SURF, SHOTS_RECT)
        DISPLAYSURF.blit(COUNTER_SURF, COUNTER_RECT)
        draw_board(main_board, revealed_tiles)
        mouse_clicked = False

        check_for_quit()
        for event in pygame.event.get():
            if event.type == MOUSEBUTTONUP:
                if HELP_RECT.collidepoint(event.pos):
                    DISPLAYSURF.fill(BGCOLOR)
                    show_help_screen()
                elif NEW_RECT.collidepoint(event.pos):
                    main()
                else:
                    mousex, mousey = event.pos
                    mouse_clicked = True
            elif event.type == MOUSEMOTION:
                mousex, mousey = event.pos

        tilex, tiley = get_tile_at_pixel(mousex, mousey)
        if tilex != None and tiley != None:
            if not revealed_tiles[tilex][tiley]:
                draw_highlight_tile(tilex, tiley)
            if not revealed_tiles[tilex][tiley] and mouse_clicked:
                reveal_tile_animation(main_board, [(tilex, tiley)])
                revealed_tiles[tilex][tiley] = True
                counter.append((tilex, tiley))

        pygame.display.update()
        FPSCLOCK.tick(FPS)

def generate_default_tiles(default_value):
    '''
    returns list of 10 x 10 tiles set to False
    '''
    default_tiles = []
    for i in range(BOARDWIDTH):
        default_tiles.append([default_value] * BOARDHEIGHT)
    return default_tiles

def blowup_animation(coord):
	'''
	coord --> tuple of tile coords to apply the blowup animation
	'''
	ex = Blowup()
	ex.is_going_on = True
	while ex.is_going_on:
		ex.update()
		ex.image = pygame.transform.scale(ex.image, (TILESIZE+10,TILESIZE+10))
		DISPLAYSURF.blit(ex.image,coord)
		pygame.display.flip()
		FPSCLOCK.tick(10)

def reveal_tile_animation(board, tile_to_reveal):
    '''
    board: list of board tiles
    tile_to_reveal: tuple of tile coords to apply the reveal animation to
    '''
    left, top = left_top_coords_tile(tile_to_reveal[0][0], tile_to_reveal[0][1])
    if board[tile_to_reveal[0][0]][tile_to_reveal[0][1]] != (None, None):
    	blowup_animation((left,top))
        pygame.draw.rect(DISPLAYSURF, SHIPCOLOR, (left, top, TILESIZE,
                                                  TILESIZE))
    else:
        pygame.draw.rect(DISPLAYSURF, BGCOLOR, (left, top, TILESIZE,
                                                TILESIZE))

    pygame.display.update()
    FPSCLOCK.tick(FPS)

def check_for_quit():
    for event in pygame.event.get(QUIT):
        pygame.quit()
        sys.exit()

def draw_board(board, revealed):
    '''
    board: list of board tiles
    revealed: list of revealed tiles
    '''
    for tilex in range(BOARDWIDTH):
        for tiley in range(BOARDHEIGHT):
            left, top = left_top_coords_tile(tilex, tiley)
            if not revealed[tilex][tiley]:
                pygame.draw.rect(DISPLAYSURF, TILECOLOR, (left, top, TILESIZE,
                                                          TILESIZE))
            else:
                if board[tilex][tiley] != (None, None):
                    pygame.draw.rect(DISPLAYSURF, SHIPCOLOR, (left, top,
                                     TILESIZE, TILESIZE))
                else:
                    pygame.draw.rect(DISPLAYSURF, BGCOLOR, (left, top,
                                     TILESIZE, TILESIZE))

    for x in range(0, BOARDWIDTH * TILESIZE, TILESIZE):
        pygame.draw.line(DISPLAYSURF, DARKGRAY, (x + XMARGIN, YMARGIN), \
                        (x + XMARGIN, WINDOWHEIGHT - YMARGIN))
    for y in range(0, BOARDHEIGHT * TILESIZE, TILESIZE):
        pygame.draw.line(DISPLAYSURF, DARKGRAY, (XMARGIN, y + YMARGIN), \
                (WINDOWWIDTH - (DISPLAYWIDTH + 20 + XMARGIN), y + YMARGIN))

def make_ships():
    '''
    returns list of tuples of ships, each section of a ship has
    the tuple (ship name, shot)
    '''
    slist = []
    # make battleship
    ship = []
    for i in range(4):
        ship.append(('battleship',False))
    slist.append(ship)
    # make destroyer
    ship = []
    for i in range(3):
        ship.append(('destroyer',False))
    slist.append(ship)
    # make submarine
    ship = []
    for i in range(3):
        ship.append(('submarine',False))
    slist.append(ship)

    return slist

def add_ships_to_board(board, ships):
    '''
    return list of board tiles with ships placed on certain tiles
    board: list of board tiles
    ships: list of ships (name, bool) to place on board
    '''
    new_board = board[:]
    for ship in ships:
        for i in range(len(ship)):
            if ship[i][0] == 'battleship':
                new_board[1][1+i] = ship[i]
            elif ship[i][0] == 'destroyer':
                new_board[3][2+i] = ship[i]
            elif ship[i][0] == 'submarine':
                new_board[3+i][8] = ship[i]
    return new_board

def left_top_coords_tile(tilex, tiley):
    '''
    returns left and top pixel coords
    tilex: int
    tiley: int
    return: tuple (int, int)
    '''
    left = tilex * TILESIZE + XMARGIN
    top = tiley * TILESIZE + YMARGIN
    return (left, top)

def get_tile_at_pixel(x, y):
    '''
    returns tile coordinates of pixel at top left, defaults to (None, None)
    x: int
    y: int
    return: tuple (tilex, tiley)
    '''
    for tilex in range(BOARDWIDTH):
        for tiley in range(BOARDHEIGHT):
            left, top = left_top_coords_tile(tilex, tiley)
            tile_rect = pygame.Rect(left, top, TILESIZE, TILESIZE)
            if tile_rect.collidepoint(x, y):
                return (tilex, tiley)
    return (None, None)

def draw_highlight_tile(tilex, tiley):
    '''
    tilex: int
    tiley: int
    '''
    left, top = left_top_coords_tile(tilex, tiley)
    pygame.draw.rect(DISPLAYSURF, HIGHLIGHTCOLOR,
                    (left, top, TILESIZE, TILESIZE), 4)

def show_help_screen():
    # display the help screen until a button is pressed
    line1_surf, line1_rect = make_text_objs('Press a key to return to the game',
                                            BASICFONT, TEXTCOLOR)
    line1_rect.topleft = (TEXT_LEFT_POSN, TEXT_HEIGHT)
    DISPLAYSURF.blit(line1_surf, line1_rect)

    line2_surf, line2_rect = make_text_objs('This is the classic game of ' \
                                            'battleship. Your objective is ' \
                                            'to sink as many ships',
                                            BASICFONT, TEXTCOLOR)
    line2_rect.topleft = (TEXT_LEFT_POSN, TEXT_HEIGHT * 3)
    DISPLAYSURF.blit(line2_surf, line2_rect)

    line3_surf, line3_rect = make_text_objs('as possible using only 20 shots.',
                                            BASICFONT, TEXTCOLOR)
    line3_rect.topleft = (TEXT_LEFT_POSN, TEXT_HEIGHT * 4)
    DISPLAYSURF.blit(line3_surf, line3_rect)

    while check_for_keypress() == None:
        pygame.display.update()
        FPSCLOCK.tick()

def check_for_keypress():
    # pulling out all KEYDOWN and KEYUP events from queue and returning any
    # KEYUP else return None
    for event in pygame.event.get([KEYDOWN, KEYUP]):
        if event.type == KEYDOWN:
            continue
        return event.key
    return None

def make_text_objs(text, font, color):
    '''
    text: string
    font: Font object
    color: tuple of color (red, green blue)
    return: surface object, rectangle object
    '''
    surf = font.render(text, True, color)
    return surf, surf.get_rect()

def show_text_screen(text):
    '''
    text: string
    '''
    DISPLAYSURF.fill(BGCOLOR)
    titleSurf, titleRect = make_text_objs(text, BIGFONT, TEXTSHADOWCOLOR)
    titleRect.center = (int(WINDOWWIDTH / 2), int(WINDOWHEIGHT / 2))
    DISPLAYSURF.blit(titleSurf, titleRect)

    titleSurf, titleRect = make_text_objs(text, BIGFONT, TEXTCOLOR)
    titleRect.center = (int(WINDOWWIDTH / 2) - 3, int(WINDOWHEIGHT / 2) - 3)
    DISPLAYSURF.blit(titleSurf, titleRect)

    pressKeySurf, pressKeyRect = make_text_objs('Press a key to play.',
                                                BASICFONT, TEXTCOLOR)
    pressKeyRect.center = (int(WINDOWWIDTH / 2), int(WINDOWHEIGHT / 2) + 100)
    DISPLAYSURF.blit(pressKeySurf, pressKeyRect)

    while check_for_keypress() == None:
        pygame.display.update()
        FPSCLOCK.tick()

if __name__ == "__main__": #This calls the game loop
    main()
<commit_msg>generated ship length randomly<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-

# Importing pygame modules
import random, sys, pygame
from pygame.locals import *
from explosion import Blowup

# Set variables, like screen width and height
# globals
FPS = 30
WINDOWWIDTH = 800
WINDOWHEIGHT = 600
TILESIZE = 40
BUTTONHEIGHT = 20
BUTTONWIDTH = 40
TEXT_HEIGHT = 25
TEXT_LEFT_POSN = 10
BOARDWIDTH = 10
BOARDHEIGHT = 10
DISPLAYWIDTH = 200
XMARGIN = int((WINDOWWIDTH - (BOARDWIDTH * TILESIZE) - (DISPLAYWIDTH + 20)) / 2)
YMARGIN = int((WINDOWHEIGHT - (BOARDHEIGHT * TILESIZE)) / 2)

BLACK   = (  0,   0,   0)
WHITE   = (255, 255, 255)
GREEN   = (  0, 204,   0)
GRAY    = ( 60,  60,  60)
BLUE    = (  0,  50, 255)
YELLOW  = (255, 255,   0)
DARKGRAY =( 40,  40,  40)

BGCOLOR = GRAY
BUTTONCOLOR = GREEN
TEXTCOLOR = WHITE
TILECOLOR = GREEN
BORDERCOLOR = BLUE
TEXTSHADOWCOLOR = BLUE
SHIPCOLOR = YELLOW
HIGHLIGHTCOLOR = BLUE

def main():
    global DISPLAYSURF, FPSCLOCK, BASICFONT, HELP_SURF, HELP_RECT, NEW_SURF, \
           NEW_RECT, SHOTS_SURF, SHOTS_RECT, BIGFONT, COUNTER_SURF, \
           COUNTER_RECT
    pygame.init()
    FPSCLOCK = pygame.time.Clock()
    DISPLAYSURF = pygame.display.set_mode((WINDOWWIDTH, WINDOWHEIGHT))
    BASICFONT = pygame.font.Font('freesansbold.ttf', 20)
    BIGFONT = pygame.font.Font('freesansbold.ttf', 100)

    # create buttons
    HELP_SURF = BASICFONT.render("HELP", True, WHITE)
    HELP_RECT = HELP_SURF.get_rect()
    HELP_RECT.topleft = (WINDOWWIDTH - 180, WINDOWHEIGHT - 350)
    NEW_SURF = BASICFONT.render("NEW GAME", True, WHITE)
    NEW_RECT = NEW_SURF.get_rect()
    NEW_RECT.topleft = (WINDOWWIDTH - 200, WINDOWHEIGHT - 200)

    # 'Shots:' label
    SHOTS_SURF = BASICFONT.render("Shots: ", True, WHITE)
    SHOTS_RECT = SHOTS_SURF.get_rect()
    SHOTS_RECT.topleft = (WINDOWWIDTH - 750, WINDOWHEIGHT - 570)

    pygame.display.set_caption('Battleship')

    while True:
        run_game()
        show_text_screen('Game Over')

def run_game():
    revealed_tiles = generate_default_tiles(False)
    main_board = generate_default_tiles((None, None))
    ship_objs = make_ships() # list of ships to be used, holds list of tuples
                             # of coords
    main_board = add_ships_to_board(main_board, ship_objs)
    mousex, mousey = 0, 0
    counter = [] # counter to track number of shots fired

    while True:
        # counter display (it needs to be here in order to refresh it)
        COUNTER_SURF = BASICFONT.render(str(len(counter)), True, WHITE)
        COUNTER_RECT = SHOTS_SURF.get_rect()
        COUNTER_RECT.topleft = (WINDOWWIDTH - 680, WINDOWHEIGHT - 570)
        # end of the counter
        DISPLAYSURF.fill(BGCOLOR)
        DISPLAYSURF.blit(HELP_SURF, HELP_RECT)
        DISPLAYSURF.blit(NEW_SURF, NEW_RECT)
        DISPLAYSURF.blit(SHOTS_SURF, SHOTS_RECT)
        DISPLAYSURF.blit(COUNTER_SURF, COUNTER_RECT)
        draw_board(main_board, revealed_tiles)
        mouse_clicked = False

        check_for_quit()
        for event in pygame.event.get():
            if event.type == MOUSEBUTTONUP:
                if HELP_RECT.collidepoint(event.pos):
                    DISPLAYSURF.fill(BGCOLOR)
                    show_help_screen()
                elif NEW_RECT.collidepoint(event.pos):
                    main()
                else:
                    mousex, mousey = event.pos
                    mouse_clicked = True
            elif event.type == MOUSEMOTION:
                mousex, mousey = event.pos

        tilex, tiley = get_tile_at_pixel(mousex, mousey)
        if tilex != None and tiley != None:
            if not revealed_tiles[tilex][tiley]:
                draw_highlight_tile(tilex, tiley)
            if not revealed_tiles[tilex][tiley] and mouse_clicked:
                reveal_tile_animation(main_board, [(tilex, tiley)])
                revealed_tiles[tilex][tiley] = True
                counter.append((tilex, tiley))

        pygame.display.update()
        FPSCLOCK.tick(FPS)

def generate_default_tiles(default_value):
    '''
    returns list of 10 x 10 tiles set to False
    '''
    default_tiles = []
    for i in range(BOARDWIDTH):
        default_tiles.append([default_value] * BOARDHEIGHT)
    return default_tiles

def blowup_animation(coord):
	'''
	coord --> tuple of tile coords to apply the blowup animation
	'''
	ex = Blowup()
	ex.is_going_on = True
	while ex.is_going_on:
		ex.update()
		ex.image = pygame.transform.scale(ex.image, (TILESIZE+10,TILESIZE+10))
		DISPLAYSURF.blit(ex.image,coord)
		pygame.display.flip()
		FPSCLOCK.tick(10)

def reveal_tile_animation(board, tile_to_reveal):
    '''
    board: list of board tiles
    tile_to_reveal: tuple of tile coords to apply the reveal animation to
    '''
    left, top = left_top_coords_tile(tile_to_reveal[0][0], tile_to_reveal[0][1])
    if board[tile_to_reveal[0][0]][tile_to_reveal[0][1]] != (None, None):
    	blowup_animation((left,top))
        pygame.draw.rect(DISPLAYSURF, SHIPCOLOR, (left, top, TILESIZE,
                                                  TILESIZE))
    else:
        pygame.draw.rect(DISPLAYSURF, BGCOLOR, (left, top, TILESIZE,
                                                TILESIZE))

    pygame.display.update()
    FPSCLOCK.tick(FPS)

def check_for_quit():
    for event in pygame.event.get(QUIT):
        pygame.quit()
        sys.exit()

def draw_board(board, revealed):
    '''
    board: list of board tiles
    revealed: list of revealed tiles
    '''
    for tilex in range(BOARDWIDTH):
        for tiley in range(BOARDHEIGHT):
            left, top = left_top_coords_tile(tilex, tiley)
            if not revealed[tilex][tiley]:
                pygame.draw.rect(DISPLAYSURF, TILECOLOR, (left, top, TILESIZE,
                                                          TILESIZE))
            else:
                if board[tilex][tiley] != (None, None):
                    pygame.draw.rect(DISPLAYSURF, SHIPCOLOR, (left, top,
                                     TILESIZE, TILESIZE))
                else:
                    pygame.draw.rect(DISPLAYSURF, BGCOLOR, (left, top,
                                     TILESIZE, TILESIZE))

    for x in range(0, BOARDWIDTH * TILESIZE, TILESIZE):
        pygame.draw.line(DISPLAYSURF, DARKGRAY, (x + XMARGIN, YMARGIN), \
                        (x + XMARGIN, WINDOWHEIGHT - YMARGIN))
    for y in range(0, BOARDHEIGHT * TILESIZE, TILESIZE):
        pygame.draw.line(DISPLAYSURF, DARKGRAY, (XMARGIN, y + YMARGIN), \
                (WINDOWWIDTH - (DISPLAYWIDTH + 20 + XMARGIN), y + YMARGIN))

def ship_length():
	'''
	return --> randomly ship length
	'''
	return random.randint(2,6)

def make_ships():
    '''
    returns list of tuples of ships, each section of a ship has
    the tuple (ship name, shot)
    '''
    slist = []

    # make battleship
    ship = [('battleship',False) for i in range(ship_length())]
    slist.append(ship)
    # make destroyer
    ship = [('destroyer',False) for i in range(ship_length())]
    slist.append(ship)
    # make submarine
    ship = [('submarine',False) for i in range(ship_length())]
    slist.append(ship)

    return slist

def add_ships_to_board(board, ships):
    '''
    return list of board tiles with ships placed on certain tiles
    board: list of board tiles
    ships: list of ships (name, bool) to place on board
    '''
    new_board = board[:]
    for ship in ships:
        for i in range(len(ship)):
            if ship[i][0] == 'battleship':
                new_board[1][1+i] = ship[i]
            elif ship[i][0] == 'destroyer':
                new_board[3][2+i] = ship[i]
            elif ship[i][0] == 'submarine':
                new_board[3+i][8] = ship[i]
    return new_board

def left_top_coords_tile(tilex, tiley):
    '''
    returns left and top pixel coords
    tilex: int
    tiley: int
    return: tuple (int, int)
    '''
    left = tilex * TILESIZE + XMARGIN
    top = tiley * TILESIZE + YMARGIN
    return (left, top)

def get_tile_at_pixel(x, y):
    '''
    returns tile coordinates of pixel at top left, defaults to (None, None)
    x: int
    y: int
    return: tuple (tilex, tiley)
    '''
    for tilex in range(BOARDWIDTH):
        for tiley in range(BOARDHEIGHT):
            left, top = left_top_coords_tile(tilex, tiley)
            tile_rect = pygame.Rect(left, top, TILESIZE, TILESIZE)
            if tile_rect.collidepoint(x, y):
                return (tilex, tiley)
    return (None, None)

def draw_highlight_tile(tilex, tiley):
    '''
    tilex: int
    tiley: int
    '''
    left, top = left_top_coords_tile(tilex, tiley)
    pygame.draw.rect(DISPLAYSURF, HIGHLIGHTCOLOR,
                    (left, top, TILESIZE, TILESIZE), 4)

def show_help_screen():
    # display the help screen until a button is pressed
    line1_surf, line1_rect = make_text_objs('Press a key to return to the game',
                                            BASICFONT, TEXTCOLOR)
    line1_rect.topleft = (TEXT_LEFT_POSN, TEXT_HEIGHT)
    DISPLAYSURF.blit(line1_surf, line1_rect)

    line2_surf, line2_rect = make_text_objs('This is the classic game of ' \
                                            'battleship. Your objective is ' \
                                            'to sink as many ships',
                                            BASICFONT, TEXTCOLOR)
    line2_rect.topleft = (TEXT_LEFT_POSN, TEXT_HEIGHT * 3)
    DISPLAYSURF.blit(line2_surf, line2_rect)

    line3_surf, line3_rect = make_text_objs('as possible using only 20 shots.',
                                            BASICFONT, TEXTCOLOR)
    line3_rect.topleft = (TEXT_LEFT_POSN, TEXT_HEIGHT * 4)
    DISPLAYSURF.blit(line3_surf, line3_rect)

    while check_for_keypress() == None:
        pygame.display.update()
        FPSCLOCK.tick()

def check_for_keypress():
    # pulling out all KEYDOWN and KEYUP events from queue and returning any
    # KEYUP else return None
    for event in pygame.event.get([KEYDOWN, KEYUP]):
        if event.type == KEYDOWN:
            continue
        return event.key
    return None

def make_text_objs(text, font, color):
    '''
    text: string
    font: Font object
    color: tuple of color (red, green blue)
    return: surface object, rectangle object
    '''
    surf = font.render(text, True, color)
    return surf, surf.get_rect()

def show_text_screen(text):
    '''
    text: string
    '''
    DISPLAYSURF.fill(BGCOLOR)
    titleSurf, titleRect = make_text_objs(text, BIGFONT, TEXTSHADOWCOLOR)
    titleRect.center = (int(WINDOWWIDTH / 2), int(WINDOWHEIGHT / 2))
    DISPLAYSURF.blit(titleSurf, titleRect)

    titleSurf, titleRect = make_text_objs(text, BIGFONT, TEXTCOLOR)
    titleRect.center = (int(WINDOWWIDTH / 2) - 3, int(WINDOWHEIGHT / 2) - 3)
    DISPLAYSURF.blit(titleSurf, titleRect)

    pressKeySurf, pressKeyRect = make_text_objs('Press a key to play.',
                                                BASICFONT, TEXTCOLOR)
    pressKeyRect.center = (int(WINDOWWIDTH / 2), int(WINDOWHEIGHT / 2) + 100)
    DISPLAYSURF.blit(pressKeySurf, pressKeyRect)

    while check_for_keypress() == None:
        pygame.display.update()
        FPSCLOCK.tick()

if __name__ == "__main__": #This calls the game loop
    main()
<|endoftext|>


<commit_before>#! /usr/bin/env python
'''  Tool for sorting imports alphabetically, and automatically separated into sections.

Copyright (C) 2013  Timothy Edmund Crosley

Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated
documentation files (the "Software"), to deal in the Software without restriction, including without limitation
the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and
to permit persons to whom the Software is furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all copies or
substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED
TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF
CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
OTHER DEALINGS IN THE SOFTWARE.

'''
from __future__ import absolute_import, division, print_function, unicode_literals

import argparse
import glob
import os
import sys

import setuptools

from isort import SortImports, __version__
from isort.settings import DEFAULT_SECTIONS, default, from_path

from .pie_slice import *

INTRO = """
/#######################################################################\\

     `sMMy`
     .yyyy-                                                      `
    ##soos##                                                    ./o.
          `     ``..-..`         ``...`.``         `   ```` ``-ssso```
     .s:-y-   .+osssssso/.     ./ossss+:so+:`    :+o-`/osso:+sssssssso/
     .s::y-   osss+.``.``     -ssss+-.`-ossso`   ssssso/::..::+ssss:::.
     .s::y-   /ssss+//:-.`   `ssss+     `ssss+   sssso`       :ssss`
     .s::y-   `-/+oossssso/  `ssss/      sssso   ssss/        :ssss`
     .y-/y-       ````:ssss`  ossso.    :ssss:   ssss/        :ssss.
     `/so:`    `-//::/osss+   `+ssss+-/ossso:    /sso-        `osssso/.
       \/      `-/oooo++/-      .:/++:/++/-`      ..           `://++/.

         isort your Python imports for you so you don't have to

                            VERSION {0}

\########################################################################/
""".format(__version__)

def iter_source_code(paths):
    """Iterate over all Python source files defined in paths."""
    for path in paths:
        if os.path.isdir(path):
            for dirpath, dirnames, filenames in os.walk(path):
                for filename in filenames:
                    if filename.endswith('.py'):
                        yield os.path.join(dirpath, filename)
        else:
            yield path

class ISortCommand(setuptools.Command):
    """The :class:`ISortCommand` class is used by setuptools to perform
    imports checks on registered modules.
    """

    description = "Run isort on modules registered in setuptools"
    user_options = []

    def initialize_options(self):
        default_settings = default.copy()
        for (key, value) in itemsview(default_settings):
            setattr(self, key, value)

    def finalize_options(self):
        "Get options from config files."
        self.arguments = {}
        computed_settings = from_path(os.getcwd())
        for (key, value) in itemsview(computed_settings):
            self.arguments[key] = value

    def distribution_files(self):
        """Find distribution packages."""
        # This is verbatim from flake8
        if self.distribution.packages:
            package_dirs = self.distribution.package_dir or {}
            for package in self.distribution.packages:
                pkg_dir = package
                if package in package_dirs:
                    pkg_dir = package_dirs[package]
                elif '' in package_dirs:
                    pkg_dir = package_dirs[''] + os.path.sep + pkg_dir
                yield pkg_dir.replace('.', os.path.sep)

        if self.distribution.py_modules:
            for filename in self.distribution.py_modules:
                yield "%s.py" % filename
        # Don't miss the setup.py file itself
        yield "setup.py"

    def run(self):
        arguments = self.arguments
        wrong_sorted_files = False
        arguments['check'] = True
        for path in self.distribution_files():
            for python_file in glob.iglob(os.path.join(path, '*.py')):
                try:
                    incorrectly_sorted = SortImports(python_file, **arguments).incorrectly_sorted
                    if incorrectly_sorted:
                        wrong_sorted_files = True
                except IOError as e:
                    print("WARNING: Unable to parse file {0} due to {1}".format(file_name, e))
        if wrong_sorted_files:
            exit(1)

def create_parser():
    parser = argparse.ArgumentParser(description='Sort Python import definitions alphabetically '
                                                 'within logical sections.')
    parser.add_argument('files', nargs='+', help='One or more Python source files that need their imports sorted.')
    parser.add_argument('-l', '--lines', help='[Deprecated] The max length of an import line (used for wrapping '
                        'long imports).',
                        dest='line_length', type=int)
    parser.add_argument('-w', '--line-width', help='The max length of an import line (used for wrapping long imports).',
                        dest='line_length', type=int)
    parser.add_argument('-s', '--skip', help='Files that sort imports should skip over. If you want to skip multiple '
                        'files you should specify twice: --skip file1 --skip file2.', dest='skip', action='append')
    parser.add_argument('-ns', '--dont-skip', help='Files that sort imports should never skip over.',
                        dest='not_skip', action='append')
    parser.add_argument('-sg', '--skip-glob', help='Files that sort imports should skip over.', dest='skip_glob',
                        action='append')
    parser.add_argument('-t', '--top', help='Force specific imports to the top of their appropriate section.',
                        dest='force_to_top', action='append')
    parser.add_argument('-f', '--future', dest='known_future_library', action='append',
                        help='Force sortImports to recognize a module as part of the future compatibility libraries.')
    parser.add_argument('-b', '--builtin', dest='known_standard_library', action='append',
                        help='Force sortImports to recognize a module as part of the python standard library.')
    parser.add_argument('-o', '--thirdparty', dest='known_third_party', action='append',
                        help='Force sortImports to recognize a module as being part of a third party library.')
    parser.add_argument('-p', '--project', dest='known_first_party', action='append',
                        help='Force sortImports to recognize a module as being part of the current python project.')
    parser.add_argument('-m', '--multi_line', dest='multi_line_output', type=int, choices=[0, 1, 2, 3, 4, 5],
                        help='Multi line output (0-grid, 1-vertical, 2-hanging, 3-vert-hanging, 4-vert-grid, '
                        '5-vert-grid-grouped).')
    parser.add_argument('-i', '--indent', help='String to place for indents defaults to "    " (4 spaces).',
                        dest='indent', type=str)
    parser.add_argument('-a', '--add_import', dest='add_imports', action='append',
                        help='Adds the specified import line to all files, '
                             'automatically determining correct placement.')
    parser.add_argument('-af', '--force_adds', dest='force_adds', action='store_true',
                        help='Forces import adds even if the original file is empty.')
    parser.add_argument('-r', '--remove_import', dest='remove_imports', action='append',
                        help='Removes the specified import from all files.')
    parser.add_argument('-ls', '--length_sort', help='Sort imports by their string length.',
                        dest='length_sort', action='store_true', default=False)
    parser.add_argument('-d', '--stdout', help='Force resulting output to stdout, instead of in-place.',
                        dest='write_to_stdout', action='store_true')
    parser.add_argument('-c', '--check-only', action='store_true', default=False, dest="check",
                        help='Checks the file for unsorted / unformatted imports and prints them to the '
                             'command line without modifying the file.')
    parser.add_argument('-sl', '--force_single_line_imports', dest='force_single_line', action='store_true',
                        help='Forces all from imports to appear on their own line')
    parser.add_argument('-sd', '--section-default', dest='default_section',
                        help='Sets the default section for imports (by default FIRSTPARTY) options: ' +
                        str(DEFAULT_SECTIONS))
    parser.add_argument('-df', '--diff', dest='show_diff', default=False, action='store_true',
                        help="Prints a diff of all the changes isort would make to a file, instead of "
                             "changing it in place")
    parser.add_argument('-e', '--balanced', dest='balanced_wrapping', action='store_true',
                        help='Balances wrapping to produce the most consistent line length possible')
    parser.add_argument('-rc', '--recursive', dest='recursive', action='store_true',
                        help='Recursively look for Python files of which to sort imports')
    parser.add_argument('-ot', '--order-by-type', dest='order_by_type',
                        action='store_true', help='Order imports by type in addition to alphabetically')
    parser.add_argument('-ac', '--atomic', dest='atomic', action='store_true',
                        help="Ensures the output doesn't save if the resulting file contains syntax errors.")
    parser.add_argument('-cs', '--combine-star', dest='combine_star', action='store_true',
                        help="Ensures that if a star import is present, nothing else is imported from that namespace.")
    parser.add_argument('-ca', '--combine-as', dest='combine_as_imports', action='store_true',
                        help="Combines as imports on the same line.")
    parser.add_argument('-tc', '--trailing-comma', dest='include_trailing_comma', action='store_true',
                        help='Includes a trailing comma on multi line imports that include parentheses.')
    parser.add_argument('-v', '--version', action='version', version='isort {0}'.format(__version__))
    parser.add_argument('-vb', '--verbose', action='store_true', dest="verbose",
                        help='Shows verbose output, such as when files are skipped or when a check is successful.')
    parser.add_argument('-q', '--quiet', action='store_true', dest="quiet",
                        help='Shows extra quiet output, only errors are outputted.')
    parser.add_argument('-sp', '--settings-path',  dest="settings_path",
                        help='Explicitly set the settings path instead of auto determining based on file location.')
    parser.add_argument('-ff', '--from-first', dest='from_first',
                        help="Switches the typical ordering preference, showing from imports first then straight ones.")
    parser.add_argument('-wl', '--wrap-length', dest='wrap_length',
                        help="Specifies how long lines that are wrapped should be, if not set line_length is used.")
    parser.add_argument('-fgw', '--force-grid-wrap',  action='store_true', dest="force_grid_wrap",
                        help='Force from imports to be grid wrapped regardless of line length')
    parser.add_argument('-fas', '--force-alphabetical-sort',  action='store_true', dest="force_alphabetical_sort",
                        help='Force all imports to be sorted as a single section')

    arguments = dict((key, value) for (key, value) in itemsview(vars(parser.parse_args())) if value)
    return arguments

def main():
    arguments = create_parser()
    file_names = arguments.pop('files', [])

    if file_names == ['-']:
        SortImports(file_contents=sys.stdin.read(), write_to_stdout=True, **arguments)
    else:
        wrong_sorted_files = False
        if arguments.get('recursive', False):
            file_names = iter_source_code(file_names)
        num_skipped = 0
        if not arguments.get('quiet', False):
            print(INTRO)
        for file_name in file_names:
            try:
                sort_attempt = SortImports(file_name, **arguments)
                incorrectly_sorted = sort_attempt.incorrectly_sorted
                if arguments.get('check', False) and incorrectly_sorted:
                    wrong_sorted_files = True
                if sort_attempt.skipped:
                    num_skipped += 1
            except IOError as e:
                print("WARNING: Unable to parse file {0} due to {1}".format(file_name, e))
        if wrong_sorted_files:
            exit(1)

        if num_skipped and not arguments.get('quiet', False):
            print("Skipped {0} files".format(num_skipped))

if __name__ == "__main__":
    main()
<commit_msg>Get rid of isort banner except for version<commit_after>#! /usr/bin/env python
'''  Tool for sorting imports alphabetically, and automatically separated into sections.

Copyright (C) 2013  Timothy Edmund Crosley

Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated
documentation files (the "Software"), to deal in the Software without restriction, including without limitation
the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and
to permit persons to whom the Software is furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all copies or
substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED
TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF
CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
OTHER DEALINGS IN THE SOFTWARE.

'''
from __future__ import absolute_import, division, print_function, unicode_literals

import argparse
import glob
import os
import sys

import setuptools

from isort import SortImports, __version__
from isort.settings import DEFAULT_SECTIONS, default, from_path

from .pie_slice import *

INTRO = """
/#######################################################################\\

     `sMMy`
     .yyyy-                                                      `
    ##soos##                                                    ./o.
          `     ``..-..`         ``...`.``         `   ```` ``-ssso```
     .s:-y-   .+osssssso/.     ./ossss+:so+:`    :+o-`/osso:+sssssssso/
     .s::y-   osss+.``.``     -ssss+-.`-ossso`   ssssso/::..::+ssss:::.
     .s::y-   /ssss+//:-.`   `ssss+     `ssss+   sssso`       :ssss`
     .s::y-   `-/+oossssso/  `ssss/      sssso   ssss/        :ssss`
     .y-/y-       ````:ssss`  ossso.    :ssss:   ssss/        :ssss.
     `/so:`    `-//::/osss+   `+ssss+-/ossso:    /sso-        `osssso/.
       \/      `-/oooo++/-      .:/++:/++/-`      ..           `://++/.

         isort your Python imports for you so you don't have to

                            VERSION {0}

\########################################################################/
""".format(__version__)

def iter_source_code(paths):
    """Iterate over all Python source files defined in paths."""
    for path in paths:
        if os.path.isdir(path):
            for dirpath, dirnames, filenames in os.walk(path):
                for filename in filenames:
                    if filename.endswith('.py'):
                        yield os.path.join(dirpath, filename)
        else:
            yield path

class ISortCommand(setuptools.Command):
    """The :class:`ISortCommand` class is used by setuptools to perform
    imports checks on registered modules.
    """

    description = "Run isort on modules registered in setuptools"
    user_options = []

    def initialize_options(self):
        default_settings = default.copy()
        for (key, value) in itemsview(default_settings):
            setattr(self, key, value)

    def finalize_options(self):
        "Get options from config files."
        self.arguments = {}
        computed_settings = from_path(os.getcwd())
        for (key, value) in itemsview(computed_settings):
            self.arguments[key] = value

    def distribution_files(self):
        """Find distribution packages."""
        # This is verbatim from flake8
        if self.distribution.packages:
            package_dirs = self.distribution.package_dir or {}
            for package in self.distribution.packages:
                pkg_dir = package
                if package in package_dirs:
                    pkg_dir = package_dirs[package]
                elif '' in package_dirs:
                    pkg_dir = package_dirs[''] + os.path.sep + pkg_dir
                yield pkg_dir.replace('.', os.path.sep)

        if self.distribution.py_modules:
            for filename in self.distribution.py_modules:
                yield "%s.py" % filename
        # Don't miss the setup.py file itself
        yield "setup.py"

    def run(self):
        arguments = self.arguments
        wrong_sorted_files = False
        arguments['check'] = True
        for path in self.distribution_files():
            for python_file in glob.iglob(os.path.join(path, '*.py')):
                try:
                    incorrectly_sorted = SortImports(python_file, **arguments).incorrectly_sorted
                    if incorrectly_sorted:
                        wrong_sorted_files = True
                except IOError as e:
                    print("WARNING: Unable to parse file {0} due to {1}".format(file_name, e))
        if wrong_sorted_files:
            exit(1)

def create_parser():
    parser = argparse.ArgumentParser(description='Sort Python import definitions alphabetically '
                                                 'within logical sections.')
    parser.add_argument('files', nargs='*', help='One or more Python source files that need their imports sorted.')
    parser.add_argument('-l', '--lines', help='[Deprecated] The max length of an import line (used for wrapping '
                        'long imports).',
                        dest='line_length', type=int)
    parser.add_argument('-w', '--line-width', help='The max length of an import line (used for wrapping long imports).',
                        dest='line_length', type=int)
    parser.add_argument('-s', '--skip', help='Files that sort imports should skip over. If you want to skip multiple '
                        'files you should specify twice: --skip file1 --skip file2.', dest='skip', action='append')
    parser.add_argument('-ns', '--dont-skip', help='Files that sort imports should never skip over.',
                        dest='not_skip', action='append')
    parser.add_argument('-sg', '--skip-glob', help='Files that sort imports should skip over.', dest='skip_glob',
                        action='append')
    parser.add_argument('-t', '--top', help='Force specific imports to the top of their appropriate section.',
                        dest='force_to_top', action='append')
    parser.add_argument('-f', '--future', dest='known_future_library', action='append',
                        help='Force sortImports to recognize a module as part of the future compatibility libraries.')
    parser.add_argument('-b', '--builtin', dest='known_standard_library', action='append',
                        help='Force sortImports to recognize a module as part of the python standard library.')
    parser.add_argument('-o', '--thirdparty', dest='known_third_party', action='append',
                        help='Force sortImports to recognize a module as being part of a third party library.')
    parser.add_argument('-p', '--project', dest='known_first_party', action='append',
                        help='Force sortImports to recognize a module as being part of the current python project.')
    parser.add_argument('-m', '--multi_line', dest='multi_line_output', type=int, choices=[0, 1, 2, 3, 4, 5],
                        help='Multi line output (0-grid, 1-vertical, 2-hanging, 3-vert-hanging, 4-vert-grid, '
                        '5-vert-grid-grouped).')
    parser.add_argument('-i', '--indent', help='String to place for indents defaults to "    " (4 spaces).',
                        dest='indent', type=str)
    parser.add_argument('-a', '--add_import', dest='add_imports', action='append',
                        help='Adds the specified import line to all files, '
                             'automatically determining correct placement.')
    parser.add_argument('-af', '--force_adds', dest='force_adds', action='store_true',
                        help='Forces import adds even if the original file is empty.')
    parser.add_argument('-r', '--remove_import', dest='remove_imports', action='append',
                        help='Removes the specified import from all files.')
    parser.add_argument('-ls', '--length_sort', help='Sort imports by their string length.',
                        dest='length_sort', action='store_true', default=False)
    parser.add_argument('-d', '--stdout', help='Force resulting output to stdout, instead of in-place.',
                        dest='write_to_stdout', action='store_true')
    parser.add_argument('-c', '--check-only', action='store_true', default=False, dest="check",
                        help='Checks the file for unsorted / unformatted imports and prints them to the '
                             'command line without modifying the file.')
    parser.add_argument('-sl', '--force_single_line_imports', dest='force_single_line', action='store_true',
                        help='Forces all from imports to appear on their own line')
    parser.add_argument('-sd', '--section-default', dest='default_section',
                        help='Sets the default section for imports (by default FIRSTPARTY) options: ' +
                        str(DEFAULT_SECTIONS))
    parser.add_argument('-df', '--diff', dest='show_diff', default=False, action='store_true',
                        help="Prints a diff of all the changes isort would make to a file, instead of "
                             "changing it in place")
    parser.add_argument('-e', '--balanced', dest='balanced_wrapping', action='store_true',
                        help='Balances wrapping to produce the most consistent line length possible')
    parser.add_argument('-rc', '--recursive', dest='recursive', action='store_true',
                        help='Recursively look for Python files of which to sort imports')
    parser.add_argument('-ot', '--order-by-type', dest='order_by_type',
                        action='store_true', help='Order imports by type in addition to alphabetically')
    parser.add_argument('-ac', '--atomic', dest='atomic', action='store_true',
                        help="Ensures the output doesn't save if the resulting file contains syntax errors.")
    parser.add_argument('-cs', '--combine-star', dest='combine_star', action='store_true',
                        help="Ensures that if a star import is present, nothing else is imported from that namespace.")
    parser.add_argument('-ca', '--combine-as', dest='combine_as_imports', action='store_true',
                        help="Combines as imports on the same line.")
    parser.add_argument('-tc', '--trailing-comma', dest='include_trailing_comma', action='store_true',
                        help='Includes a trailing comma on multi line imports that include parentheses.')
    parser.add_argument('-v', '--version', action='store_true', dest='show_version')
    parser.add_argument('-vb', '--verbose', action='store_true', dest="verbose",
                        help='Shows verbose output, such as when files are skipped or when a check is successful.')
    parser.add_argument('-q', '--quiet', action='store_true', dest="quiet",
                        help='Shows extra quiet output, only errors are outputted.')
    parser.add_argument('-sp', '--settings-path',  dest="settings_path",
                        help='Explicitly set the settings path instead of auto determining based on file location.')
    parser.add_argument('-ff', '--from-first', dest='from_first',
                        help="Switches the typical ordering preference, showing from imports first then straight ones.")
    parser.add_argument('-wl', '--wrap-length', dest='wrap_length',
                        help="Specifies how long lines that are wrapped should be, if not set line_length is used.")
    parser.add_argument('-fgw', '--force-grid-wrap',  action='store_true', dest="force_grid_wrap",
                        help='Force from imports to be grid wrapped regardless of line length')
    parser.add_argument('-fas', '--force-alphabetical-sort',  action='store_true', dest="force_alphabetical_sort",
                        help='Force all imports to be sorted as a single section')

    arguments = dict((key, value) for (key, value) in itemsview(vars(parser.parse_args())) if value)
    return arguments

def main():
    arguments = create_parser()
    if arguments.get('show_version'):
        print(INTRO)
        return

    file_names = arguments.pop('files', [])
    if not file_names:
        sys.stderr.write('isort: error: the following arguments are required: files\n')
        sys.exit(1)

    if file_names == ['-']:
        SortImports(file_contents=sys.stdin.read(), write_to_stdout=True, **arguments)
    else:
        wrong_sorted_files = False
        if arguments.get('recursive', False):
            file_names = iter_source_code(file_names)
        num_skipped = 0
        if arguments.get('verbose', False):
            print(INTRO)
        for file_name in file_names:
            try:
                sort_attempt = SortImports(file_name, **arguments)
                incorrectly_sorted = sort_attempt.incorrectly_sorted
                if arguments.get('check', False) and incorrectly_sorted:
                    wrong_sorted_files = True
                if sort_attempt.skipped:
                    num_skipped += 1
            except IOError as e:
                print("WARNING: Unable to parse file {0} due to {1}".format(file_name, e))
        if wrong_sorted_files:
            exit(1)

        if num_skipped and not arguments.get('quiet', False):
            print("Skipped {0} files".format(num_skipped))

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


<commit_before>import logging

from datetime import timedelta
import time
import hashlib
import pickle

from redis import StrictRedis, WatchError

log = logging.getLogger(__name__)

class Task(object):
    """
    Represents a single crawl task, corresponding to a future HTTP request to a
    particular URL, with some additional properties governing request
    parameters and the scheduling of that crawl.
    """
    DEFAULT_REFERER = object()
    DEFAULT_DOCUMENT_TYPE = object()

    @staticmethod
    def to_integer_seconds(val):
        if isinstance(val, timedelta):
            return val.total_seconds()
        else:
            return val

    def __init__(self, url, method='GET', data=None,
                 document_type=DEFAULT_DOCUMENT_TYPE, referer=DEFAULT_REFERER,
                 repeat_after=None, scheduled_timestamp=0, high_priority=False,
                 min_age=3600):

        if high_priority and scheduled_timestamp:
            raise ValueError('task cannot be high priority and also deferred')

        self.url = url
        self.document_type = document_type
        self.referer = referer
        self.high_priority = high_priority
        self.scheduled_timestamp = scheduled_timestamp
        self.repeat_after = Task.to_integer_seconds(repeat_after)
        self.min_age = Task.to_integer_seconds(min_age)

    def set_originating_task(self, task):
        if self.referer == self.DEFAULT_REFERER:
            self.referer = task.referer
        if self.document_type == self.DEFAULT_DOCUMENT_TYPE:
            self.document_type = task.document_type

    def __repr__(self):
        return '<Task %s [%s]%s>' % (self.url, self.document_type,
                                     ' HP' if self.high_priority else '')

class Empty(Exception):
    """
    Raised when there are no tasks scheduled for now left in the queue.
    """
    pass

class Queue(object):
    """
    Implements a sort-of-queue of crawl tasks.
    """
    def __init__(self, name, url_canonicalizer=None):
        self.name = name
        self.canonicalize_url = url_canonicalizer or (lambda url: url)
        self.redis = StrictRedis(host='localhost', port=6379, db=0)
        # XXX Temporary
        self.redis.flushall()

    def serialize(self, task):
        return pickle.dumps(task)

    def deserialize(self, s):
        return pickle.loads(s)

    def prefix_redis_key(self, prefix, key):
        return ':'.join([self.name, prefix, key])

    def pick_queue_key(self, high_priority):
        return self.prefix_redis_key('todo', 'hp' if high_priority else 'nn')

    def hash_url(self, url):
        return hashlib.md5(url).hexdigest()

    def record_crawl_timestamp(self, url, now):
        self.redis.set(self.prefix_redis_key('urlts', self.hash_url(url)),
                       '%d' % now)

    def get_crawl_timestamp(self, url):
        s = self.redis.get(self.prefix_redis_key('urlts', self.hash_url(url)))
        if s:
            return int(s)

    def get_existing_task_id(self, url):
        s = self.redis.get(self.prefix_redis_key('taskbyurl',
                                                 self.hash_url(url)))
        if s:
            task_id, high_priority_bool = s.split(',')
            return task_id, high_priority_bool == '1'
        else:
            return None, None

    def get_task(self, task_id):
        s = self.redis.get(task_id)
        return self.deserialize(s)

    def drop_task(self, task_id, high_priority):
        self.redis.zrem(self.pick_queue_key(high_priority), task_id)
        self.redis.delete(task_id)

    def push(self, task):
        """
        Schedule a new crawl task, which is an instance of ``Task``.

        If ``high_priority`` is set, drop all other tasks for the same
        canonicalized URL and enqueue to a separate high priority queue.

        If ``min_age`` is set and this task URL has been crawled before, check
        the last time it was crawled, and reschedule it to a timestamp which is
        ``min_age`` after the last crawl time.

        If a different crawl task is already scheduled for this URL, check the
        scheduled time. If it has an intentionally delayed crawl time and that
        scheduled time is *after* this task, drop it. Otherwise, drop this
        task. Either way, the earlier of the two tasks should be kept.
        """
        # XXX FIXME
        # We really should take a lock on this URL, to prevent race conditions
        # in the conditional logic below.

        existing_task_id, existing_is_hp = self.get_existing_task_id(task.url)

        if task.high_priority:
            # If this task is HP, drop any other task for this URL..
            if existing_task_id:
                self.drop_task(existing_task_id, existing_is_hp)

        elif existing_is_hp:
            # If there's an existing high priority task, drop this one.
            return

        else:
            # If min_age is set, check for previous crawls and further delay
            # this task in order to ensure that min_age is not violated.
            if task.min_age:
                last_timestamp = self.get_crawl_timestamp(task.url)
                if last_timestamp:
                    task.scheduled_timestamp = max(last_timestamp +
                                                   task.min_age,
                                                   task.scheduled_timestamp)

            # If there is an existing task, check if it's scheduled earlier.
            if existing_task_id:
                scheduled_ts = self.redis.zscore(
                    self.pick_queue_key(existing_is_hp),
                    existing_task_id)

                if scheduled_ts > task.scheduled_timestamp:
                    # If the existing task is later, drop it.
                    self.drop_task(existing_task_id, existing_is_hp)
                else:
                    # If this task is later, drop it and keep the existing
                    # task.
                    return

        # At this point we're ready to enqueue.
        s = self.serialize(task)
        task_id = hashlib.md5(s).hexdigest()

        # Store the task by ID.
        self.redis.set(task_id, s)

        # Finally, enqueue the task ID to the appropriate sorted set.
        self.redis.zadd(self.pick_queue_key(task.high_priority),
                        float(task.scheduled_timestamp), task_id)

    def pop_queue(self, high_priority, cutoff):
        key = self.pick_queue_key(high_priority)
        with self.redis.pipeline() as pipe:
            while True:
                try:
                    pipe.watch(key)
                    elements = pipe.zrangebyscore(key, min=0, max=cutoff,
                                                  start=0, num=1)
                    if elements:
                        task_id = elements[0]
                        pipe.multi()
                        pipe.zrem(key, task_id)
                        pipe.execute()
                        return task_id
                    else:
                        return
                except WatchError:
                    continue

    def pop(self):
        """
        Get the next scheduled crawl task, or raise Empty() if there is nothing
        to do.
        """
        now = time.time()
        for high_priority in (True, False):
            next_task_id = self.pop_queue(high_priority, now)
            if next_task_id:
                break

        # Nothing to do?
        if not next_task_id:
            raise Empty()

        # Otherwise, fetch the task and deserialize it.
        s = self.redis.get(next_task_id)
        self.redis.delete(next_task_id)
        task = self.deserialize(s)

        # Record a crawl timestamp for this URL.
        self.record_crawl_timestamp(task.url, now)

        # If the task has a repeat interval, reschedule it.
        if task.repeat_after:
            repeat_task = Task(url=task.url,
                               document_type=task.document_type,
                               repeat_after=task.repeat_after,
                               scheduled_timestamp=now + task.repeat_after)
            self.push(repeat_task)

        return task
<commit_msg>Add a queue.count() method to view the number and type of pending tasks<commit_after>import logging

from datetime import timedelta
import time
import hashlib
import pickle

from redis import StrictRedis, WatchError

log = logging.getLogger(__name__)

class Task(object):
    """
    Represents a single crawl task, corresponding to a future HTTP request to a
    particular URL, with some additional properties governing request
    parameters and the scheduling of that crawl.
    """
    DEFAULT_REFERER = object()
    DEFAULT_DOCUMENT_TYPE = object()

    @staticmethod
    def to_integer_seconds(val):
        if isinstance(val, timedelta):
            return val.total_seconds()
        else:
            return val

    def __init__(self, url, method='GET', data=None,
                 document_type=DEFAULT_DOCUMENT_TYPE, referer=DEFAULT_REFERER,
                 repeat_after=None, scheduled_timestamp=0, high_priority=False,
                 min_age=3600):

        if high_priority and scheduled_timestamp:
            raise ValueError('task cannot be high priority and also deferred')

        self.url = url
        self.document_type = document_type
        self.referer = referer
        self.high_priority = high_priority
        self.scheduled_timestamp = scheduled_timestamp
        self.repeat_after = Task.to_integer_seconds(repeat_after)
        self.min_age = Task.to_integer_seconds(min_age)

    def set_originating_task(self, task):
        if self.referer == self.DEFAULT_REFERER:
            self.referer = task.referer
        if self.document_type == self.DEFAULT_DOCUMENT_TYPE:
            self.document_type = task.document_type

    def __repr__(self):
        return '<Task %s [%s]%s>' % (self.url, self.document_type,
                                     ' HP' if self.high_priority else '')

class Empty(Exception):
    """
    Raised when there are no tasks scheduled for now left in the queue.
    """
    pass

class Queue(object):
    """
    Implements a sort-of-queue of crawl tasks.
    """
    def __init__(self, name, url_canonicalizer=None):
        self.name = name
        self.canonicalize_url = url_canonicalizer or (lambda url: url)
        self.redis = StrictRedis(host='localhost', port=6379, db=0)
        # XXX Temporary
        self.redis.flushall()

    def serialize(self, task):
        return pickle.dumps(task)

    def deserialize(self, s):
        return pickle.loads(s)

    def prefix_redis_key(self, prefix, key):
        return ':'.join([self.name, prefix, key])

    def pick_queue_key(self, high_priority):
        return self.prefix_redis_key('todo', 'hp' if high_priority else 'nn')

    def hash_url(self, url):
        return hashlib.md5(url).hexdigest()

    def record_crawl_timestamp(self, url, now):
        self.redis.set(self.prefix_redis_key('urlts', self.hash_url(url)),
                       '%d' % now)

    def get_crawl_timestamp(self, url):
        s = self.redis.get(self.prefix_redis_key('urlts', self.hash_url(url)))
        if s:
            return int(s)

    def get_existing_task_id(self, url):
        s = self.redis.get(self.prefix_redis_key('taskbyurl',
                                                 self.hash_url(url)))
        if s:
            task_id, high_priority_bool = s.split(',')
            return task_id, high_priority_bool == '1'
        else:
            return None, None

    def get_task(self, task_id):
        s = self.redis.get(task_id)
        return self.deserialize(s)

    def drop_task(self, task_id, high_priority):
        self.redis.zrem(self.pick_queue_key(high_priority), task_id)
        self.redis.delete(task_id)

    def push(self, task):
        """
        Schedule a new crawl task, which is an instance of ``Task``.

        If ``high_priority`` is set, drop all other tasks for the same
        canonicalized URL and enqueue to a separate high priority queue.

        If ``min_age`` is set and this task URL has been crawled before, check
        the last time it was crawled, and reschedule it to a timestamp which is
        ``min_age`` after the last crawl time.

        If a different crawl task is already scheduled for this URL, check the
        scheduled time. If it has an intentionally delayed crawl time and that
        scheduled time is *after* this task, drop it. Otherwise, drop this
        task. Either way, the earlier of the two tasks should be kept.
        """
        # XXX FIXME
        # We really should take a lock on this URL, to prevent race conditions
        # in the conditional logic below.

        existing_task_id, existing_is_hp = self.get_existing_task_id(task.url)

        if task.high_priority:
            # If this task is HP, drop any other task for this URL..
            if existing_task_id:
                self.drop_task(existing_task_id, existing_is_hp)

        elif existing_is_hp:
            # If there's an existing high priority task, drop this one.
            return

        else:
            # If min_age is set, check for previous crawls and further delay
            # this task in order to ensure that min_age is not violated.
            if task.min_age:
                last_timestamp = self.get_crawl_timestamp(task.url)
                if last_timestamp:
                    task.scheduled_timestamp = max(last_timestamp +
                                                   task.min_age,
                                                   task.scheduled_timestamp)

            # If there is an existing task, check if it's scheduled earlier.
            if existing_task_id:
                scheduled_ts = self.redis.zscore(
                    self.pick_queue_key(existing_is_hp),
                    existing_task_id)

                if scheduled_ts > task.scheduled_timestamp:
                    # If the existing task is later, drop it.
                    self.drop_task(existing_task_id, existing_is_hp)
                else:
                    # If this task is later, drop it and keep the existing
                    # task.
                    return

        # At this point we're ready to enqueue.
        s = self.serialize(task)
        task_id = hashlib.md5(s).hexdigest()

        # Store the task by ID.
        self.redis.set(task_id, s)

        # Finally, enqueue the task ID to the appropriate sorted set.
        self.redis.zadd(self.pick_queue_key(task.high_priority),
                        float(task.scheduled_timestamp), task_id)

    def pop_queue(self, high_priority, cutoff):
        key = self.pick_queue_key(high_priority)
        with self.redis.pipeline() as pipe:
            while True:
                try:
                    pipe.watch(key)
                    elements = pipe.zrangebyscore(key, min=0, max=cutoff,
                                                  start=0, num=1)
                    if elements:
                        task_id = elements[0]
                        pipe.multi()
                        pipe.zrem(key, task_id)
                        pipe.execute()
                        return task_id
                    else:
                        return
                except WatchError:
                    continue

    def pop(self):
        """
        Get the next scheduled crawl task, or raise Empty() if there is nothing
        to do.
        """
        now = time.time()
        for high_priority in (True, False):
            next_task_id = self.pop_queue(high_priority, now)
            if next_task_id:
                break

        # Nothing to do?
        if not next_task_id:
            raise Empty()

        # Otherwise, fetch the task and deserialize it.
        s = self.redis.get(next_task_id)
        self.redis.delete(next_task_id)
        task = self.deserialize(s)

        # Record a crawl timestamp for this URL.
        self.record_crawl_timestamp(task.url, now)

        # If the task has a repeat interval, reschedule it.
        if task.repeat_after:
            repeat_task = Task(url=task.url,
                               document_type=task.document_type,
                               repeat_after=task.repeat_after,
                               scheduled_timestamp=now + task.repeat_after)
            self.push(repeat_task)

        return task

    def count(self):
        """
        Return the number of currently enqueued tasks.
        """
        counts = []
        for high_priority in (True, False):
            key = self.pick_queue_key(high_priority)
            counts.append(self.redis.zcard(key))

        return {'high': counts[0],
                'normal': counts[1]}
<|endoftext|>


<commit_before># Copyright (C) 2012 Google Inc. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are
# met:
#
#     * Redistributions of source code must retain the above copyright
# notice, this list of conditions and the following disclaimer.
#     * Redistributions in binary form must reproduce the above
# copyright notice, this list of conditions and the following disclaimer
# in the documentation and/or other materials provided with the
# distribution.
#     * Neither the name of Google Inc. nor the names of its
# contributors may be used to endorse or promote products derived from
# this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
# A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
# OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
# THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

import logging
import optparse
import signal
import traceback

from webkitpy.common.host import Host
from webkitpy.layout_tests.models import test_expectations
from webkitpy.layout_tests.port import platform_options

# This mirrors what the shell normally does.
INTERRUPTED_EXIT_STATUS = signal.SIGINT + 128

# This is a randomly chosen exit code that can be tested against to
# indicate that an unexpected exception occurred.
EXCEPTIONAL_EXIT_STATUS = 254

_log = logging.getLogger(__name__)

def lint(host, options, logging_stream):
    logger = logging.getLogger()
    logger.setLevel(logging.INFO)
    handler = logging.StreamHandler(logging_stream)
    logger.addHandler(handler)

    try:
        ports_to_lint = [host.port_factory.get(name) for name in host.port_factory.all_port_names(options.platform)]
        files_linted = set()
        lint_failed = False

        for port_to_lint in ports_to_lint:
            expectations_dict = port_to_lint.expectations_dict()

            # FIXME: This won't work if multiple ports share a TestExpectations file but support different modifiers in the file.
            for expectations_file in expectations_dict.keys():
                if expectations_file in files_linted:
                    continue

                try:
                    test_expectations.TestExpectations(port_to_lint,
                        expectations_to_lint={expectations_file: expectations_dict[expectations_file]})
                except test_expectations.ParseError as e:
                    lint_failed = True
                    _log.error('')
                    for warning in e.warnings:
                        _log.error(warning)
                    _log.error('')
                files_linted.add(expectations_file)

        if lint_failed:
            _log.error('Lint failed.')
            return -1

        _log.info('Lint succeeded.')
        return 0
    finally:
        logger.removeHandler(handler)

def main(argv, _, stderr):
    parser = optparse.OptionParser(option_list=platform_options(use_globs=True))
    options, _ = parser.parse_args(argv)

    if options.platform and 'test' in options.platform:
        # It's a bit lame to import mocks into real code, but this allows the user
        # to run tests against the test platform interactively, which is useful for
        # debugging test failures.
        from webkitpy.common.host_mock import MockHost
        host = MockHost()
    else:
        host = Host()

    try:
        exit_status = lint(host, options, stderr)
    except KeyboardInterrupt:
        exit_status = INTERRUPTED_EXIT_STATUS
    except Exception as e:
        print >> stderr, '\n%s raised: %s' % (e.__class__.__name__, str(e))
        traceback.print_exc(file=stderr)
        exit_status = EXCEPTIONAL_EXIT_STATUS

    return exit_status
<commit_msg>make run-webkit-tests --chromium --lint-test-files work properly<commit_after># Copyright (C) 2012 Google Inc. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are
# met:
#
#     * Redistributions of source code must retain the above copyright
# notice, this list of conditions and the following disclaimer.
#     * Redistributions in binary form must reproduce the above
# copyright notice, this list of conditions and the following disclaimer
# in the documentation and/or other materials provided with the
# distribution.
#     * Neither the name of Google Inc. nor the names of its
# contributors may be used to endorse or promote products derived from
# this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
# A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
# OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
# THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

import logging
import optparse
import signal
import traceback

from webkitpy.common.host import Host
from webkitpy.layout_tests.models import test_expectations
from webkitpy.layout_tests.port import platform_options

# This mirrors what the shell normally does.
INTERRUPTED_EXIT_STATUS = signal.SIGINT + 128

# This is a randomly chosen exit code that can be tested against to
# indicate that an unexpected exception occurred.
EXCEPTIONAL_EXIT_STATUS = 254

_log = logging.getLogger(__name__)

def lint(host, options, logging_stream):
    logger = logging.getLogger()
    logger.setLevel(logging.INFO)
    handler = logging.StreamHandler(logging_stream)
    logger.addHandler(handler)

    try:
        # FIXME: Remove this when we remove the --chromium flag (crbug.com/245504).
        if options.platform == 'chromium':
            options.platform = None

        ports_to_lint = [host.port_factory.get(name) for name in host.port_factory.all_port_names(options.platform)]
        files_linted = set()
        lint_failed = False

        for port_to_lint in ports_to_lint:
            expectations_dict = port_to_lint.expectations_dict()

            # FIXME: This won't work if multiple ports share a TestExpectations file but support different modifiers in the file.
            for expectations_file in expectations_dict.keys():
                if expectations_file in files_linted:
                    continue

                try:
                    test_expectations.TestExpectations(port_to_lint,
                        expectations_to_lint={expectations_file: expectations_dict[expectations_file]})
                except test_expectations.ParseError as e:
                    lint_failed = True
                    _log.error('')
                    for warning in e.warnings:
                        _log.error(warning)
                    _log.error('')
                files_linted.add(expectations_file)

        if lint_failed:
            _log.error('Lint failed.')
            return -1

        _log.info('Lint succeeded.')
        return 0
    finally:
        logger.removeHandler(handler)

def main(argv, _, stderr):
    parser = optparse.OptionParser(option_list=platform_options(use_globs=True))
    options, _ = parser.parse_args(argv)

    if options.platform and 'test' in options.platform:
        # It's a bit lame to import mocks into real code, but this allows the user
        # to run tests against the test platform interactively, which is useful for
        # debugging test failures.
        from webkitpy.common.host_mock import MockHost
        host = MockHost()
    else:
        host = Host()

    try:
        exit_status = lint(host, options, stderr)
    except KeyboardInterrupt:
        exit_status = INTERRUPTED_EXIT_STATUS
    except Exception as e:
        print >> stderr, '\n%s raised: %s' % (e.__class__.__name__, str(e))
        traceback.print_exc(file=stderr)
        exit_status = EXCEPTIONAL_EXIT_STATUS

    return exit_status
<|endoftext|>


<commit_before># the following specifies the current package version
__version__ = "0.1.1"
<commit_msg>Bumped version<commit_after># the following specifies the current package version
__version__ = "0.1.2"
<|endoftext|>


<commit_before>"""
Filename:     regrid.py
Author:       Damien Irving, irving.damien@gmail.com
Description:  Regrid input data
"""

# Import general Python modules

import sys, os, pdb
import argparse
import numpy
import iris
import cmdline_provenance as cmdprov

# Import my modules

cwd = os.getcwd()
repo_dir = '/'
for directory in cwd.split('/')[1:]:
    repo_dir = os.path.join(repo_dir, directory)
    if directory == 'ocean-analysis':
        break

modules_dir = os.path.join(repo_dir, 'modules')
sys.path.append(modules_dir)

try:
    import timeseries
    import grids
    import general_io as gio
except ImportError:
    raise ImportError('Must run this script from anywhere within the ocean-analysis git repo')

# Define functions

def check_inputs(inargs):
    """Check inputs"""

    if inargs.lats:
        assert inargs.lons, 'must provide new lats and lons'

    if inargs.lons:
        assert inargs.lats, 'must provide new lats and lons'

def get_dim_vals(user_input, bounds=False):
    """Get the list of values for a data dimension/axis.

    Args:
      bounds (bool): input represents the bounds, not centre points

    """

    vals = user_input
    if user_input:
        if len(user_input) == 3:
            start, stop, step = user_input
            vals = list(numpy.arange(start, stop + step, step))
        if bounds:
            vals = numpy.array(vals)
            vals = (vals[1:] + vals[:-1]) / 2

    return list(vals)

def get_sample_points(cube, dim_vals):
    """Define the sample points for interpolation"""

    sample_points = []
    coord_names = [coord.name() for coord in cube.dim_coords]
    if 'time' in coord_names:
        coord_names.remove('time')

    for coord in coord_names:
        if dim_vals[coord]:
            sample_points.append((coord, dim_vals[coord]))
        else:
            sample_points.append((coord, cube.coord(coord).points))

    return sample_points

def get_depth_bounds(depth_intervals):
    """Create an array of bounds pairs"""

    depth_bounds = []
    for index in range(len(depth_intervals) - 1):
        depth_bounds.append([depth_intervals[index], depth_intervals[index + 1]])

    return numpy.array(depth_bounds)

def remove_nans(cube):
    """Turn NaN masked fill value."""

    nan_locations = numpy.argwhere(numpy.isnan(cube.data))
    for loc in nan_locations:
        cube.data[tuple(loc)] = cube.data.fill_value
        cube.data.mask[tuple(loc)] = True

    print('There was %s NaNs' %(str(len(nan_locations))) )

    return cube

def guess_depth_bounds(cube):
    """Guess the depth coordinate bounds."""

    cube.coord('depth').guess_bounds()
    bounds = cube.coord('depth').bounds.copy()
    if bounds[0][0] < 0.0:
        bounds[0][0] = 0.0
        cube.coord('depth').bounds = bounds

    return cube

def main(inargs):
    """Run the program."""

    cube = iris.load_cube(inargs.infile, inargs.var)
    cube = gio.check_time_units(cube)
    cube = timeseries.convert_to_annual(cube)
    log = cmdprov.new_log(infile_history={inargs.infile: cube.attributes['history']}, git_repo=repo_dir)

    dim_vals = {}
    dim_vals['latitude'] = get_dim_vals(inargs.lats)
    dim_vals['longitude'] = get_dim_vals(inargs.lons)
    if inargs.levs:
        dim_vals['depth'] = get_dim_vals(inargs.levs)
    else:
        dim_vals['depth'] = get_dim_vals(inargs.depth_bnds, bounds=True)

    # Regrid from curvilinear to rectilinear if necessary
    regrid_status = False
    if inargs.lats:
        horizontal_grid = grids.make_grid(dim_vals['latitude'], dim_vals['longitude'])
        cube, coord_names, regrid_status = grids.curvilinear_to_rectilinear(cube, target_grid_cube=horizontal_grid)

    # Regrid to new grid
    if dim_vals['depth'] or not regrid_status:
        sample_points = get_sample_points(cube, dim_vals)
        cube = cube.interpolate(sample_points, iris.analysis.Linear())
        cube.coord('latitude').guess_bounds()
        cube.coord('longitude').guess_bounds()
        if inargs.levs:
            cube = guess_depth_bounds(cube)
        else:
            cube.coord('depth').bounds = get_depth_bounds(inargs.depth_bnds)

    if numpy.isnan(numpy.min(cube.data)):
        cube = remove_nans(cube)

    cube.attributes['history'] = log
    iris.save(cube, inargs.outfile)

if __name__ == '__main__':

    extra_info ="""
example:

author:
    Damien Irving, irving.damien@gmail.com

"""

    description='Apply annual timescale smoothing'
    parser = argparse.ArgumentParser(description=description,
                                     epilog=extra_info,
                                     argument_default=argparse.SUPPRESS,
                                     formatter_class=argparse.RawDescriptionHelpFormatter)

    parser.add_argument("infile", type=str, help="Input file")
    parser.add_argument("var", type=str, help="Variable name")
    parser.add_argument("outfile", type=str, help="Output file")

    parser.add_argument("--lats", type=float, nargs='*', default=None,
                        help="list all latitude points of new grid or three values: start, stop, interval")
    parser.add_argument("--lons", type=float, nargs='*', default=None,
                        help="list all longitude points of new grid or three values: start, stop, interval")
    parser.add_argument("--levs", type=float, nargs='*', default=None,
                        help="list all depth points of new grid or three values: start, stop, interval")
    parser.add_argument("--depth_bnds", type=float, nargs='*', default=None,
                        help="list of the depth bounds of new grid or three values: start, stop, interval")

    parser.add_argument("--annual", action="store_true", default=False,
                        help="Convert data to annual timescale [default: False]")

    args = parser.parse_args()
    check_inputs(args)

main(args)
<commit_msg>Added chunking option<commit_after>"""
Filename:     regrid.py
Author:       Damien Irving, irving.damien@gmail.com
Description:  Regrid input data
"""

# Import general Python modules

import sys, os, pdb
import argparse
import numpy
import iris
import cmdline_provenance as cmdprov

# Import my modules

cwd = os.getcwd()
repo_dir = '/'
for directory in cwd.split('/')[1:]:
    repo_dir = os.path.join(repo_dir, directory)
    if directory == 'ocean-analysis':
        break

modules_dir = os.path.join(repo_dir, 'modules')
sys.path.append(modules_dir)

try:
    import timeseries
    import grids
    import general_io as gio
except ImportError:
    raise ImportError('Must run this script from anywhere within the ocean-analysis git repo')

# Define functions

def check_inputs(inargs):
    """Check inputs"""

    if inargs.lats:
        assert inargs.lons, 'must provide new lats and lons'

    if inargs.lons:
        assert inargs.lats, 'must provide new lats and lons'

def get_dim_vals(user_input, bounds=False):
    """Get the list of values for a data dimension/axis.

    Args:
      bounds (bool): input represents the bounds, not centre points

    """

    vals = user_input
    if user_input:
        if len(user_input) == 3:
            start, stop, step = user_input
            vals = list(numpy.arange(start, stop + step, step))
        if bounds:
            vals = numpy.array(vals)
            vals = (vals[1:] + vals[:-1]) / 2

    return list(vals)

def get_sample_points(cube, dim_vals):
    """Define the sample points for interpolation"""

    sample_points = []
    coord_names = [coord.name() for coord in cube.dim_coords]
    if 'time' in coord_names:
        coord_names.remove('time')

    for coord in coord_names:
        if dim_vals[coord]:
            sample_points.append((coord, dim_vals[coord]))
        else:
            sample_points.append((coord, cube.coord(coord).points))

    return sample_points

def get_depth_bounds(depth_intervals):
    """Create an array of bounds pairs"""

    depth_bounds = []
    for index in range(len(depth_intervals) - 1):
        depth_bounds.append([depth_intervals[index], depth_intervals[index + 1]])

    return numpy.array(depth_bounds)

def remove_nans(cube):
    """Turn NaN masked fill value."""

    nan_locations = numpy.argwhere(numpy.isnan(cube.data))
    for loc in nan_locations:
        cube.data[tuple(loc)] = cube.data.fill_value
        cube.data.mask[tuple(loc)] = True

    print('There was %s NaNs' %(str(len(nan_locations))) )

    return cube

def guess_depth_bounds(cube):
    """Guess the depth coordinate bounds."""

    cube.coord('depth').guess_bounds()
    bounds = cube.coord('depth').bounds.copy()
    if bounds[0][0] < 0.0:
        bounds[0][0] = 0.0
        cube.coord('depth').bounds = bounds

    return cube

def main(inargs):
    """Run the program."""

    cube = iris.load_cube(inargs.infile, inargs.var)
    cube = gio.check_time_units(cube)
    cube = timeseries.convert_to_annual(cube, chunk=inargs.chunk)
    log = cmdprov.new_log(infile_history={inargs.infile: cube.attributes['history']}, git_repo=repo_dir)

    dim_vals = {}
    dim_vals['latitude'] = get_dim_vals(inargs.lats)
    dim_vals['longitude'] = get_dim_vals(inargs.lons)
    if inargs.levs:
        dim_vals['depth'] = get_dim_vals(inargs.levs)
    else:
        dim_vals['depth'] = get_dim_vals(inargs.depth_bnds, bounds=True)

    # Regrid from curvilinear to rectilinear if necessary
    regrid_status = False
    if inargs.lats:
        horizontal_grid = grids.make_grid(dim_vals['latitude'], dim_vals['longitude'])
        cube, coord_names, regrid_status = grids.curvilinear_to_rectilinear(cube, target_grid_cube=horizontal_grid)

    # Regrid to new grid
    if dim_vals['depth'] or not regrid_status:
        sample_points = get_sample_points(cube, dim_vals)
        cube = cube.interpolate(sample_points, iris.analysis.Linear())
        cube.coord('latitude').guess_bounds()
        cube.coord('longitude').guess_bounds()
        if inargs.levs:
            cube = guess_depth_bounds(cube)
        else:
            cube.coord('depth').bounds = get_depth_bounds(inargs.depth_bnds)

    if numpy.isnan(numpy.min(cube.data)):
        cube = remove_nans(cube)

    cube.attributes['history'] = log
    iris.save(cube, inargs.outfile)

if __name__ == '__main__':

    extra_info ="""
example:

author:
    Damien Irving, irving.damien@gmail.com

"""

    description='Apply annual timescale smoothing'
    parser = argparse.ArgumentParser(description=description,
                                     epilog=extra_info,
                                     argument_default=argparse.SUPPRESS,
                                     formatter_class=argparse.RawDescriptionHelpFormatter)

    parser.add_argument("infile", type=str, help="Input file")
    parser.add_argument("var", type=str, help="Variable name")
    parser.add_argument("outfile", type=str, help="Output file")

    parser.add_argument("--lats", type=float, nargs='*', default=None,
                        help="list all latitude points of new grid or three values: start, stop, interval")
    parser.add_argument("--lons", type=float, nargs='*', default=None,
                        help="list all longitude points of new grid or three values: start, stop, interval")
    parser.add_argument("--levs", type=float, nargs='*', default=None,
                        help="list all depth points of new grid or three values: start, stop, interval")
    parser.add_argument("--depth_bnds", type=float, nargs='*', default=None,
                        help="list of the depth bounds of new grid or three values: start, stop, interval")

    parser.add_argument("--annual", action="store_true", default=False,
                        help="Convert data to annual timescale [default: False]")
    parser.add_argument("--chunk", action="store_true", default=False,
                        help="Chunk annual timescale conversion to avoid memory errors [default: False]")

    args = parser.parse_args()
    check_inputs(args)

main(args)
<|endoftext|>


<commit_before>import calendar
import datetime

from elasticutils import S

from es import ES

def fit_time_range(start, end):
    """
    Keeps an empty time range empty. Stretches an incomplete
    time range.
    """
    if start == None and end == None:
        return (None, None)
    if start == None:
        start = datetime.datetime(year=1970, month=1, day=1)
    if end == None:
        end = datetime.datetime.now()

    return (start, end)

def apply_time_filter(query, start, end):
    start, end = fit_time_range(start, end)
    if start and end:
        return query.filter(created_at__range=(start, end))
    return query

def past_n_months(index, query, n):
    """
    Maps a query over time intervals corresponding to the past n months.
    Returns a list of objects like {'month': name, 'data': query_data}.
    """
    data = []
    now = datetime.datetime.now()
    year = now.year
    month = now.month
    for i in range(n):
        month = month - 1
        # handle first months of the year special case
        if  month == 0:
            month = 12
            year -= 1
        start = datetime.date(year=year, month=month, day=1)
        last_day = calendar.monthrange(year, month)[1]
        end = datetime.date(year=year, month=month, day=last_day)

        month_data = {
                'month': start.strftime('%B'),
                'value': query(index=index, start=start, end=end)
        }
        data.append(month_data)
    return data

def facet_counts_all(query, field):
    """
    Returns the facet counts for the equivalent .facet(field) query
    but grabs all the results (size = big).
    """
    # https://github.com/mozilla/fjord/blob/master/fjord/analytics/views.py#L527
    ALL = 2**31 - 1 # es/java maxint
    facet = {
        'terms': {
            'field': field,
            'size': ALL
        }
    }
    q = query._build_query()
    if 'filter' in q:
        facet['facet_filter'] = q['filter']

    return query.facet_raw(f=facet).facet_counts()['f']

def most_active_people(index, start=None, end=None):
    """
    Finds the most active users - as actors in all the events.

    Returns a list of dicts like:
    {'count': N, 'term': NAME}
    """
    q = S().indexes(index)
    q = apply_time_filter(q, start, end)
    return q.facet('actor.login').facet_counts()['actor.login']

def total_events(index, start=None, end=None):
    """
    Returns the number of total events for the given time
    interval, or for all the data if no interval is given.
    """
    q = S().indexes(index)
    q = apply_time_filter(q, start, end)
    return q.count()

def most_active_issues(index, start=None, end=None):
    """
    Finds the most active issues - by total number of events.
    """
    q = S().indexes(index).doctypes('IssuesEvent', 'IssueCommentEvent')
    q = apply_time_filter(q, start, end)
    return q.facet('payload.issue.number').facet_counts()['payload.issue.number']

def issues_without_comments(index):
    """
    Open issues with no comments.
    """
    # todo - facet count needs to be bigger than 10
    issues = open_issues(index)

    with_comments = S().indexes(index).doctypes('IssueCommentEvent')
    with_comments = facet_counts_all(with_comments, 'payload.issue.number')
    with_comments = [r['term'] for r in with_comments]

    return list(set(issues) - set(with_comments))

def issues_assigned_to(index, login):
    """
    List of open issues assigned to {{ login }}.
    """
    q = S().indexes(index).doctypes('IssuesEvent', 'IssueCommentEvent')
    assigned = q.filter(**{'payload.issue.assignee.login': login})
    assigned = facet_counts_all(assigned, 'payload.issue.number')
    assigned = [r['term'] for r in assigned]

    # keep open ones
    issues = open_issues(index)
    return list(set(assigned).intersection(set(issues)))

def recent_events(index, count=200, starting_from=0):
    """
    Returns the <count> most recent events, starting from
    index <starting_from>.
    """
    q = S().indexes(index).order_by('-created_at')
    q = q[starting_from : starting_from + count]
    q = q.values_dict()
    return list(q)

def available_repos():
    EXCLUDED = ['kibana-int']
    # get all indices
    indices = ES.indices.get_aliases().keys()
    # filter out those that don't look as repos
    indices = [i for i in indices if len(i.split('-')) == 2]
    # there are some that still sneak by
    indices = [i for i in indices if i not in EXCLUDED]

    # format them as github repos (slashes)
    formatted = [i.replace('-', '/') for i in indices]
    return formatted

def issue_events_count(index, action, start=None, end=None):
    # action can be 'opened' or 'closed'
    q = S().indexes(index).doctypes('IssuesEvent')
    q = apply_time_filter(q, start, end)
    q = q.filter(**{'payload.action': action})
    return q.count()

def issues_count(index, state):
    # state can be 'open' or 'closed'
    q = S().indexes(index).doctypes('IssueData')
    q = q.filter(state=state)
    return q.count()

def pulls_count(index):
    # we only have open pull requests
    q = S().indexes(index).doctypes('PullRequestData')
    return q.count()
<commit_msg>reverse past_months results for more intuitive reading<commit_after>import calendar
import datetime

from elasticutils import S

from es import ES

def fit_time_range(start, end):
    """
    Keeps an empty time range empty. Stretches an incomplete
    time range.
    """
    if start == None and end == None:
        return (None, None)
    if start == None:
        start = datetime.datetime(year=1970, month=1, day=1)
    if end == None:
        end = datetime.datetime.now()

    return (start, end)

def apply_time_filter(query, start, end):
    start, end = fit_time_range(start, end)
    if start and end:
        return query.filter(created_at__range=(start, end))
    return query

def past_n_months(index, query, n):
    """
    Maps a query over time intervals corresponding to the past n months.
    Returns a list of objects like {'month': name, 'data': query_data}.
    """
    data = []
    now = datetime.datetime.now()
    year = now.year
    month = now.month
    for i in range(n):
        month = month - 1
        # handle first months of the year special case
        if  month == 0:
            month = 12
            year -= 1
        start = datetime.date(year=year, month=month, day=1)
        last_day = calendar.monthrange(year, month)[1]
        end = datetime.date(year=year, month=month, day=last_day)

        month_data = {
                'month': start.strftime('%B'),
                'value': query(index=index, start=start, end=end)
        }
        data.append(month_data)
    data.reverse()
    return data

def facet_counts_all(query, field):
    """
    Returns the facet counts for the equivalent .facet(field) query
    but grabs all the results (size = big).
    """
    # https://github.com/mozilla/fjord/blob/master/fjord/analytics/views.py#L527
    ALL = 2**31 - 1 # es/java maxint
    facet = {
        'terms': {
            'field': field,
            'size': ALL
        }
    }
    q = query._build_query()
    if 'filter' in q:
        facet['facet_filter'] = q['filter']

    return query.facet_raw(f=facet).facet_counts()['f']

def most_active_people(index, start=None, end=None):
    """
    Finds the most active users - as actors in all the events.

    Returns a list of dicts like:
    {'count': N, 'term': NAME}
    """
    q = S().indexes(index)
    q = apply_time_filter(q, start, end)
    return q.facet('actor.login').facet_counts()['actor.login']

def total_events(index, start=None, end=None):
    """
    Returns the number of total events for the given time
    interval, or for all the data if no interval is given.
    """
    q = S().indexes(index)
    q = apply_time_filter(q, start, end)
    return q.count()

def most_active_issues(index, start=None, end=None):
    """
    Finds the most active issues - by total number of events.
    """
    q = S().indexes(index).doctypes('IssuesEvent', 'IssueCommentEvent')
    q = apply_time_filter(q, start, end)
    return q.facet('payload.issue.number').facet_counts()['payload.issue.number']

def issues_without_comments(index):
    """
    Open issues with no comments.
    """
    # todo - facet count needs to be bigger than 10
    issues = open_issues(index)

    with_comments = S().indexes(index).doctypes('IssueCommentEvent')
    with_comments = facet_counts_all(with_comments, 'payload.issue.number')
    with_comments = [r['term'] for r in with_comments]

    return list(set(issues) - set(with_comments))

def issues_assigned_to(index, login):
    """
    List of open issues assigned to {{ login }}.
    """
    q = S().indexes(index).doctypes('IssuesEvent', 'IssueCommentEvent')
    assigned = q.filter(**{'payload.issue.assignee.login': login})
    assigned = facet_counts_all(assigned, 'payload.issue.number')
    assigned = [r['term'] for r in assigned]

    # keep open ones
    issues = open_issues(index)
    return list(set(assigned).intersection(set(issues)))

def recent_events(index, count=200, starting_from=0):
    """
    Returns the <count> most recent events, starting from
    index <starting_from>.
    """
    q = S().indexes(index).order_by('-created_at')
    q = q[starting_from : starting_from + count]
    q = q.values_dict()
    return list(q)

def available_repos():
    EXCLUDED = ['kibana-int']
    # get all indices
    indices = ES.indices.get_aliases().keys()
    # filter out those that don't look as repos
    indices = [i for i in indices if len(i.split('-')) == 2]
    # there are some that still sneak by
    indices = [i for i in indices if i not in EXCLUDED]

    # format them as github repos (slashes)
    formatted = [i.replace('-', '/') for i in indices]
    return formatted

def issue_events_count(index, action, start=None, end=None):
    # action can be 'opened' or 'closed'
    q = S().indexes(index).doctypes('IssuesEvent')
    q = apply_time_filter(q, start, end)
    q = q.filter(**{'payload.action': action})
    return q.count()

def issues_count(index, state):
    # state can be 'open' or 'closed'
    q = S().indexes(index).doctypes('IssueData')
    q = q.filter(state=state)
    return q.count()

def pulls_count(index):
    # we only have open pull requests
    q = S().indexes(index).doctypes('PullRequestData')
    return q.count()
<|endoftext|>