<commit_before>from __future__ import print_function

import os
import time
import subprocess
import sys

from colorama import Fore, Style
from watchdog.events import (
    FileSystemEventHandler, FileModifiedEvent, FileCreatedEvent,
    FileMovedEvent, FileDeletedEvent)
from watchdog.observers import Observer
from watchdog.observers.polling import PollingObserver

from .spooler import EventSpooler

EVENT_NAMES = {
    FileModifiedEvent: 'modified',
    FileCreatedEvent: 'created',
    FileMovedEvent: 'moved',
    FileDeletedEvent: 'deleted',
}
WATCHED_EVENTS = list(EVENT_NAMES)
DEFAULT_EXTENSIONS = ['.py']
CLEAR_COMMAND = 'cls' if os.name == 'nt' else 'clear'
BEEP_CHARACTER = '\a'
STYLE_NORMAL = Fore.RESET
STYLE_HIGHLIGHT = Fore.CYAN + Style.NORMAL + Style.BRIGHT

class ChangeHandler(FileSystemEventHandler):
    """
    Listens for changes to files and re-runs tests after each change.
    """
    def __init__(self, auto_clear=False, beep_on_failure=True, onpass=None,
                 onfail=None, onsigint=None, runner=None, beforerun=None,
                 extensions=[], args=None, spool=True, verbose=False,
                 quiet=False):
        super(ChangeHandler, self).__init__()
        self.auto_clear = auto_clear
        self.beep_on_failure = beep_on_failure
        self.onpass = onpass
        self.onfail = onfail
        self.onsigint = onsigint
        self.runner = runner
        self.beforerun = beforerun
        self.extensions = extensions or DEFAULT_EXTENSIONS
        self.args = args or []
        self.spooler = None
        if spool:
            self.spooler = EventSpooler(0.2, self.on_queued_events)
        self.verbose = verbose
        self.quiet = quiet

    def on_queued_events(self, events):
        summary = []
        for event in events:
            paths = [event.src_path]
            if isinstance(event, FileMovedEvent):
                paths.append(event.dest_path)
            event_name = EVENT_NAMES[type(event)]
            paths = tuple(map(os.path.relpath, paths))
            if any(os.path.splitext(path)[1].lower() in self.extensions
                   for path in paths):
                summary.append((event_name, paths))
        if summary:
            self.run(sorted(set(summary)))

    def on_any_event(self, event):
        if isinstance(event, tuple(WATCHED_EVENTS)):
            if self.spooler is not None:
                self.spooler.enqueue(event)
            else:
                self.on_queued_events([event])

    def run(self, summary=None):
        """
        Called when a file is changed to re-run the tests with py.test.
        """
        runner = self.runner or 'py.test'
        argv = runner.split(' ') + self.args

        # Prepare status update
        if self.auto_clear:
            subprocess.call(CLEAR_COMMAND, shell=True)
        elif summary:
            print()

        # Print status
        if not self.quiet:
            command = ' '.join(argv)
            highlight = lambda arg: STYLE_HIGHLIGHT + arg + STYLE_NORMAL
            msg = 'Running: {}'.format(highlight(command))
            if summary:
                if self.verbose:
                    file_lines = ['    {:9s}'.format(event_name + ':') + ' ' +
                                  ' -> '.join(map(highlight, paths))
                                  for event_name, paths in summary]
                    msg = ('Changes detected in files:\n{}\n\nRerunning: {}'
                           .format('\n'.join(file_lines), highlight(command)))
                else:
                    msg = ('Changes detected, rerunning: {}'
                           .format(highlight(command)))
            print(STYLE_NORMAL + msg + Fore.RESET + Style.NORMAL)

        # Run custom command
        if self.beforerun:
            os.system(self.beforerun)

        # Run py.test or py.test runner
        exit_code = subprocess.call(argv, shell=(sys.platform == 'win32'))

        if runner == 'py.test':
            # py.test returns exit code 5 in case no tests are run/collected.
            # This can happen with tools like pytest-testmon.
            passed = exit_code in [0, 5]
        else:
            passed = exit_code == 0
        failed = not passed
        interrupted = exit_code == 2

        # Beep if failed
        if failed and self.beep_on_failure:
            sys.stdout.write(BEEP_CHARACTER)
            sys.stdout.flush()

        # Run custom commands
        if passed and self.onpass:
            os.system(self.onpass)
        elif failed and self.onfail:
            os.system(self.onfail)
        elif interrupted and self.onsigint:
            os.system(self.onsigint)

def watch(directories=[], ignore=[], auto_clear=False, beep_on_failure=True,
          onpass=None, onfail=None, runner=None, beforerun=None, onexit=None,
          onsigint=None, poll=False, extensions=[], args=[], spool=True,
          verbose=False, quiet=False):
    if not directories:
        directories = ['.']
    directories = [os.path.abspath(directory) for directory in directories]
    for directory in directories:
        if not os.path.isdir(directory):
            raise ValueError('Directory not found: ' + directory)

    if ignore:
        recursive_dirs, non_recursive_dirs = split_recursive(
            directories, ignore)
    else:
        recursive_dirs = directories
        non_recursive_dirs = []

    # Initial run
    event_handler = ChangeHandler(
        auto_clear, beep_on_failure, onpass, onfail, onsigint, runner,
        beforerun, extensions, args, spool, verbose, quiet)
    event_handler.run()

    # Setup watchdog
    observer = PollingObserver() if poll else Observer()
    for directory in recursive_dirs:
        observer.schedule(event_handler, path=directory, recursive=True)
    for directory in non_recursive_dirs:
        observer.schedule(event_handler, path=directory, recursive=False)

    # Watch and run tests until interrupted by user
    try:
        observer.start()
        while True:
            time.sleep(1)
        observer.join()
    except KeyboardInterrupt:
        observer.stop()
        if onsigint:
            os.system(onsigint)
    else:
        if onexit:
            os.system(onexit)

def samepath(left, right):
    return (os.path.abspath(os.path.normcase(left)) ==
            os.path.abspath(os.path.normcase(right)))

def split_recursive(directories, ignore):
    non_recursive_dirs = []
    recursive_dirs = []
    for directory in directories:
        subdirs = [os.path.join(directory, d)
                   for d in os.listdir(directory)
                   if os.path.isdir(d)]
        filtered = [subdir for subdir in subdirs
                    if not any(samepath(os.path.join(directory, d), subdir)
                               for d in ignore)]
        if len(subdirs) == len(filtered):
            recursive_dirs.append(directory)
        else:
            non_recursive_dirs.append(directory)
            recursive_dirs.extend(filtered)

    return sorted(set(recursive_dirs)), sorted(set(non_recursive_dirs))
<commit_msg>squash! squash! Add --onsigint handler, run on KeyboardInterrupt<commit_after>from __future__ import print_function

import os
import time
import subprocess
import sys

from colorama import Fore, Style
from watchdog.events import (
    FileSystemEventHandler, FileModifiedEvent, FileCreatedEvent,
    FileMovedEvent, FileDeletedEvent)
from watchdog.observers import Observer
from watchdog.observers.polling import PollingObserver

from .spooler import EventSpooler

EVENT_NAMES = {
    FileModifiedEvent: 'modified',
    FileCreatedEvent: 'created',
    FileMovedEvent: 'moved',
    FileDeletedEvent: 'deleted',
}
WATCHED_EVENTS = list(EVENT_NAMES)
DEFAULT_EXTENSIONS = ['.py']
CLEAR_COMMAND = 'cls' if os.name == 'nt' else 'clear'
BEEP_CHARACTER = '\a'
STYLE_NORMAL = Fore.RESET
STYLE_HIGHLIGHT = Fore.CYAN + Style.NORMAL + Style.BRIGHT

class ChangeHandler(FileSystemEventHandler):
    """
    Listens for changes to files and re-runs tests after each change.
    """
    def __init__(self, auto_clear=False, beep_on_failure=True, onpass=None,
                 onfail=None, onsigint=None, runner=None, beforerun=None,
                 extensions=[], args=None, spool=True, verbose=False,
                 quiet=False):
        super(ChangeHandler, self).__init__()
        self.auto_clear = auto_clear
        self.beep_on_failure = beep_on_failure
        self.onpass = onpass
        self.onfail = onfail
        self.onsigint = onsigint
        self.runner = runner
        self.beforerun = beforerun
        self.extensions = extensions or DEFAULT_EXTENSIONS
        self.args = args or []
        self.spooler = None
        if spool:
            self.spooler = EventSpooler(0.2, self.on_queued_events)
        self.verbose = verbose
        self.quiet = quiet

    def on_queued_events(self, events):
        summary = []
        for event in events:
            paths = [event.src_path]
            if isinstance(event, FileMovedEvent):
                paths.append(event.dest_path)
            event_name = EVENT_NAMES[type(event)]
            paths = tuple(map(os.path.relpath, paths))
            if any(os.path.splitext(path)[1].lower() in self.extensions
                   for path in paths):
                summary.append((event_name, paths))
        if summary:
            self.run(sorted(set(summary)))

    def on_any_event(self, event):
        if isinstance(event, tuple(WATCHED_EVENTS)):
            if self.spooler is not None:
                self.spooler.enqueue(event)
            else:
                self.on_queued_events([event])

    def run(self, summary=None):
        """
        Called when a file is changed to re-run the tests with py.test.
        """
        runner = self.runner or 'py.test'
        argv = runner.split(' ') + self.args

        # Prepare status update
        if self.auto_clear:
            subprocess.call(CLEAR_COMMAND, shell=True)
        elif summary:
            print()

        # Print status
        if not self.quiet:
            command = ' '.join(argv)
            highlight = lambda arg: STYLE_HIGHLIGHT + arg + STYLE_NORMAL
            msg = 'Running: {}'.format(highlight(command))
            if summary:
                if self.verbose:
                    file_lines = ['    {:9s}'.format(event_name + ':') + ' ' +
                                  ' -> '.join(map(highlight, paths))
                                  for event_name, paths in summary]
                    msg = ('Changes detected in files:\n{}\n\nRerunning: {}'
                           .format('\n'.join(file_lines), highlight(command)))
                else:
                    msg = ('Changes detected, rerunning: {}'
                           .format(highlight(command)))
            print(STYLE_NORMAL + msg + Fore.RESET + Style.NORMAL)

        # Run custom command
        if self.beforerun:
            os.system(self.beforerun)

        # Run py.test or py.test runner
        exit_code = subprocess.call(argv, shell=(sys.platform == 'win32'))

        if runner == 'py.test':
            # py.test returns exit code 5 in case no tests are run/collected.
            # This can happen with tools like pytest-testmon.
            passed = exit_code in [0, 5]
        else:
            passed = exit_code == 0
        failed = not passed
        interrupted = exit_code == 2

        # Beep if failed
        if failed and self.beep_on_failure:
            sys.stdout.write(BEEP_CHARACTER)
            sys.stdout.flush()

        # Run custom commands
        if interrupted and self.onsigint:
            os.system(self.onsigint)
        if passed and self.onpass:
            os.system(self.onpass)
        elif failed and self.onfail:
            os.system(self.onfail)

def watch(directories=[], ignore=[], auto_clear=False, beep_on_failure=True,
          onpass=None, onfail=None, runner=None, beforerun=None, onexit=None,
          onsigint=None, poll=False, extensions=[], args=[], spool=True,
          verbose=False, quiet=False):
    if not directories:
        directories = ['.']
    directories = [os.path.abspath(directory) for directory in directories]
    for directory in directories:
        if not os.path.isdir(directory):
            raise ValueError('Directory not found: ' + directory)

    if ignore:
        recursive_dirs, non_recursive_dirs = split_recursive(
            directories, ignore)
    else:
        recursive_dirs = directories
        non_recursive_dirs = []

    # Initial run
    event_handler = ChangeHandler(
        auto_clear, beep_on_failure, onpass, onfail, onsigint, runner,
        beforerun, extensions, args, spool, verbose, quiet)
    event_handler.run()

    # Setup watchdog
    observer = PollingObserver() if poll else Observer()
    for directory in recursive_dirs:
        observer.schedule(event_handler, path=directory, recursive=True)
    for directory in non_recursive_dirs:
        observer.schedule(event_handler, path=directory, recursive=False)

    # Watch and run tests until interrupted by user
    try:
        observer.start()
        while True:
            time.sleep(1)
        observer.join()
    except KeyboardInterrupt:
        observer.stop()
        if onsigint:
            os.system(onsigint)
    else:
        if onexit:
            os.system(onexit)

