<commit_before>import json
import pytest

from virtool.dispatcher import Dispatcher

@pytest.fixture(scope="function")
def empty_dispatcher(called_tester):
    return Dispatcher(called_tester())

@pytest.fixture(scope="function")
def dispatcher(empty_dispatcher, mock_interface, mock_settings):
    empty_dispatcher.add_interface("test", mock_interface, mock_settings)
    return empty_dispatcher

@pytest.fixture(params=["open", "on_message", "on_close", "write_message"])
def bad_connection(request, mock_connection):
    def create(settings):
        socket = mock_connection(settings)
        setattr(socket, request.param, None)

        return socket

    return create

class TestInit:

    def test_ping(self, dispatcher):

        callback = dispatcher.add_periodic_callback

        assert callback.was_called
        assert callback.with_args[0].__name__ == "ping"
        assert callback.with_args[1] == 10000

class TestInterfaces:

    def test_add(self, recwarn, empty_dispatcher, mock_interface, mock_settings):
        # Add and interface that has a method decorated with exposed_method.
        empty_dispatcher.add_interface("test", mock_interface, mock_settings)

        # Make sure the 'no exposed methods' warning was not raised.
        assert len(recwarn) == 0

    def test_add_empty(self, recwarn, empty_dispatcher, empty_interface, mock_settings):
        # Make sure warning is raised if no exposed methods are found in interface object.
        with pytest.warns(Warning):
            empty_dispatcher.add_interface("test", empty_interface, mock_settings)

        # Double check that ONE warning was raised.
        assert len(recwarn) == 1

    def test_add_existing(self, recwarn, dispatcher, mock_interface, mock_settings):
        with pytest.raises(ValueError) as inst:
            dispatcher.add_interface("test", mock_interface, mock_settings)

    def test_collection(self, recwarn, dispatcher, mock_interface, mock_settings):
        assert "test" not in dispatcher.collections

        dispatcher.add_interface("coll", mock_interface, mock_settings, is_collection=True)

        assert "coll" in dispatcher.collections

class TestConnections:

    def test_add_connection(self, dispatcher, mock_connection):

        conn = mock_connection(dispatcher)

        assert dispatcher.connections == []

        conn.open()

        assert dispatcher.connections == [conn]

    def test_add_bad_connection(self, dispatcher, bad_connection):

        conn = bad_connection(dispatcher)

        for method_name in ["open", "on_message", "on_close", "write_message"]:
            if getattr(conn, method_name) is None:
                with pytest.raises(AttributeError):
                    dispatcher.add_connection(conn)

    def test_remove_connection(self, dispatcher, mock_connection):

        conn = mock_connection(dispatcher)

        conn.open()

        assert dispatcher.connections == [conn]

        dispatcher.remove_connection(conn)

        assert dispatcher.connections == []

    def test_close_connection(self, dispatcher, mock_connection):

        conn = mock_connection(dispatcher)

        conn.open()

        assert dispatcher.connections == [conn]

        conn.close()

        assert dispatcher.connections == []

class TestDispatch:

    def test_dispatch_authorized(self, dispatcher, mock_connection):
        """
        Test the both of two AUTHORIZED connections can have messages dispatched through them using their
        ``write_message`` methods.

        """
        conn1 = mock_connection(dispatcher)
        conn2 = mock_connection(dispatcher)

        conn1.open()
        conn2.open()

        message = {
            "operation": "test",
            "collection_name": "test",
            "data": dict(message="This is a test")
        }

        dispatcher.dispatch(message)

        assert conn1.messages[0] == message
        assert conn2.messages[0] == message

    def test_dispatch_unauthorized(self, dispatcher, mock_connection):
        """
        Test the neither of two unauthorized connections have their ``write_message`` methods called during a dispatch.

        """
        conn1 = mock_connection(dispatcher, authorized=False)
        conn2 = mock_connection(dispatcher, authorized=False)

        conn1.open()
        conn2.open()

        dispatcher.dispatch({
            "operation": "test",
            "collection_name": "test",
            "data": dict(message="This is a test")
        })

        assert conn1.messages == []
        assert conn2.messages == []

    def test_dispatch_either(self, dispatcher, mock_connection):
        """
        Test the only the authorized connection has its ``write_message`` method called when an one authorized and one
        unauthorized connection are managed by the dispatcher.

        """
        conn1 = mock_connection(dispatcher, authorized=False)
        conn2 = mock_connection(dispatcher, authorized=True)

        conn1.open()
        conn2.open()

        message = {
            "operation": "test",
            "collection_name": "test",
            "data": dict(message="Only connection 2 should see this")
        }

        dispatcher.dispatch(message)

        assert conn1.messages == []
        assert conn2.messages == [message]

    def test_dispatch_specific(self, dispatcher, mock_connection):
        """
        Test that only the connection passed in the keyword argument ``connections`` has its ``write_message``
        method called when a dispatch occurs.

        """
        conn1 = mock_connection(dispatcher)
        conn2 = mock_connection(dispatcher)

        conn1.open()
        conn2.open()

        message = {
            "operation": "test",
            "collection_name": "test",
            "data": dict(message="This is a test")
        }

        dispatcher.dispatch(message, connections=[conn2])

        assert conn1.messages == []
        assert conn2.messages[0] == message

class TestHandle:

    @pytest.mark.gen_test
    def test_unknown_interface(self, caplog, dispatcher, mock_connection):
        """
        Test that a warning is logged when an unknown dispatcher interface is requested.

        """
        conn = mock_connection(permissions="all")

        code = yield dispatcher.handle(json.dumps({
            "tid": 9029401982,
            "interface": "foo",
            "method": "test_exposed_method",
            "data": {
                "value": "foobar"
            }
        }), conn)

        assert not code

        last_message = list(caplog.records)[-1].message

        assert last_message == "User 'test' requested unknown interface foo"

    @pytest.mark.gen_test
    def test_unknown_method(self, caplog, dispatcher, mock_connection):
        """
        Test that a warning is logged when an unknown interface method is requested.

        """
        conn = mock_connection(dispatcher)

        code = yield dispatcher.handle(json.dumps({
            "tid": 9029401982,
            "interface": "test",
            "method": "non_existent_method",
            "data": {
                "value": "foobar"
            }
        }), conn)

        assert not code

        last_message = list(caplog.records)[-1].message

        assert last_message == "User 'test' requested unknown interface method test.non_existent_method"

    @pytest.mark.gen_test
    def test_unexposed_method(self, caplog, dispatcher, mock_connection):
        """
        Test that a warning is logged when an unknown exposed method is requested.

        """
        conn = mock_connection(dispatcher)

        code = yield dispatcher.handle(json.dumps({
            "tid": 9029401983,
            "interface": "test",
            "method": "test_unexposed_method",
            "data": {
                "value": "foobar"
            }
        }), conn)

        assert not code

        last_message = list(caplog.records)[-1].message

        assert last_message == "User 'test' attempted to call unexposed method test.test_unexposed_method"

    @pytest.mark.gen_test
    def test_unauthorized(self, caplog, dispatcher, mock_connection):
        """
        Test that a warning is logged when an protected exposed method is requested by an unauthorized connection.

        """
        conn = mock_connection(dispatcher, authorized=False)

        code = yield dispatcher.handle(json.dumps({
            "tid": 9029401983,
            "interface": "test",
            "method": "test_exposed_method",
            "data": {
                "value": "foobar"
            }
        }), conn)

        assert not code

        last_message = list(caplog.records)[-1].message

        assert last_message == "Unauthorized connection at 8.8.8.8 attempted to call" \
                               " protected method test.test_exposed_method"

    @pytest.mark.gen_test
    def test_unprotected(self, dispatcher, mock_connection):
        """
        Test that unauthorized connections can successfully call unprotected methods.

        """
        conn = mock_connection(dispatcher, authorized=False)

        code = yield dispatcher.handle(json.dumps({
            "tid": 9029401983,
            "interface": "test",
            "method": "test_unprotected_method",
            "data": {
                "value": "foobar"
            }
        }), conn)

        assert code

        message_interfaces = [message["collection_name"] for message in conn.messages]

        assert "transaction" in message_interfaces
        assert "test" not in message_interfaces

    @pytest.mark.gen_test
    def test_authorized(self, dispatcher, mock_connection):
        """
        Test that unauthorized connections can successfully call unprotected methods.

        """
        conn = mock_connection(dispatcher, authorized=True)

        code = yield dispatcher.handle(json.dumps({
            "tid": 9029401983,
            "interface": "test",
            "method": "test_exposed_method",
            "data": {
                "value": "foobar"
            }
        }), conn)

        assert code

        message_interfaces = [message["collection_name"] for message in conn.messages]

        assert "transaction" in message_interfaces
        assert "test" not in message_interfaces

    @pytest.mark.gen_test
    def test_not_permitted(self, caplog, dispatcher, mock_connection):
        """
        Test that unauthorized connections can successfully call unprotected methods.

        """
        conn = mock_connection(dispatcher, permissions=[], authorized=True)

        code = yield dispatcher.handle(json.dumps({
            "tid": 9029401983,
            "interface": "test",
            "method": "test_permissions_exposed_method",
            "data": {
                "value": "foobar"
            }
        }), conn)

        last_message = list(caplog.records)[-1].message

        assert not code

        assert last_message == "User test attempted to call exposed method test.test_permissions_exposed_method" \
                               " without the required permissions"

    @pytest.mark.gen_test
    def test_is_permitted(self, dispatcher, mock_connection):
        """
        Test that unauthorized connections can successfully call unprotected methods.

        """
        conn = mock_connection(dispatcher, permissions=["modify_options"], authorized=True)

        code = yield dispatcher.handle(json.dumps({
            "tid": 9029401983,
            "interface": "test",
            "method": "test_permissions_exposed_method",
            "data": {
                "value": "foobar"
            }
        }), conn)

        assert code

    @pytest.mark.gen_test
    def test_no_transaction(self, dispatcher, mock_connection):
        """
        Test that exposed methods with no transaction argument raise a TypeError.

        """
        conn = mock_connection(dispatcher, authorized=True)

        with pytest.raises(TypeError) as inst:
            yield dispatcher.handle(json.dumps({
                "tid": 9029401983,
                "interface": "test",
                "method": "test_no_transaction_method",
                "data": {
                    "value": "foobar"
                }
            }), conn)
