<commit_before># Copyright 2017-2019 TensorHub, Inc.
#
# 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
from __future__ import division

import csv
import logging
import os
import sys

from guild import cli
from guild import config
from guild import index2 as indexlib
from guild import guildfile
from guild import op_util
from guild import query
from guild import run as runlib
from guild import tabview
from guild import util
from guild import var

from . import runs_impl

log = logging.getLogger("guild")

BASE_COLS = ".run, .operation, .started, .time, .status, .label"
MIN_COLS = ".run"
DEFAULT_COMPARE = ["loss step as step", "loss"]

NO_RUNS_CAPTION = "no runs"

def main(args):
    if args.print_scalars:
        _print_scalars(args)
    elif args.format == "csv":
        _print_csv(args)
    elif args.format == "table":
        _print_table(args)
    else:
        _tabview(args)

def _print_scalars(args):
    runs = runs_impl.runs_for_args(args)
    index = indexlib.RunIndex()
    index.refresh(runs, ["scalar"])
    for run in runs:
        cli.out("[%s] %s" % (run.short_id, op_util.format_op_desc(run)))
        for s in index.run_scalars(run):
            prefix = s["prefix"]
            if prefix:
                cli.out("  %s#%s" % (s["prefix"], s["tag"]))
            else:
                cli.out("  %s" % s["tag"])

def _print_csv(args):
    data = _get_data(args, format_cells=False, skip_header_if_empty=True)
    if not data:
        return
    writer = csv.writer(sys.stdout)
    for row in data:
        writer.writerow(row)

def _get_data(args, format_cells=True, skip_header_if_empty=False):
    index = indexlib.RunIndex()
    cb = _get_data_cb(
        args, index,
        format_cells=format_cells,
        skip_header_if_empty=skip_header_if_empty)
    data, logs = cb()
    for record in logs:
        log.handle(record)
    return data

def _print_table(args):
    data = _get_data(args, skip_header_if_empty=True)
    if not data:
        return
    cols = data[0]
    col_indexes = list(zip(cols, range(len(cols))))
    def format_row(row):
        return {
            col_name: row[i]
            for col_name, i in col_indexes
        }
    heading = {
        col_name: col_name
        for col_name in cols
    }
    data = [heading] + [format_row(row) for row in data[1:]]
    cli.table(data, cols)

def _tabview(args):
    config.set_log_output(True)
    index = indexlib.RunIndex()
    tabview.view_runs(
        _get_data_cb(args, index),
        _get_run_detail_cb(index),
        _tabview_actions())

def _get_data_cb(args, index, format_cells=True, skip_header_if_empty=False):
    def f():
        _try_init_tf_logging()
        log_capture = util.LogCapture()
        with log_capture:
            runs = _runs_for_args(args)
            cols_table = _cols_table(runs, args)
            index.refresh(runs, _refresh_types(cols_table))
            table = _resolve_table_cols(cols_table, index)
            header = _table_header(table)
            if not header:
                if skip_header_if_empty:
                    return [], []
                header = [NO_RUNS_CAPTION]
            rows = _table_rows(table, header)
            if format_cells:
                _format_cells(rows, header, runs)
            log = log_capture.get_all()
            return [header] + rows, log
    return f

def _try_init_tf_logging():
    """Load TensorFlow, forcing init of TF logging.

    This is part of our handing of logging, which can interfere with
    the curses display used by tabview. By forcing an init of TF logs,
    we can patch loggers with LogCapture (see guild.tabview module)
    for display in a curses window.
    """
    try:
        import tensorflow as _
    except ImportError:
        pass

def _runs_for_args(args):
    runs = runs_impl.runs_for_args(args)
    if not args.include_batch:
        runs = [run for run in runs if not _is_batch(run)]
    return runs

def _is_batch(run):
    return os.path.exists(run.guild_path("proto"))

def _cols_table(runs, args):
    parse_cache = {}
    return [(run, _run_cols(run, args, parse_cache)) for run in runs]

def _run_cols(run, args, parse_cache):
    return (
        _core_cols(args, parse_cache),
        _op_cols(run, args, parse_cache),
        _user_cols(args, parse_cache)
    )

def _core_cols(args, parse_cache):
    return _colspec_cols(_core_colspec(args), parse_cache)

def _core_colspec(args):
    if args.skip_core:
        return MIN_COLS
    else:
        return BASE_COLS

def _colspec_cols(colspec, parse_cache):
    try:
        return parse_cache[colspec]
    except KeyError:
        try:
            cols = query.parse_colspec(colspec).cols
        except query.ParseError as e:
            log.warning("error parsing %r: %s", colspec, e)
            cols = []
        parse_cache[colspec] = cols
        return cols

def _op_cols(run, args, parse_cache):
    if args.skip_op_cols:
        return []
    compare = _run_op_compare(run)
    if not compare:
        return []
    cols = []
    for colspec in compare:
        cols.extend(_colspec_cols(colspec, parse_cache))
    return cols

def _run_op_compare(run):
    """Returns compare cols for run.

    If we can get the current compare cols for the run op source
    definition (in the Guild file) we use that, otherwise we use the
    run "compare" attr.
    """
    compare = _try_guildfile_compare(run)
    if compare is not None:
        return compare
    return run.get("compare", DEFAULT_COMPARE)

def _try_guildfile_compare(run):
    """Returns the current compare for run op if available."""
    if run.opref.pkg_type != "guildfile":
        return None
    gf_path = run.opref.pkg_name
    if not os.path.exists(gf_path):
        return None
    try:
        gf = guildfile.from_dir(gf_path)
    except guildfile.NoModels:
        return None
    else:
        return _try_guildfile_op_compare(
            gf, run.opref.model_name, run.opref.op_name)

def _try_guildfile_op_compare(gf, model_name, op_name):
    try:
        m = gf.models[model_name]
    except KeyError:
        return None
    else:
        op = m.get_operation(op_name)
        return op.compare if op else None

def _user_cols(args, parse_cache):
    if not args.columns:
        return []
    return _colspec_cols(args.columns, parse_cache)

def _refresh_types(cols_table):
    """Returns a set of types to refresh for cols.

    This scheme is used to avoid refreshing types that aren't needed
    (e.g. scalars are expensive to refresh).
    """
    types = set()
    for _run, sections in cols_table:
        for section in sections:
            for col in section:
                if isinstance(col, query.Flag):
                    types.add("flag")
                elif isinstance(col, query.Attr):
                    types.add("attr")
                elif isinstance(col, query.Scalar):
                    types.add("scalar")
    return types

def _resolve_table_cols(table, index):
    return [
        _resolve_table_section(run, section, index)
        for run, section in table
    ]

def _resolve_table_section(run, section, index):
    return tuple([_resolve_run_cols(run, cols, index) for cols in section])

def _resolve_run_cols(run, cols, index):
    return [(col.header, _col_data(run, col, index)) for col in cols]

def _col_data(run, col, index):
    if isinstance(col, query.Flag):
        return index.run_flag(run, col.name)
    elif isinstance(col, query.Attr):
        return index.run_attr(run, col.name)
    elif isinstance(col, query.Scalar):
        prefix, tag = col.split_key()
        return index.run_scalar(run, prefix, tag, col.qualifier, col.step)
    else:
        assert False, col

def _table_header(table):
    header = []
    # Build up headers from oldest run (bottom of table) first to
    # ensure consistency in col ordering as new runs appear.
    reversed_table = reversed(table)
    for section in zip(*reversed_table):
        for row in section:
            for name, _val in row:
                if name not in header:
                    header.append(name)
    return header

def _table_rows(table, header):
    return [_table_row(row, header) for row in table]

def _table_row(row, header):
    return [_row_val(col_name, row) for col_name in header]

def _row_val(col_name, sections):
    val = None
    for cells in sections:
        for cell_name, cell_val in reversed(cells):
            if cell_name == col_name and cell_val is not None:
                val = cell_val
    return val

def _format_cells(rows, col_names, runs):
    runs_lookup = {run.short_id: run for run in runs}
    for row in rows:
        for i, (name, val) in enumerate(zip(col_names, row)):
            if name == "operation":
                run_id = row[0]
                run = runs_lookup[run_id]
                if run.get("marked"):
                    row[i] += " [marked]"
            elif val is None:
                row[i] = ""
            elif isinstance(val, float):
                row[i] = "%0.6f" % val
            else:
                row[i] = str(val)

def _get_run_detail_cb(index):
    def f(data, y, _x):
        run_short_id = data[y][0]
        if run_short_id == NO_RUNS_CAPTION:
            return (
                "\nPress 'r' in the main screen to refresh the list.",
                "There are no matching runs currently")
        title = "Run {}".format(run_short_id)
        try:
            run_id, path = next(var.find_runs(run_short_id))
        except StopIteration:
            return "This run no longer exists.", title
        else:
            run = runlib.Run(run_id, path)
            index.refresh([run], ["scalar"])
            detail = _format_run_detail(run, index)
            return detail, title
    return f

def _format_run_detail(run, index):
    lines = [
        "Id: %s" % run.id,
        "Operation: %s" % op_util.format_op_desc(run),
        "Status: %s" % run.status,
        "Started: %s" % util.format_timestamp(run.get("started")),
        "Stopped: %s" % util.format_timestamp(run.get("stopped")),
        "Label: %s" % (run.get("label") or ""),
    ]
    flags = run.get("flags")
    if flags:
        lines.append("Flags:")
        for name, val in sorted(flags.items()):
            val = val if val is not None else ""
            lines.append("  {}: {}".format(name, val))
    scalars = list(index.run_scalars(run))
    if scalars:
        lines.append("Scalars:")
        for s in scalars:
            prefix = s["prefix"]
            if prefix:
                lines.append("  %s#%s" % (s["prefix"], s["tag"]))
            else:
                lines.append("  %s" % s["tag"])
    return "\n".join(lines)

def _tabview_actions():
    return [
        (("m", "Mark run"), _mark),
        (("u", "Unmark run"), _unmark),
    ]

def _mark(view):
    _mark_impl(view, True)

def _unmark(view):
    _mark_impl(view, False)

def _mark_impl(view, flag):
    if not view.data:
        return
    run_short_id = view.data[view.y][0]
    try:
        run_id, path = next(var.find_runs(run_short_id))
    except StopIteration:
        view.text_box(
            "This run no longer exists.\n"
            "\n"
            "Press 'q' to exist this screen and then press 'r' to "
            "refresh the list.",
            run_short_id)
    else:
        run = runlib.Run(run_id, path)
        _mark_run(run, flag)
        _try_mark_operation(view, flag)

def _mark_run(run, flag):
    if flag:
        run.write_attr("marked", True)
    else:
        run.del_attr("marked")

def _try_mark_operation(view, flag):
    try:
        op_index = view.header.index("operation")
    except ValueError:
        pass
    else:
        op = view.data[view.y][op_index]
        if flag:
            new_op = _ensure_marked(op)
        else:
            new_op = _strip_marked(op)
        view.data[view.y][op_index] = new_op
        view.column_width[op_index] = max(
            view.column_width[op_index],
            len(new_op))