def samepath(left, right):
    return (os.path.abspath(os.path.normcase(left)) ==
            os.path.abspath(os.path.normcase(right)))

def split_recursive(directories, ignore):
    non_recursive_dirs = []
    recursive_dirs = []
    for directory in directories:
        subdirs = [os.path.join(directory, d)
                   for d in os.listdir(directory)
                   if os.path.isdir(d)]
        filtered = [subdir for subdir in subdirs
                    if not any(samepath(os.path.join(directory, d), subdir)
                               for d in ignore)]
        if len(subdirs) == len(filtered):
            recursive_dirs.append(directory)
        else:
            non_recursive_dirs.append(directory)
            recursive_dirs.extend(filtered)

    return sorted(set(recursive_dirs)), sorted(set(non_recursive_dirs))
<|endoftext|>


<commit_before># Copyright (c) 2015 IBM Corporation and others.
#
# Licensed 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.
from __future__ import absolute_import

import json
import io
import uuid
import os
import inspect
import fnmatch
import jpype
import sys

import brunel.brunelWidgets as brunelWidgets
import brunel.brunel_util as brunel_util
import jinja2 as jin
from IPython.display import Javascript, HTML
from IPython.display import display as ipydisplay

# JS & HTML Files.
D3_TEMPLATE_FILE ="D3_Template.js"
D3_TEMPLATE_HTML_FILE =  "D3_Template.html"

# Jinja templates
templateLoader = jin.PackageLoader("brunel", "")
templateEnv = jin.Environment(loader=templateLoader)
D3_TEMPLATE = templateEnv.get_template(D3_TEMPLATE_FILE)
D3_TEMPLATE_HTML = templateEnv.get_template(D3_TEMPLATE_HTML_FILE)

#Directory containing the Brunel .jar files
lib_dir = os.path.join(os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe()))), "lib")

def display(brunel, data, width=800, height=600, output='d3'):

    csv = None
    if data is not None:
        csv = to_csv(data)

    # unique identifier for HTML tags
    visid = "visid" + str(uuid.uuid1())

    # D3 is currently the only supported renderer
    if output == 'd3':
        result = brunel_jpype_call(csv, brunel, width, height, visid)
        return d3_output(result, visid, width, height)
    else:
        raise ValueError("Valid Output Choices Are:   d3")

def to_csv(df):
        # CSV to pass to service
        # Code is different in python 2 vs. 3
        if sys.version_info < (3,0):
            import StringIO
            csvIO = StringIO.StringIO()
            df.to_csv(csvIO, index=False)
            csv = csvIO.getvalue()
            return csv
        else:
            csvIO = io.StringIO()
            df.to_csv(csvIO, index=False)
            csv = csvIO.getvalue()
            return csv

# Uses jpype to call the main Brunel D3 integration method
def brunel_jpype_call(data, brunel_src, width, height, visid):
    try:
        return brunel_util_java.D3Integration.createBrunelJSON(data, brunel_src, int(width), int(height), visid, None)
    except jpype.JavaException as exception:
        raise ValueError(exception.message())

def get_dataset_names(brunel_src):
    return brunel_util_java.D3Integration.getDatasetNames(brunel_src)

def cacheData(data_key, data):
    brunel_util_java.D3Integration.cacheData(data_key, data)

# D3 response should contain the D3 JS and D3 CSS
def d3_output(response, visid, width, height):
    results = json.loads(response)
    d3js = results["js"]
    d3css = results["css"]
    controls = results["controls"]
    html = D3_TEMPLATE_HTML.render({'jsloc': brunel_util.JS_LOC, 'd3css': d3css, 'visId': visid, 'width': width, 'height': height})
    # side effect pushes required D3 HTML to the client
    ipydisplay(HTML(html))
    widgets = brunelWidgets.build_widgets(controls, visid)
    if (widgets is not None):
        # Push widgets & D3 JS
        js = D3_TEMPLATE.render({'jsloc': brunel_util.JS_LOC, 'd3loc': brunel_util.D3_LOC, 'd3js': d3js, 'controls': widgets['wire_code']})
        return ipydisplay(widgets['widget_box'], Javascript(js))
    else:
        js = D3_TEMPLATE.render({'jsloc': brunel_util.JS_LOC, 'd3loc': brunel_util.D3_LOC, 'd3js': d3js, 'controls': ""})
        return Javascript(js)

#File search given a path.  Used to find the JVM if needed
def find_file(pattern, path):
    result = []
    for root, dirs, files in os.walk(path):
        for name in files:
            if fnmatch.fnmatch(name, pattern):
                result.append(os.path.join(root, name))
    return result

#Start the JVM using jpype
def start_JVM():

    #only start JVM once since it is expensive
    if (not jpype.isJVMStarted()):
        #Use Brunel .jar files
        lib_ext = "-Djava.ext.dirs=" + lib_dir

        # headless execution of java is needed due to
        # calls to AWT.  See git issue #70
        headless = "-Djava.awt.headless=true"

        try:
            #First use explicit path if provided
            if brunel_util.JVM_PATH != "":
                jpype.startJVM(brunel_util.JVM_PATH, headless, lib_ext)
            else:
                # Try jpype's default way
                jpype.startJVM(jpype.getDefaultJVMPath(), headless, lib_ext)
        except:
            #jpype could not find JVM (this happens currently for IBM JDK)
            #Try to find the JVM starting from JAVA_HOME either as a .dll or a .so
            jvms = find_file('jvm.dll', os.environ['JAVA_HOME'])
            if (not jvms):
                jvms = find_file('libjvm.so', os.environ['JAVA_HOME'])
            if (not jvms):
                raise ValueError(
                    "No JVM was found.  First be sure the JAVA_HOME environment variable has been properly "
                    "set before starting IPython.  If it still fails, try to manually set the JVM using:  "
                    "brunel.brunel_util.JVM_PATH=[path]. Where 'path' is the location of the JVM file (not "
                    "directory). Typically this is the full path to 'jvm.dll' on Windows or 'libjvm.so' on Unix ")
            jpype.startJVM(jvms[0], headless, lib_ext)

#Take the JVM startup hit once
start_JVM()
brunel_util_java = jpype.JPackage("org.brunel.util")<commit_msg>Git issue #71.  Exports the index as a field--if it has actually has a name<commit_after># Copyright (c) 2015 IBM Corporation and others.
#
# Licensed 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.
from __future__ import absolute_import

import json
import io
import uuid
import os
import inspect
import fnmatch
import jpype
import sys

import brunel.brunelWidgets as brunelWidgets
import brunel.brunel_util as brunel_util
import jinja2 as jin
from IPython.display import Javascript, HTML
from IPython.display import display as ipydisplay

# JS & HTML Files.
D3_TEMPLATE_FILE ="D3_Template.js"
D3_TEMPLATE_HTML_FILE =  "D3_Template.html"

# Jinja templates
templateLoader = jin.PackageLoader("brunel", "")
templateEnv = jin.Environment(loader=templateLoader)
D3_TEMPLATE = templateEnv.get_template(D3_TEMPLATE_FILE)
D3_TEMPLATE_HTML = templateEnv.get_template(D3_TEMPLATE_HTML_FILE)

#Directory containing the Brunel .jar files
lib_dir = os.path.join(os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe()))), "lib")

def display(brunel, data, width=800, height=600, output='d3'):

    csv = None
    if data is not None:
        csv = to_csv(data)

    # unique identifier for HTML tags
    visid = "visid" + str(uuid.uuid1())

    # D3 is currently the only supported renderer
    if output == 'd3':
        result = brunel_jpype_call(csv, brunel, width, height, visid)
        return d3_output(result, visid, width, height)
    else:
        raise ValueError("Valid Output Choices Are:   d3")

def to_csv(df):

        #If user has done something to cause a named Index, preserve it
        use_index = False
        if  df.index.name is not None:
            use_index=True

        # CSV to pass to service
        # Code is different in python 2 vs. 3
        if sys.version_info < (3,0):
            import StringIO
            csvIO = StringIO.StringIO()
            df.to_csv(csvIO, index=use_index)
            csv = csvIO.getvalue()
            return csv
        else:
            csvIO = io.StringIO()
            df.to_csv(csvIO, index=use_index)
            csv = csvIO.getvalue()
            return csv

# Uses jpype to call the main Brunel D3 integration method
def brunel_jpype_call(data, brunel_src, width, height, visid):
    try:
        return brunel_util_java.D3Integration.createBrunelJSON(data, brunel_src, int(width), int(height), visid, None)
    except jpype.JavaException as exception:
        raise ValueError(exception.message())

def get_dataset_names(brunel_src):
    return brunel_util_java.D3Integration.getDatasetNames(brunel_src)

def cacheData(data_key, data):
    brunel_util_java.D3Integration.cacheData(data_key, data)

# D3 response should contain the D3 JS and D3 CSS
def d3_output(response, visid, width, height):
    results = json.loads(response)
    d3js = results["js"]
    d3css = results["css"]
    controls = results["controls"]
    html = D3_TEMPLATE_HTML.render({'jsloc': brunel_util.JS_LOC, 'd3css': d3css, 'visId': visid, 'width': width, 'height': height})
    # side effect pushes required D3 HTML to the client
    ipydisplay(HTML(html))
    widgets = brunelWidgets.build_widgets(controls, visid)
    if (widgets is not None):
        # Push widgets & D3 JS
        js = D3_TEMPLATE.render({'jsloc': brunel_util.JS_LOC, 'd3loc': brunel_util.D3_LOC, 'd3js': d3js, 'controls': widgets['wire_code']})
        return ipydisplay(widgets['widget_box'], Javascript(js))
    else:
        js = D3_TEMPLATE.render({'jsloc': brunel_util.JS_LOC, 'd3loc': brunel_util.D3_LOC, 'd3js': d3js, 'controls': ""})
        return Javascript(js)

#File search given a path.  Used to find the JVM if needed
def find_file(pattern, path):
    result = []
    for root, dirs, files in os.walk(path):
        for name in files:
            if fnmatch.fnmatch(name, pattern):
                result.append(os.path.join(root, name))
    return result

#Start the JVM using jpype
def start_JVM():

    #only start JVM once since it is expensive
    if (not jpype.isJVMStarted()):
        #Use Brunel .jar files
        lib_ext = "-Djava.ext.dirs=" + lib_dir

        # headless execution of java is needed due to
        # calls to AWT.  See git issue #70
        headless = "-Djava.awt.headless=true"

        try:
            #First use explicit path if provided
            if brunel_util.JVM_PATH != "":
                jpype.startJVM(brunel_util.JVM_PATH, headless, lib_ext)
            else:
                # Try jpype's default way
                jpype.startJVM(jpype.getDefaultJVMPath(), headless, lib_ext)
        except:
            #jpype could not find JVM (this happens currently for IBM JDK)
            #Try to find the JVM starting from JAVA_HOME either as a .dll or a .so
            jvms = find_file('jvm.dll', os.environ['JAVA_HOME'])
            if (not jvms):
                jvms = find_file('libjvm.so', os.environ['JAVA_HOME'])
            if (not jvms):
                raise ValueError(
                    "No JVM was found.  First be sure the JAVA_HOME environment variable has been properly "
                    "set before starting IPython.  If it still fails, try to manually set the JVM using:  "
                    "brunel.brunel_util.JVM_PATH=[path]. Where 'path' is the location of the JVM file (not "
                    "directory). Typically this is the full path to 'jvm.dll' on Windows or 'libjvm.so' on Unix ")
            jpype.startJVM(jvms[0], headless, lib_ext)

