<commit_before>"""
Graph node classification using GraphSAGE.
Requires a EPGM graph as input.
This currently is only tested on the CORA dataset.

Example usage:
python epgm-example.py -g ../../tests/resources/data/cora/cora.epgm -l 50 50 -s 20 10 -e 20 -d 0.5 -r 0.01

usage: epgm-example.py [-h] [-c [CHECKPOINT]] [-n BATCH_SIZE] [-e EPOCHS]
                       [-s [NEIGHBOUR_SAMPLES [NEIGHBOUR_SAMPLES ...]]]
                       [-l [LAYER_SIZE [LAYER_SIZE ...]]] [-g GRAPH]
                       [-f FEATURES] [-t TARGET]

optional arguments:
  -h, --help            show this help message and exit
  -c [CHECKPOINT], --checkpoint [CHECKPOINT]
                        Load a saved checkpoint file
  -n BATCH_SIZE, --batch_size BATCH_SIZE
                        Load a save checkpoint file
  -e EPOCHS, --epochs EPOCHS
                        Number of epochs to train for
  -s [NEIGHBOUR_SAMPLES [NEIGHBOUR_SAMPLES ...]], --neighbour_samples [NEIGHBOUR_SAMPLES [NEIGHBOUR_SAMPLES ...]]
                        The number of nodes sampled at each layer
  -l [LAYER_SIZE [LAYER_SIZE ...]], --layer_size [LAYER_SIZE [LAYER_SIZE ...]]
                        The number of hidden features at each layer
  -g GRAPH, --graph GRAPH
                        The graph stored in EPGM format.
  -f FEATURES, --features FEATURES
                        The node features to use, stored as a pickled numpy
                        array.
  -t TARGET, --target TARGET
                        The target node attribute
"""
import os
import argparse
import random
import numpy as np
import networkx as nx
from typing import AnyStr, Any, List, Tuple, Dict, Optional

import keras
from keras import optimizers, losses, layers, metrics
from keras.utils.np_utils import to_categorical

from stellar.data.epgm import EPGM
from stellar.data.node_splitter import NodeSplitter
from stellar.data.explorer import SampledBreadthFirstWalk

from stellar.layer.graphsage import GraphSAGE, MeanAggregator
from stellar.mapper.node_mappers import GraphSAGENodeMapper

def read_epgm_graph(
    graph_file,
    dataset_name=None,
    node_type=None,
    target_attribute=None,
    ignored_attributes=[],
    target_type=None,
    remove_converted_attrs=False,
):
    G_epgm = EPGM(graph_file)
    graphs = G_epgm.G["graphs"]

    # if dataset_name is not given, use the name of the 1st graph head
    if not dataset_name:
        dataset_name = graphs[0]["meta"]["label"]
        print(
            "WARNING: dataset name not specified, using dataset '{}' in the 1st graph head".format(
                dataset_name
            )
        )

    graph_id = None
    for g in graphs:
        if g["meta"]["label"] == dataset_name:
            graph_id = g["id"]

    if node_type is None:
        node_type = G_epgm.node_types(graph_id)[0]

    g_nx = G_epgm.to_nx(graph_id)

    # Find target and predicted attributes from attribute set
    node_attributes = set(G_epgm.node_attributes(graph_id, node_type))
    pred_attr = node_attributes.difference(
        set(ignored_attributes).union([target_attribute])
    )
    converted_attr = pred_attr.union([target_attribute])

    # Index nodes in graph
    for ii, v in enumerate(g_nx.nodes()):
        g_nx.node[v]["id"] = ii

    # Enumerate attributes to give numerical index
    g_nx.pred_map = {a: ii for ii, a in enumerate(pred_attr)}

    # Store feature size in graph [??]
    g_nx.feature_size = len(g_nx.pred_map)

    # How do we map target attributes to numerical values?
    g_nx.target_category_values = None
    if target_type is None:
        target_value_function = lambda x: x

    elif target_type == "categorical":
        g_nx.target_category_values = list(
            set([g_nx.node[n][target_attribute] for n in g_nx.nodes()])
        )
        target_value_function = lambda x: g_nx.target_category_values.index(x)

    elif target_type == "1hot":
        g_nx.target_category_values = list(
            set([g_nx.node[n][target_attribute] for n in g_nx.nodes()])
        )
        target_value_function = lambda x: to_categorical(
            g_nx.target_category_values.index(x), len(g_nx.target_category_values)
        )

    else:
        raise ValueError("Target type '{}' is not supported.".format(target_type))

    for v, vdata in g_nx.nodes(data=True):
        # Decode attributes to a feature array
        attr_array = np.zeros(g_nx.feature_size)
        for attr_name, attr_value in vdata.items():
            col = g_nx.pred_map.get(attr_name)
            if col:
                attr_array[col] = attr_value

        # Replace with feature array
        vdata["feature"] = attr_array

        # Decode target attribute to target array
        vdata["target"] = target_value_function(vdata.get(target_attribute))

        # Remove attributes
        if remove_converted_attrs:
            for attr_name in converted_attr:
                if attr_name in vdata:
                    del vdata[attr_name]

    print(
        "Graph statistics: {} nodes, {} edges".format(
            g_nx.number_of_nodes(), g_nx.number_of_edges()
        )
    )
    return g_nx

class GraphsageSampler:
    def __init__(self, G):
        self.G = G

    def _sample_neighbourhood(self, head_nodes, samples_per_hop):
        if len(samples_per_hop) < 1:
            return [head_nodes]

        num_for_hop, next_samples = samples_per_hop[0], samples_per_hop[1:]
        hop_samples = []
        for hn in head_nodes:
            neighs = list(self.G.neighbors(hn))
            if len(neighs) < num_for_hop:
                neigh_sampled = random.choices(neighs, k=num_for_hop)
            else:
                neigh_sampled = random.sample(neighs, num_for_hop)
            hop_samples.extend(neigh_sampled)

        return [head_nodes] + self._sample_neighbourhood(hop_samples, next_samples)

    def run(self, nodes=[], n=1, n_size=[]):

        return self._sample_neighbourhood(nodes, n_size)

def train(
    G,
    layer_size: List[int],
    num_samples: List[int],
    batch_size: int = 100,
    num_epochs: int = 10,
    learning_rate: float = 0.005,
    dropout: float = 0.0,
):
    # Split head nodes into train/test
    splitter = NodeSplitter()
    graph_nodes = np.array(
        [(v, vdata.get("subject")) for v, vdata in G.nodes(data=True)]
    )
    train_nodes, val_nodes, test_nodes, _ = splitter.train_test_split(
        y=graph_nodes, p=300, test_size=500
    )
    train_ids = [v[0] for v in train_nodes]
    test_ids = list(G.nodes())
    val_ids = [v[0] for v in val_nodes]

    # Sampler chooses random sampled subgraph for each head node
    sampler = SampledBreadthFirstWalk(G)

    # Mapper feeds data from sampled subgraph to GraphSAGE model
    train_mapper = GraphSAGENodeMapper(
        G, train_ids, sampler, batch_size, num_samples, target_id="target", name="train"
    )
    test_mapper = GraphSAGENodeMapper(
        G, test_ids, sampler, batch_size, num_samples, target_id="target", name="test"
    )

    # GraphSAGE model
    model = GraphSAGE(
        output_dims=layer_size,
        n_samples=num_samples,
        input_dim=G.feature_size,
        bias=True,
        dropout=dropout,
    )
    x_inp, x_out = model.default_model(flatten_output=True)

    # Final estimator layer
    prediction = layers.Dense(
        units=len(G.target_category_values), activation="softmax"
    )(x_out)

    # Create Keras model for training
    model = keras.Model(inputs=x_inp, outputs=prediction)
    model.compile(
        optimizer=optimizers.Adam(lr=learning_rate),
        loss=losses.categorical_crossentropy,
        metrics=[metrics.categorical_accuracy],
    )

    # Train model
    history = model.fit_generator(
        train_mapper, epochs=num_epochs, verbose=2, shuffle=True
    )

    # Evaluate and print metrics
    test_metrics = model.evaluate_generator(test_mapper)

    print("\nTest Evaluation:")
    for name, val in zip(model.metrics_names, test_metrics):
        print("\t{}: {:0.4f}".format(name, val))

    # Save model
    str_numsamp = "_".join([str(x) for x in num_samples])
    str_layer = "_".join([str(x) for x in layer_size])
    model.save(
        "graphsage_n{}_l{}_d{}_i{}.h5".format(
            str_numsamp, str_layer, dropout, G.feature_size
        )
    )