<commit_msg>Remove unneccessary recwarn funcargs<commit_after>import json
import pytest

from virtool.dispatcher import Dispatcher

@pytest.fixture(scope="function")
def empty_dispatcher(called_tester):
    return Dispatcher(called_tester())

@pytest.fixture(scope="function")
def dispatcher(empty_dispatcher, mock_interface, mock_settings):
    empty_dispatcher.add_interface("test", mock_interface, mock_settings)
    return empty_dispatcher

@pytest.fixture(params=["open", "on_message", "on_close", "write_message"])
def bad_connection(request, mock_connection):
    def create(settings):
        socket = mock_connection(settings)
        setattr(socket, request.param, None)

        return socket

    return create

class TestInit:

    def test_ping(self, dispatcher):

        callback = dispatcher.add_periodic_callback

        assert callback.was_called
        assert callback.with_args[0].__name__ == "ping"
        assert callback.with_args[1] == 10000

class TestInterfaces:

    def test_add(self, recwarn, empty_dispatcher, mock_interface, mock_settings):
        # Add and interface that has a method decorated with exposed_method.
        empty_dispatcher.add_interface("test", mock_interface, mock_settings)

        # Make sure the 'no exposed methods' warning was not raised.
        assert len(recwarn) == 0

    def test_add_empty(self, recwarn, empty_dispatcher, empty_interface, mock_settings):
        # Make sure warning is raised if no exposed methods are found in interface object.
        with pytest.warns(Warning):
            empty_dispatcher.add_interface("test", empty_interface, mock_settings)

        # Double check that ONE warning was raised.
        assert len(recwarn) == 1

    def test_add_existing(self, dispatcher, mock_interface, mock_settings):
        with pytest.raises(ValueError) as inst:
            dispatcher.add_interface("test", mock_interface, mock_settings)

    def test_collection(self, dispatcher, mock_interface, mock_settings):
        assert "test" not in dispatcher.collections

        dispatcher.add_interface("coll", mock_interface, mock_settings, is_collection=True)

        assert "coll" in dispatcher.collections

class TestConnections:

    def test_add_connection(self, dispatcher, mock_connection):

        conn = mock_connection(dispatcher)

        assert dispatcher.connections == []

        conn.open()

        assert dispatcher.connections == [conn]

    def test_add_bad_connection(self, dispatcher, bad_connection):

        conn = bad_connection(dispatcher)

        for method_name in ["open", "on_message", "on_close", "write_message"]:
            if getattr(conn, method_name) is None:
                with pytest.raises(AttributeError):
                    dispatcher.add_connection(conn)

    def test_remove_connection(self, dispatcher, mock_connection):

        conn = mock_connection(dispatcher)

        conn.open()

        assert dispatcher.connections == [conn]

        dispatcher.remove_connection(conn)

        assert dispatcher.connections == []

    def test_close_connection(self, dispatcher, mock_connection):

        conn = mock_connection(dispatcher)

        conn.open()

        assert dispatcher.connections == [conn]

        conn.close()

        assert dispatcher.connections == []

class TestDispatch:

    def test_dispatch_authorized(self, dispatcher, mock_connection):
        """
        Test the both of two AUTHORIZED connections can have messages dispatched through them using their
        ``write_message`` methods.

        """
        conn1 = mock_connection(dispatcher)
        conn2 = mock_connection(dispatcher)

        conn1.open()
        conn2.open()

        message = {
            "operation": "test",
            "collection_name": "test",
            "data": dict(message="This is a test")
        }

        dispatcher.dispatch(message)

        assert conn1.messages[0] == message
        assert conn2.messages[0] == message

    def test_dispatch_unauthorized(self, dispatcher, mock_connection):
        """
        Test the neither of two unauthorized connections have their ``write_message`` methods called during a dispatch.

        """
        conn1 = mock_connection(dispatcher, authorized=False)
        conn2 = mock_connection(dispatcher, authorized=False)

        conn1.open()
        conn2.open()

        dispatcher.dispatch({
            "operation": "test",
            "collection_name": "test",
            "data": dict(message="This is a test")
        })

        assert conn1.messages == []
        assert conn2.messages == []

    def test_dispatch_either(self, dispatcher, mock_connection):
        """
        Test the only the authorized connection has its ``write_message`` method called when an one authorized and one
        unauthorized connection are managed by the dispatcher.

        """
        conn1 = mock_connection(dispatcher, authorized=False)
        conn2 = mock_connection(dispatcher, authorized=True)

        conn1.open()
        conn2.open()

        message = {
            "operation": "test",
            "collection_name": "test",
            "data": dict(message="Only connection 2 should see this")
        }

        dispatcher.dispatch(message)

        assert conn1.messages == []
        assert conn2.messages == [message]

    def test_dispatch_specific(self, dispatcher, mock_connection):
        """
        Test that only the connection passed in the keyword argument ``connections`` has its ``write_message``
        method called when a dispatch occurs.

        """
        conn1 = mock_connection(dispatcher)
        conn2 = mock_connection(dispatcher)

        conn1.open()
        conn2.open()

        message = {
            "operation": "test",
            "collection_name": "test",
            "data": dict(message="This is a test")
        }

        dispatcher.dispatch(message, connections=[conn2])

        assert conn1.messages == []
        assert conn2.messages[0] == message

class TestHandle:

    @pytest.mark.gen_test
    def test_unknown_interface(self, caplog, dispatcher, mock_connection):
        """
        Test that a warning is logged when an unknown dispatcher interface is requested.

        """
        conn = mock_connection(permissions="all")

        code = yield dispatcher.handle(json.dumps({
            "tid": 9029401982,
            "interface": "foo",
            "method": "test_exposed_method",
            "data": {
                "value": "foobar"
            }
        }), conn)

        assert not code

        last_message = list(caplog.records)[-1].message

        assert last_message == "User 'test' requested unknown interface foo"

    @pytest.mark.gen_test
    def test_unknown_method(self, caplog, dispatcher, mock_connection):
        """
        Test that a warning is logged when an unknown interface method is requested.

        """
        conn = mock_connection(dispatcher)

        code = yield dispatcher.handle(json.dumps({
            "tid": 9029401982,
            "interface": "test",
            "method": "non_existent_method",
            "data": {
                "value": "foobar"
            }
        }), conn)

        assert not code

        last_message = list(caplog.records)[-1].message

        assert last_message == "User 'test' requested unknown interface method test.non_existent_method"

    @pytest.mark.gen_test
    def test_unexposed_method(self, caplog, dispatcher, mock_connection):
        """
        Test that a warning is logged when an unknown exposed method is requested.

        """
        conn = mock_connection(dispatcher)

        code = yield dispatcher.handle(json.dumps({
            "tid": 9029401983,
            "interface": "test",
            "method": "test_unexposed_method",
            "data": {
                "value": "foobar"
            }
        }), conn)

        assert not code

        last_message = list(caplog.records)[-1].message

        assert last_message == "User 'test' attempted to call unexposed method test.test_unexposed_method"

    @pytest.mark.gen_test
    def test_unauthorized(self, caplog, dispatcher, mock_connection):
        """
        Test that a warning is logged when an protected exposed method is requested by an unauthorized connection.

        """
        conn = mock_connection(dispatcher, authorized=False)

        code = yield dispatcher.handle(json.dumps({
            "tid": 9029401983,
            "interface": "test",
            "method": "test_exposed_method",
            "data": {
                "value": "foobar"
            }
        }), conn)

        assert not code

        last_message = list(caplog.records)[-1].message

        assert last_message == "Unauthorized connection at 8.8.8.8 attempted to call" \
                               " protected method test.test_exposed_method"

    @pytest.mark.gen_test
    def test_unprotected(self, dispatcher, mock_connection):
        """
        Test that unauthorized connections can successfully call unprotected methods.

        """
        conn = mock_connection(dispatcher, authorized=False)

        code = yield dispatcher.handle(json.dumps({
            "tid": 9029401983,
            "interface": "test",
            "method": "test_unprotected_method",
            "data": {
                "value": "foobar"
            }
        }), conn)

        assert code

        message_interfaces = [message["collection_name"] for message in conn.messages]

        assert "transaction" in message_interfaces
        assert "test" not in message_interfaces

    @pytest.mark.gen_test
    def test_authorized(self, dispatcher, mock_connection):
        """
        Test that unauthorized connections can successfully call unprotected methods.

        """
        conn = mock_connection(dispatcher, authorized=True)

        code = yield dispatcher.handle(json.dumps({
            "tid": 9029401983,
            "interface": "test",
            "method": "test_exposed_method",
            "data": {
                "value": "foobar"
            }
        }), conn)

        assert code

        message_interfaces = [message["collection_name"] for message in conn.messages]

        assert "transaction" in message_interfaces
        assert "test" not in message_interfaces

    @pytest.mark.gen_test
    def test_not_permitted(self, caplog, dispatcher, mock_connection):
        """
        Test that unauthorized connections can successfully call unprotected methods.

        """
        conn = mock_connection(dispatcher, permissions=[], authorized=True)

        code = yield dispatcher.handle(json.dumps({
            "tid": 9029401983,
            "interface": "test",
            "method": "test_permissions_exposed_method",
            "data": {
                "value": "foobar"
            }
        }), conn)

        last_message = list(caplog.records)[-1].message

        assert not code

        assert last_message == "User test attempted to call exposed method test.test_permissions_exposed_method" \
                               " without the required permissions"

    @pytest.mark.gen_test
    def test_is_permitted(self, dispatcher, mock_connection):
        """
        Test that unauthorized connections can successfully call unprotected methods.

        """
        conn = mock_connection(dispatcher, permissions=["modify_options"], authorized=True)

        code = yield dispatcher.handle(json.dumps({
            "tid": 9029401983,
            "interface": "test",
            "method": "test_permissions_exposed_method",
            "data": {
                "value": "foobar"
            }
        }), conn)

        assert code

    @pytest.mark.gen_test
    def test_no_transaction(self, dispatcher, mock_connection):
        """
        Test that exposed methods with no transaction argument raise a TypeError.

        """
        conn = mock_connection(dispatcher, authorized=True)

        with pytest.raises(TypeError) as inst:
            yield dispatcher.handle(json.dumps({
                "tid": 9029401983,
                "interface": "test",
                "method": "test_no_transaction_method",
                "data": {
                    "value": "foobar"
                }
            }), conn)
