<commit_before>"""A simple game of chess
TODO:
        check for check
        add win conditions
        restrict each side to only being able to move pieces of their own colour
        handle user input more elegantly
        handle bad user input (i.e., non-ints, or giving coordinates not on the board)
        add special cases (castling, promotion, en passant)

        bugs: currently valid to select a square of the board with no pieces on it
"""
import sys
import math

board_size = 8
empty_square = "-"
move_piece_located_at = [0,0]
move_piece_to = [0,0]
pieces_in_play = []
captured_pieces = []
board = [[empty_square for j in range(board_size)] for i in range(board_size)]

def update_board():
    """Check if there is a piece located at each square on the board. if there is, add the piece's icon there"""
    for piece in pieces_in_play:
        for row in range(board_size):
            for column in range(board_size):
                if piece.row == row and piece.column == column:
                    board[row][column] = piece.icon
    print_board()

def print_board():
    """Prints out the board 1 square at a time, with row and column numbering"""
    for row in range(board_size):
        print(row, end = " ")
        for column in range(board_size):
            print(board[row][column], end = " ")
        print()
    print(" ", end = " ")
    for column in range(board_size):
        print(column, end = " ")
    print()

def piece_on_square(row,column):
    """If there is a piece on the square of the board, return true"""

    for piece in pieces_in_play:
        if piece.row == row and piece.column == column:
            return True
    else:
        return False

def get_piece_on_square(row, column):
    """Returns piece on a particular square """
    # need to reuse this more often

    for piece in pieces_in_play:
        if piece.row == row and piece.column == column:
            return piece

def capture_piece(row,column):
    """Finds the piece to be captured, removes it from play, and places it in another list for easy printing later on """
    # needs checks for colours!!
    # needs exception for kings!!

    piece = get_piece_on_square(row, column)
    captured_pieces.append(piece)
    pieces_in_play.remove(piece)
    print("Captured ", captured_pieces[-1].icon)

class Piece:
    """Pieces are initialised with a name, colour and icon.
    Additionally they have attributes row and column that correspond to their
    current location, and the attributes proposed_row and proposed_column which
    correspond to coordinates of their proposed move (and are initialisd at 0,0)
    """
    proposed_row = 0
    proposed_column = 0
    def __init__(self,name,colour,icon,row,column,):
        self.name = name
        self.colour = colour
        self.icon = icon
        self.row = row
        self.column = column

def create_pawns():
    """Places two rows of pawns; one for each colour"""
    for i in range(board_size):
        pieces_in_play.append(Piece("pawn","white","♙",6,i,))
        pieces_in_play.append(Piece("pawn","black","♟",1,i,))

def create_other_pieces():
    """Places all pieces which appear twice on the board. Their placement is
    mirrored across the columns of the board"""
    for i in range(2):
        pieces_in_play.append(Piece("rook","white","♖",7,0 + (i * (board_size - 1)),))
        pieces_in_play.append(Piece("rook","black","♜",0,0 + (i * (board_size - 1)),))
        pieces_in_play.append(Piece("knight","white","♘",7,1 + (i * (board_size - 3)),))
        pieces_in_play.append(Piece("knight","black","♞",0,1 + (i * (board_size - 3)),))
        pieces_in_play.append(Piece("bishop","white","♗",7,2 + (i * (board_size - 5)),))
        pieces_in_play.append(Piece("bishop","black","♝",0,2 + (i * (board_size - 5)),))

def create_royalty():
    """Places kings and queens in specified squares"""
    pieces_in_play.append(Piece("king","white","♔",7,4,))
    pieces_in_play.append(Piece("king","black","♚",0,4,))
    pieces_in_play.append(Piece("queen","white","♕",7,3,))
    pieces_in_play.append(Piece("queen","black","♛",0,3,))

def create_pieces():
    """Calls all the functions required to populate the board"""
    create_pawns()
    create_other_pieces()
    create_royalty()

def move_piece():
    """Decides whether or not to move the piece, based on the coordinates
    inputted by the user and move_legal and collision_detected. In the future,
    this will also depend on a check_for_check function.
    If a move is made then the old square is made empty"""
    get_user_input()
    piece = get_piece_on_square(move_piece_located_at[0], move_piece_located_at[1])
    piece.proposed_row = move_piece_to[0]
    piece.proposed_column = move_piece_to[1]
    if move_legal(piece) and not collision_detected(piece):
        board[piece.row][piece.column] = empty_square
        if piece_on_square(piece.proposed_row,piece.proposed_column):
            capture_piece(piece.proposed_row,piece.proposed_column)
        piece.row = piece.proposed_row
        piece.column = piece.proposed_column
    else:
        print("Sorry, that move is invalid")
        move_piece()

def get_user_input():
    """Takes user input, assuming a "x,y" format. Stores the 4 coordinates in 2 universal variables"""
    # doesn't currently validate the string supplied
    user_choice = input("Choose a piece (row,column)")
    if "quit" in user_choice.lower():
        sys.exit("Exiting...")
    move_piece_located_at[0] = int(user_choice[0])
    move_piece_located_at[1] = int(user_choice[2])
    user_move = input("Choose where to move the piece (row,column)")
    if "quit" in user_move.lower():
        sys.exit("Exiting...2")
    move_piece_to[0] = int(user_move[0])
    move_piece_to[1] = int(user_move[2])

def move_legal(piece):
    """Calls relevant methods for determining legality"""
    if piece.name == "bishop":
        return diagonal_move_legal(piece)
    elif piece.name == "rook":
        return lateral_move_legal(piece)
    elif piece.name == "queen":
        # If the queen behaves like either a bishop or a rook, return true
        return (lateral_move_legal(piece) or diagonal_move_legal(piece))
    elif piece.name == "pawn":
        return pawn_move_legal(piece)
    elif piece.name == "knight":
        return knight_move_legal(piece)
    elif piece.name == "king":
        return king_move_legal(piece)
    else:
        return False

def diagonal_move_legal(piece):
    """A bishop has to move along diagonals, so the abs deltaX and the abs deltaY should be equal """
    if abs(piece.row - piece.proposed_row) == abs(piece.column - piece.proposed_column):
        return True

def lateral_move_legal(piece):
    """A rook moves either along a row or a column, but not both """
    if ((piece.row - piece.proposed_row == 0) and (abs(piece.column - piece.proposed_column) > 0)) or ((abs(piece.row - piece.proposed_row) > 0) and (piece.column - piece.proposed_column == 0)):
        return True

def pawn_move_legal(piece):
    """Determines legality of pawn moves, taking into account their ability to
    move two rows from their starting square and their need to move diagonally
    to capture a piece. Does not currently handle en passant or promotion."""
    # could probably get rid of black/white checking with some additional maths. not sure if worth it
    if piece.colour == "black":
        if piece.row == 1 and piece.proposed_row == 3 and piece.column == piece.proposed_column:
            # collision code
            if (not piece_on_square(2,piece.column)) or (not piece_on_square(3,piece.column)):
                return True
        elif piece.row == piece.proposed_row - 1 and piece.column == piece.proposed_column:
            if (not piece_on_square(piece.proposed_row,piece.proposed_column)):
                return True
        elif piece.row == piece.proposed_row - 1 and abs(piece.column - piece.proposed_column) == 1:
            # we're capturing a piece here
            # and sort of doing collision detection. may want to re-think structure
            if piece_on_square(piece.proposed_row,piece.proposed_column) and get_piece_on_square(piece.proposed_row, piece.proposed_column).colour != piece.colour:
                return True
        else:
            return False

    elif piece.colour == "white":
        if piece.row == 6 and piece.proposed_row == 4 and piece.column == piece.proposed_column:
            if (not piece_on_square(5,piece.column)) or (not piece_on_square(4,piece.column)):
                return True
        elif piece.row == piece.proposed_row + 1 and piece.column == piece.proposed_column:
            if (not piece_on_square(piece.proposed_row,piece.proposed_column)):
                return True
        elif piece.row == piece.proposed_row + 1 and abs(piece.column - piece.proposed_column) == 1:
            if piece_on_square(piece.proposed_row,piece.proposed_column) and get_piece_on_square(piece.proposed_row, piece.proposed_column).colour != piece.colour:
                return True
        else:
            return False

def knight_move_legal(piece):
    """Determines knight move legality"""
    if (abs(piece.row - piece.proposed_row) == 2 and abs(piece.column - piece.proposed_column) == 1) or (abs(piece.row - piece.proposed_row) == 1 and abs(piece.column - piece.proposed_column) == 2):
        return True
    else:
        return False

def king_move_legal(piece):
    """Determines king move legality. The check status of the king is not
    currently taken into account"""
    if (abs(piece.row - piece.proposed_row) == 1 and abs(piece.column - piece.proposed_column) == 0) or (abs(piece.row - piece.proposed_row) == 1 and abs(piece.column - piece.proposed_column) == 1) or (abs(piece.row - piece.proposed_row) == 0 and abs(piece.column - piece.proposed_column) == 1):
        return True
    else:
        return False

def collision_detected(piece):
    """Calls the relevant methods to check for collision of different types
    of pieces"""
    # want to check all squares between start and finish - 1
    # also check final square to see if piece is of SAME colour and return False
    # pawn collision currently implemented in pawn_move_legal()
    if piece.name == "rook":
        return rook_collision_detected(piece)
    elif piece.name == "bishop":
        return bishop_collision_detected(piece)
    elif piece.name == "queen":
        # queens behave either as rooks or bishops
        return (rook_collision_detected(piece) or bishop_collision_detected(piece))
    elif piece.name == "knight":
        # knights only care about the square they land on
        return landing_square_collision_detected(piece)
    elif piece.name == "king":
        # kings care about the square they land on, and not moving into check
        # checking for check is not currently implemented
        return landing_square_collision_detected(piece)
    else:
        return False

def rook_collision_detected(piece):
    """Checks for rook collisions, taking into account that they can move in
    one of four directions"""
    # might be able to use max() and min() here rather than .copysign(). not sure if that would come in handy for bishop collision detection
    signed_step_column = int(math.copysign(1,piece.proposed_column - piece.column))
    signed_step_row = int(math.copysign(1,piece.proposed_row - piece.row))

    # check all the squares between start and destination (exclusive). return True if any piece is on these squares
    if piece.row - piece.proposed_row == 0:
        for column in range(piece.column + signed_step_column, piece.proposed_column, signed_step_column):
            if piece_on_square(piece.row, column):
                return True
    elif piece.column - piece.proposed_column == 0:
        for row in range(piece.row + signed_step_row, piece.proposed_row, signed_step_row):
            if piece_on_square(row, piece.column):
                return True
    else:
        landing_square_collision_detected(piece)