def test(G, model_file: AnyStr, batch_size: int):
    model = keras.models.load_model(
        model_file, custom_objects={"MeanAggregator": MeanAggregator}
    )

    # Get required input shapes from model
    num_samples = [
        int(model.input_shape[ii + 1][1] / model.input_shape[ii][1])
        for ii in range(len(model.input_shape) - 1)
    ]

    # Split head nodes into train/test
    splitter = NodeSplitter()
    all_ids = list(G.nodes())

    # Sampler chooses random sampled subgraph for each head node
    sampler = SampledBreadthFirstWalk(G)

    # Mapper feeds data from sampled subgraph to GraphSAGE model
    test_mapper = GraphSAGENodeMapper(
        G, all_ids, sampler, batch_size, num_samples, target_id="target", name="test"
    )

    # Evaluate and print metrics
    test_metrics = model.evaluate_generator(test_mapper)

    print("\nTest Evaluation:")
    for name, val in zip(model.metrics_names, test_metrics):
        print("\t{}: {:0.4f}".format(name, val))

if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Graph node classification using GraphSAGE"
    )
    parser.add_argument(
        "-c",
        "--checkpoint",
        nargs="?",
        type=str,
        default=None,
        help="Load a saved checkpoint file",
    )
    parser.add_argument(
        "-n", "--batch_size", type=int, default=500, help="Load a save checkpoint file"
    )
    parser.add_argument(
        "-e", "--epochs", type=int, default=10, help="Number of epochs to train for"
    )
    parser.add_argument(
        "-d", "--dropout", type=float, default=0.0, help="Dropout for GraphSAGE model"
    )
    parser.add_argument(
        "-r",
        "--learningrate",
        type=float,
        default=0.0005,
        help="Learning rate for training model",
    )
    parser.add_argument(
        "-s",
        "--neighbour_samples",
        type=int,
        nargs="*",
        default=[30, 10],
        help="The number of nodes sampled at each layer",
    )
    parser.add_argument(
        "-l",
        "--layer_size",
        type=int,
        nargs="*",
        default=[50, 50],
        help="The number of hidden features at each layer",
    )
    parser.add_argument(
        "-g", "--graph", type=str, default=None, help="The graph stored in EPGM format."
    )
    parser.add_argument(
        "-f",
        "--features",
        type=str,
        default=None,
        help="The node features to use, stored as a pickled numpy array.",
    )
    parser.add_argument(
        "-t",
        "--target",
        type=str,
        default="subject",
        help="The target node attribute (categorical)",
    )
    args, cmdline_args = parser.parse_known_args()

    graph_loc = os.path.expanduser(args.graph)
    G = read_epgm_graph(
        graph_loc,
        target_attribute=args.target,
        target_type="1hot",
        remove_converted_attrs=False,
    )

    if args.checkpoint is None:
        train(
            G,
            args.layer_size,
            args.neighbour_samples,
            args.batch_size,
            args.epochs,
            args.learningrate,
            args.dropout,
        )
    else:
        test(G, args.checkpoint, args.batch_size)
<commit_msg>Removed old sampler class from demo<commit_after>"""
Graph node classification using GraphSAGE.
Requires a EPGM graph as input.
This currently is only tested on the CORA dataset.

Example usage:
python epgm-example.py -g ../../tests/resources/data/cora/cora.epgm -l 50 50 -s 20 10 -e 20 -d 0.5 -r 0.01

usage: epgm-example.py [-h] [-c [CHECKPOINT]] [-n BATCH_SIZE] [-e EPOCHS]
                       [-s [NEIGHBOUR_SAMPLES [NEIGHBOUR_SAMPLES ...]]]
                       [-l [LAYER_SIZE [LAYER_SIZE ...]]] [-g GRAPH]
                       [-f FEATURES] [-t TARGET]

optional arguments:
  -h, --help            show this help message and exit
  -c [CHECKPOINT], --checkpoint [CHECKPOINT]
                        Load a saved checkpoint file
  -n BATCH_SIZE, --batch_size BATCH_SIZE
                        Load a save checkpoint file
  -e EPOCHS, --epochs EPOCHS
                        Number of epochs to train for
  -s [NEIGHBOUR_SAMPLES [NEIGHBOUR_SAMPLES ...]], --neighbour_samples [NEIGHBOUR_SAMPLES [NEIGHBOUR_SAMPLES ...]]
                        The number of nodes sampled at each layer
  -l [LAYER_SIZE [LAYER_SIZE ...]], --layer_size [LAYER_SIZE [LAYER_SIZE ...]]
                        The number of hidden features at each layer
  -g GRAPH, --graph GRAPH
                        The graph stored in EPGM format.
  -f FEATURES, --features FEATURES
                        The node features to use, stored as a pickled numpy
                        array.
  -t TARGET, --target TARGET
                        The target node attribute
"""
import os
import argparse
import random
import numpy as np
import networkx as nx
from typing import AnyStr, Any, List, Tuple, Dict, Optional

import keras
from keras import optimizers, losses, layers, metrics
from keras.utils.np_utils import to_categorical

from stellar.data.epgm import EPGM
from stellar.data.node_splitter import NodeSplitter
from stellar.data.explorer import SampledBreadthFirstWalk

from stellar.layer.graphsage import GraphSAGE, MeanAggregator
from stellar.mapper.node_mappers import GraphSAGENodeMapper

def read_epgm_graph(
    graph_file,
    dataset_name=None,
    node_type=None,
    target_attribute=None,
    ignored_attributes=[],
    target_type=None,
    remove_converted_attrs=False,
):
    G_epgm = EPGM(graph_file)
    graphs = G_epgm.G["graphs"]

    # if dataset_name is not given, use the name of the 1st graph head
    if not dataset_name:
        dataset_name = graphs[0]["meta"]["label"]
        print(
            "WARNING: dataset name not specified, using dataset '{}' in the 1st graph head".format(
                dataset_name
            )
        )

    graph_id = None
    for g in graphs:
        if g["meta"]["label"] == dataset_name:
            graph_id = g["id"]

    if node_type is None:
        node_type = G_epgm.node_types(graph_id)[0]

    g_nx = G_epgm.to_nx(graph_id)

    # Find target and predicted attributes from attribute set
    node_attributes = set(G_epgm.node_attributes(graph_id, node_type))
    pred_attr = node_attributes.difference(
        set(ignored_attributes).union([target_attribute])
    )
    converted_attr = pred_attr.union([target_attribute])

    # Index nodes in graph
    for ii, v in enumerate(g_nx.nodes()):
        g_nx.node[v]["id"] = ii

    # Enumerate attributes to give numerical index
    g_nx.pred_map = {a: ii for ii, a in enumerate(pred_attr)}

    # Store feature size in graph [??]
    g_nx.feature_size = len(g_nx.pred_map)

    # How do we map target attributes to numerical values?
    g_nx.target_category_values = None
    if target_type is None:
        target_value_function = lambda x: x

    elif target_type == "categorical":
        g_nx.target_category_values = list(
            set([g_nx.node[n][target_attribute] for n in g_nx.nodes()])
        )
        target_value_function = lambda x: g_nx.target_category_values.index(x)

    elif target_type == "1hot":
        g_nx.target_category_values = list(
            set([g_nx.node[n][target_attribute] for n in g_nx.nodes()])
        )
        target_value_function = lambda x: to_categorical(
            g_nx.target_category_values.index(x), len(g_nx.target_category_values)
        )

    else:
        raise ValueError("Target type '{}' is not supported.".format(target_type))

    for v, vdata in g_nx.nodes(data=True):
        # Decode attributes to a feature array
        attr_array = np.zeros(g_nx.feature_size)
        for attr_name, attr_value in vdata.items():
            col = g_nx.pred_map.get(attr_name)
            if col:
                attr_array[col] = attr_value

        # Replace with feature array
        vdata["feature"] = attr_array

        # Decode target attribute to target array
        vdata["target"] = target_value_function(vdata.get(target_attribute))

        # Remove attributes
        if remove_converted_attrs:
            for attr_name in converted_attr:
                if attr_name in vdata:
                    del vdata[attr_name]

    print(
        "Graph statistics: {} nodes, {} edges".format(
            g_nx.number_of_nodes(), g_nx.number_of_edges()
        )
    )
    return g_nx