<|endoftext|>


<commit_before>#-------------------------------------------------------------------------------------------------------------
## Import Statements
#-------------------------------------------------------------------------------------------------------------
from __future__ import print_function

import sys

from .utils import get_builds, get_ggd_channels, get_species

SPECIES_LIST =  sorted(get_species(update_files=False))
GENOME_BUILDS = sorted(get_builds("*"))
CHANNEL_LIST = [x.encode('ascii') for x in get_ggd_channels()]

#-------------------------------------------------------------------------------------------------------------
## Argument Parser
#-------------------------------------------------------------------------------------------------------------
def add_search(p):
    import argparse
    c = p.add_parser("search", help="Search for a ggd data package", description="Search for available ggd data packages. Results are filtered by match score from high to low. (Only 5 results will be reported unless the -dn flag is changed)")
    c.add_argument("search_term", nargs="+", help="**Required** The term(s) to search for. Multiple terms can be used. Example: 'ggd search reference genome'")
    c.add_argument("-g", "--genome-build",default=[],action="append", choices=[str(x) for x in GENOME_BUILDS], help="(Optional) Filter results by the genome build of the desired recipe")
    c.add_argument("-s", "--species", default=[], action="append", help="(Optional) Filter results by the species for the desired recipe", choices=[str(x) for x in SPECIES_LIST])
    c.add_argument("-dn","--display-number", default=5, help="(Optional) The number of search results to display. (Default = 5)")
    c.add_argument("-m", "--match-score", default="75", help="(Optional) A score between 0 and 100 to use to filter the results by. (Default = 75). The lower the number the more results will be output")
    c.add_argument("-c", "--channel", help="(Optional) The ggd channel to search. (Default = genomics)", choices=[x.decode('ascii') for x in CHANNEL_LIST],
                    default="genomics")
    c.set_defaults(func=search)

#-------------------------------------------------------------------------------------------------------------
## Functions/Methods
#-------------------------------------------------------------------------------------------------------------

def load_json(jfile):
    """Method to load a json file into a dictionary

    load_json
    =========
    Method to load a json file

    Parameters:
    ---------
    1) jfile: The path to the json file

    Returns:
    1) A dictionary of a json object
    """
    import json

    with open(jfile) as jsonFile:
       return(json.load(jsonFile))

def load_json_from_url(json_url):
    """Method to load a json file from a url

    load_json_from_url
    ==================
    Method to load a json file  from a url. Uses the requests module
     to get the json file from the url.

    Parameters:
    ---------
    1) json_url: The url to the json path

    Returns:
    ++++++++
    1) A dictionary of a json object
    """
    import json
    import requests
    import traceback

    try:
        return(requests.get(json_url).json())
    except ValueError as e:
        sys.stderr.write("\n:ggd:search: !!ERROR!! in loading json file from url")
        sys.stderr.write("\n\t Invalid URL: %s" %json_url)
        sys.stderr.write(str(e))
        sys.stderr.write(traceback.format_exc())
        sys.exit(1)

def search_packages(json_dict,search_terms,score_cutoff=50):
    """Method to search for ggd packages in the ggd channeldata.json metadata file based on user provided search terms

    search_packages
    ===============
    Method to search for ggd packages/recipes
     contaning specific search terms

     NOTE: Both the package name and the package keywords are searched

    Parameters:
    ---------
    1) json_dict: A json file loaded into a dictionary. (The file to search)
                  the load_json_from_url() method creates the dictionary
    2) search_terms: A list of terms representing package names or keywords to search for
    3) score_cutoff: A number betwnee 0 and 100 that represent which matches to return
                        (Default = 50)

    Returns:
    ++++++++
    1) A list of pkg names who's either name or keyword match score reached the score cutoff
    """
    from fuzzywuzzy import fuzz
    from fuzzywuzzy import process

    pkg_score = {}
    for term in search_terms:
        for pkg in json_dict["packages"].keys():
            ## Get match score between name and term
            score = fuzz.partial_ratio(term.lower(),pkg.lower())
            ## Get a list of match score for the keywrods
            additional_scores = process.extract(term, json_dict["packages"][pkg]["keywords"])
            ## Get the max score from all scores found
            max_score = max([x[1] for x in additional_scores] + [score])

            ## Set max score in dict
            if pkg in pkg_score:
                pkg_score[pkg] = max(pkg_score[pkg], max_score)
            else:
                pkg_score[pkg] = max_score

    ## Get a final list of pkg names
    temp_list = sorted([[pkg,int(max_score)] for pkg,max_score in pkg_score.items() if max_score >= score_cutoff], key = lambda x: x[1], reverse = True)
    final_list = [pkg_list[0] for pkg_list in temp_list]

    return(final_list)

def check_installed(ggd_recipe,ggd_jdict):
    """Method to check if the recipe has already been installed and is in the conda ggd storage path.

    check_if_installed
    ==================
    This method is used to check if the ggd package has been installed and is located in the ggd storage path.
    """
    import glob
    import os
    from .utils import conda_root

    species = ggd_jdict["packages"][ggd_recipe]["identifiers"]["species"]
    build = ggd_jdict["packages"][ggd_recipe]["identifiers"]["genome-build"]
    version = ggd_jdict["packages"][ggd_recipe]["version"]

    CONDA_ROOT = conda_root()
    path = os.path.join(CONDA_ROOT,"share","ggd",species,build,ggd_recipe,version)
    recipe_exists = glob.glob(path)
    if recipe_exists:
        return(True,path)

    else:
        return(False,None)

def filter_by_identifiers(iden_keys,json_dict,filter_terms):
    """Method to filter a dictionary by an identifier field for the certain package.

    filter_by_identifiers
    =====================
    A method used to filter the list of data packages by information in the
     identifiers field in the channeldata.json file

    Parameters:
    ----------
    1) iden_keys: A list of he identifiers keys. Example = ["species","genome-build"]
    2) json_dict: The json dictionary craeted from laod_json()
    3) filter_terms: A list of the term(s) to filter by. Example: ["Homo_sapiens","hg19"]

    NOTE: List order of iden_keys should match list order of filter_terms

    Returns:
    ++++++++
    1) Updated/filtered json_dict
    """
    import copy

    key_count = len(json_dict["packages"].keys())
    keys = json_dict["packages"].keys()

    keys_to_keep = set()
    if len(iden_keys) > 0 and len(iden_keys) == len(filter_terms):
        for key in keys:
            for i,iden_key in enumerate(iden_keys):
                if iden_key in json_dict["packages"][key]["identifiers"]:
                    if len(filter_terms[i]) == 0:
                        continue
                    if filter_terms[i] in json_dict["packages"][key]["identifiers"][iden_key]:
                        keys_to_keep.add(key)

    new_json_dict = copy.deepcopy(json_dict)
    ## Remove packages
    if len(keys_to_keep) > 0:
        for key in keys:
            if key not in keys_to_keep:
                del new_json_dict["packages"][key]

    if len(new_json_dict["packages"].keys()) == key_count:
        ## If unable to return a filtered set return the original match list
        print("\n:ggd:search: WARNING: Unable to filter packages using: '%s'" %", ".join(filter_terms))
        print("\tThe un-filtered list will be used\n")

    return(new_json_dict)

def print_summary(search_terms,json_dict,match_list,installed_pkgs,installed_paths):
    """ Method to print the summary/results of the search

    print_summary
    ============
    Method used to print out the final set of searched packages

    Parameters:
    ---------
    1) search_terms: The search terms from the user
    2) json_dict: The json dictionary from the load_json() method
    3) match_list: The filtered and final set of searched recipes
    4) isntalled_pkgs: A set of pkg names that are installed
    5) installed_paths: A dictionary with keys = pkg names, values = installed paths

    Returns:
    +++++++
    1) True if print summary printed out successfully
    """

    dash = "     " + "-" * 100

    if len(match_list) < 1:
        print("\n:ggd:search: No results for %s. Update your search term(s) and try again" %", ".join(search_terms))
        sys.exit()
    print("\n",dash)
    for pkg in match_list:
        results = []
        if pkg in json_dict["packages"]:
            #results.append("\n\t{} {}\n".format(("\033[1m" + "GGD Package:" + "\033[0m"), pkg))
            results.append("\n\t{}\n\t{}".format(("\033[1m" + pkg + "\033[0m"),"="*len(pkg)))
            if "summary" in json_dict["packages"][pkg] and json_dict["packages"][pkg]["summary"]:
                results.append("\t{} {}".format(("\033[1m" + "Summary:" + "\033[0m"), json_dict["packages"][pkg]["summary"]))
            if "identifiers" in json_dict["packages"][pkg] and json_dict["packages"][pkg]["identifiers"]:
                results.append("\t{} {}".format(("\033[1m" + "Species:" + "\033[0m"), json_dict["packages"][pkg]["identifiers"]["species"]))
                results.append("\t{} {}".format(("\033[1m" + "Genome Build:" + "\033[0m"), json_dict["packages"][pkg]["identifiers"]["genome-build"]))
            if "keywords" in json_dict["packages"][pkg] and json_dict["packages"][pkg]["keywords"]:
                results.append("\t{} {}".format(("\033[1m" + "Keywords:" + "\033[0m"), ", ".join(json_dict["packages"][pkg]["keywords"])))
            if "tags" in json_dict["packages"][pkg] and json_dict["packages"][pkg]["tags"]:
                if "cache" in json_dict["packages"][pkg]["tags"]:
                    results.append("\t{} {}".format(("\033[1m" + "Cached:" + "\033[0m"), json_dict["packages"][pkg]["tags"]["cached"]))
                if "data-provider" in json_dict["packages"][pkg]["tags"]:
                    results.append("\t{} {}".format(("\033[1m" + "Data Provider" + "\033[0m"), json_dict["packages"][pkg]["tags"]["data-provider"]))
                if "data-version" in json_dict["packages"][pkg]["tags"]:
                    results.append("\t{} {}".format(("\033[1m" + "Data Version:" + "\033[0m"), json_dict["packages"][pkg]["tags"]["data-version"]))
                if "file-type" in json_dict["packages"][pkg]["tags"]:
                    results.append("\t{} {}".format(("\033[1m" + "File type(s):" + "\033[0m"), ", ".join(json_dict["packages"][pkg]["tags"]["file-type"])))
                if "genomic-coordinate-base" in json_dict["packages"][pkg]["tags"]:
                    results.append("\t{} {}".format(("\033[1m" + "Data file coordinate base:" + "\033[0m"),json_dict["packages"][pkg]["tags"]["genomic-coordinate-base"]))
                if "final-files" in json_dict["packages"][pkg]["tags"]:
                    results.append("\t{} {}".format(("\033[1m" + "Included Data Files:" + "\033[0m"), "\n\t\t"+"\n\t\t".join(json_dict["packages"][pkg]["tags"]["final-files"])))

            if pkg in installed_pkgs: ## IF installed
                results.append("\n\tThis pacakge is already installed on your system.\n\t  You can find the installed data files here:  %s" %installed_paths[pkg])
            else:
                results.append("\n\tTo install run:\n\t\tggd install %s" %pkg)
        print("\n\n".join(results))
        print("\n",dash)

    return(True)

