<commit_before>#--
# Copyright (c) 2008-2013 Net-ng.
# All rights reserved.
#
# This software is licensed under the BSD License, as described in
# the file LICENSE.txt, which you should have received as part of
# this distribution.
#--

"""The ``shell`` and ``batch`` administrative commands

The ``shell`` command launches an interactive Python shell.
The ``batch`` command execute Python statements from a file.

In both cases:

  - the global variable ``apps`` is a dictionary of application name -> application object
  - the global variable ``session`` is the database session
  - the metadata of the applications are activated
"""

import __builtin__
import sys
import os
import code
import pkg_resources

from nagare import database, log, local
from nagare.admin import util, reference, command

def activate_applications(cfgfiles, debug, error):
    """Initialize applications

    In:
      - ``cfgfile`` -- paths to application configuration files or names of
        registered applications
      - ``debug`` -- enable the display of the generated SQL statements
      - ``error`` -- the function to call in case of configuration errors

    Return:
      - database session
      - {application name -> application object}
    """
    # Configure the local service
    local.worker = local.Process()
    local.request = local.Process()

    configs = []

    # Merge all the ``[logging]]`` section of all the applications
    for cfgfile in cfgfiles:
        # Read the configuration file of the application
        (cfgfile, app, dist, aconf) = util.read_application(cfgfile, error)
        configs.append((cfgfile, app, dist, aconf))

        log.configure(aconf['logging'].dict(), aconf['application']['name'])

    # Configure the logging service
    log.activate()

    apps = {}

    # For each application given, activate its metadata and its logger
    for (cfgfile, app, dist, aconf) in configs:
        name = aconf['application']['name']

        log.set_logger('nagare.application.' + name)

        requirement = None if not dist else pkg_resources.Requirement.parse(dist.project_name)
        data_path = None if not requirement else pkg_resources.resource_filename(requirement, '/data')

        (apps[name], databases) = util.activate_WSGIApp(app, cfgfile, aconf, error, data_path=data_path, debug=debug)

        for (database_settings, populate) in databases:
            database.set_metadata(*database_settings)

    session = database.session
    session.begin()
    return (session, apps)

def activate_application(cfgfile, debug, error):
    """Initialize one application

    In:
      - ``cfgfile`` -- path to an application configuration file or name of
        a registered application
      - ``debug`` -- enable the display of the generated SQL statements
      - ``error`` -- the function to call in case of configuration errors

    Return:
      - database session
      - application name
      - application object
    """
    session, apps = activate_applications([cfgfile], debug, error)
    return (session,) + apps.items()[0]

def create_globals(cfgfiles, debug, error):
    """Return a namespace with the initialized applications

    In:
      - ``cfgfile`` -- paths to application configuration files or names of
        registered applications
      - ``debug`` -- enable the display of the generated SQL statements
      - ``error`` -- the function to call in case of configuration errors

    Return:
      - the namespace with the ``apps`` and ``session`` variables defined
    """
    session, apps = activate_applications(cfgfiles, debug, error)
    return {'session': session, 'apps': apps}

# -----------------------------------------------------------------------------

def set_shell_options(optparser):
    optparser.usage += ' [application ...]'

    optparser.add_option('-d', '--debug', action='store_const', const=True, default=False, dest='debug', help='debug mode for the database engine')
    optparser.add_option('--plain', action='store_const', const=False, default=True, dest='ipython', help='launch a plain Python interpreter instead of IPython')

class IPythonShellV1(object):
    """A IPython < 0.11 interpreter
    """
    def __init__(self, ipython, banner, app_names, ns):
        """Initialisation

        In:
          - ``ipython`` -- the ``IPython`` module
          - ``banner`` -- banner to display
          - ``app_names`` -- names of the activated applications
          - ``ns`` -- the namespace with the ``apps`` and ``session`` variables defined
        """
        prompt = '[%s]' % app_names[0] if len(app_names) == 1 else ''
        self.shell = ipython.Shell.IPShell(argv=['--prompt_in1=Nagare%s [\\#]: ' % prompt], user_ns=ns)
        self.banner = banner

    def __call__(self):
        """Launch the interpreter
        """
        self.shell.mainloop(1, self.banner)

class IPythonShellV2(object):
    """A IPython >= 0.11 interpreter
    """
    def __init__(self, ipython, banner, app_names, ns):
        """Initialisation

        In:
          - ``ipython`` -- the ``IPython`` module
          - ``banner`` -- banner to display
          - ``app_names`` -- names of the activated applications
          - ``ns`` -- the namespace with the ``apps`` and ``session`` variables defined
        """
        config = ipython.config.loader.Config()
        prompt = '[%s]' % app_names[0] if len(app_names) == 1 else ''
        config.PromptManager.in_template = 'Nagare%s [\\#]: ' % prompt

        self.shell = ipython.frontend.terminal.embed.InteractiveShellEmbed(config=config, user_ns=ns, banner1=banner)
        self.shell.confirm_exit = False

    def __call__(self):
        """Launch the interpreter
        """
        self.shell()

class PythonShell(code.InteractiveConsole):
    """A plain Python interpreter
    """
    def __init__(self, banner, app_names, ns):
        """Initialisation

        In:
          - ``banner`` -- banner to display
          - ``app_names`` -- names of the activated applications
          - ``ns`` -- the namespace with the ``apps`` and ``session`` variables defined
        """
        code.InteractiveConsole.__init__(self, ns)
        self.banner = banner
        self.prompt = '[%s]' % app_names[0] if len(app_names) == 1 else ''

    def raw_input(self, prompt):
        return code.InteractiveConsole.raw_input(self, 'nagare' + self.prompt + prompt)

    def __call__(self):
        self.interact(self.banner)

class PythonShellWithHistory(PythonShell):
    """A plain Python interpreter with a readline history
    """
    def __call__(self, readline):
        """Launch the interpreter

        In:
          - ``readline`` -- the ``readline`` module
          - ``banner`` -- banner to display
        """
        # Set completion on TAB and a dedicated commands history file
        readline.parse_and_bind('tab: complete')

        history_path = os.path.expanduser('~/.nagarehistory')

        if os.path.exists(history_path):
            readline.read_history_file(history_path)

        readline.set_history_length(200)

        PythonShell.__call__(self)

        readline.write_history_file(history_path)

def create_python_shell(ipython, banner, app_names, ns):
    """Shell factory

    Create a shell according to the installed modules (``readline`` and ``ipython``)

    In:
      - ``ipython`` -- does the user want a IPython shell?
      - ``banner`` -- banner to display
      - ``app_names`` -- names of the activated applications
      - ``ns`` -- the namespace with the ``apps`` and ``session`` variables defined
    """
    if ipython:
        try:
            import IPython
        except ImportError:
            pass
        else:
            shell_factory = IPythonShellV1 if map(int, IPython.__version__.split('.')) < [0, 11] else IPythonShellV2
            shell_factory(IPython, banner, app_names, ns)()
            return

    try:
        import readline

        PythonShellWithHistory(banner, app_names, ns)(readline)
    except ImportError:
        PythonShell(banner, app_names, ns)()

def shell(parser, options, args):
    """Launch an interactive shell

    In:
      - ``parser`` -- the ``optparse.OptParser`` object used to parse the configuration file
      - ``options`` -- options in the command line
      - ``args`` -- arguments in the command line

    The arguments are a list of names of registered applications
    or paths to applications configuration files.
    """
    ns = create_globals(args, options.debug, parser.error)
    ns['__name__'] = '__console__'

    banner = "Python %s on %s\nNagare %s\n\nVariables 'apps' and 'session' are available" % (sys.version,
                sys.platform,
                pkg_resources.get_distribution('nagare').version)

    create_python_shell(options.ipython, banner, [app.name for app in ns['apps'].values()], ns)

class Shell(command.Command):
    desc = 'Launch a shell'

    set_options = staticmethod(set_shell_options)
    run = staticmethod(shell)

# -----------------------------------------------------------------------------

def set_batch_options(optparser):
    optparser.usage += ' <application> <file.py>'

    optparser.add_option('-d', '--debug', action='store_const', const=True, default=False, dest='debug', help='debug mode for the database engine')

    if not sys.argv[3:]:
        return None

    for (i, option) in enumerate(sys.argv[3:]):
        if not option.startswith('-'):
            break

    return sys.argv[:i + 4]

def batch(parser, options, args):
    """Execute Python statements a file

    In:
      - ``parser`` -- the ``optparse.OptParser`` object used to parse the configuration file
      - ``options`` -- options in the command lines
      - ``args`` -- arguments in the command lines

    The arguments are the name of a registered applications, or the path to
    an applications configuration file, followed by the paths of a file to
    execute
    """
    if not args:
        parser.error('No application given')

    if len(args) == 1:
        parser.error('No file given')

    for (i, option) in enumerate(sys.argv[3:]):
        if not option.startswith('-'):
            break

    del sys.argv[:i + 3]

    ns = create_globals(args[:1], options.debug, parser.error)
    __builtin__.__dict__.update(ns)

    reference.load_file(args[1], None)

class Batch(command.Command):
    desc = 'Execute Python statements from a file'

    set_options = staticmethod(set_batch_options)
    run = staticmethod(batch)
<commit_msg>Fixes #1533<commit_after>#--
# Copyright (c) 2008-2013 Net-ng.
# All rights reserved.
#
# This software is licensed under the BSD License, as described in
# the file LICENSE.txt, which you should have received as part of
# this distribution.
#--

"""The ``shell`` and ``batch`` administrative commands

The ``shell`` command launches an interactive Python shell.
The ``batch`` command execute Python statements from a file.

In both cases:

  - the global variable ``apps`` is a dictionary of application name -> application object
  - the global variable ``session`` is the database session
  - the metadata of the applications are activated
"""