def _ensure_marked(op):
    if not op.endswith(" [marked]"):
        op += " [marked]"
    return op

def _strip_marked(op):
    if op.endswith(" [marked]"):
        op = op[:-9]
    return op
<commit_msg>Include accuracy in default compare cols<commit_after># Copyright 2017-2019 TensorHub, Inc.
#
# 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
from __future__ import division

import csv
import logging
import os
import sys

from guild import cli
from guild import config
from guild import index2 as indexlib
from guild import guildfile
from guild import op_util
from guild import query
from guild import run as runlib
from guild import tabview
from guild import util
from guild import var

from . import runs_impl

log = logging.getLogger("guild")

BASE_COLS = ".run, .operation, .started, .time, .status, .label"
MIN_COLS = ".run"
DEFAULT_COMPARE = [
    "loss step as step",
    "loss",
    "acc",
]

NO_RUNS_CAPTION = "no runs"

def main(args):
    if args.print_scalars:
        _print_scalars(args)
    elif args.format == "csv":
        _print_csv(args)
    elif args.format == "table":
        _print_table(args)
    else:
        _tabview(args)

def _print_scalars(args):
    runs = runs_impl.runs_for_args(args)
    index = indexlib.RunIndex()
    index.refresh(runs, ["scalar"])
    for run in runs:
        cli.out("[%s] %s" % (run.short_id, op_util.format_op_desc(run)))
        for s in index.run_scalars(run):
            prefix = s["prefix"]
            if prefix:
                cli.out("  %s#%s" % (s["prefix"], s["tag"]))
            else:
                cli.out("  %s" % s["tag"])

def _print_csv(args):
    data = _get_data(args, format_cells=False, skip_header_if_empty=True)
    if not data:
        return
    writer = csv.writer(sys.stdout)
    for row in data:
        writer.writerow(row)

def _get_data(args, format_cells=True, skip_header_if_empty=False):
    index = indexlib.RunIndex()
    cb = _get_data_cb(
        args, index,
        format_cells=format_cells,
        skip_header_if_empty=skip_header_if_empty)
    data, logs = cb()
    for record in logs:
        log.handle(record)
    return data

def _print_table(args):
    data = _get_data(args, skip_header_if_empty=True)
    if not data:
        return
    cols = data[0]
    col_indexes = list(zip(cols, range(len(cols))))
    def format_row(row):
        return {
            col_name: row[i]
            for col_name, i in col_indexes
        }
    heading = {
        col_name: col_name
        for col_name in cols
    }
    data = [heading] + [format_row(row) for row in data[1:]]
    cli.table(data, cols)

def _tabview(args):
    config.set_log_output(True)
    index = indexlib.RunIndex()
    tabview.view_runs(
        _get_data_cb(args, index),
        _get_run_detail_cb(index),
        _tabview_actions())

def _get_data_cb(args, index, format_cells=True, skip_header_if_empty=False):
    def f():
        _try_init_tf_logging()
        log_capture = util.LogCapture()
        with log_capture:
            runs = _runs_for_args(args)
            cols_table = _cols_table(runs, args)
            index.refresh(runs, _refresh_types(cols_table))
            table = _resolve_table_cols(cols_table, index)
            header = _table_header(table)
            if not header:
                if skip_header_if_empty:
                    return [], []
                header = [NO_RUNS_CAPTION]
            rows = _table_rows(table, header)
            if format_cells:
                _format_cells(rows, header, runs)
            log = log_capture.get_all()
            return [header] + rows, log
    return f

def _try_init_tf_logging():
    """Load TensorFlow, forcing init of TF logging.

    This is part of our handing of logging, which can interfere with
    the curses display used by tabview. By forcing an init of TF logs,
    we can patch loggers with LogCapture (see guild.tabview module)
    for display in a curses window.
    """
    try:
        import tensorflow as _
    except ImportError:
        pass

def _runs_for_args(args):
    runs = runs_impl.runs_for_args(args)
    if not args.include_batch:
        runs = [run for run in runs if not _is_batch(run)]
    return runs

def _is_batch(run):
    return os.path.exists(run.guild_path("proto"))

def _cols_table(runs, args):
    parse_cache = {}
    return [(run, _run_cols(run, args, parse_cache)) for run in runs]

def _run_cols(run, args, parse_cache):
    return (
        _core_cols(args, parse_cache),
        _op_cols(run, args, parse_cache),
        _user_cols(args, parse_cache)
    )

def _core_cols(args, parse_cache):
    return _colspec_cols(_core_colspec(args), parse_cache)

def _core_colspec(args):
    if args.skip_core:
        return MIN_COLS
    else:
        return BASE_COLS

def _colspec_cols(colspec, parse_cache):
    try:
        return parse_cache[colspec]
    except KeyError:
        try:
            cols = query.parse_colspec(colspec).cols
        except query.ParseError as e:
            log.warning("error parsing %r: %s", colspec, e)
            cols = []
        parse_cache[colspec] = cols
        return cols

def _op_cols(run, args, parse_cache):
    if args.skip_op_cols:
        return []
    compare = _run_op_compare(run)
    if not compare:
        return []
    cols = []
    for colspec in compare:
        cols.extend(_colspec_cols(colspec, parse_cache))
    return cols

def _run_op_compare(run):
    """Returns compare cols for run.

    If we can get the current compare cols for the run op source
    definition (in the Guild file) we use that, otherwise we use the
    run "compare" attr.
    """
    compare = _try_guildfile_compare(run)
    if compare is not None:
        return compare
    return run.get("compare", DEFAULT_COMPARE)

def _try_guildfile_compare(run):
    """Returns the current compare for run op if available."""
    if run.opref.pkg_type != "guildfile":
        return None
    gf_path = run.opref.pkg_name
    if not os.path.exists(gf_path):
        return None
    try:
        gf = guildfile.from_dir(gf_path)
    except guildfile.NoModels:
        return None
    else:
        return _try_guildfile_op_compare(
            gf, run.opref.model_name, run.opref.op_name)

def _try_guildfile_op_compare(gf, model_name, op_name):
    try:
        m = gf.models[model_name]
    except KeyError:
        return None
    else:
        op = m.get_operation(op_name)
        return op.compare if op else None

def _user_cols(args, parse_cache):
    if not args.columns:
        return []
    return _colspec_cols(args.columns, parse_cache)

def _refresh_types(cols_table):
    """Returns a set of types to refresh for cols.

    This scheme is used to avoid refreshing types that aren't needed
    (e.g. scalars are expensive to refresh).
    """
    types = set()
    for _run, sections in cols_table:
        for section in sections:
            for col in section:
                if isinstance(col, query.Flag):
                    types.add("flag")
                elif isinstance(col, query.Attr):
                    types.add("attr")
                elif isinstance(col, query.Scalar):
                    types.add("scalar")
    return types

def _resolve_table_cols(table, index):
    return [
        _resolve_table_section(run, section, index)
        for run, section in table
    ]

def _resolve_table_section(run, section, index):
    return tuple([_resolve_run_cols(run, cols, index) for cols in section])

def _resolve_run_cols(run, cols, index):
    return [(col.header, _col_data(run, col, index)) for col in cols]

def _col_data(run, col, index):
    if isinstance(col, query.Flag):
        return index.run_flag(run, col.name)
    elif isinstance(col, query.Attr):
        return index.run_attr(run, col.name)
    elif isinstance(col, query.Scalar):
        prefix, tag = col.split_key()
        return index.run_scalar(run, prefix, tag, col.qualifier, col.step)
    else:
        assert False, col

def _table_header(table):
    header = []
    # Build up headers from oldest run (bottom of table) first to
    # ensure consistency in col ordering as new runs appear.
    reversed_table = reversed(table)
    for section in zip(*reversed_table):
        for row in section:
            for name, _val in row:
                if name not in header:
                    header.append(name)
    return header

def _table_rows(table, header):
    return [_table_row(row, header) for row in table]

def _table_row(row, header):
    return [_row_val(col_name, row) for col_name in header]

def _row_val(col_name, sections):
    val = None
    for cells in sections:
        for cell_name, cell_val in reversed(cells):
            if cell_name == col_name and cell_val is not None:
                val = cell_val
    return val

def _format_cells(rows, col_names, runs):
    runs_lookup = {run.short_id: run for run in runs}
    for row in rows:
        for i, (name, val) in enumerate(zip(col_names, row)):
            if name == "operation":
                run_id = row[0]
                run = runs_lookup[run_id]
                if run.get("marked"):
                    row[i] += " [marked]"
            elif val is None:
                row[i] = ""
            elif isinstance(val, float):
                row[i] = "%0.6f" % val
            else:
                row[i] = str(val)

def _get_run_detail_cb(index):
    def f(data, y, _x):
        run_short_id = data[y][0]
        if run_short_id == NO_RUNS_CAPTION:
            return (
                "\nPress 'r' in the main screen to refresh the list.",
                "There are no matching runs currently")
        title = "Run {}".format(run_short_id)
        try:
            run_id, path = next(var.find_runs(run_short_id))
        except StopIteration:
            return "This run no longer exists.", title
        else:
            run = runlib.Run(run_id, path)
            index.refresh([run], ["scalar"])
            detail = _format_run_detail(run, index)
            return detail, title
    return f

def _format_run_detail(run, index):
    lines = [
        "Id: %s" % run.id,
        "Operation: %s" % op_util.format_op_desc(run),
        "Status: %s" % run.status,
        "Started: %s" % util.format_timestamp(run.get("started")),
        "Stopped: %s" % util.format_timestamp(run.get("stopped")),
        "Label: %s" % (run.get("label") or ""),
    ]
    flags = run.get("flags")
    if flags:
        lines.append("Flags:")
        for name, val in sorted(flags.items()):
            val = val if val is not None else ""
            lines.append("  {}: {}".format(name, val))
    scalars = list(index.run_scalars(run))
    if scalars:
        lines.append("Scalars:")
        for s in scalars:
            prefix = s["prefix"]
            if prefix:
                lines.append("  %s#%s" % (s["prefix"], s["tag"]))
            else:
                lines.append("  %s" % s["tag"])
    return "\n".join(lines)

def _tabview_actions():
    return [
        (("m", "Mark run"), _mark),
        (("u", "Unmark run"), _unmark),
    ]

def _mark(view):
    _mark_impl(view, True)

def _unmark(view):
    _mark_impl(view, False)

def _mark_impl(view, flag):
    if not view.data:
        return
    run_short_id = view.data[view.y][0]
    try:
        run_id, path = next(var.find_runs(run_short_id))
    except StopIteration:
        view.text_box(
            "This run no longer exists.\n"
            "\n"
            "Press 'q' to exist this screen and then press 'r' to "
            "refresh the list.",
            run_short_id)
    else:
        run = runlib.Run(run_id, path)
        _mark_run(run, flag)
        _try_mark_operation(view, flag)

def _mark_run(run, flag):
    if flag:
        run.write_attr("marked", True)
    else:
        run.del_attr("marked")