def train(
    G,
    layer_size: List[int],
    num_samples: List[int],
    batch_size: int = 100,
    num_epochs: int = 10,
    learning_rate: float = 0.005,
    dropout: float = 0.0,
):
    # Split head nodes into train/test
    splitter = NodeSplitter()
    graph_nodes = np.array(
        [(v, vdata.get("subject")) for v, vdata in G.nodes(data=True)]
    )
    train_nodes, val_nodes, test_nodes, _ = splitter.train_test_split(
        y=graph_nodes, p=300, test_size=500
    )
    train_ids = [v[0] for v in train_nodes]
    test_ids = list(G.nodes())
    val_ids = [v[0] for v in val_nodes]

    # Sampler chooses random sampled subgraph for each head node
    sampler = SampledBreadthFirstWalk(G)

    # Mapper feeds data from sampled subgraph to GraphSAGE model
    train_mapper = GraphSAGENodeMapper(
        G, train_ids, sampler, batch_size, num_samples, target_id="target", name="train"
    )
    test_mapper = GraphSAGENodeMapper(
        G, test_ids, sampler, batch_size, num_samples, target_id="target", name="test"
    )

    # GraphSAGE model
    model = GraphSAGE(
        output_dims=layer_size,
        n_samples=num_samples,
        input_dim=G.feature_size,
        bias=True,
        dropout=dropout,
    )
    x_inp, x_out = model.default_model(flatten_output=True)

    # Final estimator layer
    prediction = layers.Dense(
        units=len(G.target_category_values), activation="softmax"
    )(x_out)

    # Create Keras model for training
    model = keras.Model(inputs=x_inp, outputs=prediction)
    model.compile(
        optimizer=optimizers.Adam(lr=learning_rate),
        loss=losses.categorical_crossentropy,
        metrics=[metrics.categorical_accuracy],
    )

    # Train model
    history = model.fit_generator(
        train_mapper, epochs=num_epochs, verbose=2, shuffle=True
    )

    # Evaluate and print metrics
    test_metrics = model.evaluate_generator(test_mapper)

    print("\nTest Evaluation:")
    for name, val in zip(model.metrics_names, test_metrics):
        print("\t{}: {:0.4f}".format(name, val))

    # Save model
    str_numsamp = "_".join([str(x) for x in num_samples])
    str_layer = "_".join([str(x) for x in layer_size])
    model.save(
        "graphsage_n{}_l{}_d{}_i{}.h5".format(
            str_numsamp, str_layer, dropout, G.feature_size
        )
    )

def test(G, model_file: AnyStr, batch_size: int):
    model = keras.models.load_model(
        model_file, custom_objects={"MeanAggregator": MeanAggregator}
    )

    # Get required input shapes from model
    num_samples = [
        int(model.input_shape[ii + 1][1] / model.input_shape[ii][1])
        for ii in range(len(model.input_shape) - 1)
    ]

    # Split head nodes into train/test
    splitter = NodeSplitter()
    all_ids = list(G.nodes())

    # Sampler chooses random sampled subgraph for each head node
    sampler = SampledBreadthFirstWalk(G)

    # Mapper feeds data from sampled subgraph to GraphSAGE model
    test_mapper = GraphSAGENodeMapper(
        G, all_ids, sampler, batch_size, num_samples, target_id="target", name="test"
    )

    # Evaluate and print metrics
    test_metrics = model.evaluate_generator(test_mapper)

    print("\nTest Evaluation:")
    for name, val in zip(model.metrics_names, test_metrics):
        print("\t{}: {:0.4f}".format(name, val))

if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Graph node classification using GraphSAGE"
    )
    parser.add_argument(
        "-c",
        "--checkpoint",
        nargs="?",
        type=str,
        default=None,
        help="Load a saved checkpoint file",
    )
    parser.add_argument(
        "-n", "--batch_size", type=int, default=500, help="Load a save checkpoint file"
    )
    parser.add_argument(
        "-e", "--epochs", type=int, default=10, help="Number of epochs to train for"
    )
    parser.add_argument(
        "-d", "--dropout", type=float, default=0.0, help="Dropout for GraphSAGE model"
    )
    parser.add_argument(
        "-r",
        "--learningrate",
        type=float,
        default=0.0005,
        help="Learning rate for training model",
    )
    parser.add_argument(
        "-s",
        "--neighbour_samples",
        type=int,
        nargs="*",
        default=[30, 10],
        help="The number of nodes sampled at each layer",
    )
    parser.add_argument(
        "-l",
        "--layer_size",
        type=int,
        nargs="*",
        default=[50, 50],
        help="The number of hidden features at each layer",
    )
    parser.add_argument(
        "-g", "--graph", type=str, default=None, help="The graph stored in EPGM format."
    )
    parser.add_argument(
        "-f",
        "--features",
        type=str,
        default=None,
        help="The node features to use, stored as a pickled numpy array.",
    )
    parser.add_argument(
        "-t",
        "--target",
        type=str,
        default="subject",
        help="The target node attribute (categorical)",
    )
    args, cmdline_args = parser.parse_known_args()

    graph_loc = os.path.expanduser(args.graph)
    G = read_epgm_graph(
        graph_loc,
        target_attribute=args.target,
        target_type="1hot",
        remove_converted_attrs=False,
    )

    if args.checkpoint is None:
        train(
            G,
            args.layer_size,
            args.neighbour_samples,
            args.batch_size,
            args.epochs,
            args.learningrate,
            args.dropout,
        )
    else:
        test(G, args.checkpoint, args.batch_size)
<|endoftext|>


<commit_before>#!/usr/bin/env python
from __future__ import print_function, division
import os, sys, logging, argparse, subprocess, shutil, json

TEMP_SCRIPT_DIRECTORY = os.path.dirname(os.path.abspath(__file__)) + "/"
sys.path.append(os.path.abspath(TEMP_SCRIPT_DIRECTORY + '../Serotyper/'))
sys.path.append(os.path.abspath(TEMP_SCRIPT_DIRECTORY + '../'))

from sharedmethods import *
from ecvalidatingfiles import *

SCRIPT_DIRECTORY = os.path.dirname(os.path.abspath(__file__))
GENOMES = {}
FILENAMES = {}

def parseCommandLine():
    """
    Initalizing the main commands of the command line for the project.
    - input: refers to the location of the file(s) that will be processed
    - out: refers to the output of the program. Default is STDOUT
    - pi: refers to the percentage of identity wanted. Default is 90%
    - pl: refers to the percentage of length wanted. Default is 90%.
    - csv: if the user wants a csv copy of the results
    - min: the minimum number of genomes containing a certain gene

    :return parser.parse_args(): Data from the commands.
    """

    parser = argparse.ArgumentParser()

    parser.add_argument("-in", "--input", help="Location of new file(s). Can be a single file or a directory", required=True)
    parser.add_argument("-out", "--output", type=argparse.FileType('w'), help="Output of the program. Default is STDOUT.", default=sys.stdout)
    parser.add_argument("-pi", "--percentIdentity", type=int, help="Percentage of identity wanted to use against the database. From 0 to 100, default is 90%.", default=90, choices=range(0,100))
    parser.add_argument("-pl", "--percentLength", type=int, help="Percentage of length wanted to use against the database. From 0 to 100, default is 90%.", default=90, choices=range(0,100))
    parser.add_argument("-csv", type=int, help="If set to 1, the results will be sent to a .csv file in the temp/Results folder. Options are 0 and 1, default is 1.", default=1, choices=[0,1])
    parser.add_argument("-min", "--minGenomes", type=int, help="Minimum number of genomes threshold for a virulence factor. Default is 1.", default=0)

    return parser.parse_args()

def initializeDB():
    """
    Generating the database if it was not already generated. The database can be found in the tmp folder.

    :return int 0 or 1: 0 being that the database was created successfully (or already existed).
    """

    REL_DIR = os.path.join(SCRIPT_DIRECTORY, '../../temp/databases/VF_Database/')

    db_in = os.path.join(SCRIPT_DIRECTORY, '../../Data/repaired_ecoli_vfs.ffn')
    db_out = os.path.join(REL_DIR, 'VirulenceFactorsDB')

    makeblastdb = ("makeblastdb", "-in", db_in, "-dbtype", "nucl",
                   "-title", "VirulenceFactorsDB", "-out", db_out)

    if os.path.isfile(REL_DIR + 'VirulenceFactorsDB.nin'):
        logging.info('Database already exists.')
        retcode = 0

    else:
        logging.info('Generating the database.')

        retcode = subprocess.call(makeblastdb)

    return retcode

