<commit_before>from sklearn import tree, svm
from src.hw1_Classification import *
from sklearn.ensemble import RandomForestClassifier

def read_from_file(path):
    feature_list = []
    label_list = []
    with open(path, 'r') as f:
        while True:
            line = f.readline()
            if not line or len(line) == 1:
                break
            info = list(map(lambda x: x.strip(' '), line.strip('\n').split(',')))
            feature_list.append(generate_feature(info))
            label_list.append(generate_value(info[14].strip('.'), income))
        f.close()
    return (feature_list, label_list)
    pass

def generate_feature(input_args):
    feature = list()
    feature.append(int(input_args[0]))
    feature.append(generate_value(input_args[1], workclass))
    feature.append(int(input_args[2]))
    feature.append(generate_value(input_args[3], education))
    feature.append(int(input_args[4]))
    feature.append(generate_value(input_args[5], marital_status))
    feature.append(generate_value(input_args[6], occupation))
    feature.append(generate_value(input_args[7], relationship))
    feature.append(generate_value(input_args[8], race))
    feature.append(generate_value(input_args[9], sex))
    feature.append(int(input_args[10]))
    feature.append(int(input_args[11]))
    feature.append(int(input_args[12]))
    feature.append(generate_value(input_args[13], native_country))
    return feature
    pass

def generate_value(name, known_list):
    if name not in known_list:
        return 0
    else:
        return known_list.index(name) + 1

if __name__ == "__main__":
    # clf = tree.DecisionTreeClassifier()
    clf = RandomForestClassifier(n_estimators=10)
    # clf = svm.SVC()

    train_set = read_from_file("../../data/hw1/adult.data.txt")
    clf = clf.fit(train_set[0], train_set[1])
    test_set = read_from_file("../../data/hw1/adult.test.txt")

    num = 0
    correct = 0
    for i in clf.predict(test_set[0]):
        if i == test_set[1][num]:
            correct += 1
        num += 1
    print(correct)
    print(num)
    print(correct / num)
    pass
<commit_msg>change vetcor<commit_after>from sklearn import tree, svm
from src.hw1_Classification import *
from sklearn.ensemble import RandomForestClassifier

def read_from_file(path):
    feature_list = []
    label_list = []
    with open(path, 'r') as f:
        while True:
            line = f.readline()
            if not line or len(line) == 1:
                break
            info = list(map(lambda x: x.strip(' '), line.strip('\n').split(',')))
            feature_list.append(generate_feature(info))
            label_list.append(generate_value(info[14].strip('.'), income))
        f.close()
    return (feature_list, label_list)
    pass

