<commit_before>import mdpopups
import sublime
import sublime_plugin
import webbrowser
import re
import html

try:
    from typing import Any, List, Dict, Optional
    assert Any and List and Dict and Optional
except ImportError:
    pass

from .core.configurations import is_supported_syntax
from .core.registry import config_for_scope, session_for_view, client_for_view
from .core.documents import get_document_position
from .core.events import global_events
from .core.protocol import Request
from .core.logging import debug
from .core.popups import popup_css, popup_class
from .core.settings import settings

class SignatureHelpListener(sublime_plugin.ViewEventListener):

    def __init__(self, view):
        self.view = view
        self._initialized = False
        self._signature_help_triggers = []  # type: List[str]
        self._visible = False
        self._language_id = ""
        self._signatures = []  # type: List[Any]
        self._active_signature = -1
        self._active_parameter = -1

    @classmethod
    def is_applicable(cls, settings):
        syntax = settings.get('syntax')
        return syntax and is_supported_syntax(syntax)

    def initialize(self):
        session = session_for_view(self.view)
        if session:
            signatureHelpProvider = session.get_capability(
                'signatureHelpProvider')
            if signatureHelpProvider:
                self._signature_help_triggers = signatureHelpProvider.get(
                    'triggerCharacters')

        config = config_for_scope(self.view)
        if config:
            self._language_id = self._view_language(self.view, config.name)

        self._initialized = True

    def on_modified_async(self):
        pos = self.view.sel()[0].begin()
        # TODO: this will fire too often, narrow down using scopes or regex
        if not self._initialized:
            self.initialize()

        if self._signature_help_triggers:
            last_char = self.view.substr(pos - 1)
            if last_char in self._signature_help_triggers:
                self.request_signature_help(pos)
            elif self._visible:
                if last_char.isspace():
                    # Peek behind to find the last non-whitespace character.
                    last_char = self.view.substr(self.view.find_by_class(pos, False, ~0) - 1)
                if last_char not in self._signature_help_triggers:
                    self.view.hide_popup()

    def request_signature_help(self, point) -> None:
        client = client_for_view(self.view)
        if client:
            global_events.publish("view.on_purge_changes", self.view)
            document_position = get_document_position(self.view, point)
            if document_position:
                client.send_request(
                    Request.signatureHelp(document_position),
                    lambda response: self.handle_response(response, point))

    def handle_response(self, response: 'Optional[Dict]', point) -> None:
        if response is not None:
            self._signatures = response.get("signatures", [])
            self._active_signature = response.get("activeSignature", -1)
            self._active_parameter = response.get("activeParameter", -1)

            if self._signatures:
                if not 0 <= self._active_signature < len(self._signatures):
                    debug("activeSignature {} not a valid index for signatures length {}".format(
                        self._active_signature, len(self._signatures)))
                    self._active_signature = 0
            else:
                if self._active_signature != -1:
                    debug("activeSignature should be -1 or null when no signatures are returned")
                    self._active_signature = -1

            if len(self._signatures) > 0:
                if self._visible:
                    self._update_popup()
                else:
                    self._show_popup(point)

    def on_query_context(self, key, _, operand, __):
        if key != "lsp.signature_help":
            return False  # Let someone else handle this keybinding.
        elif not self._visible:
            if operand == 0:
                self.request_signature_help(self.view.sel()[0].begin())
                return True
            else:
                return False  # Let someone else handle this keybinding.
        elif len(self._signatures) < 2:
            return False  # Let someone else handle this keybinding.
        else:
            # We use the "operand" for the number -1 or +1. See the keybindings.
            new_index = self._active_signature + operand

            # clamp signature index
            new_index = max(0, min(new_index, len(self._signatures) - 1))

            # only update when changed
            if new_index != self._active_signature:
                self._active_signature = new_index
                self._update_popup()

            return True  # We handled this keybinding.

    def _show_popup(self, point: int) -> None:
        mdpopups.show_popup(self.view,
                            self._build_popup_content(),
                            css=popup_css,
                            md=True,
                            flags=sublime.HIDE_ON_MOUSE_MOVE_AWAY,
                            location=point,
                            wrapper_class=popup_class,
                            max_width=800,
                            on_hide=self._on_hide,
                            on_navigate=self._on_hover_navigate)
        self._visible = True

    def _update_popup(self) -> None:
        mdpopups.update_popup(self.view,
                              self._build_popup_content(),
                              css=popup_css,
                              md=True,
                              wrapper_class=popup_class)

    def _build_popup_content(self) -> str:
        if settings.highlight_active_signature_parameter:
            return self._build_popup_content_style_vscode()
        else:
            # Default to "sublime".
            return self._build_popup_content_style_sublime()

    def _view_language(self, view: sublime.View, config_name: str) -> 'Optional[str]':
        languages = view.settings().get('lsp_language')
        return languages.get(config_name) if languages else None

    def _on_hide(self):
        self._visible = False

    def _on_hover_navigate(self, href):
        webbrowser.open_new_tab(href)

    def _build_overload_selector(self) -> str:
        return "**{}** of **{}** overloads (use the ↑ ↓ keys to navigate):\n".format(
            str(self._active_signature + 1), str(len(self._signatures)))

    def _build_popup_content_style_sublime(self) -> str:
        signature = self._signatures[self._active_signature]
        formatted = []

        if len(self._signatures) > 1:
            formatted.append(self._build_overload_selector())

        signature_label = signature.get('label')
        if len(signature_label) > 400:
            label = "```{} ...```".format(signature_label[0:400])  # long code blocks = hangs
        else:
            label = "```{}\n{}\n```\n".format(self._language_id, signature_label)
        formatted.append(label)

        params = signature.get('parameters')
        if params:
            for parameter in params:
                paramDocs = parameter.get('documentation', None)
                if paramDocs:
                    formatted.append("**{}**\n".format(parameter.get('label')))
                    formatted.append("* *{}*\n".format(paramDocs))
        sigDocs = signature.get('documentation', None)
        if sigDocs:
            formatted.append(sigDocs)
        return "\n".join(formatted)

    def _build_popup_content_style_vscode(self) -> str:
        # Fetch all the relevant data.
        signature_label = ""
        signature_documentation = ""
        parameter_label = ""
        parameter_documentation = ""
        if self._active_signature in range(0, len(self._signatures)):
            signature = self._signatures[self._active_signature]
            signature_label = html.escape(signature["label"], quote=False)
            signature_documentation = signature.get("documentation", "")  # Optional.
            parameters = signature.get("parameters", None)
            if parameters and self._active_parameter in range(0, len(parameters)):
                parameter = parameters[self._active_parameter]
                parameter_label = html.escape(parameter["label"], quote=False)
                parameter_documentation = parameter.get("documentation", "")  # Optional.

        formatted = []

        if len(self._signatures) > 1:
            formatted.append(self._build_overload_selector())

        # Write the active signature and give special treatment to the active parameter (if found).
        # Note that this
 class and the extra
 are copied from mdpopups' HTML output. When mdpopups changes
        # its output style, we must update this literal string accordingly.
        formatted.append('

')
        if parameter_label:
            signature_label = self._replace_active_parameter(signature_label, parameter_label)
        formatted.append(signature_label)
        formatted.append("

")

        if parameter_documentation:
            formatted.append(parameter_documentation)

        if signature_documentation:
            formatted.append(signature_documentation)

        return "\n".join(formatted)

    def _replace_active_parameter(self, signature: str, parameter: str) -> str:
        if parameter[0].isalnum() and parameter[-1].isalnum():
            pattern = r'\b{}\b'.format(re.escape(parameter))
        else:
            # If the left or right boundary of the parameter string is not an alphanumeric character, the \b check will
            # never match. In this case, it's probably safe to assume the parameter string itself will be a good pattern
            # to search for.
            pattern = re.escape(parameter)
        replacement = '{}'.format(parameter)
        # FIXME: This is somewhat language-specific to look for an opening parenthesis. Most languages use parentheses
        # for their parameter lists though.
        start_of_param_list_pos = signature.find('(')
        # Note that this works even when we don't find an opening parenthesis, because .find returns -1 in that case.
        start_of_param_list = signature[start_of_param_list_pos + 1:]
        return signature[:start_of_param_list_pos + 1] + re.sub(pattern, replacement, start_of_param_list, 1)
<commit_msg>Account for 'MarkupContent' (intead of just 'str') in signature help (#520)<commit_after>import mdpopups
import sublime
import sublime_plugin
import webbrowser
import re
import html

try:
    from typing import Any, List, Dict, Optional
    assert Any and List and Dict and Optional
except ImportError:
    pass

from .core.configurations import is_supported_syntax
from .core.registry import config_for_scope, session_for_view, client_for_view
from .core.documents import get_document_position
from .core.events import global_events
from .core.protocol import Request
from .core.logging import debug
from .core.popups import popup_css, popup_class
from .core.settings import settings

def get_documentation(d: 'Dict[str, Any]') -> 'Optional[str]':
    docs = d.get('documentation', None)
    if docs is None:
        return None
    elif isinstance(docs, str):
        # In older version of the protocol, documentation was just a string.
        return docs
    elif isinstance(docs, dict):
        # This can be either "plaintext" or "markdown" format. For now, we can dump it into the popup box. It would
        # be nice to handle the markdown in a special way.
        return docs.get('value', None)
    else:
        debug('unknown documentation type:', str(d))
        return None

class SignatureHelpListener(sublime_plugin.ViewEventListener):

    def __init__(self, view):
        self.view = view
        self._initialized = False
        self._signature_help_triggers = []  # type: List[str]
        self._visible = False
        self._language_id = ""
        self._signatures = []  # type: List[Any]
        self._active_signature = -1
        self._active_parameter = -1

    @classmethod
    def is_applicable(cls, settings):
        syntax = settings.get('syntax')
        return syntax and is_supported_syntax(syntax)

    def initialize(self):
        session = session_for_view(self.view)
        if session:
            signatureHelpProvider = session.get_capability(
                'signatureHelpProvider')
            if signatureHelpProvider:
                self._signature_help_triggers = signatureHelpProvider.get(
                    'triggerCharacters')

        config = config_for_scope(self.view)
        if config:
            self._language_id = self._view_language(self.view, config.name)

        self._initialized = True

    def on_modified_async(self):
        pos = self.view.sel()[0].begin()
        # TODO: this will fire too often, narrow down using scopes or regex
        if not self._initialized:
            self.initialize()

        if self._signature_help_triggers:
            last_char = self.view.substr(pos - 1)
            if last_char in self._signature_help_triggers:
                self.request_signature_help(pos)
            elif self._visible:
                if last_char.isspace():
                    # Peek behind to find the last non-whitespace character.
                    last_char = self.view.substr(self.view.find_by_class(pos, False, ~0) - 1)
                if last_char not in self._signature_help_triggers:
                    self.view.hide_popup()

    def request_signature_help(self, point) -> None:
        client = client_for_view(self.view)
        if client:
            global_events.publish("view.on_purge_changes", self.view)
            document_position = get_document_position(self.view, point)
            if document_position:
                client.send_request(
                    Request.signatureHelp(document_position),
                    lambda response: self.handle_response(response, point))

    def handle_response(self, response: 'Optional[Dict]', point) -> None:
        if response is not None:
            self._signatures = response.get("signatures", [])
            self._active_signature = response.get("activeSignature", -1)
            self._active_parameter = response.get("activeParameter", -1)

            if self._signatures:
                if not 0 <= self._active_signature < len(self._signatures):
                    debug("activeSignature {} not a valid index for signatures length {}".format(
                        self._active_signature, len(self._signatures)))
                    self._active_signature = 0
            else:
                if self._active_signature != -1:
                    debug("activeSignature should be -1 or null when no signatures are returned")
                    self._active_signature = -1

            if len(self._signatures) > 0:
                if self._visible:
                    self._update_popup()
                else:
                    self._show_popup(point)

    def on_query_context(self, key, _, operand, __):
        if key != "lsp.signature_help":
            return False  # Let someone else handle this keybinding.
        elif not self._visible:
            if operand == 0:
                self.request_signature_help(self.view.sel()[0].begin())
                return True
            else:
                return False  # Let someone else handle this keybinding.
        elif len(self._signatures) < 2:
            return False  # Let someone else handle this keybinding.
        else:
            # We use the "operand" for the number -1 or +1. See the keybindings.
            new_index = self._active_signature + operand

            # clamp signature index
            new_index = max(0, min(new_index, len(self._signatures) - 1))

            # only update when changed
            if new_index != self._active_signature:
                self._active_signature = new_index
                self._update_popup()

            return True  # We handled this keybinding.

    def _show_popup(self, point: int) -> None:
        mdpopups.show_popup(self.view,
                            self._build_popup_content(),
                            css=popup_css,
                            md=True,
                            flags=sublime.HIDE_ON_MOUSE_MOVE_AWAY,
                            location=point,
                            wrapper_class=popup_class,
                            max_width=800,
                            on_hide=self._on_hide,
                            on_navigate=self._on_hover_navigate)
        self._visible = True

    def _update_popup(self) -> None:
        mdpopups.update_popup(self.view,
                              self._build_popup_content(),
                              css=popup_css,
                              md=True,
                              wrapper_class=popup_class)

    def _build_popup_content(self) -> str:
        if settings.highlight_active_signature_parameter:
            return self._build_popup_content_style_vscode()
        else:
            # Default to "sublime".
            return self._build_popup_content_style_sublime()

    def _view_language(self, view: sublime.View, config_name: str) -> 'Optional[str]':
        languages = view.settings().get('lsp_language')
        return languages.get(config_name) if languages else None

    def _on_hide(self):
        self._visible = False

    def _on_hover_navigate(self, href):
        webbrowser.open_new_tab(href)

    def _build_overload_selector(self) -> str:
        return "**{}** of **{}** overloads (use the ↑ ↓ keys to navigate):\n".format(
            str(self._active_signature + 1), str(len(self._signatures)))

    def _build_popup_content_style_sublime(self) -> str:
        signature = self._signatures[self._active_signature]
        formatted = []

        if len(self._signatures) > 1:
            formatted.append(self._build_overload_selector())

        signature_label = signature.get('label')
        if len(signature_label) > 400:
            label = "```{} ...```".format(signature_label[0:400])  # long code blocks = hangs
        else:
            label = "```{}\n{}\n```\n".format(self._language_id, signature_label)
        formatted.append(label)

        params = signature.get('parameters')
        if params:
            for parameter in params:
                param_docs = get_documentation(parameter)
                if param_docs:
                    formatted.append("**{}**\n".format(parameter.get('label')))
                    formatted.append("* *{}*\n".format(param_docs))
        sigDocs = signature.get('documentation', None)
        if sigDocs:
            formatted.append(sigDocs)
        return "\n".join(formatted)

    def _build_popup_content_style_vscode(self) -> str:
        # Fetch all the relevant data.
        signature_label = ""
        signature_documentation = ""  # type: Optional[str]
        parameter_label = ""
        parameter_documentation = ""  # type: Optional[str]
        if self._active_signature in range(0, len(self._signatures)):
            signature = self._signatures[self._active_signature]
            signature_label = html.escape(signature["label"], quote=False)
            signature_documentation = get_documentation(signature)
            parameters = signature.get("parameters", None)
            if parameters and self._active_parameter in range(0, len(parameters)):
                parameter = parameters[self._active_parameter]
                parameter_label = html.escape(parameter["label"], quote=False)
                parameter_documentation = get_documentation(parameter)

        formatted = []

        if len(self._signatures) > 1:
            formatted.append(self._build_overload_selector())

        # Write the active signature and give special treatment to the active parameter (if found).
        # Note that this
 class and the extra
 are copied from mdpopups' HTML output. When mdpopups changes
        # its output style, we must update this literal string accordingly.
        formatted.append('

')
        if parameter_label:
            signature_label = self._replace_active_parameter(signature_label, parameter_label)
        formatted.append(signature_label)
        formatted.append("

")

        if parameter_documentation:
            formatted.append(parameter_documentation)

        if signature_documentation:
            formatted.append(signature_documentation)

        return "\n".join(formatted)

    def _replace_active_parameter(self, signature: str, parameter: str) -> str:
        if parameter[0].isalnum() and parameter[-1].isalnum():
            pattern = r'\b{}\b'.format(re.escape(parameter))
        else:
            # If the left or right boundary of the parameter string is not an alphanumeric character, the \b check will
            # never match. In this case, it's probably safe to assume the parameter string itself will be a good pattern
            # to search for.
            pattern = re.escape(parameter)
        replacement = '{}'.format(parameter)
        # FIXME: This is somewhat language-specific to look for an opening parenthesis. Most languages use parentheses
        # for their parameter lists though.
        start_of_param_list_pos = signature.find('(')
        # Note that this works even when we don't find an opening parenthesis, because .find returns -1 in that case.
        start_of_param_list = signature[start_of_param_list_pos + 1:]
        return signature[:start_of_param_list_pos + 1] + re.sub(pattern, replacement, start_of_param_list, 1)
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
from haystack import indexes

from .models import Album, Single

class ReleaseIndex(indexes.SearchIndex, indexes.SearchIndex):
    text = indexes.NgramField(document=True, use_template=True)
    model = indexes.CharField(model_attr='_meta__verbose_name_plural', faceted=True)
    romanized_name = indexes.CharField(model_attr='romanized_name')
    name = indexes.CharField(model_attr='name')

    class Meta:
        abstract = True

    def index_queryset(self, using=None):
        qs = super(ReleaseIndex, self).index_queryset(using=using)
        qs = qs.prefetch_related('participating_idols', 'participating_groups')
        return qs

class AlbumIndex(ReleaseIndex, indexes.Indexable):
    def get_model(self):
        return Album

class SingleIndex(ReleaseIndex, indexes.Indexable):
    def get_model(self):
        return Single
<commit_msg>Whoooops.<commit_after># -*- coding: utf-8 -*-
from haystack import indexes

from .models import Album, Single

class ReleaseIndex(indexes.SearchIndex):
    text = indexes.NgramField(document=True, use_template=True)
    model = indexes.CharField(model_attr='_meta__verbose_name_plural', faceted=True)
    romanized_name = indexes.CharField(model_attr='romanized_name')
    name = indexes.CharField(model_attr='name')

    class Meta:
        abstract = True

    def index_queryset(self, using=None):
        qs = super(ReleaseIndex, self).index_queryset(using=using)
        qs = qs.prefetch_related('participating_idols', 'participating_groups')
        return qs

class AlbumIndex(ReleaseIndex, indexes.Indexable):
    def get_model(self):
        return Album

class SingleIndex(ReleaseIndex, indexes.Indexable):
    def get_model(self):
        return Single
<|endoftext|>


<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-

import io
import sys
import codecs
import argparse

from .learn_bpe import learn_bpe
from .apply_bpe import BPE, read_vocabulary
from .get_vocab import get_vocab
from .learn_joint_bpe_and_vocab import learn_joint_bpe_and_vocab

from .learn_bpe import create_parser as create_learn_bpe_parser
from .apply_bpe import create_parser as create_apply_bpe_parser
from .get_vocab import create_parser as create_get_vocab_parser
from .learn_joint_bpe_and_vocab import create_parser as create_learn_joint_bpe_and_vocab_parser

# hack for python2/3 compatibility
argparse.open = io.open

def main():
    parser = argparse.ArgumentParser(
        formatter_class=argparse.RawTextHelpFormatter,
        description="subword-nmt: unsupervised word segmentation for neural machine translation and text generation ")
    subparsers = parser.add_subparsers(dest='command',
                                       help="""command to run. Run one of the commands with '-h' for more info.

learn-bpe: learn BPE merge operations on input text.
apply-bpe: apply given BPE operations to input text.
get-vocab: extract vocabulary and word frequencies from input text.
learn-joint-bpe-and-vocab: executes recommended workflow for joint BPE.""")

    learn_bpe_parser = create_learn_bpe_parser(subparsers)
    apply_bpe_parser = create_apply_bpe_parser(subparsers)
    get_vocab_parser = create_get_vocab_parser(subparsers)
    learn_joint_bpe_and_vocab_parser = create_learn_joint_bpe_and_vocab_parser(subparsers)

    args = parser.parse_args()

    if args.command == 'learn-bpe':
        # read/write files as UTF-8
        if args.input.name != '<stdin>':
            args.input = codecs.open(args.input.name, encoding='utf-8')
        if args.output.name != '<stdout>':
            args.output = codecs.open(args.output.name, 'w', encoding='utf-8')

        learn_bpe(args.input, args.output, args.symbols, args.min_frequency, args.verbose,
                  is_dict=args.dict_input, total_symbols=args.total_symbols)
    elif args.command == 'apply-bpe':
        # read/write files as UTF-8
        args.codes = codecs.open(args.codes.name, encoding='utf-8')
        if args.input.name != '<stdin>':
            args.input = codecs.open(args.input.name, encoding='utf-8')
        if args.output.name != '<stdout>':
            args.output = codecs.open(args.output.name, 'w', encoding='utf-8')
        if args.vocabulary:
            args.vocabulary = codecs.open(args.vocabulary.name, encoding='utf-8')

        if args.vocabulary:
            vocabulary = read_vocabulary(args.vocabulary, args.vocabulary_threshold)
        else:
            vocabulary = None

        if sys.version_info < (3, 0):
            args.separator = args.separator.decode('UTF-8')
            if args.glossaries:
                args.glossaries = [g.decode('UTF-8') for g in args.glossaries]

        bpe = BPE(args.codes, args.merges, args.separator, vocabulary, args.glossaries)

        for line in args.input:
            args.output.write(bpe.process_line(line))

    elif args.command == 'get-vocab':
        if args.input.name != '<stdin>':
            args.input = codecs.open(args.input.name, encoding='utf-8')
        if args.output.name != '<stdout>':
            args.output = codecs.open(args.output.name, 'w', encoding='utf-8')
        get_vocab(args.input, args.output)
    elif args.command == 'learn-joint-bpe-and-vocab':
        learn_joint_bpe_and_vocab(args)
        if sys.version_info < (3, 0):
            args.separator = args.separator.decode('UTF-8')
    else:
        raise Exception('Invalid command provided')

# python 2/3 compatibility
if sys.version_info < (3, 0):
    sys.stderr = codecs.getwriter('UTF-8')(sys.stderr)
    sys.stdout = codecs.getwriter('UTF-8')(sys.stdout)
    sys.stdin = codecs.getreader('UTF-8')(sys.stdin)
else:
    sys.stderr = codecs.getwriter('UTF-8')(sys.stderr.buffer)
    sys.stdout = codecs.getwriter('UTF-8')(sys.stdout.buffer)
    sys.stdin = codecs.getreader('UTF-8')(sys.stdin.buffer)
<commit_msg>apply bpe-dropout in subword-nmt cli mode<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-

import io
import sys
import codecs
import argparse

from .learn_bpe import learn_bpe
from .apply_bpe import BPE, read_vocabulary
from .get_vocab import get_vocab
from .learn_joint_bpe_and_vocab import learn_joint_bpe_and_vocab

from .learn_bpe import create_parser as create_learn_bpe_parser
from .apply_bpe import create_parser as create_apply_bpe_parser
from .get_vocab import create_parser as create_get_vocab_parser
from .learn_joint_bpe_and_vocab import create_parser as create_learn_joint_bpe_and_vocab_parser

# hack for python2/3 compatibility
argparse.open = io.open

def main():
    parser = argparse.ArgumentParser(
        formatter_class=argparse.RawTextHelpFormatter,
        description="subword-nmt: unsupervised word segmentation for neural machine translation and text generation ")
    subparsers = parser.add_subparsers(dest='command',
                                       help="""command to run. Run one of the commands with '-h' for more info.

learn-bpe: learn BPE merge operations on input text.
apply-bpe: apply given BPE operations to input text.
get-vocab: extract vocabulary and word frequencies from input text.
learn-joint-bpe-and-vocab: executes recommended workflow for joint BPE.""")

    learn_bpe_parser = create_learn_bpe_parser(subparsers)
    apply_bpe_parser = create_apply_bpe_parser(subparsers)
    get_vocab_parser = create_get_vocab_parser(subparsers)
    learn_joint_bpe_and_vocab_parser = create_learn_joint_bpe_and_vocab_parser(subparsers)

    args = parser.parse_args()

    if args.command == 'learn-bpe':
        # read/write files as UTF-8
        if args.input.name != '<stdin>':
            args.input = codecs.open(args.input.name, encoding='utf-8')
        if args.output.name != '<stdout>':
            args.output = codecs.open(args.output.name, 'w', encoding='utf-8')

        learn_bpe(args.input, args.output, args.symbols, args.min_frequency, args.verbose,
                  is_dict=args.dict_input, total_symbols=args.total_symbols)
    elif args.command == 'apply-bpe':
        # read/write files as UTF-8
        args.codes = codecs.open(args.codes.name, encoding='utf-8')
        if args.input.name != '<stdin>':
            args.input = codecs.open(args.input.name, encoding='utf-8')
        if args.output.name != '<stdout>':
            args.output = codecs.open(args.output.name, 'w', encoding='utf-8')
        if args.vocabulary:
            args.vocabulary = codecs.open(args.vocabulary.name, encoding='utf-8')

        if args.vocabulary:
            vocabulary = read_vocabulary(args.vocabulary, args.vocabulary_threshold)
        else:
            vocabulary = None

        if sys.version_info < (3, 0):
            args.separator = args.separator.decode('UTF-8')
            if args.glossaries:
                args.glossaries = [g.decode('UTF-8') for g in args.glossaries]

        bpe = BPE(args.codes, args.merges, args.separator, vocabulary, args.glossaries)

        for line in args.input:
            args.output.write(bpe.process_line(line, args.dropout))

    elif args.command == 'get-vocab':
        if args.input.name != '<stdin>':
            args.input = codecs.open(args.input.name, encoding='utf-8')
        if args.output.name != '<stdout>':
            args.output = codecs.open(args.output.name, 'w', encoding='utf-8')
        get_vocab(args.input, args.output)
    elif args.command == 'learn-joint-bpe-and-vocab':
        learn_joint_bpe_and_vocab(args)
        if sys.version_info < (3, 0):
            args.separator = args.separator.decode('UTF-8')
    else:
        raise Exception('Invalid command provided')

# python 2/3 compatibility
if sys.version_info < (3, 0):
    sys.stderr = codecs.getwriter('UTF-8')(sys.stderr)
    sys.stdout = codecs.getwriter('UTF-8')(sys.stdout)
    sys.stdin = codecs.getreader('UTF-8')(sys.stdin)
else:
    sys.stderr = codecs.getwriter('UTF-8')(sys.stderr.buffer)
    sys.stdout = codecs.getwriter('UTF-8')(sys.stdout.buffer)
    sys.stdin = codecs.getreader('UTF-8')(sys.stdin.buffer)
<|endoftext|>


<commit_before>import gzip
import json
from collections import namedtuple

class Trait:
    def __init__(self, name):
        self.name = name
        self.xref_set = set()
        self.count = 1

    # def __eq__(self, other):
    #     if isinstance(other, Trait):
    #         return self.name == other.name and set(self.xref_set) == set(other.xref_set)
    #     else:
    #         return False
    #
    # def __hash__(self):
    #     return hash((self.name, self.xref_set))
    #
    # def __ne__(self, other):
    #     return not (self == other)

    @property
    def xref_string(self):
        return "|".join(["{}/{}".format(xref.db, xref.id_)
                         for xref in self.xref_set if xref.status.lower() == "current"])

    def __str__(self):
        return "{}\t{}\t{}".format(self.name, self.xref_string, self.count)

TraitXref = namedtuple("TraitXref", ["db", "id_", "status"])

def clinvar_jsons(filepath):
    with gzip.open(filepath, "rt") as f:
        for line in f:
            line = line.rstrip()
            yield json.loads(line)

def get_traits_from_json(clinvar_json, trait_dict):

    # This if-else block is due to the change in the format of the CellBase JSON that holds the
    # ClinVar data. Originally "clinvarSet" was the top level, but this level was removed and
    # referenceClinVarAssertion is now the top level.
    if "clinvarSet" in clinvar_json:
        trait_set = clinvar_json["clinvarSet"]["referenceClinVarAssertion"]["traitSet"]
    else:
        trait_set = clinvar_json["referenceClinVarAssertion"]["traitSet"]
    for trait_doc in trait_set['trait']:
        preferred_trait_name = None
        non_preferred_names = []
        for name in trait_doc['name']:
            # First trait name in the list will always be the "Preferred" one
            if name['elementValue']['type'].lower() == 'preferred':
                preferred_trait_name = name['elementValue']['value']
                break
            elif name['elementValue']['type'] in ["EFO URL", "EFO id", "EFO name"]:
                continue  # if the trait name not originally from clinvar
            else:
                non_preferred_names.append(name['elementValue']['value'])
        if preferred_trait_name is None:
            preferred_trait_name = non_preferred_names[0]

        if preferred_trait_name in trait_dict:
            trait = trait_dict[preferred_trait_name]
            trait.count += 1
        else:
            trait = Trait(preferred_trait_name)

        if "xref" in trait_doc:
            for xref in trait_doc["xref"]:
                trait.xref_set.add(TraitXref(xref["db"], xref["id"], xref["status"]))

        trait_dict[preferred_trait_name] = trait

    return trait_dict

def has_allowed_clinical_significance(clinvar_json):
    if "description" in clinvar_json["clinvarSet"]["referenceClinVarAssertion"]["clinicalSignificance"]:
        if clinvar_json["clinvarSet"]["referenceClinVarAssertion"]["clinicalSignificance"]["description"].lower() \
                in ["pathogenic", "likely pathogenic", "protective", "association", "risk_factor", "affects", "drug response"]:
            return True
    return False

<commit_msg>get_trait_set to separate function<commit_after>import gzip
import json
from collections import namedtuple

class Trait:
    def __init__(self, name):
        self.name = name
        self.xref_set = set()
        self.count = 1

    # def __eq__(self, other):
    #     if isinstance(other, Trait):
    #         return self.name == other.name and set(self.xref_set) == set(other.xref_set)
    #     else:
    #         return False
    #
    # def __hash__(self):
    #     return hash((self.name, self.xref_set))
    #
    # def __ne__(self, other):
    #     return not (self == other)

    @property
    def xref_string(self):
        return "|".join(["{}/{}".format(xref.db, xref.id_)
                         for xref in self.xref_set if xref.status.lower() == "current"])

    def __str__(self):
        return "{}\t{}\t{}".format(self.name, self.xref_string, self.count)

TraitXref = namedtuple("TraitXref", ["db", "id_", "status"])

def clinvar_jsons(filepath):
    with gzip.open(filepath, "rt") as f:
        for line in f:
            line = line.rstrip()
            yield json.loads(line)

def get_trait_set(clinvar_json):
    # This if-else block is due to the change in the format of the CellBase JSON that holds the
    # ClinVar data. Originally "clinvarSet" was the top level, but this level was removed and
    # referenceClinVarAssertion is now the top level.
    if "clinvarSet" in clinvar_json:
        trait_set = clinvar_json["clinvarSet"]["referenceClinVarAssertion"]["traitSet"]
    else:
        trait_set = clinvar_json["referenceClinVarAssertion"]["traitSet"]
    return trait_set

def get_traits_from_json(clinvar_json, trait_dict):
    trait_set = get_trait_set(clinvar_json)
    for trait_doc in trait_set['trait']:
        preferred_trait_name = None
        non_preferred_names = []
        for name in trait_doc['name']:
            # First trait name in the list will always be the "Preferred" one
            if name['elementValue']['type'].lower() == 'preferred':
                preferred_trait_name = name['elementValue']['value']
                break
            elif name['elementValue']['type'] in ["EFO URL", "EFO id", "EFO name"]:
                continue  # if the trait name not originally from clinvar
            else:
                non_preferred_names.append(name['elementValue']['value'])
        if preferred_trait_name is None:
            preferred_trait_name = non_preferred_names[0]

        if preferred_trait_name in trait_dict:
            trait = trait_dict[preferred_trait_name]
            trait.count += 1
        else:
            trait = Trait(preferred_trait_name)

        if "xref" in trait_doc:
            for xref in trait_doc["xref"]:
                trait.xref_set.add(TraitXref(xref["db"], xref["id"], xref["status"]))

        trait_dict[preferred_trait_name] = trait

    return trait_dict

def has_allowed_clinical_significance(clinvar_json):
    if "description" in clinvar_json["clinvarSet"]["referenceClinVarAssertion"]["clinicalSignificance"]:
        if clinvar_json["clinvarSet"]["referenceClinVarAssertion"]["clinicalSignificance"]["description"].lower() \
                in ["pathogenic", "likely pathogenic", "protective", "association", "risk_factor", "affects", "drug response"]:
            return True
    return False

<|endoftext|>


<commit_before># -*- coding: utf-8 -*-

## This file is part of Invenio.
## Copyright (C) 2011, 2012 CERN.
##
## Invenio is free software; you can redistribute it and/or
## modify it under the terms of the GNU General Public License as
## published by the Free Software Foundation; either version 2 of the
## License, or (at your option) any later version.
##
## Invenio is distributed in the hope that it will be useful, but
## WITHOUT ANY WARRANTY; without even the implied warranty of
## MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
## General Public License for more details.
##
## You should have received a copy of the GNU General Public License
## along with Invenio; if not, write to the Free Software Foundation, Inc.,
## 59 Temple Place, Suite 330, Boston, MA 02111-1307, USA.

"""dbquery Regression Test Suite."""

__revision__ = "$Id$"

import unittest

from invenio import dbquery
from invenio.testutils import make_test_suite, run_test_suite

class RunSqlReturnListOfDictionaries(unittest.TestCase):
    """Test run_sql behavior when with_dict parameter is provided"""

    def test_select_simple_columns_query(self):
        """dbquery - select simple columns query"""
        res = dbquery.run_sql("SELECT id,name,dbquery FROM collection WHERE id<3", with_dict=True)
        self.assertEqual(res, ([{'dbquery': None, 'id': 1, 'name': 'Atlantis Institute of Fictive Science'},
                                {'dbquery': '980:"PREPRINT"', 'id': 2, 'name': 'Preprints'}]))

    def test_select_date_format_column_query(self):
        """dbquery - select date format column query"""
        import time
        year = time.localtime().tm_year
        res = dbquery.run_sql("SELECT DATE_FORMAT(creation_date, '%Y') FROM bibrec WHERE id<3", with_dict=True)
        self.assertEqual(res, ([{"DATE_FORMAT(creation_date, '%Y')": str(year)},
                                {"DATE_FORMAT(creation_date, '%Y')": str(year)}]))

    def test_select_sum_columns_query(self):
        """dbquery - select sum columns query"""
        res = dbquery.run_sql("SELECT id+nbrecs,name FROM collection WHERE id>1 AND id<5", with_dict=True)
        self.assertEqual(res, ([{'id+nbrecs': 39L, 'name': 'Preprints'},
                                {'id+nbrecs': 17L, 'name': 'Books'},
                                {'id+nbrecs': 13L, 'name': 'Theses'}]))

    def test_select_all_columns_query(self):
        """dbquery - select all columns query"""
        res = dbquery.run_sql("SELECT * FROM collection WHERE id=2", with_dict=True) # kwalitee: disable=sql
        self.assertEqual(res, ([{'dbquery': '980:"PREPRINT"',
                                  'id': 2,
                                  'name': 'Preprints',
                                  'nbrecs': 37L,
                                  'reclist': 'x\x9cc\xf8\xcf\xc8\xc0\xf0\xe3\xff\x7ff\x066E\x16\x06\x04\x00\x00N\xbd\x04%'}]))

TEST_SUITE = make_test_suite(RunSqlReturnListOfDictionaries,)

if __name__ == "__main__":
    run_test_suite(TEST_SUITE, warn_user=True)
<commit_msg>dbquery: regression tests fix<commit_after># -*- coding: utf-8 -*-

## This file is part of Invenio.
## Copyright (C) 2011, 2012, 2013 CERN.
##
## Invenio is free software; you can redistribute it and/or
## modify it under the terms of the GNU General Public License as
## published by the Free Software Foundation; either version 2 of the
## License, or (at your option) any later version.
##
## Invenio is distributed in the hope that it will be useful, but
## WITHOUT ANY WARRANTY; without even the implied warranty of
## MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
## General Public License for more details.
##
## You should have received a copy of the GNU General Public License
## along with Invenio; if not, write to the Free Software Foundation, Inc.,
## 59 Temple Place, Suite 330, Boston, MA 02111-1307, USA.

"""dbquery Regression Test Suite."""

__revision__ = "$Id$"

from invenio import dbquery
from invenio.testutils import InvenioTestCase, make_test_suite, run_test_suite

class RunSqlReturnListOfDictionaries(InvenioTestCase):
    """Test run_sql behavior when with_dict parameter is provided"""

    def test_select_simple_columns_query(self):
        """dbquery - select simple columns query"""
        res = dbquery.run_sql("SELECT id,name,dbquery FROM collection WHERE id<3", with_dict=True)
        self.assertEqual(res, ([{'dbquery': None, 'id': 1, 'name': 'Atlantis Institute of Fictive Science'},
                                {'dbquery': 'collection:PREPRINT', 'id': 2, 'name': 'Preprints'}]))

    def test_select_date_format_column_query(self):
        """dbquery - select date format column query"""
        import time
        year = time.localtime().tm_year
        res = dbquery.run_sql("SELECT DATE_FORMAT(creation_date, '%Y') FROM bibrec WHERE id<3", with_dict=True)
        self.assertEqual(res, ([{"DATE_FORMAT(creation_date, '%Y')": str(year)},
                                {"DATE_FORMAT(creation_date, '%Y')": str(year)}]))

    def test_select_sum_columns_query(self):
        """dbquery - select sum columns query"""
        res = dbquery.run_sql("SELECT id+nbrecs,name FROM collection WHERE id>1 AND id<5", with_dict=True)
        self.assertEqual(res, ([{'id+nbrecs': 39L, 'name': 'Preprints'},
                                {'id+nbrecs': 17L, 'name': 'Books'},
                                {'id+nbrecs': 13L, 'name': 'Theses'}]))

    def test_select_all_columns_query(self):
        """dbquery - select all columns query"""
        res = dbquery.run_sql("SELECT * FROM collection WHERE id=2", with_dict=True) # kwalitee: disable=sql
        self.assertEqual(res, ([{'dbquery': 'collection:PREPRINT',
                                  'id': 2,
                                  'name': 'Preprints',
                                  'nbrecs': 37L,
                                  'reclist': 'x\x9cc\xf8\xcf\xc8\xc0\xf0\xe3\xff\x7ff\x066E\x16\x06\x04\x00\x00N\xbd\x04%'}]))

TEST_SUITE = make_test_suite(RunSqlReturnListOfDictionaries,)

if __name__ == "__main__":
    run_test_suite(TEST_SUITE, warn_user=True)
<|endoftext|>


<commit_before>#!/usr/bin/env python
"""
# gpuspec.py

Python script to produce spectra from GuppiRaw data files.

This script requires PyCuda and scikit-cuda to be installed and a GPU card
to go along with it.

"""

import numpy as np
import os
import sys

import h5py

try:
	import pycuda.driver as cuda
	import pycuda.autoinit
	import pycuda.gpuarray as gpuarray
except ImportError:
	raise RuntimeError("could not load pycuda. Please check your install.")

try:
	from skcuda.fft import Plan
	from skcuda.fft import fft as cufft
	from skcuda.misc import multiply as cumultiply
except ImportError:
	raise RuntimeError("Could not load scikit-cuda. Please check your install")

import pprint
import pylab as plt
import time

from guppi import GuppiRaw
from filterbank import Filterbank

def gpuspec(raw, n_int, f_avg, blank_dc_bin):
	header = raw.read_first_header()
	f_xx_avg = np.zeros(header["BLOCSIZE"] / dec_fac / 4)
	f_yy_avg = np.zeros(header["BLOCSIZE"] / dec_fac / 4)
	f_xy_avg = np.zeros(header["BLOCSIZE"] / dec_fac / 4, dtype='complex64')

	fft_plan = None
	df_gpu   = None

	t00 = time.time()
	for ii in range(n_int):
		print "\nIntegration ID:	 %i" % ii

		t1 = time.time()
		header, data = raw.read_next_data_block()
		t2 = time.time()
		print "(Data load:		 %2.2fs)" % (t2 - t1)

		d_xx = data[..., 0]
		d_yy = data[..., 1]

		if not fft_plan:
			t1 = time.time()
			#		  Plan(N_fft,				 input_dtype,  output_dtype, batch_size
			fft_plan = Plan(d_xx.shape[1]/dec_fac, np.complex64, np.complex64, d_xx.shape[0])
			t2 = time.time()
			print "FFT Plan:		   %2.2fs" % (t2 - t1)

		t1 = time.time()

		#print d_xx.dtype, d_xx.shape

		# Malloc on GPU
		if not df_gpu:
			print d_xx.shape[0], d_xx.shape[1]/dec_fac
			df_gpu = gpuarray.empty((d_xx.shape[0], d_xx.shape[1]/dec_fac), np.complex64)

		## XX POL
		d_gpu  = gpuarray.to_gpu(d_xx)
		cufft(d_gpu, df_gpu, fft_plan)
		d_xx_fft  = df_gpu.get()
		f_xx  = cumultiply(df_gpu, df_gpu).view('float32')[::2].get()

		## YY POL
		d_gpu  = gpuarray.to_gpu(d_yy)
		cufft(d_gpu, df_gpu, fft_plan)
		d_yy_fft  = df_gpu.get()
                f_yy = cumultiply(df_gpu, df_gpu).view('float32')[::2].get()

		## XY CROSS POL
		# Reuse d_gpu
		d_gpu = gpuarray.to_gpu(d_xx_fft)
		f_xy = cumultiply(d_gpu, df_gpu.conj()).get()

		t2 = time.time()
		print "cuFFT:			  %2.2fs" % (t2 - t1)

		t1 = time.time()

		#f_xx = np.abs(d_xx_fft)
		#f_yy = np.abs(d_yy_fft)
		#f_xy = d_xx_fft * d_yy_fft.conj()
		#print np.allclose(f_xy, f_xy2)

		if blank_dc_bin:
			print "BLANKING"
			f_xx[:, 0] = (f_xx[:, 1] + f_xx[:, -1]) / 2
			f_yy[:, 0] = (f_yy[:, 1] + f_yy[:, -1]) / 2
			f_xy[:, 0] = (f_xy[:, 1] + f_xy[:, -1]) / 2

		f_xx = np.fft.fftshift(f_xx)
		f_xx_avg += f_xx.flatten()

		f_yy = np.fft.fftshift(f_yy)
		f_yy_avg += f_yy.flatten()

		f_xy = np.fft.fftshift(f_xy)
		f_xy_avg += f_xy.flatten()

		t2 = time.time()
		print "Square + integrate: %2.2fs" % (t2 - t1)

	t01 = time.time()
	print "\nTotal time:		 %2.2fs" % (t01 - t00)

	f_xx_avg = f_xx_avg.reshape(f_xx_avg.shape[0] / f_avg, f_avg).mean(axis=1)
	f_yy_avg = f_yy_avg.reshape(f_yy_avg.shape[0] / f_avg, f_avg).mean(axis=1)
	f_xy_avg = f_xy_avg.reshape(f_xy_avg.shape[0] / f_avg, f_avg).mean(axis=1)

	# A roll is required for some reason, probably to do with the FFTshift
	f_xx_avg = np.roll(f_xx_avg, f_xx_avg.shape[0]/2)
	f_yy_avg = np.roll(f_yy_avg, f_yy_avg.shape[0]/2)
	f_xy_avg = np.roll(f_xy_avg, f_xy_avg.shape[0]/2)
	return (f_xx_avg, f_yy_avg, f_xy_avg)

if __name__ == "__main__":

	from argparse import ArgumentParser

	parser = ArgumentParser(description="Command line utility for creating \
										 spectra from GuppiRaw files.")

	parser.add_argument('filename', type=str,
						help='Name of file to read')
	parser.add_argument('-n', action='store', default=None, dest='n_int', type=int,
						help='Number of blocks to read from raw file')
	parser.add_argument('-f', action='store', default=None, dest='f_avg', type=int,
						help='Number of channels to average together after FFT')
	parser.add_argument('-b', action='store_false', default=True, dest='blank_dc_bin',
						help='Turn off blanking DC bins of coarse channels')

	args = parser.parse_args()

	# Open filterbank data
	filename = args.filename
	raw = GuppiRaw(filename)

	if args.n_int:
		n_int = args.n_int
	else:
		n_int = raw.n_blocks

	if args.f_avg:
		f_avg = args.f_avg
	else:
		f_avg = 1

	blank_dc_bin = args.blank_dc_bin

	pprint.pprint(raw.read_header()[0])

	print "Num. blocks: %s " % raw.n_blocks

	dec_fac = 1  		# This may be bad to use

	print "Decimation factor: %i" % dec_fac
	print "Num. integrations: %i" % n_int

	# Compute spectrum from raw file
	(xx, yy, xy) = gpuspec(raw, n_int, f_avg, blank_dc_bin)
	print xx.shape, yy.shape, xy.shape

	fil_header = raw.generate_filterbank_header(nchans=xx.shape[0])
	fil_data = np.row_stack((xx, yy)).reshape(2, 1, xx.shape[0])

	#fb = Filterbank(filename='test.h5', header_dict=fil_header, data_array=fil_data)
	#fb.write_to_hdf5('test_gpuspec.h5')

	# plot data
	print "Plotting..."
	plt.plot(10*np.log10(xx))
	plt.plot(10*np.log10(yy))

	plt.show()
<commit_msg>Updates for GPU calculations where possible<commit_after>#!/usr/bin/env python
"""
# gpuspec.py

Python script to produce spectra from GuppiRaw data files.

This script requires PyCuda and scikit-cuda to be installed and a GPU card
to go along with it.

"""

import numpy as np
import os
import sys

import h5py

try:
	import pycuda.driver as cuda
	import pycuda.autoinit
	import pycuda.gpuarray as gpuarray
except ImportError:
	raise RuntimeError("could not load pycuda. Please check your install.")

try:
	from skcuda.fft import Plan
	from skcuda.fft import fft as cufft
	from skcuda.misc import multiply as cumultiply
except ImportError:
	raise RuntimeError("Could not load scikit-cuda. Please check your install")

import pprint
import pylab as plt
import time

from guppi import GuppiRaw
from filterbank import Filterbank

def gpuspec(raw, n_int, f_avg, blank_dc_bin):
	header = raw.read_first_header()
	f_xx_avg = np.zeros(header["BLOCSIZE"] / 4)
	f_yy_avg = np.zeros(header["BLOCSIZE"] / 4)
	f_xy_avg = np.zeros(header["BLOCSIZE"] / 4, dtype='complex64')

	fft_plan = None
	df_gpu   = None

	t00 = time.time()
	for ii in range(n_int):
		print "\nIntegration ID:	 %i" % ii

		t1 = time.time()
		header, data = raw.read_next_data_block()
		t2 = time.time()
		print "(Data load:		 %2.2fs)" % (t2 - t1)

		d_xx = data[..., 0]
		d_yy = data[..., 1]

		if not fft_plan:
			t1 = time.time()
			#		  Plan(N_fft,				 input_dtype,  output_dtype, batch_size
			fft_plan = Plan(d_xx.shape[1], np.complex64, np.complex64, d_xx.shape[0])
			t2 = time.time()
			print "FFT Plan:		   %2.2fs" % (t2 - t1)

		t1 = time.time()

		#print d_xx.dtype, d_xx.shape

		# Malloc on GPU
		if not df_gpu:
			df_gpu = gpuarray.empty((d_xx.shape[0], d_xx.shape[1]), np.complex64)

		## XX POL
		d_gpu  = gpuarray.to_gpu(d_xx)
		cufft(d_gpu, df_gpu, fft_plan)
		d_xx_fft  = df_gpu.get()
		f_xx  = cumultiply(df_gpu, df_gpu).get().real

		## YY POL
		d_gpu  = gpuarray.to_gpu(d_yy)
		cufft(d_gpu, df_gpu, fft_plan)
		d_yy_fft  = df_gpu.get()
                f_yy = cumultiply(df_gpu, df_gpu).get().real

		## XY CROSS POL
		# Reuse d_gpu
		d_gpu = gpuarray.to_gpu(d_xx_fft)
		f_xy = cumultiply(d_gpu, df_gpu.conj()).get()

		t2 = time.time()
		print "cuFFT:			  %2.2fs" % (t2 - t1)

		#f_xx = np.abs(d_xx_fft)
		#f_yy = np.abs(d_yy_fft)
		#f_xy = d_xx_fft * d_yy_fft.conj()
		#print np.allclose(f_xy, f_xy2)

		if blank_dc_bin:
			t1 = time.time()
			f_xx[:, 0] = (f_xx[:, 1] + f_xx[:, -1]) / 2
			f_yy[:, 0] = (f_yy[:, 1] + f_yy[:, -1]) / 2
			f_xy[:, 0] = (f_xy[:, 1] + f_xy[:, -1]) / 2
			t2 = time.time()
			print "DC blanking:              %2.2fs" % (t2 - t1)

		t1 = time.time()
		f_xx = np.fft.fftshift(f_xx)
		f_yy = np.fft.fftshift(f_yy)
		f_xy = np.fft.fftshift(f_xy)
		t2 = time.time()

		print "FFT shift:                %2.2fs" % (t2 - t1)
		t1 = time.time()
		f_xx_avg += f_xx.flatten()
		f_yy_avg += f_yy.flatten()
		f_xy_avg += f_xy.flatten()
		t2 = time.time()
		print "Accumulate:               %2.2fs" % (t2 - t1)

	t01 = time.time()
	print "\nTotal time:		 %2.2fs" % (t01 - t00)

	f_xx_avg = f_xx_avg.reshape(f_xx_avg.shape[0] / f_avg, f_avg).mean(axis=1)
	f_yy_avg = f_yy_avg.reshape(f_yy_avg.shape[0] / f_avg, f_avg).mean(axis=1)
	f_xy_avg = f_xy_avg.reshape(f_xy_avg.shape[0] / f_avg, f_avg).mean(axis=1)

	# A roll is required for some reason, probably to do with the FFTshift
	f_xx_avg = np.roll(f_xx_avg, f_xx_avg.shape[0]/2)
	f_yy_avg = np.roll(f_yy_avg, f_yy_avg.shape[0]/2)
	f_xy_avg = np.roll(f_xy_avg, f_xy_avg.shape[0]/2)
	return (f_xx_avg, f_yy_avg, f_xy_avg)

if __name__ == "__main__":

	from argparse import ArgumentParser

	parser = ArgumentParser(description="Command line utility for creating \
										 spectra from GuppiRaw files.")

	parser.add_argument('filename', type=str,
						help='Name of file to read')
	parser.add_argument('-n', action='store', default=None, dest='n_int', type=int,
						help='Number of blocks to read from raw file')
	parser.add_argument('-f', action='store', default=None, dest='f_avg', type=int,
						help='Number of channels to average together after FFT')
	parser.add_argument('-b', action='store_false', default=True, dest='blank_dc_bin',
						help='Turn off blanking DC bins of coarse channels')
	parser.add_argument('-N', action='store', default=1, dest='n_int', type=int,
						help='number of integrations per dump')
	args = parser.parse_args()

	# Open filterbank data
	filename = args.filename
	raw = GuppiRaw(filename)

	if args.n_int:
		n_int = args.n_int
	else:
		n_int = raw.n_blocks

	if args.f_avg:
		f_avg = args.f_avg
	else:
		f_avg = 1

	blank_dc_bin = args.blank_dc_bin

	pprint.pprint(raw.read_header()[0])

	print "Num. blocks: %s " % raw.n_blocks
	print "Num. integrations: %i" % n_int

	# Compute spectrum from raw file
	(xx, yy, xy) = gpuspec(raw, n_int, f_avg, blank_dc_bin)
	print xx.shape, yy.shape, xy.shape

	fil_header = raw.generate_filterbank_header(nchans=xx.shape[0])
	fil_data = np.row_stack((xx, yy)).reshape(2, 1, xx.shape[0])

	#fb = Filterbank(filename='test.h5', header_dict=fil_header, data_array=fil_data)
	#fb.write_to_hdf5('test_gpuspec.h5')

	# plot data
	print "Plotting..."
	plt.plot(10*np.log10(xx))
	plt.plot(10*np.log10(yy))

	plt.show()
<|endoftext|>


<commit_before>import ciso8601
import xml.etree.cElementTree as ET
from django.contrib.gis.geos import Point
from isodate import parse_duration
from ..import_live_vehicles import ImportLiveVehiclesCommand
from ...models import Vehicle, VehicleLocation, Operator, Service

NS = {'siri': 'http://www.siri.org.uk/siri'}

def get_latlong(mvj):
    vl = mvj.find('siri:VehicleLocation', NS)
    lat = vl.find('siri:Latitude', NS).text
    long = vl.find('siri:Longitude', NS).text
    return Point(float(long), float(lat))

def items_from_response(response):
    try:
        items = ET.fromstring(response.text)
    except ET.ParseError:
        print(response)
        return ()
    return items.findall('siri:ServiceDelivery/siri:VehicleMonitoringDelivery/siri:VehicleActivity', NS)

class Command(ImportLiveVehiclesCommand):
    source_name = 'sirivm'
    url = 'sslink/SSLinkHTTP'

    operators = {
        'ENS': ('ENSB',),
        'HO': ('HEDO',),
        'SE': ('SESX',),
        'SE': ('SESX',),
        'FE': ('FESX',),
        'AKE': ('ARHE',),
        'SQ': ('BLUS', 'SVCT', 'UNIL', 'SWWD', 'DAMY', 'TDTR', 'TOUR', 'WDBC'),
        'FH': ('FHAM',),
        'RL': ('RLNE',),
        'FT': ('FTVA',),
        'FD': ('FDOR',),
    }

    def get_items(self):
        for subdomain in ('essex', 'southampton', 'slough'):
            data = """
                <Siri xmlns="http://www.siri.org.uk/siri">
                    <ServiceRequest><VehicleMonitoringRequest/></ServiceRequest>
                </Siri>
            """
            response = self.session.post('http://{}.jmwrti.co.uk:8080/RTI-SIRI-Server/SIRIHandler'.format(subdomain),
                                         data=data, timeout=5)
            for item in items_from_response(response):
                yield item

        data = """
            <Siri xmlns="http://www.siri.org.uk/siri">
                <ServiceRequest>
                    <RequestorRef>torbaydevon_siri_traveline</RequestorRef>
                    <VehicleMonitoringRequest/>
                </ServiceRequest>
            </Siri>
        """
        response = self.session.post('http://data.icarus.cloudamber.com/VehicleMonitoringRequest.ashx',
                                     data=data, timeout=5)
        for item in items_from_response(response):
            yield item

    def get_vehicle_and_service(self, item):
        mvj = item.find('siri:MonitoredVehicleJourney', NS)
        operator_ref = mvj.find('siri:OperatorRef', NS).text
        operator = None
        operator_options = None

        service = mvj.find('siri:LineRef', NS).text

        try:
            if operator_ref and not (operator_ref in {'TV', 'RB'} or operator_ref == 'TV' and not service):
                operator_options = self.operators.get(operator_ref)
                if operator_options:
                    operator = Operator.objects.get(id=operator_options[0])
                else:
                    operator = Operator.objects.get(id=operator_ref)
        except (Operator.MultipleObjectsReturned, Operator.DoesNotExist) as e:
            print(e, operator_ref, service)

        vehicle_code = mvj.find('siri:VehicleRef', NS).text
        if operator_ref and vehicle_code.startswith(operator_ref + '-'):
            vehicle_code = vehicle_code[len(operator_ref) + 1:]

        vehicle, created = Vehicle.objects.get_or_create(
            {'operator': operator},
            source=self.source,
            code=vehicle_code
        )

        # TODO: use ServiceCodes for this
        if service == 'QC':
            service = 'QuayConnect'
        elif service == 'FLCN':
            service = 'FALCON'
        elif service == 'P&R' and operator_ref == 'AKE':
            service = 'Colchester Park & Ride'
        elif service == '700' and operator_ref == 'FE':
            service = 'Sandon Park & Ride'
        elif service == '701' and operator_ref == 'FE':
            service = 'Chelmsford Park & Ride'
        elif service and service[:3] == 'BOB':
            service = service[:3] + ' ' + service[3] + ' ' + service[4:]

        services = Service.objects.filter(line_name=service, current=True)
        if operator_options:
            services = services.filter(operator__in=operator_options)
        elif operator:
            services = services.filter(operator=operator)
        else:
            return vehicle, created, None

        if services.count() > 1:
            latlong = get_latlong(mvj)
            services = services.filter(geometry__bboverlaps=latlong.buffer(0.1))

        try:
            service = services.get()
        except (Service.MultipleObjectsReturned, Service.DoesNotExist) as e:
            print(e, operator_ref, service, services, get_latlong(mvj))
            service = None

        return vehicle, created, service

    def create_vehicle_location(self, item, vehicle, service):
        datetime = item.find('siri:RecordedAtTime', NS).text
        mvj = item.find('siri:MonitoredVehicleJourney', NS)
        latlong = get_latlong(mvj)
        heading = mvj.find('siri:Bearing', NS)
        if heading is not None:
            heading = int(heading.text)
            if heading == -1:
                heading = None
        delay = mvj.find('siri:Delay', NS)
        if (delay is not None) and delay.text:
            try:
                delay = parse_duration(delay.text)
                early = -delay.total_seconds()  # "Early times are shown as negative values."
            except ValueError:
                early = None
        else:
            early = None
        return VehicleLocation(
            datetime=ciso8601.parse_datetime(datetime),
            latlong=latlong,
            heading=heading,
            early=early
        )
<commit_msg>SIRI-VM: cope with multiple vehicles having same fleet number<commit_after>import ciso8601
import xml.etree.cElementTree as ET
from django.contrib.gis.geos import Point
from isodate import parse_duration
from ..import_live_vehicles import ImportLiveVehiclesCommand
from ...models import Vehicle, VehicleLocation, Operator, Service

NS = {'siri': 'http://www.siri.org.uk/siri'}

def get_latlong(mvj):
    vl = mvj.find('siri:VehicleLocation', NS)
    lat = vl.find('siri:Latitude', NS).text
    long = vl.find('siri:Longitude', NS).text
    return Point(float(long), float(lat))

def items_from_response(response):
    try:
        items = ET.fromstring(response.text)
    except ET.ParseError:
        print(response)
        return ()
    return items.findall('siri:ServiceDelivery/siri:VehicleMonitoringDelivery/siri:VehicleActivity', NS)

class Command(ImportLiveVehiclesCommand):
    source_name = 'sirivm'
    url = 'sslink/SSLinkHTTP'

    operators = {
        'ENS': ('ENSB',),
        'HO': ('HEDO',),
        'SE': ('SESX',),
        'SE': ('SESX',),
        'FE': ('FESX',),
        'AKE': ('ARHE',),
        'SQ': ('BLUS', 'SVCT', 'UNIL', 'SWWD', 'DAMY', 'TDTR', 'TOUR', 'WDBC'),
        'FH': ('FHAM',),
        'RL': ('RLNE',),
        'FT': ('FTVA',),
        'FD': ('FDOR',),
    }

    def get_items(self):
        for subdomain in ('essex', 'southampton', 'slough'):
            data = """
                <Siri xmlns="http://www.siri.org.uk/siri">
                    <ServiceRequest><VehicleMonitoringRequest/></ServiceRequest>
                </Siri>
            """
            response = self.session.post('http://{}.jmwrti.co.uk:8080/RTI-SIRI-Server/SIRIHandler'.format(subdomain),
                                         data=data, timeout=5)
            for item in items_from_response(response):
                yield item

        data = """
            <Siri xmlns="http://www.siri.org.uk/siri">
                <ServiceRequest>
                    <RequestorRef>torbaydevon_siri_traveline</RequestorRef>
                    <VehicleMonitoringRequest/>
                </ServiceRequest>
            </Siri>
        """
        response = self.session.post('http://data.icarus.cloudamber.com/VehicleMonitoringRequest.ashx',
                                     data=data, timeout=5)
        for item in items_from_response(response):
            yield item

    def get_vehicle_and_service(self, item):
        mvj = item.find('siri:MonitoredVehicleJourney', NS)
        operator_ref = mvj.find('siri:OperatorRef', NS).text
        operator = None
        operator_options = None

        service = mvj.find('siri:LineRef', NS).text

        try:
            if operator_ref and not (operator_ref in {'TV', 'RB'} or operator_ref == 'TV' and not service):
                operator_options = self.operators.get(operator_ref)
                if operator_options:
                    operator = Operator.objects.get(id=operator_options[0])
                else:
                    operator = Operator.objects.get(id=operator_ref)
        except (Operator.MultipleObjectsReturned, Operator.DoesNotExist) as e:
            print(e, operator_ref, service)

        vehicle_code = mvj.find('siri:VehicleRef', NS).text
        if operator_ref and vehicle_code.startswith(operator_ref + '-'):
            vehicle_code = vehicle_code[len(operator_ref) + 1:]

        vehicle, created = Vehicle.objects.get_or_create(
            operator=operator,
            source=self.source,
            code=vehicle_code
        )

        # TODO: use ServiceCodes for this
        if service == 'QC':
            service = 'QuayConnect'
        elif service == 'FLCN':
            service = 'FALCON'
        elif service == 'P&R' and operator_ref == 'AKE':
            service = 'Colchester Park & Ride'
        elif service == '700' and operator_ref == 'FE':
            service = 'Sandon Park & Ride'
        elif service == '701' and operator_ref == 'FE':
            service = 'Chelmsford Park & Ride'
        elif service and service[:3] == 'BOB':
            service = service[:3] + ' ' + service[3] + ' ' + service[4:]

        services = Service.objects.filter(line_name=service, current=True)
        if operator_options:
            services = services.filter(operator__in=operator_options)
        elif operator:
            services = services.filter(operator=operator)
        else:
            return vehicle, created, None

        if services.count() > 1:
            latlong = get_latlong(mvj)
            services = services.filter(geometry__bboverlaps=latlong.buffer(0.1))

        try:
            service = services.get()
        except (Service.MultipleObjectsReturned, Service.DoesNotExist) as e:
            print(e, operator_ref, service, services, get_latlong(mvj))
            service = None

        return vehicle, created, service

    def create_vehicle_location(self, item, vehicle, service):
        datetime = item.find('siri:RecordedAtTime', NS).text
        mvj = item.find('siri:MonitoredVehicleJourney', NS)
        latlong = get_latlong(mvj)
        heading = mvj.find('siri:Bearing', NS)
        if heading is not None:
            heading = int(heading.text)
            if heading == -1:
                heading = None
        delay = mvj.find('siri:Delay', NS)
        if (delay is not None) and delay.text:
            try:
                delay = parse_duration(delay.text)
                early = -delay.total_seconds()  # "Early times are shown as negative values."
            except ValueError:
                early = None
        else:
            early = None
        return VehicleLocation(
            datetime=ciso8601.parse_datetime(datetime),
            latlong=latlong,
            heading=heading,
            early=early
        )
<|endoftext|>


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

"""
the gpustat script :)

@author Jongwook Choi
@url https://github.com/wookayin/gpustat
"""

from __future__ import print_function
from subprocess import check_output, CalledProcessError
from datetime import datetime
from collections import OrderedDict, defaultdict
try:
    from cStringIO import StringIO
except ImportError:
    from io import StringIO
import sys
import locale
import platform

__version__ = '0.2.0.dev'

class ANSIColors:
    RESET   = '\033[0m'
    WHITE   = '\033[1m'
    RED     = '\033[0;31m'
    GREEN   = '\033[0;32m'
    YELLOW  = '\033[0;33m'
    BLUE    = '\033[0;34m'
    MAGENTA = '\033[0;35m'
    CYAN    = '\033[0;36m'
    GRAY        = '\033[1;30m'
    BOLD_RED    = '\033[1;31m'
    BOLD_GREEN  = '\033[1;32m'
    BOLD_YELLOW = '\033[1;33m'

    @staticmethod
    def wrap(color, msg):
        return (color + msg + ANSIColors.RESET)

class GPUStat(object):

    def __init__(self, entry):
        if not isinstance(entry, dict):
            raise TypeError('entry should be a dict, {} given'.format(type(entry)))
        self.entry = entry
        self.processes = []

        # Handle '[Not Supported] for old GPU cards (#6)
        for k in self.entry.keys():
            if 'Not Supported' in self.entry[k]:
                self.entry[k] = None

        if self.entry['utilization.gpu'] is None:
            self.entry['utilization.gpu'] = '??'

    def __repr__(self):
        return self.print_to(StringIO()).getvalue()

    def print_to(self, fp,
                 with_colors=True,
                 show_cmd=False,
                 show_user=False,
                 show_pid=False,
                 gpuname_width=16
                 ):
        # color settings
        colors = {}
        def _conditional(cond_fn, true_value, false_value,
                         error_value=ANSIColors.GRAY):
            try:
                if cond_fn(): return true_value
                else: return false_value
            except:
                return error_value

        colors['C0'] = ANSIColors.RESET
        colors['C1'] = ANSIColors.CYAN
        colors['CName'] = ANSIColors.BLUE
        colors['CTemp'] = _conditional(lambda: int(self.entry['temperature.gpu']) < 50,
                                       ANSIColors.RED, ANSIColors.BOLD_RED)
        colors['CMemU'] = ANSIColors.BOLD_YELLOW
        colors['CMemT'] = ANSIColors.YELLOW
        colors['CMemP'] = ANSIColors.YELLOW
        colors['CUser'] = ANSIColors.GRAY
        colors['CUtil'] = _conditional(lambda: int(self.entry['utilization.gpu']) < 30,
                                       ANSIColors.GREEN, ANSIColors.BOLD_GREEN)

        if not with_colors:
            for k in list(colors.keys()):
                colors[k] = ''

        # build one-line display information
        reps = ("%(C1)s[{entry[index]}]%(C0)s %(CName)s{entry[name]:{gpuname_width}}%(C0)s |" +
                "%(CTemp)s{entry[temperature.gpu]:>3}'C%(C0)s, %(CUtil)s{entry[utilization.gpu]:>3} %%%(C0)s | " +
                "%(C1)s%(CMemU)s{entry[memory.used]:>5}%(C0)s / %(CMemT)s{entry[memory.total]:>5}%(C0)s MB"
                ) % colors
        reps = reps.format(entry=self.entry,
                           gpuname_width=gpuname_width)
        reps += " |"

        def _repr(v, none_value='???'):
            if v is None: return none_value
            else: return str(v)

        def process_repr(p):
            r = ''
            if not show_cmd or show_user:
                r += "{CUser}{}{C0}".format(_repr(p['user'], '--'), **colors)
            if show_cmd:
                if r: r += ':'
                r += "{C1}{}{C0}".format(_repr(p.get('comm', p['pid']), '--'), **colors)

            if show_pid:
                r += ("/%s" % _repr(p['pid'], '--'))
            r += '({CMemP}{}M{C0})'.format(_repr(p['used_memory'], '?'), **colors)
            return r

        for p in self.processes:
            reps += ' ' + process_repr(p)

        fp.write(reps)
        return fp

    @property
    def uuid(self):
        return self.entry['uuid']

    def add_process(self, p):
        self.processes.append(p)
        return self

class GPUStatCollection(object):

    def __init__(self, gpu_list):
        self.gpus = OrderedDict()
        for g in gpu_list:
            self.gpus[g.uuid] = g

        # attach process information (owner, pid, etc.)
        self.update_process_information()

        # attach additional system information
        self.hostname = platform.node()
        self.query_time = datetime.now()

    @staticmethod
    def new_query():
        # 1. get the list of gpu and status
        gpu_query_columns = ('index', 'uuid', 'name', 'temperature.gpu',
                             'utilization.gpu', 'memory.used', 'memory.total')
        gpu_list = []

        smi_output = check_output(
            r'nvidia-smi --query-gpu={query_cols} --format=csv,noheader,nounits'.format(
                query_cols=','.join(gpu_query_columns)
            ), shell=True).decode().strip()

        for line in smi_output.split('\n'):
            if not line: continue
            query_results = line.split(',')

            g = GPUStat({col_name: col_value.strip() for
                         (col_name, col_value) in zip(gpu_query_columns, query_results)
                         })
            gpu_list.append(g)

        return GPUStatCollection(gpu_list)

    @staticmethod
    def running_processes():
        # 1. collect all running GPU processes
        gpu_query_columns = ('gpu_uuid', 'pid', 'used_memory')
        smi_output = check_output(
            r'nvidia-smi --query-compute-apps={query_cols} --format=csv,noheader,nounits'.format(
                query_cols=','.join(gpu_query_columns)
            ), shell=True).decode()

        process_entries = []
        for line in smi_output.split('\n'):
            if not line: continue
            query_results = line.split(',')
            process_entry = dict({col_name: col_value.strip() for
                                  (col_name, col_value) in zip(gpu_query_columns, query_results)
                                  })
            process_entries.append(process_entry)

        pid_map = {int(e['pid']) : None for e in process_entries
                   if not 'Not Supported' in e['pid']}

        # 2. map pid to username, etc.
        if pid_map:
            pid_output = check_output('ps -o {} -p {}'.format(
                'pid,user:16,comm',
                ','.join(map(str, pid_map.keys()))
            ), shell=True).decode().strip()
            for line in pid_output.split('\n'):
                if (not line) or 'PID' in line: continue
                pid, user, comm = line.split()[:3]
                pid_map[int(pid)] = {
                    'user' : user,
                    'comm' : comm
                }

        # 3. add some process information to each process_entry
        for process_entry in process_entries[:]:

            if 'Not Supported' in process_entry['pid']:
                # TODO move this stuff into somewhere appropriate
                # such as running_processes(): process_entry = ...
                # or introduce Process class to elegantly handle it
                process_entry['user'] = None
                process_entry['comm'] = None
                process_entry['pid'] = None
                process_entry['used_memory'] = None
                continue

            pid = int(process_entry['pid'])

            if pid_map[pid] is None:
                # !?!? this pid is listed up in nvidia-smi's query result,
                # but actually seems not to be a valid running process. ignore!
                process_entries.remove(process_entry)
                continue

            process_entry.update(pid_map[pid])

        return process_entries

    def update_process_information(self):
        processes = self.running_processes()
        for p in processes:
            try:
                g = self.gpus[p['gpu_uuid']]
            except KeyError:
                # ignore?
                pass
            g.add_process(p)
        return self

    def __repr__(self):
        s = 'GPUStatCollection(host=%s, [\n' % self.hostname
        s += '\n'.join('  ' + str(g) for g in self.gpus)
        s += '\n])'
        return s

    def __len__(self):
        return len(self.gpus)

    def __iter__(self):
        return iter(self.gpus.values())

    def __getitem__(self, index):
        return list(self.gpus.values())[index]

    def print_formatted(self, fp=sys.stdout, no_color=False,
                        show_cmd=False, show_user=False, show_pid=False,
                        gpuname_width=16,
                        ):
        # header
        time_format = locale.nl_langinfo(locale.D_T_FMT)
        header_msg = '%(WHITE)s{hostname}%(RESET)s  {timestr}'.format(**{
            'hostname' : self.hostname,
            'timestr' : self.query_time.strftime(time_format)

        }) % (defaultdict(str) if no_color else ANSIColors.__dict__)

        print(header_msg)

        # body
        gpuname_width = max([gpuname_width] + [len(g.entry['name']) for g in self])
        for g in self:
            g.print_to(fp,
                       with_colors=not no_color,
                       show_cmd=show_cmd,
                       show_user=show_user,
                       show_pid=show_pid,
                       gpuname_width=gpuname_width)
            fp.write('\n')

        fp.flush()

def self_test():
    gpu_stats = GPUStatCollection.new_query()
    print('# of GPUS:', len(gpu_stats))
    for g in gpu_stats:
        print(g)

    process_entries = GPUStatCollection.running_processes()
    print('---Entries---')
    print(process_entries)

    print('-------------')

def new_query():
    '''
    Obtain a new GPUStatCollection instance by querying nvidia-smi
    to get the list of GPUs and running process information.
    '''
    return GPUStatCollection.new_query()

def print_gpustat(**args):
    '''
    Display the GPU query results into standard output.
    '''
    try:
        gpu_stats = GPUStatCollection.new_query()
    except CalledProcessError:
        sys.stderr.write('Error on calling nvidia-smi\n')
        sys.exit(1)

    gpu_stats.print_formatted(sys.stdout, **args)

def main():
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('--no-color', action='store_true',
                        help='Suppress colored output')
    parser.add_argument('-c', '--show-cmd', action='store_true',
                        help='Display cmd name of running process')
    parser.add_argument('-u', '--show-user', action='store_true',
                        help='Display username of running process')
    parser.add_argument('-p', '--show-pid', action='store_true',
                        help='Display PID of running process')
    parser.add_argument('--gpuname-width', type=int, default=16,
                        help='The minimum column width of GPU names, defaults to 16')
    parser.add_argument('-v', '--version', action='version',
                        version=__version__)
    args = parser.parse_args()

    print_gpustat(**vars(args))

if __name__ == '__main__':
    main()
<commit_msg>Version display is prefixed with 'gpustat'<commit_after>#!/usr/bin/env python

"""
the gpustat script :)

@author Jongwook Choi
@url https://github.com/wookayin/gpustat
"""

from __future__ import print_function
from subprocess import check_output, CalledProcessError
from datetime import datetime
from collections import OrderedDict, defaultdict
try:
    from cStringIO import StringIO
except ImportError:
    from io import StringIO
import sys
import locale
import platform

__version__ = '0.2.0.dev'

class ANSIColors:
    RESET   = '\033[0m'
    WHITE   = '\033[1m'
    RED     = '\033[0;31m'
    GREEN   = '\033[0;32m'
    YELLOW  = '\033[0;33m'
    BLUE    = '\033[0;34m'
    MAGENTA = '\033[0;35m'
    CYAN    = '\033[0;36m'
    GRAY        = '\033[1;30m'
    BOLD_RED    = '\033[1;31m'
    BOLD_GREEN  = '\033[1;32m'
    BOLD_YELLOW = '\033[1;33m'

    @staticmethod
    def wrap(color, msg):
        return (color + msg + ANSIColors.RESET)

class GPUStat(object):

    def __init__(self, entry):
        if not isinstance(entry, dict):
            raise TypeError('entry should be a dict, {} given'.format(type(entry)))
        self.entry = entry
        self.processes = []

        # Handle '[Not Supported] for old GPU cards (#6)
        for k in self.entry.keys():
            if 'Not Supported' in self.entry[k]:
                self.entry[k] = None

        if self.entry['utilization.gpu'] is None:
            self.entry['utilization.gpu'] = '??'

    def __repr__(self):
        return self.print_to(StringIO()).getvalue()

    def print_to(self, fp,
                 with_colors=True,
                 show_cmd=False,
                 show_user=False,
                 show_pid=False,
                 gpuname_width=16
                 ):
        # color settings
        colors = {}
        def _conditional(cond_fn, true_value, false_value,
                         error_value=ANSIColors.GRAY):
            try:
                if cond_fn(): return true_value
                else: return false_value
            except:
                return error_value

        colors['C0'] = ANSIColors.RESET
        colors['C1'] = ANSIColors.CYAN
        colors['CName'] = ANSIColors.BLUE
        colors['CTemp'] = _conditional(lambda: int(self.entry['temperature.gpu']) < 50,
                                       ANSIColors.RED, ANSIColors.BOLD_RED)
        colors['CMemU'] = ANSIColors.BOLD_YELLOW
        colors['CMemT'] = ANSIColors.YELLOW
        colors['CMemP'] = ANSIColors.YELLOW
        colors['CUser'] = ANSIColors.GRAY
        colors['CUtil'] = _conditional(lambda: int(self.entry['utilization.gpu']) < 30,
                                       ANSIColors.GREEN, ANSIColors.BOLD_GREEN)

        if not with_colors:
            for k in list(colors.keys()):
                colors[k] = ''

        # build one-line display information
        reps = ("%(C1)s[{entry[index]}]%(C0)s %(CName)s{entry[name]:{gpuname_width}}%(C0)s |" +
                "%(CTemp)s{entry[temperature.gpu]:>3}'C%(C0)s, %(CUtil)s{entry[utilization.gpu]:>3} %%%(C0)s | " +
                "%(C1)s%(CMemU)s{entry[memory.used]:>5}%(C0)s / %(CMemT)s{entry[memory.total]:>5}%(C0)s MB"
                ) % colors
        reps = reps.format(entry=self.entry,
                           gpuname_width=gpuname_width)
        reps += " |"

        def _repr(v, none_value='???'):
            if v is None: return none_value
            else: return str(v)

        def process_repr(p):
            r = ''
            if not show_cmd or show_user:
                r += "{CUser}{}{C0}".format(_repr(p['user'], '--'), **colors)
            if show_cmd:
                if r: r += ':'
                r += "{C1}{}{C0}".format(_repr(p.get('comm', p['pid']), '--'), **colors)

            if show_pid:
                r += ("/%s" % _repr(p['pid'], '--'))
            r += '({CMemP}{}M{C0})'.format(_repr(p['used_memory'], '?'), **colors)
            return r

        for p in self.processes:
            reps += ' ' + process_repr(p)

        fp.write(reps)
        return fp

    @property
    def uuid(self):
        return self.entry['uuid']

    def add_process(self, p):
        self.processes.append(p)
        return self

class GPUStatCollection(object):

    def __init__(self, gpu_list):
        self.gpus = OrderedDict()
        for g in gpu_list:
            self.gpus[g.uuid] = g

        # attach process information (owner, pid, etc.)
        self.update_process_information()

        # attach additional system information
        self.hostname = platform.node()
        self.query_time = datetime.now()

    @staticmethod
    def new_query():
        # 1. get the list of gpu and status
        gpu_query_columns = ('index', 'uuid', 'name', 'temperature.gpu',
                             'utilization.gpu', 'memory.used', 'memory.total')
        gpu_list = []

        smi_output = check_output(
            r'nvidia-smi --query-gpu={query_cols} --format=csv,noheader,nounits'.format(
                query_cols=','.join(gpu_query_columns)
            ), shell=True).decode().strip()

        for line in smi_output.split('\n'):
            if not line: continue
            query_results = line.split(',')

            g = GPUStat({col_name: col_value.strip() for
                         (col_name, col_value) in zip(gpu_query_columns, query_results)
                         })
            gpu_list.append(g)

        return GPUStatCollection(gpu_list)

    @staticmethod
    def running_processes():
        # 1. collect all running GPU processes
        gpu_query_columns = ('gpu_uuid', 'pid', 'used_memory')
        smi_output = check_output(
            r'nvidia-smi --query-compute-apps={query_cols} --format=csv,noheader,nounits'.format(
                query_cols=','.join(gpu_query_columns)
            ), shell=True).decode()

        process_entries = []
        for line in smi_output.split('\n'):
            if not line: continue
            query_results = line.split(',')
            process_entry = dict({col_name: col_value.strip() for
                                  (col_name, col_value) in zip(gpu_query_columns, query_results)
                                  })
            process_entries.append(process_entry)

        pid_map = {int(e['pid']) : None for e in process_entries
                   if not 'Not Supported' in e['pid']}

        # 2. map pid to username, etc.
        if pid_map:
            pid_output = check_output('ps -o {} -p {}'.format(
                'pid,user:16,comm',
                ','.join(map(str, pid_map.keys()))
            ), shell=True).decode().strip()
            for line in pid_output.split('\n'):
                if (not line) or 'PID' in line: continue
                pid, user, comm = line.split()[:3]
                pid_map[int(pid)] = {
                    'user' : user,
                    'comm' : comm
                }

        # 3. add some process information to each process_entry
        for process_entry in process_entries[:]:

            if 'Not Supported' in process_entry['pid']:
                # TODO move this stuff into somewhere appropriate
                # such as running_processes(): process_entry = ...
                # or introduce Process class to elegantly handle it
                process_entry['user'] = None
                process_entry['comm'] = None
                process_entry['pid'] = None
                process_entry['used_memory'] = None
                continue

            pid = int(process_entry['pid'])

            if pid_map[pid] is None:
                # !?!? this pid is listed up in nvidia-smi's query result,
                # but actually seems not to be a valid running process. ignore!
                process_entries.remove(process_entry)
                continue

            process_entry.update(pid_map[pid])

        return process_entries

    def update_process_information(self):
        processes = self.running_processes()
        for p in processes:
            try:
                g = self.gpus[p['gpu_uuid']]
            except KeyError:
                # ignore?
                pass
            g.add_process(p)
        return self

    def __repr__(self):
        s = 'GPUStatCollection(host=%s, [\n' % self.hostname
        s += '\n'.join('  ' + str(g) for g in self.gpus)
        s += '\n])'
        return s

    def __len__(self):
        return len(self.gpus)

    def __iter__(self):
        return iter(self.gpus.values())

    def __getitem__(self, index):
        return list(self.gpus.values())[index]

    def print_formatted(self, fp=sys.stdout, no_color=False,
                        show_cmd=False, show_user=False, show_pid=False,
                        gpuname_width=16,
                        ):
        # header
        time_format = locale.nl_langinfo(locale.D_T_FMT)
        header_msg = '%(WHITE)s{hostname}%(RESET)s  {timestr}'.format(**{
            'hostname' : self.hostname,
            'timestr' : self.query_time.strftime(time_format)

        }) % (defaultdict(str) if no_color else ANSIColors.__dict__)

        print(header_msg)

        # body
        gpuname_width = max([gpuname_width] + [len(g.entry['name']) for g in self])
        for g in self:
            g.print_to(fp,
                       with_colors=not no_color,
                       show_cmd=show_cmd,
                       show_user=show_user,
                       show_pid=show_pid,
                       gpuname_width=gpuname_width)
            fp.write('\n')

        fp.flush()

def self_test():
    gpu_stats = GPUStatCollection.new_query()
    print('# of GPUS:', len(gpu_stats))
    for g in gpu_stats:
        print(g)

    process_entries = GPUStatCollection.running_processes()
    print('---Entries---')
    print(process_entries)

    print('-------------')

def new_query():
    '''
    Obtain a new GPUStatCollection instance by querying nvidia-smi
    to get the list of GPUs and running process information.
    '''
    return GPUStatCollection.new_query()

def print_gpustat(**args):
    '''
    Display the GPU query results into standard output.
    '''
    try:
        gpu_stats = GPUStatCollection.new_query()
    except CalledProcessError:
        sys.stderr.write('Error on calling nvidia-smi\n')
        sys.exit(1)

    gpu_stats.print_formatted(sys.stdout, **args)

def main():
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('--no-color', action='store_true',
                        help='Suppress colored output')
    parser.add_argument('-c', '--show-cmd', action='store_true',
                        help='Display cmd name of running process')
    parser.add_argument('-u', '--show-user', action='store_true',
                        help='Display username of running process')
    parser.add_argument('-p', '--show-pid', action='store_true',
                        help='Display PID of running process')
    parser.add_argument('--gpuname-width', type=int, default=16,
                        help='The minimum column width of GPU names, defaults to 16')
    parser.add_argument('-v', '--version', action='version',
                        version=('gpustat %s' % __version__))
    args = parser.parse_args()

    print_gpustat(**vars(args))

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