import __builtin__
import sys
import os
import code
import pkg_resources

from nagare import database, log, local
from nagare.admin import util, reference, command

def activate_applications(cfgfiles, debug, error):
    """Initialize applications

    In:
      - ``cfgfile`` -- paths to application configuration files or names of
        registered applications
      - ``debug`` -- enable the display of the generated SQL statements
      - ``error`` -- the function to call in case of configuration errors

    Return:
      - database session
      - {application name -> application object}
    """
    # Configure the local service
    local.worker = local.Process()
    local.request = local.Process()

    configs = []

    # Merge all the ``[logging]]`` section of all the applications
    for cfgfile in cfgfiles:
        # Read the configuration file of the application
        (cfgfile, app, dist, aconf) = util.read_application(cfgfile, error)
        configs.append((cfgfile, app, dist, aconf))

        log.configure(aconf['logging'].dict(), aconf['application']['name'])

    # Configure the logging service
    log.activate()

    apps = {}

    # For each application given, activate its metadata and its logger
    for (cfgfile, app, dist, aconf) in configs:
        name = aconf['application']['name']

        log.set_logger('nagare.application.' + name)

        requirement = None if not dist else pkg_resources.Requirement.parse(dist.project_name)
        data_path = None if not requirement else pkg_resources.resource_filename(requirement, '/data')

        (apps[name], databases) = util.activate_WSGIApp(app, cfgfile, aconf, error, data_path=data_path, debug=debug)

        for (database_settings, populate) in databases:
            database.set_metadata(*database_settings)

    session = database.session
    session.begin()
    return (session, apps)

def activate_application(cfgfile, debug, error):
    """Initialize one application

    In:
      - ``cfgfile`` -- path to an application configuration file or name of
        a registered application
      - ``debug`` -- enable the display of the generated SQL statements
      - ``error`` -- the function to call in case of configuration errors

    Return:
      - database session
      - application name
      - application object
    """
    session, apps = activate_applications([cfgfile], debug, error)
    return (session,) + apps.items()[0]

def create_globals(cfgfiles, debug, error):
    """Return a namespace with the initialized applications

    In:
      - ``cfgfile`` -- paths to application configuration files or names of
        registered applications
      - ``debug`` -- enable the display of the generated SQL statements
      - ``error`` -- the function to call in case of configuration errors

    Return:
      - the namespace with the ``apps`` and ``session`` variables defined
    """
    session, apps = activate_applications(cfgfiles, debug, error)
    return {'session': session, 'apps': apps}

# -----------------------------------------------------------------------------

def set_shell_options(optparser):
    optparser.usage += ' [application ...]'

    optparser.add_option('-d', '--debug', action='store_const', const=True, default=False, dest='debug', help='debug mode for the database engine')
    optparser.add_option('--plain', action='store_const', const=False, default=True, dest='ipython', help='launch a plain Python interpreter instead of IPython')

class IPythonShellV1(object):
    """A IPython < 0.11 interpreter
    """
    def __init__(self, ipython, banner, app_names, ns):
        """Initialisation

        In:
          - ``ipython`` -- the ``IPython`` module
          - ``banner`` -- banner to display
          - ``app_names`` -- names of the activated applications
          - ``ns`` -- the namespace with the ``apps`` and ``session`` variables defined
        """
        prompt = '[%s]' % app_names[0] if len(app_names) == 1 else ''
        self.shell = ipython.Shell.IPShell(argv=['--prompt_in1=Nagare%s [\\#]: ' % prompt], user_ns=ns)
        self.banner = banner

    def __call__(self):
        """Launch the interpreter
        """
        self.shell.mainloop(1, self.banner)

class IPythonShellV2(object):
    """A IPython >= 0.11 interpreter
    """
    def __init__(self, ipython, banner, app_names, ns):
        """Initialisation

        In:
          - ``ipython`` -- the ``IPython`` module
          - ``banner`` -- banner to display
          - ``app_names`` -- names of the activated applications
          - ``ns`` -- the namespace with the ``apps`` and ``session`` variables defined
        """
        config = ipython.config.loader.Config()
        prompt = '[%s]' % app_names[0] if len(app_names) == 1 else ''
        config.PromptManager.in_template = 'Nagare%s [\\#]: ' % prompt

        self.shell = getattr(ipython, 'frontend', ipython).terminal.embed.InteractiveShellEmbed(config=config, user_ns=ns, banner1=banner)
        self.shell.confirm_exit = False

    def __call__(self):
        """Launch the interpreter
        """
        self.shell()

class PythonShell(code.InteractiveConsole):
    """A plain Python interpreter
    """
    def __init__(self, banner, app_names, ns):
        """Initialisation

        In:
          - ``banner`` -- banner to display
          - ``app_names`` -- names of the activated applications
          - ``ns`` -- the namespace with the ``apps`` and ``session`` variables defined
        """
        code.InteractiveConsole.__init__(self, ns)
        self.banner = banner
        self.prompt = '[%s]' % app_names[0] if len(app_names) == 1 else ''

    def raw_input(self, prompt):
        return code.InteractiveConsole.raw_input(self, 'nagare' + self.prompt + prompt)

    def __call__(self):
        self.interact(self.banner)

class PythonShellWithHistory(PythonShell):
    """A plain Python interpreter with a readline history
    """
    def __call__(self, readline):
        """Launch the interpreter

        In:
          - ``readline`` -- the ``readline`` module
          - ``banner`` -- banner to display
        """
        # Set completion on TAB and a dedicated commands history file
        readline.parse_and_bind('tab: complete')

        history_path = os.path.expanduser('~/.nagarehistory')

        if os.path.exists(history_path):
            readline.read_history_file(history_path)

        readline.set_history_length(200)

        PythonShell.__call__(self)

        readline.write_history_file(history_path)

def create_python_shell(ipython, banner, app_names, ns):
    """Shell factory

    Create a shell according to the installed modules (``readline`` and ``ipython``)

    In:
      - ``ipython`` -- does the user want a IPython shell?
      - ``banner`` -- banner to display
      - ``app_names`` -- names of the activated applications
      - ``ns`` -- the namespace with the ``apps`` and ``session`` variables defined
    """
    if ipython:
        try:
            import IPython
        except ImportError:
            pass
        else:
            shell_factory = IPythonShellV1 if map(int, IPython.__version__.split('.')) < [0, 11] else IPythonShellV2
            shell_factory(IPython, banner, app_names, ns)()
            return

    try:
        import readline

        PythonShellWithHistory(banner, app_names, ns)(readline)
    except ImportError:
        PythonShell(banner, app_names, ns)()

def shell(parser, options, args):
    """Launch an interactive shell

    In:
      - ``parser`` -- the ``optparse.OptParser`` object used to parse the configuration file
      - ``options`` -- options in the command line
      - ``args`` -- arguments in the command line

    The arguments are a list of names of registered applications
    or paths to applications configuration files.
    """
    ns = create_globals(args, options.debug, parser.error)
    ns['__name__'] = '__console__'

    banner = "Python %s on %s\nNagare %s\n\nVariables 'apps' and 'session' are available" % (sys.version,
                sys.platform,
                pkg_resources.get_distribution('nagare').version)

    create_python_shell(options.ipython, banner, [app.name for app in ns['apps'].values()], ns)

class Shell(command.Command):
    desc = 'Launch a shell'

    set_options = staticmethod(set_shell_options)
    run = staticmethod(shell)

# -----------------------------------------------------------------------------

def set_batch_options(optparser):
    optparser.usage += ' <application> <file.py>'

    optparser.add_option('-d', '--debug', action='store_const', const=True, default=False, dest='debug', help='debug mode for the database engine')

    if not sys.argv[3:]:
        return None

    for (i, option) in enumerate(sys.argv[3:]):
        if not option.startswith('-'):
            break

    return sys.argv[:i + 4]

def batch(parser, options, args):
    """Execute Python statements a file

    In:
      - ``parser`` -- the ``optparse.OptParser`` object used to parse the configuration file
      - ``options`` -- options in the command lines
      - ``args`` -- arguments in the command lines

    The arguments are the name of a registered applications, or the path to
    an applications configuration file, followed by the paths of a file to
    execute
    """
    if not args:
        parser.error('No application given')

    if len(args) == 1:
        parser.error('No file given')

    for (i, option) in enumerate(sys.argv[3:]):
        if not option.startswith('-'):
            break

    del sys.argv[:i + 3]

    ns = create_globals(args[:1], options.debug, parser.error)
    __builtin__.__dict__.update(ns)

    reference.load_file(args[1], None)

class Batch(command.Command):
    desc = 'Execute Python statements from a file'

    set_options = staticmethod(set_batch_options)
    run = staticmethod(batch)
<|endoftext|>


<commit_before>import re
import string
from metaphone import doublemetaphone
import math

######### TOKEN CONFIG #########

# some config on how to split a string into tokens
# - set characters to split on?
# - set characters to append at the end of tokens?

######### LABEL CONFIG #########

# these are the labels for tagging tokens
LABELS = [
'PrefixMarital',
'PrefixOther',
'GivenName',
'FirstInitial',
'MiddleName',
'MiddleInitial',
'Surname',
'LastInitial',
'SuffixGenerational',
'SuffixOther',
'Nickname',
'And'
]

NULL_LABEL = 'Null'

######## XML CONFIG ############

# this is the tag for each string
PARENT_LABEL = 'Name'

# this is the tag for a group of strings
GROUP_LABEL = 'NameCollection'

######## FEATURE CONFIG ########

VOWELS_Y = tuple('aeiouy')