def bishop_collision_detected(piece):
    """Checks for bishop collision, taking into account that they can move in
    one of four directions"""
    signed_step_column = int(math.copysign(1,piece.proposed_column - piece.column))
    signed_step_row = int(math.copysign(1,piece.proposed_row - piece.row))

    for row in range(piece.row + signed_step_row, piece.proposed_row, signed_step_row):
        for column in range(piece.column + signed_step_column, piece.proposed_column, signed_step_column):
            if abs(piece.row - row) == abs(piece.column - column):
                if piece_on_square(row, column):
                    return True
    else:
        landing_square_collision_detected(piece)

def landing_square_collision_detected(piece):
    """Determines whether there is a piece on the final square of the move,
    and allows the move only if the piece is of the opposite colour. Does
    not currently allow for check."""
    if not piece_on_square(piece.proposed_row, piece.proposed_column):
        return False
    # allow movement to a square only with a piece of opposite colour on it
    # need to check for kings here
    elif get_piece_on_square(piece.proposed_row, piece.proposed_column).colour != piece.colour:
        return False
    else:
        return True

def main():
    """Calls the relevant methods to set up the board, and then loops each move
    of the game. Win conditions are not currently implemented."""
    game_over = False
    turn_counter = 1
    print("Remember to set your terminal to black text on a white background!\n")
    create_pieces()
    update_board()

    while not game_over:
        if turn_counter % 2 == 1:
            print("White's turn")
        else:
            print("Black's turn")
        move_piece()
        update_board()
        turn_counter += 1

if __name__ == "__main__":
    main()
<commit_msg>small coding style cleanups<commit_after>"""A simple game of chess
TODO:
        check for check
        add win conditions
        restrict each side to only being able to move pieces of their own colour
        handle user input more elegantly
        handle bad user input (i.e., non-ints, or giving coordinates not on the board)
        add special cases (castling, promotion, en passant)

        bugs: currently valid to select a square of the board with no pieces on it
"""
import sys
import math

board_size = 8
empty_square = "-"
move_piece_located_at = [0, 0]
move_piece_to = [0, 0]
pieces_in_play = []
captured_pieces = []
board = [[empty_square for j in range(board_size)] for i in range(board_size)]

def update_board():
    """Check if there is a piece located at each square on the board. if there is, add the piece's icon there"""
    for piece in pieces_in_play:
        for row in range(board_size):
            for column in range(board_size):
                if piece.row == row and piece.column == column:
                    board[row][column] = piece.icon
    print_board()

def print_board():
    """Prints out the board 1 square at a time, with row and column numbering"""
    for row in range(board_size):
        print(row, end = " ")
        for column in range(board_size):
            print(board[row][column], end = " ")
        print()
    print(" ", end = " ")
    for column in range(board_size):
        print(column, end = " ")
    print()

def piece_on_square(row, column):
    """If there is a piece on the square of the board, return true"""
    for piece in pieces_in_play:
        if piece.row == row and piece.column == column:
            return True
    else:
        return False

def get_piece_on_square(row, column):
    """Returns piece on a particular square """
    # need to reuse this more often
    for piece in pieces_in_play:
        if piece.row == row and piece.column == column:
            return piece

def capture_piece(row, column):
    """Finds the piece to be captured, removes it from play, and places it in another list for easy printing later on """
    # needs checks for colours!!
    # needs exception for kings!!
    piece = get_piece_on_square(row, column)
    captured_pieces.append(piece)
    pieces_in_play.remove(piece)
    print("Captured %s" % captured_pieces[-1].icon)

class Piece:
    """Pieces are initialised with a name, colour and icon.
    Additionally they have attributes row and column that correspond to their
    current location, and the attributes proposed_row and proposed_column which
    correspond to coordinates of their proposed move (and are initialisd at 0,0)
    """
    proposed_row = 0
    proposed_column = 0
    def __init__(self, name, colour, icon, row, column,):
        self.name = name
        self.colour = colour
        self.icon = icon
        self.row = row
        self.column = column

def create_pawns():
    """Places two rows of pawns; one for each colour"""
    for i in range(board_size):
        pieces_in_play.append(Piece("pawn", "white", "♙", 6, i))
        pieces_in_play.append(Piece("pawn", "black", "♟", 1, i))

def create_other_pieces():
    """Places all pieces which appear twice on the board. Their placement is
    mirrored across the columns of the board"""
    for i in range(2):
        pieces_in_play.append(Piece("rook", "white", "♖", 7, 0 + (i * (board_size - 1)),))
        pieces_in_play.append(Piece("rook", "black", "♜", 0, 0 + (i * (board_size - 1)),))
        pieces_in_play.append(Piece("knight", "white", "♘", 7, 1 + (i * (board_size - 3)),))
        pieces_in_play.append(Piece("knight", "black", "♞", 0, 1 + (i * (board_size - 3)),))
        pieces_in_play.append(Piece("bishop", "white", "♗", 7, 2 + (i * (board_size - 5)),))
        pieces_in_play.append(Piece("bishop", "black", "♝", 0, 2 + (i * (board_size - 5)),))

def create_royalty():
    """Places kings and queens in specified squares"""
    pieces_in_play.append(Piece("king", "white", "♔", 7, 4,))
    pieces_in_play.append(Piece("king", "black", "♚", 0, 4,))
    pieces_in_play.append(Piece("queen", "white", "♕", 7, 3,))
    pieces_in_play.append(Piece("queen", "black", "♛", 0, 3,))

def create_pieces():
    """Calls all the functions required to populate the board"""
    create_pawns()
    create_other_pieces()
    create_royalty()

def move_piece():
    """Decides whether or not to move the piece, based on the coordinates
    inputted by the user and move_legal and collision_detected. In the future,
    this will also depend on a check_for_check function.
    If a move is made then the old square is made empty"""
    get_user_input()
    piece = get_piece_on_square(move_piece_located_at[0], move_piece_located_at[1])
    piece.proposed_row = move_piece_to[0]
    piece.proposed_column = move_piece_to[1]
    if move_legal(piece) and not collision_detected(piece):
        board[piece.row][piece.column] = empty_square
        if piece_on_square(piece.proposed_row, piece.proposed_column):
            capture_piece(piece.proposed_row, piece.proposed_column)
        piece.row = piece.proposed_row
        piece.column = piece.proposed_column
    else:
        print("Sorry, that move is invalid")
        move_piece()

def get_user_input():
    """Takes user input, assuming a "x,y" format. Stores the 4 coordinates in 2 universal variables"""
    # doesn't currently validate the string supplied
    user_choice = input("Choose a piece (row,column)")
    if "quit" in user_choice.lower():
        sys.exit("Exiting...")
    move_piece_located_at[0] = int(user_choice[0])
    move_piece_located_at[1] = int(user_choice[2])
    user_move = input("Choose where to move the piece (row,column)")
    if "quit" in user_move.lower():
        sys.exit("Exiting...2")
    move_piece_to[0] = int(user_move[0])
    move_piece_to[1] = int(user_move[2])

def move_legal(piece):
    """Calls relevant functions for determining legality"""
    if piece.name == "bishop":
        return diagonal_move_legal(piece)
    elif piece.name == "rook":
        return lateral_move_legal(piece)
    elif piece.name == "queen":
        # If the queen behaves like either a bishop or a rook, return true
        return (lateral_move_legal(piece) or diagonal_move_legal(piece))
    elif piece.name == "pawn":
        return pawn_move_legal(piece)
    elif piece.name == "knight":
        return knight_move_legal(piece)
    elif piece.name == "king":
        return king_move_legal(piece)
    else:
        return False

def diagonal_move_legal(piece):
    """A bishop has to move along diagonals, so the abs deltaX and the abs deltaY should be equal """
    if abs(piece.row - piece.proposed_row) == abs(piece.column - piece.proposed_column):
        return True

def lateral_move_legal(piece):
    """A rook moves either along a row or a column, but not both """
    if ((piece.row - piece.proposed_row == 0) and (abs(piece.column - piece.proposed_column) > 0)) or ((abs(piece.row - piece.proposed_row) > 0) and (piece.column - piece.proposed_column == 0)):
        return True

def pawn_move_legal(piece):
    """Determines legality of pawn moves, taking into account their ability to
    move two rows from their starting square and their need to move diagonally
    to capture a piece. Does not currently handle en passant or promotion."""
    # could probably get rid of black/white checking with some additional maths. not sure if worth it
    if piece.colour == "black":
        if piece.row == 1 and piece.proposed_row == 3 and piece.column == piece.proposed_column:
            # collision code
            if (not piece_on_square(2, piece.column)) or (not piece_on_square(3, piece.column)):
                return True
        elif piece.row == piece.proposed_row - 1 and piece.column == piece.proposed_column:
            if not piece_on_square(piece.proposed_row, piece.proposed_column):
                return True
        elif piece.row == piece.proposed_row - 1 and abs(piece.column - piece.proposed_column) == 1:
            # we're capturing a piece here
            # and sort of doing collision detection. may want to re-think structure
            if piece_on_square(piece.proposed_row, piece.proposed_column) and get_piece_on_square(piece.proposed_row, piece.proposed_column).colour != piece.colour:
                return True
        else:
            return False
    elif piece.colour == "white":
        if piece.row == 6 and piece.proposed_row == 4 and piece.column == piece.proposed_column:
            if (not piece_on_square(5, piece.column)) or (not piece_on_square(4, piece.column)):
                return True
        elif piece.row == piece.proposed_row + 1 and piece.column == piece.proposed_column:
            if not piece_on_square(piece.proposed_row, piece.proposed_column):
                return True
        elif piece.row == piece.proposed_row + 1 and abs(piece.column - piece.proposed_column) == 1:
            if piece_on_square(piece.proposed_row, piece.proposed_column) and get_piece_on_square(piece.proposed_row, piece.proposed_column).colour != piece.colour:
                return True
        else:
            return False

def knight_move_legal(piece):
    """Determines knight move legality"""
    if (abs(piece.row - piece.proposed_row) == 2 and abs(piece.column - piece.proposed_column) == 1) or (abs(piece.row - piece.proposed_row) == 1 and abs(piece.column - piece.proposed_column) == 2):
        return True
    else:
        return False

def king_move_legal(piece):
    """Determines king move legality. The check status of the king is not
    currently taken into account"""
    if (abs(piece.row - piece.proposed_row) == 1 and abs(piece.column - piece.proposed_column) == 0) or (abs(piece.row - piece.proposed_row) == 1 and abs(piece.column - piece.proposed_column) == 1) or (abs(piece.row - piece.proposed_row) == 0 and abs(piece.column - piece.proposed_column) == 1):
        return True
    else:
        return False

