<commit_before># Copyright 2012 Lars Hupfeldt Nielsen, Hupfeldt IT
# This code is free for anybody to use

from collections import Sequence, OrderedDict
from envs import Env
from attribute_collector import AttributeCollector
from config_errors import ConfigException, NoAttributeException

class _ConfigBase(object):
    nested = []

    # Decoration attributes
    _deco_named_as = None
    _deco_repeatable = False
    _deco_nested_repeatables = []
    _deco_required_attributes = []
    _deco_required_if_attributes = (None, ())

    def __init__(self, **attr):
        self._debug_exc = True

        # Object linking
        self._root_conf = None

        # Dict of dicts: attributes['a'] = dict(Env('prod')=1, EnvGroup('dev')=0)
        self._attributes = OrderedDict()
        for key in self.__class__._deco_nested_repeatables:
            self._attributes[key] = OrderedDict()

        self._defaults = attr
        self._frozen = False
        self._may_freeze = True

    def named_as(self):
        if self.__class__._deco_named_as:
            return self.__class__._deco_named_as
        if self.__class__._deco_repeatable:
            return self.__class__.__name__ + 's'
        return self.__class__.__name__

    def irepr(self, indent_level):
        """Indented repr"""
        indent1 = '  ' * indent_level
        indent2 =  indent1 + '     '
        # + ':' + self.__class__.__name__
        not_frozen_msg = "" if self._frozen else " not-frozen, defaults: " + repr(self._defaults)
        return self.named_as() + not_frozen_msg + ' {\n' + \
            ''.join([indent2 + name + ': ' + repr(value) + ',\n' for name, value in self.iteritems()]) + \
            indent1 + '}'

    def __repr__(self):
        return self.irepr(len(self.__class__.nested) -1)

    def __enter__(self):
        assert not self._frozen
        self._may_freeze = False
        return self

    def freeze_validate_required(self):
        missing = []
        for req in self.__class__._deco_required_attributes:
            if not req in self._attributes:
                missing.append(req)
        if missing:
            raise ConfigException("No value given for required attributes: " + repr(missing))

    def freeze_validate_required_if(self):
        required_if_key = self.__class__._deco_required_if_attributes[0]
        if not required_if_key:
            return

        try:
            required_if = self._attributes[required_if_key].value()
            if not required_if:
                return
        except KeyError:
            return
        except NoAttributeException:
            return

        missing = []
        for req in self.__class__._deco_required_if_attributes[1]:
            if not req in self._attributes:
                missing.append(req)
        if missing:
            raise ConfigException("Missing required_if attributes. Condition attribute: " + repr(required_if_key) + "==" +  repr(required_if) + ", missing: " + repr(missing))

    def freeze_validation(self):
        """Override this method if you need special checks"""
        self.freeze_validate_required()
        self.freeze_validate_required_if()

    def freeze(self):
        assert not self._frozen
        assert self._may_freeze

        # Collect remaining default values
        for name in list(self._defaults):
            AttributeCollector(name, self)()
        self._frozen = True

        try:
            self.freeze_validation()
            return self
        except ConfigException as ex:
            if self._debug_exc:
                raise
            # Strip stack
            raise ex
        except NoAttributeException as ex:
            if self._debug_exc:
                raise
            raise ex

    def __exit__(self, exc_type, exc_value, traceback):
        if exc_type:
            return None

        try:
            self._may_freeze = True
            self.freeze()
        except ConfigException as ex:
            if self._debug_exc:
                raise
            raise ex
        except NoAttributeException as ex:
            if self._debug_exc:
                raise
            raise ex

    def __setattr__(self, name, value):
        if name[0] == '_':
            # Needed to set private values in __init__
            super(_ConfigBase, self).__setattr__(name, value)
            return
        raise ConfigException("Trying to set a property " + repr(name) + " on a config item")

    def _check_valid_env(self, env, valid_envs):
        if type(env) != Env:
            raise ConfigException(self.__class__.__name__ + ': env must be of type ' + repr(Env.__name__) + '; found type ' + repr(env.__class__.__name__))

        for valid_env in valid_envs:
            if env in valid_env:
                return
        raise ConfigException("The env " + repr(env) + " must be in the (nested) list of valid_envs " + repr(valid_envs))

    def _env_specific_value(self, attr_name, attr_coll, env):
        if isinstance(attr_coll, ConfigItem) or isinstance(attr_coll, OrderedDict):
            return attr_coll

        try:
            return attr_coll.env_values[env]
        except KeyError:
            self._check_valid_env(env, self._root_conf._valid_envs)
            raise NoAttributeException("Attribute " + repr(attr_name) + " undefined for env " + repr(env))

    def getattr_env(self, name, env):
        try:
            attr_coll = self._attributes[name]
        except KeyError:
            error_message = ""
            repeatable_name = name + 's'
            if self._attributes.get(repeatable_name):
                error_message = ", but found attribute " + repr(repeatable_name)
            raise ConfigException(repr(self) + " has no attribute " + repr(name) + error_message)

        return self._env_specific_value(name, attr_coll, env)

    def __getattr__(self, name):
        if not self._frozen and self._may_freeze:
            self.freeze()

        if not self._frozen:
            try:
                # Return existing collector/dict of nested items if any
                coll = self._attributes[name]
                if type(coll) in (AttributeCollector, OrderedDict):
                    return coll
                if isinstance(coll, _ConfigItem):
                    return AttributeCollector(name, self)
                raise Exception("Internal error, unexpected type: " + repr(type(coll)) + ':' + repr(coll))
            except KeyError:
                return AttributeCollector(name, self)

        try:
            return self.getattr_env(name, self._root_conf._selected_env)
        except ConfigException as ex:
            if self._debug_exc:
                raise
            raise ex

    def iteritems(self):
        for key, value in self._attributes.iteritems():
            try:
                yield key, self._env_specific_value(key, value, self._root_conf._selected_env)
            except NoAttributeException:
                # This should only happen in case of  a conditional attribute
                pass

    def items(self):
        return list(self.iteritems())

    @property
    def contained_in(self):
        return self._contained_in

    @property
    def root_conf(self):
        return self._root_conf

    @property
    def attributes(self):
        return self._attributes

class _ConfigItem(_ConfigBase):
    def _validate(self):
        self.validate()
        for _child_name, child_value in self.iteritems():
            if isinstance(child_value, OrderedDict):
                for dict_entry in child_value.values():
                    dict_entry._validate()

            if isinstance(child_value, _ConfigItem):
                child_value._validate()

    def validate(self):
        """Can be overridden to provide post-frozen validation"""
        pass

    def __enter__(self):
        super(_ConfigItem, self).__enter__()
        self.__class__.nested.append(self)
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        self.__class__.nested.pop()
        super(_ConfigItem, self).__exit__(exc_type, exc_value, traceback)

class ConfigRoot(_ConfigItem):
    def __init__(self, selected_env, valid_envs, **attr):
        if not isinstance(valid_envs, Sequence) or isinstance(valid_envs, str):
            raise ConfigException(self.__class__.__name__ + ": valid_envs arg must be a 'Sequence'; found type " + repr(valid_envs.__class__.__name__))

        for env in valid_envs:
            if not isinstance(env, Env):
                raise ConfigException(self.__class__.__name__ + ": valid_envs items must be instance of 'Env'; found a " + repr(env.__class__.__name__))

        self._check_valid_env(selected_env, valid_envs)

        self._selected_env = selected_env
        self._valid_envs = valid_envs
        super(ConfigRoot, self).__init__(**attr)
        self._root_conf = self
        self._contained_in = None

    def __exit__(self, exc_type, exc_value, traceback):
        super(ConfigRoot, self).__exit__(exc_type, exc_value, traceback)
        self._validate()

    @property
    def selected_env(self):
        return self._selected_env

class ConfigItem(_ConfigItem):
    def __init__(self, **attr):
        super(ConfigItem, self).__init__(**attr)

        # Automatic Nested Insert in parent
        if not self.__class__.nested:
            raise ConfigException(self.__class__.__name__ + " object must be nested (indirectly) in a " + repr(ConfigRoot.__name__))

        # Set back reference to containing Item and root item
        self._contained_in = self.__class__.nested[-1]
        self._root_conf = self._contained_in.root_conf

        # Insert self in containing Item's attributes
        my_key = self.named_as()

        if self.__class__._deco_repeatable:
            # Validate that the containing item has specified this item as repeatable
            if not my_key in self._contained_in.__class__._deco_nested_repeatables:
                msg = repr(my_key) + ': ' + repr(self) + ' is defined as repeatable, but this is not defined as a repeatable item in the containing class: ' + \
                    repr(self._contained_in.named_as())
                raise ConfigException(msg)
                # TODO?: type check of list items (instanceof(ConfigItem). Same type?

            if my_key in self._contained_in.attributes:
                # Validate that an attribute value of parent is not overridden by nested item (self),
                if not isinstance(self._contained_in.attributes[my_key], OrderedDict):
                    msg = repr(my_key) + ' is defined both as simple value and a contained item: ' + repr(self)
                    raise ConfigException(msg)
                # TODO?: type check of list items (instanceof(ConfigItem). Same type?

            # Insert in Ordered dict by 'id' or 'name', 'id' is preferred if given
            try:
                try:
                    obj_key = self._defaults['id']
                except KeyError:
                    obj_key = self._defaults['name']

                # Check that we are no replacing an object with the same id/name
                if self._contained_in.attributes[my_key].get(obj_key):
                    raise ConfigException("Re-used id/name " + repr(obj_key) + " in nested objects")
            except KeyError:
                obj_key = id(self)

            self._contained_in.attributes[my_key][obj_key] = self
            return

        if my_key in self._contained_in.attributes:
            if isinstance(self._contained_in.attributes[my_key], ConfigItem):
                raise ConfigException("Repeated non repeatable conf item: " + repr(my_key))
            if isinstance(self._contained_in.attributes[my_key], OrderedDict):
                msg = repr(my_key) + ': ' + repr(self) + ' is defined as non-repeatable, but the containing object has repatable items with the same name: ' + \
                    repr(self._contained_in)
                raise ConfigException(msg)
            raise ConfigException(repr(my_key) + ' is defined both as simple value and a contained item: ' + repr(self))

        self._contained_in.attributes[my_key] = self

class ConfigBuilder(_ConfigBase):
    def __init__(self, **attr):
        super(ConfigBuilder, self).__init__(**attr)

        # Set back reference to containing Item and root item
        self._contained_in = self.__class__.nested[-1]
        self._root_conf = self._contained_in.root_conf

    def __exit__(self, exc_type, exc_value, traceback):
        super(ConfigBuilder, self).__exit__(exc_type, exc_value, traceback)
        self.build()

    def build(self):
        """Override this in derived classes. This is where child ConfigItems are declared"""
        raise ConfigException("'build' must be overridded")

    def override(self, config_item):
        """Assign attributes from builder to child Item"""
        for key, value in self.attributes.iteritems():
            config_item_attr = config_item.attributes.get(key)
            if config_item_attr:
                config_item_attr.override(value)
                continue
            config_item.attributes[key] = value