def tokenFeatures(token) :

    if token in (u'&') :
        token_clean = token_abbrev = token

    else :
        token_clean = re.sub(r'(^[\W]*)|([^.\w]*$)', u'', token)
        token_abbrev = re.sub(r'\W', u'', token_clean.lower())

    metaphone = doublemetaphone(token_abbrev)

    features = {'nopunc' : token_abbrev,
                'abbrev' : token_clean.endswith('.'),
                'comma'  : token.endswith(','),
                'hyphenated' : '-' in token_clean,
                'contracted' : "'" in token_clean,
                'bracketed' : bool(re.match(r'["(\']\w+[")\']', token)),
                'case' : casing(token_clean),
                'length' : len(token_abbrev),
                'initial' : len(token_abbrev) == 1 and token_abbrev.isalpha(),
                'has.vowels'  : bool(set(token_abbrev[1:]) & set('aeiouy')),
                'roman' : set('xvi').issuperset(token_abbrev),
                'endswith.vowel' : token_abbrev.endswith(VOWELS_Y),
                'metaphone1' : metaphone[0],
                'more.vowels' : vowelRatio(token_abbrev)
                }

    if metaphone[1] :
        features['metaphone2'] = metaphone[1]

    reversed_token = token_abbrev[::-1]
    for i in range(1, len(token_abbrev)) :
        features['prefix_%s' % i] = token_abbrev[:i]
        features['suffix_%s' % i] = reversed_token[:i][::-1]
        if i > 4 :
            break

    return features

def casing(token) :
    if token.isupper() :
        return 'upper'
    elif token.islower() :
        return 'lower'
    elif token.istitle() :
        return 'title'
    elif token.isalpha() :
        return 'mixed'
    else :
        return False

def vowelRatio(token) :
    n_chars = len(token)
    if n_chars > 1:
        n_vowels = sum(token.count(c) for c in VOWELS_Y)
        return n_vowels/float(n_chars)
    else :
        return False

<commit_msg>tweak metaphone 2<commit_after>import re
import string
from metaphone import doublemetaphone
import math

######### TOKEN CONFIG #########

# some config on how to split a string into tokens
# - set characters to split on?
# - set characters to append at the end of tokens?

######### LABEL CONFIG #########

# these are the labels for tagging tokens
LABELS = [
'PrefixMarital',
'PrefixOther',
'GivenName',
'FirstInitial',
'MiddleName',
'MiddleInitial',
'Surname',
'LastInitial',
'SuffixGenerational',
'SuffixOther',
'Nickname',
'And'
]

NULL_LABEL = 'Null'

######## XML CONFIG ############

# this is the tag for each string
PARENT_LABEL = 'Name'

# this is the tag for a group of strings
GROUP_LABEL = 'NameCollection'

######## FEATURE CONFIG ########

VOWELS_Y = tuple('aeiouy')

def tokenFeatures(token) :

    if token in (u'&') :
        token_clean = token_abbrev = token

    else :
        token_clean = re.sub(r'(^[\W]*)|([^.\w]*$)', u'', token)
        token_abbrev = re.sub(r'\W', u'', token_clean.lower())

    metaphone = doublemetaphone(token_abbrev)

    features = {'nopunc' : token_abbrev,
                'abbrev' : token_clean.endswith('.'),
                'comma'  : token.endswith(','),
                'hyphenated' : '-' in token_clean,
                'contracted' : "'" in token_clean,
                'bracketed' : bool(re.match(r'["(\']\w+[")\']', token)),
                'case' : casing(token_clean),
                'length' : len(token_abbrev),
                'initial' : len(token_abbrev) == 1 and token_abbrev.isalpha(),
                'has.vowels'  : bool(set(token_abbrev[1:]) & set('aeiouy')),
                'roman' : set('xvi').issuperset(token_abbrev),
                'endswith.vowel' : token_abbrev.endswith(VOWELS_Y),
                'metaphone1' : metaphone[0],
                'metaphone2' : (metaphone[1] if metaphone[1] else metaphone[0])
                'more.vowels' : vowelRatio(token_abbrev)
                }

    reversed_token = token_abbrev[::-1]
    for i in range(1, len(token_abbrev)) :
        features['prefix_%s' % i] = token_abbrev[:i]
        features['suffix_%s' % i] = reversed_token[:i][::-1]
        if i > 4 :
            break

    return features

def casing(token) :
    if token.isupper() :
        return 'upper'
    elif token.islower() :
        return 'lower'
    elif token.istitle() :
        return 'title'
    elif token.isalpha() :
        return 'mixed'
    else :
        return False

def vowelRatio(token) :
    n_chars = len(token)
    if n_chars > 1:
        n_vowels = sum(token.count(c) for c in VOWELS_Y)
        return n_vowels/float(n_chars)
    else :
        return False

<|endoftext|>


<commit_before>import functools

import numpy as np
import torch
import torch.nn.functional
from nara_wpe.wpe import segment_axis

def torch_moveaxis(x: torch.tensor, source, destination):
    """

    >>> torch_moveaxis(torch.ones(2, 25), 1, 0).shape
    torch.Size([25, 2])
    >>> torch_moveaxis(torch.ones(2, 25), -1, -2).shape
    torch.Size([25, 2])
    >>> torch_moveaxis(torch.ones(2, 25), 0, 1).shape
    torch.Size([25, 2])
    >>> torch_moveaxis(torch.ones(2, 25), -2, -1).shape
    torch.Size([25, 2])
    >>> torch_moveaxis(torch.ones(2, 25) + 1j, -2, -1).shape
    torch.Size([25, 2])
    """
    ndim = len(x.shape)
    permutation = list(range(ndim))
    source = permutation.pop(source)
    permutation.insert(destination % ndim, source)
    return x.permute(*permutation)

def build_y_tilde(Y, taps, delay):
    """

    Note: The returned y_tilde consumes a similar amount of memory as Y, because
        of tricks with strides. Usually the memory consumprion is K times
        smaller than the memory consumprion of a contignous array,

    >>> T, D = 20, 2
    >>> Y = torch.arange(start=1, end=T * D + 1).reshape([T, D]).t()
    >>> # Y = torch.arange(start=1, end=T * D + 1).to(dtype=torch.complex128).reshape([T, D]).t()
    >>> print(Y.numpy())
    [[ 1  3  5  7  9 11 13 15 17 19 21 23 25 27 29 31 33 35 37 39]
     [ 2  4  6  8 10 12 14 16 18 20 22 24 26 28 30 32 34 36 38 40]]
    >>> taps, delay = 4, 2
    >>> Y_tilde = build_y_tilde(Y, taps, delay)
    >>> print(Y_tilde.shape, (taps*D, T))
    torch.Size([8, 20]) (8, 20)
    >>> print(Y_tilde.numpy())
    [[ 0  0  0  0  0  1  3  5  7  9 11 13 15 17 19 21 23 25 27 29]
     [ 0  0  0  0  0  2  4  6  8 10 12 14 16 18 20 22 24 26 28 30]
     [ 0  0  0  0  1  3  5  7  9 11 13 15 17 19 21 23 25 27 29 31]
     [ 0  0  0  0  2  4  6  8 10 12 14 16 18 20 22 24 26 28 30 32]
     [ 0  0  0  1  3  5  7  9 11 13 15 17 19 21 23 25 27 29 31 33]
     [ 0  0  0  2  4  6  8 10 12 14 16 18 20 22 24 26 28 30 32 34]
     [ 0  0  1  3  5  7  9 11 13 15 17 19 21 23 25 27 29 31 33 35]
     [ 0  0  2  4  6  8 10 12 14 16 18 20 22 24 26 28 30 32 34 36]]
    >>> Y_tilde = build_y_tilde(Y, taps, 0)
    >>> print(Y_tilde.shape, (taps*D, T), Y_tilde.stride())
    torch.Size([8, 20]) (8, 20) (1, 2)
    >>> print('Pseudo size:', np.prod(Y_tilde.size()) * Y_tilde.element_size())
    Pseudo size: 1280
    >>> print('Real size:', Y_tilde.storage().size() * Y_tilde.storage().element_size())
    Real size: 368
    >>> print(Y_tilde.numpy())
    [[ 0  0  0  1  3  5  7  9 11 13 15 17 19 21 23 25 27 29 31 33]
     [ 0  0  0  2  4  6  8 10 12 14 16 18 20 22 24 26 28 30 32 34]
     [ 0  0  1  3  5  7  9 11 13 15 17 19 21 23 25 27 29 31 33 35]
     [ 0  0  2  4  6  8 10 12 14 16 18 20 22 24 26 28 30 32 34 36]
     [ 0  1  3  5  7  9 11 13 15 17 19 21 23 25 27 29 31 33 35 37]
     [ 0  2  4  6  8 10 12 14 16 18 20 22 24 26 28 30 32 34 36 38]
     [ 1  3  5  7  9 11 13 15 17 19 21 23 25 27 29 31 33 35 37 39]
     [ 2  4  6  8 10 12 14 16 18 20 22 24 26 28 30 32 34 36 38 40]]

    >>> print(Y_tilde.shape, Y_tilde.stride())
    torch.Size([8, 20]) (1, 2)
    >>> print(Y_tilde[::3].shape, Y_tilde[::3].stride())
    torch.Size([3, 20]) (3, 2)
    >>> print(Y_tilde[::3].shape, Y_tilde[::3].contiguous().stride())
    torch.Size([3, 20]) (20, 1)
    >>> print(Y_tilde[::3].numpy())
    [[ 0  0  0  1  3  5  7  9 11 13 15 17 19 21 23 25 27 29 31 33]
     [ 0  0  2  4  6  8 10 12 14 16 18 20 22 24 26 28 30 32 34 36]
     [ 1  3  5  7  9 11 13 15 17 19 21 23 25 27 29 31 33 35 37 39]]

    The first columns are zero because of the delay.

    """
    S = Y.shape[:-2]
    D = Y.shape[-2]
    T = Y.shape[-1]

    def pad(x, axis=-1, pad_width=taps + delay - 1):
        npad = np.zeros([x.ndimension(), 2], dtype=np.int)
        npad[axis, 0] = pad_width
        # x_np = (np.pad(x.numpy(),
        #            pad_width=npad,
        #            mode='constant',
        #            constant_values=0))
        x = torch.nn.functional.pad(
            x,
            pad=npad[::-1].ravel().tolist(),
            mode='constant',
            value=0,
        )
        # assert x_np.shape == x.shape, (x_np.shape, x.shape)
        return x

    # Y_ = segment_axis(pad(Y), K, 1, axis=-1)
    # Y_ = np.flip(Y_, axis=-1)
    # if delay > 0:
    #     Y_ = Y_[..., :-delay, :]
    # # Y_: ... x D x T x K
    # Y_ = np.moveaxis(Y_, -1, -3)
    # # Y_: ... x K x D x T
    # Y_ = np.reshape(Y_, [*S, K * D, T])
    # # Y_: ... x KD x T

    # ToDo: write the shape
    Y_ = pad(Y)
    Y_ = torch_moveaxis(Y_, -1, -2)
    Y_ = torch.flip(Y_, dims=[-1 % Y_.ndimension()])
    Y_ = Y_.contiguous()  # Y_ = np.ascontiguousarray(Y_)
    Y_ = torch.flip(Y_, dims=[-1 % Y_.ndimension()])
    Y_ = torch_segment_axis(Y_, taps, 1, axis=-2)

    # Pytorch does not support negative strides.
    # Without this flip, the output of this function does not match the
    # analytical form, but the output of WPE will be equal.
    # Y_ = torch.flip(Y_, dims=[-2 % Y_.ndimension()])

    if delay > 0:
        Y_ = Y_[..., :-delay, :, :]
    Y_ = torch.reshape(Y_, list(S) + [T, taps * D])
    Y_ = torch_moveaxis(Y_, -2, -1)

    return Y_