def collision_detected(piece):
    """Calls the relevant functions to check for collision of different types
    of pieces"""
    # want to check all squares between start and finish - 1
    # also check final square to see if piece is of SAME colour and return False
    # pawn collision currently implemented in pawn_move_legal()
    if piece.name == "rook":
        return rook_collision_detected(piece)
    elif piece.name == "bishop":
        return bishop_collision_detected(piece)
    elif piece.name == "queen":
        # queens behave either as rooks or bishops
        return (rook_collision_detected(piece) or bishop_collision_detected(piece))
    elif piece.name == "knight":
        # knights only care about the square they land on
        return landing_square_collision_detected(piece)
    elif piece.name == "king":
        # kings care about the square they land on, and not moving into check
        # checking for check is not currently implemented
        return landing_square_collision_detected(piece)
    else:
        return False

def rook_collision_detected(piece):
    """Checks for rook collisions, taking into account that they can move in
    one of four directions"""
    # might be able to use max() and min() here rather than .copysign(). not sure if that would come in handy for bishop collision detection
    signed_step_column = int(math.copysign(1, piece.proposed_column - piece.column))
    signed_step_row = int(math.copysign(1, piece.proposed_row - piece.row))

    # check all the squares between start and destination (exclusive). return True if any piece is on these squares
    if piece.row - piece.proposed_row == 0:
        for column in range(piece.column + signed_step_column, piece.proposed_column, signed_step_column):
            if piece_on_square(piece.row, column):
                return True
    elif piece.column - piece.proposed_column == 0:
        for row in range(piece.row + signed_step_row, piece.proposed_row, signed_step_row):
            if piece_on_square(row, piece.column):
                return True
    else:
        landing_square_collision_detected(piece)

def bishop_collision_detected(piece):
    """Checks for bishop collision, taking into account that they can move in
    one of four directions"""
    signed_step_column = int(math.copysign(1, piece.proposed_column - piece.column))
    signed_step_row = int(math.copysign(1, piece.proposed_row - piece.row))
    for row in range(piece.row + signed_step_row, piece.proposed_row, signed_step_row):
        for column in range(piece.column + signed_step_column, piece.proposed_column, signed_step_column):
            if abs(piece.row - row) == abs(piece.column - column):
                if piece_on_square(row, column):
                    return True
    else:
        landing_square_collision_detected(piece)

def landing_square_collision_detected(piece):
    """Determines whether there is a piece on the final square of the move,
    and allows the move only if the piece is of the opposite colour. Does
    not currently allow for check."""
    if not piece_on_square(piece.proposed_row, piece.proposed_column):
        return False
    # allow movement to a square only with a piece of opposite colour on it
    # need to check for kings here
    elif get_piece_on_square(piece.proposed_row, piece.proposed_column).colour != piece.colour:
        return False
    else:
        return True

def main():
    """Calls the relevant functions to set up the board, and then loops each move
    of the game. Win conditions are not currently implemented."""
    game_over = False
    turn_counter = 1
    print("Remember to set your terminal to black text on a white background!\n")
    create_pieces()
    update_board()
    while not game_over:
        if turn_counter % 2 == 1:
            print("White's turn")
        else:
            print("Black's turn")
        move_piece()
        update_board()
        turn_counter += 1

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


<commit_before>import click

from commands.initialize import init_handler

@click.group()
def cli():
    pass

@cli.command()
def init():
    """Initialize the application"""
    init_handler()

<commit_msg>Make main script non-executable<commit_after>import click

from commands.initialize import init_handler

@click.group()
def cli():
    pass

@cli.command()
def init():
    """Initialize the application"""
    init_handler()

<|endoftext|>


<commit_before>from importlib import import_module
from itertools import chain
import json
from os import path, getcwd
import os
import subprocess
import inspect

from ..fileop import IOHelper
from ...util import Utility
from ..exechelper import func_exec_run

# import_module can't load the following modules in the NGINX server
# while running in 'immediate' mode. The early imports here are needed.
# This needs to be fixed, otherwise dynamic loading will not work.
try:
    import app.biowl.libraries.galaxy.adapter
except:
    pass

try:
    import app.biowl.libraries.bwa.adapter
    import app.biowl.libraries.pysam.adapter
    import app.biowl.libraries.fastqc.adapter
    import app.biowl.libraries.flash.adapter
    import app.biowl.libraries.hadoop.adapter
    import app.biowl.libraries.pear.adapter
    import app.biowl.libraries.seqtk.adapter
    import app.biowl.libraries.usearch.adapter
    import app.biowl.libraries.vsearch.adapter
except:
    pass

def load_module(modulename):
    '''
    Load a module dynamically from a string module name.
    It was first implemented with __import__, but later
    replaced by importlib.import_module.
    :param modulename:
    '''
    #if modulename not in sys.modules:
    #name = "package." + modulename
    #return __import__(modulename, fromlist=[''])
    return import_module(modulename)

class Function():
    def __init__(self, name, internal, package = None, module = None, params = [], example = None, desc = None, runmode = None, level = 0, group = None):
        self.name = name
        self.internal = internal
        self.package = package
        self.module = module
        self.params = params
        self.example = example
        self.desc = desc
        self.runmode = runmode
        self.level = level
        self.group = group

class Library():
    def __init__(self, funcs = {}):
        self.funcs = funcs
        self.tasks = {}
        self.localdir = path.join(path.abspath(path.dirname(__file__)), 'storage')

    def add_task(self, name, expr):
        self.tasks[name] = expr

    def run_task(self, name, args, dotaskstmt):
        if name in self.tasks:
            return dotaskstmt(self.tasks[name][1:], args)

    def code_run_task(self, name, args, dotaskstmt):
        if name in self.tasks:
            return dotaskstmt(self.tasks[name], args), set()

    @staticmethod
    def load(library_def_file):
        library = Library()
        library.funcs = Library.load_funcs_recursive(library_def_file)
        return library

    @staticmethod
    def load_funcs_recursive(library_def_file):
        if os.path.isfile(library_def_file):
            return Library.load_funcs(library_def_file)

        all_funcs = {}
        for f in os.listdir(library_def_file):
            funcs = Library.load_funcs_recursive(os.path.join(library_def_file, f))
            for k,v in funcs.items():
                if k in all_funcs:
                    all_funcs[k].extend(v)
                else:
                    all_funcs[k] = v if isinstance(v, list) else [v]
        return all_funcs

    @staticmethod
    def load_funcs(library_def_file):
        funcs = {}
        try:
            if not os.path.isfile(library_def_file) or not library_def_file.endswith(".json"):
                return funcs

            with open(library_def_file, 'r') as json_data:
                d = json.load(json_data)
                libraries = d["functions"]
                libraries = sorted(libraries, key = lambda k : k['package'].lower())
                for f in libraries:
                    name = f["name"] if f.get("name") else f["internal"]
                    internal = f["internal"] if f.get("internal") else f["name"]
                    module = f["module"] if f.get("module") else None
                    package = f["package"] if f.get("package") else ""
                    example = f["example"] if f.get("example") else ""
                    desc = f["desc"] if f.get("desc") else ""
                    runmode = f["runmode"] if f.get("runmode") else ""
                    level = int(f["level"]) if f.get("level") else 0
                    group = f["group"] if f.get("group") else ""
                    params = []
                    if f.get("params"):
                        for param in f["params"]:
                            params.append(param)
                    func = Function(name, internal, package, module, params, example, desc, runmode, level, group)
                    if name.lower() in funcs:
                        funcs[name.lower()].extend([func])
                    else:
                        funcs[name.lower()] = [func]
        finally:
            return funcs

    def func_to_internal_name(self, funcname):
        for f in self.funcs:
            if f.get("name") and self.iequal(f["name"], funcname):
                return f["internal"]

    def get_function(self, name, package = None):
        if package is not None:
            for func in self.funcs[name.lower()]:
                if func.package == package:
                    return [func]
        else:
            return self.funcs[name.lower()]

    def check_function(self, name, package = None):
        functions = self.get_function(name, package)
        return functions is not None and len(functions) > 0

    def funcs_flat(self):
        funcs = []
        for v in self.funcs.values():
            funcs.extend(v)
        return funcs

    @staticmethod
    def split_args(arguments):
        args = []
        kwargs = {}
        for arg in arguments:
            if isinstance(arg, tuple):
                kwargs[arg[0]] = arg[1]
            else:
                args.append(arg)
        return args, kwargs

    def call_func(self, context, package, function, args):
        '''
        Call a function from a module.
        :param context: The context for output and error
        :param package: The name of the package. If it's empty, local function is called
        :param function: Name of the function
        :param args: The arguments for the function
        '''
        arguments, kwargs = Library.split_args(args)

        if not package or package == "None":
            if function.lower() == "print":
                return context.write(*arguments)
            elif function.lower() == "range":
                return range(*arguments)
            elif function.lower() == "read":
                if not arguments:
                    raise ValueError("Read must have one argument.")
                fs = Utility.fs_by_prefix(arguments[0])
                return fs.read(arguments[0])
            elif function.lower() == "write":
                if len(arguments) < 2:
                    raise ValueError("Write must have two arguments.")
                fs = Utility.fs_by_prefix(arguments[0])
                return fs.write(arguments[0], arguments[1])
            elif function.lower() == "getfiles":
                fs = Utility.fs_by_prefix(arguments[0])
                return fs.get_files(arguments[0])
            elif function.lower() == "getfolders":
                fs = Utility.fs_by_prefix(arguments[0])
                return fs.get_folders(arguments[0])
            elif function.lower() == "createfolder":
                fs = Utility.fs_by_prefix(arguments[0])
                return fs.create_folder(arguments[0])
            elif function.lower() == "remove":
                fs = Utility.fs_by_prefix(arguments[0])
                return fs.remove(arguments[0])
            elif function.lower() == "makedirs":
                fs = Utility.fs_by_prefix(arguments[0])
                return fs.makedirs(arguments[0])
            elif function.lower() == "getcwd":
                return getcwd()
            elif function.lower() == "len":
                return len(arguments[0])
            elif function.lower() == "exec":
                return func_exec_run(arguments[0], *arguments[1:])
            #    return func_exec(arguments[0], *arguments[1:])