def searchDB(genomesList):
    """
    Querying the database with a single file combining all the genomes.

    :param genomesList: files that will be used against the database
    :return new_filename: .xml file containing the results from querying the database.
    """

    REL_DIR = os.path.join(SCRIPT_DIRECTORY, '../../temp/xml/')

    if len(genomesList) > 1:
        combined_genomes = os.path.join(SCRIPT_DIRECTORY,
                                        '../../temp/Uploads/combined_genomesVF.fasta')

        #Copying the content of every fasta file into one file to simplify the database search
        with open(combined_genomes, 'wb') as outfile:
            for file in genomesList:
                with open(file, 'rb') as fastafile:
                    shutil.copyfileobj(fastafile, outfile,1024*1024*10)

        new_filename = os.path.abspath(os.path.join(REL_DIR, 'combined_genomesVF.xml'))

    else:
        filename = os.path.basename(genomesList[0])
        filename, ext = os.path.splitext(filename)
        combined_genomes = genomesList[0]
        new_filename = os.path.abspath(os.path.join(REL_DIR, filename + '.xml'))

    #Querying the database

    db_path = os.path.join(REL_DIR, '../databases/VF_Database/VirulenceFactorsDB')

    blastn_cline = NcbiblastnCommandline(cmd="blastn", query=combined_genomes,
                                         db=db_path, outfmt=5, out=new_filename)
    stdout, stderr = blastn_cline()

    logging.info("Searched the database.")
    return new_filename

def parseFile(result_file, perc_len, perc_id):
    """
    Going through each record in the result file and filter them by the percentages provided.

    :param result_file: file containing all the results from querying the database.
    :param perc_len: percentage length cutoff
    :param perc_id: percentage identity cutoff
    """

    global GENOMES
    global FILENAMES
    logging.info("Parsing results from " + str(result_file))

    result_handle = open(result_file)
    blast_records = NCBIXML.parse(result_handle)

    perc_len = perc_len / 100
    perc_id = perc_id / 100

    for blast_record in blast_records:
        filename = FILENAMES[blast_record.query]
        genome_name = getGenomeName(blast_record.query, filename)
        alignmentsDict = {}
        if genome_name in GENOMES:
            alignmentsDict = dict(GENOMES[genome_name])

        contig_accession = blast_record.query.split(' ')[0]
        # intend to make it: {filename{contig_accession{'START','STOP','ORIENTATION','GENE_NAME'}}}
        alignmentsDict[contig_accession] = {}

        #Find the gene name + compare with the identity and length cutoffs.
        for alignment in blast_record.alignments:
            match = re.search('(\s\(.*\)\s)', alignment.title)
            match = match.group()
            match = match.split('(')[1]
            match = match.split(')')[0]
            align_title = str(match)

            alignmentsDict[contig_accession]['GENE_NAME'] = align_title

            hsp = alignment.hsps[0]

            tmp_perc_len = abs(1 - (1 - hsp.positives) / alignment.length)

            tmp_perc_id = hsp.positives / hsp.align_length

            if tmp_perc_len > perc_len and tmp_perc_id > perc_id and align_title not in alignmentsDict.keys():
                alignmentsDict[align_title] = 1

                # check if the subject (what we're looking for) was found in forward or reverse
                # this means found in forward
                if hsp.sbjct_start < hsp.sbjct_end:
                    alignmentsDict[contig_accession]['ORIENTATION'] = '+'
                else:
                    alignmentsDict[contig_accession]['ORIENTATION'] = '-'

                alignmentsDict[contig_accession]['START']=hsp.query_start
                alignmentsDict[contig_accession]['STOP']=hsp.query_end

            else:
                if align_title not in alignmentsDict.keys() or alignmentsDict[align_title]!=1:
                    alignmentsDict[align_title] = 0

            GENOMES[genome_name] = alignmentsDict

if __name__=='__main__':

    log_path = os.path.join(SCRIPT_DIRECTORY, 'virulencefactors.log')
    logging.basicConfig(filename=log_path,level=logging.INFO)

    args = parseCommandLine()
    createDirs()

    logging.info('Starting Virulence Factors tool')

    roughGenomesList = getFilesList(args.input)
    genomesList = checkFiles(roughGenomesList)

    GENOMES = getGENOMES()
    FILENAMES = getFILENAMES()
    clearGlobalDicts()

    if isinstance(genomesList, list):

        if initializeDB() == 0:

            results_file = searchDB(genomesList)
            parseFile(results_file, args.percentLength, args.percentIdentity)
            resultsDict = filterGenes(GENOMES, args.minGenomes)

            if args.csv == 1:
                toTSV(resultsDict, 'VF_Results')

            json.dump(resultsDict, sys.stdout)
            logging.info('Program ended successfully.')

    else:
        print(genomesList)
<commit_msg>VF returns okay now in new dict format<commit_after>#!/usr/bin/env python
from __future__ import print_function, division
import os, sys, logging, argparse, subprocess, shutil, json

TEMP_SCRIPT_DIRECTORY = os.path.dirname(os.path.abspath(__file__)) + "/"
sys.path.append(os.path.abspath(TEMP_SCRIPT_DIRECTORY + '../Serotyper/'))
sys.path.append(os.path.abspath(TEMP_SCRIPT_DIRECTORY + '../'))

from sharedmethods import *
from ecvalidatingfiles import *

SCRIPT_DIRECTORY = os.path.dirname(os.path.abspath(__file__))
GENOMES = {}
FILENAMES = {}

def parseCommandLine():
    """
    Initalizing the main commands of the command line for the project.
    - input: refers to the location of the file(s) that will be processed
    - out: refers to the output of the program. Default is STDOUT
    - pi: refers to the percentage of identity wanted. Default is 90%
    - pl: refers to the percentage of length wanted. Default is 90%.
    - csv: if the user wants a csv copy of the results
    - min: the minimum number of genomes containing a certain gene

    :return parser.parse_args(): Data from the commands.
    """

    parser = argparse.ArgumentParser()

    parser.add_argument("-in", "--input", help="Location of new file(s). Can be a single file or a directory", required=True)
    parser.add_argument("-out", "--output", type=argparse.FileType('w'), help="Output of the program. Default is STDOUT.", default=sys.stdout)
    parser.add_argument("-pi", "--percentIdentity", type=int, help="Percentage of identity wanted to use against the database. From 0 to 100, default is 90%.", default=90, choices=range(0,100))
    parser.add_argument("-pl", "--percentLength", type=int, help="Percentage of length wanted to use against the database. From 0 to 100, default is 90%.", default=90, choices=range(0,100))
    parser.add_argument("-csv", type=int, help="If set to 1, the results will be sent to a .csv file in the temp/Results folder. Options are 0 and 1, default is 1.", default=1, choices=[0,1])
    parser.add_argument("-min", "--minGenomes", type=int, help="Minimum number of genomes threshold for a virulence factor. Default is 1.", default=0)

    return parser.parse_args()

def initializeDB():
    """
    Generating the database if it was not already generated. The database can be found in the tmp folder.

    :return int 0 or 1: 0 being that the database was created successfully (or already existed).
    """

    REL_DIR = os.path.join(SCRIPT_DIRECTORY, '../../temp/databases/VF_Database/')

    db_in = os.path.join(SCRIPT_DIRECTORY, '../../Data/repaired_ecoli_vfs.ffn')
    db_out = os.path.join(REL_DIR, 'VirulenceFactorsDB')

    makeblastdb = ("makeblastdb", "-in", db_in, "-dbtype", "nucl",
                   "-title", "VirulenceFactorsDB", "-out", db_out)

    if os.path.isfile(REL_DIR + 'VirulenceFactorsDB.nin'):
        logging.info('Database already exists.')
        retcode = 0

    else:
        logging.info('Generating the database.')

        retcode = subprocess.call(makeblastdb)

    return retcode

def searchDB(genomesList):
    """
    Querying the database with a single file combining all the genomes.

    :param genomesList: files that will be used against the database
    :return new_filename: .xml file containing the results from querying the database.
    """

    REL_DIR = os.path.join(SCRIPT_DIRECTORY, '../../temp/xml/')

    if len(genomesList) > 1:
        combined_genomes = os.path.join(SCRIPT_DIRECTORY,
                                        '../../temp/Uploads/combined_genomesVF.fasta')

        #Copying the content of every fasta file into one file to simplify the database search
        with open(combined_genomes, 'wb') as outfile:
            for file in genomesList:
                with open(file, 'rb') as fastafile:
                    shutil.copyfileobj(fastafile, outfile,1024*1024*10)

        new_filename = os.path.abspath(os.path.join(REL_DIR, 'combined_genomesVF.xml'))

    else:
        filename = os.path.basename(genomesList[0])
        filename, ext = os.path.splitext(filename)
        combined_genomes = genomesList[0]
        new_filename = os.path.abspath(os.path.join(REL_DIR, filename + '.xml'))

    #Querying the database

    db_path = os.path.join(REL_DIR, '../databases/VF_Database/VirulenceFactorsDB')

    blastn_cline = NcbiblastnCommandline(cmd="blastn", query=combined_genomes,
                                         db=db_path, outfmt=5, out=new_filename)
    stdout, stderr = blastn_cline()

    logging.info("Searched the database.")
    return new_filename