def search(parser, args):
    """Main method for ggd search.

    search
    =====
    Main method for running a recipe/package search

    Parameters:
    ----------
    1) parser
    2) args
    """
    from .utils import get_channeldata_url, get_builds

    ## load the channeldata.json file
    j_dict = load_json_from_url(get_channeldata_url(args.channel))

    ## identify if search_terms have any species or genome build in them
    species_lower = {x.lower():x for x in SPECIES_LIST}
    gb_lower = {x.lower():x for x in GENOME_BUILDS}
    filtered_search_terms = []
    for term in args.search_term:
        if term.lower() in species_lower.keys():
            if species_lower[term.lower()] not in args.species:
                args.species.append(species_lower[term.lower()])
        elif term.lower() in gb_lower.keys():
            if gb_lower[term.lower()] not in args.genome_build:
                args.genome_build.append(gb_lower[term.lower()])
        else:
            ## Only use search terms that are not used to filter the results by identifiers
            filtered_search_terms.append(term)

    ## genome_build takes precedence over species (If genome build provided, species is implied)
    final_species_list = args.species
    for species in args.species:
        build = get_builds(species)
        if [x for x in build if x in args.genome_build]:
            final_species_list.remove(species)
    args.species = final_species_list

    ## Filter the json dict by species or genome build if applicable
    if args.genome_build or args.species:
        j_dict = filter_by_identifiers(["species"]*len(args.species)+["genome-build"]*len(args.genome_build),j_dict,args.species+args.genome_build)

    ## Search pkg names and keywords
    match_results = search_packages(j_dict,filtered_search_terms,int(args.match_score))

    ## Get installed paths
    installed_dict = {}
    installed_set = set()
    for pkg in match_results:
        isinstalled, path = check_installed(pkg,j_dict)
        if isinstalled:
            installed_dict[pkg] = path
            installed_set.add(pkg)

    ## Print search results
    if len(match_results) >= int(args.display_number):
        subset_match_results = match_results[0:int(args.display_number)]
    else:
        subset_match_results = match_results

    return(print_summary(args.search_term,j_dict,subset_match_results,installed_set,installed_dict))
<commit_msg>Update stdout to inform user if not all results are being displayed<commit_after>#-------------------------------------------------------------------------------------------------------------
## Import Statements
#-------------------------------------------------------------------------------------------------------------
from __future__ import print_function

import sys

from .utils import get_builds, get_ggd_channels, get_species

SPECIES_LIST =  sorted(get_species(update_files=False))
GENOME_BUILDS = sorted(get_builds("*"))
CHANNEL_LIST = [x.encode('ascii') for x in get_ggd_channels()]

#-------------------------------------------------------------------------------------------------------------
## Argument Parser
#-------------------------------------------------------------------------------------------------------------
def add_search(p):
    import argparse
    c = p.add_parser("search", help="Search for a ggd data package", description="Search for available ggd data packages. Results are filtered by match score from high to low. (Only 5 results will be reported unless the -dn flag is changed)")
    c.add_argument("search_term", nargs="+", help="**Required** The term(s) to search for. Multiple terms can be used. Example: 'ggd search reference genome'")
    c.add_argument("-g", "--genome-build",default=[],action="append", choices=[str(x) for x in GENOME_BUILDS], help="(Optional) Filter results by the genome build of the desired recipe")
    c.add_argument("-s", "--species", default=[], action="append", help="(Optional) Filter results by the species for the desired recipe", choices=[str(x) for x in SPECIES_LIST])
    c.add_argument("-dn","--display-number", default=5, help="(Optional) The number of search results to display. (Default = 5)")
    c.add_argument("-m", "--match-score", default="75", help="(Optional) A score between 0 and 100 to use to filter the results by. (Default = 75). The lower the number the more results will be output")
    c.add_argument("-c", "--channel", help="(Optional) The ggd channel to search. (Default = genomics)", choices=[x.decode('ascii') for x in CHANNEL_LIST],
                    default="genomics")
    c.set_defaults(func=search)

#-------------------------------------------------------------------------------------------------------------
## Functions/Methods
#-------------------------------------------------------------------------------------------------------------

def load_json(jfile):
    """Method to load a json file into a dictionary

    load_json
    =========
    Method to load a json file

    Parameters:
    ---------
    1) jfile: The path to the json file

    Returns:
    1) A dictionary of a json object
    """
    import json

    with open(jfile) as jsonFile:
       return(json.load(jsonFile))

def load_json_from_url(json_url):
    """Method to load a json file from a url

    load_json_from_url
    ==================
    Method to load a json file  from a url. Uses the requests module
     to get the json file from the url.

    Parameters:
    ---------
    1) json_url: The url to the json path

    Returns:
    ++++++++
    1) A dictionary of a json object
    """
    import json
    import requests
    import traceback

    try:
        return(requests.get(json_url).json())
    except ValueError as e:
        sys.stderr.write("\n:ggd:search: !!ERROR!! in loading json file from url")
        sys.stderr.write("\n\t Invalid URL: %s" %json_url)
        sys.stderr.write(str(e))
        sys.stderr.write(traceback.format_exc())
        sys.exit(1)

def search_packages(json_dict,search_terms,score_cutoff=50):
    """Method to search for ggd packages in the ggd channeldata.json metadata file based on user provided search terms

    search_packages
    ===============
    Method to search for ggd packages/recipes
     contaning specific search terms

     NOTE: Both the package name and the package keywords are searched

    Parameters:
    ---------
    1) json_dict: A json file loaded into a dictionary. (The file to search)
                  the load_json_from_url() method creates the dictionary
    2) search_terms: A list of terms representing package names or keywords to search for
    3) score_cutoff: A number betwnee 0 and 100 that represent which matches to return
                        (Default = 50)

    Returns:
    ++++++++
    1) A list of pkg names who's either name or keyword match score reached the score cutoff
    """
    from fuzzywuzzy import fuzz
    from fuzzywuzzy import process

    pkg_score = {}
    for term in search_terms:
        for pkg in json_dict["packages"].keys():
            ## Get match score between name and term
            score = fuzz.partial_ratio(term.lower(),pkg.lower())
            ## Get a list of match score for the keywrods
            additional_scores = process.extract(term, json_dict["packages"][pkg]["keywords"])
            ## Get the max score from all scores found
            max_score = max([x[1] for x in additional_scores] + [score])

            ## Set max score in dict
            if pkg in pkg_score:
                pkg_score[pkg] = max(pkg_score[pkg], max_score)
            else:
                pkg_score[pkg] = max_score

    ## Get a final list of pkg names
    temp_list = sorted([[pkg,int(max_score)] for pkg,max_score in pkg_score.items() if max_score >= score_cutoff], key = lambda x: x[1], reverse = True)
    final_list = [pkg_list[0] for pkg_list in temp_list]

    return(final_list)

def check_installed(ggd_recipe,ggd_jdict):
    """Method to check if the recipe has already been installed and is in the conda ggd storage path.

    check_if_installed
    ==================
    This method is used to check if the ggd package has been installed and is located in the ggd storage path.
    """
    import glob
    import os
    from .utils import conda_root

    species = ggd_jdict["packages"][ggd_recipe]["identifiers"]["species"]
    build = ggd_jdict["packages"][ggd_recipe]["identifiers"]["genome-build"]
    version = ggd_jdict["packages"][ggd_recipe]["version"]

    CONDA_ROOT = conda_root()
    path = os.path.join(CONDA_ROOT,"share","ggd",species,build,ggd_recipe,version)
    recipe_exists = glob.glob(path)
    if recipe_exists:
        return(True,path)

    else:
        return(False,None)