def generate_feature(input_args):
    feature = list()
    feature.append(int(input_args[0]) // 10)
    feature.append(generate_value(input_args[1], workclass))
    feature.append(int(input_args[2]) // 10000)
    feature.append(generate_value(input_args[3], education))
    feature.append(int(input_args[4]) // 4)
    feature.append(generate_value(input_args[5], marital_status))
    feature.append(generate_value(input_args[6], occupation))
    feature.append(generate_value(input_args[7], relationship))
    feature.append(generate_value(input_args[8], race))
    feature.append(generate_value(input_args[9], sex))
    feature.append(int(input_args[10]) // 1000)
    feature.append(int(input_args[11]))
    feature.append(int(input_args[12]) // 10)
    feature.append(generate_value(input_args[13], native_country))
    return feature
    pass

def generate_value(name, known_list):
    if name not in known_list:
        return 0
    else:
        return known_list.index(name) + 1

if __name__ == "__main__":
    # clf = tree.DecisionTreeClassifier()
    # clf = RandomForestClassifier(n_estimators=10)
    clf = svm.SVC()

    train_set = read_from_file("../../data/hw1/adult.data.txt")
    clf = clf.fit(train_set[0], train_set[1])
    test_set = read_from_file("../../data/hw1/adult.test.txt")

    num = 0
    correct = 0
    for i in clf.predict(test_set[0]):
        if i == test_set[1][num]:
            correct += 1
        num += 1
    print(correct)
    print(num)
    print(correct / num)
    pass
<|endoftext|>


<commit_before>import os
import sys
import subprocess
import logging
import time

from sqlalchemy.orm.exc import NoResultFound

from Pegasus.service import app, request
from Pegasus.service.ensembles.models import Ensembles, EnsembleStates, EnsembleWorkflowStates
from Pegasus.db.schema.stampede_dashboard_schema import DashboardWorkflow, DashboardWorkflowstate

log = logging.getLogger(__name__)

class EMException(Exception): pass

def pathfind(exe):
    PATH = os.getenv("PATH","/bin:/usr/bin:/usr/local/bin")
    PATH = PATH.split(":")
    for prefix in PATH:
        exepath = os.path.join(prefix, exe)
        if os.path.isfile(exepath):
            return exepath
    raise EMException("%s not found on PATH" % exe)

def get_bin(name, exe):
    # Try to find NAME/bin using 1) NAME env var, 2) NAME config
    # variable, 3) PATH env var
    exepath = None

    HOME = os.getenv(name, app.config.get(name, None))

    if HOME is not None:
        if not os.path.isdir(HOME):
            raise EMException("%s is not a directory: %s" % (name, HOME))
        BIN = os.path.join(HOME, "bin")
        if not os.path.isdir(BIN):
            raise EMException("%s/bin is not a directory: %s" % (name, BIN))
        exepath = os.path.join(BIN, exe)

    exepath = exepath or pathfind(exe)

    if not os.path.isfile(exepath):
        raise EMException("%s not found: %s" % (exe, exepath))

    BIN = os.path.dirname(exepath)

    return BIN

def get_pegasus_bin():
    return get_bin("PEGASUS_HOME", "pegasus-plan")

def get_condor_bin():
    return get_bin("CONDOR_HOME", "condor_submit_dag")

def check_environment():
    # Make sure Pegasus and Condor are on the PATH, or PEGASUS_HOME
    # and CONDOR_HOME are set in the environment or configuration
    get_pegasus_bin()
    get_condor_bin()

def get_script_env():
    PEGASUS_BIN = get_pegasus_bin()
    CONDOR_BIN = get_condor_bin()

    # Add pegasus bin dirs to PATH
    env = dict(os.environ)
    PATH = env.get("PATH", "/bin:/usr/bin:/usr/local/bin")
    PATH = PEGASUS_BIN + ":" + CONDOR_BIN + ":" + PATH
    env["PATH"] = PATH

    return env

def runscript(script, cwd=None, env=None):
    # Make sure the cwd is OK
    if cwd is not None and not os.path.isdir(cwd):
        raise EMException("Working directory does not exist: %s" % cwd)

    if env is None:
        env = dict(os.environ)

    p = subprocess.Popen(script, shell=True, env=env, cwd=cwd)

    rc = p.wait()

    if rc != 0:
        raise EMException("Script failed with exitcode %d" % rc)

def forkscript(script, pidfile=None, cwd=None, env=None):
    # This does a double fork to detach the process from the python
    # interpreter so that we don't have to call wait() on it

    # Make sure the cwd is OK
    if cwd is not None and not os.path.isdir(cwd):
        raise EMException("Working directory does not exist: %s" % cwd)

    if env is None:
        env = dict(os.environ)

    # This is just to ensure we get an exception if there is
    # something wrong with the pidfile
    if pidfile is not None:
        try:
            open(pidfile, "w").close()
        except:
            raise EMException("Unable to write pidfile: %s" % pidfile)

    pid1 = os.fork()
    if pid1 == 0:
        if cwd is not None:
            os.chdir(cwd)

        pid2 = os.fork()
        if pid2 == 0:
            os.execve("/bin/sh", ["/bin/sh","-c",script], env)
            os._exit(255)

        if pidfile is not None:
            f = open(pidfile, "w")
            f.write("%d\n" % pid2)
            f.close()

        os._exit(0)

    pid, exitcode = os.waitpid(pid1, 0)
    if exitcode != 0:
        raise EMException("Non-zero exitcode launching script: %d" % exitcode)

class WorkflowProcessor:
    def __init__(self, db, workflow):
        self.db = db
        self.workflow = workflow

    def get_file(self, filename):
        return os.path.join(self.workflow.basedir, filename)

    def get_bundledir(self):
        return os.path.join(self.workflow.basedir, "bundle")

    def get_pidfile(self):
        return self.get_file("planner.pid")

    def get_resultfile(self):
        return self.get_file("planner.result")

    def get_logfile(self):
        return self.get_file("planner.log")

    def get_planfile(self):
        return self.get_file("plan.sh")

    def plan(self):
        "Launch the pegasus planner"
        workdir = self.get_bundledir()
        pidfile = self.get_pidfile()
        logfile = self.get_logfile()
        planfile = self.get_planfile()
        resultfile = self.get_resultfile()

        if os.path.isfile(pidfile) and self.planning():
            raise EMException("Planner already running")

        # When we re-plan, we need to remove all the old
        # files so that the ensemble manager doesn't get
        # confused.
        files = [
            logfile,
            resultfile,
            pidfile
        ]
        for f in files:
            if os.path.isfile(f):
                os.remove(f)

        script = "%s >%s 2>&1; /bin/echo $? >%s" % (planfile, logfile, resultfile)
        forkscript(script, cwd=workdir, pidfile=pidfile, env=get_script_env())

    def planning(self):
        "Check pidfile to see if the planner is still running"
        pidfile = self.get_pidfile()

        if not os.path.exists(pidfile):
            raise EMException("pidfile missing")

        pid = int(open(pidfile,"r").read())

        try:
            os.kill(pid, 0)
            # If that succeeds, the process is still running
            return True
        except OSError, e:
            # errno 3 is No Such Process
            if e.errno != 3:
                raise

        return False

    def planning_successful(self):
        "Check to see if planning was successful"
        resultfile = self.get_resultfile()

        if not os.path.exists(resultfile):
            raise EMException("Result file not found: %s" % resultfile)

        exitcode = int(open(resultfile, "r").read())

        return exitcode == 0

    def get_submitdir(self):
        "Get the workflow submitdir from the workflow log"
        logfile = self.get_logfile()

        if not os.path.isfile(logfile):
            raise EMException("Workflow log file not found: %s" % logfile)

        submitdir = None

        f = open(logfile, "r")
        try:
            for l in f:
                if l.startswith("pegasus-run"):
                    submitdir = l.split()[1]
        finally:
            f.close()

        if submitdir is None:
            raise EMException("No pegasus-run found in the workflow log: %s" % logfile)

        return submitdir

    def get_wf_uuid(self):
        "Get the workflow UUID from the braindump file"
        submitdir = self.get_submitdir()

        braindump = os.path.join(submitdir, "braindump.txt")

        if not os.path.isfile(braindump):
            raise EMException("braindump.txt not found")

        wf_uuid = None

        f = open(braindump, "r")
        try:
            for l in f:
                if l.startswith("wf_uuid"):
                    wf_uuid = l.split()[1]
        finally:
            f.close()

        if wf_uuid is None:
            raise EMException("wf_uuid not found in braindump.txt")

        return wf_uuid

    def run(self):
        "Run the workflow using pegasus-run"
        submitdir = self.workflow.submitdir

        if submitdir is None:
            raise EMException("Workflow submitdir not set")

        if not os.path.isdir(submitdir):
            raise EMException("Workflow submit dir does not exist: %s" % submitdir)

        logfile = self.get_logfile()

        runscript("pegasus-run %s >>%s 2>&1" % (submitdir,logfile), env=get_script_env())

    def get_dashboard(self):
        "Get the dashboard record for the workflow"
        wf_uuid = self.workflow.wf_uuid
        if wf_uuid is None:
            raise EMException("wf_uuid is none")

        try:
            w = self.db.session.query(DashboardWorkflow)\
                    .filter_by(wf_uuid=str(wf_uuid))\
                    .one()
            return w
        except NoResultFound:
            return None

    def get_dashboard_state(self):
        "Get the latest state of the workflow from the dashboard tables"
        w = self.get_dashboard()
        if w is None:
            raise EMException("Dashboard workflow not found")

        try:
            ws = self.db.session.query(DashboardWorkflowstate)\
                    .filter_by(wf_id=w.wf_id)\
                    .order_by("timestamp desc")\
                    .first()
        except NoResultFound:
            raise EMException("Dashboard workflow state not found")

        return ws

    def pending(self):
        "The workflow is pending if there is no dashboard record"
        w = self.get_dashboard()
        return w is None

    def running(self):
        "Is the workflow running"
        ws = self.get_dashboard_state()
        return ws.state == "WORKFLOW_STARTED"

    def running_successful(self):
        "Assuming the workflow is done running, did it finish successfully?"
        ws = self.get_dashboard_state()
        if ws.state == "WORKFLOW_STARTED":
            raise EMException("Workflow is running")
        return ws.status == 0

class EnsembleProcessor:
    Processor = WorkflowProcessor

    def __init__(self, db, ensemble):
        self.db = db
        self.ensemble = ensemble
        self.active = ensemble.state == EnsembleStates.ACTIVE
        self.max_running = ensemble.max_running
        self.max_planning = ensemble.max_planning
        self.running = 0
        self.planning = 0

        # We are only interested in workflows that are in one
        # of the following states. The EM doesn't handle other
        # states.
        active_states = set((
            EnsembleWorkflowStates.READY,
            EnsembleWorkflowStates.PLANNING,
            EnsembleWorkflowStates.QUEUED,
            EnsembleWorkflowStates.RUNNING
        ))

        # We need a copy of the list so we can sort it
        self.workflows = [w for w in ensemble.workflows if w.state in active_states]

        # Sort workflows by priority
        def get_priority(w):
            return w.priority

        self.workflows.sort(key=get_priority)

        # Count the number of planning and running workflows
        for w in self.workflows:
            if w.state == EnsembleWorkflowStates.PLANNING:
                self.planning += 1
            elif w.state == EnsembleWorkflowStates.RUNNING:
                self.running += 1

    def can_plan(self):
        return self.active and self.planning < self.max_planning

    def plan_workflow(self, workflow):
        log.info("Planning %s" % workflow.name)
        self.planning += 1
        workflow.set_state(EnsembleWorkflowStates.PLANNING)
        workflow.set_updated()

        self.db.session.flush()

        # Fork planning task
        p = self.Processor(self.db, workflow)
        try:
            p.plan()
        except Exception, e:
            log.exception(e)
            workflow.set_state(EnsembleWorkflowStates.PLAN_FAILED)
            workflow.set_updated()

        self.db.session.commit()

    def handle_ready(self, workflow):
        if self.can_plan():
            self.plan_workflow(workflow)

    def handle_planning(self, workflow):
        p = self.Processor(self.db, workflow)

        if p.planning():
            return

        self.planning -= 1

        # Planning failed
        if not p.planning_successful():
            workflow.set_state(EnsembleWorkflowStates.PLAN_FAILED)
            workflow.set_updated()
            self.db.session.commit()
            return

        log.info("Queueing workflow %s" % workflow.name)

        # Planning succeeded, get uuid and queue workflow
        workflow.set_wf_uuid(p.get_wf_uuid())
        workflow.set_submitdir(p.get_submitdir())
        workflow.set_state(EnsembleWorkflowStates.QUEUED)
        workflow.set_updated()
        self.db.session.commit()

        # Go ahead and handle the queued state now
        self.handle_queued(workflow)

    def can_run(self):
        return self.active and self.running < self.max_running

    def run_workflow(self, workflow):
        log.info("Running workflow %s" % workflow.name)

        self.running += 1
        workflow.set_state(EnsembleWorkflowStates.RUNNING)
        workflow.set_updated()
        self.db.session.flush()

        p = self.Processor(self.db, workflow)
        try:
            p.run()
        except Exception, e:
            log.exception(e)
            workflow.set_state(EnsembleWorkflowStates.RUN_FAILED)
            workflow.set_updated()

        self.db.session.commit()

    def handle_queued(self, workflow):
        if self.can_run():
            self.run_workflow(workflow)

    def handle_running(self, workflow):
        p = self.Processor(self.db, workflow)

        if p.pending() or p.running():
            return

        self.running -= 1

        if p.running_successful():
            workflow.set_state(EnsembleWorkflowStates.SUCCESSFUL)
        else:
            workflow.set_state(EnsembleWorkflowStates.FAILED)

        workflow.set_updated()

        self.db.session.commit()

    def handle_workflow(self, w):
        if w.state == EnsembleWorkflowStates.READY:
            self.handle_ready(w)
        elif w.state == EnsembleWorkflowStates.PLANNING:
            self.handle_planning(w)
        elif w.state == EnsembleWorkflowStates.QUEUED:
            self.handle_queued(w)
        elif w.state == EnsembleWorkflowStates.RUNNING:
            self.handle_running(w)

    def run(self):
        log.info("Processing %d ensemble workflows..." % len(self.workflows))
        for w in self.workflows:
            try:
                self.handle_workflow(w)
            except Exception, e:
                log.error("Processing workflow %s of ensemble %s" % (w.name, self.ensemble.name))
                log.exception(e)
                self.db.session.rollback()

class EnsembleManager(object):
    Processor = EnsembleProcessor

    def __init__(self, interval=None):
        if interval is None:
            interval = float(app.config["EM_INTERVAL"])
        self.interval = interval

    def start(self):
        pid = os.fork()
        if pid == 0:
            self.run()
        else:
            return pid

    def run(self):
        log.info("Ensemble Manager starting")
        self.loop_forever()

    def loop_forever(self):
        while True:
            self.loop_once()
            time.sleep(self.interval)

    def get_active_users(self):
        # TODO Identify the active users
        return [request.get_user_by_uid(os.getuid())]

    def loop_once(self):
        for u in self.get_active_users():

            # Switch to user
            if os.getuid() != u.uid:
                if os.getuid() != 0:
                    log.error("Pegasus service must run as root to process ensembles for %s", u.username)
                    continue

                os.setegid(u.gid)
                os.seteuid(u.uid)

            log.info("Processing ensembles for %s", u.username)

            db = Ensembles(u.get_master_db_url())
            for e in db.list_actionable_ensembles():
                log.info("Processing ensemble %s", e.name)
                p = self.Processor(db, e)
                p.run()

<commit_msg>Reduce logging output<commit_after>import os
import sys
import subprocess
import logging
import time

from sqlalchemy.orm.exc import NoResultFound

from Pegasus.service import app, request
from Pegasus.service.ensembles.models import Ensembles, EnsembleStates, EnsembleWorkflowStates
from Pegasus.db.schema.stampede_dashboard_schema import DashboardWorkflow, DashboardWorkflowstate

log = logging.getLogger(__name__)

class EMException(Exception): pass

def pathfind(exe):
    PATH = os.getenv("PATH","/bin:/usr/bin:/usr/local/bin")
    PATH = PATH.split(":")
    for prefix in PATH:
        exepath = os.path.join(prefix, exe)
        if os.path.isfile(exepath):
            return exepath
    raise EMException("%s not found on PATH" % exe)

def get_bin(name, exe):
    # Try to find NAME/bin using 1) NAME env var, 2) NAME config
    # variable, 3) PATH env var
    exepath = None

    HOME = os.getenv(name, app.config.get(name, None))

    if HOME is not None:
        if not os.path.isdir(HOME):
            raise EMException("%s is not a directory: %s" % (name, HOME))
        BIN = os.path.join(HOME, "bin")
        if not os.path.isdir(BIN):
            raise EMException("%s/bin is not a directory: %s" % (name, BIN))
        exepath = os.path.join(BIN, exe)

    exepath = exepath or pathfind(exe)

    if not os.path.isfile(exepath):
        raise EMException("%s not found: %s" % (exe, exepath))

    BIN = os.path.dirname(exepath)

    return BIN

def get_pegasus_bin():
    return get_bin("PEGASUS_HOME", "pegasus-plan")

def get_condor_bin():
    return get_bin("CONDOR_HOME", "condor_submit_dag")

def check_environment():
    # Make sure Pegasus and Condor are on the PATH, or PEGASUS_HOME
    # and CONDOR_HOME are set in the environment or configuration
    get_pegasus_bin()
    get_condor_bin()

def get_script_env():
    PEGASUS_BIN = get_pegasus_bin()
    CONDOR_BIN = get_condor_bin()

    # Add pegasus bin dirs to PATH
    env = dict(os.environ)
    PATH = env.get("PATH", "/bin:/usr/bin:/usr/local/bin")
    PATH = PEGASUS_BIN + ":" + CONDOR_BIN + ":" + PATH
    env["PATH"] = PATH

    return env

def runscript(script, cwd=None, env=None):
    # Make sure the cwd is OK
    if cwd is not None and not os.path.isdir(cwd):
        raise EMException("Working directory does not exist: %s" % cwd)

    if env is None:
        env = dict(os.environ)

    p = subprocess.Popen(script, shell=True, env=env, cwd=cwd)

    rc = p.wait()

    if rc != 0:
        raise EMException("Script failed with exitcode %d" % rc)

def forkscript(script, pidfile=None, cwd=None, env=None):
    # This does a double fork to detach the process from the python
    # interpreter so that we don't have to call wait() on it

    # Make sure the cwd is OK
    if cwd is not None and not os.path.isdir(cwd):
        raise EMException("Working directory does not exist: %s" % cwd)

    if env is None:
        env = dict(os.environ)

    # This is just to ensure we get an exception if there is
    # something wrong with the pidfile
    if pidfile is not None:
        try:
            open(pidfile, "w").close()
        except:
            raise EMException("Unable to write pidfile: %s" % pidfile)

    pid1 = os.fork()
    if pid1 == 0:
        if cwd is not None:
            os.chdir(cwd)

        pid2 = os.fork()
        if pid2 == 0:
            os.execve("/bin/sh", ["/bin/sh","-c",script], env)
            os._exit(255)

        if pidfile is not None:
            f = open(pidfile, "w")
            f.write("%d\n" % pid2)
            f.close()

        os._exit(0)

    pid, exitcode = os.waitpid(pid1, 0)
    if exitcode != 0:
        raise EMException("Non-zero exitcode launching script: %d" % exitcode)

class WorkflowProcessor:
    def __init__(self, db, workflow):
        self.db = db
        self.workflow = workflow

    def get_file(self, filename):
        return os.path.join(self.workflow.basedir, filename)

    def get_bundledir(self):
        return os.path.join(self.workflow.basedir, "bundle")

    def get_pidfile(self):
        return self.get_file("planner.pid")

    def get_resultfile(self):
        return self.get_file("planner.result")

    def get_logfile(self):
        return self.get_file("planner.log")

    def get_planfile(self):
        return self.get_file("plan.sh")

    def plan(self):
        "Launch the pegasus planner"
        workdir = self.get_bundledir()
        pidfile = self.get_pidfile()
        logfile = self.get_logfile()
        planfile = self.get_planfile()
        resultfile = self.get_resultfile()

        if os.path.isfile(pidfile) and self.planning():
            raise EMException("Planner already running")

        # When we re-plan, we need to remove all the old
        # files so that the ensemble manager doesn't get
        # confused.
        files = [
            logfile,
            resultfile,
            pidfile
        ]
        for f in files:
            if os.path.isfile(f):
                os.remove(f)

        script = "%s >%s 2>&1; /bin/echo $? >%s" % (planfile, logfile, resultfile)
        forkscript(script, cwd=workdir, pidfile=pidfile, env=get_script_env())

    def planning(self):
        "Check pidfile to see if the planner is still running"
        pidfile = self.get_pidfile()

        if not os.path.exists(pidfile):
            raise EMException("pidfile missing")

        pid = int(open(pidfile,"r").read())

        try:
            os.kill(pid, 0)
            # If that succeeds, the process is still running
            return True
        except OSError, e:
            # errno 3 is No Such Process
            if e.errno != 3:
                raise

        return False

    def planning_successful(self):
        "Check to see if planning was successful"
        resultfile = self.get_resultfile()

        if not os.path.exists(resultfile):
            raise EMException("Result file not found: %s" % resultfile)

        exitcode = int(open(resultfile, "r").read())

        return exitcode == 0

    def get_submitdir(self):
        "Get the workflow submitdir from the workflow log"
        logfile = self.get_logfile()

        if not os.path.isfile(logfile):
            raise EMException("Workflow log file not found: %s" % logfile)

        submitdir = None

        f = open(logfile, "r")
        try:
            for l in f:
                if l.startswith("pegasus-run"):
                    submitdir = l.split()[1]
        finally:
            f.close()

        if submitdir is None:
            raise EMException("No pegasus-run found in the workflow log: %s" % logfile)

        return submitdir

    def get_wf_uuid(self):
        "Get the workflow UUID from the braindump file"
        submitdir = self.get_submitdir()

        braindump = os.path.join(submitdir, "braindump.txt")

        if not os.path.isfile(braindump):
            raise EMException("braindump.txt not found")

        wf_uuid = None

        f = open(braindump, "r")
        try:
            for l in f:
                if l.startswith("wf_uuid"):
                    wf_uuid = l.split()[1]
        finally:
            f.close()

        if wf_uuid is None:
            raise EMException("wf_uuid not found in braindump.txt")

        return wf_uuid

    def run(self):
        "Run the workflow using pegasus-run"
        submitdir = self.workflow.submitdir

        if submitdir is None:
            raise EMException("Workflow submitdir not set")

        if not os.path.isdir(submitdir):
            raise EMException("Workflow submit dir does not exist: %s" % submitdir)

        logfile = self.get_logfile()

        runscript("pegasus-run %s >>%s 2>&1" % (submitdir,logfile), env=get_script_env())

    def get_dashboard(self):
        "Get the dashboard record for the workflow"
        wf_uuid = self.workflow.wf_uuid
        if wf_uuid is None:
            raise EMException("wf_uuid is none")

        try:
            w = self.db.session.query(DashboardWorkflow)\
                    .filter_by(wf_uuid=str(wf_uuid))\
                    .one()
            return w
        except NoResultFound:
            return None

    def get_dashboard_state(self):
        "Get the latest state of the workflow from the dashboard tables"
        w = self.get_dashboard()
        if w is None:
            raise EMException("Dashboard workflow not found")

        try:
            ws = self.db.session.query(DashboardWorkflowstate)\
                    .filter_by(wf_id=w.wf_id)\
                    .order_by("timestamp desc")\
                    .first()
        except NoResultFound:
            raise EMException("Dashboard workflow state not found")

        return ws

    def pending(self):
        "The workflow is pending if there is no dashboard record"
        w = self.get_dashboard()
        return w is None

    def running(self):
        "Is the workflow running"
        ws = self.get_dashboard_state()
        return ws.state == "WORKFLOW_STARTED"

    def running_successful(self):
        "Assuming the workflow is done running, did it finish successfully?"
        ws = self.get_dashboard_state()
        if ws.state == "WORKFLOW_STARTED":
            raise EMException("Workflow is running")
        return ws.status == 0

class EnsembleProcessor:
    Processor = WorkflowProcessor

    def __init__(self, db, ensemble):
        self.db = db
        self.ensemble = ensemble
        self.active = ensemble.state == EnsembleStates.ACTIVE
        self.max_running = ensemble.max_running
        self.max_planning = ensemble.max_planning
        self.running = 0
        self.planning = 0

        # We are only interested in workflows that are in one
        # of the following states. The EM doesn't handle other
        # states.
        active_states = set((
            EnsembleWorkflowStates.READY,
            EnsembleWorkflowStates.PLANNING,
            EnsembleWorkflowStates.QUEUED,
            EnsembleWorkflowStates.RUNNING
        ))

        # We need a copy of the list so we can sort it
        self.workflows = [w for w in ensemble.workflows if w.state in active_states]

        # Sort workflows by priority
        def get_priority(w):
            return w.priority

        self.workflows.sort(key=get_priority)

        # Count the number of planning and running workflows
        for w in self.workflows:
            if w.state == EnsembleWorkflowStates.PLANNING:
                self.planning += 1
            elif w.state == EnsembleWorkflowStates.RUNNING:
                self.running += 1

    def can_plan(self):
        return self.active and self.planning < self.max_planning

    def plan_workflow(self, workflow):
        log.info("Planning %s" % workflow.name)
        self.planning += 1
        workflow.set_state(EnsembleWorkflowStates.PLANNING)
        workflow.set_updated()

        self.db.session.flush()

        # Fork planning task
        p = self.Processor(self.db, workflow)
        try:
            p.plan()
        except Exception, e:
            log.exception(e)
            workflow.set_state(EnsembleWorkflowStates.PLAN_FAILED)
            workflow.set_updated()

        self.db.session.commit()

    def handle_ready(self, workflow):
        if self.can_plan():
            self.plan_workflow(workflow)

    def handle_planning(self, workflow):
        p = self.Processor(self.db, workflow)

        if p.planning():
            return

        self.planning -= 1

        # Planning failed
        if not p.planning_successful():
            workflow.set_state(EnsembleWorkflowStates.PLAN_FAILED)
            workflow.set_updated()
            self.db.session.commit()
            return

        log.info("Queueing workflow %s" % workflow.name)

        # Planning succeeded, get uuid and queue workflow
        workflow.set_wf_uuid(p.get_wf_uuid())
        workflow.set_submitdir(p.get_submitdir())
        workflow.set_state(EnsembleWorkflowStates.QUEUED)
        workflow.set_updated()
        self.db.session.commit()

        # Go ahead and handle the queued state now
        self.handle_queued(workflow)

    def can_run(self):
        return self.active and self.running < self.max_running

    def run_workflow(self, workflow):
        log.info("Running workflow %s" % workflow.name)

        self.running += 1
        workflow.set_state(EnsembleWorkflowStates.RUNNING)
        workflow.set_updated()
        self.db.session.flush()

        p = self.Processor(self.db, workflow)
        try:
            p.run()
        except Exception, e:
            log.exception(e)
            workflow.set_state(EnsembleWorkflowStates.RUN_FAILED)
            workflow.set_updated()

        self.db.session.commit()

    def handle_queued(self, workflow):
        if self.can_run():
            self.run_workflow(workflow)

    def handle_running(self, workflow):
        p = self.Processor(self.db, workflow)

        if p.pending() or p.running():
            return

        self.running -= 1

        if p.running_successful():
            workflow.set_state(EnsembleWorkflowStates.SUCCESSFUL)
        else:
            workflow.set_state(EnsembleWorkflowStates.FAILED)

        workflow.set_updated()

        self.db.session.commit()

    def handle_workflow(self, w):
        if w.state == EnsembleWorkflowStates.READY:
            self.handle_ready(w)
        elif w.state == EnsembleWorkflowStates.PLANNING:
            self.handle_planning(w)
        elif w.state == EnsembleWorkflowStates.QUEUED:
            self.handle_queued(w)
        elif w.state == EnsembleWorkflowStates.RUNNING:
            self.handle_running(w)

    def run(self):
        log.info("Processing %d ensemble workflows..." % len(self.workflows))
        for w in self.workflows:
            try:
                self.handle_workflow(w)
            except Exception, e:
                log.error("Processing workflow %s of ensemble %s" % (w.name, self.ensemble.name))
                log.exception(e)
                self.db.session.rollback()

class EnsembleManager(object):
    Processor = EnsembleProcessor

    def __init__(self, interval=None):
        if interval is None:
            interval = float(app.config["EM_INTERVAL"])
        self.interval = interval

    def start(self):
        pid = os.fork()
        if pid == 0:
            self.run()
        else:
            return pid

    def run(self):
        log.info("Ensemble Manager starting")
        self.loop_forever()

    def loop_forever(self):
        while True:
            self.loop_once()
            time.sleep(self.interval)

    def get_active_users(self):
        # TODO Identify the active users
        return [request.get_user_by_uid(os.getuid())]

    def loop_once(self):
        for u in self.get_active_users():

            # Switch to user
            if os.getuid() != u.uid:
                if os.getuid() != 0:
                    log.error("Pegasus service must run as root to process ensembles for %s", u.username)
                    continue

                os.setegid(u.gid)
                os.seteuid(u.uid)

            db = Ensembles(u.get_master_db_url())

            actionable = db.list_actionable_ensembles()
            if len(actionable) == 0:
                continue

            log.info("Processing ensembles for %s", u.username)
            for e in actionable:
                log.info("Processing ensemble %s", e.name)
                p = self.Processor(db, e)
                p.run()

<|endoftext|>


<commit_before>import pytest
import salt.states.pkg as pkg
from salt.ext.six.moves import zip
from tests.support.mock import MagicMock, patch

@pytest.fixture
def configure_loader_modules():
    return {pkg: {"__grains__": {"os": "CentOS"}}}

@pytest.fixture(scope="module")
def pkgs():
    return {
        "pkga": {"old": "1.0.1", "new": "2.0.1"},
        "pkgb": {"old": "1.0.2", "new": "2.0.2"},
        "pkgc": {"old": "1.0.3", "new": "2.0.3"},
    }

def test_uptodate_with_changes():
    """
    Test pkg.uptodate with simulated changes
    """

    pkgs = {
        "pkga": {"old": "1.0.1", "new": "2.0.1"},
        "pkgb": {"old": "1.0.2", "new": "2.0.2"},
        "pkgc": {"old": "1.0.3", "new": "2.0.3"},
    }

    list_upgrades = MagicMock(
        return_value={pkgname: pkgver["new"] for pkgname, pkgver in pkgs.items()}
    )
    upgrade = MagicMock(return_value=pkgs)
    version = MagicMock(side_effect=lambda pkgname, **_: pkgs[pkgname]["old"])

    with patch.dict(
        pkg.__salt__,
        {
            "pkg.list_upgrades": list_upgrades,
            "pkg.upgrade": upgrade,
            "pkg.version": version,
        },
    ):

        # Run state with test=false
        with patch.dict(pkg.__opts__, {"test": False}):
            ret = pkg.uptodate("dummy", test=True)
            assert ret["result"]
            assert ret["changes"] == pkgs

        # Run state with test=true
        with patch.dict(pkg.__opts__, {"test": True}):
            ret = pkg.uptodate("dummy", test=True)
            assert ret["result"] is None
            assert ret["changes"] == pkgs

def test_uptodate_with_pkgs_with_changes():
    """
    Test pkg.uptodate with simulated changes
    """

    pkgs = {
        "pkga": {"old": "1.0.1", "new": "2.0.1"},
        "pkgb": {"old": "1.0.2", "new": "2.0.2"},
        "pkgc": {"old": "1.0.3", "new": "2.0.3"},
    }

    list_upgrades = MagicMock(
        return_value={pkgname: pkgver["new"] for pkgname, pkgver in pkgs.items()}
    )
    upgrade = MagicMock(return_value=pkgs)
    version = MagicMock(side_effect=lambda pkgname, **_: pkgs[pkgname]["old"])

    with patch.dict(
        pkg.__salt__,
        {
            "pkg.list_upgrades": list_upgrades,
            "pkg.upgrade": upgrade,
            "pkg.version": version,
        },
    ):
        # Run state with test=false
        with patch.dict(pkg.__opts__, {"test": False}):
            ret = pkg.uptodate("dummy", test=True, pkgs=[pkgname for pkgname in pkgs],)
            assert ret["result"]
            assert ret["changes"] == pkgs

        # Run state with test=true
        with patch.dict(pkg.__opts__, {"test": True}):
            ret = pkg.uptodate("dummy", test=True, pkgs=[pkgname for pkgname in pkgs],)
            assert ret["result"] is None
            assert ret["changes"] == pkgs

def test_uptodate_no_changes():
    """
    Test pkg.uptodate with no changes
    """
    list_upgrades = MagicMock(return_value={})
    upgrade = MagicMock(return_value={})

    with patch.dict(
        pkg.__salt__, {"pkg.list_upgrades": list_upgrades, "pkg.upgrade": upgrade}
    ):

        # Run state with test=false
        with patch.dict(pkg.__opts__, {"test": False}):

            ret = pkg.uptodate("dummy", test=True)
            assert ret["result"]
            assert ret["changes"] == {}

        # Run state with test=true
        with patch.dict(pkg.__opts__, {"test": True}):
            ret = pkg.uptodate("dummy", test=True)
            assert ret["result"]
            assert ret["changes"] == {}

def test_uptodate_with_pkgs_no_changes():
    """
    Test pkg.uptodate with no changes
    """

    pkgs = {
        "pkga": {"old": "1.0.1", "new": "2.0.1"},
        "pkgb": {"old": "1.0.2", "new": "2.0.2"},
        "pkgc": {"old": "1.0.3", "new": "2.0.3"},
    }

    list_upgrades = MagicMock(return_value={})
    upgrade = MagicMock(return_value={})

    with patch.dict(
        pkg.__salt__, {"pkg.list_upgrades": list_upgrades, "pkg.upgrade": upgrade}
    ):
        # Run state with test=false
        with patch.dict(pkg.__opts__, {"test": False}):
            ret = pkg.uptodate("dummy", test=True, pkgs=[pkgname for pkgname in pkgs],)
            assert ret["result"]
            assert ret["changes"] == {}

        # Run state with test=true
        with patch.dict(pkg.__opts__, {"test": True}):
            ret = pkg.uptodate("dummy", test=True, pkgs=[pkgname for pkgname in pkgs],)
            assert ret["result"]
            assert ret["changes"] == {}

def test_uptodate_with_failed_changes():
    """
    Test pkg.uptodate with simulated failed changes
    """

    pkgs = {
        "pkga": {"old": "1.0.1", "new": "2.0.1"},
        "pkgb": {"old": "1.0.2", "new": "2.0.2"},
        "pkgc": {"old": "1.0.3", "new": "2.0.3"},
    }

    list_upgrades = MagicMock(
        return_value={pkgname: pkgver["new"] for pkgname, pkgver in pkgs.items()}
    )
    upgrade = MagicMock(return_value={})
    version = MagicMock(side_effect=lambda pkgname, **_: pkgs[pkgname]["old"])

    with patch.dict(
        pkg.__salt__,
        {
            "pkg.list_upgrades": list_upgrades,
            "pkg.upgrade": upgrade,
            "pkg.version": version,
        },
    ):
        # Run state with test=false
        with patch.dict(pkg.__opts__, {"test": False}):
            ret = pkg.uptodate("dummy", test=True, pkgs=[pkgname for pkgname in pkgs],)
            assert not ret["result"]
            assert ret["changes"] == {}

        # Run state with test=true
        with patch.dict(pkg.__opts__, {"test": True}):
            ret = pkg.uptodate("dummy", test=True, pkgs=[pkgname for pkgname in pkgs],)
            assert ret["result"] is None
            assert ret["changes"] == pkgs

def test_parse_version_string():
    test_parameters = [
        (
            "> 1.0.0, < 15.0.0, != 14.0.1",
            [(">", "1.0.0"), ("<", "15.0.0"), ("!=", "14.0.1")],
        ),
        (
            "> 1.0.0,< 15.0.0,!= 14.0.1",
            [(">", "1.0.0"), ("<", "15.0.0"), ("!=", "14.0.1")],
        ),
        (">= 1.0.0, < 15.0.0", [(">=", "1.0.0"), ("<", "15.0.0")]),
        (">=1.0.0,< 15.0.0", [(">=", "1.0.0"), ("<", "15.0.0")]),
        ("< 15.0.0", [("<", "15.0.0")]),
        ("<15.0.0", [("<", "15.0.0")]),
        ("15.0.0", [("==", "15.0.0")]),
        ("", []),
    ]
    for version_string, expected_version_conditions in test_parameters:
        version_conditions = pkg._parse_version_string(version_string)
        assert len(expected_version_conditions) == len(version_conditions)
        for expected_version_condition, version_condition in zip(
            expected_version_conditions, version_conditions
        ):
            assert expected_version_condition[0] == version_condition[0]
            assert expected_version_condition[1] == version_condition[1]

def test_fulfills_version_string():
    test_parameters = [
        ("> 1.0.0, < 15.0.0, != 14.0.1", [], False),
        ("> 1.0.0, < 15.0.0, != 14.0.1", ["1.0.0"], False),
        ("> 1.0.0, < 15.0.0, != 14.0.1", ["14.0.1"], False),
        ("> 1.0.0, < 15.0.0, != 14.0.1", ["16.0.0"], False),
        ("> 1.0.0, < 15.0.0, != 14.0.1", ["2.0.0"], True),
        ("> 1.0.0, < 15.0.0, != 14.0.1", ["1.0.0", "14.0.1", "16.0.0", "2.0.0"], True,),
        ("> 15.0.0", [], False),
        ("> 15.0.0", ["1.0.0"], False),
        ("> 15.0.0", ["16.0.0"], True),
        ("15.0.0", [], False),
        ("15.0.0", ["15.0.0"], True),
        # No version specified, whatever version installed. This is threated like ANY version installed fulfills.
        ("", ["15.0.0"], True),
        # No version specified, no version installed.
        ("", [], False),
    ]
    for version_string, installed_versions, expected_result in test_parameters:
        msg = "version_string: {}, installed_versions: {}, expected_result: {}".format(
            version_string, installed_versions, expected_result
        )
        assert expected_result == pkg._fulfills_version_string(
            installed_versions, version_string
        ), msg

def test_fulfills_version_spec():
    test_parameters = [
        (["1.0.0", "14.0.1", "16.0.0", "2.0.0"], "==", "1.0.0", True),
        (["1.0.0", "14.0.1", "16.0.0", "2.0.0"], ">=", "1.0.0", True),
        (["1.0.0", "14.0.1", "16.0.0", "2.0.0"], ">", "1.0.0", True),
        (["1.0.0", "14.0.1", "16.0.0", "2.0.0"], "<", "2.0.0", True),
        (["1.0.0", "14.0.1", "16.0.0", "2.0.0"], "<=", "2.0.0", True),
        (["1.0.0", "14.0.1", "16.0.0", "2.0.0"], "!=", "1.0.0", True),
        (["1.0.0", "14.0.1", "16.0.0", "2.0.0"], "==", "17.0.0", False),
        (["1.0.0"], "!=", "1.0.0", False),
        ([], "==", "17.0.0", False),
    ]
    for installed_versions, operator, version, expected_result in test_parameters:
        msg = "installed_versions: {}, operator: {}, version: {}, expected_result: {}".format(
            installed_versions, operator, version, expected_result
        )
        assert expected_result == pkg._fulfills_version_spec(
            installed_versions, operator, version
        ), msg
<commit_msg>Utilize pytest.mark.parametrize fixture<commit_after>import pytest
import salt.states.pkg as pkg
from tests.support.mock import MagicMock, patch

@pytest.fixture
def configure_loader_modules():
    return {pkg: {"__grains__": {"os": "CentOS"}}}

@pytest.fixture(scope="module")
def pkgs():
    return {
        "pkga": {"old": "1.0.1", "new": "2.0.1"},
        "pkgb": {"old": "1.0.2", "new": "2.0.2"},
        "pkgc": {"old": "1.0.3", "new": "2.0.3"},
    }

def test_uptodate_with_changes(pkgs):
    """
    Test pkg.uptodate with simulated changes
    """

    list_upgrades = MagicMock(
        return_value={pkgname: pkgver["new"] for pkgname, pkgver in pkgs.items()}
    )
    upgrade = MagicMock(return_value=pkgs)
    version = MagicMock(side_effect=lambda pkgname, **_: pkgs[pkgname]["old"])

    with patch.dict(
        pkg.__salt__,
        {
            "pkg.list_upgrades": list_upgrades,
            "pkg.upgrade": upgrade,
            "pkg.version": version,
        },
    ):

        # Run state with test=false
        with patch.dict(pkg.__opts__, {"test": False}):
            ret = pkg.uptodate("dummy", test=True)
            assert ret["result"]
            assert ret["changes"] == pkgs

        # Run state with test=true
        with patch.dict(pkg.__opts__, {"test": True}):
            ret = pkg.uptodate("dummy", test=True)
            assert ret["result"] is None
            assert ret["changes"] == pkgs

def test_uptodate_with_pkgs_with_changes(pkgs):
    """
    Test pkg.uptodate with simulated changes
    """

    list_upgrades = MagicMock(
        return_value={pkgname: pkgver["new"] for pkgname, pkgver in pkgs.items()}
    )
    upgrade = MagicMock(return_value=pkgs)
    version = MagicMock(side_effect=lambda pkgname, **_: pkgs[pkgname]["old"])

    with patch.dict(
        pkg.__salt__,
        {
            "pkg.list_upgrades": list_upgrades,
            "pkg.upgrade": upgrade,
            "pkg.version": version,
        },
    ):
        # Run state with test=false
        with patch.dict(pkg.__opts__, {"test": False}):
            ret = pkg.uptodate("dummy", test=True, pkgs=[pkgname for pkgname in pkgs],)
            assert ret["result"]
            assert ret["changes"] == pkgs

        # Run state with test=true
        with patch.dict(pkg.__opts__, {"test": True}):
            ret = pkg.uptodate("dummy", test=True, pkgs=[pkgname for pkgname in pkgs],)
            assert ret["result"] is None
            assert ret["changes"] == pkgs

def test_uptodate_no_changes():
    """
    Test pkg.uptodate with no changes
    """
    list_upgrades = MagicMock(return_value={})
    upgrade = MagicMock(return_value={})

    with patch.dict(
        pkg.__salt__, {"pkg.list_upgrades": list_upgrades, "pkg.upgrade": upgrade}
    ):

        # Run state with test=false
        with patch.dict(pkg.__opts__, {"test": False}):

            ret = pkg.uptodate("dummy", test=True)
            assert ret["result"]
            assert ret["changes"] == {}

        # Run state with test=true
        with patch.dict(pkg.__opts__, {"test": True}):
            ret = pkg.uptodate("dummy", test=True)
            assert ret["result"]
            assert ret["changes"] == {}

def test_uptodate_with_pkgs_no_changes(pkgs):
    """
    Test pkg.uptodate with no changes
    """

    list_upgrades = MagicMock(return_value={})
    upgrade = MagicMock(return_value={})

    with patch.dict(
        pkg.__salt__, {"pkg.list_upgrades": list_upgrades, "pkg.upgrade": upgrade}
    ):
        # Run state with test=false
        with patch.dict(pkg.__opts__, {"test": False}):
            ret = pkg.uptodate("dummy", test=True, pkgs=[pkgname for pkgname in pkgs],)
            assert ret["result"]
            assert ret["changes"] == {}

        # Run state with test=true
        with patch.dict(pkg.__opts__, {"test": True}):
            ret = pkg.uptodate("dummy", test=True, pkgs=[pkgname for pkgname in pkgs],)
            assert ret["result"]
            assert ret["changes"] == {}

def test_uptodate_with_failed_changes(pkgs):
    """
    Test pkg.uptodate with simulated failed changes
    """

    list_upgrades = MagicMock(
        return_value={pkgname: pkgver["new"] for pkgname, pkgver in pkgs.items()}
    )
    upgrade = MagicMock(return_value={})
    version = MagicMock(side_effect=lambda pkgname, **_: pkgs[pkgname]["old"])

    with patch.dict(
        pkg.__salt__,
        {
            "pkg.list_upgrades": list_upgrades,
            "pkg.upgrade": upgrade,
            "pkg.version": version,
        },
    ):
        # Run state with test=false
        with patch.dict(pkg.__opts__, {"test": False}):
            ret = pkg.uptodate("dummy", test=True, pkgs=[pkgname for pkgname in pkgs],)
            assert not ret["result"]
            assert ret["changes"] == {}

        # Run state with test=true
        with patch.dict(pkg.__opts__, {"test": True}):
            ret = pkg.uptodate("dummy", test=True, pkgs=[pkgname for pkgname in pkgs],)
            assert ret["result"] is None
            assert ret["changes"] == pkgs

@pytest.mark.parametrize(
    "version_string, expected_version_conditions",
    [
        (
            "> 1.0.0, < 15.0.0, != 14.0.1",
            [(">", "1.0.0"), ("<", "15.0.0"), ("!=", "14.0.1")],
        ),
        (
            "> 1.0.0,< 15.0.0,!= 14.0.1",
            [(">", "1.0.0"), ("<", "15.0.0"), ("!=", "14.0.1")],
        ),
        (">= 1.0.0, < 15.0.0", [(">=", "1.0.0"), ("<", "15.0.0")]),
        (">=1.0.0,< 15.0.0", [(">=", "1.0.0"), ("<", "15.0.0")]),
        ("< 15.0.0", [("<", "15.0.0")]),
        ("<15.0.0", [("<", "15.0.0")]),
        ("15.0.0", [("==", "15.0.0")]),
        ("", []),
    ],
)
def test_parse_version_string(version_string, expected_version_conditions):
    version_conditions = pkg._parse_version_string(version_string)
    assert len(expected_version_conditions) == len(version_conditions)
    for expected_version_condition, version_condition in zip(
        expected_version_conditions, version_conditions
    ):
        assert expected_version_condition[0] == version_condition[0]
        assert expected_version_condition[1] == version_condition[1]

@pytest.mark.parametrize(
    "version_string, installed_versions, expected_result",
    [
        ("> 1.0.0, < 15.0.0, != 14.0.1", [], False),
        ("> 1.0.0, < 15.0.0, != 14.0.1", ["1.0.0"], False),
        ("> 1.0.0, < 15.0.0, != 14.0.1", ["14.0.1"], False),
        ("> 1.0.0, < 15.0.0, != 14.0.1", ["16.0.0"], False),
        ("> 1.0.0, < 15.0.0, != 14.0.1", ["2.0.0"], True),
        ("> 1.0.0, < 15.0.0, != 14.0.1", ["1.0.0", "14.0.1", "16.0.0", "2.0.0"], True,),
        ("> 15.0.0", [], False),
        ("> 15.0.0", ["1.0.0"], False),
        ("> 15.0.0", ["16.0.0"], True),
        ("15.0.0", [], False),
        ("15.0.0", ["15.0.0"], True),
        # No version specified, whatever version installed. This is threated like ANY version installed fulfills.
        ("", ["15.0.0"], True),
        # No version specified, no version installed.
        ("", [], False),
    ],
)
def test_fulfills_version_string(version_string, installed_versions, expected_result):
    msg = "version_string: {}, installed_versions: {}, expected_result: {}".format(
        version_string, installed_versions, expected_result
    )
    assert expected_result == pkg._fulfills_version_string(
        installed_versions, version_string
    ), msg

@pytest.mark.parametrize(
    "installed_versions, operator, version, expected_result",
    [
        (["1.0.0", "14.0.1", "16.0.0", "2.0.0"], "==", "1.0.0", True),
        (["1.0.0", "14.0.1", "16.0.0", "2.0.0"], ">=", "1.0.0", True),
        (["1.0.0", "14.0.1", "16.0.0", "2.0.0"], ">", "1.0.0", True),
        (["1.0.0", "14.0.1", "16.0.0", "2.0.0"], "<", "2.0.0", True),
        (["1.0.0", "14.0.1", "16.0.0", "2.0.0"], "<=", "2.0.0", True),
        (["1.0.0", "14.0.1", "16.0.0", "2.0.0"], "!=", "1.0.0", True),
        (["1.0.0", "14.0.1", "16.0.0", "2.0.0"], "==", "17.0.0", False),
        (["1.0.0"], "!=", "1.0.0", False),
        ([], "==", "17.0.0", False),
    ],
)
def test_fulfills_version_spec(installed_versions, operator, version, expected_result):
    msg = "installed_versions: {}, operator: {}, version: {}, expected_result: {}".format(
        installed_versions, operator, version, expected_result
    )
    assert expected_result == pkg._fulfills_version_spec(
        installed_versions, operator, version
    ), msg
<|endoftext|>


<commit_before>from req import Service
from service.base import BaseService
import config
import os

class Problem(BaseService):
    def get_problem_list(self):
        res = yield self.db.execute("SELECT * FROM problems ORDER BY id ASC")
        res = res.fetchall()
        return (None, res)

    def get_problem(self, data={}):
        required_args = [{
            'name': '+id',
            'type': int,
        }]
        err = self.form_validation(data, required_args)
        if err: return (err, None)
        res = yield self.db.execute("SELECT * FROM problems WHERE id=%s", (data['id'],))
        res = res.fetchone()
        return (None, res)

    def post_problem(self, data={}):
        required_args = [{
            'name': '+title',
            'type': str,
        }, {
            'name': '+score_type',
            'type': int
        }, {
            'name': '+pdf',
        }]
        err = self.form_validation(data, required_args)
        if err: return (err, None)
        pdf = data.pop('pdf')
        sql, param = self.gen_insert_sql('problems', data)
        res = yield self.db.execute(sql, param)
        res = res.fetchone()
        self.log(res)
        folder = '%s/data/problems/%s'%(config.DATA_ROOT, str(res['id']))
        file_path = '%s/pdf.pdf'%(folder)
        try: os.makedirs(folder)
        except: pass
        with open(file_path, 'wb+') as f:
            f.write(pdf['body'])
        return (None, res)

    def put_problem(self, data):
        required_args = [{
            'name': '+id',
            'type': int
        }, {
            'name': '+title',
            'type': str,
        }, {
            'name': '+score_type',
            'type': int
        }, {
            'name': 'pdf',
        }]
        err = self.form_validation(data, required_args)
        if err: return (err, None)
        if 'pdf' in data:
            pdf = data.pop('pdf')
        else:
            pdf = None
        id = data.pop('id')
        sql, param = self.gen_update_sql('problems', data)
        res = yield self.db.execute(sql + ' WHERE id = %s', param + (id,))
        if pdf:
            folder = '%s/data/problems/%s'%(config.DATA_ROOT, str(id))
            file_path = '%s/pdf.pdf'%(folder)
            try: os.makedirs(folder)
            except: pass
            with open(file_path, 'wb+') as f:
                f.write(pdf['body'])
        return (None, None)

    def get_problem_execute(self, data={}):
        required_args = [{
            'name': '+id',
            'type': int
        }]
        err = self.form_validation(data, required_args)
        if err: return (err, None)
        res = yield self.db.execute("SELECT e.id, e.description FROM map_problem_execute as me, execute_types as e WHERE e.id=me.execute_type_id and me.problem_id=%s", (data['id'],))
        res = res.fetchall()
        return (None, res)

    def put_problem_execute(self, data={}):
        required_args = [{
            'name': '+id',
            'type': int
        }, {
            'name': '+executes',
            'type': list
        }]
        self.log(data)
        err = self.form_validation(data, required_args)
        if err: return (err, None)
        yield self.db.execute("DELETE FROM map_problem_execute WHERE problem_id=%s", (data['id'],))
        for x in data['executes']:
            try:
                yield self.db.execute("INSERT INTO map_problem_execute (problem_id, execute_type_id) VALUES (%s, %s)", (data['id'], x,))
            except:
                pass
        return (None, None)
<commit_msg>return executes in problem<commit_after>from req import Service
from service.base import BaseService
import config
import os

class Problem(BaseService):
    def get_problem_list(self):
        res = yield self.db.execute("SELECT * FROM problems ORDER BY id ASC")
        res = res.fetchall()
        return (None, res)

    def get_problem(self, data={}):
        required_args = [{
            'name': '+id',
            'type': int,
        }]
        err = self.form_validation(data, required_args)
        if err: return (err, None)
        res = yield self.db.execute("SELECT * FROM problems WHERE id=%s", (data['id'],))
        res = res.fetchone()
        res['executes'] = yield from self.get_problem_execute(data)
        return (None, res)

    def post_problem(self, data={}):
        required_args = [{
            'name': '+title',
            'type': str,
        }, {
            'name': '+score_type',
            'type': int
        }, {
            'name': '+pdf',
        }]
        err = self.form_validation(data, required_args)
        if err: return (err, None)
        pdf = data.pop('pdf')
        sql, param = self.gen_insert_sql('problems', data)
        res = yield self.db.execute(sql, param)
        res = res.fetchone()
        self.log(res)
        folder = '%s/data/problems/%s'%(config.DATA_ROOT, str(res['id']))
        file_path = '%s/pdf.pdf'%(folder)
        try: os.makedirs(folder)
        except: pass
        with open(file_path, 'wb+') as f:
            f.write(pdf['body'])
        return (None, res)

    def put_problem(self, data):
        required_args = [{
            'name': '+id',
            'type': int
        }, {
            'name': '+title',
            'type': str,
        }, {
            'name': '+score_type',
            'type': int
        }, {
            'name': 'pdf',
        }]
        err = self.form_validation(data, required_args)
        if err: return (err, None)
        if 'pdf' in data:
            pdf = data.pop('pdf')
        else:
            pdf = None
        id = data.pop('id')
        sql, param = self.gen_update_sql('problems', data)
        res = yield self.db.execute(sql + ' WHERE id = %s', param + (id,))
        if pdf:
            folder = '%s/data/problems/%s'%(config.DATA_ROOT, str(id))
            file_path = '%s/pdf.pdf'%(folder)
            try: os.makedirs(folder)
            except: pass
            with open(file_path, 'wb+') as f:
                f.write(pdf['body'])
        return (None, None)

    def get_problem_execute(self, data={}):
        required_args = [{
            'name': '+id',
            'type': int
        }]
        err = self.form_validation(data, required_args)
        if err: return (err, None)
        res = yield self.db.execute("SELECT e.id, e.description FROM map_problem_execute as me, execute_types as e WHERE e.id=me.execute_type_id and me.problem_id=%s", (data['id'],))
        res = res.fetchall()
        return (None, res)

    def put_problem_execute(self, data={}):
        required_args = [{
            'name': '+id',
            'type': int
        }, {
            'name': '+executes',
            'type': list
        }]
        self.log(data)
        err = self.form_validation(data, required_args)
        if err: return (err, None)
        yield self.db.execute("DELETE FROM map_problem_execute WHERE problem_id=%s", (data['id'],))
        for x in data['executes']:
            try:
                yield self.db.execute("INSERT INTO map_problem_execute (problem_id, execute_type_id) VALUES (%s, %s)", (data['id'], x,))
            except:
                pass
        return (None, None)
<|endoftext|>


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

##PACKAGES##
from __future__ import division
import sys
import warnings
import re
import itertools
import seaborn as sns
import pandas as pd
from matplotlib.lines import Line2D
import matplotlib.animation as animation
import matplotlib.pyplot as plt
from .._shared.helpers import *
from ..tools.cluster import cluster
from ..tools.df2mat import df2mat
from ..tools.reduce import reduce as reducer
from ..tools.normalize import normalize as normalizer
from ..tools.align import align as aligner
from .draw import draw

def plot(x, fmt=None, marker=None, markers=None, linestyle=None,
         linestyles=None, color=None, colors=None, palette='hls', group=None,
         labels=None, legend=None, title=None, elev=10, azim=-60, ndims=None,
         model='IncrementalPCA', model_params={}, align=False, normalize=False,
         n_clusters=None, save_path=None, animate=False, duration=30, tail_duration=2,
         rotations=2, zoom=1, chemtrails=False, precog=False, bullettime=False,
         frame_rate=50, explore=False, show=True):
    """
    Plots dimensionality reduced data and parses plot arguments

    Parameters
    ----------
    x : Numpy array, DataFrame or list of arrays/dfs
        Data for the plot. The form should be samples (rows) by features (cols).

    fmt : str or list of strings
        A list of format strings.  All matplotlib format strings are supported.

    linestyle(s) : str or list of str
        A list of line styles

    marker(s) : str or list of str
        A list of marker types

    color(s) : str or list of str
        A list of marker types

    palette : str
        A matplotlib or seaborn color palette

    group : str/int/float or list
        A list of group labels. Length must match the number of rows in your
        dataset. If the data type is numerical, the values will be mapped to
        rgb values in the specified palette. If the data type is strings,
        the points will be labeled categorically. To label a subset of points,
        use None (i.e. ['a', None, 'b','a']).

    labels : list
        A list of labels for each point. Must be dimensionality of data (x).
        If no label is wanted for a particular point, input None.

    legend : list or bool
        If set to True, legend is implicitly computed from data. Passing a
        list will add string labels to the legend (one for each list item).

    title : str
        A title for the plot

    ndims : int
        An `int` representing the number of dims to reduce the data x
        to. If ndims > 3, will plot in 3 dimensions but return the higher
        dimensional data. Default is None, which will plot data in 3
        dimensions and return the data with the same number of dimensions
        possibly normalized and/or aligned according to normalize/align
        kwargs.

    model : str
        Decomposition/manifold learning model to use.  Models supported: PCA,
        IncrementalPCA, SparsePCA, MiniBatchSparsePCA, KernelPCA, FastICA,
        FactorAnalysis, TruncatedSVD, DictionaryLearning, MiniBatchDictionaryLearning,
        TSNE, Isomap, SpectralEmbedding, LocallyLinearEmbedding, and MDS.

    model_params : dict
        Optional dictionary of scikit-learn parameters to pass to reduction model.
        See scikit-learn specific model docs for details.

    align : bool
        If set to True, data will be run through the ``hyperalignment''
        algorithm implemented in hypertools.tools.align (default: False).

    normalize : str or False
        If set to 'across', the columns of the input data will be z-scored
        across lists (default). If set to 'within', the columns will be
        z-scored within each list that is passed. If set to 'row', each row of
        the input data will be z-scored. If set to False, the input data will
        be returned (default is False).

    n_clusters : int
        If n_clusters is passed, HyperTools will perform k-means clustering
        with the k parameter set to n_clusters. The resulting clusters will
        be plotted in different colors according to the color palette.

    save_path str :
        Path to save the image/movie. Must include the file extension in the
        save path (i.e. save_path='/path/to/file/image.png'). NOTE: If saving
        an animation, FFMPEG must be installed (this is a matplotlib req).
        FFMPEG can be easily installed on a mac via homebrew brew install
        ffmpeg or linux via apt-get apt-get install ffmpeg. If you don't
        have homebrew (mac only), you can install it like this:
        /usr/bin/ruby -e "$(curl -fsSL
        https://raw.githubusercontent.com/Homebrew/install/master/install)".

    animate : bool, 'parallel' or 'spin'
        If True or 'parallel', plots the data as an animated trajectory, with
        each dataset plotted simultaneously. If 'spin', all the data is plotted
        at once but the camera spins around the plot (default: False).

    duration (animation only) : float
        Length of the animation in seconds (default: 30 seconds)

    tail_duration (animation only) : float
        Sets the length of the tail of the data (default: 2 seconds)

    rotations (animation only) : float
        Number of rotations around the box (default: 2)

    zoom (animation only) : float
        How far to zoom into the plot, positive numbers will zoom in (default: 0)

    chemtrails (animation only) : bool
        A low-opacity trail is left behind the trajectory (default: False).

    precog (animation only) : bool
        A low-opacity trail is plotted ahead of the trajectory (default: False).

    bullettime (animation only) : bool
        A low-opacity trail is plotted ahead and behind the trajectory
        (default: False).

    frame_rate (animation only) : int or float
        Frame rate for animation (default: 50)

    explore : bool
        Displays user defined labels will appear on hover. If no labels are
        passed, the point index and coordinate will be plotted. To use,
        set explore=True. Note: Explore mode is currently only supported
        for 3D static plots, and is an experimental feature (i.e it may not yet
        work properly).

    show : bool
        If set to False, the figure will not be displayed, but the figure,
        axis and data objects will still be returned (default: True).

    Returns
    ----------
    fig, ax, data, line_ani : matplotlib.figure.figure, matplotlib.axis.axes, numpy.array, matplotlib.animation.funcanimation
        The plot function outputs a figure handle, axis handle, data, and line
        animation object.  The line animation object is None if animation=False.

    """

    # turn data into common format - a list of arrays
    x = format_data(x)

    # normalize
    x = normalizer(x, normalize=normalize, internal=True)

    # reduce data to ndims
    if ndims is not None:
        x = reducer(x, ndims=ndims, model=model, model_params=model_params,
                    internal=True)

    # align data
    if align:
        if len(x) == 1:
            warn('Data in list of length 1 can not be aligned. '
                 'Skipping the alignment.')
        else:
            x = aligner(x)
    # Return data that has been normalized and possibly reduced and/or aligned
    return_data = x

    # catch all matplotlib kwargs here to pass on
    mpl_kwargs = {}

    # handle color (to be passed onto matplotlib)
    if color is not None:
        mpl_kwargs['color'] = color
        if colors is not None:
            mpl_kwargs['color'] = colors
            warnings.warn('Both color and colors defined: color will be ignored \
                          in favor of colors.')

    # handle linestyle (to be passed onto matplotlib)
    if linestyle is not None:
        mpl_kwargs['linestyle'] = linestyle
        if linestyles is not None:
            mpl_kwargs['linestyle'] = linestyles
            warnings.warn('Both linestyle and linestyles defined: linestyle  \
                          will be ignored in favor of linestyles.')

    # handle marker (to be passed onto matplotlib)
    if marker is not None:
        mpl_kwargs['marker'] = marker
        if markers is not None:
            mpl_kwargs['marker'] = markers
            warnings.warn('Both marker and markers defined: marker will be \
                          ignored in favor of markers.')

    # reduce data to 3 dims for plotting, if ndims is None, return this
    if (ndims and ndims > 3) or (ndims is None and x[0].shape[1] > 3):
        x = reducer(x, ndims=3, model=model, model_params=model_params, internal=True)

    # find cluster and reshape if n_clusters
    if n_clusters is not None:
        cluster_labels = cluster(x, n_clusters=n_clusters, ndims=ndims)
        x = reshape_data(x, cluster_labels)
        if group:
            warnings.warn('n_clusters overrides group, ignoring group.')

    # group data if there is a grouping var
    if group is not None:

        if color is not None:
            warnings.warn("Using group, color keyword will be ignored.")

        # if list of lists, unpack
        if any(isinstance(el, list) for el in group):
            group = list(itertools.chain(*group))

        # if all of the elements are numbers, map them to colors
        if all(isinstance(el, int) or isinstance(el, float) for el in group):
            group = vals2bins(group)
        elif all(isinstance(el, str) for el in group):
            group = group_by_category(group)

        # reshape the data according to group
        x = reshape_data(x, group)

        # interpolate lines if they are grouped
        if is_line(fmt):
            x = patch_lines(x)

    # handle legend
    if legend is not None:
        if legend is False:
            legend = None
        elif legend is True and group is not None:
            legend = [item for item in sorted(set(group), key=list(group).index)]
        elif legend is True and group is None:
            legend = [i + 1 for i in range(len(x))]

        mpl_kwargs['label'] = legend

    # interpolate if its a line plot
    if fmt is None or type(fmt) is str:
        if is_line(fmt):
            if x[0].shape[0] > 1:
                x = interp_array_list(x, interp_val=frame_rate*duration/(x[0].shape[0] - 1))
    elif type(fmt) is list:
        for idx, xi in enumerate(x):
            if is_line(fmt[idx]):
                if xi.shape[0] > 1:
                    x[idx] = interp_array_list(xi, interp_val=frame_rate*duration/(xi.shape[0] - 1))

    # handle explore flag
    if explore:
        assert x[0].shape[1] is 3, "Explore mode is currently only supported for 3D plots."
        mpl_kwargs['picker']=True

    # center
    x = center(x)

    # scale
    x = scale(x)

    # handle palette with seaborn
    sns.set_palette(palette=palette, n_colors=len(x))
    sns.set_style(style='whitegrid')

    # turn kwargs into a list
    kwargs_list = parse_kwargs(x, mpl_kwargs)

    # handle format strings
    if fmt is not None:
        if type(fmt) is not list:
            fmt = [fmt for i in x]

    # draw the plot
    fig, ax, data, line_ani = draw(x, fmt=fmt,
                            kwargs_list=kwargs_list,
                            labels=labels,
                            legend=legend,
                            title=title,
                            animate=animate,
                            duration=duration,
                            tail_duration=tail_duration,
                            rotations=rotations,
                            zoom=zoom,
                            chemtrails=chemtrails,
                            precog=precog,
                            bullettime=bullettime,
                            frame_rate=frame_rate,
                            elev=elev,
                            azim=azim,
                            explore=explore,
                            show=show)

    # tighten layout
    plt.tight_layout()

    # save
    if save_path is not None:
        if animate:
            Writer = animation.writers['ffmpeg']
            writer = Writer(fps=frame_rate, bitrate=1800)
            line_ani.save(save_path, writer=writer)

        else:
            plt.savefig(save_path)

    # show the plot
    if show:
        plt.show()

    return fig, ax, return_data, line_ani
<commit_msg>fixed bug with second plot call showing two plots<commit_after>#!/usr/bin/env python

##PACKAGES##
from __future__ import division
import sys
import warnings
import re
import itertools
import seaborn as sns
import pandas as pd
from matplotlib.lines import Line2D
import matplotlib.animation as animation
import matplotlib.pyplot as plt
from .._shared.helpers import *
from ..tools.cluster import cluster
from ..tools.df2mat import df2mat
from ..tools.reduce import reduce as reducer
from ..tools.normalize import normalize as normalizer
from ..tools.align import align as aligner
from .draw import draw

def plot(x, fmt=None, marker=None, markers=None, linestyle=None,
         linestyles=None, color=None, colors=None, palette='hls', group=None,
         labels=None, legend=None, title=None, elev=10, azim=-60, ndims=None,
         model='IncrementalPCA', model_params={}, align=False, normalize=False,
         n_clusters=None, save_path=None, animate=False, duration=30, tail_duration=2,
         rotations=2, zoom=1, chemtrails=False, precog=False, bullettime=False,
         frame_rate=50, explore=False, show=True):
    """
    Plots dimensionality reduced data and parses plot arguments

    Parameters
    ----------
    x : Numpy array, DataFrame or list of arrays/dfs
        Data for the plot. The form should be samples (rows) by features (cols).

    fmt : str or list of strings
        A list of format strings.  All matplotlib format strings are supported.

    linestyle(s) : str or list of str
        A list of line styles

    marker(s) : str or list of str
        A list of marker types

    color(s) : str or list of str
        A list of marker types

    palette : str
        A matplotlib or seaborn color palette

    group : str/int/float or list
        A list of group labels. Length must match the number of rows in your
        dataset. If the data type is numerical, the values will be mapped to
        rgb values in the specified palette. If the data type is strings,
        the points will be labeled categorically. To label a subset of points,
        use None (i.e. ['a', None, 'b','a']).

    labels : list
        A list of labels for each point. Must be dimensionality of data (x).
        If no label is wanted for a particular point, input None.

    legend : list or bool
        If set to True, legend is implicitly computed from data. Passing a
        list will add string labels to the legend (one for each list item).

    title : str
        A title for the plot

    ndims : int
        An `int` representing the number of dims to reduce the data x
        to. If ndims > 3, will plot in 3 dimensions but return the higher
        dimensional data. Default is None, which will plot data in 3
        dimensions and return the data with the same number of dimensions
        possibly normalized and/or aligned according to normalize/align
        kwargs.

    model : str
        Decomposition/manifold learning model to use.  Models supported: PCA,
        IncrementalPCA, SparsePCA, MiniBatchSparsePCA, KernelPCA, FastICA,
        FactorAnalysis, TruncatedSVD, DictionaryLearning, MiniBatchDictionaryLearning,
        TSNE, Isomap, SpectralEmbedding, LocallyLinearEmbedding, and MDS.

    model_params : dict
        Optional dictionary of scikit-learn parameters to pass to reduction model.
        See scikit-learn specific model docs for details.

    align : bool
        If set to True, data will be run through the ``hyperalignment''
        algorithm implemented in hypertools.tools.align (default: False).

    normalize : str or False
        If set to 'across', the columns of the input data will be z-scored
        across lists (default). If set to 'within', the columns will be
        z-scored within each list that is passed. If set to 'row', each row of
        the input data will be z-scored. If set to False, the input data will
        be returned (default is False).

    n_clusters : int
        If n_clusters is passed, HyperTools will perform k-means clustering
        with the k parameter set to n_clusters. The resulting clusters will
        be plotted in different colors according to the color palette.

    save_path str :
        Path to save the image/movie. Must include the file extension in the
        save path (i.e. save_path='/path/to/file/image.png'). NOTE: If saving
        an animation, FFMPEG must be installed (this is a matplotlib req).
        FFMPEG can be easily installed on a mac via homebrew brew install
        ffmpeg or linux via apt-get apt-get install ffmpeg. If you don't
        have homebrew (mac only), you can install it like this:
        /usr/bin/ruby -e "$(curl -fsSL
        https://raw.githubusercontent.com/Homebrew/install/master/install)".

    animate : bool, 'parallel' or 'spin'
        If True or 'parallel', plots the data as an animated trajectory, with
        each dataset plotted simultaneously. If 'spin', all the data is plotted
        at once but the camera spins around the plot (default: False).

    duration (animation only) : float
        Length of the animation in seconds (default: 30 seconds)

    tail_duration (animation only) : float
        Sets the length of the tail of the data (default: 2 seconds)

    rotations (animation only) : float
        Number of rotations around the box (default: 2)

    zoom (animation only) : float
        How far to zoom into the plot, positive numbers will zoom in (default: 0)

    chemtrails (animation only) : bool
        A low-opacity trail is left behind the trajectory (default: False).

    precog (animation only) : bool
        A low-opacity trail is plotted ahead of the trajectory (default: False).

    bullettime (animation only) : bool
        A low-opacity trail is plotted ahead and behind the trajectory
        (default: False).

    frame_rate (animation only) : int or float
        Frame rate for animation (default: 50)

    explore : bool
        Displays user defined labels will appear on hover. If no labels are
        passed, the point index and coordinate will be plotted. To use,
        set explore=True. Note: Explore mode is currently only supported
        for 3D static plots, and is an experimental feature (i.e it may not yet
        work properly).

    show : bool
        If set to False, the figure will not be displayed, but the figure,
        axis and data objects will still be returned (default: True).

    Returns
    ----------
    fig, ax, data, line_ani : matplotlib.figure.figure, matplotlib.axis.axes, numpy.array, matplotlib.animation.funcanimation
        The plot function outputs a figure handle, axis handle, data, and line
        animation object.  The line animation object is None if animation=False.

    """

    # turn data into common format - a list of arrays
    x = format_data(x)

    # normalize
    x = normalizer(x, normalize=normalize, internal=True)

    # reduce data to ndims
    if ndims is not None:
        x = reducer(x, ndims=ndims, model=model, model_params=model_params,
                    internal=True)

    # align data
    if align:
        if len(x) == 1:
            warn('Data in list of length 1 can not be aligned. '
                 'Skipping the alignment.')
        else:
            x = aligner(x)
    # Return data that has been normalized and possibly reduced and/or aligned
    return_data = x

    # catch all matplotlib kwargs here to pass on
    mpl_kwargs = {}

    # handle color (to be passed onto matplotlib)
    if color is not None:
        mpl_kwargs['color'] = color
        if colors is not None:
            mpl_kwargs['color'] = colors
            warnings.warn('Both color and colors defined: color will be ignored \
                          in favor of colors.')

    # handle linestyle (to be passed onto matplotlib)
    if linestyle is not None:
        mpl_kwargs['linestyle'] = linestyle
        if linestyles is not None:
            mpl_kwargs['linestyle'] = linestyles
            warnings.warn('Both linestyle and linestyles defined: linestyle  \
                          will be ignored in favor of linestyles.')

    # handle marker (to be passed onto matplotlib)
    if marker is not None:
        mpl_kwargs['marker'] = marker
        if markers is not None:
            mpl_kwargs['marker'] = markers
            warnings.warn('Both marker and markers defined: marker will be \
                          ignored in favor of markers.')

    # reduce data to 3 dims for plotting, if ndims is None, return this
    if (ndims and ndims > 3) or (ndims is None and x[0].shape[1] > 3):
        x = reducer(x, ndims=3, model=model, model_params=model_params, internal=True)

    # find cluster and reshape if n_clusters
    if n_clusters is not None:
        cluster_labels = cluster(x, n_clusters=n_clusters, ndims=ndims)
        x = reshape_data(x, cluster_labels)
        if group:
            warnings.warn('n_clusters overrides group, ignoring group.')

    # group data if there is a grouping var
    if group is not None:

        if color is not None:
            warnings.warn("Using group, color keyword will be ignored.")

        # if list of lists, unpack
        if any(isinstance(el, list) for el in group):
            group = list(itertools.chain(*group))

        # if all of the elements are numbers, map them to colors
        if all(isinstance(el, int) or isinstance(el, float) for el in group):
            group = vals2bins(group)
        elif all(isinstance(el, str) for el in group):
            group = group_by_category(group)

        # reshape the data according to group
        x = reshape_data(x, group)

        # interpolate lines if they are grouped
        if is_line(fmt):
            x = patch_lines(x)

    # handle legend
    if legend is not None:
        if legend is False:
            legend = None
        elif legend is True and group is not None:
            legend = [item for item in sorted(set(group), key=list(group).index)]
        elif legend is True and group is None:
            legend = [i + 1 for i in range(len(x))]

        mpl_kwargs['label'] = legend

    # interpolate if its a line plot
    if fmt is None or type(fmt) is str:
        if is_line(fmt):
            if x[0].shape[0] > 1:
                x = interp_array_list(x, interp_val=frame_rate*duration/(x[0].shape[0] - 1))
    elif type(fmt) is list:
        for idx, xi in enumerate(x):
            if is_line(fmt[idx]):
                if xi.shape[0] > 1:
                    x[idx] = interp_array_list(xi, interp_val=frame_rate*duration/(xi.shape[0] - 1))

    # handle explore flag
    if explore:
        assert x[0].shape[1] is 3, "Explore mode is currently only supported for 3D plots."
        mpl_kwargs['picker']=True

    # center
    x = center(x)

    # scale
    x = scale(x)

    # handle palette with seaborn
    sns.set_palette(palette=palette, n_colors=len(x))
    sns.set_style(style='whitegrid')

    # turn kwargs into a list
    kwargs_list = parse_kwargs(x, mpl_kwargs)

    # handle format strings
    if fmt is not None:
        if type(fmt) is not list:
            fmt = [fmt for i in x]

    # draw the plot
    fig, ax, data, line_ani = draw(x, fmt=fmt,
                            kwargs_list=kwargs_list,
                            labels=labels,
                            legend=legend,
                            title=title,
                            animate=animate,
                            duration=duration,
                            tail_duration=tail_duration,
                            rotations=rotations,
                            zoom=zoom,
                            chemtrails=chemtrails,
                            precog=precog,
                            bullettime=bullettime,
                            frame_rate=frame_rate,
                            elev=elev,
                            azim=azim,
                            explore=explore,
                            show=show)

    # tighten layout
    plt.tight_layout()

    # save
    if save_path is not None:
        if animate:
            Writer = animation.writers['ffmpeg']
            writer = Writer(fps=frame_rate, bitrate=1800)
            line_ani.save(save_path, writer=writer)

        else:
            plt.savefig(save_path)

    # show the plot
    if show:
        plt.show()
    else:
        plt.close()

    return fig, ax, return_data, line_ani
<|endoftext|>