<commit_msg>Fixed spelling error in error message<commit_after># Copyright 2012 Lars Hupfeldt Nielsen, Hupfeldt IT
# This code is free for anybody to use

from collections import Sequence, OrderedDict
from envs import Env
from attribute_collector import AttributeCollector
from config_errors import ConfigException, NoAttributeException

class _ConfigBase(object):
    nested = []

    # Decoration attributes
    _deco_named_as = None
    _deco_repeatable = False
    _deco_nested_repeatables = []
    _deco_required_attributes = []
    _deco_required_if_attributes = (None, ())

    def __init__(self, **attr):
        self._debug_exc = True

        # Object linking
        self._root_conf = None

        # Dict of dicts: attributes['a'] = dict(Env('prod')=1, EnvGroup('dev')=0)
        self._attributes = OrderedDict()
        for key in self.__class__._deco_nested_repeatables:
            self._attributes[key] = OrderedDict()

        self._defaults = attr
        self._frozen = False
        self._may_freeze = True

    def named_as(self):
        if self.__class__._deco_named_as:
            return self.__class__._deco_named_as
        if self.__class__._deco_repeatable:
            return self.__class__.__name__ + 's'
        return self.__class__.__name__

    def irepr(self, indent_level):
        """Indented repr"""
        indent1 = '  ' * indent_level
        indent2 =  indent1 + '     '
        # + ':' + self.__class__.__name__
        not_frozen_msg = "" if self._frozen else " not-frozen, defaults: " + repr(self._defaults)
        return self.named_as() + not_frozen_msg + ' {\n' + \
            ''.join([indent2 + name + ': ' + repr(value) + ',\n' for name, value in self.iteritems()]) + \
            indent1 + '}'

    def __repr__(self):
        return self.irepr(len(self.__class__.nested) -1)

    def __enter__(self):
        assert not self._frozen
        self._may_freeze = False
        return self

    def freeze_validate_required(self):
        missing = []
        for req in self.__class__._deco_required_attributes:
            if not req in self._attributes:
                missing.append(req)
        if missing:
            raise ConfigException("No value given for required attributes: " + repr(missing))

    def freeze_validate_required_if(self):
        required_if_key = self.__class__._deco_required_if_attributes[0]
        if not required_if_key:
            return

        try:
            required_if = self._attributes[required_if_key].value()
            if not required_if:
                return
        except KeyError:
            return
        except NoAttributeException:
            return

        missing = []
        for req in self.__class__._deco_required_if_attributes[1]:
            if not req in self._attributes:
                missing.append(req)
        if missing:
            raise ConfigException("Missing required_if attributes. Condition attribute: " + repr(required_if_key) + "==" +  repr(required_if) + ", missing: " + repr(missing))

    def freeze_validation(self):
        """Override this method if you need special checks"""
        self.freeze_validate_required()
        self.freeze_validate_required_if()

    def freeze(self):
        assert not self._frozen
        assert self._may_freeze

        # Collect remaining default values
        for name in list(self._defaults):
            AttributeCollector(name, self)()
        self._frozen = True

        try:
            self.freeze_validation()
            return self
        except ConfigException as ex:
            if self._debug_exc:
                raise
            # Strip stack
            raise ex
        except NoAttributeException as ex:
            if self._debug_exc:
                raise
            raise ex

    def __exit__(self, exc_type, exc_value, traceback):
        if exc_type:
            return None

        try:
            self._may_freeze = True
            self.freeze()
        except ConfigException as ex:
            if self._debug_exc:
                raise
            raise ex
        except NoAttributeException as ex:
            if self._debug_exc:
                raise
            raise ex

    def __setattr__(self, name, value):
        if name[0] == '_':
            # Needed to set private values in __init__
            super(_ConfigBase, self).__setattr__(name, value)
            return
        raise ConfigException("Trying to set a property " + repr(name) + " on a config item")

    def _check_valid_env(self, env, valid_envs):
        if type(env) != Env:
            raise ConfigException(self.__class__.__name__ + ': env must be of type ' + repr(Env.__name__) + '; found type ' + repr(env.__class__.__name__))

        for valid_env in valid_envs:
            if env in valid_env:
                return
        raise ConfigException("The env " + repr(env) + " must be in the (nested) list of valid_envs " + repr(valid_envs))

    def _env_specific_value(self, attr_name, attr_coll, env):
        if isinstance(attr_coll, ConfigItem) or isinstance(attr_coll, OrderedDict):
            return attr_coll

        try:
            return attr_coll.env_values[env]
        except KeyError:
            self._check_valid_env(env, self._root_conf._valid_envs)
            raise NoAttributeException("Attribute " + repr(attr_name) + " undefined for env " + repr(env))

    def getattr_env(self, name, env):
        try:
            attr_coll = self._attributes[name]
        except KeyError:
            error_message = ""
            repeatable_name = name + 's'
            if self._attributes.get(repeatable_name):
                error_message = ", but found attribute " + repr(repeatable_name)
            raise ConfigException(repr(self) + " has no attribute " + repr(name) + error_message)

        return self._env_specific_value(name, attr_coll, env)

    def __getattr__(self, name):
        if not self._frozen and self._may_freeze:
            self.freeze()

        if not self._frozen:
            try:
                # Return existing collector/dict of nested items if any
                coll = self._attributes[name]
                if type(coll) in (AttributeCollector, OrderedDict):
                    return coll
                if isinstance(coll, _ConfigItem):
                    return AttributeCollector(name, self)
                raise Exception("Internal error, unexpected type: " + repr(type(coll)) + ':' + repr(coll))
            except KeyError:
                return AttributeCollector(name, self)

        try:
            return self.getattr_env(name, self._root_conf._selected_env)
        except ConfigException as ex:
            if self._debug_exc:
                raise
            raise ex

    def iteritems(self):
        for key, value in self._attributes.iteritems():
            try:
                yield key, self._env_specific_value(key, value, self._root_conf._selected_env)
            except NoAttributeException:
                # This should only happen in case of  a conditional attribute
                pass

    def items(self):
        return list(self.iteritems())

    @property
    def contained_in(self):
        return self._contained_in

    @property
    def root_conf(self):
        return self._root_conf

    @property
    def attributes(self):
        return self._attributes

class _ConfigItem(_ConfigBase):
    def _validate(self):
        self.validate()
        for _child_name, child_value in self.iteritems():
            if isinstance(child_value, OrderedDict):
                for dict_entry in child_value.values():
                    dict_entry._validate()

            if isinstance(child_value, _ConfigItem):
                child_value._validate()

    def validate(self):
        """Can be overridden to provide post-frozen validation"""
        pass

    def __enter__(self):
        super(_ConfigItem, self).__enter__()
        self.__class__.nested.append(self)
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        self.__class__.nested.pop()
        super(_ConfigItem, self).__exit__(exc_type, exc_value, traceback)

class ConfigRoot(_ConfigItem):
    def __init__(self, selected_env, valid_envs, **attr):
        if not isinstance(valid_envs, Sequence) or isinstance(valid_envs, str):
            raise ConfigException(self.__class__.__name__ + ": valid_envs arg must be a 'Sequence'; found type " + repr(valid_envs.__class__.__name__))

        for env in valid_envs:
            if not isinstance(env, Env):
                raise ConfigException(self.__class__.__name__ + ": valid_envs items must be instance of 'Env'; found a " + repr(env.__class__.__name__))

        self._check_valid_env(selected_env, valid_envs)

        self._selected_env = selected_env
        self._valid_envs = valid_envs
        super(ConfigRoot, self).__init__(**attr)
        self._root_conf = self
        self._contained_in = None

    def __exit__(self, exc_type, exc_value, traceback):
        super(ConfigRoot, self).__exit__(exc_type, exc_value, traceback)
        self._validate()

    @property
    def selected_env(self):
        return self._selected_env

class ConfigItem(_ConfigItem):
    def __init__(self, **attr):
        super(ConfigItem, self).__init__(**attr)

        # Automatic Nested Insert in parent
        if not self.__class__.nested:
            raise ConfigException(self.__class__.__name__ + " object must be nested (indirectly) in a " + repr(ConfigRoot.__name__))

        # Set back reference to containing Item and root item
        self._contained_in = self.__class__.nested[-1]
        self._root_conf = self._contained_in.root_conf

        # Insert self in containing Item's attributes
        my_key = self.named_as()

        if self.__class__._deco_repeatable:
            # Validate that the containing item has specified this item as repeatable
            if not my_key in self._contained_in.__class__._deco_nested_repeatables:
                msg = repr(my_key) + ': ' + repr(self) + ' is defined as repeatable, but this is not defined as a repeatable item in the containing class: ' + \
                    repr(self._contained_in.named_as())
                raise ConfigException(msg)
                # TODO?: type check of list items (instanceof(ConfigItem). Same type?

            if my_key in self._contained_in.attributes:
                # Validate that an attribute value of parent is not overridden by nested item (self),
                if not isinstance(self._contained_in.attributes[my_key], OrderedDict):
                    msg = repr(my_key) + ' is defined both as simple value and a contained item: ' + repr(self)
                    raise ConfigException(msg)
                # TODO?: type check of list items (instanceof(ConfigItem). Same type?

            # Insert in Ordered dict by 'id' or 'name', 'id' is preferred if given
            try:
                try:
                    obj_key = self._defaults['id']
                except KeyError:
                    obj_key = self._defaults['name']

                # Check that we are no replacing an object with the same id/name
                if self._contained_in.attributes[my_key].get(obj_key):
                    raise ConfigException("Re-used id/name " + repr(obj_key) + " in nested objects")
            except KeyError:
                obj_key = id(self)

            self._contained_in.attributes[my_key][obj_key] = self
            return

        if my_key in self._contained_in.attributes:
            if isinstance(self._contained_in.attributes[my_key], ConfigItem):
                raise ConfigException("Repeated non repeatable conf item: " + repr(my_key))
            if isinstance(self._contained_in.attributes[my_key], OrderedDict):
                msg = repr(my_key) + ': ' + repr(self) + ' is defined as non-repeatable, but the containing object has repeatable items with the same name: ' + \
                    repr(self._contained_in)
                raise ConfigException(msg)
            raise ConfigException(repr(my_key) + ' is defined both as simple value and a contained item: ' + repr(self))

        self._contained_in.attributes[my_key] = self