def filter_by_identifiers(iden_keys,json_dict,filter_terms):
    """Method to filter a dictionary by an identifier field for the certain package.

    filter_by_identifiers
    =====================
    A method used to filter the list of data packages by information in the
     identifiers field in the channeldata.json file

    Parameters:
    ----------
    1) iden_keys: A list of he identifiers keys. Example = ["species","genome-build"]
    2) json_dict: The json dictionary craeted from laod_json()
    3) filter_terms: A list of the term(s) to filter by. Example: ["Homo_sapiens","hg19"]

    NOTE: List order of iden_keys should match list order of filter_terms

    Returns:
    ++++++++
    1) Updated/filtered json_dict
    """
    import copy

    key_count = len(json_dict["packages"].keys())
    keys = json_dict["packages"].keys()

    keys_to_keep = set()
    if len(iden_keys) > 0 and len(iden_keys) == len(filter_terms):
        for key in keys:
            for i,iden_key in enumerate(iden_keys):
                if iden_key in json_dict["packages"][key]["identifiers"]:
                    if len(filter_terms[i]) == 0:
                        continue
                    if filter_terms[i] in json_dict["packages"][key]["identifiers"][iden_key]:
                        keys_to_keep.add(key)

    new_json_dict = copy.deepcopy(json_dict)
    ## Remove packages
    if len(keys_to_keep) > 0:
        for key in keys:
            if key not in keys_to_keep:
                del new_json_dict["packages"][key]

    if len(new_json_dict["packages"].keys()) == key_count:
        ## If unable to return a filtered set return the original match list
        print("\n:ggd:search: WARNING: Unable to filter packages using: '%s'" %", ".join(filter_terms))
        print("\tThe un-filtered list will be used\n")

    return(new_json_dict)

def print_summary(search_terms,json_dict,match_list,installed_pkgs,installed_paths):
    """ Method to print the summary/results of the search

    print_summary
    ============
    Method used to print out the final set of searched packages

    Parameters:
    ---------
    1) search_terms: The search terms from the user
    2) json_dict: The json dictionary from the load_json() method
    3) match_list: The filtered and final set of searched recipes
    4) isntalled_pkgs: A set of pkg names that are installed
    5) installed_paths: A dictionary with keys = pkg names, values = installed paths

    Returns:
    +++++++
    1) True if print summary printed out successfully
    """

    dash = "     " + "-" * 100

    if len(match_list) < 1:
        print("\n:ggd:search: No results for %s. Update your search term(s) and try again" %", ".join(search_terms))
        sys.exit()
    print("\n",dash)
    for pkg in match_list:
        results = []
        if pkg in json_dict["packages"]:
            #results.append("\n\t{} {}\n".format(("\033[1m" + "GGD Package:" + "\033[0m"), pkg))
            results.append("\n\t{}\n\t{}".format(("\033[1m" + pkg + "\033[0m"),"="*len(pkg)))
            if "summary" in json_dict["packages"][pkg] and json_dict["packages"][pkg]["summary"]:
                results.append("\t{} {}".format(("\033[1m" + "Summary:" + "\033[0m"), json_dict["packages"][pkg]["summary"]))
            if "identifiers" in json_dict["packages"][pkg] and json_dict["packages"][pkg]["identifiers"]:
                results.append("\t{} {}".format(("\033[1m" + "Species:" + "\033[0m"), json_dict["packages"][pkg]["identifiers"]["species"]))
                results.append("\t{} {}".format(("\033[1m" + "Genome Build:" + "\033[0m"), json_dict["packages"][pkg]["identifiers"]["genome-build"]))
            if "keywords" in json_dict["packages"][pkg] and json_dict["packages"][pkg]["keywords"]:
                results.append("\t{} {}".format(("\033[1m" + "Keywords:" + "\033[0m"), ", ".join(json_dict["packages"][pkg]["keywords"])))
            if "tags" in json_dict["packages"][pkg] and json_dict["packages"][pkg]["tags"]:
                if "cache" in json_dict["packages"][pkg]["tags"]:
                    results.append("\t{} {}".format(("\033[1m" + "Cached:" + "\033[0m"), json_dict["packages"][pkg]["tags"]["cached"]))
                if "data-provider" in json_dict["packages"][pkg]["tags"]:
                    results.append("\t{} {}".format(("\033[1m" + "Data Provider" + "\033[0m"), json_dict["packages"][pkg]["tags"]["data-provider"]))
                if "data-version" in json_dict["packages"][pkg]["tags"]:
                    results.append("\t{} {}".format(("\033[1m" + "Data Version:" + "\033[0m"), json_dict["packages"][pkg]["tags"]["data-version"]))
                if "file-type" in json_dict["packages"][pkg]["tags"]:
                    results.append("\t{} {}".format(("\033[1m" + "File type(s):" + "\033[0m"), ", ".join(json_dict["packages"][pkg]["tags"]["file-type"])))
                if "genomic-coordinate-base" in json_dict["packages"][pkg]["tags"]:
                    results.append("\t{} {}".format(("\033[1m" + "Data file coordinate base:" + "\033[0m"),json_dict["packages"][pkg]["tags"]["genomic-coordinate-base"]))
                if "final-files" in json_dict["packages"][pkg]["tags"]:
                    results.append("\t{} {}".format(("\033[1m" + "Included Data Files:" + "\033[0m"), "\n\t\t"+"\n\t\t".join(json_dict["packages"][pkg]["tags"]["final-files"])))

            if pkg in installed_pkgs: ## IF installed
                results.append("\n\tThis pacakge is already installed on your system.\n\t  You can find the installed data files here:  %s" %installed_paths[pkg])
            else:
                results.append("\n\tTo install run:\n\t\tggd install %s" %pkg)
        print("\n\n".join(results))
        print("\n",dash)

    return(True)

def search(parser, args):
    """Main method for ggd search.

    search
    =====
    Main method for running a recipe/package search

    Parameters:
    ----------
    1) parser
    2) args
    """
    from .utils import get_channeldata_url, get_builds

    ## load the channeldata.json file
    j_dict = load_json_from_url(get_channeldata_url(args.channel))

    ## identify if search_terms have any species or genome build in them
    species_lower = {x.lower():x for x in SPECIES_LIST}
    gb_lower = {x.lower():x for x in GENOME_BUILDS}
    filtered_search_terms = []
    for term in args.search_term:
        if term.lower() in species_lower.keys():
            if species_lower[term.lower()] not in args.species:
                args.species.append(species_lower[term.lower()])
        elif term.lower() in gb_lower.keys():
            if gb_lower[term.lower()] not in args.genome_build:
                args.genome_build.append(gb_lower[term.lower()])
        else:
            ## Only use search terms that are not used to filter the results by identifiers
            filtered_search_terms.append(term)

    ## genome_build takes precedence over species (If genome build provided, species is implied)
    final_species_list = args.species
    for species in args.species:
        build = get_builds(species)
        if [x for x in build if x in args.genome_build]:
            final_species_list.remove(species)
    args.species = final_species_list

    ## Filter the json dict by species or genome build if applicable
    if args.genome_build or args.species:
        j_dict = filter_by_identifiers(["species"]*len(args.species)+["genome-build"]*len(args.genome_build),j_dict,args.species+args.genome_build)

    ## Search pkg names and keywords
    match_results = search_packages(j_dict,filtered_search_terms,int(args.match_score))

    ## Get installed paths
    installed_dict = {}
    installed_set = set()
    for pkg in match_results:
        isinstalled, path = check_installed(pkg,j_dict)
        if isinstalled:
            installed_dict[pkg] = path
            installed_set.add(pkg)

    ## Print search results
    match_result_num = str(len(match_results))
    if int(match_result_num) >= int(args.display_number):
        subset_match_results = match_results[0:int(args.display_number)]
    else:
        subset_match_results = match_results

    ## Print search results to STDOUT
    printed = print_summary(args.search_term,j_dict,subset_match_results,installed_set,installed_dict)

    ## Add a comment if a subset of search results are provided
    if int(match_result_num) > int(args.display_number):
        print("\n\n:ggd:search: NOTE  Only showing results for top {d} of {m} matches.".format(d = str(args.display_number), m = match_result_num))

        print(":ggd:search: To display all matches append your search command with '-dn {m}'".format(m = match_result_num))
        print("\n\t ggd search {t} -dn {m}\n".format(t = " ".join(args.search_term), m = match_result_num))

    ## Return result of print_summary
    return(printed)

<|endoftext|>


<commit_before>import yaml
import os
import collections
from bcbio import utils
import subprocess
import functools
import stat
import difflib
from itertools import repeat
from functools import reduce

def in2out(in_file, word, transform=True, out_dir=None):
    """
    creates an out_file name from an in_file name by adding word to the
    in_file name and writing it to the output_directory if specified.
    if transform is True, the word replaces the extension of in_file.
    if False it adds it to the stem
    transform: in2out("test.bam", "sam", transform=True) -> "test.sam"
    non transform: in2out("test.bam", "sam", transform=False) -> "test_bam.sam"

    """
    (base, ext) = os.path.splitext(in_file)
    if out_dir:
        base = os.path.join(out_dir, os.path.basename(base))

    if transform:
        return "%s.%s" % (base, word)

    else:
        return "%s_%s.%s" % (base, ext, word)

def get_stem(filename):
    return os.path.basename(remove_suffix(filename))

def combine_pairs(input_files):
    """ calls files pairs if they are completely the same except
    for one has _1 and the other has _2 returns a list of tuples
    of pairs or singles """
    PAIR_FILE_IDENTIFIERS = ["1", "2"]

    pairs = []
    used = []
    for in_file in input_files:
        if in_file in used:
            continue
        for comp_file in input_files:
            if comp_file in used:
                continue
            s = difflib.SequenceMatcher(a=in_file, b=comp_file)
            blocks = s.get_matching_blocks()
            # length 3 means on match in the middle of the string
            if len(s.get_matching_blocks()) is not 3:
                continue
            if comp_file[blocks[0][2]] in PAIR_FILE_IDENTIFIERS:
                used.append(in_file)
                used.append(comp_file)
                pairs.append([in_file, comp_file])
                break
        if in_file not in used:
            pairs.append([in_file])
            used.append(in_file)

    return pairs

def freeze_files(files, directory):
    """ freezes a list of files, copying them to directory and
    setting them to be read only """
    new_files = [os.path.join(directory, x) for x in
                 map(os.path.basename, files)]
    [os.chmod(x, stat.S_IREAD | stat.S_IRGRP) for x in new_files]
    return new_files