def _try_mark_operation(view, flag):
    try:
        op_index = view.header.index("operation")
    except ValueError:
        pass
    else:
        op = view.data[view.y][op_index]
        if flag:
            new_op = _ensure_marked(op)
        else:
            new_op = _strip_marked(op)
        view.data[view.y][op_index] = new_op
        view.column_width[op_index] = max(
            view.column_width[op_index],
            len(new_op))

def _ensure_marked(op):
    if not op.endswith(" [marked]"):
        op += " [marked]"
    return op

def _strip_marked(op):
    if op.endswith(" [marked]"):
        op = op[:-9]
    return op
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
<commit_msg>Add arg normalizing function for labels and index<commit_after># -*- coding: utf-8 -*-

import numpy

import dask.array

from . import _compat

def _norm_input_labels_index(input, labels=None, index=None):
    """
    Normalize arguments to a standard form.
    """

    input = _compat._asarray(input)

    if labels is None:
        labels = (input != 0).astype(numpy.int64)
        index = None

    if index is None:
        labels = (labels > 0).astype(numpy.int64)
        index = dask.array.ones(tuple(), dtype=numpy.int64, chunks=tuple())

    labels = _compat._asarray(labels)
    index = _compat._asarray(index)

    # SciPy transposes these for some reason.
    # So we do the same thing here.
    # This only matters if index is some array.
    index = index.T

    if input.shape != labels.shape:
        raise ValueError("The input and labels arrays must be the same shape.")

    return (input, labels, index)
<|endoftext|>


<commit_before>import gevent
import gevent.baseserver
import gevent.event
import gevent.pool
import gevent.util
from gservice.util import defaultproperty

import functools

NOT_READY = 1

def require_ready(func):
    @functools.wraps(func)
    def wrapped(*args, **kwargs):
        assert args[0].ready, "Service must be ready to call this method."
        func(*args, **kwargs)
    return wrapped

class NamedService(object):
    def __init__(self, name, use_dict):
        print "ctor namedService", name, use_dict
        self.name = name
        self.use_dict = use_dict

    def __get__(self, instance, owner):
        return self.value

    def __set__(self, instance, value):
        self.value = value

    @property
    def value(self):
        return Service._get_named_service(self.name, self.use_dict)

    @value.setter
    def setvalue(self, value):
        Service.register_named_service(self.name, value, self.use_dict)

    def __str__(self):
        return str(self.value)

class Service(object):
    """Service base class for creating standalone or composable services

    A service is a container for two things: other services and greenlets. It
    then provides a common interface for starting and stopping them, based on
    a subset of the gevent.baseserver interface. This way you can include
    StreamServer or WSGIServer as child services.

    Service also lets you catch exceptions in the greenlets started from this
    service and introduces the concept of `ready`, letting us block :meth:`start`
    until the service is actually ready.

    """
    stop_timeout = defaultproperty(int, 1)
    ready_timeout = defaultproperty(int, 2)
    started = defaultproperty(bool, False)

    _children = defaultproperty(list)
    _stopped_event = defaultproperty(gevent.event.Event)
    _ready_event = defaultproperty(gevent.event.Event)
    _greenlets = defaultproperty(gevent.pool.Group)
    _error_handlers = defaultproperty(dict)

    # main services dictionary for looking up named services
    _main_services = {}

    @classmethod
    def register_named_service(cls, name, service, use_dict=_main_services):
        use_dict[name] = service

    @classmethod
    def _get_named_service(cls, name, use_dict=_main_services):
        return use_dict.get(name, None)

    def __new__(cls, *args, **kwargs):
        """
        Allow for Service('name') to lookup named global services.
        """
        if cls == Service:
            if 'mock_dict' in kwargs:
                use_dict = kwargs['mock_dict']
            else:
                use_dict = Service._main_services

            if 'name' in kwargs:
                name = kwargs['name']
            else:
                name = args[0]
            return NamedService(name, use_dict=use_dict)
        else:
            return super(Service, cls).__new__(cls)

    @property
    def ready(self):
        """This property returns whether this service is ready for business"""
        return self._ready_event.isSet()

    def set_ready(self):
        """Internal convenience function to proclaim readiness"""
        self._ready_event.set()

    def add_service(self, service):
        """Add a child service to this service

        The service added will be started when this service starts, before
        its :meth:`_start` method is called. It will also be stopped when this
        service stops, before its :meth:`_stop` method is called.

        """
        if isinstance(service, gevent.baseserver.BaseServer):
            service = ServiceWrapper(service)
        self._children.append(service)

    def remove_service(self, service):
        """Remove a child service from this service"""
        self._children.remove(service)

    def _wrap_errors(self, func):
        """Wrap a callable for triggering error handlers

        This is used by the greenlet spawn methods so you can handle known
        exception cases instead of gevent's default behavior of just printing
        a stack trace for exceptions running in parallel greenlets.

        """
        @functools.wraps(func)
        def wrapped_f(*args, **kwargs):
            exceptions = tuple(self._error_handlers.keys())
            try:
                return func(*args, **kwargs)
            except exceptions, exception:
                for type in self._error_handlers:
                    if isinstance(exception, type):
                        handler, greenlet = self._error_handlers[type]
                        self._wrap_errors(handler)(exception, greenlet)
                return exception
        return wrapped_f

    def catch(self, type, handler):
        """Set an error handler for exceptions.

        Catches exceptions of `type` raised in greenlets for this service
        and recursively any existing child services.
        """
        self._error_handlers[type] = (handler, gevent.getcurrent())
        for child in self._children:
            child.catch(type, handler)

    def spawn(self, func, *args, **kwargs):
        """Spawn a greenlet under this service"""
        func_wrap = self._wrap_errors(func)
        return self._greenlets.spawn(func_wrap, *args, **kwargs)

    def spawn_later(self, seconds, func, *args, **kwargs):
        """Spawn a greenlet in the future under this service"""
        group = self._greenlets
        func_wrap = self._wrap_errors(func)
        g = group.greenlet_class(func_wrap, *args, **kwargs)
        g.start_later(seconds)
        group.add(g)
        return g

    def start(self, block_until_ready=True):
        """Public interface for starting this service and children. By default it blocks until ready."""
        assert not self.started, '%s already started' % self.__class__.__name__
        self._stopped_event.clear()
        self._ready_event.clear()
        try:
            self.pre_start()
            for child in self._children:
                if isinstance(child, Service):
                    if not child.started:
                        child.start(block_until_ready)
                elif isinstance(child, gevent.baseserver.BaseServer):
                    if not child.started:
                        child.start()
            ready = self.do_start()
            if ready == NOT_READY and block_until_ready is True:
                self._ready_event.wait(self.ready_timeout)
            elif ready != NOT_READY:
                self._ready_event.set()
            self.started = True
            self.post_start()
        except:
            #stop may raise, print the current exception before trying
            import traceback
            traceback.print_exc()
            self.stop()
            raise

    def pre_start(self):
        pass

    def post_start(self):
        pass

    def do_start(self):
        """Empty implementation of service start. Implement me!

        Return `service.NOT_READY` to block until :meth:`set_ready` is
        called (or `ready_timeout` is reached).

        """
        return

    def stop(self, timeout=None):
        """Stop this service and child services

        If the server uses a pool to spawn the requests, then :meth:`stop` also waits
        for all the handlers to exit. If there are still handlers executing after *timeout*
        has expired (default 1 second), then the currently running handlers in the pool are killed."""
        if gevent.getcurrent() in self._greenlets:
            return gevent.spawn(self.stop)
        self.started = False
        try:
            self.pre_stop()
            for child in reversed(self._children):
                #iterate over children in reverse order, in case dependancies
                # were implied by the starting order
                if child.started:
                    child.stop()
            self.do_stop()
        finally:
            if timeout is None:
                timeout = self.stop_timeout
            if self._greenlets:
                self._greenlets.join(timeout=timeout)
                self._greenlets.kill(block=True, timeout=1)
            self._ready_event.clear()
            self._stopped_event.set()
            self.post_stop()

    def pre_stop(self):
        pass

    def post_stop(self):
        pass

    def do_stop(self):
        """Empty implementation of service stop. Implement me!"""
        return

    def reload(self):
        for child in self._children:
            child.reload()
        self.do_reload()

    def do_reload(self):
        """Empty implementation of service reload. Implement me!"""
        pass

    def serve_forever(self, stop_timeout=None, ready_callback=None):
        """Start the service if it hasn't been already started and wait until it's stopped."""
        if not self.started:
            self.start()
            if ready_callback:
                ready_callback()
        try:
            self._stopped_event.wait()
        except:
            #stop may raise, print the current exception before trying
            import traceback
            traceback.print_exc()

            self.stop(timeout=stop_timeout)
            raise

    def __enter__(self):
        self.start()
        return self

    def __exit__(self, type, value, traceback):
        self.stop()

class ServiceWrapper(Service):
    """Wrapper for gevent servers that are based on gevent.baseserver.BaseServer

    Although the Service class mostly looks like the BaseServer interface,
    there are certain extra methods (like exception catching) that are assumed
    to be available. This class allows us to wrap gevent servers so they actually
    are a Service. Ideally, you would never use this directly. For example, it's
    being used in Service.add_service to automatically wrap gevent servers passed in.
    """
    def __init__(self, klass_or_server, *args, **kwargs):
        super(ServiceWrapper, self).__init__()
        if isinstance(klass_or_server, gevent.baseserver.BaseServer):
            self.wrapped = klass_or_server
        else:
            self.wrapped = klass(*args, **kwargs)

    def do_start(self):
        self.spawn(self.wrapped.start)

    def do_stop(self):
        self.wrapped.stop()
<commit_msg>removing errant print statement.<commit_after>import gevent
import gevent.baseserver
import gevent.event
import gevent.pool
import gevent.util
from gservice.util import defaultproperty

import functools

NOT_READY = 1

def require_ready(func):
    @functools.wraps(func)
    def wrapped(*args, **kwargs):
        assert args[0].ready, "Service must be ready to call this method."
        func(*args, **kwargs)
    return wrapped

class NamedService(object):
    def __init__(self, name, use_dict):
        self.name = name
        self.use_dict = use_dict

    def __get__(self, instance, owner):
        return self.value

    def __set__(self, instance, value):
        self.value = value

    @property
    def value(self):
        return Service._get_named_service(self.name, self.use_dict)

    @value.setter
    def setvalue(self, value):
        Service.register_named_service(self.name, value, self.use_dict)

    def __str__(self):
        return str(self.value)