class ConfigBuilder(_ConfigBase):
    def __init__(self, **attr):
        super(ConfigBuilder, self).__init__(**attr)

        # Set back reference to containing Item and root item
        self._contained_in = self.__class__.nested[-1]
        self._root_conf = self._contained_in.root_conf

    def __exit__(self, exc_type, exc_value, traceback):
        super(ConfigBuilder, self).__exit__(exc_type, exc_value, traceback)
        self.build()

    def build(self):
        """Override this in derived classes. This is where child ConfigItems are declared"""
        raise ConfigException("'build' must be overridded")

    def override(self, config_item):
        """Assign attributes from builder to child Item"""
        for key, value in self.attributes.iteritems():
            config_item_attr = config_item.attributes.get(key)
            if config_item_attr:
                config_item_attr.override(value)
                continue
            config_item.attributes[key] = value

<|endoftext|>


<commit_before>"""
A small OOP-oriented wrapper for pyspeckit,
for extended flexibility with input guesses,
model selection, and multiple component fits.
"""
import numpy as np
import pyspeckit

# TODO: make informative log/on-screen messages
#       about what's being done to the subcubes

class SubCube(pyspeckit.Cube):
    """
    An extention of Cube, tinkered to be an instance of MultiCube, from which
    it receives references to instances of pyspeckit.Cube that do not depend
    on a spectral model chosen (so that parent MultiCube doesn't weigh so much)

    Is designed to have methods that operate within a single spectral model.
    """
    def __init___(self, *args, **kwargs):
        super(pyspeckit.Cube, self).__init__(*args, **kwargs)
        # so I either define some things as `None`
        # or I'll have to call hasattr or them...
        # TODO: which is a more Pythonic approach?
        # A: see here  http://programmers.stackexchange.com/questions/
        #             254576/is-it-a-good-practice-to-declare-instance
        #                   -variables-as-none-in-a-class-in-python
        self.guess_grid = None
        self.model_grid = None

    def update_model(self, fit_type='gaussian'):
        """
        Tie a model to a SubCube. Didn't test it
        on anything but gaussian fitter so far.
        """
        try:
            allowed_fitters = self.specfit.Registry.multifitters
            self.specfit.fitter = allowed_fitters[fit_type]
        except KeyError:
            raise ValueError('Unknown fit type: %s\n'
                             'Choose one from %s'
                             % (fit_type, allowed_fitters.keys()))
        self.specfit.fittype = fit_type
        self.fittype = fit_type

    def generate_model(self, guess_grid=None):
        """
        Generates a grid of spectral models matching the
        shape of the input guess_grid array. Can take the
        following numpy arrays as an input:

        Parameters
        ----------
        guess_grid : numpy.array
                     A grid of input parameters.
                     Can be one of the following:
                     1) An (M,)-shaped 1d array of model parameters
                        (see pyspeckit docs for details)
                     2) An (N, M) array to compute N models
                        for M sets of parameters
                     3) A guess cube of (Y, X, M) size
                     4) An (N, Y, X, M)-shaped array, to
                        iterate over cubes of guesses.

        WARNING: (3) and (4) aren't implemented yet.

        Returns
        -------
        model_grid : a grid of spectral models, following the
                     shape of guess_grid. Also saved as an
                     instance under SubCube.model_grid
        """
        # NOTE: 2016.03.22: currently works for (a) and (c), see below

        # TODO: test for all posible model grid sizes,
        #       and write up proper error handling!
        #       things you should be able receive:
        #        a) a single par-iterable
        #        b) an XY collection of pars
        #        c) an N*par collection of different guesses
        #        d) an N*XY*par mess: should not proceed without
        #           vectorization! it will simply be too slow :/
        if guess_grid is None:
            # TODO: what to do if the guess grid is not in place?
            guess_grid = self.guess_grid

        # TODO: this if/elif Christmas tree is ugly,
        #       there should be a smarter way to do this
        if len(guess_grid.shape)==4 \
           and self.cube.shape[1] in guess_grid.shape \
           and self.cube.shape[1] in guess_grid.shape:
            # FIXME: this allows for [100,2,4] guess_grid to pass
            #        for an  [xarr.size, 100, 100] cube. Do it right!
            #
            # TODO: implement cube-like guessing grids!

            # this is a good place to re-implement get_modelcube
            # as a method that takes modelcube as a function,
            # or just bait and switch guesses into inherited
            # Cube.parcube and call get_modelcube . . .

            # Will do the former for now, but it may clash later
            # when I will finish up MultiCube.multiplot()

            raise NotImplementedError('Someone should remind me '
                                      'to write this up. Please?')
        if len(guess_grid.shape)==3 \
           and self.cube.shape[1] in guess_grid.shape \
           and self.cube.shape[1] in guess_grid.shape:
            # FIXME: this allows for [100,2,4] guess_grid to pass
            #        for an [xarr.size, 100, 100] cube. Do it right!
            #
            yy, xx = np.indices(self.cube.shape[1:])
            model_grid = np.zeros_like(self.cube)
            # TODO: vectorize this please please please?
            for x,y in zip(xx.flat,yy.flat):
                model_grid[:,y,x] = \
                       self.specfit.get_full_model(pars=self.guess_grid[:,y,x])

        elif len(guess_grid.shape)==2:
            # set up the modelled spectrum grid
            model_grid = np.empty(shape=(guess_grid.shape[0],
                                         self.xarr.size      ))
            for i, par in enumerate(guess_grid):
                model_grid[i] = self.specfit.get_full_model(pars=par)
        elif len(guess_grid.shape)==1:
            par = guess_grid
            model_grid = self.specfit.get_full_model(pars=par)
        else:
            raise IndexError('Guess grid size can not be matched'
                             ' to either cube or spectrum size. ')

        self.model_grid = model_grid
        return model_grid

    def best_guess(self, model_grid=None, xy_list=None, sn_cut=None):
        """
        For a grid of intitial guesses, determine the optimal one based
        on the preliminary residual of the specified spectral model.

        Parameters
        ----------
        model_grid : numpy.array
                     A model grid to choose from.

        use_cube : boolean
                   If true, every xy-slice of a cube will be
                   compared to every model from the model_grid.
                   sn_cut (see below) is still applied.

        xy_list : iterable
                  A collection of positions on the data cube
                  which to check for the lowest residuals.
                  Ignored if use_cube was set to `True`
                  Actually, I guess I should use a mask for this...

        sn_cut : float
                 Ignore items on xy_list if the corresponding
                 spectra have too low signal-to-noise ratios.

        """
        if model_grid is None:
            if self.model_grid is None:
                raise TypeError('sooo the model_grid is empty, '
                                'did you run generate_model()?')
            model_grid = self.model_grid
        # TODO: scale this up later for MultiCube.judge() method
        #       to include the deviance information criterion, DIC
        #       (c.f Kunz et al. 2006 and Sebastian's IMPRS slides)

        # TODO break-down:
        #
        # + First milestone  : make the blind iteration over the spectra,
        #                      have some winner-takes-it-all selection in place
        # - Second milestone : implement the (still broken) snr mask
        # - Third milestone  : xy_list functionality, pixel coordinates only
        # - Fourth milestone : extend to all possible model_grid shapes
        #
        if model_grid.shape!=self.cube.shape:
            raise NotImplementedError
        else:
            # commence the invasion!
            residual = (self.cube - model_grid).std(axis=0)
            self.residual = residual
            vmin, (xmin,ymin) = residual.min(), np.where(residual==residual.min())
            print "Best guess at %.2f on (%i, %i)" % (vmin, xmin, ymin)
            return vmin, (xmin, ymin)

    def get_snr_map(self):
        """
        Calculates S/N ratio for the cube.

        Returns
        -------
        snr_map : numpy.array
                  Also stored under SubCube.snr_map
        """
        snr_map = self.get_signal_map(self) / self.get_rms_map(self)
        self._snr_map = snr_map
        return snr_map

    def get_rms_map(self, noise_heads=None, noise_tails=None, unit='km/s'):
        """
        Make an rms estimate, will try to find the noise channels in
        the input values or in class instances. If noise channels are
        not identified, defaults to calculating rms of all channels.
        """
        # TODO: implement sigal masks!
        #       for a gaussian, they can be
        #       calculated from guess_grids!

        rms_map = self.cube.std(axis=0) # this is so wroooooooong I can't even
        self._rms_map = rms_map
        return rms_map

    def get_signal_map(self, signal_heads=None,
                       signal_tails=None, unit='km/s'):
        """
        Make a signal strength estimate. Currently just selects maximum
        value of the all channels, mimicking pyspeckit approach.
        Will try to improve it soon.
        """
        # TODO: implement sigal masks!
        #       for a gaussian, they can be
        #       calculated from guess_grids!
        signal_map = self.cube.max(axis=0)
        self._signal_map = signal_map
        return signal_map

class MultiCube:
    def __init__(self, *args):
        """
        A collection of Specfit objects mapped to SubCubes
        by the mapper method. Includes* methods to fit multiple
        guess grids for different models, and the means to
        decide between the results of those fits.

        *In a distant[citation needed] future.

        Input parameters: see ~pyspeckit.Cube
        """

        # parent cube, used for attribute propagation
        self.supercube = pyspeckit.Cube(*args)

        # making a bunch of references to make our life easier
        self.cube = self.SuperCube.cube
        self.xarr = self.SuperCube.xarr
        self.header = self.SuperCube.header

        # FIXME: rewrite mapplot to include the mapper/judge methods!
        #        doesn't work in its current implementation, will need
        #        to rewire different params and errors, for variable
        #        number of parameters across different models
        self.multiplot = self.SuperCube.mapplot

        # MultiCube's own instances:
        self.multigrid = {}
        self.tesseract = {}

    def __repr__(self):
        return ('Parent: MultiCube with TODO models\n'
                'Child: %s' % self.SuperCube.__repr__())

    def spawn(self, model, guesses=None):
        """
        Add a new model and a SubCube for it thorugh Cube()
        The idea is to pass a reference to large data instances
        of SuperCube to avoid excessive memory usage.

        Not implemented yet.
        """
        self.tesseract[model]=SubCube()
        raise NotImplementedError

    def mapper(model):
        """
        Returns a list of SubCubes for a given model?
        """
        raise NotImplementedError

    def judge_multimodel(subcubes, model, method):
        """
        Decide which model to use.
        First milestone: have one component added only
                         when residual has SNR>3
        Actual goal: proper model selection via DIC.
        """
        raise NotImplementedError

    def multifit(self, multigrid=None):
        """
        Fit the optimized guesses. This should be delegated
        to SubCubes maybe? MultiCube should only call the
        judge function.

        Not really, this approach would allow to juggle all
        the SubCubes defined! In this case, multifit is a
        wrapper for SubCube.fiteach() method. This will do.
        """
        raise NotImplementedError