def parseFile(result_file, perc_len, perc_id):
    """
    Going through each record in the result file and filter them by the percentages provided.

    :param result_file: file containing all the results from querying the database.
    :param perc_len: percentage length cutoff
    :param perc_id: percentage identity cutoff
    """

    global GENOMES
    global FILENAMES
    logging.info("Parsing results from " + str(result_file))

    result_handle = open(result_file)
    blast_records = NCBIXML.parse(result_handle)

    perc_len = perc_len / 100
    perc_id = perc_id / 100

    for blast_record in blast_records:
        filename = FILENAMES[blast_record.query]
        genome_name = getGenomeName(blast_record.query, filename)
        alignmentsDict = {}
        if genome_name in GENOMES:
            alignmentsDict = dict(GENOMES[genome_name])

        #Find the gene name + compare with the identity and length cutoffs.
        for alignment in blast_record.alignments:
            match = re.search('(\s\(.*\)\s)', alignment.title)
            match = match.group()
            match = match.split('(')[1]
            match = match.split(')')[0]
            align_title = str(match)

            hsp = alignment.hsps[0]

            tmp_perc_len = abs(1 - (1 - hsp.positives) / alignment.length)

            tmp_perc_id = hsp.positives / hsp.align_length

            if tmp_perc_len > perc_len and tmp_perc_id > perc_id and align_title not in alignmentsDict.keys():
                contig_accession = blast_record.query.split(' ')[0]
                # intend to make it: {filename{contig_accession{'START','STOP','ORIENTATION','GENE_NAME'}}}
                alignmentsDict[contig_accession] = {}
                alignmentsDict[contig_accession]['GENE_NAME'] = align_title
                # check if the subject (what we're looking for) was found in forward or reverse
                # this means found in forward
                if hsp.sbjct_start < hsp.sbjct_end:
                    alignmentsDict[contig_accession]['ORIENTATION'] = '+'
                else:
                    alignmentsDict[contig_accession]['ORIENTATION'] = '-'

                alignmentsDict[contig_accession]['START']=hsp.query_start
                alignmentsDict[contig_accession]['STOP']=hsp.query_end

            GENOMES[genome_name] = alignmentsDict

if __name__=='__main__':

    log_path = os.path.join(SCRIPT_DIRECTORY, 'virulencefactors.log')
    logging.basicConfig(filename=log_path,level=logging.INFO)

    args = parseCommandLine()
    createDirs()

    logging.info('Starting Virulence Factors tool')

    roughGenomesList = getFilesList(args.input)
    genomesList = checkFiles(roughGenomesList)

    GENOMES = getGENOMES()
    FILENAMES = getFILENAMES()
    clearGlobalDicts()

    if isinstance(genomesList, list):

        if initializeDB() == 0:

            results_file = searchDB(genomesList)
            parseFile(results_file, args.percentLength, args.percentIdentity)
            #resultsDict = filterGenes(GENOMES, args.minGenomes)
            resultsDict =  GENOMES

            if args.csv == 1:
                toTSV(resultsDict, 'VF_Results')

            json.dump(resultsDict, sys.stdout)
            logging.info('Program ended successfully.')

    else:
        print(genomesList)
<|endoftext|>


<commit_before>import sublime
import json
import os.path
from sublime_plugin import TextCommand
from sublime_plugin import EventListener

def plugin_loaded():
    # update settings.
    SettingStore().update_settings()

    # update config.
    SettingStore().update_config()

    # add settings modified listener.
    SettingStore().get_settings().add_on_change(
        'reload', SettingStore().update_settings)

class TabControlListener(EventListener):

    def on_window_command(self, window, command, option):
        def _new_window():
            TabStatusStore().set_active_window_status()

        def _toggle_tabs(command, option):
            if(option == 'sidebar_separator' and SettingStore().get_auto_hide_option(
            ) and TabStatusStore().get_show_tab_status()):
                # set tab hide flag.
                TabStatusStore().set_show_tab_status(False)
                return True
            elif(not SettingStore().get_auto_hide_option()):
                TabStatusStore().toggle_show_tab_status()
                return True
            elif(SettingStore().get_auto_hide_option()):
                return False
            return False

        if(command == 'toggle_tabs'):
            if (not _toggle_tabs(command, option)):
                return ('None')
        elif(command == 'new_window'):
            _new_window()

class TabStatusStore():
    # singleton instance.
    __instance = None

    def __new__(cls, *args, **kwargs):
        if not cls.__instance:
            cls.__instance = super(TabStatusStore, cls).__new__(
                cls, *args, **kwargs)

            # show_tabs parameter from Session.sublime_session.
            cls.__show_tab_status = {}

            # just before active window status.
            cls.__active_window_status = None

        return cls.__instance

    def set_active_window_status(self):
        self.__active_window_status = self.get_show_tab_status()

    def get_active_window_status(self):
        return self.__active_window_status

    def get_show_tab_status(self):
        # get active window_id
        window_id = sublime.active_window().id()

        if(not window_id in self.__show_tab_status):
            # set just before active window id.
            self.__show_tab_status[window_id] = self.__active_window_status

        return self.__show_tab_status[window_id]

    def set_show_tab_status(self, status):
        # get active window_id
        window_id = sublime.active_window().id()

        # set show_tab_status.
        self.__show_tab_status[window_id] = status

    def toggle_show_tab_status(self):
        # get active window_id
        window_id = sublime.active_window().id()

        self.__show_tab_status[
            window_id] = not self.get_show_tab_status()

class SettingStore():
    # singleton instance.
    __instance = None

    def __new__(cls, *args, **kwargs):
        if not cls.__instance:
            cls.__instance = super(SettingStore, cls).__new__(
                cls, *args, **kwargs)

            # initial value of the setting file.
            cls.__settings = None

            # initial value of the configuration file.
            cls.__config = None

        return cls.__instance

    def get_config(self):
        return self.__config

    def set_config(self, config):
        self.__config = config

    def get_settings(self):
        return self.__settings

    def set_settings(self, settings):
        self.__settings = settings

    def update_config(self):
        def _load_config():
            def _parse_json(config_file):
                # parse json from config file.
                opened_file = open(config_file, 'r', encoding="utf8")
                return json.loads(opened_file.read(), strict=False)

            path = sublime.packages_path().replace('Packages', '')

            if(sublime.platform() == 'windows'):
                # for windows.
                config_files = (path + '\Local\Auto Save Session.sublime_session',
                                path + '\Local\Session.sublime_session')
            else:
                # for mac/linux.
                config_files = (path + '/Local/Auto Save Session.sublime_session',
                                path + '/Local/Session.sublime_session')

            # preferentially set automatic writing file.
            if(os.path.isfile(config_files[0])):
                config_file = config_files[0]
            else:
                config_file = config_files[1]

            return _parse_json(config_file)

        if(SettingStore().get_config() is None):
            SettingStore().set_config(_load_config())

        TabStatusStore().set_show_tab_status(
            SettingStore().get_tab_visibility_option())

    def update_settings(self):
        # load settings value from setting file.
        self.set_settings(sublime.load_settings(
            'sidebar_separator.sublime-settings'))

    def get_auto_hide_option(self):
        # get auto_tab_hide option and return the flag.
        auto_hide_flag = self.get_settings().get('auto_tab_hide', True)

        return auto_hide_flag

    def get_tab_visibility_option(self):
        config = self.get_config()
        return config['windows'][0]['show_tabs']

class SidebarSeparator(TextCommand):

    def run(self, edit):
        # create separate file.
        self.create_separater()

        # auto tab hide.
        self.hide_tab_bar()

    def create_separater(self):
         # create separate file.
        separate_file = sublime.active_window().new_file()

        # set buffer name.
        separate_file.set_name(self.get_separate_value())

        # set not save as separate file propertie.
        separate_file.set_scratch(True)

        # set read only propertie.
        separate_file.set_read_only(True)

    def hide_tab_bar(self):
        # controlling the tabs when the flag is true.
        if(SettingStore().get_auto_hide_option() and TabStatusStore().get_show_tab_status()):

            active_window = sublime.active_window()
            active_window.run_command('toggle_tabs', 'sidebar_separator')

    def get_separate_value(self):
        # get separate value and create separater.
        value = SettingStore().get_settings().get('separate_value', '-')
        count = SettingStore().get_settings().get('separate_count', 100)

        return value * count