class Service(object):
    """Service base class for creating standalone or composable services

    A service is a container for two things: other services and greenlets. It
    then provides a common interface for starting and stopping them, based on
    a subset of the gevent.baseserver interface. This way you can include
    StreamServer or WSGIServer as child services.

    Service also lets you catch exceptions in the greenlets started from this
    service and introduces the concept of `ready`, letting us block :meth:`start`
    until the service is actually ready.

    """
    stop_timeout = defaultproperty(int, 1)
    ready_timeout = defaultproperty(int, 2)
    started = defaultproperty(bool, False)

    _children = defaultproperty(list)
    _stopped_event = defaultproperty(gevent.event.Event)
    _ready_event = defaultproperty(gevent.event.Event)
    _greenlets = defaultproperty(gevent.pool.Group)
    _error_handlers = defaultproperty(dict)

    # main services dictionary for looking up named services
    _main_services = {}

    @classmethod
    def register_named_service(cls, name, service, use_dict=_main_services):
        use_dict[name] = service

    @classmethod
    def _get_named_service(cls, name, use_dict=_main_services):
        return use_dict.get(name, None)

    def __new__(cls, *args, **kwargs):
        """
        Allow for Service('name') to lookup named global services.
        """
        if cls == Service:
            if 'mock_dict' in kwargs:
                use_dict = kwargs['mock_dict']
            else:
                use_dict = Service._main_services

            if 'name' in kwargs:
                name = kwargs['name']
            else:
                name = args[0]
            return NamedService(name, use_dict=use_dict)
        else:
            return super(Service, cls).__new__(cls)

    @property
    def ready(self):
        """This property returns whether this service is ready for business"""
        return self._ready_event.isSet()

    def set_ready(self):
        """Internal convenience function to proclaim readiness"""
        self._ready_event.set()

    def add_service(self, service):
        """Add a child service to this service

        The service added will be started when this service starts, before
        its :meth:`_start` method is called. It will also be stopped when this
        service stops, before its :meth:`_stop` method is called.

        """
        if isinstance(service, gevent.baseserver.BaseServer):
            service = ServiceWrapper(service)
        self._children.append(service)

    def remove_service(self, service):
        """Remove a child service from this service"""
        self._children.remove(service)

    def _wrap_errors(self, func):
        """Wrap a callable for triggering error handlers

        This is used by the greenlet spawn methods so you can handle known
        exception cases instead of gevent's default behavior of just printing
        a stack trace for exceptions running in parallel greenlets.

        """
        @functools.wraps(func)
        def wrapped_f(*args, **kwargs):
            exceptions = tuple(self._error_handlers.keys())
            try:
                return func(*args, **kwargs)
            except exceptions, exception:
                for type in self._error_handlers:
                    if isinstance(exception, type):
                        handler, greenlet = self._error_handlers[type]
                        self._wrap_errors(handler)(exception, greenlet)
                return exception
        return wrapped_f

    def catch(self, type, handler):
        """Set an error handler for exceptions.

        Catches exceptions of `type` raised in greenlets for this service
        and recursively any existing child services.
        """
        self._error_handlers[type] = (handler, gevent.getcurrent())
        for child in self._children:
            child.catch(type, handler)

    def spawn(self, func, *args, **kwargs):
        """Spawn a greenlet under this service"""
        func_wrap = self._wrap_errors(func)
        return self._greenlets.spawn(func_wrap, *args, **kwargs)

    def spawn_later(self, seconds, func, *args, **kwargs):
        """Spawn a greenlet in the future under this service"""
        group = self._greenlets
        func_wrap = self._wrap_errors(func)
        g = group.greenlet_class(func_wrap, *args, **kwargs)
        g.start_later(seconds)
        group.add(g)
        return g

    def start(self, block_until_ready=True):
        """Public interface for starting this service and children. By default it blocks until ready."""
        assert not self.started, '%s already started' % self.__class__.__name__
        self._stopped_event.clear()
        self._ready_event.clear()
        try:
            self.pre_start()
            for child in self._children:
                if isinstance(child, Service):
                    if not child.started:
                        child.start(block_until_ready)
                elif isinstance(child, gevent.baseserver.BaseServer):
                    if not child.started:
                        child.start()
            ready = self.do_start()
            if ready == NOT_READY and block_until_ready is True:
                self._ready_event.wait(self.ready_timeout)
            elif ready != NOT_READY:
                self._ready_event.set()
            self.started = True
            self.post_start()
        except:
            #stop may raise, print the current exception before trying
            import traceback
            traceback.print_exc()
            self.stop()
            raise

    def pre_start(self):
        pass

    def post_start(self):
        pass

    def do_start(self):
        """Empty implementation of service start. Implement me!

        Return `service.NOT_READY` to block until :meth:`set_ready` is
        called (or `ready_timeout` is reached).

        """
        return

    def stop(self, timeout=None):
        """Stop this service and child services

        If the server uses a pool to spawn the requests, then :meth:`stop` also waits
        for all the handlers to exit. If there are still handlers executing after *timeout*
        has expired (default 1 second), then the currently running handlers in the pool are killed."""
        if gevent.getcurrent() in self._greenlets:
            return gevent.spawn(self.stop)
        self.started = False
        try:
            self.pre_stop()
            for child in reversed(self._children):
                #iterate over children in reverse order, in case dependancies
                # were implied by the starting order
                if child.started:
                    child.stop()
            self.do_stop()
        finally:
            if timeout is None:
                timeout = self.stop_timeout
            if self._greenlets:
                self._greenlets.join(timeout=timeout)
                self._greenlets.kill(block=True, timeout=1)
            self._ready_event.clear()
            self._stopped_event.set()
            self.post_stop()

    def pre_stop(self):
        pass

    def post_stop(self):
        pass

    def do_stop(self):
        """Empty implementation of service stop. Implement me!"""
        return

    def reload(self):
        for child in self._children:
            child.reload()
        self.do_reload()

    def do_reload(self):
        """Empty implementation of service reload. Implement me!"""
        pass

    def serve_forever(self, stop_timeout=None, ready_callback=None):
        """Start the service if it hasn't been already started and wait until it's stopped."""
        if not self.started:
            self.start()
            if ready_callback:
                ready_callback()
        try:
            self._stopped_event.wait()
        except:
            #stop may raise, print the current exception before trying
            import traceback
            traceback.print_exc()

            self.stop(timeout=stop_timeout)
            raise

    def __enter__(self):
        self.start()
        return self

    def __exit__(self, type, value, traceback):
        self.stop()

class ServiceWrapper(Service):
    """Wrapper for gevent servers that are based on gevent.baseserver.BaseServer

    Although the Service class mostly looks like the BaseServer interface,
    there are certain extra methods (like exception catching) that are assumed
    to be available. This class allows us to wrap gevent servers so they actually
    are a Service. Ideally, you would never use this directly. For example, it's
    being used in Service.add_service to automatically wrap gevent servers passed in.
    """
    def __init__(self, klass_or_server, *args, **kwargs):
        super(ServiceWrapper, self).__init__()
        if isinstance(klass_or_server, gevent.baseserver.BaseServer):
            self.wrapped = klass_or_server
        else:
            self.wrapped = klass(*args, **kwargs)

    def do_start(self):
        self.spawn(self.wrapped.start)

    def do_stop(self):
        self.wrapped.stop()
<|endoftext|>


<commit_before>#!/usr/bin/env python
#
# portable serial port access with python
#
# This is a module that gathers a list of serial ports including details on OSX
#
# code originally from https://github.com/makerbot/pyserial/tree/master/serial/tools
# with contributions from cibomahto, dgs3, FarMcKon, tedbrandston
# and modifications by cliechti, hoihu, hardkrash
#
# (C) 2013-2015
#
# SPDX-License-Identifier:    BSD-3-Clause

# List all of the callout devices in OS/X by querying IOKit.

# See the following for a reference of how to do this:
# http://developer.apple.com/library/mac/#documentation/DeviceDrivers/Conceptual/WorkingWSerial/WWSerial_SerialDevs/SerialDevices.html#//apple_ref/doc/uid/TP30000384-CIHGEAFD

# More help from darwin_hid.py

# Also see the 'IORegistryExplorer' for an idea of what we are actually searching

import ctypes
from serial.tools import list_ports_common

iokit = ctypes.cdll.LoadLibrary(ctypes.util.find_library('IOKit'))
cf = ctypes.cdll.LoadLibrary(ctypes.util.find_library('CoreFoundation'))

kIOMasterPortDefault = ctypes.c_void_p.in_dll(iokit, "kIOMasterPortDefault")
kCFAllocatorDefault = ctypes.c_void_p.in_dll(cf, "kCFAllocatorDefault")

kCFStringEncodingMacRoman = 0

iokit.IOServiceMatching.restype = ctypes.c_void_p

iokit.IOServiceGetMatchingServices.argtypes = [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p]
iokit.IOServiceGetMatchingServices.restype = ctypes.c_void_p

iokit.IORegistryEntryGetParentEntry.argtypes = [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p]

iokit.IORegistryEntryCreateCFProperty.argtypes = [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p, ctypes.c_uint32]
iokit.IORegistryEntryCreateCFProperty.restype = ctypes.c_void_p

iokit.IORegistryEntryGetPath.argtypes = [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p]
iokit.IORegistryEntryGetPath.restype = ctypes.c_void_p

iokit.IORegistryEntryGetName.argtypes = [ctypes.c_void_p, ctypes.c_void_p]
iokit.IORegistryEntryGetName.restype = ctypes.c_void_p

iokit.IOObjectGetClass.argtypes = [ctypes.c_void_p, ctypes.c_void_p]
iokit.IOObjectGetClass.restype = ctypes.c_void_p

iokit.IOObjectRelease.argtypes = [ctypes.c_void_p]

cf.CFStringCreateWithCString.argtypes = [ctypes.c_void_p, ctypes.c_char_p, ctypes.c_int32]
cf.CFStringCreateWithCString.restype = ctypes.c_void_p

cf.CFStringGetCStringPtr.argtypes = [ctypes.c_void_p, ctypes.c_uint32]
cf.CFStringGetCStringPtr.restype = ctypes.c_char_p

cf.CFNumberGetValue.argtypes = [ctypes.c_void_p, ctypes.c_uint32, ctypes.c_void_p]
cf.CFNumberGetValue.restype = ctypes.c_void_p

# void CFRelease ( CFTypeRef cf );
cf.CFRelease.argtypes = [ctypes.c_void_p]
cf.CFRelease.restype = None

# CFNumber type defines
kCFNumberSInt8Type = 1
kCFNumberSInt16Type = 2
kCFNumberSInt32Type = 3
kCFNumberSInt64Type = 4

def get_string_property(device_type, property):
    """
    Search the given device for the specified string property

    @param device_type Type of Device
    @param property String to search for
    @return Python string containing the value, or None if not found.
    """
    key = cf.CFStringCreateWithCString(
            kCFAllocatorDefault,
            property.encode("mac_roman"),
            kCFStringEncodingMacRoman)

    CFContainer = iokit.IORegistryEntryCreateCFProperty(
            device_type,
            key,
            kCFAllocatorDefault,
            0)
    output = None

    if CFContainer:
        output = cf.CFStringGetCStringPtr(CFContainer, 0).decode('mac_roman')
        cf.CFRelease(CFContainer)
    return output

def get_int_property(device_type, property, cf_number_type):
    """
    Search the given device for the specified string property

    @param device_type Device to search
    @param property String to search for
    @param cf_number_type CFType number

    @return Python string containing the value, or None if not found.
    """
    key = cf.CFStringCreateWithCString(
            kCFAllocatorDefault,
            property.encode("mac_roman"),
            kCFStringEncodingMacRoman)

    CFContainer = iokit.IORegistryEntryCreateCFProperty(
            device_type,
            key,
            kCFAllocatorDefault,
            0)

    if CFContainer:
        if (cf_number_type == kCFNumberSInt32Type):
            number = ctypes.c_uint32()
        elif (cf_number_type == kCFNumberSInt16Type):
            number = ctypes.c_uint16()
        cf.CFNumberGetValue(CFContainer, cf_number_type, ctypes.byref(number))
        cf.CFRelease(CFContainer)
        return number.value
    return None

