<commit_before>import pytest
from os import path
import numpy as np
import yass
from yass import preprocess, detect, cluster, templates, deconvolute
from util import ReferenceTesting

def test_decovnolution(path_to_threshold_config,
                       make_tmp_folder):
    yass.set_config('tests/config_nnet.yaml')

    # FIXME: hacky solution for the test to pass, i need to re-train the
    # triage network
    CONFIG = yass.read_config()
    d = CONFIG.detect._data
    d['neural_network_triage']['threshold_collision'] = 0
    CONFIG._set_param('detect', d)

    (standarized_path,
     standarized_params,
     whiten_filter) = preprocess.run(output_directory=make_tmp_folder)

    (score,
     spike_index_clear,
     spike_index_all) = detect.run(standarized_path,
                                   standarized_params,
                                   whiten_filter,
                                   output_directory=make_tmp_folder)

    spike_train_clear, tmp_loc, vbParam = cluster.run(
        spike_index_clear,
        output_directory=make_tmp_folder)

    (templates_, spike_train,
     groups, idx_good_templates) = templates.run(
        spike_train_clear, tmp_loc,
        output_directory=make_tmp_folder)

    deconvolute.run(spike_index_all, templates_,
                    output_directory=make_tmp_folder)

@pytest.mark.xfail
def test_deconvolution_returns_expected_results(path_to_threshold_config,
                                                path_to_output_reference,
                                                make_tmp_folder):
    np.random.seed(0)

    yass.set_config(path_to_threshold_config)

    (standarized_path,
        standarized_params,
        whiten_filter) = preprocess.run(output_directory=make_tmp_folder)

    (score, spike_index_clear,
     spike_index_all) = detect.run(standarized_path,
                                   standarized_params,
                                   whiten_filter,
                                   output_directory=make_tmp_folder)

    (spike_train_clear,
        tmp_loc,
        vbParam) = cluster.run(spike_index_clear,
                               output_directory=make_tmp_folder)

    (templates_, spike_train,
     groups,
     idx_good_templates) = templates.run(spike_train_clear, tmp_loc,
                                         output_directory=make_tmp_folder)

    spike_train = deconvolute.run(spike_index_all, templates_,
                                  output_directory=make_tmp_folder)

    path_to_spike_train = path.join(path_to_output_reference,
                                    'spike_train.npy')

    ReferenceTesting.assert_array_equal(spike_train, path_to_spike_train)
<commit_msg>fixes deconv test<commit_after>import pytest
from os import path
import numpy as np
import yass
from yass import preprocess, detect, cluster, templates, deconvolute
from util import ReferenceTesting

def test_decovnolution(path_to_nnet_config,
                       make_tmp_folder):
    yass.set_config(path_to_nnet_config, make_tmp_folder)

    # FIXME: hacky solution for the test to pass, i need to re-train the
    # triage network
    CONFIG = yass.read_config()
    d = CONFIG.detect._data
    d['neural_network_triage']['threshold_collision'] = 0
    CONFIG._set_param('detect', d)

    (standarized_path,
     standarized_params,
     whiten_filter) = preprocess.run(output_directory=make_tmp_folder)

    (score,
     spike_index_clear,
     spike_index_all) = detect.run(standarized_path,
                                   standarized_params,
                                   whiten_filter,
                                   output_directory=make_tmp_folder)

    spike_train_clear, tmp_loc, vbParam = cluster.run(
        spike_index_clear,
        output_directory=make_tmp_folder)

    (templates_, spike_train,
     groups, idx_good_templates) = templates.run(
        spike_train_clear, tmp_loc,
        output_directory=make_tmp_folder)

    deconvolute.run(spike_index_all, templates_,
                    output_directory=make_tmp_folder)

@pytest.mark.xfail
def test_deconvolution_returns_expected_results(path_to_threshold_config,
                                                path_to_output_reference,
                                                make_tmp_folder):
    np.random.seed(0)

    yass.set_config(path_to_threshold_config)

    (standarized_path,
        standarized_params,
        whiten_filter) = preprocess.run(output_directory=make_tmp_folder)

    (score, spike_index_clear,
     spike_index_all) = detect.run(standarized_path,
                                   standarized_params,
                                   whiten_filter,
                                   output_directory=make_tmp_folder)

    (spike_train_clear,
        tmp_loc,
        vbParam) = cluster.run(spike_index_clear,
                               output_directory=make_tmp_folder)

    (templates_, spike_train,
     groups,
     idx_good_templates) = templates.run(spike_train_clear, tmp_loc,
                                         output_directory=make_tmp_folder)

    spike_train = deconvolute.run(spike_index_all, templates_,
                                  output_directory=make_tmp_folder)

    path_to_spike_train = path.join(path_to_output_reference,
                                    'spike_train.npy')

    ReferenceTesting.assert_array_equal(spike_train, path_to_spike_train)
<|endoftext|>


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

import argparse
import sys

class Error(Exception):
    """Base class for custom exceptions."""
    pass

class GridStringError(Error):
    """Exception raised for errors in the string representation of a Sudoku grid."""
    pass

class VertexValueError(Error):
    """Exception raised for invalid values on the Sudoku grid."""
    pass