#Take the JVM startup hit once
start_JVM()
brunel_util_java = jpype.JPackage("org.brunel.util")<|endoftext|>


<commit_before># Copyright (c) 2015-2017 Vector 35 LLC
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to
# deal in the Software without restriction, including without limitation the
# rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
# sell copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
# IN THE SOFTWARE.

import ctypes

# Binary Ninja components
import _binaryninjacore as core
from .enums import PluginType, PluginUpdateStatus
import startup

class RepoPlugin(object):
	"""
	``RepoPlugin` is mostly read-only, however you can install/uninstall enable/disable plugins. RepoPlugins are
	created by parsing the plugins.json in a plugin repository.
	"""
	def __init__(self, handle):
		self.handle = core.handle_of_type(handle, core.BNRepoPlugin)

	def __del__(self):
		core.BNFreePlugin(self.handle)

	def __repr__(self):
		return "<{} {}/{}>".format(self.path, "installed" if self.installed else "not-installed", "enabled" if self.enabled else "disabled")

	@property
	def path(self):
		"""Relative path from the base of the repository to the actual plugin"""
		return core.BNPluginGetPath(self.handle)

	@property
	def installed(self):
		"""Boolean True if the plugin is installed, False otherwise"""
		return core.BNPluginIsInstalled(self.handle)

	@installed.setter
	def installed(self, state):
		if state:
			return core.BNPluginInstall(self.handle)
		else:
			return core.BNPluginUninstall(self.handle)

	@property
	def enabled(self):
		"""Boolean True if the plugin is currently enabled, False otherwise"""
		return core.BNPluginIsEnabled(self.handle)

	@enabled.setter
	def enabled(self, state):
		if state:
			return core.BNPluginEnable(self.handle)
		else:
			return core.BNPluginDisable(self.handle)

	@property
	def api(self):
		"""string indicating the api used by the plugin"""
		return core.BNPluginGetApi(self.handle)

	@property
	def description(self):
		"""String short description of the plugin"""
		return core.BNPluginGetDescription(self.handle)

	@property
	def license(self):
		"""String short license description (ie MIT, BSD, GPLv2, etc)"""
		return core.BNPluginGetLicense(self.handle)

	@property
	def license_text(self):
		"""String complete license text for the given plugin"""
		return core.BNPluginGetLicenseText(self.handle)

	@property
	def long_description(self):
		"""String long description of the plugin"""
		return core.BNPluginGetLongdescription(self.handle)

	@property
	def minimum_version(self):
		"""String minimum version the plugin was tested on"""
		return core.BNPluginGetMinimimVersions(self.handle)

	@property
	def name(self):
		"""String name of the plugin"""
		return core.BNPluginGetName(self.handle)

	@property
	def plugin_types(self):
		"""List of PluginType enumeration objects indicating the plugin type(s)"""
		result = []
		count = ctypes.c_ulonglong(0)
		plugintypes = core.BNPluginGetPluginTypes(self.handle, count)
		for i in xrange(count.value):
			result.append(PluginType(plugintypes[i]))
		core.BNFreePluginTypes(plugintypes)
		return result

	@property
	def url(self):
		"""String url of the plugin's git repository"""
		return core.BNPluginGetUrl(self.handle)

	@property
	def version(self):
		"""String version of the plugin"""
		return core.BNPluginGetVersion(self.handle)

	@property
	def update_status(self):
		"""PluginUpdateStatus enumeration indicating if the plugin is up to date or not"""
		return PluginUpdateStatus(core.BNPluginGetPluginUpdateStatus(self.handle))

class Repository(object):
	"""
	``Repository`` is a read-only class. Use RepositoryManager to Enable/Disable/Install/Uninstall plugins.
	"""
	def __init__(self, handle):
		self.handle = core.handle_of_type(handle, core.BNRepository)

	def __del__(self):
		core.BNFreeRepository(self.handle)

	def __repr__(self):
		return "<{} - {}/{}>".format(self.path, self.remote_reference, self.local_reference)

	def __getitem__(self, plugin_path):
		for plugin in self.plugins:
			if plugin_path == plugin.path:
				return plugin
		raise KeyError()

	@property
	def url(self):
		"""String url of the git repository where the plugin repository's are stored"""
		return core.BNRepositoryGetUrl(self.handle)

	@property
	def path(self):
		"""String local path to store the given plugin repository"""
		return core.BNRepositoryGetRepoPath(self.handle)

	@property
	def full_path(self):
		"""String full path the repository"""
		return core.BNRepositoryGetPluginsPath(self.handle)

	@property
	def local_reference(self):
		"""String for the local git reference (ie 'master')"""
		return core.BNRepositoryGetLocalReference(self.handle)

	@property
	def remote_reference(self):
		"""String for the remote git reference (ie 'origin')"""
		return core.BNRepositoryGetRemoteReference(self.handle)

	@property
	def plugins(self):
		"""List of RepoPlugin objects contained within this repository"""
		pluginlist = []
		count = ctypes.c_ulonglong(0)
		result = core.BNRepositoryGetPlugins(self.handle, count)
		for i in xrange(count.value):
			pluginlist.append(RepoPlugin(handle=result[i]))
		core.BNFreeRepositoryPluginList(result, count.value)
		del result
		return pluginlist

	@property
	def initialized(self):
		"""Boolean True when the repository has been initialized"""
		return core.BNRepositoryIsInitialized(self.handle)

class RepositoryManager(object):
	"""
	``RepositoryManager`` Keeps track of all the repositories and keeps the enabled_plugins.json file coherent with
	the plugins that are installed/unstalled enabled/disabled
	"""
	def __init__(self, handle=None):
		self.handle = core.BNGetRepositoryManager()

	def __getitem__(self, repo_path):
		for repo in self.repositories:
			if repo_path == repo.path:
				return repo
		raise KeyError()

	def check_for_updates(self):
		"""Check for updates for all managed Repository objects"""
		return core.BNRepositoryManagerCheckForUpdates(self.handle)

	@property
	def repositories(self):
		"""List of Repository objects being managed"""
		result = []
		count = ctypes.c_ulonglong(0)
		repos = core.BNRepositoryManagerGetRepositories(self.handle, count)
		for i in xrange(count.value):
			result.append(Repository(handle=repos[i]))
		core.BNFreeRepositoryManagerRepositoriesList(repos)
		return result

	@property
	def plugins(self):
		"""List of all RepoPlugins in each repository"""
		plgs = {}
		for repo in self.repositories:
			plgs[repo.path] = repo.plugins
		return plgs

	@property
	def default_repository(self):
		"""Gets the default Repository"""
		startup._init_plugins()
		return Repository(handle=core.BNRepositoryManagerGetDefaultRepository(self.handle))

	def enable_plugin(self, plugin, install=True, repo=None):
		"""
		``enable_plugin`` Enables the installed plugin 'plugin', optionally installing the plugin if `install` is set to
		True (default), and optionally using the non-default repository.

		:param str name: Name of the plugin to enable
		:param Boolean install: Optionally install the repo, defaults to True.
		:param str repo: Optional, specify a repository other than the default repository.
		:return: Boolean value True if the plugin was successfully enabled, False otherwise
		:rtype: Boolean
		:Example:

			>>> mgr = RepositoryManager()
			>>> mgr.enable_plugin('binaryninja-bookmarks')
			True
			>>>
		"""
		if install:
			if not self.install_plugin(plugin, repo):
				return False

		if repo is None:
			repo = self.default_repository
		repopath = repo
		pluginpath = plugin
		if not isinstance(repo, str):
			repopath = repo.path
		if not isinstance(plugin, str):
			pluginpath = plugin.path
		return core.BNRepositoryManagerEnablePlugin(self.handle, repopath, pluginpath)

	def disable_plugin(self, plugin, repo=None):
		"""
		``disable_plugin`` Disable the specified plugin, pluginpath

		:param Repository or str repo: Repository containing the plugin to disable
		:param RepoPlugin or str plugin: RepoPlugin to disable
		:return: Boolean value True if the plugin was successfully disabled, False otherwise
		:rtype: Boolean
		:Example:

			>>> mgr = RepositoryManager()
			>>> mgr.disable_plugin('binaryninja-bookmarks')
			True
			>>>
		"""
		if repo is None:
			repo = self.default_repository
		repopath = repo
		pluginpath = plugin
		if not isinstance(repo, str):
			repopath = repo.path
		if not isinstance(plugin, str):
			pluginpath = plugin.path
		return core.BNRepositoryManagerDisablePlugin(self.handle, repopath, pluginpath)

	def install_plugin(self, plugin, repo=None):
		"""
		``install_plugin`` install the specified plugin, pluginpath

		:param Repository or str repo: Repository containing the plugin to install
		:param RepoPlugin or str plugin: RepoPlugin to install
		:return: Boolean value True if the plugin was successfully installed, False otherwise
		:rtype: Boolean
		:Example:

			>>> mgr = RepositoryManager()
			>>> mgr.install_plugin('binaryninja-bookmarks')
			True
			>>>
		"""
		if repo is None:
			repo = self.default_repository
		repopath = repo
		pluginpath = plugin
		if not isinstance(repo, str):
			repopath = repo.path
		if not isinstance(plugin, str):
			pluginpath = plugin.path
		return core.BNRepositoryManagerInstallPlugin(self.handle, repopath, pluginpath)

	def uninstall_plugin(self, plugin, repo=None):
		"""
		``uninstall_plugin`` uninstall the specified plugin, pluginpath

		:param Repository or str repo: Repository containing the plugin to uninstall
		:param RepoPlugin or str plugin: RepoPlugin to uninstall
		:return: Boolean value True if the plugin was successfully uninstalled, False otherwise
		:rtype: Boolean
		:Example:

			>>> mgr = RepositoryManager()
			>>> mgr.uninstall_plugin('binaryninja-bookmarks')
			True
			>>>
		"""
		if repo is None:
			repo = self.default_repository
		repopath = repo
		pluginpath = plugin
		if not isinstance(repo, str):
			repopath = repo.path
		if not isinstance(plugin, str):
			pluginpath = plugin.path
		return core.BNRepositoryManagerUninstallPlugin(self.handle, repopath, pluginpath)

	def update_plugin(self, plugin, repo=None):
		"""
		``update_plugin`` update the specified plugin, pluginpath

		:param Repository or str repo: Repository containing the plugin to update
		:param RepoPlugin or str plugin: RepoPlugin to update
		:return: Boolean value True if the plugin was successfully updated, False otherwise
		:rtype: Boolean
		:Example:

			>>> mgr = RepositoryManager()
			>>> mgr.update_plugin('binaryninja-bookmarks')
			True
			>>>
		"""
		if repo is None:
			repo = self.default_repository
		repopath = repo
		pluginpath = plugin
		if not isinstance(repo, str):
			repopath = repo.path
		if not isinstance(plugin, str):
			pluginpath = plugin.path
		return core.BNRepositoryManagerUpdatePlugin(self.handle, repopath, pluginpath)

	def add_repository(self, url=None, repopath=None, localreference="master", remotereference="origin"):
		"""
		``add_repository`` adds a new plugin repository for the manager to track.

		:param str url: Url to the git repository where the plugins are stored.
		:param str repopath: path to where the repository will be stored on disk locally
		:param str localreference: Optional reference to the local tracking branch typically "master"
		:param str remotereference: Optional reference to the remote tracking branch typically "origin"
		:return: Boolean value True if the repository was successfully added, False otherwise.
		:rtype: Boolean
		:Example:

			>>> mgr = RepositoryManager()
			>>> mgr.add_repository(url="https://github.com/vector35/community-plugins.git",
			                       repopath="myrepo",
			                       repomanifest="plugins",
			                       localreference="master", remotereference="origin")
			True
			>>>
		"""
		if not (isinstance(url, str) and isinstance(repopath, str) and
			isinstance(localreference, str) and isinstance(remotereference, str)):
			raise ValueError("Parameter is incorrect type")

		return core.BNRepositoryManagerAddRepository(self.handle, url, repopath, localreference, remotereference)