def get_power_inverse(signal, psd_context=0):
    """
    Assumes single frequency bin with shape (D, T).

    >>> s = 1 / torch.tensor([np.arange(1, 6).astype(np.complex128)]*3)
    >>> get_power_inverse(s).numpy()
    array([ 1.,  4.,  9., 16., 25.])

    # >>> get_power_inverse(s * 0 + 1, 1).numpy()
    # array([1., 1., 1., 1., 1.])
    # >>> get_power_inverse(s, 1).numpy()
    # array([ 1.6       ,  2.20408163,  7.08196721, 14.04421326, 19.51219512])
    # >>> get_power_inverse(s, np.inf).numpy()
    # array([3.41620801, 3.41620801, 3.41620801, 3.41620801, 3.41620801])
    """
    power = torch.mean(torch.abs(signal)**2, dim=-2)

    if np.isposinf(psd_context):
        raise NotImplementedError(psd_context)
        # power = torch.broadcast_to(torch.mean(power, dim=-1, keepdims=True), power.shape)
    elif psd_context > 0:
        raise NotImplementedError(psd_context)
        #     assert int(psd_context) == psd_context, psd_context
        #     psd_context = int(psd_context)
        #     # import bottleneck as bn
        #     # Handle the corner case correctly (i.e. sum() / count)
        #     # Use bottleneck when only left context is requested
        #     # power = bn.move_mean(power, psd_context*2+1, min_count=1)
        #     power = window_mean(power, (psd_context, psd_context))
    elif psd_context == 0:
        pass
    else:
        raise ValueError(psd_context)
    eps = 1e-10 * torch.max(power)
    inverse_power = 1 / torch.max(power, eps)
    return inverse_power

def hermite(x):
    return x.transpose(-2, -1)  #.conj()

def wpe_v6(Y, taps=10, delay=3, iterations=3, psd_context=0, statistics_mode='full'):
    """
    Short of wpe_v7 with no extern references.
    Applicable in for-loops.

    >>> T = np.random.randint(100, 120)
    >>> D = np.random.randint(2, 8)
    >>> K = np.random.randint(3, 5)
    >>> delay = np.random.randint(0, 2)

    # Real test:
    >>> Y = np.random.normal(size=(D, T))
    >>> from nara_wpe import wpe as np_wpe
    >>> desired = np_wpe.wpe_v6(Y, K, delay, statistics_mode='full')
    >>> actual = wpe_v6(torch.tensor(Y), K, delay, statistics_mode='full').numpy()
    >>> np.testing.assert_allclose(actual, desired, atol=1e-6)

    # Complex test:
    >>> Y = np.random.normal(size=(D, T)) + 1j * np.random.normal(size=(D, T))
    >>> from nara_wpe import wpe as np_wpe
    >>> desired = np_wpe.wpe_v6(Y, K, delay, statistics_mode='full')
    >>> actual = wpe_v6(torch.tensor(Y), K, delay, statistics_mode='full').numpy()
    >>> np.testing.assert_allclose(actual, desired, atol=1e-6)
    """

    if statistics_mode == 'full':
        s = Ellipsis
    elif statistics_mode == 'valid':
        s = (Ellipsis, slice(delay + taps - 1, None))
    else:
        raise ValueError(statistics_mode)

    X = torch.clone(Y)
    Y_tilde = build_y_tilde(Y, taps, delay)
    for iteration in range(iterations):
        inverse_power = get_power_inverse(X, psd_context=psd_context)
        Y_tilde_inverse_power = Y_tilde * inverse_power[..., None, :]
        R = torch.matmul(Y_tilde_inverse_power[s], hermite(Y_tilde[s]))
        P = torch.matmul(Y_tilde_inverse_power[s], hermite(Y[s]))
        # G = _stable_solve(R, P)
        G, _ = torch.solve(P, R)
        X = Y - torch.matmul(hermite(G), Y_tilde)

    return X
<commit_msg>Add missing conj in hermite, with larger torch versions, the code now works for complex numbers<commit_after>import functools

import numpy as np
import torch
import torch.nn.functional
from nara_wpe.wpe import segment_axis

def torch_moveaxis(x: torch.tensor, source, destination):
    """

    >>> torch_moveaxis(torch.ones(2, 25), 1, 0).shape
    torch.Size([25, 2])
    >>> torch_moveaxis(torch.ones(2, 25), -1, -2).shape
    torch.Size([25, 2])
    >>> torch_moveaxis(torch.ones(2, 25), 0, 1).shape
    torch.Size([25, 2])
    >>> torch_moveaxis(torch.ones(2, 25), -2, -1).shape
    torch.Size([25, 2])
    >>> torch_moveaxis(torch.ones(2, 25) + 1j, -2, -1).shape
    torch.Size([25, 2])
    """
    ndim = len(x.shape)
    permutation = list(range(ndim))
    source = permutation.pop(source)
    permutation.insert(destination % ndim, source)
    return x.permute(*permutation)