#             else:
#                 raise ValueError("{0} function not implemented".format(function))
    #             possibles = globals().copy()
    #             possibles.update(locals())
    #             function = possibles.get(function)
    #             return function(*arguments)
        func = self.get_function(function, package)
        module_obj = load_module(func[0].module)
        function = getattr(module_obj, func[0].internal)
        if func[0].runmode == 'dist':
            arguments = context.get_activedci() + arguments

        # special handling for galaxy if history_id is not given, use history id from symbol table
        if (func[0].module == 'app.biowl.libraries.galaxy.adapter'):
            if not 'history_id' in kwargs and context.var_exists('history_id'):
                fullargspec = inspect.getfullargspec(function)
                if fullargspec.varkw:
                    kwargs['history_id'] = context.get_var('history_id')
        return function(*arguments, **kwargs)

    def code_func(self, context, package, function, arguments):
        '''
        Call a function from a module.
        :param context: The context for output and error
        :param package: The name of the package. If it's empty, local function is called
        :param function: Name of the function
        :param arguments: The arguments for the function
        '''
        imports = set()
        args = ','.join(arguments)
        code = ''
        if not package or package == "None":
            if function.lower() == "print":
                code = "print({0})".format(args)
            elif function.lower() == "range":
                code = "range({0})".format(args)
            elif function.lower() == "read":
                imports.add("from fileop import IOHelper")
                code = "IOHelper.read({0})".format(args)
            elif function.lower() == "write":
                imports.add("from fileop import IOHelper")
                code = "IOHelper.write({0})".format(args)
            elif function.lower() == "getfiles":
                imports.add("from fileop import IOHelper")
                code = "IOHelper.getfiles({0})".format(args)
            elif function.lower() == "getfolders":
                imports.add("from fileop import IOHelper")
                code = "IOHelper.getfolders({0})".format(args)
            elif function.lower() == "remove":
                imports.add("from fileop import IOHelper")
                code = "IOHelper.remove({0})".format(args)
            elif function.lower() == "createfolder":
                imports.add("from fileop import IOHelper")
                code = "IOHelper.makedirs({0})".format(args)
            elif function.lower() == "getcwd":
                imports.add("import os")
                code = "os.getcwd()"
            elif function.lower() == "len":
                code = "len({0})".format(arguments[0])
            elif function.lower() == "exec":
                imports.add("import subprocess")
                code =  "func_exec_run({0}, {1})".format(arguments[0], arguments[1])

        if code:
            return code, imports

        imports.add("from importlib import import_module")
        func = self.get_function(function, package)
        code = "module_obj = load_module({0})\n".format(func[0].module)
        code += "function = getattr(module_obj, {0})\n".format(func[0].internal)
        if context.dci and context.dci[-1] and func.runmode == 'distibuted':
            args = [context.dci[-1]] + args
        code += "function({0})".format(args)
        return code, imports

    def __repr__(self):
        return "Library: " + repr(self.funcs)
    def __getitem__(self, key):
        return self.funcs[key]
    def __setitem__(self, key, val):
        self.funcs[key] = val
    def __delitem__(self, key):
        del self.funcs[key]
    def __contains__(self, key):
        return key in self.funcs
    def __iter__(self):
        return iter(self.funcs.keys())

    def __str__(self):
        funcs = self.funcs_flat();
        #funcs = [num for elem in funcs for num in elem]

        if len(funcs) > 0:
            mod_name = "Module name"
            mod_len = max(max(len(i.module) if i.module is not None else 0 for i in funcs), len(mod_name))

            internal_name = "Internal Name"
            internal_len = max(max(len(i.internal) for i in funcs), len(internal_name))

            func_name = "Function Name"
            func_len = max(max(len(i.name) for i in funcs), len(func_name))

            param_names = "Parameters"
            param_len = len(param_names)
            l = 0
            for a in funcs:
                for v in a.params:
                    l += len(v)
                param_len = max(param_len, l)

            # print table header for vars
            display = "\n\n{0:3s} | {1:^{2}s} | {3:^{4}s} | {5:^{6}s} | {7:^{8}s}".format(" No", mod_name, mod_len, internal_name, internal_len, func_name, func_len, param_names, param_len)
            display += ("\n-------------------" + "-" * (mod_len + internal_len + func_len + param_len))
            # print symbol table

            i = 1
            for v in funcs:
                module = v.module if v.module is not None else "None"
                parameters = ""
                for p in v.params:
                    if parameters == "":
                        parameters = "{0}".format(p)
                    else:
                        parameters += ", {0}".format(p)
                display += "\n{0:3d} | {1:^{2}s} | {3:^{4}s} | {5:^{6}s} | {7:^{8}s}".format(i, module, mod_len, v.internal, internal_len, v.name, func_len, parameters, param_len)

                i += 1

        return display

<commit_msg>bowtie2<commit_after>from importlib import import_module
from itertools import chain
import json
from os import path, getcwd
import os
import subprocess
import inspect

from ..fileop import IOHelper
from ...util import Utility
from ..exechelper import func_exec_run

# import_module can't load the following modules in the NGINX server
# while running in 'immediate' mode. The early imports here are needed.
# This needs to be fixed, otherwise dynamic loading will not work.
try:
    import app.biowl.libraries.galaxy.adapter
except:
    pass

try:
    import app.biowl.libraries.bowtie2.adapter
    import app.biowl.libraries.bwa.adapter
    import app.biowl.libraries.pysam.adapter
    import app.biowl.libraries.fastqc.adapter
    import app.biowl.libraries.flash.adapter
    import app.biowl.libraries.hadoop.adapter
    import app.biowl.libraries.pear.adapter
    import app.biowl.libraries.seqtk.adapter
    import app.biowl.libraries.usearch.adapter
    import app.biowl.libraries.vsearch.adapter
except:
    pass

def load_module(modulename):
    '''
    Load a module dynamically from a string module name.
    It was first implemented with __import__, but later
    replaced by importlib.import_module.
    :param modulename:
    '''
    #if modulename not in sys.modules:
    #name = "package." + modulename
    #return __import__(modulename, fromlist=[''])
    return import_module(modulename)

class Function():
    def __init__(self, name, internal, package = None, module = None, params = [], example = None, desc = None, runmode = None, level = 0, group = None):
        self.name = name
        self.internal = internal
        self.package = package
        self.module = module
        self.params = params
        self.example = example
        self.desc = desc
        self.runmode = runmode
        self.level = level
        self.group = group

class Library():
    def __init__(self, funcs = {}):
        self.funcs = funcs
        self.tasks = {}
        self.localdir = path.join(path.abspath(path.dirname(__file__)), 'storage')

    def add_task(self, name, expr):
        self.tasks[name] = expr

    def run_task(self, name, args, dotaskstmt):
        if name in self.tasks:
            return dotaskstmt(self.tasks[name][1:], args)

    def code_run_task(self, name, args, dotaskstmt):
        if name in self.tasks:
            return dotaskstmt(self.tasks[name], args), set()

    @staticmethod
    def load(library_def_file):
        library = Library()
        library.funcs = Library.load_funcs_recursive(library_def_file)
        return library

    @staticmethod
    def load_funcs_recursive(library_def_file):
        if os.path.isfile(library_def_file):
            return Library.load_funcs(library_def_file)

        all_funcs = {}
        for f in os.listdir(library_def_file):
            funcs = Library.load_funcs_recursive(os.path.join(library_def_file, f))
            for k,v in funcs.items():
                if k in all_funcs:
                    all_funcs[k].extend(v)
                else:
                    all_funcs[k] = v if isinstance(v, list) else [v]
        return all_funcs

    @staticmethod
    def load_funcs(library_def_file):
        funcs = {}
        try:
            if not os.path.isfile(library_def_file) or not library_def_file.endswith(".json"):
                return funcs

            with open(library_def_file, 'r') as json_data:
                d = json.load(json_data)
                libraries = d["functions"]
                libraries = sorted(libraries, key = lambda k : k['package'].lower())
                for f in libraries:
                    name = f["name"] if f.get("name") else f["internal"]
                    internal = f["internal"] if f.get("internal") else f["name"]
                    module = f["module"] if f.get("module") else None
                    package = f["package"] if f.get("package") else ""
                    example = f["example"] if f.get("example") else ""
                    desc = f["desc"] if f.get("desc") else ""
                    runmode = f["runmode"] if f.get("runmode") else ""
                    level = int(f["level"]) if f.get("level") else 0
                    group = f["group"] if f.get("group") else ""
                    params = []
                    if f.get("params"):
                        for param in f["params"]:
                            params.append(param)
                    func = Function(name, internal, package, module, params, example, desc, runmode, level, group)
                    if name.lower() in funcs:
                        funcs[name.lower()].extend([func])
                    else:
                        funcs[name.lower()] = [func]
        finally:
            return funcs

    def func_to_internal_name(self, funcname):
        for f in self.funcs:
            if f.get("name") and self.iequal(f["name"], funcname):
                return f["internal"]

    def get_function(self, name, package = None):
        if package is not None:
            for func in self.funcs[name.lower()]:
                if func.package == package:
                    return [func]
        else:
            return self.funcs[name.lower()]

    def check_function(self, name, package = None):
        functions = self.get_function(name, package)
        return functions is not None and len(functions) > 0

    def funcs_flat(self):
        funcs = []
        for v in self.funcs.values():
            funcs.extend(v)
        return funcs

    @staticmethod
    def split_args(arguments):
        args = []
        kwargs = {}
        for arg in arguments:
            if isinstance(arg, tuple):
                kwargs[arg[0]] = arg[1]
            else:
                args.append(arg)
        return args, kwargs

    def call_func(self, context, package, function, args):
        '''
        Call a function from a module.
        :param context: The context for output and error
        :param package: The name of the package. If it's empty, local function is called
        :param function: Name of the function
        :param args: The arguments for the function
        '''
        arguments, kwargs = Library.split_args(args)

        if not package or package == "None":
            if function.lower() == "print":
                return context.write(*arguments)
            elif function.lower() == "range":
                return range(*arguments)
            elif function.lower() == "read":
                if not arguments:
                    raise ValueError("Read must have one argument.")
                fs = Utility.fs_by_prefix(arguments[0])
                return fs.read(arguments[0])
            elif function.lower() == "write":
                if len(arguments) < 2:
                    raise ValueError("Write must have two arguments.")
                fs = Utility.fs_by_prefix(arguments[0])
                return fs.write(arguments[0], arguments[1])
            elif function.lower() == "getfiles":
                fs = Utility.fs_by_prefix(arguments[0])
                return fs.get_files(arguments[0])
            elif function.lower() == "getfolders":
                fs = Utility.fs_by_prefix(arguments[0])
                return fs.get_folders(arguments[0])
            elif function.lower() == "createfolder":
                fs = Utility.fs_by_prefix(arguments[0])
                return fs.create_folder(arguments[0])
            elif function.lower() == "remove":
                fs = Utility.fs_by_prefix(arguments[0])
                return fs.remove(arguments[0])
            elif function.lower() == "makedirs":
                fs = Utility.fs_by_prefix(arguments[0])
                return fs.makedirs(arguments[0])
            elif function.lower() == "getcwd":
                return getcwd()
            elif function.lower() == "len":
                return len(arguments[0])
            elif function.lower() == "exec":
                return func_exec_run(arguments[0], *arguments[1:])
            #    return func_exec(arguments[0], *arguments[1:])