<commit_msg>Fix PluginManager Doc Example (#885)<commit_after># Copyright (c) 2015-2017 Vector 35 LLC
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to
# deal in the Software without restriction, including without limitation the
# rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
# sell copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
# IN THE SOFTWARE.

import ctypes

# Binary Ninja components
import _binaryninjacore as core
from .enums import PluginType, PluginUpdateStatus
import startup

class RepoPlugin(object):
	"""
	``RepoPlugin` is mostly read-only, however you can install/uninstall enable/disable plugins. RepoPlugins are
	created by parsing the plugins.json in a plugin repository.
	"""
	def __init__(self, handle):
		self.handle = core.handle_of_type(handle, core.BNRepoPlugin)

	def __del__(self):
		core.BNFreePlugin(self.handle)

	def __repr__(self):
		return "<{} {}/{}>".format(self.path, "installed" if self.installed else "not-installed", "enabled" if self.enabled else "disabled")

	@property
	def path(self):
		"""Relative path from the base of the repository to the actual plugin"""
		return core.BNPluginGetPath(self.handle)

	@property
	def installed(self):
		"""Boolean True if the plugin is installed, False otherwise"""
		return core.BNPluginIsInstalled(self.handle)

	@installed.setter
	def installed(self, state):
		if state:
			return core.BNPluginInstall(self.handle)
		else:
			return core.BNPluginUninstall(self.handle)

	@property
	def enabled(self):
		"""Boolean True if the plugin is currently enabled, False otherwise"""
		return core.BNPluginIsEnabled(self.handle)

	@enabled.setter
	def enabled(self, state):
		if state:
			return core.BNPluginEnable(self.handle)
		else:
			return core.BNPluginDisable(self.handle)

	@property
	def api(self):
		"""string indicating the api used by the plugin"""
		return core.BNPluginGetApi(self.handle)

	@property
	def description(self):
		"""String short description of the plugin"""
		return core.BNPluginGetDescription(self.handle)

	@property
	def license(self):
		"""String short license description (ie MIT, BSD, GPLv2, etc)"""
		return core.BNPluginGetLicense(self.handle)

	@property
	def license_text(self):
		"""String complete license text for the given plugin"""
		return core.BNPluginGetLicenseText(self.handle)

	@property
	def long_description(self):
		"""String long description of the plugin"""
		return core.BNPluginGetLongdescription(self.handle)

	@property
	def minimum_version(self):
		"""String minimum version the plugin was tested on"""
		return core.BNPluginGetMinimimVersions(self.handle)

	@property
	def name(self):
		"""String name of the plugin"""
		return core.BNPluginGetName(self.handle)

	@property
	def plugin_types(self):
		"""List of PluginType enumeration objects indicating the plugin type(s)"""
		result = []
		count = ctypes.c_ulonglong(0)
		plugintypes = core.BNPluginGetPluginTypes(self.handle, count)
		for i in xrange(count.value):
			result.append(PluginType(plugintypes[i]))
		core.BNFreePluginTypes(plugintypes)
		return result

	@property
	def url(self):
		"""String url of the plugin's git repository"""
		return core.BNPluginGetUrl(self.handle)

	@property
	def version(self):
		"""String version of the plugin"""
		return core.BNPluginGetVersion(self.handle)

	@property
	def update_status(self):
		"""PluginUpdateStatus enumeration indicating if the plugin is up to date or not"""
		return PluginUpdateStatus(core.BNPluginGetPluginUpdateStatus(self.handle))

class Repository(object):
	"""
	``Repository`` is a read-only class. Use RepositoryManager to Enable/Disable/Install/Uninstall plugins.
	"""
	def __init__(self, handle):
		self.handle = core.handle_of_type(handle, core.BNRepository)

	def __del__(self):
		core.BNFreeRepository(self.handle)

	def __repr__(self):
		return "<{} - {}/{}>".format(self.path, self.remote_reference, self.local_reference)

	def __getitem__(self, plugin_path):
		for plugin in self.plugins:
			if plugin_path == plugin.path:
				return plugin
		raise KeyError()

	@property
	def url(self):
		"""String url of the git repository where the plugin repository's are stored"""
		return core.BNRepositoryGetUrl(self.handle)

	@property
	def path(self):
		"""String local path to store the given plugin repository"""
		return core.BNRepositoryGetRepoPath(self.handle)

	@property
	def full_path(self):
		"""String full path the repository"""
		return core.BNRepositoryGetPluginsPath(self.handle)

	@property
	def local_reference(self):
		"""String for the local git reference (ie 'master')"""
		return core.BNRepositoryGetLocalReference(self.handle)

	@property
	def remote_reference(self):
		"""String for the remote git reference (ie 'origin')"""
		return core.BNRepositoryGetRemoteReference(self.handle)

	@property
	def plugins(self):
		"""List of RepoPlugin objects contained within this repository"""
		pluginlist = []
		count = ctypes.c_ulonglong(0)
		result = core.BNRepositoryGetPlugins(self.handle, count)
		for i in xrange(count.value):
			pluginlist.append(RepoPlugin(handle=result[i]))
		core.BNFreeRepositoryPluginList(result, count.value)
		del result
		return pluginlist

	@property
	def initialized(self):
		"""Boolean True when the repository has been initialized"""
		return core.BNRepositoryIsInitialized(self.handle)

class RepositoryManager(object):
	"""
	``RepositoryManager`` Keeps track of all the repositories and keeps the enabled_plugins.json file coherent with
	the plugins that are installed/unstalled enabled/disabled
	"""
	def __init__(self, handle=None):
		self.handle = core.BNGetRepositoryManager()

	def __getitem__(self, repo_path):
		for repo in self.repositories:
			if repo_path == repo.path:
				return repo
		raise KeyError()

	def check_for_updates(self):
		"""Check for updates for all managed Repository objects"""
		return core.BNRepositoryManagerCheckForUpdates(self.handle)

	@property
	def repositories(self):
		"""List of Repository objects being managed"""
		result = []
		count = ctypes.c_ulonglong(0)
		repos = core.BNRepositoryManagerGetRepositories(self.handle, count)
		for i in xrange(count.value):
			result.append(Repository(handle=repos[i]))
		core.BNFreeRepositoryManagerRepositoriesList(repos)
		return result

	@property
	def plugins(self):
		"""List of all RepoPlugins in each repository"""
		plgs = {}
		for repo in self.repositories:
			plgs[repo.path] = repo.plugins
		return plgs

	@property
	def default_repository(self):
		"""Gets the default Repository"""
		startup._init_plugins()
		return Repository(handle=core.BNRepositoryManagerGetDefaultRepository(self.handle))

	def enable_plugin(self, plugin, install=True, repo=None):
		"""
		``enable_plugin`` Enables the installed plugin 'plugin', optionally installing the plugin if `install` is set to
		True (default), and optionally using the non-default repository.

		:param str name: Name of the plugin to enable
		:param Boolean install: Optionally install the repo, defaults to True.
		:param str repo: Optional, specify a repository other than the default repository.
		:return: Boolean value True if the plugin was successfully enabled, False otherwise
		:rtype: Boolean
		:Example:

			>>> mgr = RepositoryManager()
			>>> mgr.enable_plugin('binaryninja-bookmarks')
			True
			>>>
		"""
		if install:
			if not self.install_plugin(plugin, repo):
				return False

		if repo is None:
			repo = self.default_repository
		repopath = repo
		pluginpath = plugin
		if not isinstance(repo, str):
			repopath = repo.path
		if not isinstance(plugin, str):
			pluginpath = plugin.path
		return core.BNRepositoryManagerEnablePlugin(self.handle, repopath, pluginpath)

	def disable_plugin(self, plugin, repo=None):
		"""
		``disable_plugin`` Disable the specified plugin, pluginpath

		:param Repository or str repo: Repository containing the plugin to disable
		:param RepoPlugin or str plugin: RepoPlugin to disable
		:return: Boolean value True if the plugin was successfully disabled, False otherwise
		:rtype: Boolean
		:Example:

			>>> mgr = RepositoryManager()
			>>> mgr.disable_plugin('binaryninja-bookmarks')
			True
			>>>
		"""
		if repo is None:
			repo = self.default_repository
		repopath = repo
		pluginpath = plugin
		if not isinstance(repo, str):
			repopath = repo.path
		if not isinstance(plugin, str):
			pluginpath = plugin.path
		return core.BNRepositoryManagerDisablePlugin(self.handle, repopath, pluginpath)

	def install_plugin(self, plugin, repo=None):
		"""
		``install_plugin`` install the specified plugin, pluginpath

		:param Repository or str repo: Repository containing the plugin to install
		:param RepoPlugin or str plugin: RepoPlugin to install
		:return: Boolean value True if the plugin was successfully installed, False otherwise
		:rtype: Boolean
		:Example:

			>>> mgr = RepositoryManager()
			>>> mgr.install_plugin('binaryninja-bookmarks')
			True
			>>>
		"""
		if repo is None:
			repo = self.default_repository
		repopath = repo
		pluginpath = plugin
		if not isinstance(repo, str):
			repopath = repo.path
		if not isinstance(plugin, str):
			pluginpath = plugin.path
		return core.BNRepositoryManagerInstallPlugin(self.handle, repopath, pluginpath)

	def uninstall_plugin(self, plugin, repo=None):
		"""
		``uninstall_plugin`` uninstall the specified plugin, pluginpath

		:param Repository or str repo: Repository containing the plugin to uninstall
		:param RepoPlugin or str plugin: RepoPlugin to uninstall
		:return: Boolean value True if the plugin was successfully uninstalled, False otherwise
		:rtype: Boolean
		:Example:

			>>> mgr = RepositoryManager()
			>>> mgr.uninstall_plugin('binaryninja-bookmarks')
			True
			>>>
		"""
		if repo is None:
			repo = self.default_repository
		repopath = repo
		pluginpath = plugin
		if not isinstance(repo, str):
			repopath = repo.path
		if not isinstance(plugin, str):
			pluginpath = plugin.path
		return core.BNRepositoryManagerUninstallPlugin(self.handle, repopath, pluginpath)

	def update_plugin(self, plugin, repo=None):
		"""
		``update_plugin`` update the specified plugin, pluginpath

		:param Repository or str repo: Repository containing the plugin to update
		:param RepoPlugin or str plugin: RepoPlugin to update
		:return: Boolean value True if the plugin was successfully updated, False otherwise
		:rtype: Boolean
		:Example:

			>>> mgr = RepositoryManager()
			>>> mgr.update_plugin('binaryninja-bookmarks')
			True
			>>>
		"""
		if repo is None:
			repo = self.default_repository
		repopath = repo
		pluginpath = plugin
		if not isinstance(repo, str):
			repopath = repo.path
		if not isinstance(plugin, str):
			pluginpath = plugin.path
		return core.BNRepositoryManagerUpdatePlugin(self.handle, repopath, pluginpath)

	def add_repository(self, url=None, repopath=None, localreference="master", remotereference="origin"):
		"""
		``add_repository`` adds a new plugin repository for the manager to track.

		:param str url: Url to the git repository where the plugins are stored.
		:param str repopath: path to where the repository will be stored on disk locally
		:param str localreference: Optional reference to the local tracking branch typically "master"
		:param str remotereference: Optional reference to the remote tracking branch typically "origin"
		:return: Boolean value True if the repository was successfully added, False otherwise.
		:rtype: Boolean
		:Example:

			>>> mgr = RepositoryManager()
			>>> mgr.add_repository(url="https://github.com/vector35/community-plugins.git",
			                       repopath="myrepo",
			                       localreference="master", remotereference="origin")
			True
			>>>
		"""
		if not (isinstance(url, str) and isinstance(repopath, str) and
			isinstance(localreference, str) and isinstance(remotereference, str)):
			raise ValueError("Parameter is incorrect type")

		return core.BNRepositoryManagerAddRepository(self.handle, url, repopath, localreference, remotereference)
<|endoftext|>


<commit_before>import json
import pymongo
import pandas as pd
import os
__all__ = ["log2json", "read_config"]

def log2json(log_file):
    s = []
    with open(log_file) as f:
        for line in f:
            x = line.replace("'", '"')
            s.append(json.loads(x))
    return s
<commit_msg>Change export<commit_after>import json
import pymongo
import pandas as pd
import os
__all__ = ["log2json"]

def log2json(log_file):
    s = []
    with open(log_file) as f:
        for line in f:
            x = line.replace("'", '"')
            s.append(json.loads(x))
    return s
<|endoftext|>


<commit_before>"""
Code to handle EC2 stand-alone instances
"""
import boto3
import os
import pywren
import base64
import logging

logger = logging.getLogger(__name__)

def b64s(string):
    """
    Base-64 encode a string and return a string
    """
    return base64.b64encode(string.encode('utf-8')).decode('ascii')

def sd(filename):
    """
    get the file in the standalone dir
    """
    return os.path.join(pywren.SOURCE_DIR,
                        'ec2_standalone_files', filename)

def create_instance_profile(instance_profile_name):
    iam = boto3.resource('iam')
    #iam.create_instance_profile(InstanceProfileName=INSTANCE_PROFILE_NAME)

    instance_profile = iam.InstanceProfile(instance_profile_name)
    #instance_profile.add_role(RoleName='pywren_exec_role_refactor8')

def launch_instances(number, tgt_ami, aws_region, my_aws_key, instance_type,
                     instance_name, instance_profile_name, sqs_queue_name,
                     default_volume_size=100,
                     max_idle_time=60, idle_terminate_granularity=600,
                     pywren_git_branch='master',
                     pywren_git_commit=None):

    logger.info("launching {} {} instances in {}".format(number, instance_type,
                                                         aws_region))
    # INSTANCE_TYPE = 'm3.xlarge'
    # instance_name = AWS_INSTANCE_NAME

    ec2 = boto3.resource('ec2', region_name=aws_region)

    BlockDeviceMappings=[
        {
            'DeviceName': '/dev/xvda',
            'Ebs': {
                'VolumeSize': default_volume_size,
                'DeleteOnTermination': True,
                'VolumeType': 'standard',
                'SnapshotId' : 'snap-c87f35ec'
            },
        },
    ]
    template_file = sd('ec2standalone.cloudinit.template')

    user_data = open(template_file, 'r').read()

    supervisord_init_script = open(sd('supervisord.init'), 'r').read()
    supervisord_init_script_64 = b64s(supervisord_init_script)

    supervisord_conf = open(sd('supervisord.conf'), 'r').read()
    logger.info("Running with idle_terminate_granularity={}".format(idle_terminate_granularity))
    supervisord_conf = supervisord_conf.format(run_dir = "/tmp/pywren.runner",
                                               sqs_queue_name=sqs_queue_name,
                                               aws_region=aws_region,
                                               max_idle_time=max_idle_time,
                                               idle_terminate_granularity=idle_terminate_granularity)
    supervisord_conf_64 = b64s(supervisord_conf)

    cloud_agent_conf = open(sd("cloudwatch-agent.config"),
                            'r').read()
    cloud_agent_conf_64 = b64s(cloud_agent_conf)

    if pywren_git_commit is not None:
        # use a git commit
        git_checkout_string = str(pywren_git_commit)
    else:
        git_checkout_string = "-b {}".format(pywren_git_branch)

    user_data = user_data.format(supervisord_init_script = supervisord_init_script_64,
                                 supervisord_conf = supervisord_conf_64,
                                 git_checkout_string = git_checkout_string,
                                 aws_region = aws_region,
                                 cloud_agent_conf = cloud_agent_conf_64)

    open("/tmp/user_data", 'w').write(user_data)

    iam = boto3.resource('iam')
    instance_profile = iam.InstanceProfile(instance_profile_name)
    instance_profile_dict =  {
                              'Name' : instance_profile.name}
    instances = ec2.create_instances(ImageId=tgt_ami, MinCount=number,
                                     MaxCount=number,
                                     KeyName=my_aws_key,
                                     InstanceType=instance_type,
                                     BlockDeviceMappings = BlockDeviceMappings,
                                     InstanceInitiatedShutdownBehavior='terminate',
                                     EbsOptimized=True,
                                     IamInstanceProfile = instance_profile_dict,
                                     UserData=user_data)

    # FIXME there's a race condition where we could end up with two
    # instances with the same name but that's ok
    existing_instance_names = [a[0] for a in list_instances(aws_region,
                                                            instance_name)]

    new_instances_with_names = []

    def generate_unique_instance_name():
        inst_pos = 0
        while True:
            name_string = "{}-{}".format(instance_name, inst_pos)
            if (name_string not in [a[0] for a in new_instances_with_names]) and \
               (name_string not in existing_instance_names):
                return name_string
            inst_pos += 1

    for inst in instances:

        unique_instance_name = generate_unique_instance_name()
        logger.info("setting instance name to {}".format(unique_instance_name))

        inst.reload()
        inst.create_tags(
            Resources=[
                inst.instance_id
            ],
            Tags=[
                {
                    'Key': 'Name',
                    'Value': unique_instance_name
                },
            ]
        )
        new_instances_with_names.append((unique_instance_name, inst))
    for inst in instances:
        inst.wait_until_running()

    return new_instances_with_names

def tags_to_dict(d):
    if d is None:
        return {}
    return {a['Key'] : a['Value'] for a in d}

def list_instances(aws_region, instance_name):
    """
    List all instances whose names match the main

    Returns [(name, instance_object)]
    """

    ec2 = boto3.resource('ec2', region_name=aws_region)
    insts = []

    for i in ec2.instances.all():
        if i.state['Name'] == 'running':
            d = tags_to_dict(i.tags)
            if 'Name' in d and instance_name in d['Name']:
                insts.append((d['Name'], i))
    return insts

def terminate_instances(instance_list):
    """
    # FIXME delete individuals
    """
    for instance_name, instance_obj in instance_list:
        instance_obj.terminate()

def prettyprint_instances(inst_list):
    for instance_name, instance_obj in inst_list:
        print(instance_name, instance_obj.public_dns_name)

<commit_msg>remove default stanpshot ID<commit_after>"""
Code to handle EC2 stand-alone instances
"""
import boto3
import os
import pywren
import base64
import logging

logger = logging.getLogger(__name__)

def b64s(string):
    """
    Base-64 encode a string and return a string
    """
    return base64.b64encode(string.encode('utf-8')).decode('ascii')

def sd(filename):
    """
    get the file in the standalone dir
    """
    return os.path.join(pywren.SOURCE_DIR,
                        'ec2_standalone_files', filename)

def create_instance_profile(instance_profile_name):
    iam = boto3.resource('iam')
    #iam.create_instance_profile(InstanceProfileName=INSTANCE_PROFILE_NAME)

    instance_profile = iam.InstanceProfile(instance_profile_name)
    #instance_profile.add_role(RoleName='pywren_exec_role_refactor8')

def launch_instances(number, tgt_ami, aws_region, my_aws_key, instance_type,
                     instance_name, instance_profile_name, sqs_queue_name,
                     default_volume_size=100,
                     max_idle_time=60, idle_terminate_granularity=600,
                     pywren_git_branch='master',
                     pywren_git_commit=None):

    logger.info("launching {} {} instances in {}".format(number, instance_type,
                                                         aws_region))
    # INSTANCE_TYPE = 'm3.xlarge'
    # instance_name = AWS_INSTANCE_NAME

    ec2 = boto3.resource('ec2', region_name=aws_region)

    BlockDeviceMappings=[
        {
            'DeviceName': '/dev/xvda',
            'Ebs': {
                'VolumeSize': default_volume_size,
                'DeleteOnTermination': True,
                'VolumeType': 'standard',
            },
        },
    ]
    template_file = sd('ec2standalone.cloudinit.template')

    user_data = open(template_file, 'r').read()

    supervisord_init_script = open(sd('supervisord.init'), 'r').read()
    supervisord_init_script_64 = b64s(supervisord_init_script)

    supervisord_conf = open(sd('supervisord.conf'), 'r').read()
    logger.info("Running with idle_terminate_granularity={}".format(idle_terminate_granularity))
    supervisord_conf = supervisord_conf.format(run_dir = "/tmp/pywren.runner",
                                               sqs_queue_name=sqs_queue_name,
                                               aws_region=aws_region,
                                               max_idle_time=max_idle_time,
                                               idle_terminate_granularity=idle_terminate_granularity)
    supervisord_conf_64 = b64s(supervisord_conf)

    cloud_agent_conf = open(sd("cloudwatch-agent.config"),
                            'r').read()
    cloud_agent_conf_64 = b64s(cloud_agent_conf)

    if pywren_git_commit is not None:
        # use a git commit
        git_checkout_string = str(pywren_git_commit)
    else:
        git_checkout_string = "-b {}".format(pywren_git_branch)

    user_data = user_data.format(supervisord_init_script = supervisord_init_script_64,
                                 supervisord_conf = supervisord_conf_64,
                                 git_checkout_string = git_checkout_string,
                                 aws_region = aws_region,
                                 cloud_agent_conf = cloud_agent_conf_64)

    open("/tmp/user_data", 'w').write(user_data)

    iam = boto3.resource('iam')
    instance_profile = iam.InstanceProfile(instance_profile_name)
    instance_profile_dict =  {
                              'Name' : instance_profile.name}
    instances = ec2.create_instances(ImageId=tgt_ami, MinCount=number,
                                     MaxCount=number,
                                     KeyName=my_aws_key,
                                     InstanceType=instance_type,
                                     BlockDeviceMappings = BlockDeviceMappings,
                                     InstanceInitiatedShutdownBehavior='terminate',
                                     EbsOptimized=True,
                                     IamInstanceProfile = instance_profile_dict,
                                     UserData=user_data)

    # FIXME there's a race condition where we could end up with two
    # instances with the same name but that's ok
    existing_instance_names = [a[0] for a in list_instances(aws_region,
                                                            instance_name)]

    new_instances_with_names = []

    def generate_unique_instance_name():
        inst_pos = 0
        while True:
            name_string = "{}-{}".format(instance_name, inst_pos)
            if (name_string not in [a[0] for a in new_instances_with_names]) and \
               (name_string not in existing_instance_names):
                return name_string
            inst_pos += 1

    for inst in instances:

        unique_instance_name = generate_unique_instance_name()
        logger.info("setting instance name to {}".format(unique_instance_name))

        inst.reload()
        inst.create_tags(
            Resources=[
                inst.instance_id
            ],
            Tags=[
                {
                    'Key': 'Name',
                    'Value': unique_instance_name
                },
            ]
        )
        new_instances_with_names.append((unique_instance_name, inst))
    for inst in instances:
        inst.wait_until_running()

    return new_instances_with_names

def tags_to_dict(d):
    if d is None:
        return {}
    return {a['Key'] : a['Value'] for a in d}

def list_instances(aws_region, instance_name):
    """
    List all instances whose names match the main

    Returns [(name, instance_object)]
    """

    ec2 = boto3.resource('ec2', region_name=aws_region)
    insts = []

    for i in ec2.instances.all():
        if i.state['Name'] == 'running':
            d = tags_to_dict(i.tags)
            if 'Name' in d and instance_name in d['Name']:
                insts.append((d['Name'], i))
    return insts

def terminate_instances(instance_list):
    """
    # FIXME delete individuals
    """
    for instance_name, instance_obj in instance_list:
        instance_obj.terminate()

def prettyprint_instances(inst_list):
    for instance_name, instance_obj in inst_list:
        print(instance_name, instance_obj.public_dns_name)

<|endoftext|>


<commit_before>import os
import re

import numpy as np
import pandas as pd
from IPython import embed
import matplotlib as mpl
#mpl.use('pdf')
from matplotlib.backends.backend_pdf import PdfPages
from mpl_toolkits.axes_grid1.inset_locator import inset_axes
import matplotlib.pyplot as plt
import seaborn as sns

from qlknn.NNDB.model import Network, NetworkJSON, PostprocessSlice
from qlknn.models.ffnn import QuaLiKizNDNN
from qlknn.dataset.data_io import sep_prefix
from qlknn.misc.analyse_names import is_transport, is_full_transport, is_pure, is_partial_heat, is_partial_particle, is_partial_rot, is_partial_diffusion, is_leading
qlknn_root = os.path.abspath('../..')

def determine_subax_loc(ax, height_perc=.35, width_perc=.35):
    cover_left = False
    cover_right = False
    xlim = ax.get_xlim()
    full = np.sum(np.abs(xlim))
    left_bound = xlim[0] + width_perc * full
    right_bound = xlim[1] - width_perc * full
    top_bound = (1 - height_perc) * ax.get_ylim()[1]
    for child in ax.get_children():
        if isinstance(child, mpl.patches.Rectangle):
            xx = child.get_x()
            too_high = child.get_height() > 0.5 * ax.get_ylim()[1]
            if child.get_height() > top_bound:
                if xx < left_bound:
                    cover_left = True
                elif xx > right_bound:
                    cover_right = True

    if not cover_right:
        loc = 1
    elif not cover_left:
        loc = 2
    else:
        loc = 9
    return loc

def plot_dataset_zoomin(store, varname, bound=0.1):
    with sns.axes_style("white"):
        df = store[sep_prefix + varname]
        df.name = varname
        df.dropna(inplace=True)
        fig = plt.figure()
        ax = sns.distplot(df.loc[df.abs() < bound],
                          hist_kws={'log': False}, kde=True)

        sns.despine(ax=ax)
    return fig

def plot_dataset_dist(store, varname, cutoff=0.01):
    with sns.axes_style("white"):
        df = store[sep_prefix + varname]
        df.name = varname
        df.dropna(inplace=True)
        fig = plt.figure()
        ax = sns.distplot(df.loc[(df.quantile(cutoff) < df) &
                                 (df < df.quantile(1 - cutoff))],
                          hist_kws={'log': False}, kde=True)
        ax.set_ylabel('Counts [#]')

        sns.despine(ax=ax)
        loc = determine_subax_loc(ax)
        subax = inset_axes(ax,
                           width="30%",
                           height="30%",
                           loc=loc)
        if is_partial_particle(df.name):
            quant_bound = .15
            low_bound = max(-1, df.quantile(quant_bound))
            high_bound = min(1, df.quantile(1 - quant_bound))
        else:
            low_bound = -1
            high_bound = 1
        sns.distplot(df.loc[(low_bound < df) & (df < high_bound)],
                     kde=True, kde_kws={'gridsize': 200},
                     ax=subax)
        if loc == 2:
            subax.yaxis.set_label_position("right")
            subax.yaxis.tick_right()
            sns.despine(ax=subax, left=True, right=False)
        else:
            sns.despine(ax=subax)
        subax.set_xlabel('')
    return fig

def generate_store_name(unstable=True, gen=3, filter_id=8, dim=7):
    store_name = 'training_gen{!s}_{!s}D_nions0_flat_filter{!s}.h5.1'.format(gen, dim, filter_id)
    if unstable:
        store_name = '_'.join(['unstable', store_name])
    return store_name

def plot_pure_network_dataset_dist(self):
    filter_id = self.filter_id
    dim = len(net.feature_names)
    #store_name = 'unstable_training_gen{!s}_{!s}D_nions0_flat_filter{!s}.h5'.format(2, filter_id, dim)
    store_name = generate_store_name(True, 3, filter_id, dim)
    store = pd.HDFStore(os.path.join(qlknn_root, store_name))
    embed()
    for train_dim in net.target_names:
        plot_dataset_dist(store, train_dim)
        deconstruct = re.split('_div_|_plus_', train_dim)
        if len(deconstruct) > 1:
            for sub_dim in deconstruct:
                plot_dataset_dist(store, sub_dim)
Network.plot_dataset_dist = plot_pure_network_dataset_dist

def generate_dataset_report(store, plot_pure=True, plot_heat=True, plot_particle=True, plot_rot=False, plot_full=False, plot_diffusion=False, plot_nonleading=False, verbose_debug=False):
    with PdfPages('multipage_pdf.pdf') as pdf:
        for varname in store:
            varname = varname.lstrip(sep_prefix)
            if verbose_debug:
                print(varname)
                print('is_full_transport', is_full_transport(varname))
                print('is_pure', is_pure(varname))
                print('is_partial_rot', is_partial_rot(varname))
                print('is_partial_heat', is_partial_heat(varname))
                print('is_partial_particle', is_partial_particle(varname))
                print('is_partial_diffusion', is_partial_diffusion(varname))
                print('is_leading', is_leading(varname))

            if (is_transport(varname) and
                ((plot_full and is_full_transport(varname)) or not is_full_transport(varname)) and
                ((plot_pure and is_pure(varname)) or not is_pure(varname)) and
                ((plot_rot and is_partial_rot(varname)) or
                 (plot_heat and is_partial_heat(varname)) or
                 (plot_particle and is_partial_particle(varname)) or
                 (plot_diffusion and is_partial_diffusion(varname)) or
                 (plot_nonleading and not is_leading(varname)) or
                  (not is_full_transport(varname) and
                   not is_partial_rot(varname) and
                   not is_partial_heat(varname) and
                   not is_partial_particle(varname) and
                   not is_partial_diffusion(varname) and
                   is_leading(varname)))):
                print(varname)
                fig = plot_dataset_dist(store, varname)
                pdf.savefig(fig)
                plt.close(fig)
                try:
                    fig = plot_dataset_zoomin(store, varname)
                except ZeroDivisionError:
                    fig = plt.figure()
                pdf.savefig(fig)
                plt.close(fig)

#net = Network.get_by_id(1409)

if __name__ == '__main__':
    store = pd.HDFStore(os.path.join(qlknn_root, generate_store_name(dim=7)), 'r')
    generate_dataset_report(store)
    embed()
<commit_msg>Popplots generalization and bugfixes..<commit_after>import os
import re

import numpy as np
import pandas as pd
from IPython import embed
import matplotlib as mpl
#mpl.use('pdf')
from matplotlib.backends.backend_pdf import PdfPages
from mpl_toolkits.axes_grid1.inset_locator import inset_axes
import matplotlib.pyplot as plt
import seaborn as sns

from qlknn.NNDB.model import Network, NetworkJSON, PostprocessSlice
from qlknn.models.ffnn import QuaLiKizNDNN
from qlknn.dataset.data_io import sep_prefix
from qlknn.misc.analyse_names import is_transport, is_full_transport, is_pure, is_partial_heat, is_partial_particle, is_partial_rot, is_partial_diffusion, is_partial_momentum, is_leading
qlknn_root = os.path.abspath('../..')

def determine_subax_loc(ax, height_perc=.35, width_perc=.35):
    cover_left = False
    cover_right = False
    xlim = ax.get_xlim()
    full = np.sum(np.abs(xlim))
    left_bound = xlim[0] + width_perc * full
    right_bound = xlim[1] - width_perc * full
    top_bound = (1 - height_perc) * ax.get_ylim()[1]
    for child in ax.get_children():
        if isinstance(child, mpl.patches.Rectangle):
            xx = child.get_x()
            too_high = child.get_height() > 0.5 * ax.get_ylim()[1]
            if child.get_height() > top_bound:
                if xx < left_bound:
                    cover_left = True
                elif xx > right_bound:
                    cover_right = True

    if not cover_right:
        loc = 1
    elif not cover_left:
        loc = 2
    else:
        loc = 9
    return loc

def plot_dataset_zoomin(store, varname, bound=0.1):
    with sns.axes_style("white"):
        df = store[sep_prefix + varname]
        df.name = varname
        df.dropna(inplace=True)
        fig = plt.figure()
        ax = sns.distplot(df.loc[df.abs() < bound],
                          hist_kws={'log': False}, kde=True)

        sns.despine(ax=ax)
    return fig

def plot_dataset_dist(store, varname, cutoff=0.01):
    with sns.axes_style("white"):
        df = store[sep_prefix + varname]
        df.name = varname
        df.dropna(inplace=True)
        fig = plt.figure()
        ax = sns.distplot(df.loc[(df.quantile(cutoff) < df) &
                                 (df < df.quantile(1 - cutoff))],
                          hist_kws={'log': False}, kde=True)
        ax.set_ylabel('density')

        sns.despine(ax=ax)
        loc = determine_subax_loc(ax)
        subax = inset_axes(ax,
                           width="30%",
                           height="30%",
                           loc=loc)
        if is_partial_particle(df.name):
            quant_bound = .15
            low_bound = max(-1, df.quantile(quant_bound))
            high_bound = min(1, df.quantile(1 - quant_bound))
        else:
            low_bound = -1
            high_bound = 1
        sns.distplot(df.loc[(low_bound < df) & (df < high_bound)],
                     kde=True, kde_kws={'gridsize': 200},
                     ax=subax)
        if loc == 2:
            subax.yaxis.set_label_position("right")
            subax.yaxis.tick_right()
            sns.despine(ax=subax, left=True, right=False)
        else:
            sns.despine(ax=subax)
        subax.set_xlabel('')
    return fig

def generate_store_name(set='training', unstable=True, gen=3, filter_id=8, dim=7):
    store_name = '{!s}_gen{!s}_{!s}D_nions0_flat_filter{!s}.h5.1'.format(set, gen, dim, filter_id)
    if unstable and set == 'training':
        store_name = '_'.join(['unstable', store_name])
    return store_name

def plot_pure_network_dataset_dist(self):
    filter_id = self.filter_id
    dim = len(net.feature_names)
    #store_name = 'unstable_training_gen{!s}_{!s}D_nions0_flat_filter{!s}.h5'.format(2, filter_id, dim)
    store_name = generate_store_name(True, 3, filter_id, dim)
    store = pd.HDFStore(os.path.join(qlknn_root, store_name))
    embed()
    for train_dim in net.target_names:
        plot_dataset_dist(store, train_dim)
        deconstruct = re.split('_div_|_plus_', train_dim)
        if len(deconstruct) > 1:
            for sub_dim in deconstruct:
                plot_dataset_dist(store, sub_dim)
Network.plot_dataset_dist = plot_pure_network_dataset_dist

def generate_dataset_report(store, plot_pure=True, plot_heat=True, plot_particle=True, plot_rot=False, plot_full=False, plot_diffusion=False, plot_nonleading=False, plot_momentum=False, verbose_debug=False):
    with PdfPages('multipage_pdf.pdf') as pdf:
        for varname in store:
            varname = varname.replace(sep_prefix, '', 1)
            if verbose_debug:
                print(varname)
                print('is_full_transport', is_full_transport(varname))
                print('is_pure', is_pure(varname))
                print('is_partial_rot', is_partial_rot(varname))
                print('is_partial_heat', is_partial_heat(varname))
                print('is_partial_particle', is_partial_particle(varname))
                print('is_partial_diffusion', is_partial_diffusion(varname))
                print('is_leading', is_leading(varname))

            if (is_transport(varname) and
                ((plot_full and is_full_transport(varname)) or not is_full_transport(varname)) and
                ((plot_pure and is_pure(varname)) or not is_pure(varname)) and
                ((plot_rot and is_partial_rot(varname)) or not is_partial_rot(varname)) and
                ((plot_heat and is_partial_heat(varname)) or not is_partial_heat(varname)) and
                ((plot_particle and is_partial_particle(varname)) or not is_partial_particle(varname)) and
                ((plot_diffusion and is_partial_diffusion(varname)) or not is_partial_diffusion(varname)) and
                ((plot_momentum and is_partial_momentum(varname)) or not is_partial_momentum(varname)) and
                ((plot_nonleading and not is_leading(varname)) or is_leading(varname))):
                  #(not is_full_transport(varname) and
                  # not is_partial_rot(varname) and
                  # not is_partial_heat(varname) and
                  # not is_partial_particle(varname) and
                  # not is_partial_diffusion(varname) and
                  # is_leading(varname))):
                print('Plotting', varname)
                fig = plot_dataset_dist(store, varname)
                pdf.savefig(fig)
                plt.close(fig)
                try:
                    fig = plot_dataset_zoomin(store, varname)
                except ZeroDivisionError:
                    fig = plt.figure()
                pdf.savefig(fig)
                plt.close(fig)

#net = Network.get_by_id(1409)

if __name__ == '__main__':
    store = pd.HDFStore(os.path.join(qlknn_root, generate_store_name(set='sane', unstable=False, dim=7)), 'r')
    generate_dataset_report(store)
    embed()
<|endoftext|>


<commit_before>#!/usr/bin/env python3
# coding=utf-8

def build_standalone_slaver(outfile_name="slaver_standalone.py"):
    import os
    import shutil
    base_path = os.path.dirname(os.path.abspath(__file__))
    output_fullpath = os.path.join(base_path, outfile_name)

    if os.path.exists(output_fullpath):
        ch = input("target file: " + output_fullpath + " already exist, overwrite it?\n(y/N)")
        if ch not in ("y", "Y", "yes"):
            print("user abort")
            return None
        else:
            os.remove(output_fullpath)

    shutil.copy(
        os.path.join(base_path, "common_func.py"),
        os.path.join(base_path, outfile_name, )
    )
    with open("slaver.py", "r", encoding='utf-8') as fr:
        with open(outfile_name, "a", encoding="utf-8") as fw:
            fw.write(fr.read())

    print("generate complete!\noutput file:", output_fullpath)

if __name__ == '__main__':
    build_standalone_slaver()
<commit_msg>alpha6<commit_after>#!/usr/bin/env python3
# coding=utf-8

def build_standalone_slaver(outfile_name="slaver_standalone.py"):
    import os
    import shutil
    base_path = os.path.dirname(os.path.abspath(__file__))
    output_fullpath = os.path.join(base_path, outfile_name)

    if os.path.exists(output_fullpath):
        ch = input("target file: " + output_fullpath + " already exist, overwrite it?\n(y/N)")
        if ch not in ("y", "Y", "yes"):
            print("user abort")
            return None
        else:
            os.remove(output_fullpath)

    shutil.copy(
        os.path.join(base_path, "common_func.py"),
        os.path.join(base_path, outfile_name, )
    )
    with open("slaver.py", "r", encoding='utf-8') as fr:
        slaver = fr.read()

    slaver = slaver.replace("from __future__ import", "# from __future__ import", 1)
    slaver = slaver.replace("from common_func import", "# from common_func import", 1)

    with open(outfile_name, "a", encoding="utf-8") as fw:
        fw.write(slaver)

    print("generate complete!\noutput file:", output_fullpath)

if __name__ == '__main__':
    build_standalone_slaver()
<|endoftext|>


<commit_before>#!/usr/bin/env python
# Copyright 2014 The Swarming Authors. All rights reserved.
# Use of this source code is governed by the Apache v2.0 license that can be
# found in the LICENSE file.

"""Integration test for the Swarming server, Swarming bot and Swarming client.

It starts both a Swarming server and a Swarming bot and triggers tasks with the
Swarming client to ensure the system works end to end.
"""

import json
import glob
import logging
import os
import shutil
import signal
import socket
import subprocess
import sys
import tempfile
import time
import unittest
import urllib

APP_DIR = os.path.dirname(os.path.abspath(__file__))
BOT_DIR = os.path.join(APP_DIR, 'swarming_bot')
CLIENT_DIR = os.path.join(APP_DIR, '..', '..', 'client')

from tools import start_bot
from tools import start_servers

sys.path.insert(0, BOT_DIR)

import os_utilities

class SwarmingClient(object):
  def __init__(self, swarming_server):
    self._swarming_server = swarming_server
    self._tmpdir = tempfile.mkdtemp(prefix='swarming_client')
    self._index = 0

  def task_trigger(self, args):
    """Triggers a task and return the task id."""
    h, tmp = tempfile.mkstemp(prefix='swarming_smoke_test', suffix='.json')
    os.close(h)
    try:
      cmd = [
        '--user', 'joe@localhost', '-d', 'cpu', 'x86', '--dump-json', tmp,
        '--raw-cmd',
      ]
      cmd.extend(args)
      assert not self._run('trigger', cmd), args
      with open(tmp, 'r') as f:
        data = json.load(f)
        return data['tasks'].popitem()[1]['task_id']
    finally:
      os.remove(tmp)

  def task_collect(self, task_id):
    """Collects the results for a task."""
    h, tmp = tempfile.mkstemp(prefix='swarming_smoke_test', suffix='.json')
    os.close(h)
    try:
      # swarming.py collect will return the exit code of the task.
      self._run('collect', ['--task-summary-json', tmp, task_id])
      with open(tmp, 'r') as f:
        return json.load(f)
    finally:
      os.remove(tmp)

  def cleanup(self):
    if self._tmpdir:
      shutil.rmtree(self._tmpdir)
      self._tmpdir = None

  def dump_log(self):
    print >> sys.stderr, '-' * 60
    print >> sys.stderr, 'Client calls'
    print >> sys.stderr, '-' * 60
    for i in xrange(self._index):
      with open(os.path.join(self._tmpdir, 'client_%d.log' % i), 'rb') as f:
        log = f.read().strip('\n')
      for l in log.splitlines():
        sys.stderr.write('  %s\n' % l)

  def _run(self, command, args):
    name = os.path.join(self._tmpdir, 'client_%d.log' % self._index)
    self._index += 1
    cmd = [
      sys.executable, 'swarming.py', command, '-S', self._swarming_server,
      '--verbose',
    ] + args
    with open(name, 'wb') as f:
      f.write('\nRunning: %s\n' % ' '.join(cmd))
      f.flush()
      p = subprocess.Popen(
          cmd, stdout=f, stderr=subprocess.STDOUT, cwd=CLIENT_DIR)
      p.communicate()
      return p.returncode

def gen_expected(**kwargs):
  expected = {
    u'abandoned_ts': None,
    u'bot_dimensions': None,
    u'bot_id': unicode(socket.getfqdn().split('.', 1)[0]),
    u'children_task_ids': [],
    u'cost_saved_usd': None,
    u'deduped_from': None,
    u'exit_codes': [0],
    u'failure': False,
    u'internal_failure': False,
    u'isolated_out': None,
    u'name': u'',
    u'outputs': [u'hi\n'],
    u'outputs_ref': None,
    u'properties_hash': None,
    u'server_versions': [u'1'],
    u'state': 0x70,  # task_result.State.COMPLETED.
    u'tags': [u'cpu:x86', u'priority:100', u'user:joe@localhost'],
    u'try_number': 1,
    u'user': u'joe@localhost',
  }
  assert set(expected).issuperset(kwargs)
  expected.update(kwargs)
  return expected

def handle_interrupt(fn):
  """Makes a unittest.TestCase test handle Ctrl-C properly."""
  def hook(self):
    try:
      fn(self)
    except KeyboardInterrupt:
      self.fail('<Ctrl-C>')
      if self._bot:
        self._bot.kill()
  return hook

class SwarmingTestCase(unittest.TestCase):
  """Test case class for Swarming integration tests."""
  def setUp(self):
    super(SwarmingTestCase, self).setUp()
    self._bot = None
    self._client = None
    self._servers = None

  def tearDown(self):
    # Kill bot, kill server, print logs if failed, delete tmpdir, call super.
    try:
      try:
        try:
          try:
            if self._bot:
              self._bot.stop()
          finally:
            if self._servers:
              self._servers.stop()
        finally:
          if self.has_failed() or is_verbose():
            # Print out the logs before deleting them.
            if self._bot:
              self._bot.dump_log()
            if self._servers:
              self._servers.dump_log()
            if self._client:
              self._client.dump_log()
      finally:
        if self._client:
          self._client.cleanup()
    finally:
      super(SwarmingTestCase, self).tearDown()

  def finish_setup(self):
    """Uploads bot_config.py and starts a bot.

    Should be called from test_* method (not from setUp), since if setUp fails
    tearDown is not getting called (and finish_setup can fail because it uses
    various server endpoints).
    """
    self._servers = start_servers.LocalServers(False)
    self._servers.start()
    self._bot = start_bot.LocalBot(self._servers.swarming_server.url)
    self._bot.start()

    # Replace bot_config.py.
    with open(os.path.join(BOT_DIR, 'bot_config.py'), 'rb') as f:
      bot_config_content = f.read() + '\n'
    # This will restart the bot.
    res = self._servers.http_client.request(
        '/restricted/upload/bot_config',
        body=urllib.urlencode({'script': bot_config_content}))
    if not res:
      raise Exception('Failed to upload bot_config')

  def has_failed(self):
    # pylint: disable=E1101
    return not self._resultForDoCleanups.wasSuccessful()

  @handle_interrupt
  def test_integration(self):
    """Runs a few task requests and wait for results."""
    self.finish_setup()
    def cmd(exit_code=0):
      return [
        'python', '-c', 'import sys; print(\'hi\'); sys.exit(%d)' % exit_code,
      ]

    # tuple(task_request, expectation)
    dimensions = os_utilities.get_dimensions()
    tasks = [
      (
        ['-T', 'simple_success', '--'] + cmd(),
        gen_expected(name=u'simple_success', bot_dimensions=dimensions),
      ),
      (
        ['-T', 'simple_failure', '--'] + cmd(1),
        gen_expected(
          name=u'simple_failure',
          bot_dimensions=dimensions,
          exit_codes=[1],
          failure=True),
      ),
      (
        ['-T', 'invalid', '--', 'unknown_invalid_command'],
        gen_expected(
            name=u'invalid',
            bot_dimensions=dimensions,
            exit_codes=[1],
            failure=True,
            outputs=[
              u'Command "unknown_invalid_command" failed to start.\n'
              u'Error: [Errno 2] No such file or directory',
            ]),
      ),
      (
        [
          # Need to flush to ensure it will be sent to the server.
          '-T', 'hard_timeout', '--hard-timeout', '1', '--',
          'python', '-c',
          'import time,sys; sys.stdout.write(\'hi\\n\'); '
            'sys.stdout.flush(); time.sleep(120)',
        ],
        gen_expected(
            name=u'hard_timeout',
            bot_dimensions=dimensions,
            exit_codes=[-signal.SIGTERM],
            failure=True,
            state=0x40),  # task_result.State.TIMED_OUT
      ),
      (
        [
          # Need to flush to ensure it will be sent to the server.
          '-T', 'io_timeout', '--io-timeout', '1', '--',
          'python', '-c',
          'import time,sys; sys.stdout.write(\'hi\\n\'); '
            'sys.stdout.flush(); time.sleep(120)',
        ],
        gen_expected(
            name=u'io_timeout',
            bot_dimensions=dimensions,
            exit_codes=[-signal.SIGTERM],
            failure=True,
            state=0x40),  # task_result.State.TIMED_OUT
      ),
      (
        ['-T', 'ending_simple_success', '--'] + cmd(),
        gen_expected(name=u'ending_simple_success', bot_dimensions=dimensions),
      ),
    ]

    # tuple(task_id, expectation)
    self._client = SwarmingClient(self._servers.swarming_server.url)
    running_tasks = [
      (self._client.task_trigger(args), expected) for args, expected in tasks
    ]

    for task_id, expectation in running_tasks:
      self.assertResults(expectation, self._client.task_collect(task_id))

  def assertResults(self, expected, result):
    self.assertEqual(['shards'], result.keys())
    self.assertEqual(1, len(result['shards']))
    result = result['shards'][0].copy()
    # These are not deterministic (or I'm too lazy to calculate the value).
    self.assertTrue(result.pop('bot_version'))
    self.assertTrue(result.pop('costs_usd'))
    self.assertTrue(result.pop('created_ts'))
    self.assertTrue(result.pop('completed_ts'))
    self.assertTrue(result.pop('durations'))
    self.assertTrue(result.pop('id'))
    self.assertTrue(result.pop('modified_ts'))
    self.assertTrue(result.pop('started_ts'))
    self.assertEqual(expected, result)

def is_verbose():
  return '-v' in sys.argv

if __name__ == '__main__':
  if is_verbose():
    logging.basicConfig(level=logging.INFO)
    unittest.TestCase.maxDiff = None
  else:
    logging.basicConfig(level=logging.ERROR)
  unittest.main()
<commit_msg>Add smoke test to ensure invalid utf-8 streams are supported.<commit_after>#!/usr/bin/env python
# Copyright 2014 The Swarming Authors. All rights reserved.
# Use of this source code is governed by the Apache v2.0 license that can be
# found in the LICENSE file.

"""Integration test for the Swarming server, Swarming bot and Swarming client.

It starts both a Swarming server and a Swarming bot and triggers tasks with the
Swarming client to ensure the system works end to end.
"""

import json
import glob
import logging
import os
import shutil
import signal
import socket
import subprocess
import sys
import tempfile
import time
import unittest
import urllib

APP_DIR = os.path.dirname(os.path.abspath(__file__))
BOT_DIR = os.path.join(APP_DIR, 'swarming_bot')
CLIENT_DIR = os.path.join(APP_DIR, '..', '..', 'client')

from tools import start_bot
from tools import start_servers

sys.path.insert(0, BOT_DIR)

import os_utilities

class SwarmingClient(object):
  def __init__(self, swarming_server):
    self._swarming_server = swarming_server
    self._tmpdir = tempfile.mkdtemp(prefix='swarming_client')
    self._index = 0

  def task_trigger(self, args):
    """Triggers a task and return the task id."""
    h, tmp = tempfile.mkstemp(prefix='swarming_smoke_test', suffix='.json')
    os.close(h)
    try:
      cmd = [
        '--user', 'joe@localhost', '-d', 'cpu', 'x86', '--dump-json', tmp,
        '--raw-cmd',
      ]
      cmd.extend(args)
      assert not self._run('trigger', cmd), args
      with open(tmp, 'r') as f:
        data = json.load(f)
        return data['tasks'].popitem()[1]['task_id']
    finally:
      os.remove(tmp)

  def task_collect(self, task_id):
    """Collects the results for a task."""
    h, tmp = tempfile.mkstemp(prefix='swarming_smoke_test', suffix='.json')
    os.close(h)
    try:
      # swarming.py collect will return the exit code of the task.
      self._run('collect', ['--task-summary-json', tmp, task_id])
      with open(tmp, 'r') as f:
        return json.load(f)
    finally:
      os.remove(tmp)

  def cleanup(self):
    if self._tmpdir:
      shutil.rmtree(self._tmpdir)
      self._tmpdir = None

  def dump_log(self):
    print >> sys.stderr, '-' * 60
    print >> sys.stderr, 'Client calls'
    print >> sys.stderr, '-' * 60
    for i in xrange(self._index):
      with open(os.path.join(self._tmpdir, 'client_%d.log' % i), 'rb') as f:
        log = f.read().strip('\n')
      for l in log.splitlines():
        sys.stderr.write('  %s\n' % l)

  def _run(self, command, args):
    name = os.path.join(self._tmpdir, 'client_%d.log' % self._index)
    self._index += 1
    cmd = [
      sys.executable, 'swarming.py', command, '-S', self._swarming_server,
      '--verbose',
    ] + args
    with open(name, 'wb') as f:
      f.write('\nRunning: %s\n' % ' '.join(cmd))
      f.flush()
      p = subprocess.Popen(
          cmd, stdout=f, stderr=subprocess.STDOUT, cwd=CLIENT_DIR)
      p.communicate()
      return p.returncode

def gen_expected(**kwargs):
  expected = {
    u'abandoned_ts': None,
    u'bot_dimensions': None,
    u'bot_id': unicode(socket.getfqdn().split('.', 1)[0]),
    u'children_task_ids': [],
    u'cost_saved_usd': None,
    u'deduped_from': None,
    u'exit_codes': [0],
    u'failure': False,
    u'internal_failure': False,
    u'isolated_out': None,
    u'name': u'',
    u'outputs': [u'hi\n'],
    u'outputs_ref': None,
    u'properties_hash': None,
    u'server_versions': [u'1'],
    u'state': 0x70,  # task_result.State.COMPLETED.
    u'tags': [u'cpu:x86', u'priority:100', u'user:joe@localhost'],
    u'try_number': 1,
    u'user': u'joe@localhost',
  }
  assert set(expected).issuperset(kwargs)
  expected.update(kwargs)
  return expected

def handle_interrupt(fn):
  """Makes a unittest.TestCase test handle Ctrl-C properly."""
  def hook(self):
    try:
      fn(self)
    except KeyboardInterrupt:
      self.fail('<Ctrl-C>')
      if self._bot:
        self._bot.kill()
  return hook

class SwarmingTestCase(unittest.TestCase):
  """Test case class for Swarming integration tests."""
  def setUp(self):
    super(SwarmingTestCase, self).setUp()
    self._bot = None
    self._client = None
    self._servers = None

  def tearDown(self):
    # Kill bot, kill server, print logs if failed, delete tmpdir, call super.
    try:
      try:
        try:
          try:
            if self._bot:
              self._bot.stop()
          finally:
            if self._servers:
              self._servers.stop()
        finally:
          if self.has_failed() or is_verbose():
            # Print out the logs before deleting them.
            if self._bot:
              self._bot.dump_log()
            if self._servers:
              self._servers.dump_log()
            if self._client:
              self._client.dump_log()
      finally:
        if self._client:
          self._client.cleanup()
    finally:
      super(SwarmingTestCase, self).tearDown()

  def finish_setup(self):
    """Uploads bot_config.py and starts a bot.

    Should be called from test_* method (not from setUp), since if setUp fails
    tearDown is not getting called (and finish_setup can fail because it uses
    various server endpoints).
    """
    self._servers = start_servers.LocalServers(False)
    self._servers.start()
    self._bot = start_bot.LocalBot(self._servers.swarming_server.url)
    self._bot.start()

    # Replace bot_config.py.
    with open(os.path.join(BOT_DIR, 'bot_config.py'), 'rb') as f:
      bot_config_content = f.read() + '\n'
    # This will restart the bot.
    res = self._servers.http_client.request(
        '/restricted/upload/bot_config',
        body=urllib.urlencode({'script': bot_config_content}))
    if not res:
      raise Exception('Failed to upload bot_config')

  def has_failed(self):
    # pylint: disable=E1101
    return not self._resultForDoCleanups.wasSuccessful()

  @handle_interrupt
  def test_integration(self):
    """Runs a few task requests and wait for results."""
    self.finish_setup()
    def cmd(exit_code=0):
      return [
        'python', '-c', 'import sys; print(\'hi\'); sys.exit(%d)' % exit_code,
      ]

    # tuple(task_request, expectation)
    dimensions = os_utilities.get_dimensions()
    # A string of a letter 'A', UTF-8 BOM then UTF-16 BOM then UTF-EDBCDIC then
    # invalid UTF-8 and the letter 'B'. It is double escaped so it can be passed
    # down the shell.
    invalid_bytes = 'A\\xEF\\xBB\\xBF\\xFE\\xFF\\xDD\\x73\\x66\\x73\\xc3\\x28B'
    tasks = [
      (
        ['-T', 'simple_success', '--'] + cmd(),
        gen_expected(name=u'simple_success', bot_dimensions=dimensions),
      ),
      (
        ['-T', 'simple_failure', '--'] + cmd(1),
        gen_expected(
          name=u'simple_failure',
          bot_dimensions=dimensions,
          exit_codes=[1],
          failure=True),
      ),
      (
        ['-T', 'invalid', '--', 'unknown_invalid_command'],
        gen_expected(
            name=u'invalid',
            bot_dimensions=dimensions,
            exit_codes=[1],
            failure=True,
            outputs=[
              u'Command "unknown_invalid_command" failed to start.\n'
              u'Error: [Errno 2] No such file or directory',
            ]),
      ),
      (
        [
          # Need to flush to ensure it will be sent to the server.
          '-T', 'hard_timeout', '--hard-timeout', '1', '--',
          'python', '-c',
          'import time,sys; sys.stdout.write(\'hi\\n\'); '
            'sys.stdout.flush(); time.sleep(120)',
        ],
        gen_expected(
            name=u'hard_timeout',
            bot_dimensions=dimensions,
            exit_codes=[-signal.SIGTERM],
            failure=True,
            state=0x40),  # task_result.State.TIMED_OUT
      ),
      (
        [
          # Need to flush to ensure it will be sent to the server.
          '-T', 'io_timeout', '--io-timeout', '1', '--',
          'python', '-c',
          'import time,sys; sys.stdout.write(\'hi\\n\'); '
            'sys.stdout.flush(); time.sleep(120)',
        ],
        gen_expected(
            name=u'io_timeout',
            bot_dimensions=dimensions,
            exit_codes=[-signal.SIGTERM],
            failure=True,
            state=0x40),  # task_result.State.TIMED_OUT
      ),
      (
        [
          '-T', 'non_utf8', '--',
          'python', '-c', 'print(\'' + invalid_bytes + '\')',
        ],
        gen_expected(
            name=u'non_utf8', bot_dimensions=dimensions,
            # The string is mostly converted to 'Replacement Character'.
            outputs=[u'A\ufeff\ufffd\ufffd\ufffdsfs\ufffd(B\n']),
      ),
      (
        ['-T', 'ending_simple_success', '--'] + cmd(),
        gen_expected(name=u'ending_simple_success', bot_dimensions=dimensions),
      ),
    ]

    # tuple(task_id, expectation)
    self._client = SwarmingClient(self._servers.swarming_server.url)
    running_tasks = [
      (self._client.task_trigger(args), expected) for args, expected in tasks
    ]

    for task_id, expectation in running_tasks:
      self.assertResults(expectation, self._client.task_collect(task_id))

  def assertResults(self, expected, result):
    self.assertEqual(['shards'], result.keys())
    self.assertEqual(1, len(result['shards']))
    result = result['shards'][0].copy()
    # These are not deterministic (or I'm too lazy to calculate the value).
    self.assertTrue(result.pop('bot_version'))
    self.assertTrue(result.pop('costs_usd'))
    self.assertTrue(result.pop('created_ts'))
    self.assertTrue(result.pop('completed_ts'))
    self.assertTrue(result.pop('durations'))
    self.assertTrue(result.pop('id'))
    self.assertTrue(result.pop('modified_ts'))
    self.assertTrue(result.pop('started_ts'))
    self.assertEqual(expected, result)

def is_verbose():
  return '-v' in sys.argv

if __name__ == '__main__':
  if is_verbose():
    logging.basicConfig(level=logging.INFO)
    unittest.TestCase.maxDiff = None
  else:
    logging.basicConfig(level=logging.ERROR)
  unittest.main()
<|endoftext|>