def build_y_tilde(Y, taps, delay):
    """

    Note: The returned y_tilde consumes a similar amount of memory as Y, because
        of tricks with strides. Usually the memory consumprion is K times
        smaller than the memory consumprion of a contignous array,

    >>> T, D = 20, 2
    >>> Y = torch.arange(start=1, end=T * D + 1).reshape([T, D]).t()
    >>> # Y = torch.arange(start=1, end=T * D + 1).to(dtype=torch.complex128).reshape([T, D]).t()
    >>> print(Y.numpy())
    [[ 1  3  5  7  9 11 13 15 17 19 21 23 25 27 29 31 33 35 37 39]
     [ 2  4  6  8 10 12 14 16 18 20 22 24 26 28 30 32 34 36 38 40]]
    >>> taps, delay = 4, 2
    >>> Y_tilde = build_y_tilde(Y, taps, delay)
    >>> print(Y_tilde.shape, (taps*D, T))
    torch.Size([8, 20]) (8, 20)
    >>> print(Y_tilde.numpy())
    [[ 0  0  0  0  0  1  3  5  7  9 11 13 15 17 19 21 23 25 27 29]
     [ 0  0  0  0  0  2  4  6  8 10 12 14 16 18 20 22 24 26 28 30]
     [ 0  0  0  0  1  3  5  7  9 11 13 15 17 19 21 23 25 27 29 31]
     [ 0  0  0  0  2  4  6  8 10 12 14 16 18 20 22 24 26 28 30 32]
     [ 0  0  0  1  3  5  7  9 11 13 15 17 19 21 23 25 27 29 31 33]
     [ 0  0  0  2  4  6  8 10 12 14 16 18 20 22 24 26 28 30 32 34]
     [ 0  0  1  3  5  7  9 11 13 15 17 19 21 23 25 27 29 31 33 35]
     [ 0  0  2  4  6  8 10 12 14 16 18 20 22 24 26 28 30 32 34 36]]
    >>> Y_tilde = build_y_tilde(Y, taps, 0)
    >>> print(Y_tilde.shape, (taps*D, T), Y_tilde.stride())
    torch.Size([8, 20]) (8, 20) (1, 2)
    >>> print('Pseudo size:', np.prod(Y_tilde.size()) * Y_tilde.element_size())
    Pseudo size: 1280
    >>> print('Real size:', Y_tilde.storage().size() * Y_tilde.storage().element_size())
    Real size: 368
    >>> print(Y_tilde.numpy())
    [[ 0  0  0  1  3  5  7  9 11 13 15 17 19 21 23 25 27 29 31 33]
     [ 0  0  0  2  4  6  8 10 12 14 16 18 20 22 24 26 28 30 32 34]
     [ 0  0  1  3  5  7  9 11 13 15 17 19 21 23 25 27 29 31 33 35]
     [ 0  0  2  4  6  8 10 12 14 16 18 20 22 24 26 28 30 32 34 36]
     [ 0  1  3  5  7  9 11 13 15 17 19 21 23 25 27 29 31 33 35 37]
     [ 0  2  4  6  8 10 12 14 16 18 20 22 24 26 28 30 32 34 36 38]
     [ 1  3  5  7  9 11 13 15 17 19 21 23 25 27 29 31 33 35 37 39]
     [ 2  4  6  8 10 12 14 16 18 20 22 24 26 28 30 32 34 36 38 40]]

    >>> print(Y_tilde.shape, Y_tilde.stride())
    torch.Size([8, 20]) (1, 2)
    >>> print(Y_tilde[::3].shape, Y_tilde[::3].stride())
    torch.Size([3, 20]) (3, 2)
    >>> print(Y_tilde[::3].shape, Y_tilde[::3].contiguous().stride())
    torch.Size([3, 20]) (20, 1)
    >>> print(Y_tilde[::3].numpy())
    [[ 0  0  0  1  3  5  7  9 11 13 15 17 19 21 23 25 27 29 31 33]
     [ 0  0  2  4  6  8 10 12 14 16 18 20 22 24 26 28 30 32 34 36]
     [ 1  3  5  7  9 11 13 15 17 19 21 23 25 27 29 31 33 35 37 39]]

    The first columns are zero because of the delay.

    """
    S = Y.shape[:-2]
    D = Y.shape[-2]
    T = Y.shape[-1]

    def pad(x, axis=-1, pad_width=taps + delay - 1):
        npad = np.zeros([x.ndimension(), 2], dtype=np.int)
        npad[axis, 0] = pad_width
        # x_np = (np.pad(x.numpy(),
        #            pad_width=npad,
        #            mode='constant',
        #            constant_values=0))
        x = torch.nn.functional.pad(
            x,
            pad=npad[::-1].ravel().tolist(),
            mode='constant',
            value=0,
        )
        # assert x_np.shape == x.shape, (x_np.shape, x.shape)
        return x

    # Y_ = segment_axis(pad(Y), K, 1, axis=-1)
    # Y_ = np.flip(Y_, axis=-1)
    # if delay > 0:
    #     Y_ = Y_[..., :-delay, :]
    # # Y_: ... x D x T x K
    # Y_ = np.moveaxis(Y_, -1, -3)
    # # Y_: ... x K x D x T
    # Y_ = np.reshape(Y_, [*S, K * D, T])
    # # Y_: ... x KD x T

    # ToDo: write the shape
    Y_ = pad(Y)
    Y_ = torch_moveaxis(Y_, -1, -2)
    Y_ = torch.flip(Y_, dims=[-1 % Y_.ndimension()])
    Y_ = Y_.contiguous()  # Y_ = np.ascontiguousarray(Y_)
    Y_ = torch.flip(Y_, dims=[-1 % Y_.ndimension()])
    Y_ = torch_segment_axis(Y_, taps, 1, axis=-2)

    # Pytorch does not support negative strides.
    # Without this flip, the output of this function does not match the
    # analytical form, but the output of WPE will be equal.
    # Y_ = torch.flip(Y_, dims=[-2 % Y_.ndimension()])

    if delay > 0:
        Y_ = Y_[..., :-delay, :, :]
    Y_ = torch.reshape(Y_, list(S) + [T, taps * D])
    Y_ = torch_moveaxis(Y_, -2, -1)

    return Y_

def get_power_inverse(signal, psd_context=0):
    """
    Assumes single frequency bin with shape (D, T).

    >>> s = 1 / torch.tensor([np.arange(1, 6).astype(np.complex128)]*3)
    >>> get_power_inverse(s).numpy()
    array([ 1.,  4.,  9., 16., 25.])

    # >>> get_power_inverse(s * 0 + 1, 1).numpy()
    # array([1., 1., 1., 1., 1.])
    # >>> get_power_inverse(s, 1).numpy()
    # array([ 1.6       ,  2.20408163,  7.08196721, 14.04421326, 19.51219512])
    # >>> get_power_inverse(s, np.inf).numpy()
    # array([3.41620801, 3.41620801, 3.41620801, 3.41620801, 3.41620801])
    """
    power = torch.mean(torch.abs(signal)**2, dim=-2)

    if np.isposinf(psd_context):
        raise NotImplementedError(psd_context)
        # power = torch.broadcast_to(torch.mean(power, dim=-1, keepdims=True), power.shape)
    elif psd_context > 0:
        raise NotImplementedError(psd_context)
        #     assert int(psd_context) == psd_context, psd_context
        #     psd_context = int(psd_context)
        #     # import bottleneck as bn
        #     # Handle the corner case correctly (i.e. sum() / count)
        #     # Use bottleneck when only left context is requested
        #     # power = bn.move_mean(power, psd_context*2+1, min_count=1)
        #     power = window_mean(power, (psd_context, psd_context))
    elif psd_context == 0:
        pass
    else:
        raise ValueError(psd_context)
    eps = 1e-10 * torch.max(power)
    inverse_power = 1 / torch.max(power, eps)
    return inverse_power

def transpose(x):
    return x.transpose(-2, -1)

def hermite(x):
    return x.transpose(-2, -1).conj()

def wpe_v6(Y, taps=10, delay=3, iterations=3, psd_context=0, statistics_mode='full'):
    """
    Short of wpe_v7 with no extern references.
    Applicable in for-loops.

    >>> T = np.random.randint(100, 120)
    >>> D = np.random.randint(2, 8)
    >>> K = np.random.randint(3, 5)
    >>> delay = np.random.randint(0, 2)

    # Real test:
    >>> Y = np.random.normal(size=(D, T))
    >>> from nara_wpe import wpe as np_wpe
    >>> desired = np_wpe.wpe_v6(Y, K, delay, statistics_mode='full')
    >>> actual = wpe_v6(torch.tensor(Y), K, delay, statistics_mode='full').numpy()
    >>> np.testing.assert_allclose(actual, desired, atol=1e-6)

    # Complex test:
    >>> Y = np.random.normal(size=(D, T)) + 1j * np.random.normal(size=(D, T))
    >>> from nara_wpe import wpe as np_wpe
    >>> desired = np_wpe.wpe_v6(Y, K, delay, statistics_mode='full')
    >>> actual = wpe_v6(torch.tensor(Y), K, delay, statistics_mode='full').numpy()
    >>> np.testing.assert_allclose(actual, desired, atol=1e-6)
    """

    if statistics_mode == 'full':
        s = Ellipsis
    elif statistics_mode == 'valid':
        s = (Ellipsis, slice(delay + taps - 1, None))
    else:
        raise ValueError(statistics_mode)

    X = torch.clone(Y)
    Y_tilde = build_y_tilde(Y, taps, delay)
    for iteration in range(iterations):
        inverse_power = get_power_inverse(X, psd_context=psd_context)
        Y_tilde_inverse_power = Y_tilde * inverse_power[..., None, :]
        R = torch.matmul(Y_tilde_inverse_power[s], hermite(Y_tilde[s]))
        P = torch.matmul(Y_tilde_inverse_power[s], hermite(Y[s]))
        # G = _stable_solve(R, P)
        G = torch.linalg.solve(R, P)
        X = Y - torch.matmul(hermite(G), Y_tilde)

    return X
<|endoftext|>


<commit_before>from .settings import *

DEBUG = True
TEMPLATE_DEBUG = True

DATABASES = {
    'default': {
        'ENGINE': 'django.db.backends.sqlite3',
        'NAME': 'db.sqlite',
    },
    'volatile': {
        'ENGINE': 'django.db.backends.sqlite3',
        'NAME': ':memory:',
    }
}

# Non root user (root under development)
RUN_AS_USER = 'nobody'

LOGGING = {
    'version': 1,
    'disable_existing_loggers': False,
    'formatters': {
        'default': {
            'format': '%(asctime)s [%(levelname)s] %(module)s - %(message)s'
        },
        'simple': {
            'format': '%(levelname)s %(message)s'
        },
    },
    'handlers': {
        'null': {
            'level': 'DEBUG',
            'class': 'logging.NullHandler',
        },
        'console': {
            'level': 'DEBUG',
            'formatter': 'default',
            'class': 'logging.StreamHandler',
        },
    },
    'loggers': {
        '': {
            'handlers': ['console'],
            'level': 'DEBUG',
            'propagate': True,
        },
        'django.db.backends': {
            'handlers': ['null'],
            'propagate': False,
            'level': 'DEBUG',
        },
    }
}
<commit_msg>Clean tests output<commit_after>from .settings import *

DEBUG = True
TEMPLATE_DEBUG = True

DATABASES = {
    'default': {
        'ENGINE': 'django.db.backends.sqlite3',
        'NAME': 'db.sqlite',
    },
    'volatile': {
        'ENGINE': 'django.db.backends.sqlite3',
        'NAME': ':memory:',
    }
}

# Non root user (root under development)
RUN_AS_USER = 'nobody'

LOGGING = {
    'version': 1,
    'disable_existing_loggers': False,
    'formatters': {
        'default': {
            'format': '%(asctime)s [%(levelname)s] %(module)s - %(message)s'
        },
        'simple': {
            'format': '%(levelname)s %(message)s'
        },
    },
    'handlers': {
        'null': {
            'level': 'DEBUG',
            'class': 'logging.NullHandler',
        },
        'console': {
            'level': 'DEBUG',
            'formatter': 'default',
            'class': 'logging.StreamHandler',
        },
    },
    'loggers': {
        '': {
            'handlers': ['null'],
            'level': 'DEBUG',
            'propagate': True,
        },
        'django.db.backends': {
            'handlers': ['null'],
            'propagate': False,
            'level': 'DEBUG',
        },
    }
}
<|endoftext|>