# NOTE: a working example, generates a grid
#       of spectra form a grid of parameters
#
# TODO: remove this, this isn't one of your km-long scripts!
def main():
    try:
        sc = SubCube('foo.fits')
    except IOError:
        from astro_toolbox import make_test_cube
        make_test_cube((100,10,10), outfile='foo.fits', sigma=(10,5))
        sc = SubCube('foo.fits')

    sc.update_model('gaussian')
    guesses = [0.5, 0.2, 0.8]

    npars = len(guesses)
    parcube_size = sc.cube.size/sc.shape[0]
    parcube_shape = (npars, sc.cube.shape[1], sc.cube.shape[2])
    parflat = np.hstack(np.repeat([guesses],parcube_size,axis=0))
    sc.guess_grid = parflat.reshape(*parcube_shape,order='F')
    sc.generate_model()
    sc.best_guess()

    #sc.fiteach(guesses=sc.guess_grid, start_from_pixel=(5,5))
    #sc.mapplot()

if __name__ == "__main__":
	main()
<commit_msg>Pinpointing issues and bugs<commit_after>"""
A small OOP-oriented wrapper for pyspeckit,
for extended flexibility with input guesses,
model selection, and multiple component fits.
"""
import numpy as np
import pyspeckit

# TODO: make informative log/on-screen messages
#       about what's being done to the subcubes

class SubCube(pyspeckit.Cube):
    """
    An extention of Cube, tinkered to be an instance of MultiCube, from which
    it receives references to instances of pyspeckit.Cube that do not depend
    on a spectral model chosen (so that parent MultiCube doesn't weigh so much)

    Is designed to have methods that operate within a single spectral model.
    """
    def __init___(self, *args):
        super(SubCube, self).__init__(*args)

        # FIXME: why isn't the block below executed?
        #        probably I don't get the super inheritance
        raise BadBadError

        # because that UnitConversionError pops up way too often
        if self.xarr.velocity_convention is None:
            self.xarr.velocity_convention = 'radio'

        # so I either define some things as `None`
        # or I'll have to call hasattr or them...
        # TODO: which is a more Pythonic approach?
        # A: probably the hasattr method, see here:
        # http://programmers.stackexchange.com/questions/
        # 254576/is-it-a-good-practice-to-declare-instance
        # -variables-as-none-in-a-class-in-python
        self.guess_grid = None
        self.model_grid = None

    def update_model(self, fit_type='gaussian'):
        """
        Tie a model to a SubCube. Didn't test it
        on anything but gaussian fitter so far.
        """
        try:
            allowed_fitters = self.specfit.Registry.multifitters
            self.specfit.fitter = allowed_fitters[fit_type]
        except KeyError:
            raise ValueError('Unknown fit type: %s\n'
                             'Choose one from %s'
                             % (fit_type, allowed_fitters.keys()))
        self.specfit.fittype = fit_type
        self.fittype = fit_type

    def generate_model(self, guess_grid=None):
        """
        Generates a grid of spectral models matching the
        shape of the input guess_grid array. Can take the
        following numpy arrays as an input:

        Parameters
        ----------
        guess_grid : numpy.array
                     A grid of input parameters.
                     Can be one of the following:
                     1) An (M,)-shaped 1d array of model parameters
                        (see pyspeckit docs for details)
                     2) An (N, M) array to compute N models
                        for M sets of parameters
                     3) A guess cube of (Y, X, M) size
                     4) An (N, Y, X, M)-shaped array, to
                        iterate over cubes of guesses.

        WARNING: (3) and (4) aren't implemented yet.

        Returns
        -------
        model_grid : a grid of spectral models, following the
                     shape of guess_grid. Also saved as an
                     instance under SubCube.model_grid
        """
        # NOTE: 2016.03.22: currently works for (a) and (c), see below

        # TODO: test for all posible model grid sizes,
        #       and write up proper error handling!
        #       things you should be able receive:
        #        a) a single par-iterable
        #        b) an XY collection of pars
        #        c) an N*par collection of different guesses
        #        d) an N*XY*par mess: should not proceed without
        #           vectorization! it will simply be too slow :/
        if guess_grid is None:
            # TODO: what to do if the guess grid is not in place?
            guess_grid = self.guess_grid

        # TODO: this if/elif Christmas tree is ugly,
        #       there should be a smarter way to do this
        if len(guess_grid.shape)==4 \
           and self.cube.shape[1] in guess_grid.shape \
           and self.cube.shape[1] in guess_grid.shape:
            # FIXME: this allows for [100,2,4] guess_grid to pass
            #        for an  [xarr.size, 100, 100] cube. Do it right!
            #
            # TODO: implement cube-like guessing grids!

            # this is a good place to re-implement get_modelcube
            # as a method that takes modelcube as a function,
            # or just bait and switch guesses into inherited
            # Cube.parcube and call get_modelcube . . .

            # Will do the former for now, but it may clash later
            # when I will finish up MultiCube.multiplot()

            raise NotImplementedError('Someone should remind me '
                                      'to write this up. Please?')
        if len(guess_grid.shape)==3 \
           and self.cube.shape[1] in guess_grid.shape \
           and self.cube.shape[1] in guess_grid.shape:
            # FIXME: this allows for [100,2,4] guess_grid to pass
            #        for an [xarr.size, 100, 100] cube. Do it right!
            #
            yy, xx = np.indices(self.cube.shape[1:])
            model_grid = np.zeros_like(self.cube)

            # TODO: vectorize this please please please?
            for x,y in zip(xx.flat,yy.flat):
                model_grid[:,y,x] = \
                       self.specfit.get_full_model(pars=self.guess_grid[:,y,x])

        elif len(guess_grid.shape)==2:
            # set up the modelled spectrum grid
            model_grid = np.empty(shape=(guess_grid.shape[0],
                                         self.xarr.size      ))
            for i, par in enumerate(guess_grid):
                model_grid[i] = self.specfit.get_full_model(pars=par)
        elif len(guess_grid.shape)==1:
            par = guess_grid
            model_grid = self.specfit.get_full_model(pars=par)
        else:
            raise IndexError('Guess grid size can not be matched'
                             ' to either cube or spectrum size. ')

        self.model_grid = model_grid
        return model_grid

    def best_guess(self, model_grid=None, xy_list=None, sn_cut=None):
        """
        For a grid of intitial guesses, determine the optimal one based
        on the preliminary residual of the specified spectral model.

        Parameters
        ----------
        model_grid : numpy.array
                     A model grid to choose from.

        use_cube : boolean
                   If true, every xy-slice of a cube will be
                   compared to every model from the model_grid.
                   sn_cut (see below) is still applied.

        xy_list : iterable
                  A collection of positions on the data cube
                  which to check for the lowest residuals.
                  Ignored if use_cube was set to `True`
                  Actually, I guess I should use a mask for this...

        sn_cut : float
                 Ignore items on xy_list if the corresponding
                 spectra have too low signal-to-noise ratios.

        """
        if model_grid is None:
            if self.model_grid is None:
                raise TypeError('sooo the model_grid is empty, '
                                'did you run generate_model()?')
            model_grid = self.model_grid
        # TODO: scale this up later for MultiCube.judge() method
        #       to include the deviance information criterion, DIC
        #       (c.f Kunz et al. 2006 and Sebastian's IMPRS slides)

        # TODO break-down:
        #
        # + First milestone  : make the blind iteration over the spectra,
        #                      have some winner-takes-it-all selection in place
        # - Second milestone : implement the (still broken) snr mask
        # - Third milestone  : xy_list functionality, pixel coordinates only
        # - Fourth milestone : extend to all possible model_grid shapes
        #
        if model_grid.shape!=self.cube.shape:
            raise NotImplementedError
        else:
            # commence the invasion!
            residual = (self.cube - model_grid).std(axis=0)
            if ~np.isnan(residual).any():
                # FIXME: throw warning for all-NaN case
                return
            self.residual = residual
            vmin, (xmin,ymin) = residual.min(), np.where(residual==residual.min())
            print "Best guess at %.2f on (%i, %i)" % (vmin, xmin, ymin)
            return vmin, (xmin, ymin)

    def get_snr_map(self):
        """
        Calculates S/N ratio for the cube.

        Returns
        -------
        snr_map : numpy.array
                  Also stored under SubCube.snr_map
        """
        snr_map = self.get_signal_map(self) / self.get_rms_map(self)
        self._snr_map = snr_map
        return snr_map

    def get_rms_map(self, noise_heads=None, noise_tails=None, unit='km/s'):
        """
        Make an rms estimate, will try to find the noise channels in
        the input values or in class instances. If noise channels are
        not identified, defaults to calculating rms of all channels.
        """
        # TODO: implement sigal masks!
        #       for a gaussian, they can be
        #       calculated from guess_grids!

        rms_map = self.cube.std(axis=0) # this is so wroooooooong I can't even
        self._rms_map = rms_map
        return rms_map

    def get_signal_map(self, signal_heads=None,
                       signal_tails=None, unit='km/s'):
        """
        Make a signal strength estimate. Currently just selects maximum
        value of the all channels, mimicking pyspeckit approach.
        Will try to improve it soon.
        """
        # TODO: implement sigal masks!
        #       for a gaussian, they can be
        #       calculated from guess_grids!
        signal_map = self.cube.max(axis=0)
        self._signal_map = signal_map
        return signal_map

class MultiCube:
    def __init__(self, *args):
        """
        A collection of Specfit objects mapped to SubCubes
        by the mapper method. Includes* methods to fit multiple
        guess grids for different models, and the means to
        decide between the results of those fits.

        *In a distant[citation needed] future.

        Input parameters: see ~pyspeckit.Cube
        """

        # parent cube, used for attribute propagation
        self.supercube = pyspeckit.Cube(*args)

        # making a bunch of references to make our life easier
        self.cube = self.SuperCube.cube
        self.xarr = self.SuperCube.xarr
        self.header = self.SuperCube.header

        # FIXME: rewrite mapplot to include the mapper/judge methods!
        #        doesn't work in its current implementation, will need
        #        to rewire different params and errors, for variable
        #        number of parameters across different models
        self.multiplot = self.SuperCube.mapplot

        # MultiCube's own instances:
        self.multigrid = {}
        self.tesseract = {}

    def __repr__(self):
        return ('Parent: MultiCube with TODO models\n'
                'Child: %s' % self.SuperCube.__repr__())

    def spawn(self, model, guesses=None):
        """
        Add a new model and a SubCube for it thorugh Cube()
        The idea is to pass a reference to large data instances
        of SuperCube to avoid excessive memory usage.

        Not implemented yet.
        """
        self.tesseract[model]=SubCube()
        raise NotImplementedError

    def mapper(model):
        """
        Returns a list of SubCubes for a given model?
        """
        raise NotImplementedError

    def judge_multimodel(subcubes, model, method):
        """
        Decide which model to use.
        First milestone: have one component added only
                         when residual has SNR>3
        Actual goal: proper model selection via DIC.
        """
        raise NotImplementedError

    def multifit(self, multigrid=None):
        """
        Fit the optimized guesses. This should be delegated
        to SubCubes maybe? MultiCube should only call the
        judge function.

        Not really, this approach would allow to juggle all
        the SubCubes defined! In this case, multifit is a
        wrapper for SubCube.fiteach() method. This will do.
        """
        raise NotImplementedError