#             else:
#                 raise ValueError("{0} function not implemented".format(function))
    #             possibles = globals().copy()
    #             possibles.update(locals())
    #             function = possibles.get(function)
    #             return function(*arguments)
        func = self.get_function(function, package)
        module_obj = load_module(func[0].module)
        function = getattr(module_obj, func[0].internal)
        if func[0].runmode == 'dist':
            arguments = context.get_activedci() + arguments

        # special handling for galaxy if history_id is not given, use history id from symbol table
        if (func[0].module == 'app.biowl.libraries.galaxy.adapter'):
            if not 'history_id' in kwargs and context.var_exists('history_id'):
                fullargspec = inspect.getfullargspec(function)
                if fullargspec.varkw:
                    kwargs['history_id'] = context.get_var('history_id')
        return function(*arguments, **kwargs)

    def code_func(self, context, package, function, arguments):
        '''
        Call a function from a module.
        :param context: The context for output and error
        :param package: The name of the package. If it's empty, local function is called
        :param function: Name of the function
        :param arguments: The arguments for the function
        '''
        imports = set()
        args = ','.join(arguments)
        code = ''
        if not package or package == "None":
            if function.lower() == "print":
                code = "print({0})".format(args)
            elif function.lower() == "range":
                code = "range({0})".format(args)
            elif function.lower() == "read":
                imports.add("from fileop import IOHelper")
                code = "IOHelper.read({0})".format(args)
            elif function.lower() == "write":
                imports.add("from fileop import IOHelper")
                code = "IOHelper.write({0})".format(args)
            elif function.lower() == "getfiles":
                imports.add("from fileop import IOHelper")
                code = "IOHelper.getfiles({0})".format(args)
            elif function.lower() == "getfolders":
                imports.add("from fileop import IOHelper")
                code = "IOHelper.getfolders({0})".format(args)
            elif function.lower() == "remove":
                imports.add("from fileop import IOHelper")
                code = "IOHelper.remove({0})".format(args)
            elif function.lower() == "createfolder":
                imports.add("from fileop import IOHelper")
                code = "IOHelper.makedirs({0})".format(args)
            elif function.lower() == "getcwd":
                imports.add("import os")
                code = "os.getcwd()"
            elif function.lower() == "len":
                code = "len({0})".format(arguments[0])
            elif function.lower() == "exec":
                imports.add("import subprocess")
                code =  "func_exec_run({0}, {1})".format(arguments[0], arguments[1])

        if code:
            return code, imports

        imports.add("from importlib import import_module")
        func = self.get_function(function, package)
        code = "module_obj = load_module({0})\n".format(func[0].module)
        code += "function = getattr(module_obj, {0})\n".format(func[0].internal)
        if context.dci and context.dci[-1] and func.runmode == 'distibuted':
            args = [context.dci[-1]] + args
        code += "function({0})".format(args)
        return code, imports

    def __repr__(self):
        return "Library: " + repr(self.funcs)
    def __getitem__(self, key):
        return self.funcs[key]
    def __setitem__(self, key, val):
        self.funcs[key] = val
    def __delitem__(self, key):
        del self.funcs[key]
    def __contains__(self, key):
        return key in self.funcs
    def __iter__(self):
        return iter(self.funcs.keys())

    def __str__(self):
        funcs = self.funcs_flat();
        #funcs = [num for elem in funcs for num in elem]

        if len(funcs) > 0:
            mod_name = "Module name"
            mod_len = max(max(len(i.module) if i.module is not None else 0 for i in funcs), len(mod_name))

            internal_name = "Internal Name"
            internal_len = max(max(len(i.internal) for i in funcs), len(internal_name))

            func_name = "Function Name"
            func_len = max(max(len(i.name) for i in funcs), len(func_name))

            param_names = "Parameters"
            param_len = len(param_names)
            l = 0
            for a in funcs:
                for v in a.params:
                    l += len(v)
                param_len = max(param_len, l)

            # print table header for vars
            display = "\n\n{0:3s} | {1:^{2}s} | {3:^{4}s} | {5:^{6}s} | {7:^{8}s}".format(" No", mod_name, mod_len, internal_name, internal_len, func_name, func_len, param_names, param_len)
            display += ("\n-------------------" + "-" * (mod_len + internal_len + func_len + param_len))
            # print symbol table

            i = 1
            for v in funcs:
                module = v.module if v.module is not None else "None"
                parameters = ""
                for p in v.params:
                    if parameters == "":
                        parameters = "{0}".format(p)
                    else:
                        parameters += ", {0}".format(p)
                display += "\n{0:3d} | {1:^{2}s} | {3:^{4}s} | {5:^{6}s} | {7:^{8}s}".format(i, module, mod_len, v.internal, internal_len, v.name, func_len, parameters, param_len)

                i += 1

        return display

<|endoftext|>


<commit_before>"""
Record collector sends produced results to kafka topic

"""

import time

from .._error import KafkaStreamsError

class DefaultStreamPartitioner:
    def __init__(self):
        pass

    def partition(self):
        return 0

class IdentitySerialiser:
    def serialise(self, value):
        return value

class RecordCollector:
    def __init__(self, _producer):
        self.producer = _producer

    def send_to_stream(self, topic, key, value, timestamp, keySerialiser , valueSerialiser,
                       *, stream_partitioner = DefaultStreamPartitioner()):

        partitions = producer.partitionsFor(topic)
        n_partitions = len(partitions)
        if n_partitions == 0:
            raise KafkaStreamsError(f"Could not get partition information for {topic}." \
                                    "This can happen if the topic does not exist.")

        self.send_to_partition(topic, key, value, timestamp, keySerialiser,
                               valueSerialiser, partition = partitioner.partition(key, value, n_partitions))

    def send_to_partition(self, topic, key, value, timestamp,
                          keySerialiser = IdentitySerialiser(), valueSerialiser = IdentitySerialiser(), *, partition = 0):
        key = keySerialiser.serialise(key)
        value = valueSerialiser.serialise(value)
        try:
            self.producer.produce(topic, value, key, partition, self.on_delivery, timestamp)
            self.producer.poll(0) # Ensure previous message's delivery reports are served
        except BufferError as be:
            # TODO: We should handle this better. The queue can be full, for example on leader election and we should retry
            log.exception(be)
        except NotImplementedError as nie:
            log.exception(nie)

    def on_delivery(self, err, msg):
        if err:
            raise KafkaStreamsError(f'Error on delivery of message {msg}')

    def flush(self):
        log.debug('Flushing producer')
        self.producer.flush()

    def close(self):
        log.debug('Closing producer')
        self.producer.close()
<commit_msg>Handle BufferError in RecordCollector<commit_after>"""
Record collector sends produced results to kafka topic

"""

import time

from .._error import KafkaStreamsError

class DefaultStreamPartitioner:
    def __init__(self):
        pass

    def partition(self):
        return 0

class IdentitySerialiser:
    def serialise(self, value):
        return value

class RecordCollector:
    def __init__(self, _producer):
        self.producer = _producer

    def send_to_stream(self, topic, key, value, timestamp, keySerialiser , valueSerialiser,
                       *, stream_partitioner = DefaultStreamPartitioner()):

        partitions = producer.partitionsFor(topic)
        n_partitions = len(partitions)
        if n_partitions == 0:
            raise KafkaStreamsError(f"Could not get partition information for {topic}." \
                                    "This can happen if the topic does not exist.")

        self.send_to_partition(topic, key, value, timestamp, keySerialiser,
                               valueSerialiser, partition = partitioner.partition(key, value, n_partitions))

    def send_to_partition(self, topic, key, value, timestamp,
                          keySerialiser = IdentitySerialiser(), valueSerialiser = IdentitySerialiser(), *, partition = 0):
        key = keySerialiser.serialise(key)
        value = valueSerialiser.serialise(value)
        produced = False
        while not produced:
            try:
                self.producer.produce(topic, value, key, partition, self.on_delivery, timestamp)
                self.producer.poll(0) # Ensure previous message's delivery reports are served
                produced = True
            except BufferError as be:
                log.exception(be)
                self.producer.poll(10) # Wait a bit longer to give buffer more time to flush
            except NotImplementedError as nie:
                log.exception(nie)
                produced = True  # should not enter infinite loop

    def on_delivery(self, err, msg):
        if err:
            raise KafkaStreamsError(f'Error on delivery of message {msg}')

    def flush(self):
        log.debug('Flushing producer')
        self.producer.flush()

    def close(self):
        log.debug('Closing producer')
        self.producer.close()
<|endoftext|>


<commit_before># vim: shiftwidth=4 tabstop=4 expandtab
import os, time, string, re, cgi
import tools, api
from entries.fileentry import FileEntry
from entries.cachedecorator import CacheDecorator
import cPickle as pickle