<commit_msg>added getter/setter decorator to functions.<commit_after>import sublime
import json
import os.path
from sublime_plugin import TextCommand
from sublime_plugin import EventListener

def plugin_loaded():
    # update settings.
    SettingStore().update_settings()

    # update config.
    SettingStore().update_config()

    # add settings modified listener.
    SettingStore().get_settings().add_on_change(
        'reload', SettingStore().update_settings)

class TabControlListener(EventListener):

    def on_window_command(self, window, command, option):
        def _new_window():
            TabStatusStore().active_window_status = TabStatusStore().get_show_tab_status()

        def _toggle_tabs(command, option):
            if(option == 'sidebar_separator' and SettingStore().get_auto_hide_option(
            ) and TabStatusStore().get_show_tab_status()):
                # set tab hide flag.
                TabStatusStore().set_show_tab_status(False)
                return True
            elif(not SettingStore().get_auto_hide_option()):
                TabStatusStore().toggle_show_tab_status()
                return True
            elif(SettingStore().get_auto_hide_option()):
                return False
            return False

        if(command == 'toggle_tabs'):
            if (not _toggle_tabs(command, option)):
                return ('None')
        elif(command == 'new_window'):
            _new_window()

class TabStatusStore():
    # singleton instance.
    __instance = None

    def __new__(cls, *args, **kwargs):
        if not cls.__instance:
            cls.__instance = super(TabStatusStore, cls).__new__(
                cls, *args, **kwargs)

            # show_tabs parameter from Session.sublime_session.
            cls.__show_tab_status = {}

            # just before active window status.
            cls.__active_window_status = None

        return cls.__instance

    @property
    def active_window_status(self):
        return self.__active_window_status

    @active_window_status.setter
    def active_window_status(self, status):
        self.__active_window_status = status

    def get_show_tab_status(self):
        # get active window_id
        window_id = sublime.active_window().id()

        if(not window_id in self.__show_tab_status):
            # set just before active window id.
            self.__show_tab_status[window_id] = self.__active_window_status

        return self.__show_tab_status[window_id]

    def set_show_tab_status(self, status):
        # get active window_id
        window_id = sublime.active_window().id()

        # set show_tab_status.
        self.__show_tab_status[window_id] = status

    def toggle_show_tab_status(self):
        # get active window_id
        window_id = sublime.active_window().id()

        self.__show_tab_status[
            window_id] = not self.get_show_tab_status()

class SettingStore():
    # singleton instance.
    __instance = None

    def __new__(cls, *args, **kwargs):
        if not cls.__instance:
            cls.__instance = super(SettingStore, cls).__new__(
                cls, *args, **kwargs)

            # initial value of the setting file.
            cls.__settings = None

            # initial value of the configuration file.
            cls.__config = None

        return cls.__instance

    def get_config(self):
        return self.__config

    def set_config(self, config):
        self.__config = config

    def get_settings(self):
        return self.__settings

    def set_settings(self, settings):
        self.__settings = settings

    def update_config(self):
        def _load_config():
            def _parse_json(config_file):
                # parse json from config file.
                opened_file = open(config_file, 'r', encoding="utf8")
                return json.loads(opened_file.read(), strict=False)

            path = sublime.packages_path().replace('Packages', '')

            if(sublime.platform() == 'windows'):
                # for windows.
                config_files = (path + '\Local\Auto Save Session.sublime_session',
                                path + '\Local\Session.sublime_session')
            else:
                # for mac/linux.
                config_files = (path + '/Local/Auto Save Session.sublime_session',
                                path + '/Local/Session.sublime_session')

            # preferentially set automatic writing file.
            if(os.path.isfile(config_files[0])):
                config_file = config_files[0]
            else:
                config_file = config_files[1]

            return _parse_json(config_file)

        if(SettingStore().get_config() is None):
            SettingStore().set_config(_load_config())

        TabStatusStore().set_show_tab_status(
            SettingStore().get_tab_visibility_option())

    def update_settings(self):
        # load settings value from setting file.
        self.set_settings(sublime.load_settings(
            'sidebar_separator.sublime-settings'))

    def get_auto_hide_option(self):
        # get auto_tab_hide option and return the flag.
        auto_hide_flag = self.get_settings().get('auto_tab_hide', True)

        return auto_hide_flag

    def get_tab_visibility_option(self):
        config = self.get_config()
        return config['windows'][0]['show_tabs']

class SidebarSeparator(TextCommand):

    def run(self, edit):
        # create separate file.
        self.create_separater()

        # auto tab hide.
        self.hide_tab_bar()

    def create_separater(self):
         # create separate file.
        separate_file = sublime.active_window().new_file()

        # set buffer name.
        separate_file.set_name(self.get_separate_value())

        # set not save as separate file propertie.
        separate_file.set_scratch(True)

        # set read only propertie.
        separate_file.set_read_only(True)

    def hide_tab_bar(self):
        # controlling the tabs when the flag is true.
        if(SettingStore().get_auto_hide_option() and TabStatusStore().get_show_tab_status()):

            active_window = sublime.active_window()
            active_window.run_command('toggle_tabs', 'sidebar_separator')

    def get_separate_value(self):
        # get separate value and create separater.
        value = SettingStore().get_settings().get('separate_value', '-')
        count = SettingStore().get_settings().get('separate_count', 100)

        return value * count
<|endoftext|>


<commit_before>from PLC.Faults import *
from PLC.Method import Method
from PLC.Parameter import Parameter, Mixed
from PLC.Sites import Site, Sites
from PLC.Auth import Auth

related_fields = Site.related_fields.keys()
can_update = lambda (field, value): field in \
             ['name', 'abbreviated_name', 'login_base',
              'is_public', 'latitude', 'longitude', 'url',
              'max_slices', 'max_slivers', 'enabled', 'ext_consortium_id'] + \
              related_fields

class UpdateSite(Method):
    """
    Updates a site. Only the fields specified in update_fields are
    updated, all other fields are left untouched.

    PIs can only update sites they are a member of. Only admins can
    update max_slices, max_slivers, and login_base.

    Returns 1 if successful, faults otherwise.
    """

    roles = ['admin', 'pi']

    site_fields = dict(filter(can_update, Site.fields.items() + Site.related_fields.items()))

    accepts = [
        Auth(),
        Mixed(Site.fields['site_id'],
              Site.fields['login_base']),
        site_fields
        ]

    returns = Parameter(int, '1 if successful')

    def call(self, auth, site_id_or_login_base, site_fields):
        site_fields = dict(filter(can_update, site_fields.items()))

        # Get site information
        sites = Sites(self.api, [site_id_or_login_base])
        if not sites:
            raise PLCInvalidArgument, "No such site"
        site = sites[0]

        if site['peer_id'] is not None:
            raise PLCInvalidArgument, "Not a local site"

        # Authenticated function
        assert self.caller is not None

        # If we are not an admin, make sure that the caller is a
        # member of the site.
        if 'admin' not in self.caller['roles']:
            if site['site_id'] not in self.caller['site_ids']:
                raise PLCPermissionDenied, "Not allowed to modify specified site"

            # Remove admin only fields
            for key in 'max_slices', 'max_slivers', 'login_base':
                if key in site_fields:
                    del site_fields[key]

        # Make requested associations
        for field in related_fields:
            if field in site_fields:
                site.associate(auth, field, site_fields[field])
                site_fields.pop(field)

        site.update(site_fields)
        site.update_last_updated(False)
        site.sync()

        # Logging variables
        self.event_objects = {'Site': [site['site_id']]}
        self.message = 'Site %d updated: %s' % \
                (site['site_id'], ", ".join(site_fields.keys()))

        return 1
<commit_msg>Update Site HRN when Site's login_base is updated<commit_after>from PLC.Faults import *
from PLC.Method import Method
from PLC.Parameter import Parameter, Mixed
from PLC.Sites import Site, Sites
from PLC.Auth import Auth

from PLC.SiteTags import SiteTags
from PLC.Methods.UpdateSiteTag import UpdateSiteTag

related_fields = Site.related_fields.keys()
can_update = lambda (field, value): field in \
             ['name', 'abbreviated_name', 'login_base',
              'is_public', 'latitude', 'longitude', 'url',
              'max_slices', 'max_slivers', 'enabled', 'ext_consortium_id'] + \
              related_fields