def IORegistryEntryGetName(device):
    pathname = ctypes.create_string_buffer(100)  # TODO: Is this ok?
    iokit.IOObjectGetClass(device, ctypes.byref(pathname))
    return pathname.value

def GetParentDeviceByType(device, parent_type):
    """ Find the first parent of a device that implements the parent_type
        @param IOService Service to inspect
        @return Pointer to the parent type, or None if it was not found.
    """
    # First, try to walk up the IOService tree to find a parent of this device that is a IOUSBDevice.
    parent_type = parent_type.encode('mac_roman')
    while IORegistryEntryGetName(device) != parent_type:
        parent = ctypes.c_void_p()
        response = iokit.IORegistryEntryGetParentEntry(
                device,
                "IOService".encode("mac_roman"),
                ctypes.byref(parent))
        # If we weren't able to find a parent for the device, we're done.
        if response != 0:
            return None
        device = parent
    return device

def GetIOServicesByType(service_type):
    """
    returns iterator over specified service_type
    """
    serial_port_iterator = ctypes.c_void_p()

    iokit.IOServiceGetMatchingServices(
            kIOMasterPortDefault,
            iokit.IOServiceMatching(service_type.encode('mac_roman')),
            ctypes.byref(serial_port_iterator))

    services = []
    while iokit.IOIteratorIsValid(serial_port_iterator):
        service = iokit.IOIteratorNext(serial_port_iterator)
        if not service:
            break
        services.append(service)
    iokit.IOObjectRelease(serial_port_iterator)
    return services

def location_to_string(location_number):
    loc = ['{}-'.format(location_number >> 24)]
    while location_number & 0xf00000:
        if len(loc) > 1:
            loc.append('.')
        loc.append('{}'.format((location_number >> 20) & 0xf))
        location_number <<= 4
    return ''.join(loc)

def comports():
    # Scan for all iokit serial ports
    services = GetIOServicesByType('IOSerialBSDClient')
    ports = []
    for service in services:
        # First, add the callout device file.
        device = get_string_property(service, "IOCalloutDevice")
        if device:
            info = list_ports_common.ListPortInfo(device)
            # If the serial port is implemented by IOUSBDevice
            usb_device = GetParentDeviceByType(service, "IOUSBDevice")
            if usb_device is not None:
                # fetch some useful informations from properties
                info.vid = get_int_property(usb_device, "idVendor", kCFNumberSInt16Type)
                info.pid = get_int_property(usb_device, "idProduct", kCFNumberSInt16Type)
                info.serial_number = get_int_property(usb_device, "iSerialNumber", kCFNumberSInt32Type)
                info.product = get_string_property(usb_device, "USB Product Name") or 'n/a'
                info.manufacturer = get_string_property(usb_device, "USB Vendor Name")
                locationID = get_int_property(usb_device, "locationID", kCFNumberSInt32Type)
                info.location = location_to_string(locationID)
                #~ bcd = get_int_property(usb_device, "bcdDevice", kCFNumberSInt16Type)
                info.hwid = info.usb_info()
            ports.append(info)
    return ports

# test
if __name__ == '__main__':
    for port, desc, hwid in sorted(comports()):
        print("%s: %s [%s]" % (port, desc, hwid))
<commit_msg>OSX: list_ports - add interface string if present<commit_after>#!/usr/bin/env python
#
# portable serial port access with python
#
# This is a module that gathers a list of serial ports including details on OSX
#
# code originally from https://github.com/makerbot/pyserial/tree/master/serial/tools
# with contributions from cibomahto, dgs3, FarMcKon, tedbrandston
# and modifications by cliechti, hoihu, hardkrash
#
# (C) 2013-2015
#
# SPDX-License-Identifier:    BSD-3-Clause

# List all of the callout devices in OS/X by querying IOKit.

# See the following for a reference of how to do this:
# http://developer.apple.com/library/mac/#documentation/DeviceDrivers/Conceptual/WorkingWSerial/WWSerial_SerialDevs/SerialDevices.html#//apple_ref/doc/uid/TP30000384-CIHGEAFD

# More help from darwin_hid.py

# Also see the 'IORegistryExplorer' for an idea of what we are actually searching

import ctypes
from ctypes import util

from serial.tools import list_ports_common

iokit = ctypes.cdll.LoadLibrary(ctypes.util.find_library('IOKit'))
cf = ctypes.cdll.LoadLibrary(ctypes.util.find_library('CoreFoundation'))

kIOMasterPortDefault = ctypes.c_void_p.in_dll(iokit, "kIOMasterPortDefault")
kCFAllocatorDefault = ctypes.c_void_p.in_dll(cf, "kCFAllocatorDefault")

kCFStringEncodingMacRoman = 0

iokit.IOServiceMatching.restype = ctypes.c_void_p

iokit.IOServiceGetMatchingServices.argtypes = [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p]
iokit.IOServiceGetMatchingServices.restype = ctypes.c_void_p

iokit.IORegistryEntryGetParentEntry.argtypes = [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p]

iokit.IORegistryEntryCreateCFProperty.argtypes = [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p, ctypes.c_uint32]
iokit.IORegistryEntryCreateCFProperty.restype = ctypes.c_void_p

iokit.IORegistryEntryGetPath.argtypes = [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p]
iokit.IORegistryEntryGetPath.restype = ctypes.c_void_p

iokit.IORegistryEntryGetName.argtypes = [ctypes.c_void_p, ctypes.c_void_p]
iokit.IORegistryEntryGetName.restype = ctypes.c_void_p

iokit.IOObjectGetClass.argtypes = [ctypes.c_void_p, ctypes.c_void_p]
iokit.IOObjectGetClass.restype = ctypes.c_void_p

iokit.IOObjectRelease.argtypes = [ctypes.c_void_p]

cf.CFStringCreateWithCString.argtypes = [ctypes.c_void_p, ctypes.c_char_p, ctypes.c_int32]
cf.CFStringCreateWithCString.restype = ctypes.c_void_p

cf.CFStringGetCStringPtr.argtypes = [ctypes.c_void_p, ctypes.c_uint32]
cf.CFStringGetCStringPtr.restype = ctypes.c_char_p

cf.CFNumberGetValue.argtypes = [ctypes.c_void_p, ctypes.c_uint32, ctypes.c_void_p]
cf.CFNumberGetValue.restype = ctypes.c_void_p

# void CFRelease ( CFTypeRef cf );
cf.CFRelease.argtypes = [ctypes.c_void_p]
cf.CFRelease.restype = None

# CFNumber type defines
kCFNumberSInt8Type = 1
kCFNumberSInt16Type = 2
kCFNumberSInt32Type = 3
kCFNumberSInt64Type = 4

def get_string_property(device_type, property):
    """
    Search the given device for the specified string property

    @param device_type Type of Device
    @param property String to search for
    @return Python string containing the value, or None if not found.
    """
    key = cf.CFStringCreateWithCString(
            kCFAllocatorDefault,
            property.encode("mac_roman"),
            kCFStringEncodingMacRoman)

    CFContainer = iokit.IORegistryEntryCreateCFProperty(
            device_type,
            key,
            kCFAllocatorDefault,
            0)
    output = None

    if CFContainer:
        output = cf.CFStringGetCStringPtr(CFContainer, 0).decode('mac_roman')
        cf.CFRelease(CFContainer)
    return output

def get_int_property(device_type, property, cf_number_type):
    """
    Search the given device for the specified string property

    @param device_type Device to search
    @param property String to search for
    @param cf_number_type CFType number

    @return Python string containing the value, or None if not found.
    """
    key = cf.CFStringCreateWithCString(
            kCFAllocatorDefault,
            property.encode("mac_roman"),
            kCFStringEncodingMacRoman)

    CFContainer = iokit.IORegistryEntryCreateCFProperty(
            device_type,
            key,
            kCFAllocatorDefault,
            0)

    if CFContainer:
        if (cf_number_type == kCFNumberSInt32Type):
            number = ctypes.c_uint32()
        elif (cf_number_type == kCFNumberSInt16Type):
            number = ctypes.c_uint16()
        cf.CFNumberGetValue(CFContainer, cf_number_type, ctypes.byref(number))
        cf.CFRelease(CFContainer)
        return number.value
    return None

def IORegistryEntryGetName(device):
    pathname = ctypes.create_string_buffer(100)  # TODO: Is this ok?
    iokit.IOObjectGetClass(device, ctypes.byref(pathname))
    return pathname.value

def GetParentDeviceByType(device, parent_type):
    """ Find the first parent of a device that implements the parent_type
        @param IOService Service to inspect
        @return Pointer to the parent type, or None if it was not found.
    """
    # First, try to walk up the IOService tree to find a parent of this device that is a IOUSBDevice.
    parent_type = parent_type.encode('mac_roman')
    while IORegistryEntryGetName(device) != parent_type:
        parent = ctypes.c_void_p()
        response = iokit.IORegistryEntryGetParentEntry(
                device,
                "IOService".encode("mac_roman"),
                ctypes.byref(parent))
        # If we weren't able to find a parent for the device, we're done.
        if response != 0:
            return None
        device = parent
    return device

def GetIOServicesByType(service_type):
    """
    returns iterator over specified service_type
    """
    serial_port_iterator = ctypes.c_void_p()

    iokit.IOServiceGetMatchingServices(
            kIOMasterPortDefault,
            iokit.IOServiceMatching(service_type.encode('mac_roman')),
            ctypes.byref(serial_port_iterator))

    services = []
    while iokit.IOIteratorIsValid(serial_port_iterator):
        service = iokit.IOIteratorNext(serial_port_iterator)
        if not service:
            break
        services.append(service)
    iokit.IOObjectRelease(serial_port_iterator)
    return services

def location_to_string(locationID):
    """
    helper to calculate port and bus number from locationID
    """
    loc = ['{}-'.format(locationID >> 24)]
    while locationID & 0xf00000:
        if len(loc) > 1:
            loc.append('.')
        loc.append('{}'.format((locationID >> 20) & 0xf))
        locationID <<= 4
    return ''.join(loc)

class SuitableSerialInterface(object):
    pass

def scan_interfaces():
    """
    helper function to scan USB interfaces
    returns a list of SuitableSerialInterface objects with name and id attributes
    """
    interfaces = []
    for service in GetIOServicesByType('IOSerialBSDClient'):
        device = get_string_property(service, "IOCalloutDevice")
        if device:
            usb_device = GetParentDeviceByType(service, "IOUSBInterface")
            if usb_device:
                name = get_string_property(usb_device, "USB Interface Name") or None
                locationID = get_int_property(usb_device, "locationID",kCFNumberSInt32Type) or ''
                i = SuitableSerialInterface()
                i.id = locationID
                i.name = name
                interfaces.append(i)
    return interfaces