class PyBlosxom:
    """
    This is the main class for PyBlosxom functionality.  It handles
    initialization, defines default behavior, and also pushes the
    request through all the steps until the output is rendered and
    we're complete.
    """
    def __init__(self, request):
        global myrequest
        myrequest = request

        self._request = request
        self.flavour = {}
        self.dayFlag = 1 # Used to print valid date footers

        registry = tools.get_registry()
        registry["request"] = request

    def startup(self):
        """
        The startup step further initializes the Request by setting
        additional information in the _data dict.
        """
        data = self._request.getData()
        pyhttp = self._request.getHttp()
        config = self._request.getConfiguration()

        data['pi_bl'] = ''

        # Get our URL and configure the base_url param
        if pyhttp.has_key('SCRIPT_NAME'):
            if not config.has_key('base_url'):
                config['base_url'] = 'http://%s%s' % (pyhttp['HTTP_HOST'], pyhttp['SCRIPT_NAME'])

            data['url'] = '%s%s' % (config['base_url'], data['pi_bl'])

    def defaultFileListHandler(self, args):
        """
        This is the default handler for getting entries.  It takes the
        request object in and figures out which entries based on the
        default behavior that we want to show and generates a list of
        EntryBase subclass objects which it returns.

        @param request: the incoming Request object
        @type request: libs.Request.Request

        @returns: the content we want to render
        @rtype: list of EntryBase objects
        """
        request = args["request"]

        data = request.getData()
        config = request.getConfiguration()
        pyhttp = request.getHttp()

        path_info = data['path_info']

        if data['bl_type'] == 'dir':
            filelist = tools.Walk(data['root_datadir'], int(config['depth']))
        elif data['bl_type'] == 'file':
            filelist = [data['root_datadir']]
        else:
            filelist = []

        entrylist = []
        for ourfile in filelist:
            entry = CacheDecorator(FileEntry(config, ourfile, data['root_datadir']))
            entrylist.append(entry)
        entrylist = tools.sortDictBy(entrylist, "mtime")

        # Match dates with files if applicable
        if data['pi_yr']:
            month = (data['pi_mo'] in tools.month2num.keys() and tools.month2num[data['pi_mo']] or data['pi_mo'])
            matchstr = "^" + data["pi_yr"] + month + data["pi_da"]
            valid_list = [x for x in entrylist if re.match(matchstr, x['fulltime'])]
        else:
            valid_list = entrylist

        return valid_list

    def run(self):
        """
        Main loop for pyblosxom.
        """
        config = self._request.getConfiguration()
        data = self._request.getData()
        pyhttp = self._request.getHttp()

        # instantiate the renderer with the current request and store it
        # in the data dict
        renderer = tools.importName('libs.renderers',
                config.get('renderer', 'blosxom')).Renderer(self._request)
        data["renderer"] = renderer

        # import entryparsers here to allow other plugins register what file
        # extensions can be used
        import libs.entryparsers.__init__
        libs.entryparsers.__init__.initialize_extensions()
        data['extensions'] = libs.entryparsers.__init__.ext

        # import plugins and allow them to register with the api
        import libs.plugins.__init__
        libs.plugins.__init__.initialize_plugins(config)

        api.fileListHandler.register(self.defaultFileListHandler, api.LAST)

        form = self._request.getHttp()["form"]
        data['flavour'] = (form.has_key('flav') and
                form['flav'].value or
                config.get('defaultFlavour', 'html'))

        path_info = []

        # Get the blog name if possible
        data['pi_yr'] = ''
        data['pi_mo'] = ''
        data['pi_da'] = ''

        if pyhttp.get('PATH_INFO', ''):
            path_info = pyhttp['PATH_INFO'].split('/')

        data['path_info'] = list(path_info)
        data['root_datadir'] = config['datadir']

        # Process path_info
        # The paths specification looks like this:
        # /cat - category
        # /2002 - year
        # /2002/Feb (or 02) - Year and Month
        # /cat/2002/Feb/31 - year and month day in category.
        # /foo.html and /cat/foo.html - file foo.* in / and /cat
        # To simplify checking, four digits directory name is not allowed.
        got_date = 0
        for path_data in path_info:
            if not path_data:
                continue
            elif len(path_data) == 4 and path_data.isdigit():
                # We got a hot date here guys :)
                got_date = 1
                break
            else:
                data['pi_bl'] = os.path.join(data['pi_bl'], path_data)

        if got_date:
            # Get messy with dates
            while not (len(path_info[0]) == 4 and path_info[0].isdigit()):
                path_info.pop(0)
            # Year
            data['pi_yr'] = path_info.pop(0)
            # Month
            if path_info and path_info[0] in tools.MONTHS:
                data['pi_mo'] = path_info.pop(0)
                # Day
                if path_info and re.match("^([0-2][0-9]|3[0-1])$", path_info[0]):
                    # Potential day here, no more processing from here
                    data['pi_da'] = path_info.pop(0)
            if path_info[0]:
                # Potential flavour after date
                filename, ext = os.path.splitext(path_info[0])
                if filename == 'index':
                    data['flavour'] = ext[1:]

        blog_result = os.path.join(config['datadir'], data['pi_bl'])

        data['bl_type'] = ''

        # If all we've got is a directory, things are simple
        if os.path.isdir(blog_result):
            if data['pi_bl'] != '':
                config['blog_title'] += ' : %s' % data['pi_bl']
            data['root_datadir'] = blog_result
            data['bl_type'] = 'dir'

        # Else we may have a file
        if not data['bl_type']:
            # Try for file
            def what_ext(path):
                for ext in data['extensions'].keys():
                    if os.path.isfile(path + '.' + ext):
                        return ext
                return None

            ext = what_ext(blog_result)
            if ext:
                config['blog_title'] += ' : %s' % data['pi_bl']
                data['bl_type'] = 'file'
                data['root_datadir'] = blog_result + '.' + ext
            else:
                # We may have flavour embedded here
                filename, ext = os.path.splitext(blog_result)
                fileext = what_ext(filename)
                dirname = os.path.dirname(filename)

                if fileext:
                    data['flavour'] = ext[1:]
                    data['root_datadir'] = filename + '.' + fileext
                    config['blog_title'] += ' : %s' % data['pi_bl']
                    data['bl_type'] = 'file'
                elif (os.path.basename(filename) == 'index' and
                        os.path.isdir(dirname)):
                    # blanket flavours?
                    data['flavour'] = ext[1:]
                    if data['pi_bl'] != '':
                        config['blog_title'] += ' : %s' % os.path.dirname(data['pi_bl'])
                    data['root_datadir'] = dirname
                    data['bl_type'] = 'dir'

        # calling fileList will generate a list of entries from the
        # api.fileListHandler
        data["entry_list"] = tools.fileList(self._request)

        # we pass the request with the entry_list through the
        # plugins giving them a chance to transform the data.
        # plugins modify the request in-place--no need to return
        # things.
        api.prepareChain.executeHandler({"request": self._request})

        # now we pass the entry_list through the renderer
        entry_list = data["entry_list"]

        if renderer and not renderer.rendered:
            if entry_list:
                renderer.setContent(entry_list)
                tools.logRequest(config.get('logfile',''), '200')
            else:
                renderer.addHeader(['Status: 404 Not Found'])
                renderer.setContent(
                    {'title': 'The page you are looking for is not available',
                     'body': 'Somehow I cannot find the page you want. ' +
                     'Go Back to %s?'
                     % (config["base_url"], config["blog_title"])})
                tools.logRequest(config.get('logfile',''), '404')
            renderer.render()

        elif not renderer:
            print "Content-Type: text/plain\n\nThere is something wrong with your setup"
<commit_msg>Fixed error that's alerted to me by Ted.<commit_after># vim: shiftwidth=4 tabstop=4 expandtab
import os, time, string, re, cgi
import tools, api
from entries.fileentry import FileEntry
from entries.cachedecorator import CacheDecorator
import cPickle as pickle

class PyBlosxom:
    """
    This is the main class for PyBlosxom functionality.  It handles
    initialization, defines default behavior, and also pushes the
    request through all the steps until the output is rendered and
    we're complete.
    """
    def __init__(self, request):
        global myrequest
        myrequest = request

        self._request = request
        self.flavour = {}
        self.dayFlag = 1 # Used to print valid date footers

        registry = tools.get_registry()
        registry["request"] = request

    def startup(self):
        """
        The startup step further initializes the Request by setting
        additional information in the _data dict.
        """
        data = self._request.getData()
        pyhttp = self._request.getHttp()
        config = self._request.getConfiguration()

        data['pi_bl'] = ''

        # Get our URL and configure the base_url param
        if pyhttp.has_key('SCRIPT_NAME'):
            if not config.has_key('base_url'):
                config['base_url'] = 'http://%s%s' % (pyhttp['HTTP_HOST'], pyhttp['SCRIPT_NAME'])

            data['url'] = '%s%s' % (config['base_url'], data['pi_bl'])

    def defaultFileListHandler(self, args):
        """
        This is the default handler for getting entries.  It takes the
        request object in and figures out which entries based on the
        default behavior that we want to show and generates a list of
        EntryBase subclass objects which it returns.

        @param request: the incoming Request object
        @type request: libs.Request.Request

        @returns: the content we want to render
        @rtype: list of EntryBase objects
        """
        request = args["request"]

        data = request.getData()
        config = request.getConfiguration()
        pyhttp = request.getHttp()

        path_info = data['path_info']

        if data['bl_type'] == 'dir':
            filelist = tools.Walk(data['root_datadir'], int(config['depth']))
        elif data['bl_type'] == 'file':
            filelist = [data['root_datadir']]
        else:
            filelist = []

        entrylist = []
        for ourfile in filelist:
            entry = CacheDecorator(FileEntry(config, ourfile, data['root_datadir']))
            entrylist.append(entry)
        entrylist = tools.sortDictBy(entrylist, "mtime")

        # Match dates with files if applicable
        if data['pi_yr']:
            month = (data['pi_mo'] in tools.month2num.keys() and tools.month2num[data['pi_mo']] or data['pi_mo'])
            matchstr = "^" + data["pi_yr"] + month + data["pi_da"]
            valid_list = [x for x in entrylist if re.match(matchstr, x['fulltime'])]
        else:
            valid_list = entrylist

        return valid_list

    def run(self):
        """
        Main loop for pyblosxom.
        """
        config = self._request.getConfiguration()
        data = self._request.getData()
        pyhttp = self._request.getHttp()

        # instantiate the renderer with the current request and store it
        # in the data dict
        renderer = tools.importName('libs.renderers',
                config.get('renderer', 'blosxom')).Renderer(self._request)
        data["renderer"] = renderer

        # import entryparsers here to allow other plugins register what file
        # extensions can be used
        import libs.entryparsers.__init__
        libs.entryparsers.__init__.initialize_extensions()
        data['extensions'] = libs.entryparsers.__init__.ext

        # import plugins and allow them to register with the api
        import libs.plugins.__init__
        libs.plugins.__init__.initialize_plugins(config)

        api.fileListHandler.register(self.defaultFileListHandler, api.LAST)

        form = self._request.getHttp()["form"]
        data['flavour'] = (form.has_key('flav') and
                form['flav'].value or
                config.get('defaultFlavour', 'html'))

        path_info = []

        # Get the blog name if possible
        data['pi_yr'] = ''
        data['pi_mo'] = ''
        data['pi_da'] = ''

        if pyhttp.get('PATH_INFO', ''):
            path_info = pyhttp['PATH_INFO'].split('/')

        data['path_info'] = list(path_info)
        data['root_datadir'] = config['datadir']

        # Process path_info
        # The paths specification looks like this:
        # /cat - category
        # /2002 - year
        # /2002/Feb (or 02) - Year and Month
        # /cat/2002/Feb/31 - year and month day in category.
        # /foo.html and /cat/foo.html - file foo.* in / and /cat
        # To simplify checking, four digits directory name is not allowed.
        got_date = 0
        for path_data in path_info:
            if not path_data:
                continue
            elif len(path_data) == 4 and path_data.isdigit():
                # We got a hot date here guys :)
                got_date = 1
                break
            else:
                data['pi_bl'] = os.path.join(data['pi_bl'], path_data)

        if got_date:
            # Get messy with dates
            while not (len(path_info[0]) == 4 and path_info[0].isdigit()):
                path_info.pop(0)
            # Year
            data['pi_yr'] = path_info.pop(0)
            # Month
            if path_info and path_info[0] in tools.MONTHS:
                data['pi_mo'] = path_info.pop(0)
                # Day
                if path_info and re.match("^([0-2][0-9]|3[0-1])$", path_info[0]):
                    # Potential day here, no more processing from here
                    data['pi_da'] = path_info.pop(0)
            if path_info and path_info[0]:
                # Potential flavour after date
                filename, ext = os.path.splitext(path_info[0])
                if filename == 'index':
                    data['flavour'] = ext[1:]

        blog_result = os.path.join(config['datadir'], data['pi_bl'])

        data['bl_type'] = ''

        # If all we've got is a directory, things are simple
        if os.path.isdir(blog_result):
            if data['pi_bl'] != '':
                config['blog_title'] += ' : %s' % data['pi_bl']
            data['root_datadir'] = blog_result
            data['bl_type'] = 'dir'

        # Else we may have a file
        if not data['bl_type']:
            # Try for file
            def what_ext(path):
                for ext in data['extensions'].keys():
                    if os.path.isfile(path + '.' + ext):
                        return ext
                return None

            ext = what_ext(blog_result)
            if ext:
                config['blog_title'] += ' : %s' % data['pi_bl']
                data['bl_type'] = 'file'
                data['root_datadir'] = blog_result + '.' + ext
            else:
                # We may have flavour embedded here
                filename, ext = os.path.splitext(blog_result)
                fileext = what_ext(filename)
                dirname = os.path.dirname(filename)

                if fileext:
                    data['flavour'] = ext[1:]
                    data['root_datadir'] = filename + '.' + fileext
                    config['blog_title'] += ' : %s' % data['pi_bl']
                    data['bl_type'] = 'file'
                elif (os.path.basename(filename) == 'index' and
                        os.path.isdir(dirname)):
                    # blanket flavours?
                    data['flavour'] = ext[1:]
                    if data['pi_bl'] != '':
                        config['blog_title'] += ' : %s' % os.path.dirname(data['pi_bl'])
                    data['root_datadir'] = dirname
                    data['bl_type'] = 'dir'

        # calling fileList will generate a list of entries from the
        # api.fileListHandler
        data["entry_list"] = tools.fileList(self._request)

        # we pass the request with the entry_list through the
        # plugins giving them a chance to transform the data.
        # plugins modify the request in-place--no need to return
        # things.
        api.prepareChain.executeHandler({"request": self._request})

        # now we pass the entry_list through the renderer
        entry_list = data["entry_list"]

        if renderer and not renderer.rendered:
            if entry_list:
                renderer.setContent(entry_list)
                tools.logRequest(config.get('logfile',''), '200')
            else:
                renderer.addHeader(['Status: 404 Not Found'])
                renderer.setContent(
                    {'title': 'The page you are looking for is not available',
                     'body': 'Somehow I cannot find the page you want. ' +
                     'Go Back to %s?'
                     % (config["base_url"], config["blog_title"])})
                tools.logRequest(config.get('logfile',''), '404')
            renderer.render()

        elif not renderer:
            print "Content-Type: text/plain\n\nThere is something wrong with your setup"