# NOTE: a working example, generates a grid
#       of spectra form a grid of parameters
#
# TODO: remove this, this isn't one of your km-long scripts!
def main():
    try:
        sc = SubCube('foo.fits')
    except IOError:
        from astro_toolbox import make_test_cube
        make_test_cube((100,10,10), outfile='foo.fits', sigma=(10,5))
        sc = SubCube('foo.fits')

    sc.update_model('gaussian')
    guesses = [0.5, 0.2, 0.8]

    npars = len(guesses)
    parcube_size = sc.cube.size/sc.shape[0]
    parcube_shape = (npars, sc.cube.shape[1], sc.cube.shape[2])
    parflat = np.hstack(np.repeat([guesses],parcube_size,axis=0))
    sc.guess_grid = parflat.reshape(*parcube_shape,order='F')

    sc.generate_model()
    sc.best_guess()

    #sc.fiteach(guesses=sc.guess_grid, start_from_pixel=(5,5))
    #sc.mapplot()

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


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

from Qt import QtWidgets
import time
from contextlib import contextmanager

def get_maya_window(cache=[]):
    '''Get Maya MainWindow as a QWidget.'''

    if cache:
        return cache[0]

    for widget in QtWidgets.QApplication.instance().topLevelWidgets():
        if widget.objectName() == 'MayaWindow':
            cache.append(widget)
            return widget

    raise RuntimeError('Could not locate MayaWindow...')

def wait(delay=1):
    '''Delay python execution for a specified amount of time'''

    app = QtWidgets.QApplication.instance()

    s = time.clock()
    while True:
        if time.clock() - s >= delay:
            return

        app.processEvents()

@contextmanager
def viewport_state(viewport, state):
    '''Sets a viewports state options for the duration of the context.

    Example:

        # Turn off the display of nurbsCurves
        viewport = Viewport.active()
        state = viewport.state()
        state['nurbsCurves'] = False

        with viewport_state(viewport, state):
            # Do something with nurbsCurves off
    '''

    previous_state = viewport.get_state()
    try:
        viewport.set_state(state)
        yield
    finally:
        viewport.set_state(previous_state)
<commit_msg>Ensure Qt imported from vendor.<commit_after># -*- coding: utf-8 -*-

import time
from contextlib import contextmanager

from .vendor.Qt import QtWidgets

def get_maya_window(cache=[]):
    '''Get Maya MainWindow as a QWidget.'''

    if cache:
        return cache[0]

    for widget in QtWidgets.QApplication.instance().topLevelWidgets():
        if widget.objectName() == 'MayaWindow':
            cache.append(widget)
            return widget

    raise RuntimeError('Could not locate MayaWindow...')

def wait(delay=1):
    '''Delay python execution for a specified amount of time'''

    app = QtWidgets.QApplication.instance()

    s = time.clock()
    while True:
        if time.clock() - s >= delay:
            return

        app.processEvents()

@contextmanager
def viewport_state(viewport, state):
    '''Sets a viewports state options for the duration of the context.

    Example:

        # Turn off the display of nurbsCurves
        viewport = Viewport.active()
        state = viewport.state()
        state['nurbsCurves'] = False

        with viewport_state(viewport, state):
            # Do something with nurbsCurves off
    '''

    previous_state = viewport.get_state()
    try:
        viewport.set_state(state)
        yield
    finally:
        viewport.set_state(previous_state)
<|endoftext|>


<commit_before>import random
import time
from collections import OrderedDict
from typing import List, Union
from weakref import WeakKeyDictionary

from twisted.plugin import IPlugin
from zope.interface import implementer

from desertbot.message import IRCMessage, TargetTypes
from desertbot.moduleinterface import IModule
from desertbot.modules.commandinterface import BotCommand
from desertbot.response import IRCResponse
from desertbot.utils import string

try:
    import re2
except ImportError:
    import re as re2

@implementer(IPlugin, IModule)
class OutOfContext(BotCommand):
    messageLimit = 50

    def actions(self):
        return super(OutOfContext, self).actions() + [("message-channel", 1, self.storeMessage),
                                                      ("action-channel", 1, self.storeMessage)]

    def triggers(self):
        return ["outofcontext"]

    def _add(self, message: IRCMessage):
        """add - Adds a quote to the OutOfContext log. The quote will be pulled from a message line buffer."""
        if len(message.parameterList) < 2:
            return IRCResponse("Add what?", message.replyTo)
        if message.targetType == TargetTypes.USER:
            return IRCResponse("You can only add messages from channels.", message.replyTo)

        regex = re2.compile(re2.escape(" ".join(message.parameterList[1:])), re2.IGNORECASE)
        if len(self.messageStore) == 0 or message.channel not in self.messageStore:
            return IRCResponse("Sorry, there are no messages in my buffer.", message.replyTo)

        matches = list(filter(regex.search, self.messageStore[message.channel]))
        if len(matches) == 0:
            return IRCResponse("Sorry, that didn't match anything in my message buffer.", message.replyTo)
        if len(matches) > 1:
            return IRCResponse("Sorry, that matches too many lines in my message buffer.", message.replyTo)

        todayDate = time.strftime("[%Y-%m-%d] [%H:%M]")
        quote = f"{todayDate} {matches[0]}"
        if message.replyTo not in self.storage:
            self.storage[message.replyTo] = []
        if len(self.storage[message.replyTo]) > 0 and self.storage[message.replyTo][-1] == quote:
            return IRCResponse("That quote has already been added to the log!", message.replyTo)
        else:
            self.storage[message.replyTo].append(quote)
            return IRCResponse(f"Quote '{quote}' was added to the log!", message.replyTo)

    def _remove(self, message: IRCMessage):
        """remove <regex> - remove a quote from the OutOfContext log."""
        if len(message.parameterList) < 2:
            return IRCResponse("Remove what?", message.replyTo)
        if len(self.storage) == 0 or message.replyTo not in self.storage:
            return IRCResponse("There are no quotes in the log.", message.replyTo)
        regex = re2.compile(" ".join(message.parameterList[1:]), re2.IGNORECASE)
        matches = list(filter(regex.search, self.storage[message.replyTo]))
        if len(matches) == 0:
            return IRCResponse("That message is not in the log.", message.replyTo)
        if len(matches) > 1:
            return IRCResponse(f"Unable to remove quote, {len(matches)} matches found.", message.replyTo)
        return self._removeQuote(message.replyTo, matches[0])

    def _removebyid(self, message: IRCMessage):
        """removebyid <quoteid> - remove the quote with the specified ID from the OutOfContext log."""
        if len(message.parameterList) < 2:
            return IRCResponse("Remove what?", message.replyTo)
        if not string.isNumber(message.parameterList[1]):
            return IRCResponse("You didn't specify a valid ID.", message.replyTo)
        if len(self.storage) == 0 or message.replyTo not in self.storage:
            return IRCResponse("There are no quotes in the log.", message.replyTo)
        index = int(message.parameterList[1]) - 1
        quotes = self.storage[message.replyTo]
        if index < len(quotes):
            return self._removeQuote(message.replyTo, quotes[index])
        return IRCResponse("That message is not in the log.", message.replyTo)

    def _list(self, message: IRCMessage):
        """list (<search/searchnick <regex>>) - post the OutOfContext log. A search regex can be provided to filter\
         the list."""
        params = [x for x in message.parameterList[1:]]
        if len(params) > 0:
            subsubcommand = params.pop(0).lower()
            if subsubcommand == "searchnick":
                return self._postList(message.replyTo, " ".join(params), True)
            elif subsubcommand == "search":
                return self._postList(message.replyTo, " ".join(params), False)
            else:
                return self._postList(message.replyTo, "", False)
        else:
            return self._postList(message.replyTo, "", False)

    def _search(self, message: IRCMessage):
        """search <regex> - look up quotes in the OutOfContext log. This search operation will look at \
        the content of the quotes."""
        if len(message.parameterList) < 2:
            return IRCResponse("Search what?", message.replyTo)
        return self._getQuote(message.replyTo, " ".join(message.parameterList[1:]), False, -1)

    def _searchnick(self, message: IRCMessage):
        """searchnick <regex> - look up quotes in the OutOfContext log. This search operation will \
        look at the nick that said the quoted line."""
        nick = message.user.nick
        if len(message.parameterList) > 1:
            nick = message.parameterList[1]
        return self._getQuote(message.replyTo, nick, True, -1)

    def _id(self, message: IRCMessage):
        """id <quoteid> - look up the quote that has the given ID."""
        if len(message.parameterList) < 2 or not string.isNumber(message.parameterList[1]):
            return IRCResponse("You didn't specify a valid ID.", message.replyTo)
        return self._getQuote(message.replyTo, "", False, int(message.parameterList[1]) - 1)

    def _random(self, message: IRCMessage):
        """random - return a random quote from the OutOfContext log."""
        return self._getQuote(message.replyTo, "", False, -1)

    def _removeQuote(self, source, quote):
        self.storage[source].remove(quote)
        return IRCResponse(f"Quote '{quote}' was removed from the log!", source)

    def _postList(self, source, searchString, searchNickname):
        if len(self.storage) == 0 or source not in self.storage:
            return IRCResponse("There are no quotes in the log.", source)
        regex = re2.compile(searchString, re2.IGNORECASE)
        matches = []
        if searchNickname:
            for x in self.storage[source]:
                if x[21] == "*":
                    match = re2.search(regex, x[:x.find(" ", 23)])
                else:
                    match = re2.search(regex, x[x.find("<") + 1:x.find(">")])
                if match:
                    matches.append(x)
        else:
            for x in self.storage[source]:
                if re2.search(regex, x[x.find(">") + 1:]):
                    matches.append(x)
        if len(matches) == 0:
            return IRCResponse(f"No matches for '{searchString}' found.", source)

        pasteLink =  self.bot.moduleHandler.runActionUntilValue('upload-dbco',
                                                                string.stripFormatting("\n".join(matches)),
                                                                10 * 60)

        return IRCResponse(f"Link posted! (Expires in 10 minutes) {pasteLink}.", source)

    def _getQuote(self, source, searchString, searchNickname, index):
        if len(self.storage) == 0 or source not in self.storage:
            return IRCResponse("There are no quotes in the log.", source)
        regex = re2.compile(searchString, re2.IGNORECASE)
        matches = []
        if searchNickname:
            for x in self.storage[source]:
                if x[21] == "*":
                    match = re2.search(regex, x[:x.find(" ", 23)])
                else:
                    match = re2.search(regex, x[x.find("<") + 1:x.find(">")])
                if match:
                    matches.append(x)
        else:
            for x in self.storage[source]:
                if re2.search(regex, x[x.find(">") + 1:]):
                    matches.append(x)
        if len(matches) == 0:
            return IRCResponse(f"No matches for '{searchString}' found.", source)
        if index < 0 or index > len(matches) - 1:
            index = random.randint(0, len(matches) - 1)
        return IRCResponse(f"Quote #{index + 1}/{len(matches)}: {matches[index]}", source)

    def _unrecognizedSubcommand(self, subCommand):
        return (f"unrecognized subcommand f'{subCommand}', "
                f"available subcommands for alias are: {', '.join(self.subCommands.keys())}")

    def onLoad(self) -> None:
        self.messageStore = WeakKeyDictionary()
        self._helpText = ("{1}outofcontext ({0}) - adds, removes or requests things from the OutOfContext quote list. "
                          "Use '{1}help outofcontext <subcommand>' for subcommand help. "
                          .format('/'.join(self.subCommands), self.bot.commandChar))

    def help(self, query: Union[List[str], None]) -> str:
        if len(query) > 1:
            subCommand = query[1].lower()
            if subCommand in self.subCommands:
                return ('{1}outofcontext {0}' .format(re2.sub(r"\s+", " ", self.subCommands[subCommand].__doc__),
                        self.bot.commandChar))
            else:
                return self._unrecognizedSubcommand(subCommand)
        else:
            return self._helpText

    def execute(self, message: IRCMessage):
        if len(message.parameterList) > 0:
            subCommand = message.parameterList[0].lower()
            if subCommand not in self.subCommands:
                return IRCResponse(self._unrecognizedSubcommand(subCommand), message.replyTo)
            return self.subCommands[subCommand](self, message)
        else:
            return IRCResponse(self._helpText, message.replyTo)

    def storeMessage(self, message: IRCMessage):
        if message.command or message.targetType == TargetTypes.USER:
            return

        if message.channel not in self.messageStore:
            self.messageStore[message.channel] = []

        if message.type == 'ACTION':
            self.messageStore[message.channel]\
                .append(f'* {message.user.nick} {string.stripFormatting(message.messageString)}')
        else:
            self.messageStore[message.channel]\
                .append(f'<{message.user.nick}> {string.stripFormatting(message.messageString)}')

        if len(self.messageStore[message.channel]) > self.messageLimit:
            self.messageStore[message.channel].pop(0)

    subCommands = OrderedDict([
         ('add', _add),
         ('remove', _remove),
         ('removebyid', _removebyid),
         ('list', _list),
         ('search', _search),
         ('searchnick', _searchnick),
         ('id', _id),
         ('random', _random)]
     )