# XXX I'm not 100% sure if this actually works, I forget use with caution
def memoize_outfile_to_dir(res_dir, ext):
    """Creates outfile from input file and ext, running if outfile not present.

    This requires a standard function usage. The first arg, or kwarg 'in_file', needs
    to be the input file that is being processed. The output name is created with the
    provided ext relative to the input. The function is only run if the created
    out_file is not present.
    Memoize_outfile was lifted from bcbio.utils (Brad Chapman), added the
    writing to a directory portion
    """
    utils.safe_makedir(res_dir)
    def decor(f):
        @functools.wraps(f)
        def wrapper(*args, **kwargs):
            if len(args) > 0:
                in_file = args[0]
            else:
                in_file = kwargs['in_file']
            out_fname = "%s%s" % (os.path.splitext(in_file)[0], ext)
            out_file = os.path.join(res_dir, out_fname)
            if not os.path.exists(out_file) or os.path.getsize(out_file) == 0:
                kwargs['out_file'] = out_file
                f(*args, **kwargs)
            return out_file
        return wrapper
    return decor

def prepare_ref_file(ref, config):
    """Get a reference file, either by URL or locally.
    Lifted from Brad Chapman
    """
    url = ref.get("url", None)
    if url:
        ref_file = _download_ref(url, config["dir"]["ref"])
    else:
        ref_file = ref.get("file", None)
    assert ref_file is not None and os.path.exists(ref_file), ref_file
    return ref_file

def _download_ref(url, ref_dir):
    #Lifted from Brad Chapman
    dl_file = os.path.basename(url)
    ref_file = None
    for supported_ext, extract_cmd in [(".gz", "gunzip"),
                                       (".tgz", ("tar", "zxvf"))]:
        if dl_file.endswith(supported_ext):
            ref_file = os.path.join(ref_dir, dl_file[:-len(supported_ext)])
            break
    assert ref_file is not None, url
    if not os.path.exists(ref_file):
        with utils.chdir(ref_dir):
            cl = ["wget", url]
            subprocess.check_call(cl)
            cl = list(flatten([extract_cmd, dl_file]))
            subprocess.check_call(cl)
    return ref_file

def build_results_dir(stage_config, config):
    outdir = os.path.join(config["dir"]["results"],
                          stage_config["name"])
    return outdir

def transform_infile(filename, stage, delim="."):
    return replace_suffix(filename, stage.get("name", "transformed"), delim)

def filter_infile(filename, stage, delim="."):
    return append_stem(filename, stage.get("name", "filtered"), delim)

def remove_suffix(filename):
    filename, extension = os.path.splitext(filename)
    return filename

def append_stem(filename, word, delim="."):
    """ returns a filename with word appended to the stem
    example: /path/to/test.run.sam -> /path/to/test.run.sorted.sam """
    dirname = os.path.dirname(filename)
    filename = os.path.basename(filename)
    fsplit = filename.split(delim)
    fsplit.insert(len(fsplit) - 1, word)
    return os.path.join(dirname, delim.join(fsplit))

def replace_suffix(filename, suffix, delim="."):
    """ returns a filename with the suffix replaced
    with a new suffix
    exampoe: /path/to/test.run.sam -> /path/to/test.run.fa"""
    dirname = os.path.dirname(filename)
    filename = os.path.basename(filename)
    fsplit = filename.split(delim)[:-1]
    fsplit.append(suffix)
    return os.path.join(dirname, delim.join(fsplit))

def flatten_options(config):
    """ returns a list of pairs of options from a config file; returns
    an empty list if no options are found """
    return(flatten(config.get("options", {}.items())))

def flatten(l):
    for el in l:
        if isinstance(el, collections.Iterable) and not isinstance(el, basestring):
            for sub in flatten(el):
                yield sub
        else:
            yield el

def rfind_key(d, item):
    for k,v in d.items():
        if k == item:
            return (k, v)
        elif isinstance(v, dict):
            rfind_key(v, item)

def rfind_value(d, item):
    for k, v in d.items():
        if v == item:
            return (k, v)
        elif isinstance(v, dict):
            rfind_value(v, item)

def validate_config(config, data=True):
    # add the full path to any executables to avoid confusion later
    for k, v in config["program"].items():
        config["program"][k] = _add_full_path(v)

    if not all([os.path.exists(x) for x in config["program"].values()]):
        print("Cannot find all of the files or directories listed "
              "in 'program'.")
        exit(-1)

    if not all(x in get_stages(config) for x in config["run"]):
        print("Cannot find all of the stages listed to run.")
        exit(-1)

    if data and not all([os.path.exists(x) for x in config["input"]]):
        print("All of the input files do not exist.")
        exit(-1)

    return(config)

def get_stages(config):
    return config["stage"].keys()

def _add_full_path(program):
    """ add the full path to any executables """
    new_path = which(program)
    if new_path:
        program_config = new_path
    return program_config

def find_one_if(predicate, coll):
    """ returns a single entry in a collection if it contains exactly
    one true instance of a predicate and None otherwise """
    results = filter(predicate, coll)
    if len(results) != 1:
        return None
    else:
        print results
        return results[0]

def load_yaml(filename):
    with open(filename) as f:
        body = yaml.load(f)
        return body

def which(program):
    """ returns the path to an executable"""

    def is_exe(fpath):
        return os.path.isfile(fpath) and os.access(fpath, os.X_OK)

    fpath, fname = os.path.split(program)
    if fpath:
        if is_exe(program):
            return program
    else:
        for path in os.environ["PATH"].split(os.pathsep):
            exe_file = os.path.join(path, program)
            if is_exe(exe_file):
                return exe_file

    return None

"""
simple file tracking to keep track of the transformations a file has made
just playing with it for now
"""

class FileWithHistory(object):
    """ this tracks transformations of an input file so we can get the
    metadata associated with it at any point """

    def __init__(self, in_file, **kwargs):
        self._name_list = [in_file]
        [setattr(self, key, value) for key, value in kwargs.items()]

    def curr_file(self):
        return self._name_list[len(self._name_list)]

    def add_file(self, in_file):
        self._name_list.append(in_file)

    def was_file(self, in_file):
        return in_file in self._name_list

def update_tracker(tracker, curr_files):
    new_tracker = [track.add_file(new_file) for (track, new_file)
                   in zip(tracker, curr_files)]
    return new_tracker

def curr_files(tracker):
    return [x.curr_file() for x in tracker]

def dict_to_vectors(d):
    """ convert a dictionary to two vectors. returns a tuple of
    (key, value). For example: {"a": [1, 2, 3], "b": [4, 5]} returns
    (["a" "a" "a" "b" "b"], [1 2 3 4 5])

    """

    ks = []
    vs = []
    for k, v in d.items():
        ks += list(repeat(k, len(v)))
        vs += v
    return (ks, vs)

class dotdict(dict):
    """
    access dictioanry items via dot syntax
    via: http://parand.com/say/index.php/2008/10/24/python-dot-notation-dictionary-access/:
    """

    def __getattr__(self, attr):
        return self.get(attr, None)
    __setattr__ = dict.__setitem__
    __delattr__ = dict.__delitem__

def nested_lookup(d, t):
    """
    look up if you can get a tuple of values from a nested dictionary,
    each item in the tuple a deeper layer
    """
    return reduce(lambda d, t: d.get(t, {}), t, d)

def get_in(d, t, default=None):
    """
    look up if you can get a tuple of values from a nested dictionary,
    each item in the tuple a deeper layer

    example: get_in({1: {2: 3}}, (1, 2)) -> 3
    example: get_in({1: {2: 3}}, (2, 3)) -> {}
    """
    result = reduce(lambda d, t: d.get(t, {}), t, d)
    if not result:
        return default
    else:
        return result
<commit_msg>Added is_sequence and is_pair utility functions.<commit_after>import yaml
import os
import collections
from bcbio import utils
import subprocess
import functools
import stat
import difflib
from itertools import repeat
from functools import reduce

def in2out(in_file, word, transform=True, out_dir=None):
    """
    creates an out_file name from an in_file name by adding word to the
    in_file name and writing it to the output_directory if specified.
    if transform is True, the word replaces the extension of in_file.
    if False it adds it to the stem
    transform: in2out("test.bam", "sam", transform=True) -> "test.sam"
    non transform: in2out("test.bam", "sam", transform=False) -> "test_bam.sam"

    """
    (base, ext) = os.path.splitext(in_file)
    if out_dir:
        base = os.path.join(out_dir, os.path.basename(base))

    if transform:
        return "%s.%s" % (base, word)

    else:
        return "%s_%s.%s" % (base, ext, word)

def get_stem(filename):
    return os.path.basename(remove_suffix(filename))

def combine_pairs(input_files):
    """ calls files pairs if they are completely the same except
    for one has _1 and the other has _2 returns a list of tuples
    of pairs or singles """
    PAIR_FILE_IDENTIFIERS = ["1", "2"]

    pairs = []
    used = []
    for in_file in input_files:
        if in_file in used:
            continue
        for comp_file in input_files:
            if comp_file in used:
                continue
            s = difflib.SequenceMatcher(a=in_file, b=comp_file)
            blocks = s.get_matching_blocks()
            # length 3 means on match in the middle of the string
            if len(s.get_matching_blocks()) is not 3:
                continue
            if comp_file[blocks[0][2]] in PAIR_FILE_IDENTIFIERS:
                used.append(in_file)
                used.append(comp_file)
                pairs.append([in_file, comp_file])
                break
        if in_file not in used:
            pairs.append([in_file])
            used.append(in_file)

    return pairs

def freeze_files(files, directory):
    """ freezes a list of files, copying them to directory and
    setting them to be read only """
    new_files = [os.path.join(directory, x) for x in
                 map(os.path.basename, files)]
    [os.chmod(x, stat.S_IREAD | stat.S_IRGRP) for x in new_files]
    return new_files