class Sudoku:
    """Sudoku grid."""

    def __init__(self, grid_string=""):
        if not isinstance(grid_string, str):
            raise GridStringError()
        self.grid = [[None for _dummy in range(9)] for _dummy in range(9)]
        if grid_string == "":
            return
        if len(grid_string) != 81:
            raise GridStringError()
        position = 0
        for (line, row) in [(ln, rw) for ln in range(9) for rw in range(9)]:
            if grid_string[position] == ".":
                pass
            elif int(grid_string[position]) in range(1, 10):
                self.grid[line][row] = int(grid_string[position])
            else:
                raise GridStringError()
            position += 1

    def verify_vertex_values(self):
        """Verify that the current grid has only valid vertex values."""
        for (line, row) in [(ln, rw) for ln in range(9) for rw in range(9)]:
            if self.grid[line][row] not in range(1, 10) and self.grid[line][row] is not None:
                raise VertexValueError()

    def __str__(self):
        # Verify correct vertex values
        self.verify_vertex_values()
        # Generate output string
        output_string = ""
        for (line, row) in [(ln, rw) for ln in range(9) for rw in range(9)]:
            if self.grid[line][row] is None:
                output_string += "."
            else:
                output_string += str(self.grid[line][row])
        return output_string

    def formatted_str(self):
        """Return formatted string for printing the grid."""
        # Verify correct vertex values
        self.verify_vertex_values()
        # Generate output string
        output_string = ""
        for line in range(9):
            for row in range(9):
                if self.grid[line][row] is None:
                    output_string += "."
                else:
                    output_string += str(self.grid[line][row])
                if row != 8:
                    output_string += " "
                elif line != 8:
                    output_string += "\n"
                if row in [2, 5]:
                    output_string += "| "
            if line in [2, 5]:
                output_string += "------+-------+------\n"
        return output_string

    def candidate_map(self):
        """Return two-dimensional list of sets composed of possible candidates for each vertex."""
        candidates = [[set(range(1, 10)) for _dummy in range(9)] for _dummy in range(9)]
        vertex_value_unknown = [[True for _dummy in range(9)] for _dummy in range(9)]
        for (line, row) in [(ln, rw) for ln in range(9) for rw in range(9)]:
            if self.grid[line][row] in range(1, 10):
                candidates[line][row] = set([self.grid[line][row]])
                vertex_value_unknown[line][row] = False
                for i in range(9):
                    if i != row:
                        candidates[line][i].discard(self.grid[line][row])
                    if i != line:
                        candidates[i][row].discard(self.grid[line][row])
                    if line - line%3 + i//3 != line or row - row%3 + i%3 != row:
                        candidates[line - line%3 + i//3][row - row%3 + i%3].discard(self.grid[line][row])
        # Further reduce candidate map
        reduce_cadidate_map_further = True
        while reduce_cadidate_map_further:
            reduce_cadidate_map_further = False
            total_number_of_candidates = sum([len(candidates[ln][rw]) for ln in range(9) for rw in range(9)])
            for number in range(1, 10):
                for i in range(9):
                    # Check for single possible vertex for *number* in candidate map line *i*
                    seen_in_j = []
                    for j in range(9):
                        if number in candidates[i][j]:
                            seen_in_j.append(j)
                    if len(seen_in_j) == 1 and vertex_value_unknown[i][seen_in_j[0]]:
                        candidates[i][seen_in_j[0]] = set([number])
                        vertex_value_unknown[i][seen_in_j[0]] = False
                        # Discard other candidates for *number* in corresponding row and subsquare
                        for j in range(9):
                            if j != i:
                                candidates[j][seen_in_j[0]].discard(number)
                            if i - i%3 + j//3 != i:
                                candidates[i - i%3 + j//3][seen_in_j[0] - seen_in_j[0]%3 + j%3].discard(number)
                    # otherwise add check wheter all candidates for *number* are in the same subsquare
                    elif 1 < len(seen_in_j) < 4:
                        subsquares = set()
                        for j in seen_in_j:
                            subsquares.add(3*(i//3) + j//3)
                        if len(subsquares) == 1:
                            subsquare = subsquares.pop()
                            for j in range(9):
                                if 3*(subsquare//3) + j//3 != i:
                                    candidates[3*(subsquare//3) + j//3][3*(subsquare%3) + j%3].discard(number)
                    # Check for single possible vertex for *number* in candidate map row *i*
                    seen_in_j = []
                    for j in range(9):
                        if number in candidates[j][i]:
                            seen_in_j.append(j)
                    if len(seen_in_j) == 1 and vertex_value_unknown[seen_in_j[0]][i]:
                        candidates[seen_in_j[0]][i] = set([number])
                        vertex_value_unknown[seen_in_j[0]][i] = False
                        # Discard other candidates for *number* in corresponding line and subsquare
                        for j in range(9):
                            if j != i:
                                candidates[seen_in_j[0]][j].discard(number)
                            if i - i%3 + j%3 != i:
                                candidates[seen_in_j[0] - seen_in_j[0]%3 + j//3][i - i%3 + j%3].discard(number)
                    # otherwise add check wheter all candidates for *number* are in the same subsquare
                    elif 1 < len(seen_in_j) < 4:
                        subsquares = set()
                        for j in seen_in_j:
                            subsquares.add(3*(j//3) + i//3)
                        if len(subsquares) == 1:
                            subsquare = subsquares.pop()
                            for j in range(9):
                                if 3*(subsquare%3) + j%3 != i:
                                    candidates[3*(subsquare//3) + j//3][3*(subsquare%3) + j%3].discard(number)
                    # Check for single possible vertex for *number* in candidate map subsquare *i*
                    seen_in_j = []
                    for j in range(9):
                        if number in candidates[3*(i//3) + j//3][3*(i%3) + j%3]:
                            seen_in_j.append(j)
                    if len(seen_in_j) == 1 and vertex_value_unknown[3*(i//3) + seen_in_j[0]//3][3*(i%3) + seen_in_j[0]%3]:
                        candidates[3*(i//3) + seen_in_j[0]//3][3*(i%3) + seen_in_j[0]%3] = set([number])
                        vertex_value_unknown[3*(i//3) + seen_in_j[0]//3][3*(i%3) + seen_in_j[0]%3] = False
                        # Discard other candidates for *number* in corresponding line and row
                        for j in range(9):
                            if j not in [3*(i%3), 3*(i%3) + 1, 3*(i%3) + 2]:
                                candidates[3*(i//3) + seen_in_j[0]//3][j].discard(number)
                            if j not in [3*(i//3), 3*(i//3) + 1, 3*(i//3) + 2]:
                                candidates[j][3*(i%3) + seen_in_j[0]%3].discard(number)
                    # otherwise add check wheter all candidates for *number* are in the same line/row
                    elif 1 < len(seen_in_j) < 4:
                        lines = set()
                        rows = set()
                        for j in seen_in_j:
                            lines.add(3*(i//3) + j//3)
                            rows.add(3*(i%3) + j%3)
                        if len(lines) == 1:
                            line = lines.pop()
                            for row in [rw for rw in range(9) if rw not in [3*(i%3), 3*(i%3) + 1, 3*(i%3) + 2]]:
                                candidates[line][row].discard(number)
                        elif len(rows) == 1:
                            row = rows.pop()
                            for line in [ln for ln in range(9) if ln not in [3*(i//3), 3*(i//3) + 1, 3*(i//3) + 2]]:
                                candidates[line][row].discard(number)
            if sum([len(candidates[ln][rw]) for ln in range(9) for rw in range(9)]) < total_number_of_candidates:
                reduce_cadidate_map_further = True
        return candidates

    def valid(self):
        """Check whether the current grid is valid and return bool."""
        # Verify correct vertex values
        self.verify_vertex_values()
        # Check for duplicate values in lines
        for line in range(9):
            seen = []
            for row in range(9):
                if self.grid[line][row] is None:
                    pass
                elif self.grid[line][row] in seen:
                    return False
                else:
                    seen.append(self.grid[line][row])
        # Check for duplicate values in rows
        for row in range(9):
            seen = []
            for line in range(9):
                if self.grid[line][row] is None:
                    pass
                elif self.grid[line][row] in seen:
                    return False
                else:
                    seen.append(self.grid[line][row])
        # Check for duplicate values in subgrids
        for (subgrid_line, subgrid_row) in [(subg_ln, subg_rw) for subg_ln in range(3) for subg_rw in range(3)]:
            seen = []
            for (line, row) in [(ln, rw) for ln in range(3) for rw in range(3)]:
                if self.grid[3*subgrid_line + line][3*subgrid_row + row] is None:
                    pass
                elif self.grid[3*subgrid_line + line][3*subgrid_row + row] in seen:
                    return False
                else:
                    seen.append(self.grid[3*subgrid_line + line][3*subgrid_row + row])
        # No duplicates found
        return True

def grids_from_stdin():
    """Read stdin, return list of Sudoku instances, one per 81 correct input characters."""
    grids = []
    filtered_stdin_data = "".join([char for char in sys.stdin.read() if char in ["1", "2", "3", "4", "5", "6", "7", "8", "9", "."]])
    while len(filtered_stdin_data) >= 81:
        grids.append(Sudoku(filtered_stdin_data[:81]))
        filtered_stdin_data = filtered_stdin_data[81:]
    return grids

def solve_puzzle(grid):
    """Solve the given puzzle, return solutions as list of Sudoku instances."""
    solutions = []
    if not grid.valid():
        return solutions
    # Backtracking, iterating over (first) smallest list of candidates for empty vertices
    candidates = grid.candidate_map()
    min_number_of_candidates = min([9] + [len(candidates[ln][rw]) for ln in range(9) for rw in range(9) if grid.grid[ln][rw] is None])
    for (line, row) in [(ln, rw) for ln in range(9) for rw in range(9)]:
        if grid.grid[line][row] is None and len(candidates[line][row]) == min_number_of_candidates:
            for guess in candidates[line][row]:
                grid.grid[line][row] = guess
                for solution in solve_puzzle(grid):
                    solutions.append(solution)
            grid.grid[line][row] = None
            break
    else:
        solutions.append(Sudoku(grid.__str__()))
    return solutions

# Main
argument_parser = argparse.ArgumentParser(description="Sudoku tool.", allow_abbrev=False)
argument_subparsers = argument_parser.add_subparsers(dest="mode", help="select program mode")
argument_subparsers.required = True
argument_parser_solve = argument_subparsers.add_parser("solve", help="solve Sudoku puzzle", description="Solve Sudoku puzzle.")
argument_parser_solve.add_argument("-H", action="store_true", help="beautify output for human parsers")
#argument_parser_solve.add_argument("-o", "--output", help="output to file OUTPUT instead of stdout")
#argument_parser_solve.add_argument("input", help="input file(s), omitting inplies stdin", nargs="*")
args = argument_parser.parse_args()

if args.mode == "solve":
    input_grids = grids_from_stdin()
    for grid in input_grids:
        solutions = solve_puzzle(grid)
        if args.H:
            print("Puzzle:")
            print(grid.formatted_str())
            print()
            if len(solutions) == 0:
                print("No solution.")
            elif len(solutions) == 1:
                print("1 solution:")
            else:
                print(len(solutions), "solutions:")
            for solution in solutions:
                print(solution.formatted_str())
                print()
        else:
            print("p", grid, sep="")
            for solution in solutions:
                print("s", solution, sep="")
<commit_msg>Add reading input from file(s)<commit_after>#!/usr/bin/env python3

import argparse
import sys

class Error(Exception):
    """Base class for custom exceptions."""
    pass

class GridStringError(Error):
    """Exception raised for errors in the string representation of a Sudoku grid."""
    pass

class VertexValueError(Error):
    """Exception raised for invalid values on the Sudoku grid."""
    pass

class Sudoku:
    """Sudoku grid."""

    def __init__(self, grid_string=""):
        if not isinstance(grid_string, str):
            raise GridStringError()
        self.grid = [[None for _dummy in range(9)] for _dummy in range(9)]
        if grid_string == "":
            return
        if len(grid_string) != 81:
            raise GridStringError()
        position = 0
        for (line, row) in [(ln, rw) for ln in range(9) for rw in range(9)]:
            if grid_string[position] == ".":
                pass
            elif int(grid_string[position]) in range(1, 10):
                self.grid[line][row] = int(grid_string[position])
            else:
                raise GridStringError()
            position += 1

    def verify_vertex_values(self):
        """Verify that the current grid has only valid vertex values."""
        for (line, row) in [(ln, rw) for ln in range(9) for rw in range(9)]:
            if self.grid[line][row] not in range(1, 10) and self.grid[line][row] is not None:
                raise VertexValueError()

    def __str__(self):
        # Verify correct vertex values
        self.verify_vertex_values()
        # Generate output string
        output_string = ""
        for (line, row) in [(ln, rw) for ln in range(9) for rw in range(9)]:
            if self.grid[line][row] is None:
                output_string += "."
            else:
                output_string += str(self.grid[line][row])
        return output_string

    def formatted_str(self):
        """Return formatted string for printing the grid."""
        # Verify correct vertex values
        self.verify_vertex_values()
        # Generate output string
        output_string = ""
        for line in range(9):
            for row in range(9):
                if self.grid[line][row] is None:
                    output_string += "."
                else:
                    output_string += str(self.grid[line][row])
                if row != 8:
                    output_string += " "
                elif line != 8:
                    output_string += "\n"
                if row in [2, 5]:
                    output_string += "| "
            if line in [2, 5]:
                output_string += "------+-------+------\n"
        return output_string

    def candidate_map(self):
        """Return two-dimensional list of sets composed of possible candidates for each vertex."""
        candidates = [[set(range(1, 10)) for _dummy in range(9)] for _dummy in range(9)]
        vertex_value_unknown = [[True for _dummy in range(9)] for _dummy in range(9)]
        for (line, row) in [(ln, rw) for ln in range(9) for rw in range(9)]:
            if self.grid[line][row] in range(1, 10):
                candidates[line][row] = set([self.grid[line][row]])
                vertex_value_unknown[line][row] = False
                for i in range(9):
                    if i != row:
                        candidates[line][i].discard(self.grid[line][row])
                    if i != line:
                        candidates[i][row].discard(self.grid[line][row])
                    if line - line%3 + i//3 != line or row - row%3 + i%3 != row:
                        candidates[line - line%3 + i//3][row - row%3 + i%3].discard(self.grid[line][row])
        # Further reduce candidate map
        reduce_cadidate_map_further = True
        while reduce_cadidate_map_further:
            reduce_cadidate_map_further = False
            total_number_of_candidates = sum([len(candidates[ln][rw]) for ln in range(9) for rw in range(9)])
            for number in range(1, 10):
                for i in range(9):
                    # Check for single possible vertex for *number* in candidate map line *i*
                    seen_in_j = []
                    for j in range(9):
                        if number in candidates[i][j]:
                            seen_in_j.append(j)
                    if len(seen_in_j) == 1 and vertex_value_unknown[i][seen_in_j[0]]:
                        candidates[i][seen_in_j[0]] = set([number])
                        vertex_value_unknown[i][seen_in_j[0]] = False
                        # Discard other candidates for *number* in corresponding row and subsquare
                        for j in range(9):
                            if j != i:
                                candidates[j][seen_in_j[0]].discard(number)
                            if i - i%3 + j//3 != i:
                                candidates[i - i%3 + j//3][seen_in_j[0] - seen_in_j[0]%3 + j%3].discard(number)
                    # otherwise add check wheter all candidates for *number* are in the same subsquare
                    elif 1 < len(seen_in_j) < 4:
                        subsquares = set()
                        for j in seen_in_j:
                            subsquares.add(3*(i//3) + j//3)
                        if len(subsquares) == 1:
                            subsquare = subsquares.pop()
                            for j in range(9):
                                if 3*(subsquare//3) + j//3 != i:
                                    candidates[3*(subsquare//3) + j//3][3*(subsquare%3) + j%3].discard(number)
                    # Check for single possible vertex for *number* in candidate map row *i*
                    seen_in_j = []
                    for j in range(9):
                        if number in candidates[j][i]:
                            seen_in_j.append(j)
                    if len(seen_in_j) == 1 and vertex_value_unknown[seen_in_j[0]][i]:
                        candidates[seen_in_j[0]][i] = set([number])
                        vertex_value_unknown[seen_in_j[0]][i] = False
                        # Discard other candidates for *number* in corresponding line and subsquare
                        for j in range(9):
                            if j != i:
                                candidates[seen_in_j[0]][j].discard(number)
                            if i - i%3 + j%3 != i:
                                candidates[seen_in_j[0] - seen_in_j[0]%3 + j//3][i - i%3 + j%3].discard(number)
                    # otherwise add check wheter all candidates for *number* are in the same subsquare
                    elif 1 < len(seen_in_j) < 4:
                        subsquares = set()
                        for j in seen_in_j:
                            subsquares.add(3*(j//3) + i//3)
                        if len(subsquares) == 1:
                            subsquare = subsquares.pop()
                            for j in range(9):
                                if 3*(subsquare%3) + j%3 != i:
                                    candidates[3*(subsquare//3) + j//3][3*(subsquare%3) + j%3].discard(number)
                    # Check for single possible vertex for *number* in candidate map subsquare *i*
                    seen_in_j = []
                    for j in range(9):
                        if number in candidates[3*(i//3) + j//3][3*(i%3) + j%3]:
                            seen_in_j.append(j)
                    if len(seen_in_j) == 1 and vertex_value_unknown[3*(i//3) + seen_in_j[0]//3][3*(i%3) + seen_in_j[0]%3]:
                        candidates[3*(i//3) + seen_in_j[0]//3][3*(i%3) + seen_in_j[0]%3] = set([number])
                        vertex_value_unknown[3*(i//3) + seen_in_j[0]//3][3*(i%3) + seen_in_j[0]%3] = False
                        # Discard other candidates for *number* in corresponding line and row
                        for j in range(9):
                            if j not in [3*(i%3), 3*(i%3) + 1, 3*(i%3) + 2]:
                                candidates[3*(i//3) + seen_in_j[0]//3][j].discard(number)
                            if j not in [3*(i//3), 3*(i//3) + 1, 3*(i//3) + 2]:
                                candidates[j][3*(i%3) + seen_in_j[0]%3].discard(number)
                    # otherwise add check wheter all candidates for *number* are in the same line/row
                    elif 1 < len(seen_in_j) < 4:
                        lines = set()
                        rows = set()
                        for j in seen_in_j:
                            lines.add(3*(i//3) + j//3)
                            rows.add(3*(i%3) + j%3)
                        if len(lines) == 1:
                            line = lines.pop()
                            for row in [rw for rw in range(9) if rw not in [3*(i%3), 3*(i%3) + 1, 3*(i%3) + 2]]:
                                candidates[line][row].discard(number)
                        elif len(rows) == 1:
                            row = rows.pop()
                            for line in [ln for ln in range(9) if ln not in [3*(i//3), 3*(i//3) + 1, 3*(i//3) + 2]]:
                                candidates[line][row].discard(number)
            if sum([len(candidates[ln][rw]) for ln in range(9) for rw in range(9)]) < total_number_of_candidates:
                reduce_cadidate_map_further = True
        return candidates

    def valid(self):
        """Check whether the current grid is valid and return bool."""
        # Verify correct vertex values
        self.verify_vertex_values()
        # Check for duplicate values in lines
        for line in range(9):
            seen = []
            for row in range(9):
                if self.grid[line][row] is None:
                    pass
                elif self.grid[line][row] in seen:
                    return False
                else:
                    seen.append(self.grid[line][row])
        # Check for duplicate values in rows
        for row in range(9):
            seen = []
            for line in range(9):
                if self.grid[line][row] is None:
                    pass
                elif self.grid[line][row] in seen:
                    return False
                else:
                    seen.append(self.grid[line][row])
        # Check for duplicate values in subgrids
        for (subgrid_line, subgrid_row) in [(subg_ln, subg_rw) for subg_ln in range(3) for subg_rw in range(3)]:
            seen = []
            for (line, row) in [(ln, rw) for ln in range(3) for rw in range(3)]:
                if self.grid[3*subgrid_line + line][3*subgrid_row + row] is None:
                    pass
                elif self.grid[3*subgrid_line + line][3*subgrid_row + row] in seen:
                    return False
                else:
                    seen.append(self.grid[3*subgrid_line + line][3*subgrid_row + row])
        # No duplicates found
        return True

def grids_from_stdin():
    """Read stdin, return list of Sudoku instances, one per 81 correct input characters."""
    grids = []
    filtered_stdin_data = "".join([char for char in sys.stdin.read() if char in ["1", "2", "3", "4", "5", "6", "7", "8", "9", "."]])
    while len(filtered_stdin_data) >= 81:
        grids.append(Sudoku(filtered_stdin_data[:81]))
        filtered_stdin_data = filtered_stdin_data[81:]
    return grids

def grids_from_files(files):
    """Take list of filenames, return list of Sudoku instances, one per 81 correct input characters in each file."""
    grids = []
    for filename in files:
        with open(filename) as f:
            filtered_file_data = "".join([char for char in f.read() if char in ["1", "2", "3", "4", "5", "6", "7", "8", "9", "."]])
        while len(filtered_file_data) >= 81:
            grids.append(Sudoku(filtered_file_data[:81]))
            filtered_file_data = filtered_file_data[81:]
    return grids

def solve_puzzle(grid):
    """Solve the given puzzle, return solutions as list of Sudoku instances."""
    solutions = []
    if not grid.valid():
        return solutions
    # Backtracking, iterating over (first) smallest list of candidates for empty vertices
    candidates = grid.candidate_map()
    min_number_of_candidates = min([9] + [len(candidates[ln][rw]) for ln in range(9) for rw in range(9) if grid.grid[ln][rw] is None])
    for (line, row) in [(ln, rw) for ln in range(9) for rw in range(9)]:
        if grid.grid[line][row] is None and len(candidates[line][row]) == min_number_of_candidates:
            for guess in candidates[line][row]:
                grid.grid[line][row] = guess
                for solution in solve_puzzle(grid):
                    solutions.append(solution)
            grid.grid[line][row] = None
            break
    else:
        solutions.append(Sudoku(grid.__str__()))
    return solutions

# Main
argument_parser = argparse.ArgumentParser(description="Sudoku tool.", allow_abbrev=False)
argument_subparsers = argument_parser.add_subparsers(dest="mode", help="select program mode")
argument_subparsers.required = True
argument_parser_solve = argument_subparsers.add_parser("solve", help="solve Sudoku puzzle", description="Solve Sudoku puzzle.")
argument_parser_solve.add_argument("-H", action="store_true", help="beautify output for human parsers")
#argument_parser_solve.add_argument("-o", "--output", help="output to file OUTPUT instead of stdout")
argument_parser_solve.add_argument("input", help="input file(s), omitting implies stdin", nargs="*")
args = argument_parser.parse_args()

if args.mode == "solve":
    if len(args.input) == 0:
        input_grids = grids_from_stdin()
    else:
        input_grids = grids_from_files(args.input)
    for grid in input_grids:
        solutions = solve_puzzle(grid)
        if args.H:
            print("Puzzle:")
            print(grid.formatted_str())
            print()
            if len(solutions) == 0:
                print("No solution.")
            elif len(solutions) == 1:
                print("1 solution:")
            else:
                print(len(solutions), "solutions:")
            for solution in solutions:
                print(solution.formatted_str())
                print()
        else:
            print("p", grid, sep="")
            for solution in solutions:
                print("s", solution, sep="")
<|endoftext|>


<commit_before>"""
Module for solving Sudoku Puzzles
"""

class Cell(object):
    def __init__(self, value=None):
        self.solved = False
        self.excludes = set()
        self.value = value if value is not None else -1
        self.row = None
        self.column = None
        self.box = None

    def exclude(self):
        # add all values in parent row, column and box to the excluded set
        pass

    def set_value(self, value, update_parents=True):
        """
        Set the value of this cell and update cells in parent boxes, rows
        and coulmuns
        """

        self.solved = True
        self.value = value
        if update_parents:
            self.update_parents()

    def update_parents(self):
        self.row.update_excludes()
        self.column.update_excludes()
        self.box.update_excludes()

class Group(object):
    """
    A group of 9 values, this could be a box or a line
    """
    def __init__(self, cells=None):
        self.cells = cells if cells is not None else []

    def solve(self):
        """
        Look for cells in this group that must have a certain value

        For all the unsolved cells in the group, go over each value and
        if the value can go in that cell, and cannot go in all other
        cells then that cell must contain that value
        """

        for index in range(9):
            others = [i for i in range(9) if i != index]
            for value in range(1,10):
                # If the cell is not already solved
                if not self.cells[index].solved:
                    # If this cell cannot hold this value, continue to next
                    # value
                    if value in self.cells[index].excludes:
                        continue
                    # If all other cells cannot hold this value, this cell
                    # must be this value. Break and go to next index
                    if all([value in self.cells[other].excludes
                            for other in others]):
                        self.cells[index].set_value(value)
                        break

    def update_excludes(self):
        """
        Update all the cells in this group based on the solved values
        """

        # collect solved values
        solved_values = set()
        for index in range(9):
            cell = self.cells[index]
            if cell.solved:
                solved_values.add(cell.value)

        # Update the cells with found values
        for index in range(9):
            self.cells[index].excludes.update(solved_values)

class Grid(object):
    def __init__(self):
        self.cells = [Cell() for _ in range(81)]
        self.rows = [Group(cells=self.cells[i*9:(i*9)+8]) for i in range(9)]
        self.columns = [Group(cells=self.cells[i:81:9]) for i in range(9)]
        self.boxes = [
            Group(cells=
                [self.cells[i], self.cells[i+1], self.cells[i+2],
                 self.cells[i+9], self.cells[i+10], self.cells[i+11],
                 self.cells[i+18], self.cells[i+19], self.cells[i+20]]
                 ) for i in [0,3,6,27,30,33,54,57,60]
            ]

    def set_test_data(self):
        for i in range(81):
            self.cells[i].set_value(i, update_parents=False)

    def display(self):
        for row in range(9):
            print "|".join(["{0:2}".format(self.cells[(row*9)+i].value) for i in range(9)])

def experiments():
    grid = Grid()
    grid.set_test_data()

    print grid.cells[0].value
    print grid.rows[0].cells[0].value
    print grid.columns[0].cells[0].value
    print grid.boxes[0].cells[0].value

    grid.cells[0].value = 99

    print grid.cells[0].value
    print grid.rows[0].cells[0].value
    print grid.columns[0].cells[0].value
    print grid.boxes[0].cells[0].value

    grid.columns[0].cells[0].value = 999

    print grid.cells[0].value
    print grid.rows[0].cells[0].value
    print grid.columns[0].cells[0].value
    print grid.boxes[0].cells[0].value

    print grid.cells[42].value
    print grid.rows[4].cells[6].value
    print grid.columns[6].cells[4].value
    print grid.boxes[5].cells[3].value

    grid.boxes[0].cells[0].value = 0

    grid.display()

    # unsolved_cell_indices = [1,2,3]
    # unsolved_values = [4,5,6]

    # for index in unsolved_cell_indices:
    #     for value in unsolved_values[:]:
    #         print unsolved_values
    #         if index == 2:
    #             if 5 in unsolved_values:
    #                 unsolved_values.remove(5)

experiments()

        # self.cell_to_noncolumn = {
        #     0: [1,2],
        #     3: [1,2],
        #     6: [1,2],
        #     1: [0,2],
        #     4: [0,2],
        #     7: [0,2],
        #     2: [0,1],
        #     5: [0,1],
        #     8: [0,1]
        #     }
        # self.cell_to_nonrow = {
        #     0: [1,2],
        #     1: [1,2],
        #     2: [1,2],
        #     3: [0,2],
        #     4: [0,2],
        #     5: [0,2],
        #     6: [0,1],
        #     7: [0,1],
        #     8: [0,1]
        #     }<commit_msg>Initial looping grid solve<commit_after>"""
Module for solving Sudoku Puzzles
"""

class Cell(object):
    def __init__(self, index, value=None):
        self.index = index
        self.solved = False
        self.excludes = set()
        self.value = value if value is not None else " "
        self.groups = []

    def set_value(self, value, update_groups=True):
        """
        Set the value of this cell and update cells in parent boxes, rows
        and coulmuns
        """
        self.solved = True
        self.value = value
        if update_groups:
            self.update_groups()

    def update_groups(self):
        for group in self.groups:
            group.update_excludes()

class Group(object):
    """
    A group of 9 values. Represents columns, rows and boxes.
    """
    def __init__(self, cells=None):
        self.cells = cells if cells is not None else []

    def solve(self):
        """
        Look for cells in this group that must have a certain value

        For all the unsolved cells in the group, go over each value and
        if the value can go in that cell, and cannot go in all other
        cells then that cell must contain that value
        """
        solved_cell = False
        for index in range(9):
            others = [i for i in range(9) if i != index]
            for value in range(1,10):
                # If the cell is not already solved
                if not self.cells[index].solved:
                    # If this cell cannot hold this value, continue to next
                    # value
                    if value in self.cells[index].excludes:
                        continue
                    # If all other cells cannot hold this value, this cell
                    # must be this value. Break and go to next index
                    # print index, value, others, len(self.cells)
                    if all([value in self.cells[other].excludes
                            for other in others]):
                        self.cells[index].set_value(value)
                        solved_cell = True
                        break
        return solved_cell

    def update_excludes(self):
        """
        Update all the cells in this group based on the solved values
        """
        # collect solved values
        solved_values = set()
        for cell in self.cells:
            if cell.solved:
                solved_values.add(cell.value)

        # Update the cells with found values
        for cell in self.cells:
            cell.excludes.update(solved_values)

    def reference_group_to_cells(self):
        """
        Add a reference to this group for all the cells it contains
        """
        for cell in self.cells:
            cell.groups.append(self)

class Grid(object):
    def __init__(self):
        self.cells = [Cell(index) for index in range(81)]
        self.rows = [Group(cells=self.cells[i*9:(i*9)+9]) for i in range(9)]
        self.columns = [Group(cells=self.cells[i:81:9]) for i in range(9)]
        self.boxes = [
            Group(cells=
                [self.cells[i], self.cells[i+1], self.cells[i+2],
                 self.cells[i+9], self.cells[i+10], self.cells[i+11],
                 self.cells[i+18], self.cells[i+19], self.cells[i+20]]
                 ) for i in [0,3,6,27,30,33,54,57,60]
            ]
        for row in self.rows:
            row.reference_group_to_cells()
        for column in self.columns:
            column.reference_group_to_cells()
        for box in self.boxes:
            box.reference_group_to_cells()

    def set_test_data(self):
        for i in range(81):
            self.cells[i].set_value(i, update_groups=False)

    def set_data_from_list(self, values):
        """
        Initiate all the cell values from a list
        """
        for index, value in enumerate(values):
            if value:
                self.cells[index].set_value(value, update_groups=False)

    def display(self):
        for row in range(9):
            print "|".join(["{0:2}".format(self.cells[(row*9)+i].value) for i in range(9)])

    def fill_excludes(self):
        for cell in self.cells:
            cell.update_groups()

    def solve(self, limit=None):
        solved_cell = True
        for _ in range(10):
            solved_cell = False
            for row in grid.rows:
                solved_cell = row.solve() or solved_cell
            for column in grid.columns:
                solved_cell = column.solve() or solved_cell
            for box in grid.boxes:
                solved_cell = box.solve() or solved_cell
            print solved_cell

def experiments():

    easy_puzzle = [
        0, 2, 0,  1, 7, 8,  0, 3, 0,
        0, 4, 0,  3, 0, 2,  0, 9, 0,
        1, 0, 0,  0, 0, 0,  0, 0, 6,

        0, 0, 8,  6, 0, 3,  5, 0, 0,
        3, 0, 0,  0, 0, 0,  0, 0, 4,
        0, 0, 6,  7, 0, 9,  2, 0, 0,

        9, 0, 0,  0, 0, 0,  0, 0, 2,
        0, 8, 0,  9, 0, 1,  0, 6, 0,
        0, 1, 0,  4, 3, 6,  0, 5, 0
        ]

    grid = Grid()
    # grid.set_test_data()
    grid.set_data_from_list(easy_puzzle)
    grid.fill_excludes()

    # for row in grid.rows:
    #     print len(row.cells)
    # for column in grid.columns:
    #     print len(column.cells)
    # for box in grid.boxes:
    #     print len(box.cells)

    grid.display()

    grid.solve()

    grid.display()

    # print grid.cells[0].value
    # print grid.rows[0].cells[0].value
    # print grid.columns[0].cells[0].value
    # print grid.boxes[0].cells[0].value

    # grid.cells[0].value = 99

    # print grid.cells[0].value
    # print grid.rows[0].cells[0].value
    # print grid.columns[0].cells[0].value
    # print grid.boxes[0].cells[0].value

    # grid.columns[0].cells[0].value = 999

    # print grid.cells[0].value
    # print grid.rows[0].cells[0].value
    # print grid.columns[0].cells[0].value
    # print grid.boxes[0].cells[0].value

    # print grid.cells[42].value
    # print grid.rows[4].cells[6].value
    # print grid.columns[6].cells[4].value
    # print grid.boxes[5].cells[3].value
    # print grid.cells[42].groups[2].cells[3].value
    # print grid.cells[42].groups[0].cells[6].value
    # print grid.cells[42].groups[1].cells[4].value
    # grid.cells[42].groups[1].cells[4].value = -1
    # print grid.cells[42].value
    # print grid.rows[4].cells[6].value
    # print grid.columns[6].cells[4].value
    # print grid.boxes[5].cells[3].value
    # print grid.cells[42].groups[2].cells[3].value
    # print grid.cells[42].groups[0].cells[6].value
    # print grid.cells[42].groups[1].cells[4].value
    # grid.boxes[0].cells[0].value = " "

    # grid.display()

    # unsolved_cell_indices = [1,2,3]
    # unsolved_values = [4,5,6]

    # for index in unsolved_cell_indices:
    #     for value in unsolved_values[:]:
    #         print unsolved_values
    #         if index == 2:
    #             if 5 in unsolved_values:
    #                 unsolved_values.remove(5)

experiments()

        # self.cell_to_noncolumn = {
        #     0: [1,2],
        #     3: [1,2],
        #     6: [1,2],
        #     1: [0,2],
        #     4: [0,2],
        #     7: [0,2],
        #     2: [0,1],
        #     5: [0,1],
        #     8: [0,1]
        #     }
        # self.cell_to_nonrow = {
        #     0: [1,2],
        #     1: [1,2],
        #     2: [1,2],
        #     3: [0,2],
        #     4: [0,2],
        #     5: [0,2],
        #     6: [0,1],
        #     7: [0,1],
        #     8: [0,1]
        #     }<|endoftext|>


<commit_before>"""
this file provides most of the functionality of the application. Reddit posts
are retrieved and converted to embeddable objects using the modules from the
providers package.
"""

import collections
import random
import string
import praw
import re

from django.conf import settings

from . import providers

reddit = praw.Reddit(user_agent='redditgag - a 9gag like interface to reddit posts by /u/niautanor', log_requests=2, cache_timeout=120)

reddit.set_oauth_app_info(
    client_id=settings.REDDIT_AUTH_CLIENT_ID,
    client_secret=settings.REDDIT_AUTH_CLIENT_SECRET,
    redirect_uri=settings.REDDIT_AUTH_REDIRECT_URI)

def get_posts(subreddit, after, show_nsfw, auth):
    """
    Requests a few (5) posts from the subreddit subreddit (frontpage if
    subreddit is None) and add the hidden attribute to nsfw posts unless
    show_nsfw is True.

    auth is a tuple consisting of a scope, an access token and a refresh token
    from the reddit API that can be used to access private subreddits / the
    personalized front page.
    """
    if auth is not None:
        creds = (auth['scope'], auth['access_token'], auth['refresh_token'])
        reddit.set_access_credentials(*creds)
    else:
        reddit.clear_authentication()

    fun = reddit.get_front_page if subreddit is None \
                                else reddit.get_subreddit(subreddit).get_hot

    submissions = list(fun(limit=5, params={'after':after}))
    last = submissions[-1].name
    # the reddit api actually returns a 'after' field to allow pagination but
    # praw won't let me access it :/

    return [ providers.get_embeddable(s, show_nsfw) for s in submissions ], last

def generate_oauth_state():
    """
    Generates a random string of 6 lower case characters to make sure we only
    finish reddit authentication requests that we started.

    See: https://github.com/reddit/reddit/wiki/OAuth2
    """
    state =  ''.join(random.choice(string.ascii_lowercase) for _ in range(6))
    return state, reddit.get_authorize_url(state, 'read', True)

def verify_auth_callback(session_state, reddit_state, code):
    assert session_state == reddit_state
    # False because we don't want to automatically bind the reddit object to
    # the retrieved credentials
    return reddit.get_access_information(code, False)
<commit_msg>Detect if the OAuth expired and report the error<commit_after>"""
this file provides most of the functionality of the application. Reddit posts
are retrieved and converted to embeddable objects using the modules from the
providers package.
"""

import collections
import random
import string
import praw
import re

from django.conf import settings
from django.urls import reverse

from praw.errors import OAuthInvalidToken

from . import providers

reddit = praw.Reddit(user_agent='redditgag - a 9gag like interface to reddit posts by /u/niautanor', log_requests=2, cache_timeout=120)

reddit.set_oauth_app_info(
    client_id=settings.REDDIT_AUTH_CLIENT_ID,
    client_secret=settings.REDDIT_AUTH_CLIENT_SECRET,
    redirect_uri=settings.REDDIT_AUTH_REDIRECT_URI)

def get_posts(subreddit, after, show_nsfw, auth):
    """
    Requests a few (5) posts from the subreddit subreddit (frontpage if
    subreddit is None) and add the hidden attribute to nsfw posts unless
    show_nsfw is True.

    auth is a tuple consisting of a scope, an access token and a refresh token
    from the reddit API that can be used to access private subreddits / the
    personalized front page.
    """
    if auth is not None:
        creds = (auth['scope'], auth['access_token'], auth['refresh_token'])
        reddit.set_access_credentials(*creds)
    else:
        reddit.clear_authentication()

    fun = reddit.get_front_page if subreddit is None \
                                else reddit.get_subreddit(subreddit).get_hot

    try:
        submissions = list(fun(limit=5, params={'after':after}))
        last = submissions[-1].name
        # the reddit api actually returns a 'after' field to allow pagination
        # but praw won't let me access it :/

        return [
            providers.get_embeddable(s, show_nsfw) for s in submissions
        ], last
    except OAuthInvalidToken:
        return [{
            'title' : 'Reddit Authentication Failure',
            'author' : 'redditgag',
            'subreddit' : 'redditgag',
            'hidden' : None,
            'kind' : 'SORRY',
            'sorrytext' : """The Reddit authentication token is invalid. This
            usually means that you have revoked this apps access in your reddit
            app preferences but haven't cleared your session cookie
            yet. If this isn't the case, I would appreciate it if you
            reported the issue on github.""" % reverse('logout'),
        }], -1

def generate_oauth_state():
    """
    Generates a random string of 6 lower case characters to make sure we only
    finish reddit authentication requests that we started.

    See: https://github.com/reddit/reddit/wiki/OAuth2
    """
    state =  ''.join(random.choice(string.ascii_lowercase) for _ in range(6))
    return state, reddit.get_authorize_url(state, 'read', True)

def verify_auth_callback(session_state, reddit_state, code):
    assert session_state == reddit_state
    # False because we don't want to automatically bind the reddit object to
    # the retrieved credentials
    return reddit.get_access_information(code, False)
<|endoftext|>


<commit_before>"""Main initialisation for extension."""

VERSION = (0, 1, 0)
__version__ = '.'.join(map(str, VERSION))

try:
    from geokey.extensions.base import register

    register(
        'geokey_webresources',
        'Web Resources',
        display_admin=True,
        superuser=False,
        version=__version__
    )
except BaseException:
    print 'Please install GeoKey first'
<commit_msg>v0.1.1<commit_after>"""Main initialisation for extension."""

VERSION = (0, 1, 1)
__version__ = '.'.join(map(str, VERSION))

try:
    from geokey.extensions.base import register

    register(
        'geokey_webresources',
        'Web Resources',
        display_admin=True,
        superuser=False,
        version=__version__
    )
except BaseException:
    print 'Please install GeoKey first'
<|endoftext|>


<commit_before>import logging
import random
import gevent
import blinker
import time

from itertools import cycle
from collections import defaultdict

from .lookupd import Lookupd
from .nsqd import Nsqd
from .backofftimer import BackoffTimer
from .states import INIT, RUNNING, BACKOFF, THROTTLED, CLOSED

from .errors import (
    NSQException,
    NSQNoConnections,
    NSQRequeueMessage
)

class Reader(object):
    def __init__(
        self,
        topic,
        channel,
        nsqd_tcp_addresses=[],
        lookupd_http_addresses=[],
        name=None,
        async=False,
        max_tries=5,
        max_in_flight=1,
        requeue_delay=0,
        lookupd_poll_interval=60,
        lookupd_poll_jitter=0.3,
        low_ready_idle_timeout=10,
        max_backoff_duration=128,
        **kwargs
    ):
        if not nsqd_tcp_addresses and not lookupd_http_addresses:
            raise ValueError('must specify at least on nsqd or lookupd')

        if isinstance(nsqd_tcp_addresses, basestring):
            self.nsqd_tcp_addresses = [nsqd_tcp_addresses]
        elif isinstance(nsqd_tcp_addresses, (list, tuple)):
            self.nsqd_tcp_addresses = set(nsqd_tcp_addresses)
        else:
            raise TypeError('nsqd_tcp_addresses must be a list, set or tuple')

        if isinstance(lookupd_http_addresses, basestring):
            lookupd_http_addresses = [lookupd_http_addresses]
        elif isinstance(lookupd_http_addresses, (list, tuple)):
            lookupd_http_addresses = list(lookupd_http_addresses)
            random.shuffle(lookupd_http_addresses)
        else:
            msg = 'lookupd_http_addresses must be a list, set or tuple'
            raise TypeError(msg)

        self.lookupds = [Lookupd(a) for a in lookupd_http_addresses]
        self.iterlookupds = cycle(self.lookupds)

        self.topic = topic
        self.channel = channel
        self.async = async
        self.max_tries = max_tries
        self.max_in_flight = max_in_flight
        self.requeue_delay = requeue_delay
        self.lookupd_poll_interval = lookupd_poll_interval
        self.lookupd_poll_jitter = lookupd_poll_jitter
        self.low_ready_idle_timeout = low_ready_idle_timeout
        self.max_backoff_duration = max_backoff_duration
        self.conn_kwargs = kwargs

        if name:
            self.name = name
        else:
            self.name = '{}.{}.{}'.format(__name__, self.topic, self.channel)

        self._need_ready_redistributed = False
        self.last_random_ready = time.time()
        self.state = INIT

        self.logger = logging.getLogger(self.name)
        self.conn_backoffs = defaultdict(self.create_backoff)
        self.backoff = self.create_backoff()

        self.on_response = blinker.Signal()
        self.on_error = blinker.Signal()
        self.on_message = blinker.Signal()
        self.on_finish = blinker.Signal()
        self.on_requeue = blinker.Signal()
        self.on_giving_up = blinker.Signal()
        self.on_exception = blinker.Signal()

        self.conns = set()
        self.pending = set()

        self.workers = []
        self.conn_workers = {}

    def start(self, block=True):
        if self.state != INIT:
            self.logger.debug('{} all ready started'.format(self.name))
            return

        self.logger.debug('starting {}...'.format(self.name))
        self.state = RUNNING
        self.query_nsqd()

        if self.lookupds:
            self.query_lookupd()
            self.workers.append(gevent.spawn(self._poll_lookupd))

        self.workers.append(gevent.spawn(self._poll_ready))

        if block:
            self.join()

    def close(self):
        if not self.is_running:
            return

        self.state = CLOSED

        for worker in self.workers:
            worker.kill()

        for conn in self.conns:
            conn.close_stream()

    def join(self, timeout=None, raise_error=False):
        gevent.joinall(self.workers, timeout, raise_error)
        gevent.joinall(self.conn_workers.values(), timeout, raise_error)

    @property
    def is_running(self):
        return self.state in (RUNNING, BACKOFF, THROTTLED)

    @property
    def is_starved(self):
        for conn in self.conns:
            if conn.in_flight >= max(conn.last_ready * 0.85, 1):
                return True
        return False

    @property
    def connection_max_in_flight(self):
        return max(1, self.max_in_flight / max(1, len(self.conns)))

    @property
    def total_ready_count(self):
        return sum(c.ready_count for c in self.conns)

    def send_ready(self, conn, count):
        if (self.total_ready_count + count) > self.max_in_flight:
            return
        conn.ready(count)

    def query_nsqd(self):
        self.logger.debug('querying nsqd...')
        for address in self.nsqd_tcp_addresses:
            address, port = address.split(':')
            conn = Nsqd(address, int(port), **self.conn_kwargs)
            self.connect_to_nsqd(conn)

    def query_lookupd(self):
        self.logger.debug('querying lookupd...')
        lookupd = self.iterlookupds.next()

        try:
            producers = lookupd.lookup(self.topic)['producers']
            self.logger.debug('found {} producers'.format(len(producers)))

        except Exception as error:
            msg = 'Failed to lookup {} on {} ({})'
            self.logger.warn(msg.format(self.topic, lookupd.address, error))
            return

        for producer in producers:
            conn = Nsqd(
                producer.get('broadcast_address') or producer['address'],
                producer['tcp_port'],
                producer['http_port'],
                **self.conn_kwargs
            )
            self.connect_to_nsqd(conn)

    def create_backoff(self):
        return BackoffTimer(max_interval=self.max_backoff_duration)

    def start_backoff(self):
        self.state = BACKOFF

        for conn in self.conns:
            conn.ready(0)

        interval = self.backoff.get_interval()
        self.logger.info('backing off for {} seconds'.format(interval))
        gevent.sleep(interval)

        self.state = THROTTLED
        if not self.conns:
            return

        conn = self.random_connection()
        self.logger.info('[{}] testing backoff state with RDY 1'.format(conn))
        self.send_ready(conn, 1)

    def complete_backoff(self):
        self.state = RUNNING
        self.logger.info('backoff complete, resuming normal operation')
        self.redistribute_ready_state()

    def _poll_lookupd(self):
        delay = self.lookupd_poll_interval * self.lookupd_poll_jitter
        gevent.sleep(random.random() * delay)

        while True:
            gevent.sleep(self.lookupd_poll_interval)
            self.query_lookupd()

    def _poll_ready(self):
        while True:
            gevent.sleep(5)
            if not self.need_ready_redistributed:
                continue
            self.redistribute_ready_state()

    @property
    def need_ready_redistributed(self):
        if self.state == BACKOFF:
            return False

        if self._need_ready_redistributed:
            return True

        if len(self.conns) > self.max_in_flight:
            return True

        if self.state == THROTTLED and len(self.conns) > 1:
            return True

    @need_ready_redistributed.setter
    def need_ready_redistributed(self, value):
        self._need_ready_redistributed = value

    def redistribute_ready_state(self):
        self.need_ready_redistributed = False

        # first set RDY 0 to all connections that have not received a message
        # within a configurable timeframe (low_ready_idle_timeout).
        for conn in self.conns:
            if conn.ready_count == 0:
                continue

            if (time.time() - conn.last_message) < self.low_ready_idle_timeout:
                continue

            msg = '[{}] idle connection, giving up RDY count'.format(conn)
            self.logger.info(msg)

            conn.ready(0)

        if self.state == THROTTLED:
            max_in_flight = 1 - self.total_ready_count
        else:
            max_in_flight = self.max_in_flight - self.total_ready_count

        if max_in_flight <= 0:
            return

        # randomly walk the list of possible connections and send RDY 1 (up to
        # our calculate "max_in_flight").  We only need to send RDY 1 because in
        # both cases described above your per connection RDY count would never
        # be higher.
        #
        # We also don't attempt to avoid the connections who previously might
        # have had RDY 1 because it would be overly complicated and not actually
        # worth it (ie. given enough redistribution rounds it doesn't matter).
        conns = list(self.conns)
        conns = random.sample(conns, min(max_in_flight, len(self.conns)))

        for conn in conns:
            self.logger.info('[{}] redistributing RDY'.format(conn))
            self.send_ready(conn, 1)

    def random_ready_conn(self, conn):
        # if all connections aren't getting RDY
        # occsionally randomize which connection gets RDY
        if len(self.conns) <= self.max_in_flight:
            return conn

        if (time.time() - self.last_random_ready) < 30:
            return conn

        self.last_random_ready = time.time()
        return random.choice([c for c in self.conns if not c.ready_count])

    def update_ready(self, conn):
        if self.state in (BACKOFF, THROTTLED):
            return

        conn = self.random_ready_conn(conn)
        if conn.ready_count < max(conn.last_ready * 0.25, 2):
            self.send_ready(conn, self.connection_max_in_flight)

    def random_connection(self):
        if not self.conns:
            return None
        return random.choice(list(self.conns))

    def publish(self, topic, message):
        conn = self.random_connection()
        if conn is None:
            raise NSQNoConnections()

        conn.publish(topic, message)

    def connect_to_nsqd(self, conn):
        if not self.is_running:
            return

        if conn in self.conns:
            self.logger.debug('[{}] already connected'.format(conn))
            return

        if conn in self.pending:
            self.logger.debug('[{}] already pending'.format(conn))
            return

        self.logger.debug('[{}] connecting...'.format(conn))

        conn.on_response.connect(self.handle_response)
        conn.on_error.connect(self.handle_error)
        conn.on_message.connect(self.handle_message)
        conn.on_finish.connect(self.handle_finish)
        conn.on_requeue.connect(self.handle_requeue)

        self.pending.add(conn)

        try:
            conn.connect()
            conn.identify()

            if conn.max_ready_count < self.max_in_flight:
                self.logger.warning(' '.join([
                    '[{}] max RDY count {} < reader max in flight {},',
                    'truncation possible'
                ]).format(conn, conn.max_ready_count, self.max_in_flight))

            conn.subscribe(self.topic, self.channel)

            # Send RDY 1 if we're not backing off and we're the first connection
            if self.state not in (BACKOFF, THROTTLED) and not self.conns:
                self.send_ready(conn, 1)

        except NSQException as error:
            msg = '[{}] connection failed ({!r})'.format(conn, error)
            self.logger.debug(msg)
            self.handle_connection_failure(conn)
            return

        finally:
            self.pending.remove(conn)

        # Check if we've closed since we started
        if not self.is_running:
            conn.close_stream()
            return

        self.logger.info('[{}] connection successful'.format(conn))
        self.handle_connection_success(conn)

    def _listen(self, conn):
        try:
            conn.listen()
        except NSQException as error:
            msg = '[{}] connection lost ({!r})'.format(conn, error)
            self.logger.warning(msg)

        self.handle_connection_failure(conn)

    def handle_connection_success(self, conn):
        self.conns.add(conn)
        self.conn_workers[conn] = gevent.spawn(self._listen, conn)

        if str(conn) not in self.nsqd_tcp_addresses:
            return

        self.conn_backoffs[conn].success()

    def handle_connection_failure(self, conn):
        self.conns.discard(conn)
        self.conn_workers.pop(conn, None)
        conn.close_stream()

        if conn.ready_count:
            self.need_ready_redistributed = True

        if str(conn) not in self.nsqd_tcp_addresses:
            return

        seconds = self.conn_backoffs[conn].failure().get_interval()
        self.logger.debug('[{}] retrying in {}s'.format(conn, seconds))
        gevent.spawn_later(seconds, self.connect_to_nsqd, conn)

    def handle_response(self, conn, response):
        self.logger.debug('[{}] response: {}'.format(conn, response))
        self.on_response.send(self, response=response)

    def handle_error(self, conn, error):
        self.logger.debug('[{}] error: {}'.format(conn, error))
        self.on_error.send(self, error=error)

    def handle_message(self, conn, message):
        self.logger.debug('[{}] got message: {}'.format(conn, message.id))

        if self.max_tries and message.attempts > self.max_tries:
            msg = "giving up on message '{}' after max tries {}"
            self.logger.warning(msg.format(message.id, self.max_tries))
            self.on_giving_up.send(self, conn, message)
            return message.finish()

        try:
            self.on_message.send(self, message=message)
            if not self.async:
                message.finish()
            return

        except NSQRequeueMessage:
            pass

        except Exception as error:
            msg = '[{}] caught exception while handling message'.format(conn)
            self.logger.exception(msg)
            self.on_exception.send(self, message=message, error=error)

        message.requeue(self.requeue_delay)

    def handle_finish(self, conn, message_id):
        self.logger.debug('[{}] finished message: {}'.format(conn, message_id))
        self.backoff.success()
        self.update_ready(conn)
        self.handle_backoff()
        self.on_finish.send(self, message_id=message_id)

    def handle_requeue(self, conn, message_id, timeout):
        msg = '[{}] requeued message: {} ({})'
        self.logger.debug(msg.format(conn, message_id, timeout))
        self.backoff.failure()
        self.update_ready(conn)
        self.handle_backoff()
        self.on_requeue.send(self, message_id=message_id, timeout=timeout)

    def handle_backoff(self):
        if self.state == BACKOFF:
            return

        if self.state == THROTTLED and self.backoff.is_reset():
            return self.complete_backoff()

        if not self.backoff.is_reset():
            return self.start_backoff()
<commit_msg>Revert "Use redistribute ready when finishing back off."<commit_after>import logging
import random
import gevent
import blinker
import time

from itertools import cycle
from collections import defaultdict

from .lookupd import Lookupd
from .nsqd import Nsqd
from .backofftimer import BackoffTimer
from .states import INIT, RUNNING, BACKOFF, THROTTLED, CLOSED

from .errors import (
    NSQException,
    NSQNoConnections,
    NSQRequeueMessage
)

class Reader(object):
    def __init__(
        self,
        topic,
        channel,
        nsqd_tcp_addresses=[],
        lookupd_http_addresses=[],
        name=None,
        async=False,
        max_tries=5,
        max_in_flight=1,
        requeue_delay=0,
        lookupd_poll_interval=60,
        lookupd_poll_jitter=0.3,
        low_ready_idle_timeout=10,
        max_backoff_duration=128,
        **kwargs
    ):
        if not nsqd_tcp_addresses and not lookupd_http_addresses:
            raise ValueError('must specify at least on nsqd or lookupd')

        if isinstance(nsqd_tcp_addresses, basestring):
            self.nsqd_tcp_addresses = [nsqd_tcp_addresses]
        elif isinstance(nsqd_tcp_addresses, (list, tuple)):
            self.nsqd_tcp_addresses = set(nsqd_tcp_addresses)
        else:
            raise TypeError('nsqd_tcp_addresses must be a list, set or tuple')

        if isinstance(lookupd_http_addresses, basestring):
            lookupd_http_addresses = [lookupd_http_addresses]
        elif isinstance(lookupd_http_addresses, (list, tuple)):
            lookupd_http_addresses = list(lookupd_http_addresses)
            random.shuffle(lookupd_http_addresses)
        else:
            msg = 'lookupd_http_addresses must be a list, set or tuple'
            raise TypeError(msg)

        self.lookupds = [Lookupd(a) for a in lookupd_http_addresses]
        self.iterlookupds = cycle(self.lookupds)

        self.topic = topic
        self.channel = channel
        self.async = async
        self.max_tries = max_tries
        self.max_in_flight = max_in_flight
        self.requeue_delay = requeue_delay
        self.lookupd_poll_interval = lookupd_poll_interval
        self.lookupd_poll_jitter = lookupd_poll_jitter
        self.low_ready_idle_timeout = low_ready_idle_timeout
        self.max_backoff_duration = max_backoff_duration
        self.conn_kwargs = kwargs

        if name:
            self.name = name
        else:
            self.name = '{}.{}.{}'.format(__name__, self.topic, self.channel)

        self._need_ready_redistributed = False
        self.last_random_ready = time.time()
        self.state = INIT

        self.logger = logging.getLogger(self.name)
        self.conn_backoffs = defaultdict(self.create_backoff)
        self.backoff = self.create_backoff()

        self.on_response = blinker.Signal()
        self.on_error = blinker.Signal()
        self.on_message = blinker.Signal()
        self.on_finish = blinker.Signal()
        self.on_requeue = blinker.Signal()
        self.on_giving_up = blinker.Signal()
        self.on_exception = blinker.Signal()

        self.conns = set()
        self.pending = set()

        self.workers = []
        self.conn_workers = {}

    def start(self, block=True):
        if self.state != INIT:
            self.logger.debug('{} all ready started'.format(self.name))
            return

        self.logger.debug('starting {}...'.format(self.name))
        self.state = RUNNING
        self.query_nsqd()

        if self.lookupds:
            self.query_lookupd()
            self.workers.append(gevent.spawn(self._poll_lookupd))

        self.workers.append(gevent.spawn(self._poll_ready))

        if block:
            self.join()

    def close(self):
        if not self.is_running:
            return

        self.state = CLOSED

        for worker in self.workers:
            worker.kill()

        for conn in self.conns:
            conn.close_stream()

    def join(self, timeout=None, raise_error=False):
        gevent.joinall(self.workers, timeout, raise_error)
        gevent.joinall(self.conn_workers.values(), timeout, raise_error)

    @property
    def is_running(self):
        return self.state in (RUNNING, BACKOFF, THROTTLED)

    @property
    def is_starved(self):
        for conn in self.conns:
            if conn.in_flight >= max(conn.last_ready * 0.85, 1):
                return True
        return False

    @property
    def connection_max_in_flight(self):
        return max(1, self.max_in_flight / max(1, len(self.conns)))

    @property
    def total_ready_count(self):
        return sum(c.ready_count for c in self.conns)

    def send_ready(self, conn, count):
        if (self.total_ready_count + count) > self.max_in_flight:
            return
        conn.ready(count)

    def query_nsqd(self):
        self.logger.debug('querying nsqd...')
        for address in self.nsqd_tcp_addresses:
            address, port = address.split(':')
            conn = Nsqd(address, int(port), **self.conn_kwargs)
            self.connect_to_nsqd(conn)

    def query_lookupd(self):
        self.logger.debug('querying lookupd...')
        lookupd = self.iterlookupds.next()

        try:
            producers = lookupd.lookup(self.topic)['producers']
            self.logger.debug('found {} producers'.format(len(producers)))

        except Exception as error:
            msg = 'Failed to lookup {} on {} ({})'
            self.logger.warn(msg.format(self.topic, lookupd.address, error))
            return

        for producer in producers:
            conn = Nsqd(
                producer.get('broadcast_address') or producer['address'],
                producer['tcp_port'],
                producer['http_port'],
                **self.conn_kwargs
            )
            self.connect_to_nsqd(conn)

    def create_backoff(self):
        return BackoffTimer(max_interval=self.max_backoff_duration)

    def start_backoff(self):
        self.state = BACKOFF

        for conn in self.conns:
            conn.ready(0)

        interval = self.backoff.get_interval()
        self.logger.info('backing off for {} seconds'.format(interval))
        gevent.sleep(interval)

        self.state = THROTTLED
        if not self.conns:
            return

        conn = self.random_connection()
        self.logger.info('[{}] testing backoff state with RDY 1'.format(conn))
        self.send_ready(conn, 1)

    def complete_backoff(self):
        self.state = RUNNING
        self.logger.info('backoff complete, resuming normal operation')

        count = self.connection_max_in_flight
        for conn in self.conns:
            self.send_ready(conn, count)

    def _poll_lookupd(self):
        delay = self.lookupd_poll_interval * self.lookupd_poll_jitter
        gevent.sleep(random.random() * delay)

        while True:
            gevent.sleep(self.lookupd_poll_interval)
            self.query_lookupd()

    def _poll_ready(self):
        while True:
            gevent.sleep(5)
            if not self.need_ready_redistributed:
                continue
            self.redistribute_ready_state()

    @property
    def need_ready_redistributed(self):
        if self.state == BACKOFF:
            return False

        if self._need_ready_redistributed:
            return True

        if len(self.conns) > self.max_in_flight:
            return True

        if self.state == THROTTLED and len(self.conns) > 1:
            return True

    @need_ready_redistributed.setter
    def need_ready_redistributed(self, value):
        self._need_ready_redistributed = value

    def redistribute_ready_state(self):
        self.need_ready_redistributed = False

        # first set RDY 0 to all connections that have not received a message
        # within a configurable timeframe (low_ready_idle_timeout).
        for conn in self.conns:
            if conn.ready_count == 0:
                continue

            if (time.time() - conn.last_message) < self.low_ready_idle_timeout:
                continue

            msg = '[{}] idle connection, giving up RDY count'.format(conn)
            self.logger.info(msg)

            conn.ready(0)

        if self.state == THROTTLED:
            max_in_flight = 1 - self.total_ready_count
        else:
            max_in_flight = self.max_in_flight - self.total_ready_count

        if max_in_flight <= 0:
            return

        # randomly walk the list of possible connections and send RDY 1 (up to
        # our calculate "max_in_flight").  We only need to send RDY 1 because in
        # both cases described above your per connection RDY count would never
        # be higher.
        #
        # We also don't attempt to avoid the connections who previously might
        # have had RDY 1 because it would be overly complicated and not actually
        # worth it (ie. given enough redistribution rounds it doesn't matter).
        conns = list(self.conns)
        conns = random.sample(conns, min(max_in_flight, len(self.conns)))

        for conn in conns:
            self.logger.info('[{}] redistributing RDY'.format(conn))
            self.send_ready(conn, 1)

    def random_ready_conn(self, conn):
        # if all connections aren't getting RDY
        # occsionally randomize which connection gets RDY
        if len(self.conns) <= self.max_in_flight:
            return conn

        if (time.time() - self.last_random_ready) < 30:
            return conn

        self.last_random_ready = time.time()
        return random.choice([c for c in self.conns if not c.ready_count])

    def update_ready(self, conn):
        if self.state in (BACKOFF, THROTTLED):
            return

        conn = self.random_ready_conn(conn)
        if conn.ready_count < max(conn.last_ready * 0.25, 2):
            self.send_ready(conn, self.connection_max_in_flight)

    def random_connection(self):
        if not self.conns:
            return None
        return random.choice(list(self.conns))

    def publish(self, topic, message):
        conn = self.random_connection()
        if conn is None:
            raise NSQNoConnections()

        conn.publish(topic, message)

    def connect_to_nsqd(self, conn):
        if not self.is_running:
            return

        if conn in self.conns:
            self.logger.debug('[{}] already connected'.format(conn))
            return

        if conn in self.pending:
            self.logger.debug('[{}] already pending'.format(conn))
            return

        self.logger.debug('[{}] connecting...'.format(conn))

        conn.on_response.connect(self.handle_response)
        conn.on_error.connect(self.handle_error)
        conn.on_message.connect(self.handle_message)
        conn.on_finish.connect(self.handle_finish)
        conn.on_requeue.connect(self.handle_requeue)

        self.pending.add(conn)

        try:
            conn.connect()
            conn.identify()

            if conn.max_ready_count < self.max_in_flight:
                self.logger.warning(' '.join([
                    '[{}] max RDY count {} < reader max in flight {},',
                    'truncation possible'
                ]).format(conn, conn.max_ready_count, self.max_in_flight))

            conn.subscribe(self.topic, self.channel)

            # Send RDY 1 if we're not backing off and we're the first connection
            if self.state not in (BACKOFF, THROTTLED) and not self.conns:
                self.send_ready(conn, 1)

        except NSQException as error:
            msg = '[{}] connection failed ({!r})'.format(conn, error)
            self.logger.debug(msg)
            self.handle_connection_failure(conn)
            return

        finally:
            self.pending.remove(conn)

        # Check if we've closed since we started
        if not self.is_running:
            conn.close_stream()
            return

        self.logger.info('[{}] connection successful'.format(conn))
        self.handle_connection_success(conn)

    def _listen(self, conn):
        try:
            conn.listen()
        except NSQException as error:
            msg = '[{}] connection lost ({!r})'.format(conn, error)
            self.logger.warning(msg)

        self.handle_connection_failure(conn)

    def handle_connection_success(self, conn):
        self.conns.add(conn)
        self.conn_workers[conn] = gevent.spawn(self._listen, conn)

        if str(conn) not in self.nsqd_tcp_addresses:
            return

        self.conn_backoffs[conn].success()

    def handle_connection_failure(self, conn):
        self.conns.discard(conn)
        self.conn_workers.pop(conn, None)
        conn.close_stream()

        if conn.ready_count:
            self.need_ready_redistributed = True

        if str(conn) not in self.nsqd_tcp_addresses:
            return

        seconds = self.conn_backoffs[conn].failure().get_interval()
        self.logger.debug('[{}] retrying in {}s'.format(conn, seconds))
        gevent.spawn_later(seconds, self.connect_to_nsqd, conn)

    def handle_response(self, conn, response):
        self.logger.debug('[{}] response: {}'.format(conn, response))
        self.on_response.send(self, response=response)

    def handle_error(self, conn, error):
        self.logger.debug('[{}] error: {}'.format(conn, error))
        self.on_error.send(self, error=error)

    def handle_message(self, conn, message):
        self.logger.debug('[{}] got message: {}'.format(conn, message.id))

        if self.max_tries and message.attempts > self.max_tries:
            msg = "giving up on message '{}' after max tries {}"
            self.logger.warning(msg.format(message.id, self.max_tries))
            self.on_giving_up.send(self, conn, message)
            return message.finish()

        try:
            self.on_message.send(self, message=message)
            if not self.async:
                message.finish()
            return

        except NSQRequeueMessage:
            pass

        except Exception as error:
            msg = '[{}] caught exception while handling message'.format(conn)
            self.logger.exception(msg)
            self.on_exception.send(self, message=message, error=error)

        message.requeue(self.requeue_delay)

    def handle_finish(self, conn, message_id):
        self.logger.debug('[{}] finished message: {}'.format(conn, message_id))
        self.backoff.success()
        self.update_ready(conn)
        self.handle_backoff()
        self.on_finish.send(self, message_id=message_id)

    def handle_requeue(self, conn, message_id, timeout):
        msg = '[{}] requeued message: {} ({})'
        self.logger.debug(msg.format(conn, message_id, timeout))
        self.backoff.failure()
        self.update_ready(conn)
        self.handle_backoff()
        self.on_requeue.send(self, message_id=message_id, timeout=timeout)

    def handle_backoff(self):
        if self.state == BACKOFF:
            return

        if self.state == THROTTLED and self.backoff.is_reset():
            return self.complete_backoff()

        if not self.backoff.is_reset():
            return self.start_backoff()
<|endoftext|>


<commit_before>from nose.tools import with_setup, raises

def _make_class(super):
    # don't use type(), it's hard to read
    from gevent_tools.util import defaultproperty
    class Foo(super):
        bar = defaultproperty(set)

    return Foo

def test_defaultproperty():
    Foo = _make_class(object)
    foo = Foo()
    assert isinstance(foo.bar, set), "foo.bar is a set"

def test_default_property_mutable():
    Foo = _make_class(object)
    foo = Foo()
    fuz = Foo()
    s = foo.bar
    s2 = fuz.bar

    s.add(3)
    s2.add(5)

    assert foo.bar == set([3]), "can mutate reference"
    assert fuz.bar == set([5]), "but not all references"

def test_default_is_iterable():
    Foo = _make_class(object)
    foo = Foo()
    assert [x for x in foo.bar] == [], "can iterate over self reference"

def test_add_as_first_dereference_works():
    Foo = _make_class(object)
    foo = Foo()
    foo.bar.add(3)
    assert foo.bar == set([3]), "previous line throws exceptino on fail"

def test_default_subclasses_work():
    Foo = _make_class(object)
    class Bar(Foo):
        pass
    bar = Bar()
    bar.bar.add(5)

    assert bar.bar == set([5]), "defaultproperties can find parent classes."
<commit_msg>better failure message in tests.<commit_after>from nose.tools import with_setup, raises

def _make_class(super):
    # don't use type(), it's hard to read
    from gevent_tools.util import defaultproperty
    class Foo(super):
        bar = defaultproperty(set)

    return Foo

def test_defaultproperty():
    Foo = _make_class(object)
    foo = Foo()
    assert isinstance(foo.bar, set), "foo.bar is a set"

def test_default_property_mutable():
    Foo = _make_class(object)
    foo = Foo()
    fuz = Foo()
    s = foo.bar
    s2 = fuz.bar

    s.add(3)
    s2.add(5)

    assert foo.bar == set([3]), "can mutate reference"
    assert fuz.bar == set([5]), "but not all references"

def test_default_is_iterable():
    Foo = _make_class(object)
    foo = Foo()
    assert [x for x in foo.bar] == [], "can iterate over iterable default"

def test_add_as_first_dereference_works():
    Foo = _make_class(object)
    foo = Foo()
    foo.bar.add(3)
    assert foo.bar == set([3]), "previous line throws exceptino on fail"

def test_default_subclasses_work():
    Foo = _make_class(object)
    class Bar(Foo):
        pass
    bar = Bar()
    bar.bar.add(5)

    assert bar.bar == set([5]), "defaultproperties can find parent classes."
<|endoftext|>


<commit_before>from unittest import TestCase, main
from svtools.vcf.file import Vcf
from svtools.vcf.variant import Variant
from svtools.vcf.genotype import Genotype
import sys

class TestVariant(TestCase):
    def setUp(self):
        header_lines = [
                '##fileformat=VCFv4.2',
                '##fileDate=20151202',
                '##INFO=<ID=SVTYPE,Number=1,Type=String,Description="Type of structural variant">',
                '##INFO=<ID=STRANDS,Number=.,Type=String,Description="Strand orientation of the adjacency in BEDPE format (DEL:+-, DUP:-+, INV:++/--)">',
                '##INFO=<ID=IMAFLAG,Number=.,Type=Flag,Description="Test Flag code">',
                '##FORMAT=<ID=GT,Number=1,Type=String,Description="Genotype">',
                '##FORMAT=<ID=SU,Number=1,Type=Integer,Description="Number of pieces of evidence supporting the variant">',
                '##FORMAT=<ID=INACTIVE,Number=1,Type=Integer,Description="A format not in use">',
                '#CHROM	POS	ID	REF	ALT	QUAL	FILTER	INFO	FORMAT	NA12878	NA0001' ]
        self.vcf = Vcf()
        self.vcf.add_header(header_lines)
        self.variant_line = '1	820915	5838_1	N	]GL000232.1:20940]N	0.00	.	SVTYPE=BND;STRANDS=-+:9;IMAFLAG	GT:SU	0/0:9	1/1:15'
        self.variant = Variant(self.variant_line.split('\t'), self.vcf)

    def test_parse_genotypes(self):
        genotype_field_strings = ['0/1:20', '0/0:15']
        parsed_dict = self.variant._parse_genotypes(genotype_field_strings)

        na12878_gt = Genotype(self.variant, genotype_field_strings[0].split(':'))
        na0001_gt = Genotype(self.variant, genotype_field_strings[1].split(':'))
        expected_genotype_dict = { 'NA12878': na12878_gt, 'NA0001': na0001_gt }

        self.assertEqual(parsed_dict, expected_genotype_dict)

    def test_set_info(self):
        self.variant.set_info('SVTYPE', 'INV')
        self.assertEqual(self.variant.info['SVTYPE'], 'INV')
        self.variant.set_info('IMAFLAG', False)
        self.assertEqual(self.variant.info['IMAFLAG'], False)
        with self.assertRaises(SystemExit) as cm:
            self.variant.set_info('SUPER', True)

    def test_get_info(self):
        self.assertEqual(self.variant.get_info('IMAFLAG'), True)
        self.assertEqual(self.variant.get_info('SVTYPE'), 'BND')
        with self.assertRaises(KeyError) as cm:
            self.variant.get_info('CALI')

    def test_get_info_string(self):
        self.assertEqual(self.variant.get_info_string(), 'SVTYPE=BND;STRANDS=-+:9;IMAFLAG')
        self.variant.set_info('IMAFLAG', False)
        self.assertEqual(self.variant.get_info_string(), 'SVTYPE=BND;STRANDS=-+:9')

    def test_get_format_string(self):
        self.assertEqual(self.variant.get_format_string(), 'GT:SU')

    def test_get_format_string_caching(self):
        header_lines = [
                '##fileformat=VCFv4.2',
                '##fileDate=20151202',
                '##INFO=<ID=SVTYPE,Number=1,Type=String,Description="Type of structural variant">',
                '##INFO=<ID=STRANDS,Number=.,Type=String,Description="Strand orientation of the adjacency in BEDPE format (DEL:+-, DUP:-+, INV:++/--)">',
                '##INFO=<ID=IMAFLAG,Number=.,Type=Flag,Description="Test Flag code">',
                '##FORMAT=<ID=GT,Number=1,Type=String,Description="Genotype">',
                '##FORMAT=<ID=SU,Number=1,Type=Integer,Description="Number of pieces of evidence supporting the variant">',
                '##FORMAT=<ID=AS,Number=1,Type=Integer,Description="Number of pieces of evidence supporting the variant">',
                '##FORMAT=<ID=INACTIVE,Number=1,Type=Integer,Description="A format not in use">',
                '#CHROM	POS	ID	REF	ALT	QUAL	FILTER	INFO	FORMAT	NA12878' ]
        vcf = Vcf()
        vcf.add_header(header_lines)
        variant_line = '1	820915	5838_1	N	]GL000232.1:20940]N	0.00	.	SVTYPE=BND;STRANDS=-+:9;IMAFLAG	GT:AS:SU	0/0:1:9'
        variant = Variant(variant_line.split('\t'), vcf)
        self.assertEqual(variant.get_format_string(), 'GT:AS:SU')

        gts = variant.genotypes()
        self.assertEqual(variant.get_format_string(), 'GT:SU:AS')

        self.assertEqual(variant.get_format_string(True), 'GT:AS:SU')

    def test_get_gt_string(self):
        self.assertEqual(self.variant.get_gt_string(), '0/0:9	1/1:15')

    def test_genotype(self):
        self.assertEqual(self.variant.genotype('NA12878').get_gt_string(), '0/0:9')

    def test_set_genotype(self):
        new_genotype = Genotype(self.variant, ['0/1', '9'])
        self.variant.set_genotype('NA12878', new_genotype)
        self.assertEqual(self.variant.genotype('NA12878').get_gt_string(), '0/1:9')

    def test_genotypes(self):
        self.assertEqual([ x.get_gt_string() for x in self.variant.genotypes() ], ['0/0:9', '1/1:15'])

    def test_var_string(self):
        self.assertEqual(self.variant.get_var_string(), self.variant_line)
        self.variant.genotype('NA12878').set_format('GT', './.')
        self.assertEqual(self.variant.get_var_string(use_cached_gt_string=True), self.variant_line)
        self.assertNotEqual(self.variant.get_var_string(), self.variant_line)

    def test_add_genotype(self):
        header_lines = [
                '##fileformat=VCFv4.2',
                '##fileDate=20151202',
                '##INFO=<ID=SVTYPE,Number=1,Type=String,Description="Type of structural variant">',
                '##INFO=<ID=STRANDS,Number=.,Type=String,Description="Strand orientation of the adjacency in BEDPE format (DEL:+-, DUP:-+, INV:++/--)">',
                '##INFO=<ID=IMAFLAG,Number=.,Type=Flag,Description="Test Flag code">',
                '##FORMAT=<ID=SU,Number=1,Type=Integer,Description="Number of pieces of evidence supporting the variant">',
                '##FORMAT=<ID=INACTIVE,Number=1,Type=Integer,Description="A format not in use">',
                '#CHROM	POS	ID	REF	ALT	QUAL	FILTER	INFO	FORMAT	NA12878' ]
        vcf = Vcf()
        vcf.add_header(header_lines)
        variant_line = '1	820915	5838_1	N	]GL000232.1:20940]N	0.00	.	SVTYPE=BND;STRANDS=-+:9;IMAFLAG	SU	9'
        variant = Variant(variant_line.split('\t'), vcf)
        self.assertEqual(variant.get_gt_string(), './.:9')

if __name__ == "__main__":
    main()

<commit_msg>test for caching issues...<commit_after>from unittest import TestCase, main
from svtools.vcf.file import Vcf
from svtools.vcf.variant import Variant
from svtools.vcf.genotype import Genotype
import sys

class TestVariant(TestCase):
    def setUp(self):
        header_lines = [
                '##fileformat=VCFv4.2',
                '##fileDate=20151202',
                '##INFO=<ID=SVTYPE,Number=1,Type=String,Description="Type of structural variant">',
                '##INFO=<ID=STRANDS,Number=.,Type=String,Description="Strand orientation of the adjacency in BEDPE format (DEL:+-, DUP:-+, INV:++/--)">',
                '##INFO=<ID=IMAFLAG,Number=.,Type=Flag,Description="Test Flag code">',
                '##FORMAT=<ID=GT,Number=1,Type=String,Description="Genotype">',
                '##FORMAT=<ID=SU,Number=1,Type=Integer,Description="Number of pieces of evidence supporting the variant">',
                '##FORMAT=<ID=INACTIVE,Number=1,Type=Integer,Description="A format not in use">',
                '#CHROM	POS	ID	REF	ALT	QUAL	FILTER	INFO	FORMAT	NA12878	NA0001' ]
        self.vcf = Vcf()
        self.vcf.add_header(header_lines)
        self.variant_line = '1	820915	5838_1	N	]GL000232.1:20940]N	0.00	.	SVTYPE=BND;STRANDS=-+:9;IMAFLAG	GT:SU	0/0:9	1/1:15'
        self.variant = Variant(self.variant_line.split('\t'), self.vcf)

    def test_parse_genotypes(self):
        genotype_field_strings = ['0/1:20', '0/0:15']
        parsed_dict = self.variant._parse_genotypes(genotype_field_strings)

        na12878_gt = Genotype(self.variant, genotype_field_strings[0].split(':'))
        na0001_gt = Genotype(self.variant, genotype_field_strings[1].split(':'))
        expected_genotype_dict = { 'NA12878': na12878_gt, 'NA0001': na0001_gt }

        self.assertEqual(parsed_dict, expected_genotype_dict)

    def test_set_info(self):
        self.variant.set_info('SVTYPE', 'INV')
        self.assertEqual(self.variant.info['SVTYPE'], 'INV')
        self.variant.set_info('IMAFLAG', False)
        self.assertEqual(self.variant.info['IMAFLAG'], False)
        with self.assertRaises(SystemExit) as cm:
            self.variant.set_info('SUPER', True)

    def test_get_info(self):
        self.assertEqual(self.variant.get_info('IMAFLAG'), True)
        self.assertEqual(self.variant.get_info('SVTYPE'), 'BND')
        with self.assertRaises(KeyError) as cm:
            self.variant.get_info('CALI')

    def test_get_info_string(self):
        self.assertEqual(self.variant.get_info_string(), 'SVTYPE=BND;STRANDS=-+:9;IMAFLAG')
        self.variant.set_info('IMAFLAG', False)
        self.assertEqual(self.variant.get_info_string(), 'SVTYPE=BND;STRANDS=-+:9')

    def test_get_format_string(self):
        self.assertEqual(self.variant.get_format_string(), 'GT:SU')

    def test_get_format_string_caching(self):
        header_lines = [
                '##fileformat=VCFv4.2',
                '##fileDate=20151202',
                '##INFO=<ID=SVTYPE,Number=1,Type=String,Description="Type of structural variant">',
                '##INFO=<ID=STRANDS,Number=.,Type=String,Description="Strand orientation of the adjacency in BEDPE format (DEL:+-, DUP:-+, INV:++/--)">',
                '##INFO=<ID=IMAFLAG,Number=.,Type=Flag,Description="Test Flag code">',
                '##FORMAT=<ID=GT,Number=1,Type=String,Description="Genotype">',
                '##FORMAT=<ID=SU,Number=1,Type=Integer,Description="Number of pieces of evidence supporting the variant">',
                '##FORMAT=<ID=AS,Number=1,Type=Integer,Description="Number of pieces of evidence supporting the variant">',
                '##FORMAT=<ID=INACTIVE,Number=1,Type=Integer,Description="A format not in use">',
                '#CHROM	POS	ID	REF	ALT	QUAL	FILTER	INFO	FORMAT	NA12878' ]
        vcf = Vcf()
        vcf.add_header(header_lines)
        variant_line = '1	820915	5838_1	N	]GL000232.1:20940]N	0.00	.	SVTYPE=BND;STRANDS=-+:9;IMAFLAG	GT:AS:SU	0/0:1:9'
        variant = Variant(variant_line.split('\t'), vcf)
        self.assertEqual(variant.get_format_string(), 'GT:AS:SU')

        gts = variant.genotypes()
        self.assertEqual(variant.get_format_string(), 'GT:SU:AS')

        self.assertEqual(variant.get_format_string(True), 'GT:AS:SU')

    def test_get_gt_string(self):
        self.assertEqual(self.variant.get_gt_string(), '0/0:9	1/1:15')

    def test_genotype(self):
        self.assertEqual(self.variant.genotype('NA12878').get_gt_string(), '0/0:9')

    def test_set_genotype(self):
        new_genotype = Genotype(self.variant, ['0/1', '9'])
        self.variant.set_genotype('NA12878', new_genotype)
        self.assertEqual(self.variant.genotype('NA12878').get_gt_string(), '0/1:9')

    def test_genotypes(self):
        self.assertEqual([ x.get_gt_string() for x in self.variant.genotypes() ], ['0/0:9', '1/1:15'])

    def test_var_string(self):
        self.assertEqual(self.variant.get_var_string(), self.variant_line)
        self.variant.genotype('NA12878').set_format('GT', './.')
        self.assertEqual(self.variant.get_var_string(use_cached_gt_string=True), self.variant_line)
        self.assertNotEqual(self.variant.get_var_string(), self.variant_line)

    def test_var_string_format_caching(self):
        header_lines = [
                '##fileformat=VCFv4.2',
                '##fileDate=20151202',
                '##INFO=<ID=SVTYPE,Number=1,Type=String,Description="Type of structural variant">',
                '##INFO=<ID=STRANDS,Number=.,Type=String,Description="Strand orientation of the adjacency in BEDPE format (DEL:+-, DUP:-+, INV:++/--)">',
                '##INFO=<ID=IMAFLAG,Number=.,Type=Flag,Description="Test Flag code">',
                '##FORMAT=<ID=GT,Number=1,Type=String,Description="Genotype">',
                '##FORMAT=<ID=SU,Number=1,Type=Integer,Description="Number of pieces of evidence supporting the variant">',
                '##FORMAT=<ID=AS,Number=1,Type=Integer,Description="Number of pieces of evidence supporting the variant">',
                '##FORMAT=<ID=INACTIVE,Number=1,Type=Integer,Description="A format not in use">',
                '#CHROM	POS	ID	REF	ALT	QUAL	FILTER	INFO	FORMAT	NA12878' ]
        vcf = Vcf()
        vcf.add_header(header_lines)
        variant_line = '1	820915	5838_1	N	]GL000232.1:20940]N	0.00	.	SVTYPE=BND;STRANDS=-+:9;IMAFLAG	GT:AS:SU	0/0:1:9'
        uncached_line = '1	820915	5838_1	N	]GL000232.1:20940]N	0.00	.	SVTYPE=BND;STRANDS=-+:9;IMAFLAG	GT:SU:AS	0/0:9:1'
        variant = Variant(variant_line.split('\t'), vcf)
        gt = variant.genotypes() #force parsing
        self.assertEqual(variant.get_var_string(), uncached_line)
        self.assertEqual(variant.get_var_string(use_cached_gt_string=True), variant_line)

    def test_add_genotype(self):
        header_lines = [
                '##fileformat=VCFv4.2',
                '##fileDate=20151202',
                '##INFO=<ID=SVTYPE,Number=1,Type=String,Description="Type of structural variant">',
                '##INFO=<ID=STRANDS,Number=.,Type=String,Description="Strand orientation of the adjacency in BEDPE format (DEL:+-, DUP:-+, INV:++/--)">',
                '##INFO=<ID=IMAFLAG,Number=.,Type=Flag,Description="Test Flag code">',
                '##FORMAT=<ID=SU,Number=1,Type=Integer,Description="Number of pieces of evidence supporting the variant">',
                '##FORMAT=<ID=INACTIVE,Number=1,Type=Integer,Description="A format not in use">',
                '#CHROM	POS	ID	REF	ALT	QUAL	FILTER	INFO	FORMAT	NA12878' ]
        vcf = Vcf()
        vcf.add_header(header_lines)
        variant_line = '1	820915	5838_1	N	]GL000232.1:20940]N	0.00	.	SVTYPE=BND;STRANDS=-+:9;IMAFLAG	SU	9'
        variant = Variant(variant_line.split('\t'), vcf)
        self.assertEqual(variant.get_gt_string(), './.:9')

if __name__ == "__main__":
    main()

<|endoftext|>