<|endoftext|>


<commit_before>from nbt import *
from nbt import _TAG_Numeric
import unittest
from StringIO import StringIO
from gzip import GzipFile

class ReadWriteTest(unittest.TestCase):     # test that we can read the test file correctly

    def testReadBig(self):
        mynbt = NBTFile("bigtest.nbt")
        self.assertTrue(mynbt.filename != None)
        self.assertEqual(len(mynbt.tags), 11)

    def testWriteBig(self):
        mynbt = NBTFile("bigtest.nbt")
        output = StringIO()
        mynbt.write_file(buffer=output)
        self.assertTrue(GzipFile("bigtest.nbt").read() == output.getvalue())

    def testWriteback(self):
        mynbt = NBTFile("bigtest.nbt")
        mynbt.write_file()

class TreeManipulationTest(unittest.TestCase):

    def setUp(self):
        self.nbtfile = NBTFile()

    def testRootNodeSetup(self):
        self.nbtfile.name = "Hello World"
        self.assertEqual(self.nbtfile.name, "Hello World")

    def testTagNumeric(self):
        for tag in TAGLIST:
            if isinstance(TAGLIST[tag], _TAG_Numeric):
                tagobj = TAGLIST[tag](name="Test", value=10)
                self.assertEqual(byte.name, "Test", "Name not set correctly for %s" % TAGLIST[tag].__class__.__name__)
                self.assertEqual(byte.value, 10, "Value not set correctly for %s" % TAGLIST[tag].__class__.__name__)
                self.nbtfile.tags.append(tagobj)

    #etcetera..... will finish later

    def tearDown(self):
        del self.nbtfile

class EmptyStringTest(unittest.TestCase):

    def setUp(self):
        self.golden_value = "\x0A\0\x04Test\x08\0\x0Cempty string\0\0\0"
        self.nbtfile = NBTFile(buffer=StringIO(self.golden_value))

    def testReadEmptyString(self):
        self.assertEqual(self.nbtfile.name, "Test")
        self.assertEqual(self.nbtfile["empty string"].value, "")

    def testWriteEmptyString(self):
        buffer = StringIO()
        self.nbtfile.write_file(buffer=buffer)
        self.assertEqual(buffer.getvalue(), self.golden_value)

if __name__ == '__main__':
    unittest.main()
<commit_msg>tests: Add test for issue 4.<commit_after>from nbt import *
from nbt import _TAG_Numeric
import unittest
from StringIO import StringIO
from gzip import GzipFile
import tempfile

class BugfixTest(unittest.TestCase):
    """
    Bugfix regression tests.

    These tend to not fit into nice categories.
    """

    def test_issue4(self):
        """
        Opening an empty file causes an uncaught exception.

        https://github.com/twoolie/NBT/issues/issue/4
        """

        temp = tempfile.NamedTemporaryFile()
        temp.flush()
        tag = NBTFile(temp.name)

class ReadWriteTest(unittest.TestCase):     # test that we can read the test file correctly

    def testReadBig(self):
        mynbt = NBTFile("bigtest.nbt")
        self.assertTrue(mynbt.filename != None)
        self.assertEqual(len(mynbt.tags), 11)

    def testWriteBig(self):
        mynbt = NBTFile("bigtest.nbt")
        output = StringIO()
        mynbt.write_file(buffer=output)
        self.assertTrue(GzipFile("bigtest.nbt").read() == output.getvalue())

    def testWriteback(self):
        mynbt = NBTFile("bigtest.nbt")
        mynbt.write_file()

class TreeManipulationTest(unittest.TestCase):

    def setUp(self):
        self.nbtfile = NBTFile()

    def testRootNodeSetup(self):
        self.nbtfile.name = "Hello World"
        self.assertEqual(self.nbtfile.name, "Hello World")

    def testTagNumeric(self):
        for tag in TAGLIST:
            if isinstance(TAGLIST[tag], _TAG_Numeric):
                tagobj = TAGLIST[tag](name="Test", value=10)
                self.assertEqual(byte.name, "Test", "Name not set correctly for %s" % TAGLIST[tag].__class__.__name__)
                self.assertEqual(byte.value, 10, "Value not set correctly for %s" % TAGLIST[tag].__class__.__name__)
                self.nbtfile.tags.append(tagobj)

    #etcetera..... will finish later

    def tearDown(self):
        del self.nbtfile

class EmptyStringTest(unittest.TestCase):

    def setUp(self):
        self.golden_value = "\x0A\0\x04Test\x08\0\x0Cempty string\0\0\0"
        self.nbtfile = NBTFile(buffer=StringIO(self.golden_value))

    def testReadEmptyString(self):
        self.assertEqual(self.nbtfile.name, "Test")
        self.assertEqual(self.nbtfile["empty string"].value, "")

    def testWriteEmptyString(self):
        buffer = StringIO()
        self.nbtfile.write_file(buffer=buffer)
        self.assertEqual(buffer.getvalue(), self.golden_value)

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


<commit_before>from Bio import Entrez
import time  # We need this if we want to filter based on age
from urllib2 import HTTPError
import xml.etree.ElementTree as ET
from settings import configuration
email = configuration["email"]

"""
Take a search term, intended to be an rs#,
and return a list of pmids
"""

def get_pmids(interm):
    # This is obsolete now, essentially, but
    # allows a user to pass a single string
    # which can be nice.
    if isinstance(interm, str):
        interm = interm + " AND pubmed_snp_cited[sb]"
        Entrez.email = email     # Always tell NCBI who you are
        search_results = Entrez.read(Entrez.esearch(db="pubmed", term=interm,
                                                    usehistory="y"))
        print("Found a total of " + search_results["Count"] + " results using search string '" + interm + "'")
        return search_results
    elif isinstance(interm, list):
        searchstring = " OR ".join(interm)
        searchstring = "(" + searchstring + ") AND pubmed_snp_cited[sb]"
        Entrez.email = email     # Always tell NCBI who you are
        search_results = Entrez.read(Entrez.esearch(db="pubmed",
                                                    term=searchstring,
                                                    usehistory="y"))
        print("Found a total of " + search_results["Count"] + " results using search string'" + searchstring + "'")
        return search_results

"""
Takes the saved list of results from get_pmids and retrieves
the raw XML to parse.  Currently goes article by article instead of
doing them in bulk.  Could go either way. Unsure which way
is better / more efficient.
"""

def get_abstracts(results):
    abstracts_list = []
    Entrez.email = email
    for start in range(0, int(results["Count"]), 1):
        end = min(int(results["Count"]), start+1)
        # print("Going to download record %i to %i" % (start+1, end))
        fetch_handle = Entrez.efetch(db="pubmed", rettype="abstract",
                                        retmode="xml", retstart=start,
                                        retmax=1,
                                        webenv=results["WebEnv"],
                                        query_key=results["QueryKey"])
        data = fetch_handle.read()
        fetch_handle.close()
        root = ET.fromstring(data)
        for abst in root.iter('Abstract'):
            for sec in abst.iter('AbstractText'):
                abstracts_list.append(sec.text)
    return abstracts_list

def get_abstracts_from_list(pmids_list):
    abstracts_list = []
    Entrez.email = email
    for each_pmid in pmids_list:
        fetch_handle = Entrez.efetch(db="pubmed", id=each_pmid, retmode='xml')
        data = fetch_handle.read()
        fetch_handle.close()
        root = ET.fromstring(data)
        for abst in root.iter('Abstract'):
            for sec in abst.iter('AbstractText'):
                abstracts_list.append(sec.text)
    return abstracts_list
<commit_msg>for get_pmids function, return a list of pmids. for getting abstracts from a list of pmids, return a dict<commit_after>from Bio import Entrez
import time  # We need this if we want to filter based on age
from urllib2 import HTTPError
import xml.etree.ElementTree as ET
from settings import configuration
email = configuration["email"]

"""
Take a search term, intended to be an rs#,
and return a list of pmids
"""

