<commit_before># Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.
# debug_options.py
import tempfile
import warnings

import torch

from .... import ortmodule
from ... import ORTModule
from ...debug_options import DebugOptions, LogLevel

# nn.Module's in this set are considered exportable to ONNX.
# For other nn.Module's, torch.onnx.export is called to check if
# they are exportable.
_force_exportable_set = set(
    [torch.nn.Linear, torch.nn.Identity, torch.nn.modules.linear.NonDynamicallyQuantizableLinear]
)

class _IteratedORTModule(torch.nn.Module):
    """
    It's possible that a module instance is called multiple times in a single forward() call with different inputs.
    If the number of inputs or the data types are different, the exported graph for a given input set cannot be used
    for others. The _IteratedORTModule class is used to handle this case. It creates multiple ORTModule instances
    for a given nn.Module instance and uses one of them for each input set.

    NOTE that we assume that for each step run, the running order of different input sets are same.
    If it's not this case (e.g., a module is used for checkpointing so that a same input set is used twice),
    this class cannot handle it. An ideal way is to maintain a map from different input sets (maybe compute a hash)
    to ORTModule instances.
    """

    def __init__(self, module, count, log_level, save_onnx, onnx_prefix):
        super(_IteratedORTModule, self).__init__()
        assert count > 1
        self._count = count
        self._it = count - 1
        self._ortmodules = []
        for idx in range(count):
            self._ortmodules.append(
                ORTModule(
                    module,
                    debug_options=DebugOptions(
                        log_level=log_level, save_onnx=save_onnx, onnx_prefix=onnx_prefix + "_it" + str(idx)
                    ),
                )
            )

    def forward(self, *inputs, **kwargs):
        self._it = (self._it + 1) % self._count
        return self._ortmodules[self._it](*inputs, **kwargs)

class HierarchicalORTModule(torch.nn.Module):
    """
    Recursively wraps submodules of `module` as ORTModule whenever possible
    Similarly to ORTModule, the actual wrapping happens in its first `forward` call during Pytorch-to-ONNX export.
    Supported computation is delegated to ONNX Runtime and unsupported computation is still done by PyTorch.
    Args:
        module (torch.nn.Module): User's PyTorch module that HierarchicalORTModule specializes.

    Example::

        import torch
        from torch.utils.checkpoint import checkpoint
        from onnxruntime.training.ortmodule.experimental.hierarchical_ortmodule import HierarchicalORTModule

        class Foo(torch.nn.Module):
            def __init__(self):
                super(Foo, self).__init__()
                self.l1 = torch.nn.Linear(2, 2)
                self.l2 = torch.nn.Linear(2, 2)
                self.l3 = torch.nn.Linear(2, 2)

            def forward(self, x):
                def custom():
                    def custom_forward(x_):
                        return self.l2(x_)
                    return custom_forward
                z = self.l1(checkpoint(custom(), self.l3(x)))
                return z

        x = torch.rand(2)
        m = HierarchicalORTModule(Foo())
        y = m(x)

    """

    def __init__(self, module, debug_options=None):
        self._initialized = False
        super(HierarchicalORTModule, self).__init__()
        self._original_module = module
        self._log_level = debug_options.logging.log_level if debug_options else LogLevel.ERROR
        self._save_onnx = debug_options.save_onnx_models.save if debug_options else False
        self._name_prefix = debug_options.save_onnx_models.name_prefix if debug_options else ""

    def _initialize(self, *args, **kwargs):
        handle_pool = []
        module_arg_pool = {}

        # A forward pre-hook to record inputs for each nn.Module.
        def record_args(module, args):
            if module in module_arg_pool:
                module_arg_pool[module].append(args)
            else:
                module_arg_pool[module] = [args]

        # Recursively hook "record_args" to module and all its sub-modules.
        # The function "record_args" records the inputs for each nn.Module and later we will try exporting
        # those nn.Module's with their recorded inputs.
        # NOTE that if a module is not called from forward(), it will fail to be captured by this hook.
        def recursive_hook(module):
            # We cannot skip module in whitelist because it's possible that a module is called multiple times
            # so that we still need to know the number of different input sets and use _IteratedORTModule to handle it.
            handle_pool.append(module.register_forward_pre_hook(record_args))
            for _, sub_module in module._modules.items():
                if isinstance(sub_module, torch.nn.ModuleList):
                    for _, sub_module_item in sub_module._modules.items():
                        recursive_hook(sub_module_item)
                else:
                    recursive_hook(sub_module)

        exportable_list = {}

        # Fill "exportable_list". exportable_list[module] = True means
        # "module" can be wrapped as ORTModule. Otherwise, "module" is
        # not exportable to ONNX.
        def check_exportable(module):
            def try_export(module, args):
                try:
                    with tempfile.NamedTemporaryFile(prefix="sub-module") as temp, torch.no_grad():
                        torch.onnx.export(
                            module,
                            args,
                            temp,
                            opset_version=ortmodule.ONNX_OPSET_VERSION,
                            do_constant_folding=False,
                            export_params=False,
                            keep_initializers_as_inputs=True,
                            training=torch.onnx.TrainingMode.TRAINING,
                        )
                except Exception as e:
                    if self._log_level <= LogLevel.WARNING:
                        warnings.warn(
                            f"Failed to export module with type {type(module).__name__}. Error message: {str(e)}",
                            UserWarning,
                        )
                    return False
                return True

            if type(module) in _force_exportable_set:
                exportable_list[module] = True
                return

            # It's possible that the model runs a module by calling some other function instead of forward()
            # so that the module is not captured by the forward pre-hook. In this case, we will treat it as
            # not exportable for now.
            module_exportable = module in module_arg_pool
            if module_exportable:
                for args in module_arg_pool[module]:
                    if not try_export(module, args):
                        module_exportable = False
                        break
            elif self._log_level <= LogLevel.WARNING:
                warnings.warn(
                    f"Module with type {type(module).__name__} is not exportable because it's not in module_arg_pool.",
                    UserWarning,
                )

            exportable_list[module] = module_exportable
            if module_exportable:
                return

            sub_module_dict = module._modules
            if not sub_module_dict:
                # No sub-module exists, so this module is a leaf
                return

            for _, sub_module in sub_module_dict.items():
                if isinstance(sub_module, torch.nn.ModuleList):
                    for _, sub_module_item in sub_module._modules.items():
                        check_exportable(sub_module_item)
                else:
                    check_exportable(sub_module)

        # Add a hook to record forward's input for all modules.
        recursive_hook(self._original_module)

        # Run forward with actual input to record all possible
        # inputs for all invoked modules.
        with torch.no_grad():
            _ = self._original_module(*args, **kwargs)

        # We already have "supported_modules" so
        # we no longer need those hooks in forward pass.
        for handle in handle_pool:
            handle.remove()

        # Try exporter on all module-input pairs. If a module can be exported with
        # all its recorded inputs, then it's exporable.
        check_exportable(self._original_module)

        # A naive way of determining if ORT can run nn.Module
        def is_supported(module):
            return module in exportable_list and exportable_list[module]

        # Top-down wrapper to replace nn.Module's with ORTModule.
        # Note that using bottom-up wrapper may lead to much
        # ORTModule instances and each ORTModule owns a much smaller graph.
        def recursive_wrap(module, save_onnx=False, onnx_prefix=""):
            sub_module_dict = module._modules
            for name, sub_module in sub_module_dict.items():
                new_prefix = onnx_prefix + "_" + name
                if isinstance(sub_module, torch.nn.ModuleList):
                    # We encounter a list of sub-modules.
                    # Let's wrap them one-by-one.
                    idx = 0
                    for item_name, sub_module_item in sub_module._modules.items():
                        # Avoid saving too many graphs.
                        new_save_onnx = save_onnx and idx == 0
                        sub_new_prefix = new_prefix + "_" + item_name
                        if is_supported(sub_module_item):
                            if sub_module_item in module_arg_pool and len(module_arg_pool[sub_module_item]) > 1:
                                sub_module._modules[item_name] = _IteratedORTModule(
                                    sub_module_item,
                                    len(module_arg_pool[sub_module_item]),
                                    self._log_level,
                                    new_save_onnx,
                                    sub_new_prefix,
                                )
                            else:
                                sub_module._modules[item_name] = ORTModule(
                                    sub_module_item,
                                    debug_options=DebugOptions(
                                        log_level=self._log_level, save_onnx=new_save_onnx, onnx_prefix=sub_new_prefix
                                    ),
                                )
                        else:
                            recursive_wrap(sub_module_item, new_save_onnx, sub_new_prefix)
                        idx += 1
                else:
                    if is_supported(sub_module):
                        # Just wrap it as ORTModule when possible.
                        if sub_module in module_arg_pool and len(module_arg_pool[sub_module]) > 1:
                            sub_module_dict[name] = _IteratedORTModule(
                                sub_module, len(module_arg_pool[sub_module]), self._log_level, save_onnx, new_prefix
                            )
                        else:
                            sub_module_dict[name] = ORTModule(
                                sub_module,
                                debug_options=DebugOptions(
                                    log_level=self._log_level, save_onnx=save_onnx, onnx_prefix=new_prefix
                                ),
                            )
                    else:
                        # This sub-module is not exportable to ONNX
                        # Let's check its sub-modules.
                        recursive_wrap(sub_module, save_onnx, new_prefix)

        if is_supported(self._original_module):
            self._original_module = ORTModule(
                self._original_module,
                debug_options=DebugOptions(
                    log_level=self._log_level, save_onnx=self._save_onnx, onnx_prefix=self._name_prefix
                ),
            )
        else:
            recursive_wrap(self._original_module, self._save_onnx, self._name_prefix)
        if self._log_level <= LogLevel.WARNING:
            warnings.warn(
                f"Wrapped module: {str(self._original_module)}.",
                UserWarning,
            )
        self._initialized = True

    def forward(self, *inputs, **kwargs):
        if not self._initialized:
            self._initialize(*inputs, **kwargs)
        # forward can be run only after initialization is done.
        assert self._initialized
        return self._original_module(*inputs, **kwargs)