def search_for_locationID_in_interfaces(serial_interfaces, locationID):
    for interface in serial_interfaces:
        if (interface.id == locationID):
            return interface.name
    return None

def comports():
    # Scan for all iokit serial ports
    services = GetIOServicesByType('IOSerialBSDClient')
    ports = []
    serial_interfaces = scan_interfaces()
    for service in services:
        # First, add the callout device file.
        device = get_string_property(service, "IOCalloutDevice")
        if device:
            info = list_ports_common.ListPortInfo(device)
            # If the serial port is implemented by IOUSBDevice
            usb_device = GetParentDeviceByType(service, "IOUSBDevice")
            if usb_device:
                # fetch some useful informations from properties
                info.vid = get_int_property(usb_device, "idVendor", kCFNumberSInt16Type)
                info.pid = get_int_property(usb_device, "idProduct", kCFNumberSInt16Type)
                info.serial_number = get_int_property(usb_device, "iSerialNumber", kCFNumberSInt32Type)
                info.product = get_string_property(usb_device, "USB Product Name") or 'n/a'
                info.manufacturer = get_string_property(usb_device, "USB Vendor Name")
                locationID = get_int_property(usb_device, "locationID", kCFNumberSInt32Type)
                info.location = location_to_string(locationID)
                info.hwid = info.usb_info()
                info.interface = search_for_locationID_in_interfaces(serial_interfaces, locationID)
            ports.append(info)
    return ports

# test
if __name__ == '__main__':
    for port, desc, hwid in sorted(comports()):
        print("%s: %s [%s]" % (port, desc, hwid))
<|endoftext|>


<commit_before><commit_msg>* oops prototype file deleted<commit_after><|endoftext|>


<commit_before>#############################################
### FOR ALL ATTENDEES
#############################################

def index():
    redirect(URL("attendees"))

#@cache(request.env.path_info,time_expire=60*5,cache_model=cache.ram)
def companies():
    if auth.has_membership(role='manager'): s=db()
    else: s=db(db.auth_user.include_in_delegate_listing==True)
    rows=s.select(db.auth_user.company_name,
                  db.auth_user.company_home_page,
                  orderby=db.auth_user.company_name,distinct=True)
    d = dict(rows=rows)
    return response.render(d)

@cache(request.env.path_info,time_expire=60*5,cache_model=cache.ram)
def attendees():
    if auth.has_membership(role='manager'): s=db(db.auth_user.attendee_type!='non_attending')
    else: s=db((db.auth_user.include_in_delegate_listing==True)&(db.auth_user.attendee_type!='non_attending')&(db.auth_user.amount_due==0.0))
    rows=s.select(db.auth_user.ALL,
                  orderby=db.auth_user.first_name|db.auth_user.last_name)
    d = dict(rows=rows)
    return response.render(d)

colors=['#ff0000','#ff0033','#ff0066','#ff0099','#ff00cc','#ff00ff',
        '#996600','#996633','#996666','#996699','#9966cc','#9966ff',
        '#669900','#669933','#669966','#669999','#6699cc','#cc99ff',
        '#33cc00','#33cc33','#33cc66','#33cc99','#33cccc','#33ccff',
        '#00ff00','#00ff33','#00ff66','#00ff99','#00ffcc','#00ffff',
        '#996600','#996633','#996666','#996699','#9966cc','#9966ff',
        '#669900','#669933','#669966','#669999','#6699cc','#cc99ff',
        '#33cc00','#33cc33','#33cc66','#33cc99','#33cccc','#33ccff',
        '#00ff00','#00ff33','#00ff66','#00ff99','#00ffcc','#00ffff',
        ] + ['#000000']*100

def barchart(data,width=400,height=15,scale=None,
             label_width=50,values_width=50):
    if not scale:
        try:
            scale=max([m for n,c,m in data])
        except ValueError:
            # empty sequence
            scale=None
    if not scale: return None
    return TABLE(_class='barchart',
           *[TR(TD(n,_width=label_width,_style="text-align: right"),
           TABLE(TR(TD(_width=int(m*width/scale),_height=height,
           _style='background-color:'+c))),TD(m,_width=values_width),
           _style="vertical-alignment: middle") for n,c,m in data])

def colorize(d,sort_key=lambda x:x):
    s=[(m,n) for n,m in d.items()]
    s.sort(key=sort_key)
    s.reverse()

    t=[(x[1],colors[i % len(colors)],x[0]) for i,x in enumerate(s)]
    return barchart(t,label_width=150)

@cache(request.env.path_info,time_expire=60*5,cache_model=cache.ram)
def charts():
    cn=[]
    if auth.has_membership(role='manager'):
        q = db.activity.id>0
    else:
        q = db.activity.status=='accepted'
        q &= db.activity.type=='talk'
    tutorials = [row.title for row in db(q).select()]
    if not is_gae:
        for k,item in enumerate(tutorials):
            m=db(db.auth_user.tutorials.like('%%|%s|%%'%item)).count()
            cn.append((item,colors[k],m))
    else:
        cn2={}
        for row in db(db.auth_user.id>0).select(db.auth_user.tutorials):
            for item in tutorials:
                    if not cn2.has_key(item): cn2[item]=0
                    if row.tutorials.find('|%s|'%item)>=0: cn2[item]+=1
        for k,item in enumerate(tutorials):
                cn.append((TUTORIALS[item],colors[k],cn2[item]))
                k+=1
    cn.sort(key=lambda x: x[-1], reverse=True)

    chart_tutorials=barchart(cn,label_width=150)

    country={}
    city={}
    state={}
    food_preference={}
    t_shirt_size={}
    attendee_type={}
    attendee_levels={}
    registration_date={}
    certificates={}
    for row in db().select(db.auth_user.ALL):
        country[row.country]=country.get(row.country,0)+1
        city[row.city.lower()]=city.get(row.city.lower(),0)+1
        state[row.state.lower()]=state.get(row.state.lower(),0)+1
        #food_preference[row.food_preference]=food_preference.get(row.food_preference,0)+1
        #t_shirt_size[row.t_shirt_size]=t_shirt_size.get(row.t_shirt_size,0)+1
        attendee_type[row.attendee_type]=attendee_type.get(row.attendee_type,0)+1
        attendee_levels[row.level]=attendee_levels.get(row.level,0)+1
        registration_date[row.created_on.date()]=registration_date.get(row.created_on.date(),0)+1
        certificates[row.certificate]=certificates.get(row.certificate,0)+1
    chart_country=colorize(country)
    chart_state=colorize(state)
    chart_city=colorize(city)
    chart_food_preference=None #colorize(food_preference)
    chart_t_shirt_size=None #colorize(t_shirt_size)
    chart_attendee_type=colorize(attendee_levels)
    chart_registration_date=colorize(registration_date,sort_key=lambda x: x[1]) #colorize(attendee_type)
    chart_certificates=colorize(certificates) #colorize(attendee_type)

    d = dict(chart_tutorials=chart_tutorials,
                chart_country=chart_country,
                chart_food_preference=chart_food_preference,
                chart_t_shirt_size=chart_t_shirt_size,
                chart_city=chart_city,
                chart_state=chart_state,
                chart_attendee_type=chart_attendee_type,
                chart_registration_date=chart_registration_date,
                chart_certificates=chart_certificates
                )
    return response.render(d)

@cache(request.env.path_info,time_expire=60*5,cache_model=cache.ram)
def brief():
    cn=[]

    activity_status={}
    activity_types={}
    activity_levels={}
    activity_categories={}
    for row in db().select(db.activity.ALL):
        status = "%s%s" % (row.status, row.confirmed and " (confirmed)" or "")
        activity_status[status]=activity_status.get(status,0)+1
        if not request.args or request.args[0] == row.status:
            activity_types[row.type]=activity_types.get(row.type,0)+1
            activity_levels[row.level]=activity_levels.get(row.level,0)+1
            for category in row.categories:
                activity_categories[category]=activity_categories.get(category,0)+1
    chart_status=colorize(activity_status)
    chart_type=colorize(activity_types)
    chart_level=colorize(activity_levels)
    chart_categories=colorize(activity_categories)

    authors_country={}
    q = db.auth_user.id==db.author.user_id
    q &=db.activity.id==db.author.activity_id
    q &= db.auth_user.country != None
    if request.args and request.args[0] in ("accepted", "rejected", "pending", "declined"):
        q &=db.activity.status==request.args[0]
    for row in db(q).select(db.auth_user.id, db.auth_user.country, groupby=db.auth_user.id):
        authors_country[row.country]=authors_country.get(row.country,0)+1
    chart_country=colorize(authors_country)

    d = dict(chart_status=chart_status,
                chart_type=chart_type,
                chart_level=chart_level,
                chart_categories=chart_categories,
                chart_country=chart_country,
                )
    return response.render(d)

##@cache(request.env.path_info,time_expire=60*5,cache_model=cache.ram)
def maps():
    rows=db(db.auth_user.id>0).select(
            db.auth_user.first_name,
            db.auth_user.last_name,
            db.auth_user.latitude,
            db.auth_user.longitude,
            db.auth_user.personal_home_page,
            db.auth_user.include_in_delegate_listing)
    x0,y0=CONFERENCE_COORDS
    d = dict(googlemap_key=GOOGLEMAP_KEY,x0=x0,y0=y0,rows=rows)
    return response.render(d)
<commit_msg>fixed attendee stat page (remove dups)<commit_after>#############################################
### FOR ALL ATTENDEES
#############################################

def index():
    redirect(URL("attendees"))

@cache(request.env.path_info,time_expire=60*5,cache_model=cache.ram)
def companies():
    if auth.has_membership(role='manager'): s=db()
    else: s=db(db.auth_user.include_in_delegate_listing==True)
    rows=s.select(db.auth_user.company_name,
                  db.auth_user.company_home_page,
                  orderby=db.auth_user.company_name,distinct=True)
    d = dict(rows=rows)
    return response.render(d)

@cache(request.env.path_info,time_expire=60*5,cache_model=cache.ram)
def attendees():
    import unicodedata
    if auth.has_membership(role='manager'): s=db(db.auth_user.attendee_type!='non_attending')
    else: s=db((db.auth_user.include_in_delegate_listing==True)&(db.auth_user.attendee_type!='non_attending')&(db.auth_user.amount_due==0.0))
    rows=s.select(db.auth_user.ALL,
                  orderby=db.auth_user.first_name|db.auth_user.last_name)
    ret = {}
    # avoid duplicates (remove spaces, acents, etc)
    for row in rows:
        name = ''.join([l.lower() for l in (row.first_name+row.last_name) if l.isalpha()])
        name = unicodedata.normalize('NFD', unicode(name, 'utf8'))
        if name in ret: continue
        ret[name] = row
    d = dict(rows=sorted(ret.values()))
    return response.render(d)