# XXX I'm not 100% sure if this actually works, I forget use with caution
def memoize_outfile_to_dir(res_dir, ext):
    """Creates outfile from input file and ext, running if outfile not present.

    This requires a standard function usage. The first arg, or kwarg 'in_file', needs
    to be the input file that is being processed. The output name is created with the
    provided ext relative to the input. The function is only run if the created
    out_file is not present.
    Memoize_outfile was lifted from bcbio.utils (Brad Chapman), added the
    writing to a directory portion
    """
    utils.safe_makedir(res_dir)
    def decor(f):
        @functools.wraps(f)
        def wrapper(*args, **kwargs):
            if len(args) > 0:
                in_file = args[0]
            else:
                in_file = kwargs['in_file']
            out_fname = "%s%s" % (os.path.splitext(in_file)[0], ext)
            out_file = os.path.join(res_dir, out_fname)
            if not os.path.exists(out_file) or os.path.getsize(out_file) == 0:
                kwargs['out_file'] = out_file
                f(*args, **kwargs)
            return out_file
        return wrapper
    return decor

def prepare_ref_file(ref, config):
    """Get a reference file, either by URL or locally.
    Lifted from Brad Chapman
    """
    url = ref.get("url", None)
    if url:
        ref_file = _download_ref(url, config["dir"]["ref"])
    else:
        ref_file = ref.get("file", None)
    assert ref_file is not None and os.path.exists(ref_file), ref_file
    return ref_file

def _download_ref(url, ref_dir):
    #Lifted from Brad Chapman
    dl_file = os.path.basename(url)
    ref_file = None
    for supported_ext, extract_cmd in [(".gz", "gunzip"),
                                       (".tgz", ("tar", "zxvf"))]:
        if dl_file.endswith(supported_ext):
            ref_file = os.path.join(ref_dir, dl_file[:-len(supported_ext)])
            break
    assert ref_file is not None, url
    if not os.path.exists(ref_file):
        with utils.chdir(ref_dir):
            cl = ["wget", url]
            subprocess.check_call(cl)
            cl = list(flatten([extract_cmd, dl_file]))
            subprocess.check_call(cl)
    return ref_file

def build_results_dir(stage_config, config):
    outdir = os.path.join(config["dir"]["results"],
                          stage_config["name"])
    return outdir

def transform_infile(filename, stage, delim="."):
    return replace_suffix(filename, stage.get("name", "transformed"), delim)

def filter_infile(filename, stage, delim="."):
    return append_stem(filename, stage.get("name", "filtered"), delim)

def remove_suffix(filename):
    filename, extension = os.path.splitext(filename)
    return filename

def append_stem(filename, word, delim="."):
    """ returns a filename with word appended to the stem
    example: /path/to/test.run.sam -> /path/to/test.run.sorted.sam """
    dirname = os.path.dirname(filename)
    filename = os.path.basename(filename)
    fsplit = filename.split(delim)
    fsplit.insert(len(fsplit) - 1, word)
    return os.path.join(dirname, delim.join(fsplit))

def replace_suffix(filename, suffix, delim="."):
    """ returns a filename with the suffix replaced
    with a new suffix
    exampoe: /path/to/test.run.sam -> /path/to/test.run.fa"""
    dirname = os.path.dirname(filename)
    filename = os.path.basename(filename)
    fsplit = filename.split(delim)[:-1]
    fsplit.append(suffix)
    return os.path.join(dirname, delim.join(fsplit))

def flatten_options(config):
    """ returns a list of pairs of options from a config file; returns
    an empty list if no options are found """
    return(flatten(config.get("options", {}.items())))

def flatten(l):
    for el in l:
        if isinstance(el, collections.Iterable) and not isinstance(el, basestring):
            for sub in flatten(el):
                yield sub
        else:
            yield el

def rfind_key(d, item):
    for k,v in d.items():
        if k == item:
            return (k, v)
        elif isinstance(v, dict):
            rfind_key(v, item)

def rfind_value(d, item):
    for k, v in d.items():
        if v == item:
            return (k, v)
        elif isinstance(v, dict):
            rfind_value(v, item)

def validate_config(config, data=True):
    # add the full path to any executables to avoid confusion later
    for k, v in config["program"].items():
        config["program"][k] = _add_full_path(v)

    if not all([os.path.exists(x) for x in config["program"].values()]):
        print("Cannot find all of the files or directories listed "
              "in 'program'.")
        exit(-1)

    if not all(x in get_stages(config) for x in config["run"]):
        print("Cannot find all of the stages listed to run.")
        exit(-1)

    if data and not all([os.path.exists(x) for x in config["input"]]):
        print("All of the input files do not exist.")
        exit(-1)

    return(config)

def get_stages(config):
    return config["stage"].keys()

def _add_full_path(program):
    """ add the full path to any executables """
    new_path = which(program)
    if new_path:
        program_config = new_path
    return program_config

def find_one_if(predicate, coll):
    """ returns a single entry in a collection if it contains exactly
    one true instance of a predicate and None otherwise """
    results = filter(predicate, coll)
    if len(results) != 1:
        return None
    else:
        print results
        return results[0]

def load_yaml(filename):
    with open(filename) as f:
        body = yaml.load(f)
        return body

def which(program):
    """ returns the path to an executable"""

    def is_exe(fpath):
        return os.path.isfile(fpath) and os.access(fpath, os.X_OK)

    fpath, fname = os.path.split(program)
    if fpath:
        if is_exe(program):
            return program
    else:
        for path in os.environ["PATH"].split(os.pathsep):
            exe_file = os.path.join(path, program)
            if is_exe(exe_file):
                return exe_file

    return None

"""
simple file tracking to keep track of the transformations a file has made
just playing with it for now
"""

class FileWithHistory(object):
    """ this tracks transformations of an input file so we can get the
    metadata associated with it at any point """

    def __init__(self, in_file, **kwargs):
        self._name_list = [in_file]
        [setattr(self, key, value) for key, value in kwargs.items()]

    def curr_file(self):
        return self._name_list[len(self._name_list)]

    def add_file(self, in_file):
        self._name_list.append(in_file)

    def was_file(self, in_file):
        return in_file in self._name_list

def update_tracker(tracker, curr_files):
    new_tracker = [track.add_file(new_file) for (track, new_file)
                   in zip(tracker, curr_files)]
    return new_tracker

def curr_files(tracker):
    return [x.curr_file() for x in tracker]

def dict_to_vectors(d):
    """ convert a dictionary to two vectors. returns a tuple of
    (key, value). For example: {"a": [1, 2, 3], "b": [4, 5]} returns
    (["a" "a" "a" "b" "b"], [1 2 3 4 5])

    """

    ks = []
    vs = []
    for k, v in d.items():
        ks += list(repeat(k, len(v)))
        vs += v
    return (ks, vs)

class dotdict(dict):
    """
    access dictioanry items via dot syntax
    via: http://parand.com/say/index.php/2008/10/24/python-dot-notation-dictionary-access/:
    """

    def __getattr__(self, attr):
        return self.get(attr, None)
    __setattr__ = dict.__setitem__
    __delattr__ = dict.__delitem__

def nested_lookup(d, t):
    """
    look up if you can get a tuple of values from a nested dictionary,
    each item in the tuple a deeper layer
    """
    return reduce(lambda d, t: d.get(t, {}), t, d)

def get_in(d, t, default=None):
    """
    look up if you can get a tuple of values from a nested dictionary,
    each item in the tuple a deeper layer

    example: get_in({1: {2: 3}}, (1, 2)) -> 3
    example: get_in({1: {2: 3}}, (2, 3)) -> {}
    """
    result = reduce(lambda d, t: d.get(t, {}), t, d)
    if not result:
        return default
    else:
        return result

def is_sequence(arg):
    """
    check if 'arg' is a sequence

    example: arg([]) -> True
    example: arg("lol") -> False

    """
    return (not hasattr(arg, "strip") and
            hasattr(arg, "__getitem__") or
            hasattr(arg, "__iter__"))

def is_pair(arg):
    """
    check if 'arg' is a two-item sequence

    """
    return is_sequence(arg) and len(arg) == 2
<|endoftext|>


<commit_before>from PyQt5.QtWidgets import QMainWindow
from PyQt5.QtWidgets import QGraphicsScene

from PyQt5.QtGui import QPen
from PyQt5.QtGui import QBrush
from PyQt5.QtGui import QColor
from PyQt5.QtCore import QTimer
from PyQt5.QtCore import pyqtSlot
#from PyQt5.QtCore import pyqtSignal

from Ui_MainWindow import Ui_MainWindow

import anthill
#from timer import Timer