outOfContext = OutOfContext()
<commit_msg>[OutOfContext] Fix help text<commit_after>import random
import time
from collections import OrderedDict
from typing import List, Union
from weakref import WeakKeyDictionary

from twisted.plugin import IPlugin
from zope.interface import implementer

from desertbot.message import IRCMessage, TargetTypes
from desertbot.moduleinterface import IModule
from desertbot.modules.commandinterface import BotCommand
from desertbot.response import IRCResponse
from desertbot.utils import string

try:
    import re2
except ImportError:
    import re as re2

@implementer(IPlugin, IModule)
class OutOfContext(BotCommand):
    messageLimit = 50

    def actions(self):
        return super(OutOfContext, self).actions() + [("message-channel", 1, self.storeMessage),
                                                      ("action-channel", 1, self.storeMessage)]

    def triggers(self):
        return ["outofcontext"]

    def _add(self, message: IRCMessage):
        """add - Adds a quote to the OutOfContext log. The quote will be pulled from a message line buffer."""
        if len(message.parameterList) < 2:
            return IRCResponse("Add what?", message.replyTo)
        if message.targetType == TargetTypes.USER:
            return IRCResponse("You can only add messages from channels.", message.replyTo)

        regex = re2.compile(re2.escape(" ".join(message.parameterList[1:])), re2.IGNORECASE)
        if len(self.messageStore) == 0 or message.channel not in self.messageStore:
            return IRCResponse("Sorry, there are no messages in my buffer.", message.replyTo)

        matches = list(filter(regex.search, self.messageStore[message.channel]))
        if len(matches) == 0:
            return IRCResponse("Sorry, that didn't match anything in my message buffer.", message.replyTo)
        if len(matches) > 1:
            return IRCResponse("Sorry, that matches too many lines in my message buffer.", message.replyTo)

        todayDate = time.strftime("[%Y-%m-%d] [%H:%M]")
        quote = f"{todayDate} {matches[0]}"
        if message.replyTo not in self.storage:
            self.storage[message.replyTo] = []
        if len(self.storage[message.replyTo]) > 0 and self.storage[message.replyTo][-1] == quote:
            return IRCResponse("That quote has already been added to the log!", message.replyTo)
        else:
            self.storage[message.replyTo].append(quote)
            return IRCResponse(f"Quote '{quote}' was added to the log!", message.replyTo)

    def _remove(self, message: IRCMessage):
        """remove <regex> - remove a quote from the OutOfContext log."""
        if len(message.parameterList) < 2:
            return IRCResponse("Remove what?", message.replyTo)
        if len(self.storage) == 0 or message.replyTo not in self.storage:
            return IRCResponse("There are no quotes in the log.", message.replyTo)
        regex = re2.compile(" ".join(message.parameterList[1:]), re2.IGNORECASE)
        matches = list(filter(regex.search, self.storage[message.replyTo]))
        if len(matches) == 0:
            return IRCResponse("That message is not in the log.", message.replyTo)
        if len(matches) > 1:
            return IRCResponse(f"Unable to remove quote, {len(matches)} matches found.", message.replyTo)
        return self._removeQuote(message.replyTo, matches[0])

    def _removebyid(self, message: IRCMessage):
        """removebyid <quoteid> - remove the quote with the specified ID from the OutOfContext log."""
        if len(message.parameterList) < 2:
            return IRCResponse("Remove what?", message.replyTo)
        if not string.isNumber(message.parameterList[1]):
            return IRCResponse("You didn't specify a valid ID.", message.replyTo)
        if len(self.storage) == 0 or message.replyTo not in self.storage:
            return IRCResponse("There are no quotes in the log.", message.replyTo)
        index = int(message.parameterList[1]) - 1
        quotes = self.storage[message.replyTo]
        if index < len(quotes):
            return self._removeQuote(message.replyTo, quotes[index])
        return IRCResponse("That message is not in the log.", message.replyTo)

    def _list(self, message: IRCMessage):
        """list (<search/searchnick <regex>>) - post the OutOfContext log. A search regex can be provided to filter\
         the list."""
        params = [x for x in message.parameterList[1:]]
        if len(params) > 0:
            subsubcommand = params.pop(0).lower()
            if subsubcommand == "searchnick":
                return self._postList(message.replyTo, " ".join(params), True)
            elif subsubcommand == "search":
                return self._postList(message.replyTo, " ".join(params), False)
            else:
                return self._postList(message.replyTo, "", False)
        else:
            return self._postList(message.replyTo, "", False)

    def _search(self, message: IRCMessage):
        """search <regex> - look up quotes in the OutOfContext log. This search operation will look at \
        the content of the quotes."""
        if len(message.parameterList) < 2:
            return IRCResponse("Search what?", message.replyTo)
        return self._getQuote(message.replyTo, " ".join(message.parameterList[1:]), False, -1)

    def _searchnick(self, message: IRCMessage):
        """searchnick <regex> - look up quotes in the OutOfContext log. This search operation will \
        look at the nick that said the quoted line."""
        nick = message.user.nick
        if len(message.parameterList) > 1:
            nick = message.parameterList[1]
        return self._getQuote(message.replyTo, nick, True, -1)

    def _id(self, message: IRCMessage):
        """id <quoteid> - look up the quote that has the given ID."""
        if len(message.parameterList) < 2 or not string.isNumber(message.parameterList[1]):
            return IRCResponse("You didn't specify a valid ID.", message.replyTo)
        return self._getQuote(message.replyTo, "", False, int(message.parameterList[1]) - 1)

    def _random(self, message: IRCMessage):
        """random - return a random quote from the OutOfContext log."""
        return self._getQuote(message.replyTo, "", False, -1)

    def _removeQuote(self, source, quote):
        self.storage[source].remove(quote)
        return IRCResponse(f"Quote '{quote}' was removed from the log!", source)

    def _postList(self, source, searchString, searchNickname):
        if len(self.storage) == 0 or source not in self.storage:
            return IRCResponse("There are no quotes in the log.", source)
        regex = re2.compile(searchString, re2.IGNORECASE)
        matches = []
        if searchNickname:
            for x in self.storage[source]:
                if x[21] == "*":
                    match = re2.search(regex, x[:x.find(" ", 23)])
                else:
                    match = re2.search(regex, x[x.find("<") + 1:x.find(">")])
                if match:
                    matches.append(x)
        else:
            for x in self.storage[source]:
                if re2.search(regex, x[x.find(">") + 1:]):
                    matches.append(x)
        if len(matches) == 0:
            return IRCResponse(f"No matches for '{searchString}' found.", source)

        pasteLink =  self.bot.moduleHandler.runActionUntilValue('upload-dbco',
                                                                string.stripFormatting("\n".join(matches)),
                                                                10 * 60)

        return IRCResponse(f"Link posted! (Expires in 10 minutes) {pasteLink}.", source)

    def _getQuote(self, source, searchString, searchNickname, index):
        if len(self.storage) == 0 or source not in self.storage:
            return IRCResponse("There are no quotes in the log.", source)
        regex = re2.compile(searchString, re2.IGNORECASE)
        matches = []
        if searchNickname:
            for x in self.storage[source]:
                if x[21] == "*":
                    match = re2.search(regex, x[:x.find(" ", 23)])
                else:
                    match = re2.search(regex, x[x.find("<") + 1:x.find(">")])
                if match:
                    matches.append(x)
        else:
            for x in self.storage[source]:
                if re2.search(regex, x[x.find(">") + 1:]):
                    matches.append(x)
        if len(matches) == 0:
            return IRCResponse(f"No matches for '{searchString}' found.", source)
        if index < 0 or index > len(matches) - 1:
            index = random.randint(0, len(matches) - 1)
        return IRCResponse(f"Quote #{index + 1}/{len(matches)}: {matches[index]}", source)

    def _unrecognizedSubcommand(self, subCommand):
        return (f"unrecognized subcommand f'{subCommand}', "
                f"available subcommands for outofcontext are: {', '.join(self.subCommands.keys())}")

    def onLoad(self) -> None:
        self.messageStore = WeakKeyDictionary()
        self._helpText = ("{1}outofcontext ({0}) - adds, removes or requests things from the OutOfContext quote list. "
                          "Use '{1}help outofcontext <subcommand>' for subcommand help. "
                          .format('/'.join(self.subCommands), self.bot.commandChar))

    def help(self, query: Union[List[str], None]) -> str:
        if len(query) > 1:
            subCommand = query[1].lower()
            if subCommand in self.subCommands:
                return ('{1}outofcontext {0}' .format(re2.sub(r"\s+", " ", self.subCommands[subCommand].__doc__),
                        self.bot.commandChar))
            else:
                return self._unrecognizedSubcommand(subCommand)
        else:
            return self._helpText

    def execute(self, message: IRCMessage):
        if len(message.parameterList) > 0:
            subCommand = message.parameterList[0].lower()
            if subCommand not in self.subCommands:
                return IRCResponse(self._unrecognizedSubcommand(subCommand), message.replyTo)
            return self.subCommands[subCommand](self, message)
        else:
            return IRCResponse(self._helpText, message.replyTo)

    def storeMessage(self, message: IRCMessage):
        if message.command or message.targetType == TargetTypes.USER:
            return

        if message.channel not in self.messageStore:
            self.messageStore[message.channel] = []

        if message.type == 'ACTION':
            self.messageStore[message.channel]\
                .append(f'* {message.user.nick} {string.stripFormatting(message.messageString)}')
        else:
            self.messageStore[message.channel]\
                .append(f'<{message.user.nick}> {string.stripFormatting(message.messageString)}')

        if len(self.messageStore[message.channel]) > self.messageLimit:
            self.messageStore[message.channel].pop(0)

    subCommands = OrderedDict([
         ('add', _add),
         ('remove', _remove),
         ('removebyid', _removebyid),
         ('list', _list),
         ('search', _search),
         ('searchnick', _searchnick),
         ('id', _id),
         ('random', _random)]
     )