colors=['#ff0000','#ff0033','#ff0066','#ff0099','#ff00cc','#ff00ff',
        '#996600','#996633','#996666','#996699','#9966cc','#9966ff',
        '#669900','#669933','#669966','#669999','#6699cc','#cc99ff',
        '#33cc00','#33cc33','#33cc66','#33cc99','#33cccc','#33ccff',
        '#00ff00','#00ff33','#00ff66','#00ff99','#00ffcc','#00ffff',
        '#996600','#996633','#996666','#996699','#9966cc','#9966ff',
        '#669900','#669933','#669966','#669999','#6699cc','#cc99ff',
        '#33cc00','#33cc33','#33cc66','#33cc99','#33cccc','#33ccff',
        '#00ff00','#00ff33','#00ff66','#00ff99','#00ffcc','#00ffff',
        ] + ['#000000']*100

def barchart(data,width=400,height=15,scale=None,
             label_width=50,values_width=50):
    if not scale:
        try:
            scale=max([m for n,c,m in data])
        except ValueError:
            # empty sequence
            scale=None
    if not scale: return None
    return TABLE(_class='barchart',
           *[TR(TD(n,_width=label_width,_style="text-align: right"),
           TABLE(TR(TD(_width=int(m*width/scale),_height=height,
           _style='background-color:'+c))),TD(m,_width=values_width),
           _style="vertical-alignment: middle") for n,c,m in data])

def colorize(d,sort_key=lambda x:x):
    s=[(m,n) for n,m in d.items()]
    s.sort(key=sort_key)
    s.reverse()

    t=[(x[1],colors[i % len(colors)],x[0]) for i,x in enumerate(s)]
    return barchart(t,label_width=150)

@cache(request.env.path_info,time_expire=60*5,cache_model=cache.ram)
def charts():
    cn=[]
    if auth.has_membership(role='manager'):
        q = db.activity.id>0
    else:
        q = db.activity.status=='accepted'
        q &= db.activity.type=='talk'
    tutorials = [row.title for row in db(q).select()]
    if not is_gae:
        for k,item in enumerate(tutorials):
            m=db(db.auth_user.tutorials.like('%%|%s|%%'%item)).count()
            cn.append((item,colors[k],m))
    else:
        cn2={}
        for row in db(db.auth_user.id>0).select(db.auth_user.tutorials):
            for item in tutorials:
                    if not cn2.has_key(item): cn2[item]=0
                    if row.tutorials.find('|%s|'%item)>=0: cn2[item]+=1
        for k,item in enumerate(tutorials):
                cn.append((TUTORIALS[item],colors[k],cn2[item]))
                k+=1
    cn.sort(key=lambda x: x[-1], reverse=True)

    chart_tutorials=barchart(cn,label_width=150)

    country={}
    city={}
    state={}
    food_preference={}
    t_shirt_size={}
    attendee_type={}
    attendee_levels={}
    registration_date={}
    certificates={}
    for row in db().select(db.auth_user.ALL):
        country[row.country]=country.get(row.country,0)+1
        city[row.city.lower()]=city.get(row.city.lower(),0)+1
        state[row.state.lower()]=state.get(row.state.lower(),0)+1
        #food_preference[row.food_preference]=food_preference.get(row.food_preference,0)+1
        #t_shirt_size[row.t_shirt_size]=t_shirt_size.get(row.t_shirt_size,0)+1
        attendee_type[row.attendee_type]=attendee_type.get(row.attendee_type,0)+1
        attendee_levels[row.level]=attendee_levels.get(row.level,0)+1
        registration_date[row.created_on.date()]=registration_date.get(row.created_on.date(),0)+1
        certificates[row.certificate]=certificates.get(row.certificate,0)+1
    chart_country=colorize(country)
    chart_state=colorize(state)
    chart_city=colorize(city)
    chart_food_preference=None #colorize(food_preference)
    chart_t_shirt_size=None #colorize(t_shirt_size)
    chart_attendee_type=colorize(attendee_levels)
    chart_registration_date=colorize(registration_date,sort_key=lambda x: x[1]) #colorize(attendee_type)
    chart_certificates=colorize(certificates) #colorize(attendee_type)

    d = dict(chart_tutorials=chart_tutorials,
                chart_country=chart_country,
                chart_food_preference=chart_food_preference,
                chart_t_shirt_size=chart_t_shirt_size,
                chart_city=chart_city,
                chart_state=chart_state,
                chart_attendee_type=chart_attendee_type,
                chart_registration_date=chart_registration_date,
                chart_certificates=chart_certificates
                )
    return response.render(d)

@cache(request.env.path_info,time_expire=60*5,cache_model=cache.ram)
def brief():
    cn=[]

    activity_status={}
    activity_types={}
    activity_levels={}
    activity_categories={}
    for row in db().select(db.activity.ALL):
        status = "%s%s" % (row.status, row.confirmed and " (confirmed)" or "")
        activity_status[status]=activity_status.get(status,0)+1
        if not request.args or request.args[0] == row.status:
            activity_types[row.type]=activity_types.get(row.type,0)+1
            activity_levels[row.level]=activity_levels.get(row.level,0)+1
            for category in row.categories:
                activity_categories[category]=activity_categories.get(category,0)+1
    chart_status=colorize(activity_status)
    chart_type=colorize(activity_types)
    chart_level=colorize(activity_levels)
    chart_categories=colorize(activity_categories)

    authors_country={}
    q = db.auth_user.id==db.author.user_id
    q &=db.activity.id==db.author.activity_id
    q &= db.auth_user.country != None
    if request.args and request.args[0] in ("accepted", "rejected", "pending", "declined"):
        q &=db.activity.status==request.args[0]
    for row in db(q).select(db.auth_user.id, db.auth_user.country, groupby=db.auth_user.id):
        authors_country[row.country]=authors_country.get(row.country,0)+1
    chart_country=colorize(authors_country)

    d = dict(chart_status=chart_status,
                chart_type=chart_type,
                chart_level=chart_level,
                chart_categories=chart_categories,
                chart_country=chart_country,
                )
    return response.render(d)

@cache(request.env.path_info,time_expire=60*5,cache_model=cache.ram)
def maps():
    rows=db(db.auth_user.id>0).select(
            db.auth_user.first_name,
            db.auth_user.last_name,
            db.auth_user.latitude,
            db.auth_user.longitude,
            db.auth_user.personal_home_page,
            db.auth_user.include_in_delegate_listing)
    x0,y0=CONFERENCE_COORDS
    d = dict(googlemap_key=GOOGLEMAP_KEY,x0=x0,y0=y0,rows=rows)
    return response.render(d)
<|endoftext|>


<commit_before># Copyright 2016 Twitter. All rights reserved.
#
# 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.

#!/usr/bin/python2.7

import argparse
import atexit
import base64
import contextlib
import glob
import logging
import logging.handlers
import os
import shutil
import sys
import subprocess
import tarfile
import tempfile

from heron.common.src.python.color import Log
from heron.proto import topology_pb2

import heron.cli.src.python.args as args
import heron.cli.src.python.execute as execute
import heron.cli.src.python.jars as jars
import heron.cli.src.python.opts as opts
import heron.cli.src.python.utils as utils

################################################################################
# Create a subparser for the submit command
################################################################################
def create_parser(subparsers):
  parser = subparsers.add_parser(
      'submit',
      help='Submit a topology',
      usage = "%(prog)s [options] cluster/[role]/[environ] " + \
              "topology-file-name topology-class-name [topology-args]",
      add_help = False
  )

  args.add_titles(parser)
  args.add_cluster_role_env(parser)
  args.add_topology_file(parser)
  args.add_topology_class(parser)
  args.add_config(parser)

  parser.add_argument(
      '--deploy-deactivated',
      metavar='(a boolean; default: "false")',
      default=False)

  parser.add_argument(
      '-D',
      default=[],
      action="append",
      dest="javaDefines",
      metavar='DEFINE',
      help='Define a system property to pass to java -D when running main.')

  args.add_verbose(parser)

  parser.set_defaults(subcommand='submit')
  return parser

################################################################################
# Launch a topology given topology jar, its definition file and configurations
################################################################################
def launch_a_topology(cl_args, tmp_dir, topology_file, topology_defn_file):

  # get the normalized path for topology.tar.gz
  topology_pkg_path = utils.normalized_class_path(os.path.join(tmp_dir, 'topology.tar.gz'))

  # TO DO - when you give a config path - the name of the directory might be
  # different - need to change the name to conf

  # create a tar package with the configuration
  config_path = cl_args['config_path']
  tar_pkg_files = [topology_file, topology_defn_file]
  utils.create_tar(topology_pkg_path, tar_pkg_files, config_path)

  # form the config overrides
  config_overrides = utils.parse_cmdline_override(cl_args)

  # pass the args to submitter main
  args = [
      "--cluster", cl_args['cluster'],
      "--role", cl_args['role'],
      "--environment", cl_args['environ'],
      "--heron_home", utils.get_heron_dir(),
      "--config_path", config_path,
      "--config_overrides", base64.b64encode(config_overrides),
      "--topology_package", topology_pkg_path,
      "--topology_defn", topology_defn_file,
      "--topology_jar", topology_file
  ]

  if opts.verbose():
    args.append("--verbose")

  lib_jars = utils.get_heron_libs(
      jars.scheduler_jars() + jars.uploader_jars() + jars.statemgr_jars() + jars.packing_jars()
  )

  # invoke the submitter to submit and launch the topology
  execute.heron_class(
      'com.twitter.heron.scheduler.SubmitterMain',
      lib_jars,
      extra_jars=[],
      args = args,
      javaDefines = cl_args['javaDefines']
  )

################################################################################
# Launch topologies
################################################################################
def launch_topologies(cl_args, topology_file, tmp_dir):

  # the submitter would have written the .defn file to the tmp_dir
  defn_files = glob.glob(tmp_dir + '/*.defn')

  if len(defn_files) == 0:
    raise Exception("No topologies found")

  try:
    for defn_file in defn_files:

      # load the topology definition from the file
      topology_defn = topology_pb2.Topology()
      try:
        f = open(defn_file, "rb")
        topology_defn.ParseFromString(f.read())
        f.close()

      except:
        raise Exception("Could not open and parse topology defn file %s" % defn_file)

      # launch the topology
      try:
        Log.info("Launching topology \'%s\'" % topology_defn.name)
        launch_a_topology(cl_args, tmp_dir, topology_file, defn_file)
        Log.info("Topology \'%s\' launched successfully" % topology_defn.name)

      except Exception as ex:
        Log.error('Failed to launch topology \'%s\' because %s' % (topology_defn.name, str(ex)))
        raise

  except:
    raise

################################################################################
# We use the packer to make a package for the jar and dump it
# to a well-known location. We then run the main method of class
# with the specified arguments. We pass arguments as heron.options.
#
# This will run the jar file with the topology_class_name. The submitter
# inside will write out the topology defn file to a location that
# we specify. Then we write the topology defn file to a well known
# location. We then write to appropriate places in zookeeper
# and launch the scheduler jobs
################################################################################
def submit_fatjar(cl_args, unknown_args, tmp_dir):

  # execute main of the topology to create the topology definition
  topology_file = cl_args['topology-file-name']
  execute.heron_class(
      cl_args['topology-class-name'],
      utils.get_heron_libs(jars.topology_jars()),
      extra_jars = [topology_file],
      args = tuple(unknown_args),
      javaDefines = cl_args['javaDefines'])

  try:
    launch_topologies(cl_args, topology_file, tmp_dir)

  except Exception as ex:
    return False

  finally:
    shutil.rmtree(tmp_dir)

  return True