def get_pmids(interm):
    # This is obsolete now, essentially, but
    # allows a user to pass a single string
    # which can be nice.
    if isinstance(interm, str):
        interm = interm + " AND pubmed_snp_cited[sb]"
        Entrez.email = email     # Always tell NCBI who you are
        search_results = Entrez.read(Entrez.esearch(db="pubmed", term=interm,
                                                    usehistory="y"))
        print("Found a total of " + search_results["Count"] + " results using search string '" + interm + "'")
        pmids_list = search_results["IdList"]
        #return search_results
        return pmids_list
    elif isinstance(interm, list):
        searchstring = " OR ".join(interm)
        searchstring = "(" + searchstring + ") AND pubmed_snp_cited[sb]"
        Entrez.email = email     # Always tell NCBI who you are
        search_results = Entrez.read(Entrez.esearch(db="pubmed",
                                                    term=searchstring,
                                                    usehistory="y"))
        print("Found a total of " + search_results["Count"] + " results using search string'" + searchstring + "'")
        pmids_list = search_results["IdList"]
        return pmids_list

"""
Takes the saved list of results from get_pmids and retrieves
the raw XML to parse.  Currently goes article by article instead of
doing them in bulk.  Could go either way. Unsure which way
is better / more efficient.
"""

def get_abstracts(results):
    abstracts_list = []
    Entrez.email = email
    for start in range(0, int(results["Count"]), 1):
        end = min(int(results["Count"]), start+1)
        # print("Going to download record %i to %i" % (start+1, end))
        fetch_handle = Entrez.efetch(db="pubmed", rettype="abstract",
                                        retmode="xml", retstart=start,
                                        retmax=1,
                                        webenv=results["WebEnv"],
                                        query_key=results["QueryKey"])
        data = fetch_handle.read()
        fetch_handle.close()
        root = ET.fromstring(data)
        for abst in root.iter('Abstract'):
            for sec in abst.iter('AbstractText'):
                abstracts_list.append(sec.text)
    return abstracts_list

def get_abstracts_from_list(pmids_list):
#    abstracts_list = []
    pmids_abstracts_dict = {}
    Entrez.email = email
    for each_pmid in pmids_list:
        fetch_handle = Entrez.efetch(db="pubmed", id=each_pmid, retmode='xml')
        data = fetch_handle.read()
        fetch_handle.close()
        root = ET.fromstring(data)
        for abst in root.iter('Abstract'):
            for sec in abst.iter('AbstractText'):
                #abstracts_list.append(sec.text)
                pmids_abstracts_dict[each_pmid] = sec.text
    print pmids_abstracts_dict
#    print abstracts_list
    #return abstracts_list
<|endoftext|>


<commit_before>#!/usr/bin/python

try:
    from lxml import etree
except ImportError:
    from xml import etree
    # if this fails, then we're screwed, so let the script die

class GraphMLWriter(object):
    def __init__(self):

        ns = {None:'http://graphml.graphdrawing.org/xmlns/graphml'}
        self._root = etree.Element('graphml', nsmap=ns)
        self._graph = etree.SubElement(self._root, 'graph', edgedefault="directed")

        # create the entry node
        self._entry = etree.SubElement(self._graph, 'node', id='0')
        etree.SubElement(self._entry, 'data', key='entry').text = 'true'

    def parsePath(self, pathFile, filename = None):
        """
        Parse .path file from klee
        \param pathFile     the .path file
        \param filename     name of the file the symbiotic ran on
                            -- in the case that we want to stick
                            only to this file in the witness
        """

        f = open(pathFile, 'r')
        last_id=1

        #ctrlelem = None
        for line in f:
            l = line.split()

            # discard invalid records
            if len(l) != 4:
                continue

            # the file name is l[2]
            if filename and l[2] != filename:
                continue

            # create new node
            etree.SubElement(self._graph, 'node', id=str(last_id))

            # create new edge
            edge = etree.SubElement(self._graph, 'edge',
                                    source=str(last_id - 1),
                                    target=str(last_id))
            etree.SubElement(edge, 'data', key='startline').text = l[3]
            etree.SubElement(edge, 'data', key='originfile').text = l[2]
            # not all of the lines are branches and also not every
            # branch evaluation corresponds to the same evaluation
            # of the source code (e.g. optimizations may negate the condition)
            # if int(l[0]) == 0:
            #     # KLEE splits the true/false inverted for some reason
            #     control = 'condition-true'
            # else:
            #     control = 'condition-false'

            # ctrlelem = etree.SubElement(edge, 'data', key='control')
            # ctrlelem.text = control

            last_id += 1

        # create the violation
        vl = etree.SubElement(self._graph, 'node', id=str(last_id))
        etree.SubElement(vl, 'data', key='violation').text = 'true'

        etree.SubElement(self._graph, 'edge',
                         source=str(last_id - 1),
                         target=str(last_id))

        # remove the control key from the last edge
        # if ctrlelem is not None:
        #     ctrlelem.getparent().remove(ctrlelem)

        f.close()

    def dump(self):
        print(etree.tostring(self._root, pretty_print=True))

    def write(self, to):
        et = etree.ElementTree(self._root)
        et.write(to, encoding='UTF-8', method="xml",
                 pretty_print=True, xml_declaration=True)

# my $output = IO::File->new(">-");
#
# my $writer = XML::Writer->new(OUTPUT => $output);
#
# $writer->xmlDecl("UTF-8", "no");
# $writer->startTag("graphml", "xmlns:xsi" => "http://www.w3.org/2001/XMLSchema-instance", "xmlns" => "http://graphml.graphdrawing.org/xmlns");
# $writer->startTag("graph");
#
# $writer->startTag("node", "id" => "A0");
# 	$writer->startTag("data", "key" => "entry");
# 		$writer->characters("true");
# 	$writer->endTag("data");
# $writer->endTag("node");
#
# my $nid = 0;
#
# while (<>) {
# 	chomp;
# 	/^[01] .* ([0-9]+)$/;
# 	my $line = $1;
#
# 	$writer->startTag("node", "id" => "A" . ($nid + 1));
# 	if (eof()) {
# 		$writer->startTag("data", "key" => "violation");
# 			$writer->characters("true");
# 		$writer->endTag("data");
# 	}
# 	$writer->endTag("node");
#
# 	$writer->startTag("edge", "source" => "A$nid",
# 			"target" => "A" . ($nid + 1));
# 		$writer->startTag("data", "key" => "startline");
# 			$writer->characters($line);
# 		$writer->endTag("data");
# 	$writer->endTag("edge");
#
# 	$nid++;
# }
#
# $writer->endTag("graph");
# $writer->endTag("graphml");
# $writer->end();
# $output->close();
#
# 1
<commit_msg>witnesses: allow run with older xml<commit_after>#!/usr/bin/python

no_lxml = False
try:
    from lxml import etree as ET
except ImportError:
    no_lxml = True

if no_lxml:
    # if this fails, then we're screwed, so let the script die
    from xml.etree import ElementTree as ET

class GraphMLWriter(object):
    def __init__(self):

	if no_lxml:
            self._root = ET.Element('graphml')
	else:
            ns = {None:'http://graphml.graphdrawing.org/xmlns/graphml'}
            self._root = ET.Element('graphml', nsmap=ns)

        self._graph = ET.SubElement(self._root, 'graph', edgedefault="directed")

        # create the entry node
        self._entry = ET.SubElement(self._graph, 'node', id='0')
        ET.SubElement(self._entry, 'data', key='entry').text = 'true'

    def parsePath(self, pathFile, filename = None):
        """
        Parse .path file from klee
        \param pathFile     the .path file
        \param filename     name of the file the symbiotic ran on
                            -- in the case that we want to stick
                            only to this file in the witness
        """

        f = open(pathFile, 'r')
        last_id=1

        #ctrlelem = None
        for line in f:
            l = line.split()

            # discard invalid records
            if len(l) != 4:
                continue

            # the file name is l[2]
            if filename and l[2] != filename:
                continue

            # create new node
            ET.SubElement(self._graph, 'node', id=str(last_id))

            # create new edge
            edge = ET.SubElement(self._graph, 'edge',
                                    source=str(last_id - 1),
                                    target=str(last_id))
            ET.SubElement(edge, 'data', key='startline').text = l[3]
            ET.SubElement(edge, 'data', key='originfile').text = l[2]
            # not all of the lines are branches and also not every
            # branch evaluation corresponds to the same evaluation
            # of the source code (e.g. optimizations may negate the condition)
            # if int(l[0]) == 0:
            #     # KLEE splits the true/false inverted for some reason
            #     control = 'condition-true'
            # else:
            #     control = 'condition-false'

            # ctrlelem = ET.SubElement(edge, 'data', key='control')
            # ctrlelem.text = control

            last_id += 1

        # create the violation
        vl = ET.SubElement(self._graph, 'node', id=str(last_id))
        ET.SubElement(vl, 'data', key='violation').text = 'true'
        ET.SubElement(self._graph, 'edge',
                         source=str(last_id - 1),
                         target=str(last_id))

        # remove the control key from the last edge
        # if ctrlelem is not None:
        #     ctrlelem.getparent().remove(ctrlelem)

        f.close()

    def dump(self):
        print(ET.tostring(self._root, pretty_print=True))

    def write(self, to):
        et = ET.ElementTree(self._root)
	if no_lxml:
            et.write(to, encoding='UTF-8', method="xml",
                     xml_declaration=True)
	else:
            et.write(to, encoding='UTF-8', method="xml",
                     pretty_print=True, xml_declaration=True)

# my $output = IO::File->new(">-");
#
# my $writer = XML::Writer->new(OUTPUT => $output);
#
# $writer->xmlDecl("UTF-8", "no");
# $writer->startTag("graphml", "xmlns:xsi" => "http://www.w3.org/2001/XMLSchema-instance", "xmlns" => "http://graphml.graphdrawing.org/xmlns");
# $writer->startTag("graph");
#
# $writer->startTag("node", "id" => "A0");
# 	$writer->startTag("data", "key" => "entry");
# 		$writer->characters("true");
# 	$writer->endTag("data");
# $writer->endTag("node");
#
# my $nid = 0;
#
# while (<>) {
# 	chomp;
# 	/^[01] .* ([0-9]+)$/;
# 	my $line = $1;
#
# 	$writer->startTag("node", "id" => "A" . ($nid + 1));
# 	if (eof()) {
# 		$writer->startTag("data", "key" => "violation");
# 			$writer->characters("true");
# 		$writer->endTag("data");
# 	}
# 	$writer->endTag("node");
#
# 	$writer->startTag("edge", "source" => "A$nid",
# 			"target" => "A" . ($nid + 1));
# 		$writer->startTag("data", "key" => "startline");
# 			$writer->characters($line);
# 		$writer->endTag("data");
# 	$writer->endTag("edge");
#
# 	$nid++;
# }
#
# $writer->endTag("graph");
# $writer->endTag("graphml");
# $writer->end();
# $output->close();
#
# 1
<|endoftext|>