outOfContext = OutOfContext()
<|endoftext|>


<commit_before># Licensed to Cloudera, Inc. under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  Cloudera, Inc. licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import absolute_import

from future import standard_library
standard_library.install_aliases()
from builtins import object
import logging
import sys

from useradmin.models import User

from desktop.auth.backend import is_admin
from desktop.conf import DEFAULT_USER, ENABLE_ORGANIZATIONS

if sys.version_info[0] > 2:
  from urllib.parse import urlparse as lib_urlparse
else:
  from urlparse import urlparse as lib_urlparse

LOG = logging.getLogger(__name__)
DEFAULT_USER = DEFAULT_USER.get()

class ProxyFS(object):

  def __init__(self, filesystems_dict, default_scheme, name='default'):
    if default_scheme not in filesystems_dict:
      raise ValueError(
        'Default scheme "%s" is not a member of provided schemes: %s' % (default_scheme, list(filesystems_dict.keys()))
      )

    self._name = name
    self._fs_dict = filesystems_dict
    self._user = {'user': None} # wrapping in an object to avoid triggering __getattr__ / __setattr__
    self._default_scheme = default_scheme
    self._default_fs = filesystems_dict[self._default_scheme](name)

  def __getattr__(self, item):
    return getattr(object.__getattribute__(self, "_default_fs"), item)

  def __setattr__(self, key, value):
    if hasattr(self, "_default_fs") and hasattr(self._default_fs, key):
      setattr(self._default_fs, key, value)
    else:
      object.__setattr__(self, key, value)

  def _get_scheme(self, path):
    scheme = None
    if path:
      split = lib_urlparse(path)
      scheme = split.scheme if split.scheme else None
    ret_scheme = scheme or self._default_scheme
    if not ret_scheme:
      raise IOError('Can not figure out scheme for path "%s"' % path)
    return ret_scheme

  def _has_access(self, fs):
    from desktop.auth.backend import rewrite_user  # Avoid cyclic loop
    try:
      filebrowser_action = fs.filebrowser_action()
      #if not filebrowser_action (hdfs) then handle permission via doas else check permission in hue
      if not filebrowser_action:
        return True
      lookup = {'email' if ENABLE_ORGANIZATIONS.get() else 'username': self.getuser()}
      user = rewrite_user(User.objects.get(**lookup))
      return user.is_authenticated() and user.is_active and (is_admin(user) or not filebrowser_action or user.has_hue_permission(action=filebrowser_action, app="filebrowser"))
    except User.DoesNotExist:
      LOG.exception('proxyfs.has_access()')
      return False

  def _get_fs(self, path):
    scheme = self._get_scheme(path)
    if self.getuser() is None:
      raise IOError('User not set')
    try:
      fs = self._fs_dict[scheme](self._name, self.getuser())
      if self._has_access(fs):
        fs.setuser(self.getuser())
        return fs
      else:
        raise IOError("Missing permissions for %s on %s" % (self.getuser(), path))
    except KeyError:
      raise IOError('Unknown scheme %s, available schemes: %s' % (scheme, list(self._fs_dict.keys())))

  def _get_fs_pair(self, src, dst):
    """
    Returns two FS for source and destination paths respectively.
    If `dst` is not self-contained path assumes it's relative path to `src`.
    """

    src_fs = self._get_fs(src)
    dst_scheme = lib_urlparse(dst).scheme
    if not dst_scheme:
      return src_fs, src_fs
    return src_fs, self._get_fs(dst)

  def setuser(self, user):
    """Set a new user. Return the past current user."""
    curr = self.getuser()
    if hasattr(user, 'username'):
      self._user['user'] = user.username
    else:
      self._user['user'] = user
    return curr

  def getuser(self):
    return self._user['user']

  def do_as_user(self, username, fn, *args, **kwargs):
    prev = self.getuser()
    try:
      self.setuser(username)
      return fn(*args, **kwargs)
    finally:
      self.setuser(prev)

  def do_as_superuser(self, fn, *args, **kwargs):
    scheme = self._get_scheme(args[0])
    fs = self._fs_dict[scheme](self._name)
    user = fs.superuser if fs.superuser else DEFAULT_USER
    return self.do_as_user(user, fn, *args, **kwargs)

  # Proxy methods to suitable filesystem
  # ------------------------------------
  def isdir(self, path):
    return self._get_fs(path).isdir(path)

  def isfile(self, path):
    return self._get_fs(path).isfile(path)

  def stats(self, path):
    return self._get_fs(path).stats(path)

  def listdir_stats(self, path, **kwargs):
    return self._get_fs(path).listdir_stats(path, **kwargs)

  def listdir(self, path, glob=None):
    return self._get_fs(path).listdir(path, glob)

  def normpath(self, path):
    return self._get_fs(path).normpath(path)

  def netnormpath(self, path):
    return self._get_fs(path).netnormpath(path)

  def open(self, path, *args, **kwargs):
    return self._get_fs(path).open(path, *args, **kwargs)

  def exists(self, path):
    return self._get_fs(path).exists(path)

  def isroot(self, path):
    return self._get_fs(path).isroot(path)

  def parent_path(self, path):
    return self._get_fs(path).parent_path(path)

  def join(self, first, *comp_list):
    return self._get_fs(first).join(first, *comp_list)

  def mkdir(self, path, *args, **kwargs):
    return self._get_fs(path).mkdir(path, *args, **kwargs)

  def read(self, path, *args, **kwargs):
    return self._get_fs(path).read(path, *args, **kwargs)

  def append(self, path, *args, **kwargs):
    return self._get_fs(path).append(path, *args, **kwargs)

  def rmtree(self, path, *args, **kwargs):
    self._get_fs(path).rmtree(path, *args, **kwargs)

  def remove(self, path, skip_trash=False):
    self._get_fs(path).remove(path, skip_trash)

  def restore(self, path):
    self._get_fs(path).restore(path)

  def create(self, path, *args, **kwargs):
    self._get_fs(path).create(path, *args, **kwargs)

  def create_home_dir(self, home_path=None):
    if home_path is None:
      home_path = self.get_home_dir()
    self._get_fs(home_path).create_home_dir(home_path)

  def chown(self, path, *args, **kwargs):
    self._get_fs(path).chown(path, *args, **kwargs)

  def chmod(self, path, *args, **kwargs):
    self._get_fs(path).chmod(path, *args, **kwargs)

  def copyFromLocal(self, local_src, remote_dst, *args, **kwargs):
    self._get_fs(remote_dst).copyFromLocal(local_src, remote_dst, *args, **kwargs)

  def mktemp(self, subdir='', prefix='tmp', basedir=None):
    fs = basedir and self._get_fs(basedir) or self.default_fs
    return fs.mktemp(subdir=subdir, prefix=prefix, basedir=basedir)

  def purge_trash(self):
    fs = self._get_fs() # Only webhdfs supports trash.
    if fs and hasattr(fs, 'purge_trash'):
      fs.purge_trash()

  # Handle file systems interactions
  # --------------------------------
  def copy(self, src, dst, *args, **kwargs):
    src_fs, dst_fs = self._get_fs_pair(src, dst)
    op = src_fs.copy if src_fs is dst_fs else self._copy_between_filesystems
    return op(src, dst, *args, **kwargs)

  def _copy_between_filesystems(self, src, dst, recursive=False, *args, **kwargs):
    raise NotImplementedError("Will be addressed in HUE-2934")

  def copyfile(self, src, dst, *args, **kwargs):
    src_fs, dst_fs = self._get_fs_pair(src, dst)
    op = src_fs.copyfile if src_fs is dst_fs else self._copyfile_between_filesystems
    return op(src, dst, *args, **kwargs)

  def _copyfile_between_filesystems(self, src, dst, *args, **kwargs):
    raise NotImplementedError("Will be addressed in HUE-2934")

  def copy_remote_dir(self, src, dst, *args, **kwargs):
    src_fs, dst_fs = self._get_fs_pair(src, dst)
    op = src_fs.copy_remote_dir if src_fs is dst_fs else self._copy_remote_dir_between_filesystems
    return op(src, dst, *args, **kwargs)

  def _copy_remote_dir_between_filesystems(self, src, dst, *args, **kwargs):
    raise NotImplementedError("Will be addressed in HUE-2934")

  def rename(self, old, new):
    old_fs, new_fs = self._get_fs_pair(old, new)
    op = old_fs.rename if old_fs is new_fs else self._rename_between_filesystems
    return op(old, new)

  def _rename_between_filesystems(self, old, new):
    raise NotImplementedError("Will be addressed in HUE-2934")

  def rename_star(self, old_dir, new_dir):
    old_fs, new_fs = self._get_fs_pair(old_dir, new_dir)
    op = old_fs.rename_star if old_fs is new_fs else self._rename_star_between_filesystems
    return op(old_dir, new_dir)

  def _rename_star_between_filesystems(self, old, new):
    raise NotImplementedError("Will be addressed in HUE-2934")

  def upload(self, file, path, *args, **kwargs):
    self._get_fs(path).upload(file, path, *args, **kwargs)

  def check_access(self, path, *args, **kwargs):
    self._get_fs(path).check_access(path, *args, **kwargs)

  def mkswap(self, filename, subdir='', suffix='swp', basedir=None):
    return self._get_fs(basedir).mkswap(filename, subdir, suffix, basedir)

  def get_upload_chuck_size(self, path):
    return self._get_fs(path).get_upload_chuck_size()