################################################################################
# Extract and execute the java files inside the tar and then add topology
# definition file created by running submitTopology
#
# We use the packer to make a package for the tar and dump it
# to a well-known location. We then run the main method of class
# with the specified arguments. We pass arguments as heron.options.
# This will run the jar file with the topology class name.
#
# The submitter inside will write out the topology defn file to a location
# that we specify. Then we write the topology defn file to a well known
# packer location. We then write to appropriate places in zookeeper
# and launch the aurora jobs
################################################################################
def submit_tar(cl_args, unknown_args, tmp_dir):

  # execute main of the topology to create the topology definition
  topology_file = cl_args['topology-file-name']
  execute.heron_tar(
      cl_args['topology-class-name'],
      topology_file,
      tuple(unknown_args),
      tmp_dir,
      cl_args['javaDefines'])

  try:
    launch_topologies(cl_args, topology_file, tmp_dir)

  except Exception as ex:
    return False

  finally:
    shutil.rmtree(tmp_dir)

  return True

################################################################################
#  Submits the topology to the scheduler
#  * Depending on the topology file name extension, we treat the file as a
#    fatjar (if the ext is .jar) or a tar file (if the ext is .tar/.tar.gz).
#  * We upload the topology file to the packer, update zookeeper and launch
#    scheduler jobs representing that topology
#  * You can see your topology in Heron UI
################################################################################
def run(command, parser, cl_args, unknown_args):

  # get the topology file name
  topology_file = cl_args['topology-file-name']

  # check to see if the topology file exists
  if not os.path.isfile(topology_file):
    Log.error("Topology jar|tar file %s does not exist" % topology_file)
    return False

  # check if it is a valid file type
  jar_type = topology_file.endswith(".jar")
  tar_type = topology_file.endswith(".tar") or topology_file.endswith(".tar.gz")
  if not jar_type and not tar_type:
    Log.error("Unknown file type. Please use .tar or .tar.gz or .jar file")
    return False

  # create a temporary directory for topology definition file
  tmp_dir = tempfile.mkdtemp()

  # if topology needs to be launched in deactivated state, do it so
  if cl_args['deploy_deactivated']:
    initial_state = topology_pb2.TopologyState.Name(topology_pb2.PAUSED)
  else:
    initial_state = topology_pb2.TopologyState.Name(topology_pb2.RUNNING)

  # set the tmp dir and deactivated state in global options
  opts.set_config('cmdline.topologydefn.tmpdirectory', tmp_dir)
  opts.set_config('cmdline.topology.initial.state', initial_state)

  # check the extension of the file name to see if it is tar/jar file.
  if jar_type:
    return submit_fatjar(cl_args, unknown_args, tmp_dir)

  elif tar_type:
    return submit_tar(cl_args, unknown_args, tmp_dir)

  return False
<commit_msg>catch exception when unable to run the topology main and print an error<commit_after># Copyright 2016 Twitter. All rights reserved.
#
# 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.

#!/usr/bin/python2.7

import argparse
import atexit
import base64
import contextlib
import glob
import logging
import logging.handlers
import os
import shutil
import sys
import subprocess
import tarfile
import tempfile

from heron.common.src.python.color import Log
from heron.proto import topology_pb2

import heron.cli.src.python.args as args
import heron.cli.src.python.execute as execute
import heron.cli.src.python.jars as jars
import heron.cli.src.python.opts as opts
import heron.cli.src.python.utils as utils

################################################################################
# Create a subparser for the submit command
################################################################################
def create_parser(subparsers):
  parser = subparsers.add_parser(
      'submit',
      help='Submit a topology',
      usage = "%(prog)s [options] cluster/[role]/[environ] " + \
              "topology-file-name topology-class-name [topology-args]",
      add_help = False
  )

  args.add_titles(parser)
  args.add_cluster_role_env(parser)
  args.add_topology_file(parser)
  args.add_topology_class(parser)
  args.add_config(parser)

  parser.add_argument(
      '--deploy-deactivated',
      metavar='(a boolean; default: "false")',
      default=False)

  parser.add_argument(
      '-D',
      default=[],
      action="append",
      dest="javaDefines",
      metavar='DEFINE',
      help='Define a system property to pass to java -D when running main.')

  args.add_verbose(parser)

  parser.set_defaults(subcommand='submit')
  return parser

################################################################################
# Launch a topology given topology jar, its definition file and configurations
################################################################################
def launch_a_topology(cl_args, tmp_dir, topology_file, topology_defn_file):

  # get the normalized path for topology.tar.gz
  topology_pkg_path = utils.normalized_class_path(os.path.join(tmp_dir, 'topology.tar.gz'))

  # TO DO - when you give a config path - the name of the directory might be
  # different - need to change the name to conf

  # create a tar package with the configuration
  config_path = cl_args['config_path']
  tar_pkg_files = [topology_file, topology_defn_file]
  utils.create_tar(topology_pkg_path, tar_pkg_files, config_path)

  # form the config overrides
  config_overrides = utils.parse_cmdline_override(cl_args)

  # pass the args to submitter main
  args = [
      "--cluster", cl_args['cluster'],
      "--role", cl_args['role'],
      "--environment", cl_args['environ'],
      "--heron_home", utils.get_heron_dir(),
      "--config_path", config_path,
      "--config_overrides", base64.b64encode(config_overrides),
      "--topology_package", topology_pkg_path,
      "--topology_defn", topology_defn_file,
      "--topology_jar", topology_file
  ]

  if opts.verbose():
    args.append("--verbose")

  lib_jars = utils.get_heron_libs(
      jars.scheduler_jars() + jars.uploader_jars() + jars.statemgr_jars() + jars.packing_jars()
  )

  # invoke the submitter to submit and launch the topology
  execute.heron_class(
      'com.twitter.heron.scheduler.SubmitterMain',
      lib_jars,
      extra_jars=[],
      args = args,
      javaDefines = cl_args['javaDefines']
  )

################################################################################
# Launch topologies
################################################################################
def launch_topologies(cl_args, topology_file, tmp_dir):

  # the submitter would have written the .defn file to the tmp_dir
  defn_files = glob.glob(tmp_dir + '/*.defn')

  if len(defn_files) == 0:
    raise Exception("No topologies found")

  try:
    for defn_file in defn_files:

      # load the topology definition from the file
      topology_defn = topology_pb2.Topology()
      try:
        f = open(defn_file, "rb")
        topology_defn.ParseFromString(f.read())
        f.close()

      except:
        raise Exception("Could not open and parse topology defn file %s" % defn_file)

      # launch the topology
      try:
        Log.info("Launching topology \'%s\'" % topology_defn.name)
        launch_a_topology(cl_args, tmp_dir, topology_file, defn_file)
        Log.info("Topology \'%s\' launched successfully" % topology_defn.name)

      except Exception as ex:
        Log.error('Failed to launch topology \'%s\' because %s' % (topology_defn.name, str(ex)))
        raise

  except:
    raise

################################################################################
# We use the packer to make a package for the jar and dump it
# to a well-known location. We then run the main method of class
# with the specified arguments. We pass arguments as heron.options.
#
# This will run the jar file with the topology_class_name. The submitter
# inside will write out the topology defn file to a location that
# we specify. Then we write the topology defn file to a well known
# location. We then write to appropriate places in zookeeper
# and launch the scheduler jobs
################################################################################
def submit_fatjar(cl_args, unknown_args, tmp_dir):

  # execute main of the topology to create the topology definition
  topology_file = cl_args['topology-file-name']
  try:
    execute.heron_class(
      cl_args['topology-class-name'],
      utils.get_heron_libs(jars.topology_jars()),
      extra_jars = [topology_file],
      args = tuple(unknown_args),
      javaDefines = cl_args['javaDefines'])

  except Exception as ex:
    Log.error("Unable to execute topology main class")
    return False

  try:
    launch_topologies(cl_args, topology_file, tmp_dir)

  except Exception as ex:
    return False

  finally:
    shutil.rmtree(tmp_dir)

  return True

################################################################################
# Extract and execute the java files inside the tar and then add topology
# definition file created by running submitTopology
#
# We use the packer to make a package for the tar and dump it
# to a well-known location. We then run the main method of class
# with the specified arguments. We pass arguments as heron.options.
# This will run the jar file with the topology class name.
#
# The submitter inside will write out the topology defn file to a location
# that we specify. Then we write the topology defn file to a well known
# packer location. We then write to appropriate places in zookeeper
# and launch the aurora jobs
################################################################################
def submit_tar(cl_args, unknown_args, tmp_dir):

  # execute main of the topology to create the topology definition
  topology_file = cl_args['topology-file-name']
  execute.heron_tar(
      cl_args['topology-class-name'],
      topology_file,
      tuple(unknown_args),
      tmp_dir,
      cl_args['javaDefines'])

  try:
    launch_topologies(cl_args, topology_file, tmp_dir)

  except Exception as ex:
    return False

  finally:
    shutil.rmtree(tmp_dir)

  return True

################################################################################
#  Submits the topology to the scheduler
#  * Depending on the topology file name extension, we treat the file as a
#    fatjar (if the ext is .jar) or a tar file (if the ext is .tar/.tar.gz).
#  * We upload the topology file to the packer, update zookeeper and launch
#    scheduler jobs representing that topology
#  * You can see your topology in Heron UI
################################################################################
def run(command, parser, cl_args, unknown_args):

  # get the topology file name
  topology_file = cl_args['topology-file-name']

  # check to see if the topology file exists
  if not os.path.isfile(topology_file):
    Log.error("Topology jar|tar file %s does not exist" % topology_file)
    return False

  # check if it is a valid file type
  jar_type = topology_file.endswith(".jar")
  tar_type = topology_file.endswith(".tar") or topology_file.endswith(".tar.gz")
  if not jar_type and not tar_type:
    Log.error("Unknown file type. Please use .tar or .tar.gz or .jar file")
    return False

  # create a temporary directory for topology definition file
  tmp_dir = tempfile.mkdtemp()

  # if topology needs to be launched in deactivated state, do it so
  if cl_args['deploy_deactivated']:
    initial_state = topology_pb2.TopologyState.Name(topology_pb2.PAUSED)
  else:
    initial_state = topology_pb2.TopologyState.Name(topology_pb2.RUNNING)

  # set the tmp dir and deactivated state in global options
  opts.set_config('cmdline.topologydefn.tmpdirectory', tmp_dir)
  opts.set_config('cmdline.topology.initial.state', initial_state)

  # check the extension of the file name to see if it is tar/jar file.
  if jar_type:
    return submit_fatjar(cl_args, unknown_args, tmp_dir)

  elif tar_type:
    return submit_tar(cl_args, unknown_args, tmp_dir)

  return False
<|endoftext|>