<commit_before>__version_info__ = (0, 1, 0, 'dev1')
__version__ = '.'.join(map(str, __version_info__))
<commit_msg>bump to dev2<commit_after>__version_info__ = (0, 1, 0, 'dev', 2)
__version__ = '.'.join(map(str, __version_info__))
<|endoftext|>


<commit_before>import modular_core.fundamental as lfu
import modular_core.data.batch_target as dba
import modular_core.settings as lset
import modular_core.parallel.dispycluster as mdcl
import modular_core.parallel.mpicluster as mmcl

import pdb,sys,time,types
import numpy as np
import multiprocessing as mp

from mpi4py import MPI

if __name__ == 'modular_core.parallel.parallelplan':
    lfu.check_gui_pack()
    lgb = lfu.gui_pack.lgb
    lgm = lfu.gui_pack.lgm
    lgd = lfu.gui_pack.lgd
if __name__ == '__main__':print 'parallelplan of modular_core'

###############################################################################
### a parallel_plan provides parallel support to an ensemble
###############################################################################

cluster_ips = [
    #'10.0.0.5',
    #'10.0.0.8',

    #'127.0.1.1',

    #'192.168.4.86',
    #'192.168.4.76',
    #'192.168.4.87',
    #'192.168.4.89',
    #'192.168.2.173',

    #'hemlock.phys.vt.edu',
        ]

class parallel_plan(lfu.plan):

    def _string(self):
        dist = str(self.distributed)
        mstr = '\tworkers : %d\n\tdistributed : %s'%(self.worker_count,dist)

    def __init__(self,*args,**kwargs):
        self._default('name','parallel plan',**kwargs)
        use_plan = lset.get_setting('multiprocessing')
        self._default('use_plan',use_plan,**kwargs)
        self._default('cluster_type','mpi',**kwargs)
        self.worker_count = lset.get_setting('worker_processes')
        self.distributed = lset.get_setting('distributed')
        self.cluster_node_ips = cluster_ips
        lfu.plan.__init__(self,*args,**kwargs)

    def _cluster(self,arc_length,work):
        ensem = self.parent
        if comm.rank == 0:print 'CLUSTERIZING...'
        if self.cluster_type == 'dispy':
            nodes = self.cluster_node_ips
            with open(ensem.mcfg_path,'r') as mh:mcfgstring = mh.read()
            modulename = ensem.module_name
            wrgs = [(mcfgstring,modulename,x) for x in range(arc_length)]
            deps = ensem.module.dependencies
            loc_0th_pools = mdcl.clusterize(nodes,work,wrgs,deps)
            if comm.rank == 0:
                zeroth = ensem.postprocess_plan.zeroth
                zcount = len(zeroth)
                for adx in range(arc_length):
                    l0p = loc_0th_pools[adx]

                    if cplan.maintain_pspmap:
                        mloc = l0p.metalocation
                        mstr = mloc.location_string
                        cplan.metamap.entries[mstr] = mloc
                        cplan.metamap.location_strings.append(mstr)

                    for zdx in range(zcount):
                        zp = zeroth[zdx]
                        zpdata = l0p.children[zdx]
                        zpdata._unfriendly()
                        zp.data._add_child(zpdata)
                        zp.data._stow_child(-1,v = False)

                if cplan.maintain_pspmap:cplan._save_metamap()
        elif self.cluster_type == 'mpi':
            comm = MPI.COMM_WORLD
            loc_0th_pools = mmcl.clusterize(ensem,arc_length)
            if comm.rank == 0:
                print 'CLUSTERIZED...'
                #pdb.set_trace()
        #return loc_0th_pools

    def _widget(self,*args,**kwargs):
        self._sanitize(*args,**kwargs)
        self.widg_templates.append(
            lgm.interface_template_gui(
                widgets = ['spin'],
                initials = [[self.worker_count]],
                minimum_values = [[1]],
                maximum_values = [[1000]],
                instances = [[self]],
                rewidget = [[True]],
                keys = [['worker_count']],
                box_labels = ['Number of worker processes']))
        self.widg_templates.append(
            lgm.interface_template_gui(
                widgets = ['check_set'],
                append_instead = [False],
                instances = [[self]],
                keys = [['distributed']],
                labels = [['run distributed']],
                box_labels = ['Cluster']))
        lfu.plan._widget(self,*args,from_sub = True)

###############################################################################
###############################################################################

<commit_msg>mpi work<commit_after>import modular_core.fundamental as lfu
import modular_core.data.batch_target as dba
import modular_core.settings as lset
import modular_core.parallel.dispycluster as mdcl
import modular_core.parallel.mpicluster as mmcl

import pdb,sys,time,types
import numpy as np
import multiprocessing as mp

from mpi4py import MPI

if __name__ == 'modular_core.parallel.parallelplan':
    lfu.check_gui_pack()
    lgb = lfu.gui_pack.lgb
    lgm = lfu.gui_pack.lgm
    lgd = lfu.gui_pack.lgd
if __name__ == '__main__':print 'parallelplan of modular_core'

###############################################################################
### a parallel_plan provides parallel support to an ensemble
###############################################################################

cluster_ips = [
    #'10.0.0.5',
    #'10.0.0.8',

    #'127.0.1.1',

    #'192.168.4.86',
    #'192.168.4.76',
    #'192.168.4.87',
    #'192.168.4.89',
    #'192.168.2.173',

    #'hemlock.phys.vt.edu',
        ]

class parallel_plan(lfu.plan):

    def _string(self):
        dist = str(self.distributed)
        mstr = '\tworkers : %d\n\tdistributed : %s'%(self.worker_count,dist)

    def __init__(self,*args,**kwargs):
        self._default('name','parallel plan',**kwargs)
        use_plan = lset.get_setting('multiprocessing')
        self._default('use_plan',use_plan,**kwargs)
        self._default('cluster_type','mpi',**kwargs)
        self.worker_count = lset.get_setting('worker_processes')
        self.distributed = lset.get_setting('distributed')
        self.cluster_node_ips = cluster_ips
        lfu.plan.__init__(self,*args,**kwargs)

    def _cluster(self,arc_length,work):
        ensem = self.parent
        comm = MPI.COMM_WORLD
        if comm.rank == 0:print 'CLUSTERIZING...'
        if self.cluster_type == 'dispy':
            nodes = self.cluster_node_ips
            with open(ensem.mcfg_path,'r') as mh:mcfgstring = mh.read()
            modulename = ensem.module_name
            wrgs = [(mcfgstring,modulename,x) for x in range(arc_length)]
            deps = ensem.module.dependencies
            loc_0th_pools = mdcl.clusterize(nodes,work,wrgs,deps)
            if comm.rank == 0:
                zeroth = ensem.postprocess_plan.zeroth
                zcount = len(zeroth)
                for adx in range(arc_length):
                    l0p = loc_0th_pools[adx]

                    if cplan.maintain_pspmap:
                        mloc = l0p.metalocation
                        mstr = mloc.location_string
                        cplan.metamap.entries[mstr] = mloc
                        cplan.metamap.location_strings.append(mstr)

                    for zdx in range(zcount):
                        zp = zeroth[zdx]
                        zpdata = l0p.children[zdx]
                        zpdata._unfriendly()
                        zp.data._add_child(zpdata)
                        zp.data._stow_child(-1,v = False)

                if cplan.maintain_pspmap:cplan._save_metamap()
        elif self.cluster_type == 'mpi':
            comm = MPI.COMM_WORLD
            loc_0th_pools = mmcl.clusterize(ensem,arc_length)
            if comm.rank == 0:
                print 'CLUSTERIZED...'
                #pdb.set_trace()
        #return loc_0th_pools

    def _widget(self,*args,**kwargs):
        self._sanitize(*args,**kwargs)
        self.widg_templates.append(
            lgm.interface_template_gui(
                widgets = ['spin'],
                initials = [[self.worker_count]],
                minimum_values = [[1]],
                maximum_values = [[1000]],
                instances = [[self]],
                rewidget = [[True]],
                keys = [['worker_count']],
                box_labels = ['Number of worker processes']))
        self.widg_templates.append(
            lgm.interface_template_gui(
                widgets = ['check_set'],
                append_instead = [False],
                instances = [[self]],
                keys = [['distributed']],
                labels = [['run distributed']],
                box_labels = ['Cluster']))
        lfu.plan._widget(self,*args,from_sub = True)

###############################################################################
###############################################################################

<|endoftext|>


<commit_before># The purpose of this script is to download information about all pull
# requests merged into the master branch of the given repository. This
# information is downloaded to a JSON file.

import os
import sys
import json
import requests

from common import get_credentials

QUERY_TEMPLATE = """
{{
  repository(owner: "{owner}", name: "{repository}") {{
    pullRequests(first:100, orderBy: {{direction: ASC, field: CREATED_AT}}, baseRefName: "master", states: MERGED{after}) {{
      edges {{
        node {{
          title
          number
          mergeCommit {{
            oid
          }}
          createdAt
          updatedAt
          mergedAt
          milestone {{
            title
          }}
          labels(first: 10) {{
            edges {{
              node {{
                name
              }}
            }}
          }}
        }}
        cursor
      }}
    }}
  }}
}}
"""

if sys.argv[1:]:
    REPOSITORY = sys.argv[1]
else:
    REPOSITORY = 'astropy/astropy'

OWNER = os.path.dirname(REPOSITORY)
NAME = os.path.basename(REPOSITORY)

print("The repository this script currently works with is '{}'.\n"
      .format(REPOSITORY))

json_filename = 'merged_pull_requests_{}.json'.format(NAME)

TOKEN = get_credentials('N/A', needs_token=True)[1]

headers = {"Authorization": f"Bearer {TOKEN}"}

cursor = None

pull_requests = {}