class UpdateSite(Method):
    """
    Updates a site. Only the fields specified in update_fields are
    updated, all other fields are left untouched.

    PIs can only update sites they are a member of. Only admins can
    update max_slices, max_slivers, and login_base.

    Returns 1 if successful, faults otherwise.
    """

    roles = ['admin', 'pi']

    site_fields = dict(filter(can_update, Site.fields.items() + Site.related_fields.items()))

    accepts = [
        Auth(),
        Mixed(Site.fields['site_id'],
              Site.fields['login_base']),
        site_fields
        ]

    returns = Parameter(int, '1 if successful')

    def call(self, auth, site_id_or_login_base, site_fields):
        site_fields = dict(filter(can_update, site_fields.items()))

        # Get site information
        sites = Sites(self.api, [site_id_or_login_base])
        if not sites:
            raise PLCInvalidArgument, "No such site"
        site = sites[0]

        if site['peer_id'] is not None:
            raise PLCInvalidArgument, "Not a local site"

        # Authenticated function
        assert self.caller is not None

        # If we are not an admin, make sure that the caller is a
        # member of the site.
        if 'admin' not in self.caller['roles']:
            if site['site_id'] not in self.caller['site_ids']:
                raise PLCPermissionDenied, "Not allowed to modify specified site"

            # Remove admin only fields
            for key in 'max_slices', 'max_slivers', 'login_base':
                if key in site_fields:
                    del site_fields[key]

        # Make requested associations
        for field in related_fields:
            if field in site_fields:
                site.associate(auth, field, site_fields[field])
                site_fields.pop(field)

        site.update(site_fields)
        site.update_last_updated(False)
        site.sync()

        # Logging variables
        self.event_objects = {'Site': [site['site_id']]}
        self.message = 'Site %d updated: %s' % \
                (site['site_id'], ", ".join(site_fields.keys()))

        # Update Site HRN if login_base changed
        if 'login_base' in site_fields:
            root_auth = self.api.config.PLC_HRN_ROOT
            tagname = 'hrn'
            tagvalue = '.'.join([root_auth, site['login_base']])
            site_tags=SiteTags(self.api,{'tagname':tagname,'site_id':site['site_id']})
            UpdateSiteTag(self.api).__call__(auth,site_tags[0]['site_tag_id'],tagvalue)

        return 1
<|endoftext|>


<commit_before>#!/usr/bin/env python3
# Based upon:
# https://github.com/voc/voctomix/tree/voctopanel/example-scripts/voctolight

from enum import Enum
import argparse
import asyncio
import configparser
import json
import os.path

class Config(configparser.ConfigParser, object):
    config_files = [
        os.path.join(os.path.dirname(os.path.realpath(__file__)),
                     'default-config.ini'),
        os.path.join(os.path.dirname(os.path.realpath(__file__)),
                     'config.ini'),
        '/etc/voctomix/voctolight.ini',
        '/etc/voctolight.ini',
        os.path.expanduser('~/.voctolight.ini'),
    ]

    def __init__(self, cmd_line_config=None):
        super().__init__()
        self.cmd_line_config = cmd_line_config
        self._read_config()

    def _read_config(self):
        self.read(self.config_files)
        if self.cmd_line_config:
            self.read_file(self.cmd_line_config)
            self.cmd_line_config.seek(0)

    def setup_with_server_config(self, server_config):
        self.clear()
        self._read_config()
        self.read_dict(server_config)

    def getlist(self, section, option):
        return [x.strip() for x in self.get(section, option).split(',')]

class Connection(asyncio.Protocol):
    def __init__(self, interpreter):
        self.interpreter = interpreter

    def connect(self, host, port=9999):
        self.loop = asyncio.get_event_loop()
        self.conn = self.loop.create_connection(lambda: self, host, port)
        self.loop.run_until_complete(conn.conn)

    def send(self, message):
        self.transport.write(message.encode())
        self.transport.write('\n'.encode())

    def connection_made(self, transport):
        self.transport = transport

    def data_received(self, data):
        # print('data received: {}'.format(data.decode()))
        lines = data.decode().rstrip('\n').split('\n')
        for line in lines:
            interpreter.handler(line)

    def onnection_lost(self, exc):
        print('server closed the connection')
        asyncio.get_event_loop().stop()

# FIXME: Duplicate from videomix.py
class CompositeModes(Enum):
        fullscreen = 0
        side_by_side_equal = 1
        side_by_side_preview = 2
        picture_in_picture = 3

class Interpreter(object):
    a_or_b = False
    primary = False
    composite_mode = CompositeModes.fullscreen

    def __init__(self, actor, config, debug=False):
        self.config = config
        self.actor = actor
        self.debug = debug
        actor.reset_led()
        if self.debug:
            print('LED has been reset to off')

    def compute_state(self):
        if self.composite_mode == CompositeModes.fullscreen:
            return self.primary
        else:
            return self.a_or_b

    def handler(self, response):
        words = response.split()
        signal = words[0]
        args = words[1:]
        try:
            handler = getattr(self, 'handle_{}'.format(signal))
        except AttributeError:
            print('Ignoring signal', signal)
        else:
            handler(args)
            enable = interpreter.compute_state()
            actor.enable_tally(enable)
            if self.debug:
                print('LED has been set to {}.'.format(
                    'on' if enable else 'off'))

    def handle_video_status(self, cams):
        mycam = self.config.get('light', 'cam')
        if mycam in cams:
            self.a_or_b = True
        else:
            self.a_or_b = False

        self.primary = (cams[0] == mycam)
        # print ('Is primary?', self.primary)

    def handle_composite_mode(self, mode_list):
        mode = mode_list[0]
        if mode == 'fullscreen':
            self.composite_mode = CompositeModes.fullscreen
        elif mode == 'side_by_side_equal':
            self.composite_mode = CompositeModes.side_by_side_equal
        elif mode == 'side_by_side_preview':
            self.composite_mode = CompositeModes.side_by_side_preview
        elif mode == 'picture_in_picture':
            self.composite_mode = CompositeModes.picture_in_picture
        else:
            print('Cannot handle', mode, 'of type', type(mode))

    def handle_server_config(self, args):
        server_config_json = ' '.join(args)
        server_config = json.loads(server_config_json)
        self.config.setup_with_server_config(server_config)

class RPiGPIODriver:
    def __init__(self, config):
        self._import()
        self.gpios = config.get('light', 'gpios').split(',')
        self.gpio_red = int(config.get('light', 'gpio_red'))
        self.reset_led()

    def _import(self):
        global GPIO
        import RPi.GPIO as GPIO

    def reset_led(self):
        GPIO.setmode(GPIO.BOARD)
        GPIO.setwarnings(False)
        for gpio in self.gpios:
            gpio = int(gpio)
            GPIO.setup(gpio, GPIO.OUT)
            GPIO.output(gpio, GPIO.HIGH)

    def enable_tally(self, enable):
        if enable:
            GPIO.output(self.gpio_red, GPIO.LOW)
        else:
            GPIO.output(self.gpio_red, GPIO.HIGH)

class DummyDriver:
    def __init__(self, config):
        pass

    def reset_led(self):
        pass

    def enable_tally(self, enable):
        pass

class SerialDTRDriver:
    def __init__(self, config):
        self.fn = config.get('light', 'port')
        self.fd = None

    def reset_led(self):
        if self.fd:
            self.fd.close()

    def enable_tally(self, enable):
        if enable:
            self.fd = open(self.fn)
        else:
            if self.fd:
                self.fd.close()
            self.fd = None

if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        description='Tallylight controlling daemon for voctomix.')
    parser.add_argument(
        '--config', type=open, help='Use a specific config file')
    parser.add_argument(
        '--debug', action='store_true',
        help='Show what would be done instead of toggling lights')
    args = parser.parse_args()
    config = Config(cmd_line_config=args.config)
    driver = config.get('light', 'driver')
    if driver == 'dummy':
        actor = DummyDriver(config)
    elif driver == 'rpi':
        actor = RPiGPIODriver(config)
    elif driver == 'serial':
        actor = SerialDTRDriver(config)
    else:
        raise Exception('Unknown driver: ' + driver)
    interpreter = Interpreter(actor, config, debug=args.debug)
    conn = Connection(interpreter)
    conn.connect(config.get('server', 'host'))
    conn.send('get_config')
    conn.send('get_composite_mode')
    conn.send('get_video')
    conn.loop.run_forever()
<commit_msg>Clean shutdown<commit_after>#!/usr/bin/env python3
# Based upon:
# https://github.com/voc/voctomix/tree/voctopanel/example-scripts/voctolight

from enum import Enum
import argparse
import asyncio
import configparser
import json
import os.path