<commit_msg>HUE-9073 [fs] Avoid missing user argument in _get_client_cached<commit_after># Licensed to Cloudera, Inc. under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  Cloudera, Inc. licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import absolute_import

from future import standard_library
standard_library.install_aliases()
from builtins import object
import logging
import sys

from useradmin.models import User

from desktop.auth.backend import is_admin
from desktop.conf import DEFAULT_USER, ENABLE_ORGANIZATIONS

if sys.version_info[0] > 2:
  from urllib.parse import urlparse as lib_urlparse
else:
  from urlparse import urlparse as lib_urlparse

LOG = logging.getLogger(__name__)
DEFAULT_USER = DEFAULT_USER.get()

class ProxyFS(object):

  def __init__(self, filesystems_dict, default_scheme, name='default'):
    if default_scheme not in filesystems_dict:
      raise ValueError(
        'Default scheme "%s" is not a member of provided schemes: %s' % (default_scheme, list(filesystems_dict.keys()))
      )

    self._name = name
    self._fs_dict = filesystems_dict
    self._user = {'user': None} # wrapping in an object to avoid triggering __getattr__ / __setattr__
    self._default_scheme = default_scheme
    self._default_fs = filesystems_dict[self._default_scheme](name, user=None)

  def __getattr__(self, item):
    return getattr(object.__getattribute__(self, "_default_fs"), item)

  def __setattr__(self, key, value):
    if hasattr(self, "_default_fs") and hasattr(self._default_fs, key):
      setattr(self._default_fs, key, value)
    else:
      object.__setattr__(self, key, value)

  def _get_scheme(self, path):
    scheme = None
    if path:
      split = lib_urlparse(path)
      scheme = split.scheme if split.scheme else None
    ret_scheme = scheme or self._default_scheme
    if not ret_scheme:
      raise IOError('Can not figure out scheme for path "%s"' % path)
    return ret_scheme

  def _has_access(self, fs):
    from desktop.auth.backend import rewrite_user  # Avoid cyclic loop
    try:
      filebrowser_action = fs.filebrowser_action()
      #if not filebrowser_action (hdfs) then handle permission via doas else check permission in hue
      if not filebrowser_action:
        return True
      lookup = {'email' if ENABLE_ORGANIZATIONS.get() else 'username': self.getuser()}
      user = rewrite_user(User.objects.get(**lookup))
      return user.is_authenticated() and user.is_active and (is_admin(user) or not filebrowser_action or user.has_hue_permission(action=filebrowser_action, app="filebrowser"))
    except User.DoesNotExist:
      LOG.exception('proxyfs.has_access()')
      return False

  def _get_fs(self, path):
    scheme = self._get_scheme(path)
    if self.getuser() is None:
      raise IOError('User not set')
    try:
      fs = self._fs_dict[scheme](self._name, self.getuser())
      if self._has_access(fs):
        fs.setuser(self.getuser())
        return fs
      else:
        raise IOError("Missing permissions for %s on %s" % (self.getuser(), path))
    except KeyError:
      raise IOError('Unknown scheme %s, available schemes: %s' % (scheme, list(self._fs_dict.keys())))

  def _get_fs_pair(self, src, dst):
    """
    Returns two FS for source and destination paths respectively.
    If `dst` is not self-contained path assumes it's relative path to `src`.
    """

    src_fs = self._get_fs(src)
    dst_scheme = lib_urlparse(dst).scheme
    if not dst_scheme:
      return src_fs, src_fs
    return src_fs, self._get_fs(dst)

  def setuser(self, user):
    """Set a new user. Return the past current user."""
    curr = self.getuser()
    if hasattr(user, 'username'):
      self._user['user'] = user.username
    else:
      self._user['user'] = user
    return curr

  def getuser(self):
    return self._user['user']

  def do_as_user(self, username, fn, *args, **kwargs):
    prev = self.getuser()
    try:
      self.setuser(username)
      return fn(*args, **kwargs)
    finally:
      self.setuser(prev)

  def do_as_superuser(self, fn, *args, **kwargs):
    scheme = self._get_scheme(args[0])
    fs = self._fs_dict[scheme](self._name)
    user = fs.superuser if fs.superuser else DEFAULT_USER
    return self.do_as_user(user, fn, *args, **kwargs)

  # Proxy methods to suitable filesystem
  # ------------------------------------
  def isdir(self, path):
    return self._get_fs(path).isdir(path)

  def isfile(self, path):
    return self._get_fs(path).isfile(path)

  def stats(self, path):
    return self._get_fs(path).stats(path)

  def listdir_stats(self, path, **kwargs):
    return self._get_fs(path).listdir_stats(path, **kwargs)

  def listdir(self, path, glob=None):
    return self._get_fs(path).listdir(path, glob)

  def normpath(self, path):
    return self._get_fs(path).normpath(path)

  def netnormpath(self, path):
    return self._get_fs(path).netnormpath(path)

  def open(self, path, *args, **kwargs):
    return self._get_fs(path).open(path, *args, **kwargs)

  def exists(self, path):
    return self._get_fs(path).exists(path)

  def isroot(self, path):
    return self._get_fs(path).isroot(path)

  def parent_path(self, path):
    return self._get_fs(path).parent_path(path)

  def join(self, first, *comp_list):
    return self._get_fs(first).join(first, *comp_list)

  def mkdir(self, path, *args, **kwargs):
    return self._get_fs(path).mkdir(path, *args, **kwargs)

  def read(self, path, *args, **kwargs):
    return self._get_fs(path).read(path, *args, **kwargs)

  def append(self, path, *args, **kwargs):
    return self._get_fs(path).append(path, *args, **kwargs)

  def rmtree(self, path, *args, **kwargs):
    self._get_fs(path).rmtree(path, *args, **kwargs)

  def remove(self, path, skip_trash=False):
    self._get_fs(path).remove(path, skip_trash)

  def restore(self, path):
    self._get_fs(path).restore(path)

  def create(self, path, *args, **kwargs):
    self._get_fs(path).create(path, *args, **kwargs)

  def create_home_dir(self, home_path=None):
    if home_path is None:
      home_path = self.get_home_dir()
    self._get_fs(home_path).create_home_dir(home_path)

  def chown(self, path, *args, **kwargs):
    self._get_fs(path).chown(path, *args, **kwargs)

  def chmod(self, path, *args, **kwargs):
    self._get_fs(path).chmod(path, *args, **kwargs)

  def copyFromLocal(self, local_src, remote_dst, *args, **kwargs):
    self._get_fs(remote_dst).copyFromLocal(local_src, remote_dst, *args, **kwargs)

  def mktemp(self, subdir='', prefix='tmp', basedir=None):
    fs = basedir and self._get_fs(basedir) or self.default_fs
    return fs.mktemp(subdir=subdir, prefix=prefix, basedir=basedir)

  def purge_trash(self):
    fs = self._get_fs() # Only webhdfs supports trash.
    if fs and hasattr(fs, 'purge_trash'):
      fs.purge_trash()

  # Handle file systems interactions
  # --------------------------------
  def copy(self, src, dst, *args, **kwargs):
    src_fs, dst_fs = self._get_fs_pair(src, dst)
    op = src_fs.copy if src_fs is dst_fs else self._copy_between_filesystems
    return op(src, dst, *args, **kwargs)

  def _copy_between_filesystems(self, src, dst, recursive=False, *args, **kwargs):
    raise NotImplementedError("Will be addressed in HUE-2934")

  def copyfile(self, src, dst, *args, **kwargs):
    src_fs, dst_fs = self._get_fs_pair(src, dst)
    op = src_fs.copyfile if src_fs is dst_fs else self._copyfile_between_filesystems
    return op(src, dst, *args, **kwargs)

  def _copyfile_between_filesystems(self, src, dst, *args, **kwargs):
    raise NotImplementedError("Will be addressed in HUE-2934")

  def copy_remote_dir(self, src, dst, *args, **kwargs):
    src_fs, dst_fs = self._get_fs_pair(src, dst)
    op = src_fs.copy_remote_dir if src_fs is dst_fs else self._copy_remote_dir_between_filesystems
    return op(src, dst, *args, **kwargs)

  def _copy_remote_dir_between_filesystems(self, src, dst, *args, **kwargs):
    raise NotImplementedError("Will be addressed in HUE-2934")

  def rename(self, old, new):
    old_fs, new_fs = self._get_fs_pair(old, new)
    op = old_fs.rename if old_fs is new_fs else self._rename_between_filesystems
    return op(old, new)

  def _rename_between_filesystems(self, old, new):
    raise NotImplementedError("Will be addressed in HUE-2934")

  def rename_star(self, old_dir, new_dir):
    old_fs, new_fs = self._get_fs_pair(old_dir, new_dir)
    op = old_fs.rename_star if old_fs is new_fs else self._rename_star_between_filesystems
    return op(old_dir, new_dir)

  def _rename_star_between_filesystems(self, old, new):
    raise NotImplementedError("Will be addressed in HUE-2934")

  def upload(self, file, path, *args, **kwargs):
    self._get_fs(path).upload(file, path, *args, **kwargs)

  def check_access(self, path, *args, **kwargs):
    self._get_fs(path).check_access(path, *args, **kwargs)

  def mkswap(self, filename, subdir='', suffix='swp', basedir=None):
    return self._get_fs(basedir).mkswap(filename, subdir, suffix, basedir)

  def get_upload_chuck_size(self, path):
    return self._get_fs(path).get_upload_chuck_size()
<|endoftext|>