try:

    while True:

        print('cursor:', cursor)

        if cursor is None:
            after = ''
        else:
            after = f', after:"{cursor}"'

        query = QUERY_TEMPLATE.format(owner=OWNER, repository=NAME, after=after)

        request = requests.post('https://api.github.com/graphql', json={'query': query}, headers=headers)

        if request.status_code != 200:
            raise Exception("Query failed")

        entries = request.json()['data']['repository']['pullRequests']['edges']

        for entry in entries:

            pr = entry['node']
            cursor = entry['cursor']

            pull_requests[str(pr['number'])] = {'milestone': pr['milestone']['title'] if pr['milestone'] else None,
                                                'title': pr['title'],
                                                'labels': [edge['node']['name'] for edge in pr['labels']['edges']],
                                                'merged': pr['mergedAt'].replace('Z', ''),
                                                'updated': pr['updatedAt'].replace('Z', ''),
                                                'created': pr['createdAt'].replace('Z', ''),
                                                'merge_commit': pr['mergeCommit']['oid'] if pr['mergeCommit'] else None}

        if len(entries) < 100:
            break

finally:

    with open(json_filename, 'w') as f:
        json.dump(pull_requests, f, sort_keys=True, indent=2)
<commit_msg>fix PR-getter to also get from main<commit_after># The purpose of this script is to download information about all pull
# requests merged into the master branch of the given repository. This
# information is downloaded to a JSON file.

import os
import sys
import json
import requests

from common import get_credentials

QUERY_TEMPLATE = """
{{
  repository(owner: "{owner}", name: "{repository}") {{
    pullRequests(first:100, orderBy: {{direction: ASC, field: CREATED_AT}}, baseRefName: "{basename}", states: MERGED{after}) {{
      edges {{
        node {{
          title
          number
          mergeCommit {{
            oid
          }}
          createdAt
          updatedAt
          mergedAt
          milestone {{
            title
          }}
          labels(first: 10) {{
            edges {{
              node {{
                name
              }}
            }}
          }}
        }}
        cursor
      }}
    }}
  }}
}}
"""

if sys.argv[1:]:
    REPOSITORY = sys.argv[1]
else:
    REPOSITORY = 'astropy/astropy'

OWNER = os.path.dirname(REPOSITORY)
NAME = os.path.basename(REPOSITORY)

print("The repository this script currently works with is '{}'.\n"
      .format(REPOSITORY))

json_filename = 'merged_pull_requests_{}.json'.format(NAME)

TOKEN = get_credentials('N/A', needs_token=True)[1]

headers = {"Authorization": f"Bearer {TOKEN}"}

cursor = None

pull_requests = {}

try:
    for basename in ('master', 'main'):
        print('Searching for PRs into branch', basename)

        entries = True
        while entries:
            print('cursor:', cursor)

            if cursor is None:
                after = ''
            else:
                after = f', after:"{cursor}"'

            query = QUERY_TEMPLATE.format(owner=OWNER, repository=NAME, after=after, basename=basename)

            request = requests.post('https://api.github.com/graphql', json={'query': query}, headers=headers)

            if request.status_code != 200:
                raise Exception("Query failed")

            entries = request.json()['data']['repository']['pullRequests']['edges']

            for entry in entries:

                pr = entry['node']
                cursor = entry['cursor']

                pull_requests[str(pr['number'])] = {'milestone': pr['milestone']['title'] if pr['milestone'] else None,
                                                    'title': pr['title'],
                                                    'labels': [edge['node']['name'] for edge in pr['labels']['edges']],
                                                    'merged': pr['mergedAt'].replace('Z', ''),
                                                    'updated': pr['updatedAt'].replace('Z', ''),
                                                    'created': pr['createdAt'].replace('Z', ''),
                                                    'merge_commit': pr['mergeCommit']['oid'] if pr['mergeCommit'] else None}

finally:
    with open(json_filename, 'w') as f:
        json.dump(pull_requests, f, sort_keys=True, indent=2)
<|endoftext|>


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

"""

Usage:
    qnib_.py [options]
    qnib_.py (-h | --help)
    qnib_.py --version

Options:
    --host <str>            neo4j hostname [default: neo4j.service.consul]

General Options:
    -h --help               Show this screen.
    --version               Show version.
    --loglevel, -L=<str>    Loglevel [default: INFO]
                            (ERROR, CRITICAL, WARN, INFO, DEBUG)
    --log2stdout, -l        Log to stdout, otherwise to logfile. [default: True]
    --logfile, -f=<path>    Logfile to log to (default: <scriptname>.log)
    --cfg, -c=<path>        Configuration file.

"""

# load librarys
import logging
import os
import re
import codecs
import ast
import sys
import json
from ConfigParser import RawConfigParser, NoOptionError
from osquery import osquery

### costum libraries
from neo4jrestclient.client import GraphDatabase, Node
from neo4jrestclient.query import QuerySequence

try:
    from docopt import docopt
except ImportError:
    HAVE_DOCOPT = False
else:
    HAVE_DOCOPT = True

__author__ = 'Christian Kniep <christian()qnib.org>'
__copyright__ = 'Copyright 2015 QNIB Solutions'
__license__ = """GPL v2 License (http://www.gnu.org/licenses/old-licenses/gpl-2.0.en.html)"""

class QnibConfig(RawConfigParser):
    """ Class to abstract config and options
    """
    specials = {
        'TRUE': True,
        'FALSE': False,
        'NONE': None,
    }

    def __init__(self, opt):
        """ init """
        RawConfigParser.__init__(self)
        if opt is None:
            self._opt = {
                "--log2stdout": False,
                "--logfile": None,
                "--loglevel": "ERROR",
            }
        else:
            self._opt = opt
            self.logformat = '%(asctime)-15s %(levelname)-5s [%(module)s] %(message)s'
            self.loglevel = opt['--loglevel']
            self.log2stdout = opt['--log2stdout']
            if self.loglevel is None and opt.get('--cfg') is None:
                print "please specify loglevel (-L)"
                sys.exit(0)
            self.eval_cfg()

        self.eval_opt()
        self.set_logging()
        logging.info("SetUp of QnibConfig is done...")

    def do_get(self, section, key, default=None):
        """ Also lent from: https://github.com/jpmens/mqttwarn
            """
        try:
            val = self.get(section, key)
            if val.upper() in self.specials:
                return self.specials[val.upper()]
            return ast.literal_eval(val)
        except NoOptionError:
            return default
        except ValueError:  # e.g. %(xxx)s in string
            return val
        except:
            raise
            return val

    def config(self, section):
        ''' Convert a whole section's options (except the options specified
                explicitly below) into a dict, turning

                    [config:mqtt]
                    host = 'localhost'
                    username = None
                    list = [1, 'aaa', 'bbb', 4]

                into

                    {u'username': None, u'host': 'localhost', u'list': [1, 'aaa', 'bbb', 4]}

                Cannot use config.items() because I want each value to be
                retrieved with g() as above
            SOURCE: https://github.com/jpmens/mqttwarn
            '''

        d = None
        if self.has_section(section):
            d = dict((key, self.do_get(section, key))
                     for (key) in self.options(section) if key not in ['targets'])
        return d

    def eval_cfg(self):
        """ eval configuration which overrules the defaults
            """
        cfg_file = self._opt.get('--cfg')
        if cfg_file is not None:
            fd = codecs.open(cfg_file, 'r', encoding='utf-8')
            self.readfp(fd)
            fd.close()
            self.__dict__.update(self.config('defaults'))

    def eval_opt(self):
        """ Updates cfg according to options """

        def handle_logfile(val):
            """ transforms logfile argument
                """
            if val is None:
                logf = os.path.splitext(os.path.basename(__file__))[0]
                self.logfile = "%s.log" % logf.lower()
            else:
                self.logfile = val

        self._mapping = {
            '--logfile': lambda val: handle_logfile(val),
        }
        for key, val in self._opt.items():
            if key in self._mapping:
                if isinstance(self._mapping[key], str):
                    self.__dict__[self._mapping[key]] = val
                else:
                    self._mapping[key](val)
                break
            else:
                if val is None:
                    continue
                mat = re.match("\-\-(.*)", key)
                if mat:
                    self.__dict__[mat.group(1)] = val
                else:
                    logging.info("Could not find opt<>cfg mapping for '%s'" % key)

    def set_logging(self):
        """ sets the logging """
        self._logger = logging.getLogger()
        self._logger.setLevel(logging.DEBUG)
        if self.log2stdout:
            hdl = logging.StreamHandler()
            hdl.setLevel(self.loglevel)
            formatter = logging.Formatter(self.logformat)
            hdl.setFormatter(formatter)
            self._logger.addHandler(hdl)
        else:
            hdl = logging.FileHandler(self.logfile)
            hdl.setLevel(self.loglevel)
            formatter = logging.Formatter(self.logformat)
            hdl.setFormatter(formatter)
            self._logger.addHandler(hdl)

    def __str__(self):
        """ print human readble """
        ret = []
        for key, val in self.__dict__.items():
            if not re.match("_.*", key):
                ret.append("%-15s: %s" % (key, val))
        return "\n".join(ret)

    def __getitem__(self, item):
        """ return item from opt or __dict__
        :param item: key to lookup
        :return: value of key
        """
        if item in self.__dict__.keys():
            return self.__dict__[item]
        else:
            return self._opt[item]