class MainWindow(QMainWindow):
    def __init__(self,path=''):
        super(MainWindow, self).__init__()

        self.ui=Ui_MainWindow()
        self.ui.setupUi(self)

        self.refreshTimer=None

        #self.ui.graphics.setOptimizationFlags(
        #    QGraphicsView.DontAdjustForAntialiasing
        #    | QGraphicsView.DontClipPainter
        #    | QGraphicsView.DontSavePainterState)
        #self.ui.graphics.setTransformationAnchor(self.ui.graphics.NoAnchor)
        #self.ui.graphics.setResizeAnchor(self.ui.graphics.NoAnchor)
        #self.on_tab_currentChanged(self.ui.tab.currentIndex())
        #self.showMaximized()

        # world/simulation parameters - hardcoded - test
        self.pixelSize=10

        self.worldWidth=50
        self.worldHeight=50

        self.simulationFramerate=10

        w=self.worldWidth * self.pixelSize
        h=self.worldHeight * self.pixelSize

        self.ui.graphicsView.setScene(QGraphicsScene(0,0,w,h))

        # configure world
        self.world=anthill.World()
        self.world.setDimensions(w,h)
        self.world.setSimulationFramerate(self.simulationFramerate)

        self.maxPheromone=1

    def __exit__(self):
        self.refreshTimer.stop()

    def drawPheromoneMap(self,map):
        t=map.getMapCopy()
        t=list(t)
        s=self.pixelSize
        scene=self.ui.graphicsView.scene()
        pbrush=QBrush(QColor(50,120,100,255))
        ppen=QPen(QBrush(QColor()),0)
        for x,a in enumerate(t):
            if(x>self.worldWidth):
                break
            for y,b in enumerate(a):
                if(y>self.worldHeight):
                    break
                self.maxPheromone=max(self.maxPheromone,b)
                pbrush=QBrush(QColor(50,120,100,abs(255*b/self.maxPheromone)))
                scene.addRect(x*s,y*s,s,s,pen=ppen,brush=pbrush)

    @pyqtSlot()
    def refresh(self):
        scene=self.ui.graphicsView.scene()
        scene.clear()
        #self.graphicsItems={}

        w=self.worldWidth * self.pixelSize
        h=self.worldHeight * self.pixelSize

        bgBrush=QBrush(QColor(50,120,50,255))
        bgPen=QPen(QBrush(QColor()),0)
        scene.addRect(0,0,w,h,pen=bgPen,brush=bgBrush)

        # symulacja tak? Kamil
        # na razie może być, później się zrobi wątek w c++. ms
        self.world.simulationStep()

        s=self.pixelSize

        # draw pheromone map
        pMaps=self.world.getPheromoneMaps()
        pMaps=[pMaps[p] for p in range(len(pMaps))]
        for m in pMaps:
            self.drawPheromoneMap(m)

        # draw ants
        # ants not working as iterable
        ants=self.world.getAnts()
        for ant in range(len(ants)):
            x=ants[ant].getLoc().posX()
            y=ants[ant].getLoc().posY()
            antPen=QPen(QBrush(QColor()),0)
            antBrush=QBrush(QColor(100,100,50,255))
            scene.addEllipse(x*s,y*s,s,s,pen=antPen,brush=antBrush)
        # draw foods
        # like ants
        foods=self.world.getFoods()
        for food in range(len(foods)):
            x=foods[food].getLoc().posX()
            y=foods[food].getLoc().posY()
            foodpen=QPen(QBrush(QColor()),0)
            foodbrush=QBrush(QColor(200,200,200,150))
            #scene.addrect(x-3,y-3,6,6,pen=foodpen,brush=foodbrush)
            scene.addRect(x*s,y*s,s,s,pen=foodpen,brush=foodbrush)
        # draw anthills
        anthills=self.world.getAnthills()
        anthills=[anthills[a] for a in range(len(anthills))]
        for anthill in anthills:
            x=anthill.getLoc().posX()
            y=anthill.getLoc().posX()
            anthillPen=QPen(QBrush(QColor()),0)
            anthillBrush=QBrush(QColor(200,20,20,255))
            scene.addEllipse(x*s,y*s,s,s,pen=anthillPen,brush=anthillBrush)

    def on_startSimulationButton_released(self):
        if self.refreshTimer is None:
            self.refresh()

            self.refreshTimer=QTimer(self)
            self.refreshTimer.timeout.connect(self.refresh)
            self.refreshTimer.start(1000/self.simulationFramerate)

            self.world.startSimulation()

    def on_stopSimulationButton_released(self):
        if self.refreshTimer:
            self.refresh()

            self.refreshTimer.stop()
            self.refreshTimer=None

            self.world.stopSimulation()

    def on_actionExit_triggered(self):
        self.close()

    def on_actionAbout_triggered(self):
        print("Anthill simulator version 0. ")
<commit_msg>better visualisation<commit_after>from PyQt5.QtWidgets import QMainWindow
from PyQt5.QtWidgets import QGraphicsView
from PyQt5.QtWidgets import QGraphicsScene

from PyQt5.QtGui import QPen
from PyQt5.QtGui import QBrush
from PyQt5.QtGui import QColor
from PyQt5.QtCore import QTimer
from PyQt5.QtCore import pyqtSlot
#from PyQt5.QtCore import pyqtSignal

from Ui_MainWindow import Ui_MainWindow

import anthill
#from timer import Timer

class MainWindow(QMainWindow):
    def __init__(self,path=''):
        super(MainWindow, self).__init__()

        self.ui=Ui_MainWindow()
        self.ui.setupUi(self)

        self.refreshTimer=None

        self.ui.graphicsView.setOptimizationFlags(
            QGraphicsView.DontAdjustForAntialiasing
            | QGraphicsView.DontClipPainter
            | QGraphicsView.DontSavePainterState)
        #self.ui.graphics.setTransformationAnchor(self.ui.graphics.NoAnchor)
        #self.ui.graphics.setResizeAnchor(self.ui.graphics.NoAnchor)
        #self.on_tab_currentChanged(self.ui.tab.currentIndex())
        #self.showMaximized()

        # world/simulation parameters - hardcoded - test
        self.pixelSize=10

        self.worldWidth=50
        self.worldHeight=50

        self.simulationFramerate=10

        w=self.worldWidth * self.pixelSize
        h=self.worldHeight * self.pixelSize

        self.ui.graphicsView.setScene(QGraphicsScene(0,0,w,h))

        # configure world
        self.world=anthill.World()
        self.world.setDimensions(w,h)
        self.world.setSimulationFramerate(self.simulationFramerate)

        self.maxPheromone=1

    def __exit__(self):
        self.refreshTimer.stop()

    def drawPheromoneMap(self,map,
            baseRGB=(50,30,100),
            baseRGB2=(10,10,10) ):
        data=map.getMapCopy()
        data=list(data)

        s=self.pixelSize
        scene=self.ui.graphicsView.scene()

        for x,a in enumerate(data):
            if(x>self.worldWidth):
                break
            for y,b in enumerate(a):
                if(y>self.worldHeight):
                    break
                alpha=255*b/self.maxPheromone
                self.maxPheromone=max(self.maxPheromone,b)
                if(alpha<0.001):
                    continue

                pbrush=QBrush(QColor(*baseRGB, alpha))
                ppen=QPen(QBrush(QColor(*baseRGB2, alpha)),0)
                scene.addRect(x*s,y*s,s,s,pen=ppen,brush=pbrush)

    def drawEntities(self,ents,qpen,qbrush,ellipse=False):
        ents=[ents[a] for a in range(len(ents))]

        s=self.pixelSize
        scene=self.ui.graphicsView.scene()

        for ent in ents:
            x=ent.getLoc().posX()
            y=ent.getLoc().posY()
            if(ellipse):
                scene.addEllipse(x*s,y*s,s,s,pen=qpen,brush=qbrush)
            else:
                scene.addRect(x*s,y*s,s,s,pen=qpen,brush=qbrush)

    @pyqtSlot()
    def refresh(self):
        scene=self.ui.graphicsView.scene()
        scene.clear()
        #self.graphicsItems={}

        w=self.worldWidth * self.pixelSize
        h=self.worldHeight * self.pixelSize

        bgBrush=QBrush(QColor(50,120,50,255))
        bgPen=QPen(QBrush(QColor()),0)
        scene.addRect(0,0,w,h,pen=bgPen,brush=bgBrush)

        # symulacja tak? Kamil
        # na razie może być, później się zrobi wątek w c++. ms
        self.world.simulationStep()

        s=self.pixelSize

        # draw pheromone map
        pMaps=self.world.getPheromoneMaps()
        pMaps=[pMaps[p] for p in range(len(pMaps))]
        for i,m in enumerate(pMaps):
            self.drawPheromoneMap(m,baseRGB=(
                    255*(i%2), 20, 200*((i+1)%2)))

        # draw foods
        foods=self.world.getFoods()
        foodPen=QPen(QBrush(QColor()),0)
        foodBrush=QBrush(QColor(200,200,200,150))
        self.drawEntities(foods,foodPen,foodBrush)

        # draw ants
        ants=self.world.getAnts()
        antPen=QPen(QBrush(QColor()),0)
        antBrush=QBrush(QColor(100,100,50,200))
        self.drawEntities(ants,antPen,antBrush,True)

        # draw anthills
        anthills=self.world.getAnthills()
        anthillPen=QPen(QBrush(QColor()),0)
        anthillBrush=QBrush(QColor(200,20,20,150))
        self.drawEntities(anthills,anthillPen,anthillBrush)

    def on_startSimulationButton_released(self):
        if self.refreshTimer is None:
            self.refresh()

            self.refreshTimer=QTimer(self)
            self.refreshTimer.timeout.connect(self.refresh)
            self.refreshTimer.start(1000/self.simulationFramerate)

            self.world.startSimulation()

    def on_stopSimulationButton_released(self):
        if self.refreshTimer:
            self.refresh()

            self.refreshTimer.stop()
            self.refreshTimer=None

            self.world.stopSimulation()

    def on_actionExit_triggered(self):
        self.close()

    def on_actionAbout_triggered(self):
        print("Anthill simulator version 0. ")
<|endoftext|>


<commit_before>"""
tests.pytests.unit.test_crypt
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

Unit tests for salt's crypt module
"""
import pytest
import salt.crypt
import salt.utils.files

def test_get_rsa_pub_key_bad_key(tmp_path):
    """
    get_rsa_pub_key raises InvalidKeyError when encoutering a bad key
    """
    key_path = tmp_path / "key"
    with salt.utils.files.fopen(key_path, "w") as fp:
        fp.write("")
    with pytest.raises(salt.crypt.InvalidKeyError):
        salt.crypt.get_rsa_pub_key(key_path)
<commit_msg>Fix test on older pythons<commit_after>"""
tests.pytests.unit.test_crypt
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

Unit tests for salt's crypt module
"""
import pytest
import salt.crypt
import salt.utils.files

def test_get_rsa_pub_key_bad_key(tmp_path):
    """
    get_rsa_pub_key raises InvalidKeyError when encoutering a bad key
    """
    key_path = str(tmp_path / "key")
    with salt.utils.files.fopen(key_path, "w") as fp:
        fp.write("")
    with pytest.raises(salt.crypt.InvalidKeyError):
        salt.crypt.get_rsa_pub_key(key_path)
<|endoftext|>