<commit_msg>Fix a PoliCheck finding in _hierarchical_ortmodule.py(#13462)<commit_after># Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.
# debug_options.py
import tempfile
import warnings

import torch

from .... import ortmodule
from ... import ORTModule
from ...debug_options import DebugOptions, LogLevel

# nn.Module's in this set are considered exportable to ONNX.
# For other nn.Module's, torch.onnx.export is called to check if
# they are exportable.
_force_exportable_set = set(
    [torch.nn.Linear, torch.nn.Identity, torch.nn.modules.linear.NonDynamicallyQuantizableLinear]
)

class _IteratedORTModule(torch.nn.Module):
    """
    It's possible that a module instance is called multiple times in a single forward() call with different inputs.
    If the number of inputs or the data types are different, the exported graph for a given input set cannot be used
    for others. The _IteratedORTModule class is used to handle this case. It creates multiple ORTModule instances
    for a given nn.Module instance and uses one of them for each input set.

    NOTE that we assume that for each step run, the running order of different input sets are same.
    If it's not this case (e.g., a module is used for checkpointing so that a same input set is used twice),
    this class cannot handle it. An ideal way is to maintain a map from different input sets (maybe compute a hash)
    to ORTModule instances.
    """

    def __init__(self, module, count, log_level, save_onnx, onnx_prefix):
        super(_IteratedORTModule, self).__init__()
        assert count > 1
        self._count = count
        self._it = count - 1
        self._ortmodules = []
        for idx in range(count):
            self._ortmodules.append(
                ORTModule(
                    module,
                    debug_options=DebugOptions(
                        log_level=log_level, save_onnx=save_onnx, onnx_prefix=onnx_prefix + "_it" + str(idx)
                    ),
                )
            )

    def forward(self, *inputs, **kwargs):
        self._it = (self._it + 1) % self._count
        return self._ortmodules[self._it](*inputs, **kwargs)

class HierarchicalORTModule(torch.nn.Module):
    """
    Recursively wraps submodules of `module` as ORTModule whenever possible
    Similarly to ORTModule, the actual wrapping happens in its first `forward` call during Pytorch-to-ONNX export.
    Supported computation is delegated to ONNX Runtime and unsupported computation is still done by PyTorch.
    Args:
        module (torch.nn.Module): User's PyTorch module that HierarchicalORTModule specializes.

    Example::

        import torch
        from torch.utils.checkpoint import checkpoint
        from onnxruntime.training.ortmodule.experimental.hierarchical_ortmodule import HierarchicalORTModule

        class Foo(torch.nn.Module):
            def __init__(self):
                super(Foo, self).__init__()
                self.l1 = torch.nn.Linear(2, 2)
                self.l2 = torch.nn.Linear(2, 2)
                self.l3 = torch.nn.Linear(2, 2)

            def forward(self, x):
                def custom():
                    def custom_forward(x_):
                        return self.l2(x_)
                    return custom_forward
                z = self.l1(checkpoint(custom(), self.l3(x)))
                return z

        x = torch.rand(2)
        m = HierarchicalORTModule(Foo())
        y = m(x)

    """

    def __init__(self, module, debug_options=None):
        self._initialized = False
        super(HierarchicalORTModule, self).__init__()
        self._original_module = module
        self._log_level = debug_options.logging.log_level if debug_options else LogLevel.ERROR
        self._save_onnx = debug_options.save_onnx_models.save if debug_options else False
        self._name_prefix = debug_options.save_onnx_models.name_prefix if debug_options else ""

    def _initialize(self, *args, **kwargs):
        handle_pool = []
        module_arg_pool = {}

        # A forward pre-hook to record inputs for each nn.Module.
        def record_args(module, args):
            if module in module_arg_pool:
                module_arg_pool[module].append(args)
            else:
                module_arg_pool[module] = [args]

        # Recursively hook "record_args" to module and all its sub-modules.
        # The function "record_args" records the inputs for each nn.Module and later we will try exporting
        # those nn.Module's with their recorded inputs.
        # NOTE that if a module is not called from forward(), it will fail to be captured by this hook.
        def recursive_hook(module):
            # We cannot skip module in allowlist because it's possible that a module is called multiple times
            # so that we still need to know the number of different input sets and use _IteratedORTModule to handle it.
            handle_pool.append(module.register_forward_pre_hook(record_args))
            for _, sub_module in module._modules.items():
                if isinstance(sub_module, torch.nn.ModuleList):
                    for _, sub_module_item in sub_module._modules.items():
                        recursive_hook(sub_module_item)
                else:
                    recursive_hook(sub_module)

        exportable_list = {}

        # Fill "exportable_list". exportable_list[module] = True means
        # "module" can be wrapped as ORTModule. Otherwise, "module" is
        # not exportable to ONNX.
        def check_exportable(module):
            def try_export(module, args):
                try:
                    with tempfile.NamedTemporaryFile(prefix="sub-module") as temp, torch.no_grad():
                        torch.onnx.export(
                            module,
                            args,
                            temp,
                            opset_version=ortmodule.ONNX_OPSET_VERSION,
                            do_constant_folding=False,
                            export_params=False,
                            keep_initializers_as_inputs=True,
                            training=torch.onnx.TrainingMode.TRAINING,
                        )
                except Exception as e:
                    if self._log_level <= LogLevel.WARNING:
                        warnings.warn(
                            f"Failed to export module with type {type(module).__name__}. Error message: {str(e)}",
                            UserWarning,
                        )
                    return False
                return True

            if type(module) in _force_exportable_set:
                exportable_list[module] = True
                return

            # It's possible that the model runs a module by calling some other function instead of forward()
            # so that the module is not captured by the forward pre-hook. In this case, we will treat it as
            # not exportable for now.
            module_exportable = module in module_arg_pool
            if module_exportable:
                for args in module_arg_pool[module]:
                    if not try_export(module, args):
                        module_exportable = False
                        break
            elif self._log_level <= LogLevel.WARNING:
                warnings.warn(
                    f"Module with type {type(module).__name__} is not exportable because it's not in module_arg_pool.",
                    UserWarning,
                )

            exportable_list[module] = module_exportable
            if module_exportable:
                return

            sub_module_dict = module._modules
            if not sub_module_dict:
                # No sub-module exists, so this module is a leaf
                return

            for _, sub_module in sub_module_dict.items():
                if isinstance(sub_module, torch.nn.ModuleList):
                    for _, sub_module_item in sub_module._modules.items():
                        check_exportable(sub_module_item)
                else:
                    check_exportable(sub_module)

        # Add a hook to record forward's input for all modules.
        recursive_hook(self._original_module)

        # Run forward with actual input to record all possible
        # inputs for all invoked modules.
        with torch.no_grad():
            _ = self._original_module(*args, **kwargs)

        # We already have "supported_modules" so
        # we no longer need those hooks in forward pass.
        for handle in handle_pool:
            handle.remove()

        # Try exporter on all module-input pairs. If a module can be exported with
        # all its recorded inputs, then it's exporable.
        check_exportable(self._original_module)

        # A naive way of determining if ORT can run nn.Module
        def is_supported(module):
            return module in exportable_list and exportable_list[module]

        # Top-down wrapper to replace nn.Module's with ORTModule.
        # Note that using bottom-up wrapper may lead to much
        # ORTModule instances and each ORTModule owns a much smaller graph.
        def recursive_wrap(module, save_onnx=False, onnx_prefix=""):
            sub_module_dict = module._modules
            for name, sub_module in sub_module_dict.items():
                new_prefix = onnx_prefix + "_" + name
                if isinstance(sub_module, torch.nn.ModuleList):
                    # We encounter a list of sub-modules.
                    # Let's wrap them one-by-one.
                    idx = 0
                    for item_name, sub_module_item in sub_module._modules.items():
                        # Avoid saving too many graphs.
                        new_save_onnx = save_onnx and idx == 0
                        sub_new_prefix = new_prefix + "_" + item_name
                        if is_supported(sub_module_item):
                            if sub_module_item in module_arg_pool and len(module_arg_pool[sub_module_item]) > 1:
                                sub_module._modules[item_name] = _IteratedORTModule(
                                    sub_module_item,
                                    len(module_arg_pool[sub_module_item]),
                                    self._log_level,
                                    new_save_onnx,
                                    sub_new_prefix,
                                )
                            else:
                                sub_module._modules[item_name] = ORTModule(
                                    sub_module_item,
                                    debug_options=DebugOptions(
                                        log_level=self._log_level, save_onnx=new_save_onnx, onnx_prefix=sub_new_prefix
                                    ),
                                )
                        else:
                            recursive_wrap(sub_module_item, new_save_onnx, sub_new_prefix)
                        idx += 1
                else:
                    if is_supported(sub_module):
                        # Just wrap it as ORTModule when possible.
                        if sub_module in module_arg_pool and len(module_arg_pool[sub_module]) > 1:
                            sub_module_dict[name] = _IteratedORTModule(
                                sub_module, len(module_arg_pool[sub_module]), self._log_level, save_onnx, new_prefix
                            )
                        else:
                            sub_module_dict[name] = ORTModule(
                                sub_module,
                                debug_options=DebugOptions(
                                    log_level=self._log_level, save_onnx=save_onnx, onnx_prefix=new_prefix
                                ),
                            )
                    else:
                        # This sub-module is not exportable to ONNX
                        # Let's check its sub-modules.
                        recursive_wrap(sub_module, save_onnx, new_prefix)

        if is_supported(self._original_module):
            self._original_module = ORTModule(
                self._original_module,
                debug_options=DebugOptions(
                    log_level=self._log_level, save_onnx=self._save_onnx, onnx_prefix=self._name_prefix
                ),
            )
        else:
            recursive_wrap(self._original_module, self._save_onnx, self._name_prefix)
        if self._log_level <= LogLevel.WARNING:
            warnings.warn(
                f"Wrapped module: {str(self._original_module)}.",
                UserWarning,
            )
        self._initialized = True

    def forward(self, *inputs, **kwargs):
        if not self._initialized:
            self._initialize(*inputs, **kwargs)
        # forward can be run only after initialization is done.
        assert self._initialized
        return self._original_module(*inputs, **kwargs)
<|endoftext|>


<commit_before>#  This file is part of the myhdl library, a Python package for using
#  Python as a Hardware Description Language.
#
#  Copyright (C) 2003 Jan Decaluwe
#
#  The myhdl library is free software; you can redistribute it and/or
#  modify it under the terms of the GNU Lesser General Public License as
#  published by the Free Software Foundation; either version 2.1 of the
#  License, or (at your option) any later version.
#
#  This library 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
#  Lesser General Public License for more details.
#
#  You should have received a copy of the GNU Lesser General Public
#  License along with this library; if not, write to the Free Software
#  Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA

""" Module that implements enum.

"""

__author__ = "Jan Decaluwe <jan@jandecaluwe.com>"
__revision__ = "$Revision$"
__date__ = "$Date$"

from types import StringType

from myhdl._bin import bin

def enum(*args, **kwargs):

    # args = args
    encoding = kwargs.get('encoding', 'binary')
    argdict = {}
    codedict = {}
    if encoding == "binary":
        nrbits = len(bin(len(args)))
    elif encoding in ("one_hot", "one_cold"):
        nrbits = len(args)
    else:
        raise ValueError("Unsupported enum encoding: %s \n    Supported encodings:" + \
                         "    'binary', 'one_hot', 'one_cold'" % encoding)

    i = 0
    for arg in args:
        if type(arg) is not StringType:
            raise TypeError
        if codedict.has_key(arg):
            raise ValueError("enum literals should be unique")
        argdict[i] = arg
        if encoding == "binary":
            code = bin(i, nrbits)
        elif encoding == "one_hot":
            code = bin(1<<i, nrbits)
        else: # one_cold
            code = bin(~(1<<i), nrbits)
        if len(code) > nrbits:
            code = code[-nrbits:]
        codedict[arg] = code
        i += 1

    class EnumItem(object):
        def __init__(self, index, arg):
            self._index = index
            self._val = codedict[arg]
            self._nrbits = nrbits
            self._nritems = len(args)
        def __repr__(self):
            return argdict[self._index]
        def __hex__(self):
            return hex(int(self._val, 2))
        __str__ = __repr__
        def _toVerilog(self, dontcare=False):
            val = self._val
            if dontcare:
                if encoding == "one_hot":
                    val = val.replace('0', '?')
                elif encoding == "one_cold":
                    val = val.replace('1', '?')
            return "%d'b%s" % (self._nrbits, val)

    class Enum(object):
        def __init__(self):
            for index, slot in enumerate(args):
                self.__dict__[slot] = EnumItem(index, slot)
            self.__dict__['_nrbits'] = nrbits
        def __setattr__(self, attr, val):
            raise AttributeError("Cannot assign to enum attributes")
        def __len__(self):
            return len(args)
        def __repr__(self):
            return "<Enum: %s>" % ", ".join(args)
        __str__ = __repr__

    return Enum()

<commit_msg>correct calculation of nrbits<commit_after>#  This file is part of the myhdl library, a Python package for using
#  Python as a Hardware Description Language.
#
#  Copyright (C) 2003 Jan Decaluwe
#
#  The myhdl library is free software; you can redistribute it and/or
#  modify it under the terms of the GNU Lesser General Public License as
#  published by the Free Software Foundation; either version 2.1 of the
#  License, or (at your option) any later version.
#
#  This library 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
#  Lesser General Public License for more details.
#
#  You should have received a copy of the GNU Lesser General Public
#  License along with this library; if not, write to the Free Software
#  Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA

""" Module that implements enum.

"""

__author__ = "Jan Decaluwe <jan@jandecaluwe.com>"
__revision__ = "$Revision$"
__date__ = "$Date$"

from types import StringType

from myhdl._bin import bin

def enum(*args, **kwargs):

    # args = args
    encoding = kwargs.get('encoding', 'binary')
    argdict = {}
    codedict = {}
    if encoding == "binary":
        nrbits = len(bin(len(args)-1))
    elif encoding in ("one_hot", "one_cold"):
        nrbits = len(args)
    else:
        raise ValueError("Unsupported enum encoding: %s \n    Supported encodings:" + \
                         "    'binary', 'one_hot', 'one_cold'" % encoding)

    i = 0
    for arg in args:
        if type(arg) is not StringType:
            raise TypeError
        if codedict.has_key(arg):
            raise ValueError("enum literals should be unique")
        argdict[i] = arg
        if encoding == "binary":
            code = bin(i, nrbits)
        elif encoding == "one_hot":
            code = bin(1<<i, nrbits)
        else: # one_cold
            code = bin(~(1<<i), nrbits)
        if len(code) > nrbits:
            code = code[-nrbits:]
        codedict[arg] = code
        i += 1

    class EnumItem(object):
        def __init__(self, index, arg):
            self._index = index
            self._val = codedict[arg]
            self._nrbits = nrbits
            self._nritems = len(args)
        def __repr__(self):
            return argdict[self._index]
        def __hex__(self):
            return hex(int(self._val, 2))
        __str__ = __repr__
        def _toVerilog(self, dontcare=False):
            val = self._val
            if dontcare:
                if encoding == "one_hot":
                    val = val.replace('0', '?')
                elif encoding == "one_cold":
                    val = val.replace('1', '?')
            return "%d'b%s" % (self._nrbits, val)

    class Enum(object):
        def __init__(self):
            for index, slot in enumerate(args):
                self.__dict__[slot] = EnumItem(index, slot)
            self.__dict__['_nrbits'] = nrbits
        def __setattr__(self, attr, val):
            raise AttributeError("Cannot assign to enum attributes")
        def __len__(self):
            return len(args)
        def __repr__(self):
            return "<Enum: %s>" % ", ".join(args)
        __str__ = __repr__

    return Enum()

<|endoftext|>


<commit_before>import numpy as np
import pandas as pd
from scipy.io import loadmat

from mypy.utils import group, extend_slice, find_range

# TODOs:
# - [ ] create_middle_events should be made more universal

def din_dataframe(eeg):
	'''
	Turns DIN1, DIN2, ... channels of an egi file
	into a dataframe with latency as index and
	DIN numbers as columns. Each row of the dataframe is
	therefore a reconstruction of the LPT digital signal

	Parameters
	----------
	eeg - mne Raw created with `read_raw_egi`

	Returns
	-------
	df : pandas DataFrame
		Dataframe with latency as index and DIN numbers as
		columns. Values of the dataframe are booleans.
		If df.loc[120, 4] is True for example - then at latency
		120 DIN4 is active.
		df.loc[250, :] on the other hand is the full binary re-
		presentation of the LPT signal at latency 250.
	'''

	 # find DIN-channels
	din_chans = [(ch, i) for i, ch in enumerate(eeg.info['ch_names'])
				 if ch.startswith('D')]

	# create a mapping: din channel -> latencies
	lat = dict()
	for ch, i in din_chans:
		digits = "".join([c for c in ch if c.isdigit()])
		lat[int(digits)] = np.where(eeg._data[i,:] > 0)[0]

	# regroup into an array:
	din_types = np.array(sorted(lat.keys()))
	evt = [[[l, t] for l in lat[t]] for t in din_types]
	evt = np.concatenate(evt, axis=0)

	# sort and add fake row to unify loop
	ind = evt[:,0].argsort(axis=0)
	evt = evt[ind,:]
	evt = np.concatenate([evt, evt[-1,:][np.newaxis, :]+100])

	# find events occuring in the same time sample
	current_row, current_latency = 0, evt[0,0]
	start_seq = 0
	latency = list()
	ifdinactive = np.zeros([evt.shape[0], len(din_types)],
		dtype='bool')

	for r in range(1, evt.shape[0]):
		if not evt[r,0] == current_latency:
			din_ind = [np.where(din_types == x)[0][0]
						for x in evt[start_seq:r, 1]]
			ifdinactive[current_row, din_ind] = True
			latency.append(current_latency)

			current_latency = evt[r,0]
			current_row += 1
			start_seq = r

	# return as a dataframe
	return pd.DataFrame(index=latency, columns=din_types,
						data=ifdinactive[:current_row, :])

def get_events_from_din(eeg):
	'''Turn DIN1, DIN2, ... etc. event channels to mne-python event array.

	Parameters
	----------
	eeg : Raw
		Instance of mne-python raw file.

	Returns
	-------
	events : array
		Numpy array of size (n_events, 3). Complies to mne-python
		convention of event arrays.
	'''
	# assert isinstance(mne.Raw) etc?
	df = din_dataframe(eeg)
	n_evnt = df.shape[0]
	events = np.zeros([n_evnt, 3], dtype='int')
	dins = np.tile(df.columns.values, [n_evnt, 1])
	events[:, 2] = np.sum(df.values * dins, axis=1)
	events[:, 0] = df.index.values
	return events

def remove_din_channels(eeg):
	'''Remove channels starting with 'D' like 'DIN2' or 'DI16'.'''
	# pick non-din channels:
	non_din_chans = [ch for ch in eeg.info['ch_names']
						if not ch.startswith('D')
						and not ch == 'STI 014']
	eeg.pick_channels(non_din_chans)

def reject_events_in_bad_segments(events, bad_segments, around_event=(-10,10),
								  remove_types=None):
	'''Remove events that coincide with bad segments.

	Parameters
	----------
	events 		 - mne events array
	bad_segments - N by 2 matrix with bad segment info
				   each row is a bad segment, first column is the
				   bad segment onset and the second column is the
				   bad segment offset (in samples)
	around_event - (lowerlim, higherlim) box limits around
				   each event. If a bad segment overlaps
				   with the box, the event is rejected.
				   Box limits are in samples.
				   (-500, 1000) means 500 samples before
				   up to 1000 samples after, but one can
				   create box that does not contain the
				   event like (250, 500)
	remove_types - list or numpy array of event types (int)
				   that should be checked for removal

	Returns
	-------
	events - corrected events array
	'''
	if remove_types is None:
		test_events = np.arange(events.shape[0])
	else:
		test_events = np.vstack([events[:,2] == x for x in remove_types])
		test_events = np.where(np.any(test_events, axis=0))

	ev = events[test_events, 0]
	ev = np.vstack([ev + x for x in around_event]).T
	remove = np.zeros(ev.shape[0], dtype='bool')
	for ii in range(ev.shape[0]):
		if np.any(np.logical_and(ev[ii,0] >= bad_segments[:,0],
			ev[ii,0] < bad_segments[:,1])):
			remove[ii] = True
			continue
		if np.any(np.logical_and(ev[ii,1] > bad_segments[:,0],
			ev[ii,1] <= bad_segments[:,1])):
			remove[ii] = True
			continue
		if np.any(np.logical_and(bad_segments[:,0] >= ev[ii,0],
			bad_segments[:,0] < ev[ii,1])):
			remove[ii] = True
	# remove events
	remove_ind = test_events[remove]
	if remove_ind.shape[0] > 0:
		events = np.delete(events, remove_ind, axis=0)
	return events

def extend_bads(rej_win, extend, copy=True):
	'''Extend rejection periods by specfied number of samples.'''
	if copy:
		rej_win = rej_win.copy()
	n_samples = rej_win.shape[1]
	bad_inds = np.any(rej_win, axis=1).nonzero()[0]
	for bad in bad_inds:
		slices = group(rej_win[bad, :], return_slice=True)
		for slc in slices:
			slc = extend_slice(slc, extend, n_samples)
			rej_win[bad, slc] = True
	return rej_win

def apply_artifacts_to_tfr(tfr, artif, orig_time, fillval=np.nan):
	'''Fill tfr data with `fillval` according to `artif`.
	Operates in-place.'''
	n_items, n_samples = artif.shape
	bad_inds = np.any(artif, axis=1).nonzero()[0]
	for bad in bad_inds:
		limits = group(artif[bad, :])
		for lim in limits:
			slc = find_range(tfr.times, list(orig_time[lim]))
			tfr.data[bad, :, :, slc] = fillval
	return tfr

# - [ ] make sure this is universal
def create_middle_events(events, min_time=14., sfreq=250):
	'''Put raven events in the trial centre,
	deleting events that begin trials shorter than
	min_time.'''
	ep_start = np.logical_or(events[:,2] == 1, events[:,2] == 0)
	ep_end = np.logical_or(events[:,2] == 10, events[:,2] == 11)
	ep_len = events[ep_end, 0] - events[ep_start, 0]
	long_enough = ep_len > int(min_time * sfreq)
	new_events = events[ep_start,:]
	new_events[:,0] += np.round(ep_len / 2).astype('int')
	return new_events[long_enough, :]

# TODO:
# - [ ] adapt this to use numpy structured datasets (?)
def read_set_events(filename, ignore_fields=None):
	'''Open set file, read epoch events and turn them into a dataframe.

	Parameters
	----------
	filename: str
		Name of the set file to read (can be absolute or relative path)
	ignore_fields: list of str | None
		Epoch event fields to ignore (these fields are note included in the df)

	Returns
	-------
	df: pandas.DataFrame
		Events read into dataframe
	'''
	EEG = loadmat(filename, uint16_codec='latin1',
				  struct_as_record=False, squeeze_me=True)['EEG']
	flds = [f for f in dir(EEG.event[0]) if not f.startswith('_')]
	events = EEG.event
	df_dict = dict()
	for f in flds:
		df = df_dict[f] = [ev.__getattribute__(f) for ev in events]
	df = pd.DataFrame(df_dict)

	# reorder columns:
	take_fields = ['epoch', 'type']
	ignore_fields = list() if ignore_fields is None else ignore_fields
	take_fields.extend([col for col in df.columns if not
					 (col in take_fields or col in ignore_fields)])
	return df.loc[:, take_fields]

def read_rej(fname, sfreq, bad_types=['reject']):
	from mne import Annotations

	rej = pd.read_table(fname)
	rej.loc[:, ['start', 'end']] = rej.loc[:, ['start', 'end']] / sfreq
	onset, duration = rej.start.values, (rej.end - rej.start).values
	description = ['BAD ' + val if val in bad_types else val
				   for val in rej.type.tolist()]
	return Annotations(onset, duration, description)

# - [ ] develop this function a little better
# - [ ] check what is used by raw.reject_bads() when no inds given
def mark_reject_peak2peak(raw, reject={'eeg': 23e-5}, window_length=1.,
						  label='bad p2p'):
	import mne
	from mne.utils import _reject_data_segments
	_, inds = _reject_data_segments(raw._data, reject, {'eeg': 0.},
									window_length, raw.info, 0.5)

	# turn inds to time, join
	time_segments = np.array(inds) / raw.info['sfreq']
	time_segments = join_segments(time_segments)

	segment_duration = np.diff(time_segments, axis=-1).ravel()
	return mne.Annotations(time_segments[:, 0], segment_duration, label)

# - [ ] maybe move to utils
def join_segments(time_segments):
	from mypy.utils import group

	# check which should be joined
	join_segments = (time_segments[:-1, 1] - time_segments[1:, 0]) == 0
	segment_groups = group(join_segments)
	segment_groups[:, 1] += 1

	# join segments
	final_segments = np.vstack([time_segments[segment_groups[:, 0], 0],
	                            time_segments[segment_groups[:, 1], 1]]).T

	# we missed single-segments, search for them
	prev = 0
	missed_segments = list()
	for row in segment_groups:
		if not row[0] == prev:
			missed_segments.extend(list(range(prev + 1, row[0])))
		prev = row[1]
	if prev < len(time_segments):
		missed_segments.extend(list(range(prev + 1, len(time_segments))))

	# append missed single segments to multi-segments and sort
	final_segments = np.vstack(
		[final_segments, time_segments[missed_segments, :]])
	final_segments = final_segments[final_segments[:, 0].argsort(), :]

	return final_segments

def get_dropped_epoch_index(epochs):
	'''get indices of dropped epochs'''
	current_epoch = 0
	removed_epochs = list()

	for info in epochs.drop_log:
	    if 'IGNORED' not in info:
	        if 'USER' in info:
	            removed_epochs.append(current_epoch)
	        current_epoch += 1

	return removed_epochs

def align_events(events1, events2):
	'''Return indices for each of the two event trains that give best match.'''
	common_events = list(set(events1).intersection(set(events2)))
	ind1, ind2 = [np.where(np.in1d(events, common_events))[0]
		      for events in [events1, events2]]

	len1, len2 = len(ind1), len(ind2)
	if len1 == len2:
		return ind1, ind2
	elif len1 > len2:
		longer, shorter = events1[ind1], events2[ind2]
	else:
		longer, shorter = events2[ind2], events1[ind1]

	# roll shorter along longer
	lng_len, shrt_len = len(longer), len(shorter)
	n_steps = lng_len - shrt_len + 1
	scores = np.zeros(n_steps, dtype='int')
	for offset in range(n_steps):
		scores[offset] = np.sum(longer[offset:offset + shrt_len] == shorter)
	best_offset = scores.argmax()

	if len1 > len2:
		ind1 = ind1[best_offset:best_offset + shrt_len]
	else:
		ind2 = ind2[best_offset:best_offset + shrt_len]
	return ind1, ind2
<commit_msg>ENH: add first version of AutoMark<commit_after>import numpy as np
import pandas as pd
from scipy.io import loadmat

from mypy.utils import group, extend_slice, find_range

# TODOs:
# - [ ] create_middle_events should be made more universal

def din_dataframe(eeg):
	'''
	Turns DIN1, DIN2, ... channels of an egi file
	into a dataframe with latency as index and
	DIN numbers as columns. Each row of the dataframe is
	therefore a reconstruction of the LPT digital signal

	Parameters
	----------
	eeg - mne Raw created with `read_raw_egi`

	Returns
	-------
	df : pandas DataFrame
		Dataframe with latency as index and DIN numbers as
		columns. Values of the dataframe are booleans.
		If df.loc[120, 4] is True for example - then at latency
		120 DIN4 is active.
		df.loc[250, :] on the other hand is the full binary re-
		presentation of the LPT signal at latency 250.
	'''

	 # find DIN-channels
	din_chans = [(ch, i) for i, ch in enumerate(eeg.info['ch_names'])
				 if ch.startswith('D')]

	# create a mapping: din channel -> latencies
	lat = dict()
	for ch, i in din_chans:
		digits = "".join([c for c in ch if c.isdigit()])
		lat[int(digits)] = np.where(eeg._data[i,:] > 0)[0]

	# regroup into an array:
	din_types = np.array(sorted(lat.keys()))
	evt = [[[l, t] for l in lat[t]] for t in din_types]
	evt = np.concatenate(evt, axis=0)

	# sort and add fake row to unify loop
	ind = evt[:,0].argsort(axis=0)
	evt = evt[ind,:]
	evt = np.concatenate([evt, evt[-1,:][np.newaxis, :]+100])

	# find events occuring in the same time sample
	current_row, current_latency = 0, evt[0,0]
	start_seq = 0
	latency = list()
	ifdinactive = np.zeros([evt.shape[0], len(din_types)],
		dtype='bool')

	for r in range(1, evt.shape[0]):
		if not evt[r,0] == current_latency:
			din_ind = [np.where(din_types == x)[0][0]
						for x in evt[start_seq:r, 1]]
			ifdinactive[current_row, din_ind] = True
			latency.append(current_latency)

			current_latency = evt[r,0]
			current_row += 1
			start_seq = r

	# return as a dataframe
	return pd.DataFrame(index=latency, columns=din_types,
						data=ifdinactive[:current_row, :])

def get_events_from_din(eeg):
	'''Turn DIN1, DIN2, ... etc. event channels to mne-python event array.

	Parameters
	----------
	eeg : Raw
		Instance of mne-python raw file.

	Returns
	-------
	events : array
		Numpy array of size (n_events, 3). Complies to mne-python
		convention of event arrays.
	'''
	# assert isinstance(mne.Raw) etc?
	df = din_dataframe(eeg)
	n_evnt = df.shape[0]
	events = np.zeros([n_evnt, 3], dtype='int')
	dins = np.tile(df.columns.values, [n_evnt, 1])
	events[:, 2] = np.sum(df.values * dins, axis=1)
	events[:, 0] = df.index.values
	return events

def remove_din_channels(eeg):
	'''Remove channels starting with 'D' like 'DIN2' or 'DI16'.'''
	# pick non-din channels:
	non_din_chans = [ch for ch in eeg.info['ch_names']
						if not ch.startswith('D')
						and not ch == 'STI 014']
	eeg.pick_channels(non_din_chans)

def reject_events_in_bad_segments(events, bad_segments, around_event=(-10,10),
								  remove_types=None):
	'''Remove events that coincide with bad segments.

	Parameters
	----------
	events 		 - mne events array
	bad_segments - N by 2 matrix with bad segment info
				   each row is a bad segment, first column is the
				   bad segment onset and the second column is the
				   bad segment offset (in samples)
	around_event - (lowerlim, higherlim) box limits around
				   each event. If a bad segment overlaps
				   with the box, the event is rejected.
				   Box limits are in samples.
				   (-500, 1000) means 500 samples before
				   up to 1000 samples after, but one can
				   create box that does not contain the
				   event like (250, 500)
	remove_types - list or numpy array of event types (int)
				   that should be checked for removal

	Returns
	-------
	events - corrected events array
	'''
	if remove_types is None:
		test_events = np.arange(events.shape[0])
	else:
		test_events = np.vstack([events[:,2] == x for x in remove_types])
		test_events = np.where(np.any(test_events, axis=0))

	ev = events[test_events, 0]
	ev = np.vstack([ev + x for x in around_event]).T
	remove = np.zeros(ev.shape[0], dtype='bool')
	for ii in range(ev.shape[0]):
		if np.any(np.logical_and(ev[ii,0] >= bad_segments[:,0],
			ev[ii,0] < bad_segments[:,1])):
			remove[ii] = True
			continue
		if np.any(np.logical_and(ev[ii,1] > bad_segments[:,0],
			ev[ii,1] <= bad_segments[:,1])):
			remove[ii] = True
			continue
		if np.any(np.logical_and(bad_segments[:,0] >= ev[ii,0],
			bad_segments[:,0] < ev[ii,1])):
			remove[ii] = True
	# remove events
	remove_ind = test_events[remove]
	if remove_ind.shape[0] > 0:
		events = np.delete(events, remove_ind, axis=0)
	return events

def extend_bads(rej_win, extend, copy=True):
	'''Extend rejection periods by specfied number of samples.'''
	if copy:
		rej_win = rej_win.copy()
	n_samples = rej_win.shape[1]
	bad_inds = np.any(rej_win, axis=1).nonzero()[0]
	for bad in bad_inds:
		slices = group(rej_win[bad, :], return_slice=True)
		for slc in slices:
			slc = extend_slice(slc, extend, n_samples)
			rej_win[bad, slc] = True
	return rej_win

def apply_artifacts_to_tfr(tfr, artif, orig_time, fillval=np.nan):
	'''Fill tfr data with `fillval` according to `artif`.
	Operates in-place.'''
	n_items, n_samples = artif.shape
	bad_inds = np.any(artif, axis=1).nonzero()[0]
	for bad in bad_inds:
		limits = group(artif[bad, :])
		for lim in limits:
			slc = find_range(tfr.times, list(orig_time[lim]))
			tfr.data[bad, :, :, slc] = fillval
	return tfr

# - [ ] make sure this is universal
def create_middle_events(events, min_time=14., sfreq=250):
	'''Put raven events in the trial centre,
	deleting events that begin trials shorter than
	min_time.'''
	ep_start = np.logical_or(events[:,2] == 1, events[:,2] == 0)
	ep_end = np.logical_or(events[:,2] == 10, events[:,2] == 11)
	ep_len = events[ep_end, 0] - events[ep_start, 0]
	long_enough = ep_len > int(min_time * sfreq)
	new_events = events[ep_start,:]
	new_events[:,0] += np.round(ep_len / 2).astype('int')
	return new_events[long_enough, :]

# TODO:
# - [ ] adapt this to use numpy structured datasets (?)
def read_set_events(filename, ignore_fields=None):
	'''Open set file, read epoch events and turn them into a dataframe.

	Parameters
	----------
	filename: str
		Name of the set file to read (can be absolute or relative path)
	ignore_fields: list of str | None
		Epoch event fields to ignore (these fields are note included in the df)

	Returns
	-------
	df: pandas.DataFrame
		Events read into dataframe
	'''
	EEG = loadmat(filename, uint16_codec='latin1',
				  struct_as_record=False, squeeze_me=True)['EEG']
	flds = [f for f in dir(EEG.event[0]) if not f.startswith('_')]
	events = EEG.event
	df_dict = dict()
	for f in flds:
		df = df_dict[f] = [ev.__getattribute__(f) for ev in events]
	df = pd.DataFrame(df_dict)

	# reorder columns:
	take_fields = ['epoch', 'type']
	ignore_fields = list() if ignore_fields is None else ignore_fields
	take_fields.extend([col for col in df.columns if not
					 (col in take_fields or col in ignore_fields)])
	return df.loc[:, take_fields]

def read_rej(fname, sfreq, bad_types=['reject']):
	from mne import Annotations

	rej = pd.read_table(fname)
	rej.loc[:, ['start', 'end']] = rej.loc[:, ['start', 'end']] / sfreq
	onset, duration = rej.start.values, (rej.end - rej.start).values
	description = ['BAD ' + val if val in bad_types else val
				   for val in rej.type.tolist()]
	return Annotations(onset, duration, description)

class AutoMark(object):
	def __init__(self):
		self.window = None
		self.step = None
		self.ranges = None
		self.variances = None

	def calculate(self, raw, window=1., step=0.25, minmax=False, variance=False,
				  reduction='mean', progressbar=True):
		if not minmax and not variance:
			raise Warning('Nothing computed. To compute variance you need to'
						  ' pass `variance=True`, to compute range you need to'
						  ' pass `minmax=True`.')
			return self

		data = raw._data
		window = int(round(window * raw.info['sfreq']))
		step = int(round(step * raw.info['sfreq']))
		self.window = window
		self.step = step
		self.sfreq = raw.info['sfreq']

		n_samples = data.shape[1]
		n_windows = int(np.floor((n_samples - window) / step))
		self.ranges = np.zeros(n_windows) if minmax else None
		self.variances = np.zeros(n_windows) if variance else None

		reduction = dict(mean=np.mean, max=np.max)[reduction]

		if progressbar:
			from tqdm import tqdm_notebook
			pbar = tqdm_notebook(total=n_windows)

		# step through data
		for window_idx in range(n_windows):
			first = window_idx * step
			last = first + window
			data_buffer = data[:, first:last]

			if minmax:
				self.ranges[window_idx] = reduction(
					data_buffer.max(axis=1) - data_buffer.min(axis=1))
			if variance:
				self.variances[window_idx] = reduction(data_buffer.var(axis=1))
			if progressbar:
				pbar.update(1)
		return self

	def reject(self, max_range=23e-5, max_variance=23e-6):
		if self.ranges is not None:
			is_bad_range = self.ranges > max_range
			# turn to annotations
			annot = self._to_annotation(is_bad_range, description='BAD_range')

		if self.variances is not None:
			is_bad_variance = self.variances > max_variance
			# turn to annotations
			if self.ranges is not None:
				annot = annot + self._to_annotation(is_bad_variance,
													description='BAD_variance')
			else:
				annot = self._to_annotation(is_bad_variance,
											description='BAD_variance')
		return annot

	def _to_annotation(self, bad_bool, description='BAD_'):
		import mne
		groups = group(bad_bool)
		if len(groups) > 0:
			onset = groups[:, 0] * (self.step / self.sfreq)
			duration = (self.window / self.sfreq) + np.diff(groups).ravel() * (self.step / self.sfreq)
			return mne.Annotations(onset, duration, [description] * groups.shape[0])
		else:
			return mne.Annotations([], [], [])

# - [ ] develop this function a little better
# - [ ] check what is used by raw.reject_bads() when no inds given
def mark_reject_peak2peak(raw, reject={'eeg': 23e-5}, window_length=1.,
						  label='bad p2p'):
	import mne
	from mne.utils import _reject_data_segments
	_, inds = _reject_data_segments(raw._data, reject, {'eeg': 0.},
									window_length, raw.info, 0.5)

	# turn inds to time, join
	time_segments = np.array(inds) / raw.info['sfreq']
	time_segments = join_segments(time_segments)

	segment_duration = np.diff(time_segments, axis=-1).ravel()
	return mne.Annotations(time_segments[:, 0], segment_duration, label)

# - [ ] maybe move to utils
def join_segments(time_segments):
	from mypy.utils import group

	# check which should be joined
	join_segments = (time_segments[:-1, 1] - time_segments[1:, 0]) == 0
	segment_groups = group(join_segments)
	segment_groups[:, 1] += 1

	# join segments
	final_segments = np.vstack([time_segments[segment_groups[:, 0], 0],
	                            time_segments[segment_groups[:, 1], 1]]).T

	# we missed single-segments, search for them
	prev = 0
	missed_segments = list()
	for row in segment_groups:
		if not row[0] == prev:
			missed_segments.extend(list(range(prev + 1, row[0])))
		prev = row[1]
	if prev < len(time_segments):
		missed_segments.extend(list(range(prev + 1, len(time_segments))))

	# append missed single segments to multi-segments and sort
	final_segments = np.vstack(
		[final_segments, time_segments[missed_segments, :]])
	final_segments = final_segments[final_segments[:, 0].argsort(), :]

	return final_segments

def get_dropped_epoch_index(epochs):
	'''get indices of dropped epochs'''
	current_epoch = 0
	removed_epochs = list()

	for info in epochs.drop_log:
	    if 'IGNORED' not in info:
	        if 'USER' in info:
	            removed_epochs.append(current_epoch)
	        current_epoch += 1

	return removed_epochs

def align_events(events1, events2):
	'''Return indices for each of the two event trains that give best match.'''
	common_events = list(set(events1).intersection(set(events2)))
	ind1, ind2 = [np.where(np.in1d(events, common_events))[0]
		      for events in [events1, events2]]

	len1, len2 = len(ind1), len(ind2)
	if len1 == len2:
		return ind1, ind2
	elif len1 > len2:
		longer, shorter = events1[ind1], events2[ind2]
	else:
		longer, shorter = events2[ind2], events1[ind1]

	# roll shorter along longer
	lng_len, shrt_len = len(longer), len(shorter)
	n_steps = lng_len - shrt_len + 1
	scores = np.zeros(n_steps, dtype='int')
	for offset in range(n_steps):
		scores[offset] = np.sum(longer[offset:offset + shrt_len] == shorter)
	best_offset = scores.argmax()

	if len1 > len2:
		ind1 = ind1[best_offset:best_offset + shrt_len]
	else:
		ind2 = ind2[best_offset:best_offset + shrt_len]
	return ind1, ind2
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
"""
Names testing module.
"""
# standard
import unittest
# package
from naming import *

class TestName(unittest.TestCase):

    def test_empty_name(self):
        n = Name()
        self.assertEqual('[base]', n.get_name())

    def test_init_name(self):
        n = Name('initname')
        self.assertEqual('initname', n.get_name())
        with self.assertRaises(NameError):
            Name(dict(my_name='dict'))
        n = Name(dict())
        self.assertFalse(None, n.get_name())

    def test_set_name(self):
        n = Name()
        n.set_name('setname')
        self.assertEqual('setname', n.get_name())

class TestPipe(unittest.TestCase):

    def test_empty_name(self):
        p = Pipe()
        self.assertEqual('[base]_[pipe]', p.get_name())
        self.assertEqual('[base].10', p.get_name(version=10))
        self.assertEqual('[base]_geo.10', p.get_name(version=10, output='geo'))
        self.assertEqual('[base]_geo.10.25', p.get_name(version=10, output='geo', frame=25))
        self.assertEqual('[base]_[output].10.25', p.get_name(version=10, frame=25))
        self.assertEqual('[base]_[output].[version].101', p.get_name(frame=101))
        self.assertEqual('[base]_cache.[version]', p.get_name(output='cache'))
        self.assertEqual('[base]', p.get_name(pipe=None))

    def test_empty_name_separator(self):
        p = Pipe()
        for sep in ' ', '.', '/', '/ .':
            p.separator = sep
            self.assertEqual(f'[base]{sep}[pipe]', p.get_name())
            self.assertEqual('[base].10', p.get_name(version=10))
            self.assertEqual(f'[base]{sep}geo.10', p.get_name(version=10, output='geo'))
            self.assertEqual(f'[base]{sep}geo.10.25', p.get_name(version=10, output='geo', frame=25))
            self.assertEqual(f'[base]{sep}[output].10.25', p.get_name(version=10, frame=25))
            self.assertEqual(f'[base]{sep}[output].[version].101', p.get_name(frame=101))
            self.assertEqual(f'[base]{sep}cache.[version]', p.get_name(output='cache'))
            self.assertEqual('[base]', p.get_name(pipe=None))

    def test_init_name(self):
        p = Pipe('initname_pipeline.0')
        self.assertEqual('initname_pipeline.0', p.get_name())
        self.assertEqual('_pipeline.0', p.pipe)
        self.assertEqual('pipeline', p.output)
        p = Pipe('initname.7')
        self.assertEqual('7', p.version)
        p = Pipe('initname_geo.0.1')
        self.assertEqual('0', p.version)
        self.assertEqual('1', p.frame)
        self.assertEqual('initname', p.nice_name)

        p = Pipe('name_with_underscores_pipeline.0')
        self.assertEqual('name_with_underscores_pipeline.0', p.get_name())
        self.assertEqual('_pipeline.0', p.pipe)
        self.assertEqual('pipeline', p.output)
        p = Pipe('name_with_underscores.7')
        self.assertEqual('7', p.version)
        p = Pipe('name_with_underscores_geo.0.1')
        self.assertEqual('0', p.version)
        self.assertEqual('1', p.frame)
        self.assertEqual('name_with_underscores', p.nice_name)

    def test_set_name(self):
        p = Pipe()
        p.set_name('setname_pipeline.0')
        self.assertEqual('_pipeline.0', p.pipe)
        self.assertEqual('pipeline', p.output)
        self.assertEqual('0', p.version)
        p.set_name('setname_pipeline.0.5')
        self.assertEqual('5', p.frame)
        self.assertEqual('setname', p.nice_name)

class TestFile(unittest.TestCase):

    def test_empty_name(self):
        f = File()
        self.assertEqual('[base].[extension]', f.get_name())
        self.assertEqual('[base].abc', f.get_name(extension='abc'))
        self.assertEqual('[base].[extension]', f.get_name(extension=''))
        with self.assertRaises(AttributeError):
            f.extension
        f.set_name('myfile.ext')
        self.assertEqual('ext', f.extension)
        self.assertEqual('myfile', f.base)
        self.assertEqual(f.get_name(), str(f.path))

class TestPipeFile(unittest.TestCase):

    def test_empty_name(self):
        f = PipeFile()
        self.assertEqual('[base]_[pipe].[extension]', f.get_name())
        self.assertEqual('[base]_[pipe].abc', f.get_name(extension='abc'))
        self.assertEqual('[base]_[pipe].[extension]', f.get_name(extension=''))
        with self.assertRaises(AttributeError):
            f.extension
        f.set_name('myfile_data.0.ext')
        self.assertEqual('ext', f.extension)
        self.assertEqual('myfile', f.base)
        self.assertEqual('_data.0', f.pipe)
        self.assertEqual('0', f.version)
        self.assertEqual(f.get_name(), str(f.path))
<commit_msg>easyname test<commit_after># -*- coding: utf-8 -*-
"""
Names testing module.
"""
# standard
import os
import unittest
from pathlib import Path
# package
from naming import *

class TestName(unittest.TestCase):

    def test_empty_name(self):
        n = Name()
        self.assertEqual('[base]', n.get_name())

    def test_init_name(self):
        n = Name('initname')
        self.assertEqual('initname', n.get_name())
        with self.assertRaises(NameError):
            Name(dict(my_name='dict'))
        n = Name(dict())
        self.assertFalse(None, n.get_name())

    def test_set_name(self):
        n = Name()
        n.set_name('setname')
        self.assertEqual('setname', n.get_name())
        extra_fields = dict(year='[0-9]{4}', username='[a-z]+', anotherfield='(constant)', lastfield='[a-zA-Z0-9]+')
        ProjectFile = type('ProjectFile', (EasyName, PipeFile), dict(config=extra_fields))
        pf = ProjectFile('project_data_name_2017_christianl_constant_iamlast_base.17.abc')
        self.assertEqual('project_data_name_2017_christianl_constant_iamlast', pf.nice_name)
        self.assertEqual('project_data_name_2017_christianl_constant_iamlast_base.17', pf.pipe_name)
        self.assertEqual('2017', pf.year)
        self.assertEqual('iamlast', pf.lastfield)
        self.assertEqual('abc', pf.extension)

class TestEasyName(unittest.TestCase):

    def test_empty_name(self):
        n = EasyName()
        self.assertEqual('[base]', n.get_name())
        self.assertEqual({}, n.get_values())

class TestPipe(unittest.TestCase):

    def test_empty_name(self):
        p = Pipe()
        self.assertEqual('[base]_[pipe]', p.get_name())
        self.assertEqual('[base].10', p.get_name(version=10))
        self.assertEqual('[base]_geo.10', p.get_name(version=10, output='geo'))
        self.assertEqual('[base]_geo.10.25', p.get_name(version=10, output='geo', frame=25))
        self.assertEqual('[base]_[output].10.25', p.get_name(version=10, frame=25))
        self.assertEqual('[base]_[output].[version].101', p.get_name(frame=101))
        self.assertEqual('[base]_cache.[version]', p.get_name(output='cache'))
        self.assertEqual('[base]', p.get_name(pipe=None))

    def test_empty_name_separator(self):
        p = Pipe()
        for sep in ' ', '.', '/', '/ .':
            p.separator = sep
            self.assertEqual(f'[base]{sep}[pipe]', p.get_name())
            self.assertEqual('[base].10', p.get_name(version=10))
            self.assertEqual(f'[base]{sep}geo.10', p.get_name(version=10, output='geo'))
            self.assertEqual(f'[base]{sep}geo.10.25', p.get_name(version=10, output='geo', frame=25))
            self.assertEqual(f'[base]{sep}[output].10.25', p.get_name(version=10, frame=25))
            self.assertEqual(f'[base]{sep}[output].[version].101', p.get_name(frame=101))
            self.assertEqual(f'[base]{sep}cache.[version]', p.get_name(output='cache'))
            self.assertEqual('[base]', p.get_name(pipe=None))

    def test_init_name(self):
        p = Pipe('initname_pipeline.0')
        self.assertEqual('initname_pipeline.0', p.get_name())
        self.assertEqual('_pipeline.0', p.pipe)
        self.assertEqual('pipeline', p.output)
        p = Pipe('initname.7')
        self.assertEqual('7', p.version)
        p = Pipe('initname_geo.0.1')
        self.assertEqual('0', p.version)
        self.assertEqual('1', p.frame)
        self.assertEqual('initname', p.nice_name)

        p = Pipe('name_with_underscores_pipeline.0')
        self.assertEqual('name_with_underscores_pipeline.0', p.get_name())
        self.assertEqual('_pipeline.0', p.pipe)
        self.assertEqual('pipeline', p.output)
        p = Pipe('name_with_underscores.7')
        self.assertEqual('7', p.version)
        p = Pipe('name_with_underscores_geo.0.1')
        self.assertEqual('0', p.version)
        self.assertEqual('1', p.frame)
        self.assertEqual('name_with_underscores', p.nice_name)

    def test_set_name(self):
        p = Pipe()
        p.set_name('setname_pipeline.0')
        self.assertEqual('_pipeline.0', p.pipe)
        self.assertEqual('pipeline', p.output)
        self.assertEqual('0', p.version)
        p.set_name('setname_pipeline.0.5')
        self.assertEqual('5', p.frame)
        self.assertEqual('setname', p.nice_name)

class TestFile(unittest.TestCase):

    def test_empty_name(self):
        f = File()
        self.assertEqual('[base].[extension]', f.get_name())
        self.assertEqual('[base].abc', f.get_name(extension='abc'))
        self.assertEqual('[base].[extension]', f.get_name(extension=''))
        with self.assertRaises(AttributeError):
            f.extension
        f.set_name('myfile.ext')
        self.assertEqual('ext', f.extension)
        self.assertEqual('myfile', f.base)
        self.assertEqual(f.get_name(), str(f.path))
        self.assertEqual(os.path.join(Path.home(), f.get_name()), str(f.full_path))

class TestPipeFile(unittest.TestCase):

    def test_empty_name(self):
        f = PipeFile()
        self.assertEqual('[base]_[pipe].[extension]', f.get_name())
        self.assertEqual('[base]_[pipe].abc', f.get_name(extension='abc'))
        self.assertEqual('[base]_[pipe].[extension]', f.get_name(extension=''))
        with self.assertRaises(AttributeError):
            f.extension
        f.set_name('myfile_data.0.ext')
        self.assertEqual('ext', f.extension)
        self.assertEqual('myfile', f.base)
        self.assertEqual('_data.0', f.pipe)
        self.assertEqual('0', f.version)
        self.assertEqual(f.get_name(), str(f.path))
<|endoftext|>


<commit_before>import pysam
import sys
import re
import subprocess
import os
from collections import defaultdict
from Bio import AlignIO

# reverse complement a sequence
complement = {'A': 'T', 'C': 'G', 'G': 'C', 'T': 'A'}
def revcomp(seq):
    reverse_complement = "".join(complement.get(base, base) for base in reversed(seq))
    return reverse_complement

# parse an LAshow read index string into a numeric ID
# the IDs are 1-based and contain commas
def lashow_idstr2idx(s):
    return int(s.replace(',', '')) - 1

# remove non-numeric characters from a string
def remove_nonnumeric(s):
    return re.sub("[^0-9]", "", s)

def remove_commas(s):
    return re.sub(",", "", s)

# parse an LAshow output file and build a map from a read index
# to the sequences that align to it
def parse_lashow(fn):
    fh = open(fn, 'r')

    out = defaultdict(list)

    # this is how to parse the LAshow output.
    pattern = re.compile("\D*(\d+)\s+(\d+)\s+(\w)\D+(\d+)\D+(\d+)\D+(\d+)\D+(\d+).*")

    for line in fh:

        # Strip commas from the numbers to make the regex easier
        line = remove_commas(line)
        m = pattern.match(line)

        if m is None:
            continue

        id1 = lashow_idstr2idx(m.group(1))
        id2 = lashow_idstr2idx(m.group(2))
        strand = m.group(3)

        #
        s = int(m.group(6))
        e = int(m.group(7))

        out[id1].append((id2, strand, s, e))
    return out

# write a fasta file for input into POA
def write_poa_input(overlaps, read_idx):
    fn = "poa.input.%d.fa" % (read_idx)
    fh = open(fn, "w")

    read_id1 = ref.references[read_idx]
    seq1 = ref.fetch(read_id1)
    fh.write(">%s\n%s\n" % ("poabaseread", seq1))

    n_reads = 0
    for o in overlaps[read_idx]:

        read_id2 = ref.references[o[0]]
        seq2 = ref.fetch(read_id2)

        # strand
        if o[1] == "c":
            seq2 = revcomp(seq2)

        # restrict to the part of the sequence that matches read1
        seq2 = seq2[o[2]:o[3]]

        fh.write(">%s\n%s\n" % (read_id2, seq2))
        n_reads += 1
    fh.close()
    return (fn, n_reads)

def clustal2consensus(fn):

    alignment = AlignIO.read(fn, "clustal")

    min_coverage = 3
    read_row = -1
    consensus_row = -1

    for (i, record) in enumerate(alignment):
        if record.id == 'poabaseread':
            read_row = i
        if record.id == 'CONSENS0':
            consensus_row = i

    if consensus_row == -1:
        return ""

    # Work out the first and last columns that contains
    # bases of the read we are correcting
    (first_col, last_col) = get_sequence_coords(alignment[read_row].seq)

    # Calculate a vector of depths along the consensus
    depths = [0] * len(alignment[read_row].seq)

    for record in alignment:
        (aln_first_col, aln_last_col) = get_sequence_coords(record.seq)
        for i in xrange(aln_first_col, aln_last_col):
            if aln_first_col >= first_col and aln_last_col <= last_col \
               and not record.id.startswith('CONSENS'):
                 depths[i] += 1

    # Change the boundaries to only include high-depth bases
    while first_col != last_col:
        if depths[first_col] >= min_coverage:
            break
        first_col += 1

    while last_col != first_col:
        if depths[last_col] >= min_coverage:
            break
        last_col -= 1

    # Extract the consensus sequence
    consensus = str(alignment[consensus_row].seq[first_col:last_col])
    consensus = consensus.replace('-', '')

    return consensus

# Return the first and last column of the multiple alignment
# that contains a base for the given sequence row
def get_sequence_coords(seq):
    first_col = -1
    last_col = -1
    for (i, s) in enumerate(seq):
        if s != '-' and first_col == -1:
            first_col = i
        if s != '-':
            last_col = i
    return (first_col, last_col)

#
def run_poa_and_consensus(overlaps, read_idx):
    (in_fn, n_reads) = write_poa_input(overlaps, read_idx)
    out_fn = "clustal-%d.out" % (read_idx)
    DEVNULL = open(os.devnull, 'wb')
    cmd = "poa -read_fasta %s -clustal %s -hb poa-blosum80.mat" % (in_fn, out_fn)
    p = subprocess.Popen(cmd, shell=True, stderr=DEVNULL)
    p.wait()
    consensus =  clustal2consensus(out_fn)

    os.remove(in_fn)
    os.remove(out_fn)
    return (consensus, n_reads)

def run_lashow(name, start, end):

    out_fn = "lashow.%s-%s.out" % (start, end)
    out_fh = open(out_fn, 'w')
    cmd = "LAshow %s.db %s.las %s-%s" % (name, name, start, end)
    p = subprocess.Popen(cmd, shell=True, stdout=out_fh)
    p.wait()
    out_fh.close()
    return out_fn

# Args
name = sys.argv[1]
read_range = sys.argv[2]

# Open reference file
ref_fn = "%s.pp.fasta" % (name)
ref = pysam.Fastafile(ref_fn)

# Parse the range of read ids to correct
(start, end) = [ int(x) for x in read_range.split(':') ]

# Generate the LAshow file indicating overlaps
lashow_fn = run_lashow(name, start, end)

# Make a dictionary of overlaps
overlaps = parse_lashow(lashow_fn)

# Correct each read with POA
for read_idx in xrange(start, end):
    (seq, n_reads) = run_poa_and_consensus(overlaps, read_idx)

    if seq != "":
       print ">%d n_reads=%d\n%s" % (read_idx, n_reads, seq)

os.remove(lashow_fn)
<commit_msg>Use blosum80 supplied with the nanocorrect<commit_after>import pysam
import sys
import re
import subprocess
import os
from collections import defaultdict
from Bio import AlignIO

# reverse complement a sequence
complement = {'A': 'T', 'C': 'G', 'G': 'C', 'T': 'A'}
def revcomp(seq):
    reverse_complement = "".join(complement.get(base, base) for base in reversed(seq))
    return reverse_complement

# parse an LAshow read index string into a numeric ID
# the IDs are 1-based and contain commas
def lashow_idstr2idx(s):
    return int(s.replace(',', '')) - 1

# remove non-numeric characters from a string
def remove_nonnumeric(s):
    return re.sub("[^0-9]", "", s)

def remove_commas(s):
    return re.sub(",", "", s)

# parse an LAshow output file and build a map from a read index
# to the sequences that align to it
def parse_lashow(fn):
    fh = open(fn, 'r')

    out = defaultdict(list)

    # this is how to parse the LAshow output.
    pattern = re.compile("\D*(\d+)\s+(\d+)\s+(\w)\D+(\d+)\D+(\d+)\D+(\d+)\D+(\d+).*")

    for line in fh:

        # Strip commas from the numbers to make the regex easier
        line = remove_commas(line)
        m = pattern.match(line)

        if m is None:
            continue

        id1 = lashow_idstr2idx(m.group(1))
        id2 = lashow_idstr2idx(m.group(2))
        strand = m.group(3)

        #
        s = int(m.group(6))
        e = int(m.group(7))

        out[id1].append((id2, strand, s, e))
    return out

# write a fasta file for input into POA
def write_poa_input(overlaps, read_idx):
    fn = "poa.input.%d.fa" % (read_idx)
    fh = open(fn, "w")

    read_id1 = ref.references[read_idx]
    seq1 = ref.fetch(read_id1)
    fh.write(">%s\n%s\n" % ("poabaseread", seq1))

    n_reads = 0
    for o in overlaps[read_idx]:

        read_id2 = ref.references[o[0]]
        seq2 = ref.fetch(read_id2)

        # strand
        if o[1] == "c":
            seq2 = revcomp(seq2)

        # restrict to the part of the sequence that matches read1
        seq2 = seq2[o[2]:o[3]]

        fh.write(">%s\n%s\n" % (read_id2, seq2))
        n_reads += 1
    fh.close()
    return (fn, n_reads)

def clustal2consensus(fn):

    alignment = AlignIO.read(fn, "clustal")

    min_coverage = 3
    read_row = -1
    consensus_row = -1

    for (i, record) in enumerate(alignment):
        if record.id == 'poabaseread':
            read_row = i
        if record.id == 'CONSENS0':
            consensus_row = i

    if consensus_row == -1:
        return ""

    # Work out the first and last columns that contains
    # bases of the read we are correcting
    (first_col, last_col) = get_sequence_coords(alignment[read_row].seq)

    # Calculate a vector of depths along the consensus
    depths = [0] * len(alignment[read_row].seq)

    for record in alignment:
        (aln_first_col, aln_last_col) = get_sequence_coords(record.seq)
        for i in xrange(aln_first_col, aln_last_col):
            if aln_first_col >= first_col and aln_last_col <= last_col \
               and not record.id.startswith('CONSENS'):
                 depths[i] += 1

    # Change the boundaries to only include high-depth bases
    while first_col != last_col:
        if depths[first_col] >= min_coverage:
            break
        first_col += 1

    while last_col != first_col:
        if depths[last_col] >= min_coverage:
            break
        last_col -= 1

    # Extract the consensus sequence
    consensus = str(alignment[consensus_row].seq[first_col:last_col])
    consensus = consensus.replace('-', '')

    return consensus

# Return the first and last column of the multiple alignment
# that contains a base for the given sequence row
def get_sequence_coords(seq):
    first_col = -1
    last_col = -1
    for (i, s) in enumerate(seq):
        if s != '-' and first_col == -1:
            first_col = i
        if s != '-':
            last_col = i
    return (first_col, last_col)

#
def run_poa_and_consensus(overlaps, read_idx):
    (in_fn, n_reads) = write_poa_input(overlaps, read_idx)
    out_fn = "clustal-%d.out" % (read_idx)
    DEVNULL = open(os.devnull, 'wb')
    # use blosum file relative to the 'nanocorrect.py'
    blosum_file = os.path.join(os.path.dirname(__file__), "poa-blosum80.mat")
    cmd = "poa -read_fasta %s -clustal %s -hb %s" % (in_fn, out_fn, blosum_file)
    p = subprocess.Popen(cmd, shell=True, stderr=DEVNULL)
    p.wait()
    consensus =  clustal2consensus(out_fn)

    os.remove(in_fn)
    os.remove(out_fn)
    return (consensus, n_reads)

def run_lashow(name, start, end):

    out_fn = "lashow.%s-%s.out" % (start, end)
    out_fh = open(out_fn, 'w')
    cmd = "LAshow %s.db %s.las %s-%s" % (name, name, start, end)
    p = subprocess.Popen(cmd, shell=True, stdout=out_fh)
    p.wait()
    out_fh.close()
    return out_fn

# Args
name = sys.argv[1]
read_range = sys.argv[2]

# Open reference file
ref_fn = "%s.pp.fasta" % (name)
ref = pysam.Fastafile(ref_fn)

# Parse the range of read ids to correct
(start, end) = [ int(x) for x in read_range.split(':') ]

# Generate the LAshow file indicating overlaps
lashow_fn = run_lashow(name, start, end)

# Make a dictionary of overlaps
overlaps = parse_lashow(lashow_fn)

# Correct each read with POA
for read_idx in xrange(start, end):
    (seq, n_reads) = run_poa_and_consensus(overlaps, read_idx)

    if seq != "":
       print ">%d n_reads=%d\n%s" % (read_idx, n_reads, seq)

os.remove(lashow_fn)
<|endoftext|>


<commit_before>import logging
from urllib.parse import urljoin

import requests
from json_encoder import json

from simple_rest_client.models import Request, Response
from simple_rest_client.decorators import handle_request_error

logger = logging.getLogger(__name__)

class BaseResource:
    def __init__(self, api_root_url=None, resource_name=None, action_urls={},
                 params={}, headers={}, timeout=None, append_slash=False,
                 json_encode_body=False):
        self.api_root_url = api_root_url
        self.resource_name = resource_name
        self.params = params
        self.headers = headers
        self.action_urls = action_urls or self.get_default_action_urls()
        self.timeout = timeout or 3
        self.append_slash = append_slash
        self.json_encode_body = json_encode_body

    def get_default_action_urls(self):
        return {
            'list': self.resource_name,
            'create': self.resource_name,
            'retrieve': self.resource_name + '/{}',
            'update': self.resource_name + '/{}',
            'partial_update': self.resource_name + '/{}',
            'destroy': self.resource_name + '/{}'
        }

    def get_full_url(self, action, *parts):
        url = self.action_urls.get(action, '').format(*parts)
        if not url:
            raise ValueError('No url match for "{}"'.format(action))
        if self.append_slash and not url.endswith('/'):
            url += '/'
        return urljoin(self.api_root_url, url)

    def build_request_instance(self, *args, action_name=None, method=None,
                               params=None, body=None, headers=None,
                               timeout=None):
        url = self.get_full_url(action_name, *args)
        if self.json_encode_body and body:
            body = json.dumps(body)
        return Request(
            url=url, method=method, params=params, body=body, headers=headers,
            timeout=timeout
        )

class Resource(BaseResource):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.session = requests.Session()

    @handle_request_error
    def make_request(self, request):
        logger.debug('operation=request_started, request={!r}'.format(request))
        method = getattr(self.session, request.method.lower())
        request.params.update(self.params)
        request.headers.update(self.headers)
        client_response = method(
            request.url,
            params=request.params,
            data=request.body,
            headers=request.headers,
            timeout=request.timeout
        )
        body = client_response.text
        content_type = client_response.headers.get('Content-Type', '')
        if 'application/json' in content_type:
            body = json.loads(body)

        response = Response(
            url=client_response.url,
            method=client_response.request.method,
            body=body,
            headers=client_response.headers,
            status_code=client_response.status_code
        )
        logger.debug(
            'operation=request_finished, request={!r}, response={!r}'.format(
                request, response
            )
        )
        return response

    def list(self, *args, params={}, headers={}):
        request = self.build_request_instance(
            *args, action_name='list', method='GET', params=params, body=None,
            headers=headers, timeout=self.timeout
        )
        return self.make_request(request)

    def create(self, *args, body=None, params={}, headers={}):
        request = self.build_request_instance(
            *args, action_name='create', method='POST', params=params,
            body=body, headers=headers, timeout=self.timeout
        )
        return self.make_request(request)

    def retrieve(self, *args, params={}, headers={}):
        request = self.build_request_instance(
            *args, action_name='retrieve', method='GET', params=params,
            body=None, headers=headers, timeout=self.timeout
        )
        return self.make_request(request)

    def update(self, *args, body=None, params={}, headers={}):
        request = self.build_request_instance(
            *args, action_name='update', method='PUT', params=params,
            body=body, headers=headers, timeout=self.timeout
        )
        return self.make_request(request)

    def partial_update(self, *args, body=None, params={}, headers={}):
        request = self.build_request_instance(
            *args, action_name='partial_update', method='PATCH', params=params,
            body=body, headers=headers, timeout=self.timeout
        )
        return self.make_request(request)

    def destroy(self, *args, params={}, headers={}):
        request = self.build_request_instance(
            *args, action_name='destroy', method='DELETE', params=params,
            body=None, headers=headers, timeout=self.timeout
        )
        return self.make_request(request)
<commit_msg>Add body param to all actions<commit_after>import logging
from urllib.parse import urljoin

import requests
from json_encoder import json

from simple_rest_client.models import Request, Response
from simple_rest_client.decorators import handle_request_error

logger = logging.getLogger(__name__)

class BaseResource:
    def __init__(self, api_root_url=None, resource_name=None, action_urls={},
                 params={}, headers={}, timeout=None, append_slash=False,
                 json_encode_body=False):
        self.api_root_url = api_root_url
        self.resource_name = resource_name
        self.params = params
        self.headers = headers
        self.action_urls = action_urls or self.get_default_action_urls()
        self.timeout = timeout or 3
        self.append_slash = append_slash
        self.json_encode_body = json_encode_body

    def get_default_action_urls(self):
        return {
            'list': self.resource_name,
            'create': self.resource_name,
            'retrieve': self.resource_name + '/{}',
            'update': self.resource_name + '/{}',
            'partial_update': self.resource_name + '/{}',
            'destroy': self.resource_name + '/{}'
        }

    def get_full_url(self, action, *parts):
        url = self.action_urls.get(action, '').format(*parts)
        if not url:
            raise ValueError('No url match for "{}"'.format(action))
        if self.append_slash and not url.endswith('/'):
            url += '/'
        return urljoin(self.api_root_url, url)

    def build_request_instance(self, *args, action_name=None, method=None,
                               params=None, body=None, headers=None,
                               timeout=None):
        url = self.get_full_url(action_name, *args)
        if self.json_encode_body and body:
            body = json.dumps(body)
        return Request(
            url=url, method=method, params=params, body=body, headers=headers,
            timeout=timeout
        )

class Resource(BaseResource):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.session = requests.Session()

    @handle_request_error
    def make_request(self, request):
        logger.debug('operation=request_started, request={!r}'.format(request))
        method = getattr(self.session, request.method.lower())
        request.params.update(self.params)
        request.headers.update(self.headers)
        client_response = method(
            request.url,
            params=request.params,
            data=request.body,
            headers=request.headers,
            timeout=request.timeout
        )
        body = client_response.text
        content_type = client_response.headers.get('Content-Type', '')
        if 'application/json' in content_type:
            body = json.loads(body)

        response = Response(
            url=client_response.url,
            method=client_response.request.method,
            body=body,
            headers=client_response.headers,
            status_code=client_response.status_code
        )
        logger.debug(
            'operation=request_finished, request={!r}, response={!r}'.format(
                request, response
            )
        )
        return response

    def list(self, *args, body=None, params={}, headers={}):
        request = self.build_request_instance(
            *args, action_name='list', method='GET', params=params, body=body,
            headers=headers, timeout=self.timeout
        )
        return self.make_request(request)

    def create(self, *args, body=None, params={}, headers={}):
        request = self.build_request_instance(
            *args, action_name='create', method='POST', params=params,
            body=body, headers=headers, timeout=self.timeout
        )
        return self.make_request(request)

    def retrieve(self, *args, body=None, params={}, headers={}):
        request = self.build_request_instance(
            *args, action_name='retrieve', method='GET', params=params,
            body=body, headers=headers, timeout=self.timeout
        )
        return self.make_request(request)

    def update(self, *args, body=None, params={}, headers={}):
        request = self.build_request_instance(
            *args, action_name='update', method='PUT', params=params,
            body=body, headers=headers, timeout=self.timeout
        )
        return self.make_request(request)

    def partial_update(self, *args, body=None, params={}, headers={}):
        request = self.build_request_instance(
            *args, action_name='partial_update', method='PATCH', params=params,
            body=body, headers=headers, timeout=self.timeout
        )
        return self.make_request(request)

    def destroy(self, *args, body=None, params={}, headers={}):
        request = self.build_request_instance(
            *args, action_name='destroy', method='DELETE', params=params,
            body=body, headers=headers, timeout=self.timeout
        )
        return self.make_request(request)
<|endoftext|>


<commit_before>#Python reminders

# Collapse the list for output
''.join(list)
#Collapse and convert items to string
''.join(map(str, lst))

#Flatten list of lists
flat_list = [item for subitem in lists for item in subitem]

#Exammple for checking list for single digit (string)
def checker(string):
	print("Boom!" if "7" in str(string) else "there is no 7 in the list")
checker([2, 55, 60, 97, 86, 3, 4, 6, 6])

#Check list for letters or numbers
letters = ''.join(x for x in item if x.isalpha())
digits = ''.join(x for x in item if x.isnumeric())

#Scan list for longest string and print max item
print(max(len(i) for i in string.split()))

#Search for highest index of substring
string.rindex(substring) -- If not found, raises exception
string.rfind(substring) -- If not found, returns -1

#Dont forget min and max
max(x) for largest
min(x) for smallest

#Count occurrences of item in list
[a,b,c,d,e].count(item)
sum([x.count(letter) for x in lst])

#Count Sublists in list
return str(lst).count('[') - 1
return sum(1 for i in lst if type(i)==list)
return sum(isinstance(i, list) for i in lst)

#Count occurrences of list 1 items in list 2
value =  len([i for i in list2 if i in list1 ])

#Put anything in square brackets to search list index
a[0]
#Reverse list of ints by converting to string and list reversing
str(num)[::-1]

#Strange list rotation
def rotate_left3(nums):
  a = nums[0]
  nums[0] = nums[1]
  nums[1] = nums[2]
  nums[2] = a
  return nums -> [1,2,3] = [2,3,1]

#Evaluator positioning
if someting:
	return (item 1 >= 10)
else:
	return (item > 1 and item <= 5)

#One line conditionals
	#get item index from lst, not there, return -1
return lst.index(item) if item in lst else -1

#Use F strings
x = "hello"
w = "world"
print(f"{x} {w}")

#Use Swapcase to flip lowercase and uppercase
string.swapcase()

#List comprehensions
	#Reverse all words >= 5 letters and print out the modified sentence
	[word[::-1] if len(word) >= 5 else word for word in words]
	#Create a dictionary of letters and their frequency of occurrence in a string
	{i : message.count(i) for i in set(message)}
	#List all indexes of a specified value
	def get_indicies(lst, el):
	[i for i in range(len(lst)) if lst[i] == el]
	#Edabit solution for odd occurrence of item in list
	[i for i in lst if lst.count(i) % 2 > 0]
	#Only return the integers in a list
	[x for x in l if type(x) == int]
	#Flatten list of lists
	flat_list = [item for subitem in lists for item in subitem]
	#Split int into list of digits
	split_number = [int(d) for d in str(n)]
	#Dictionary values for item
	value = [str(dictionary[i]) for i in string]
	#list character index from left to right, right to left
	[mylist.index(char), mylist.rindex(char)]

#Letter replacements/translations
def correct(string):
    return string.translate(str.maketrans("501", "SOI"))

#Dictionary reminders
dictionary_name.values() - list all values
keys() - list all keys
get() - access value for key
use Dict[1] or Dict[name] - access item via key

#Simple hard coded cipher
def move(word):
	return ''.join(chr(ord(i) + 1) for i in word)

#Collapse multiple inputs into one
	#from this
	def predict_age(age_1, age_2, age_3, age_4, age_5, age_6, age_7, age_8):
	#to this
	def predict_age(*ages)

#Remove char
def remove_exclamation_marks(s):
    return s.replace('!', '')

#Remove duplicate words, keep index position
def remove_duplicate_words(s):
    return ' '.join(sorted(set(s.split()), key = s.index))

#Quick Lambda implementation on multiple args - with f"string"
def newvibe(a,b,c):
	var1 = lambda x: (x * 5)
	print(f"This is a: {var1(a)}, This is b: {var1(b)}, This is c: {var1(c)}")
<commit_msg>Update REMINDERS.py<commit_after>#Python reminders

# Collapse the list for output
''.join(list)
#Collapse and convert items to string
''.join(map(str, lst))

#Flatten list of lists
flat_list = [item for subitem in lists for item in subitem]

#Exammple for checking list for single digit (string)
def checker(string):
	print("Boom!" if "7" in str(string) else "there is no 7 in the list")
checker([2, 55, 60, 97, 86, 3, 4, 6, 6])

#Check list for letters or numbers
letters = ''.join(x for x in item if x.isalpha())
digits = ''.join(x for x in item if x.isnumeric())

#Scan list for longest string and print max item
print(max(len(i) for i in string.split()))

#Search for highest index of substring
string.rindex(substring) -- If not found, raises exception
string.rfind(substring) -- If not found, returns -1

#Dont forget min and max
max(x) for largest
min(x) for smallest

#Count occurrences of item in list
[a,b,c,d,e].count(item)
sum([x.count(letter) for x in lst])

#Count Sublists in list
return str(lst).count('[') - 1
return sum(1 for i in lst if type(i)==list)
return sum(isinstance(i, list) for i in lst)

#Count occurrences of list 1 items in list 2
value =  len([i for i in list2 if i in list1 ])

#Put anything in square brackets to search list index
a[0]
#Reverse list of ints by converting to string and list reversing
str(num)[::-1]

#Strange list rotation
def rotate_left3(nums):
  a = nums[0]
  nums[0] = nums[1]
  nums[1] = nums[2]
  nums[2] = a
  return nums -> [1,2,3] = [2,3,1]

#Evaluator positioning
if someting:
	return (item 1 >= 10)
else:
	return (item > 1 and item <= 5)

#One line conditionals
	#get item index from lst, not there, return -1
return lst.index(item) if item in lst else -1

#Use F strings
x = "hello"
w = "world"
print(f"{x} {w}")

#Aligning text
print('{:<10}'.format('X')) # left align
X
print('{:>10}'.format('X')) # right align
         X
print('{:^10}'.format('X')) # center
    X
print('{:?^10}'.format('X')) # add a fill character
????X?????

#Using commas as a delimiter
print('{:,}'.format(2783727282727))
2,783,727,282,727

print('{:.2%}'.format(90.60/100))
90.60%

#Use Swapcase to flip lowercase and uppercase
string.swapcase()

#List comprehensions
	#Reverse all words >= 5 letters and print out the modified sentence
	[word[::-1] if len(word) >= 5 else word for word in words]
	#Create a dictionary of letters and their frequency of occurrence in a string
	{i : message.count(i) for i in set(message)}
	#List all indexes of a specified value
	def get_indicies(lst, el):
	[i for i in range(len(lst)) if lst[i] == el]
	#Edabit solution for odd occurrence of item in list
	[i for i in lst if lst.count(i) % 2 > 0]
	#Only return the integers in a list
	[x for x in l if type(x) == int]
	#Flatten list of lists
	flat_list = [item for subitem in lists for item in subitem]
	#Split int into list of digits
	split_number = [int(d) for d in str(n)]
	#Dictionary values for item
	value = [str(dictionary[i]) for i in string]
	#list character index from left to right, right to left
	[mylist.index(char), mylist.rindex(char)]

#Letter replacements/translations
def correct(string):
    return string.translate(str.maketrans("501", "SOI"))

#Dictionary reminders
dictionary_name.values() - list all values
keys() - list all keys
get() - access value for key
use Dict[1] or Dict[name] - access item via key

#Simple hard coded cipher
def move(word):
	return ''.join(chr(ord(i) + 1) for i in word)

#Collapse multiple inputs into one
	#from this
	def predict_age(age_1, age_2, age_3, age_4, age_5, age_6, age_7, age_8):
	#to this
	def predict_age(*ages)

#Remove char
def remove_exclamation_marks(s):
    return s.replace('!', '')

#Remove duplicate words, keep index position
def remove_duplicate_words(s):
    return ' '.join(sorted(set(s.split()), key = s.index))

#Quick Lambda implementation on multiple args - with f"string"
def newvibe(a,b,c):
	var1 = lambda x: (x * 5)
	print(f"This is a: {var1(a)}, This is b: {var1(b)}, This is c: {var1(c)}")
<|endoftext|>