class InventoryClass(object):
    """ Class to push local inventory into Neo4j
    """

    def __init__(self, cfg):
        """ Init of instance
        """
        self._cfg = cfg
        self.con_gdb()
        self._osq = osquery()

    def run(self):
        """ does sth
        """
        #self.push_pkg()
        self.push_groups()
        self.push_users()

    def push_pkg(self):
        """ fetch installed system packages and push it to Neo4j
        """
        self._pkg = self._gdb.labels.create("Pkg")
        if os.path.exists("/etc/redhat-release"):
            self.push_rpm()

    def push_rpm(self):
        """ push rpm information
        """
        query ="SELECT name, version, arch  FROM rpm_packages;"
        for item in json.loads(self._osq.setOutputMode("--json").query(query)):
            new_node = self._gdb.nodes.create(**item)
            self._pkg.add(new_node)

    def push_users(self):
        """ push users information
        """
        self._users = self._gdb.labels.create("Users")
        query ="SELECT uid,gid,username,directory,shell FROM users;"
        for item in json.loads(self._osq.setOutputMode("--json").query(query)):
	    print item
            ##new_node = self._gdb.nodes.create(**item)
            ##self._users.add(new_node)
	    query = "MATCH (g:Groups {gid:{gid}})"
	    query += " MERGE (u:Users {name:{username},directory:{directory},shell:{shell},uid:{uid}, gid:{gid}})"
	    query += " MERGE (u)-[:MEMBER]->(g)"
            new_relation = self._gdb.query(query, params=item)

    def push_groups(self):
        """ push groups information
        """
        self._groups = self._gdb.labels.create("Groups")
        query ="SELECT gid,groupname AS name FROM groups;"
        for item in json.loads(self._osq.setOutputMode("--json").query(query)):
	    print item
            new_node = self._gdb.nodes.create(**item)
            self._groups.add(new_node)

    def con_gdb(self):
        """ connect to neo4j
        """
        url = "http://%(--host)s:7474" % self._cfg
        try:
            self._gdb = GraphDatabase(url)
        except ConnectionError:
            time.sleep(3)
            self.con_gdb()

def main():
    """ main function """
    options = None
    if HAVE_DOCOPT:
        options = docopt(__doc__, version='Test Script 0.1')
    qcfg = QnibConfig(options)
    inv = InventoryClass(qcfg)
    inv.run()

if __name__ == "__main__":
    main()
<commit_msg>Uncommented pacage push<commit_after>#! /usr/bin/env python
# -*- coding: utf-8 -*-

"""

Usage:
    qnib_.py [options]
    qnib_.py (-h | --help)
    qnib_.py --version

Options:
    --host <str>            neo4j hostname [default: neo4j.service.consul]

General Options:
    -h --help               Show this screen.
    --version               Show version.
    --loglevel, -L=<str>    Loglevel [default: INFO]
                            (ERROR, CRITICAL, WARN, INFO, DEBUG)
    --log2stdout, -l        Log to stdout, otherwise to logfile. [default: True]
    --logfile, -f=<path>    Logfile to log to (default: <scriptname>.log)
    --cfg, -c=<path>        Configuration file.

"""

# load librarys
import logging
import os
import re
import codecs
import ast
import sys
import json
from ConfigParser import RawConfigParser, NoOptionError
from osquery import osquery

### costum libraries
from neo4jrestclient.client import GraphDatabase, Node
from neo4jrestclient.query import QuerySequence

try:
    from docopt import docopt
except ImportError:
    HAVE_DOCOPT = False
else:
    HAVE_DOCOPT = True

__author__ = 'Christian Kniep <christian()qnib.org>'
__copyright__ = 'Copyright 2015 QNIB Solutions'
__license__ = """GPL v2 License (http://www.gnu.org/licenses/old-licenses/gpl-2.0.en.html)"""

class QnibConfig(RawConfigParser):
    """ Class to abstract config and options
    """
    specials = {
        'TRUE': True,
        'FALSE': False,
        'NONE': None,
    }

    def __init__(self, opt):
        """ init """
        RawConfigParser.__init__(self)
        if opt is None:
            self._opt = {
                "--log2stdout": False,
                "--logfile": None,
                "--loglevel": "ERROR",
            }
        else:
            self._opt = opt
            self.logformat = '%(asctime)-15s %(levelname)-5s [%(module)s] %(message)s'
            self.loglevel = opt['--loglevel']
            self.log2stdout = opt['--log2stdout']
            if self.loglevel is None and opt.get('--cfg') is None:
                print "please specify loglevel (-L)"
                sys.exit(0)
            self.eval_cfg()

        self.eval_opt()
        self.set_logging()
        logging.info("SetUp of QnibConfig is done...")

    def do_get(self, section, key, default=None):
        """ Also lent from: https://github.com/jpmens/mqttwarn
            """
        try:
            val = self.get(section, key)
            if val.upper() in self.specials:
                return self.specials[val.upper()]
            return ast.literal_eval(val)
        except NoOptionError:
            return default
        except ValueError:  # e.g. %(xxx)s in string
            return val
        except:
            raise
            return val

    def config(self, section):
        ''' Convert a whole section's options (except the options specified
                explicitly below) into a dict, turning

                    [config:mqtt]
                    host = 'localhost'
                    username = None
                    list = [1, 'aaa', 'bbb', 4]

                into

                    {u'username': None, u'host': 'localhost', u'list': [1, 'aaa', 'bbb', 4]}

                Cannot use config.items() because I want each value to be
                retrieved with g() as above
            SOURCE: https://github.com/jpmens/mqttwarn
            '''

        d = None
        if self.has_section(section):
            d = dict((key, self.do_get(section, key))
                     for (key) in self.options(section) if key not in ['targets'])
        return d

    def eval_cfg(self):
        """ eval configuration which overrules the defaults
            """
        cfg_file = self._opt.get('--cfg')
        if cfg_file is not None:
            fd = codecs.open(cfg_file, 'r', encoding='utf-8')
            self.readfp(fd)
            fd.close()
            self.__dict__.update(self.config('defaults'))

    def eval_opt(self):
        """ Updates cfg according to options """

        def handle_logfile(val):
            """ transforms logfile argument
                """
            if val is None:
                logf = os.path.splitext(os.path.basename(__file__))[0]
                self.logfile = "%s.log" % logf.lower()
            else:
                self.logfile = val

        self._mapping = {
            '--logfile': lambda val: handle_logfile(val),
        }
        for key, val in self._opt.items():
            if key in self._mapping:
                if isinstance(self._mapping[key], str):
                    self.__dict__[self._mapping[key]] = val
                else:
                    self._mapping[key](val)
                break
            else:
                if val is None:
                    continue
                mat = re.match("\-\-(.*)", key)
                if mat:
                    self.__dict__[mat.group(1)] = val
                else:
                    logging.info("Could not find opt<>cfg mapping for '%s'" % key)

    def set_logging(self):
        """ sets the logging """
        self._logger = logging.getLogger()
        self._logger.setLevel(logging.DEBUG)
        if self.log2stdout:
            hdl = logging.StreamHandler()
            hdl.setLevel(self.loglevel)
            formatter = logging.Formatter(self.logformat)
            hdl.setFormatter(formatter)
            self._logger.addHandler(hdl)
        else:
            hdl = logging.FileHandler(self.logfile)
            hdl.setLevel(self.loglevel)
            formatter = logging.Formatter(self.logformat)
            hdl.setFormatter(formatter)
            self._logger.addHandler(hdl)

    def __str__(self):
        """ print human readble """
        ret = []
        for key, val in self.__dict__.items():
            if not re.match("_.*", key):
                ret.append("%-15s: %s" % (key, val))
        return "\n".join(ret)

    def __getitem__(self, item):
        """ return item from opt or __dict__
        :param item: key to lookup
        :return: value of key
        """
        if item in self.__dict__.keys():
            return self.__dict__[item]
        else:
            return self._opt[item]

class InventoryClass(object):
    """ Class to push local inventory into Neo4j
    """

    def __init__(self, cfg):
        """ Init of instance
        """
        self._cfg = cfg
        self.con_gdb()
        self._osq = osquery()

    def run(self):
        """ does sth
        """
        self.push_pkg()
        self.push_groups()
        self.push_users()

    def push_pkg(self):
        """ fetch installed system packages and push it to Neo4j
        """
        self._pkg = self._gdb.labels.create("Pkg")
        if os.path.exists("/etc/redhat-release"):
            self.push_rpm()

    def push_rpm(self):
        """ push rpm information
        """
        query ="SELECT name, version, arch  FROM rpm_packages;"
        for item in json.loads(self._osq.setOutputMode("--json").query(query)):
            new_node = self._gdb.nodes.create(**item)
            self._pkg.add(new_node)

    def push_users(self):
        """ push users information
        """
        self._users = self._gdb.labels.create("Users")
        query ="SELECT uid,gid,username,directory,shell FROM users;"
        for item in json.loads(self._osq.setOutputMode("--json").query(query)):
	    print item
            ##new_node = self._gdb.nodes.create(**item)
            ##self._users.add(new_node)
	    query = "MATCH (g:Groups {gid:{gid}})"
	    query += " MERGE (u:Users {name:{username},directory:{directory},shell:{shell},uid:{uid}, gid:{gid}})"
	    query += " MERGE (u)-[:MEMBER]->(g)"
            new_relation = self._gdb.query(query, params=item)

    def push_groups(self):
        """ push groups information
        """
        self._groups = self._gdb.labels.create("Groups")
        query ="SELECT gid,groupname AS name FROM groups;"
        for item in json.loads(self._osq.setOutputMode("--json").query(query)):
	    print item
            new_node = self._gdb.nodes.create(**item)
            self._groups.add(new_node)

    def con_gdb(self):
        """ connect to neo4j
        """
        url = "http://%(--host)s:7474" % self._cfg
        try:
            self._gdb = GraphDatabase(url)
        except ConnectionError:
            time.sleep(3)
            self.con_gdb()

def main():
    """ main function """
    options = None
    if HAVE_DOCOPT:
        options = docopt(__doc__, version='Test Script 0.1')
    qcfg = QnibConfig(options)
    inv = InventoryClass(qcfg)
    inv.run()

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