class Config(configparser.ConfigParser, object):
    config_files = [
        os.path.join(os.path.dirname(os.path.realpath(__file__)),
                     'default-config.ini'),
        os.path.join(os.path.dirname(os.path.realpath(__file__)),
                     'config.ini'),
        '/etc/voctomix/voctolight.ini',
        '/etc/voctolight.ini',
        os.path.expanduser('~/.voctolight.ini'),
    ]

    def __init__(self, cmd_line_config=None):
        super().__init__()
        self.cmd_line_config = cmd_line_config
        self._read_config()

    def _read_config(self):
        self.read(self.config_files)
        if self.cmd_line_config:
            self.read_file(self.cmd_line_config)
            self.cmd_line_config.seek(0)

    def setup_with_server_config(self, server_config):
        self.clear()
        self._read_config()
        self.read_dict(server_config)

    def getlist(self, section, option):
        return [x.strip() for x in self.get(section, option).split(',')]

class Connection(asyncio.Protocol):
    def __init__(self, interpreter):
        self.interpreter = interpreter

    def connect(self, host, port=9999):
        self.loop = asyncio.get_event_loop()
        self.conn = self.loop.create_connection(lambda: self, host, port)
        self.loop.run_until_complete(conn.conn)

    def send(self, message):
        self.transport.write(message.encode())
        self.transport.write('\n'.encode())

    def connection_made(self, transport):
        self.transport = transport

    def data_received(self, data):
        # print('data received: {}'.format(data.decode()))
        lines = data.decode().rstrip('\n').split('\n')
        for line in lines:
            interpreter.handler(line)

    def connection_lost(self, exc):
        print('server closed the connection')
        asyncio.get_event_loop().stop()

# FIXME: Duplicate from videomix.py
class CompositeModes(Enum):
        fullscreen = 0
        side_by_side_equal = 1
        side_by_side_preview = 2
        picture_in_picture = 3

class Interpreter(object):
    a_or_b = False
    primary = False
    composite_mode = CompositeModes.fullscreen

    def __init__(self, actor, config, debug=False):
        self.config = config
        self.actor = actor
        self.debug = debug
        actor.reset_led()
        if self.debug:
            print('LED has been reset to off')

    def compute_state(self):
        if self.composite_mode == CompositeModes.fullscreen:
            return self.primary
        else:
            return self.a_or_b

    def handler(self, response):
        words = response.split()
        signal = words[0]
        args = words[1:]
        try:
            handler = getattr(self, 'handle_{}'.format(signal))
        except AttributeError:
            print('Ignoring signal', signal)
        else:
            handler(args)
            enable = interpreter.compute_state()
            actor.enable_tally(enable)
            if self.debug:
                print('LED has been set to {}.'.format(
                    'on' if enable else 'off'))

    def handle_video_status(self, cams):
        mycam = self.config.get('light', 'cam')
        if mycam in cams:
            self.a_or_b = True
        else:
            self.a_or_b = False

        self.primary = (cams[0] == mycam)
        # print ('Is primary?', self.primary)

    def handle_composite_mode(self, mode_list):
        mode = mode_list[0]
        if mode == 'fullscreen':
            self.composite_mode = CompositeModes.fullscreen
        elif mode == 'side_by_side_equal':
            self.composite_mode = CompositeModes.side_by_side_equal
        elif mode == 'side_by_side_preview':
            self.composite_mode = CompositeModes.side_by_side_preview
        elif mode == 'picture_in_picture':
            self.composite_mode = CompositeModes.picture_in_picture
        else:
            print('Cannot handle', mode, 'of type', type(mode))

    def handle_server_config(self, args):
        server_config_json = ' '.join(args)
        server_config = json.loads(server_config_json)
        self.config.setup_with_server_config(server_config)

class RPiGPIODriver:
    def __init__(self, config):
        self._import()
        self.gpios = config.get('light', 'gpios').split(',')
        self.gpio_red = int(config.get('light', 'gpio_red'))
        self.reset_led()

    def _import(self):
        global GPIO
        import RPi.GPIO as GPIO

    def reset_led(self):
        GPIO.setmode(GPIO.BOARD)
        GPIO.setwarnings(False)
        for gpio in self.gpios:
            gpio = int(gpio)
            GPIO.setup(gpio, GPIO.OUT)
            GPIO.output(gpio, GPIO.HIGH)

    def enable_tally(self, enable):
        if enable:
            GPIO.output(self.gpio_red, GPIO.LOW)
        else:
            GPIO.output(self.gpio_red, GPIO.HIGH)

class DummyDriver:
    def __init__(self, config):
        pass

    def reset_led(self):
        pass

    def enable_tally(self, enable):
        pass

class SerialDTRDriver:
    def __init__(self, config):
        self.fn = config.get('light', 'port')
        self.fd = None

    def reset_led(self):
        if self.fd:
            self.fd.close()

    def enable_tally(self, enable):
        if enable:
            self.fd = open(self.fn)
        else:
            if self.fd:
                self.fd.close()
            self.fd = None

if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        description='Tallylight controlling daemon for voctomix.')
    parser.add_argument(
        '--config', type=open, help='Use a specific config file')
    parser.add_argument(
        '--debug', action='store_true',
        help='Show what would be done instead of toggling lights')
    args = parser.parse_args()
    config = Config(cmd_line_config=args.config)
    driver = config.get('light', 'driver')
    if driver == 'dummy':
        actor = DummyDriver(config)
    elif driver == 'rpi':
        actor = RPiGPIODriver(config)
    elif driver == 'serial':
        actor = SerialDTRDriver(config)
    else:
        raise Exception('Unknown driver: ' + driver)
    interpreter = Interpreter(actor, config, debug=args.debug)
    conn = Connection(interpreter)
    conn.connect(config.get('server', 'host'))
    conn.send('get_config')
    conn.send('get_composite_mode')
    conn.send('get_video')
    try:
        conn.loop.run_forever()
    except KeyboardInterrupt:
        print('Quitting (Ctrl-C)')
    actor.reset_led()
<|endoftext|>


<commit_before># Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

import pytest

from tests.common.custom_cluster_test_suite import CustomClusterTestSuite
from tests.common.skip import SkipIfBuildType

@SkipIfBuildType.not_dev_build
class TestAllocFail(CustomClusterTestSuite):
  """Tests for handling malloc() failure for UDF/UDA"""

  @classmethod
  def get_workload(self):
    return 'functional-query'

  @pytest.mark.execute_serially
  @CustomClusterTestSuite.with_args("--stress_free_pool_alloc=1")
  def test_alloc_fail_init(self, vector):
    self.run_test_case('QueryTest/alloc-fail-init', vector)

  @pytest.mark.execute_serially
  @CustomClusterTestSuite.with_args("--stress_free_pool_alloc=3")
  def test_alloc_fail_update(self, vector):
    # Note that this test relies on pre-aggregation to exercise the Serialize() path so
    # query option 'num_nodes' must not be 1. CustomClusterTestSuite.add_test_dimensions()
    # already filters out vectors with 'num_nodes' != 0 so just assert it here.
    assert vector.get_value('exec_option')['num_nodes'] == 0
    self.run_test_case('QueryTest/alloc-fail-update', vector)
<commit_msg>IMPALA-2925: Mark test_alloc_update as xfail.<commit_after># Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

import pytest

from tests.common.custom_cluster_test_suite import CustomClusterTestSuite
from tests.common.skip import SkipIfBuildType

@SkipIfBuildType.not_dev_build
class TestAllocFail(CustomClusterTestSuite):
  """Tests for handling malloc() failure for UDF/UDA"""

  @classmethod
  def get_workload(self):
    return 'functional-query'

  @pytest.mark.execute_serially
  @CustomClusterTestSuite.with_args("--stress_free_pool_alloc=1")
  def test_alloc_fail_init(self, vector):
    self.run_test_case('QueryTest/alloc-fail-init', vector)

  @pytest.mark.execute_serially
  @CustomClusterTestSuite.with_args("--stress_free_pool_alloc=3")
  def test_alloc_fail_update(self, vector):
    # TODO: Rewrite or remove the non-deterministic test.
    pytest.xfail("IMPALA-2925: the execution is not deterministic so some "
                 "tests sometimes don't fail as expected")
    # Note that this test relies on pre-aggregation to exercise the Serialize() path so
    # query option 'num_nodes' must not be 1. CustomClusterTestSuite.add_test_dimensions()
    # already filters out vectors with 'num_nodes' != 0 so just assert it here.
    assert vector.get_value('exec_option')['num_nodes'] == 0
    self.run_test_case('QueryTest/alloc-fail-update', vector)
<|endoftext|>