<commit_before>from .proto import xenon_pb2
from .server import __server__
from .exceptions import make_exception
import grpc

try:
    from enum import Enum
except ImportError as e:
    use_enum = False

    class Enum(object):
        """Minimal Enum replacement."""
        def __init__(self, name, items):
            for k, v in items:
                setattr(self, k, v)

else:
    use_enum = True

try:
    from inspect import (Signature, Parameter, signature)
except ImportError as e:
    use_signature = False
else:
    use_signature = True

def mirror_enum(parent, name):
    grpc_enum = getattr(parent, name)
    return Enum(name, grpc_enum.items())

def to_camel_case(name):
    return ''.join(w.title() for w in name.split('_'))

def to_lower_camel_case(name):
    words = name.split('_')
    return words[0] + ''.join(w.title() for w in words[1:])

def get_fields(msg_type):
    """Get a list of field names for Grpc message."""
    return list(f.name for f in msg_type.DESCRIPTOR.fields)

def get_field_type(f):
    """Obtain the type name of a GRPC Message field."""
    types = (t[5:] for t in dir(f) if t[:4] == 'TYPE' and
             getattr(f, t) == f.type)
    return next(types)

def get_field_description(f):
    """Get the type description of a GRPC Message field."""
    type_name = get_field_type(f)
    if type_name == 'MESSAGE' and \
            {sf.name for sf in f.message_type.fields} == {'key', 'value'}:
        return 'map<string, string>'
    elif type_name == 'MESSAGE':
        return f.message_type.full_name
    elif type_name == 'ENUM':
        return f.enum_type.full_name
    else:
        return type_name.lower()

def list_attributes(msg_type):
    """List all attributes with type description of a GRPC Message class."""
    return [(f.name, get_field_description(f))
            for f in msg_type.fields]

class GrpcMethod:
    """Data container for a GRPC method.

    :ivar name: underscore style of method name
    :ivar uses_request: wether this method has a request, if this value is
        `True`, the name is generated from `name`, if it is a string the
        contents of this string are used.
    :ivar field_name: name of `self` within the request.
    :ivar input_transform: custom method to generate a request from the
        method's arguments.
    :ivar output_transform: custom method to extract the return value from
        the return value.
    """
    def __init__(self, name, uses_request=False, field_name=None,
                 input_transform=None, output_transform=None,
                 static=False):
        self.name = name
        self.uses_request = uses_request
        self.field_name = field_name
        self.input_transform = input_transform
        self.output_transform = output_transform
        self.static = static

    @property
    def is_simple(self):
        return not self.uses_request and not self.input_transform \
            and not self.static

    @property
    def request_name(self):
        """Generate the name of the request."""
        if self.static and not self.uses_request:
            return 'Empty'

        if not self.uses_request:
            return None

        if isinstance(self.uses_request, str):
            return self.uses_request

        return to_camel_case(self.name) + "Request"

    @property
    def request_type(self):
        """Retrieve the type of the request, by fetching it from
        `xenon.proto.xenon_pb2`."""
        if self.static and not self.uses_request:
            return getattr(xenon_pb2, 'Empty')

        if not self.uses_request:
            return None

        return getattr(xenon_pb2, self.request_name)

    # python 3 only
    @property
    def signature(self):
        """Create a signature for this method, only in Python > 3.4"""
        if not use_signature:
            raise NotImplementedError("Python 3 only.")

        if self.static:
            parameters = \
                (Parameter(name='cls',
                           kind=Parameter.POSITIONAL_ONLY),)

        else:
            parameters = \
                (Parameter(name='self',
                           kind=Parameter.POSITIONAL_ONLY),)

        if self.input_transform:
            return signature(self.input_transform)

        if self.uses_request:
            fields = get_fields(self.request_type)
            if not self.static:
                if self.field_name not in fields:
                    raise NameError("field '{}' not found in {}".format(
                        self.field_name, self.request_name))
                fields.remove(self.field_name)
            parameters += tuple(
                Parameter(name=name, kind=Parameter.POSITIONAL_OR_KEYWORD,
                          default=None)
                for name in fields)

        return Signature(parameters)

    # TODO extend documentation rendered from proto
    def docstring(self, servicer):
        """Generate a doc-string."""
        s = getattr(servicer, to_lower_camel_case(self.name)).__doc__ \
            or "TODO: no docstring in .proto file"

        if self.uses_request:
            s += "\n"
            for field in get_fields(self.request_type):
                if field != self.field_name:
                    s += "    :param {}: {}\n".format(field, field)
                    s += "    :type {0}: {1}\n".format(
                        field, get_field_description(
                            self.request_type.DESCRIPTOR.fields_by_name[field]))

        return s

def unwrap(arg):
    if hasattr(arg, '__is_proxy__'):
        return arg.__wrapped__
    else:
        return arg

def make_static_request(method, *args, **kwargs):
    """Creates a request from a static method function call."""
    if args and not use_signature:
        raise NotImplementedError("Only keyword arguments allowed in Python2")

    if use_signature:
        new_kwargs = {kw: unwrap(value) for kw, value in kwargs.items()}
        new_args = tuple(unwrap(value) for value in args)
        bound_args = method.signature.bind(
            None, *new_args, **new_kwargs).arguments

        # if we encounter any Enum arguments, replace them with their value
        for k in bound_args:
            if isinstance(bound_args[k], Enum):
                bound_args[k] = bound_args[k].value

        new_kwargs = {kw: v for kw, v in bound_args.items() if kw != 'cls'}

    else:
        new_kwargs = {kw: unwrap(value) for kw, value in kwargs.items()}

    return method.request_type(**new_kwargs)

def make_request(self, method, *args, **kwargs):
    """Creates a request from a method function call."""
    if args and not use_signature:
        raise NotImplementedError("Only keyword arguments allowed in Python2")

    new_kwargs = {kw: unwrap(value) for kw, value in kwargs.items()}

    if use_signature:
        new_args = tuple(unwrap(value) for value in args)
        bound_args = method.signature.bind(
            unwrap(self), *new_args, **new_kwargs).arguments

        # if we encounter any Enum arguments, replace them with their value
        def translate_enum(arg):
            return arg.value if isinstance(arg, Enum) else arg

        for k in bound_args:
            if isinstance(bound_args[k], str):
                continue
            if isinstance(bound_args[k], dict):
                continue

            try:
                x = [translate_enum(arg) for arg in bound_args[k]]
                bound_args[k] = x
            except TypeError:
                bound_args[k] = translate_enum(bound_args[k])

        # replace `self` with the correct keyword
        new_kwargs = {(kw if kw != 'self' else method.field_name): v
                      for kw, v in bound_args.items()}
        # args = tuple(x.value if isinstance(x, Enum) else x for x in args)

    else:
        new_kwargs[self.field_name] = unwrap(self)

    return method.request_type(**new_kwargs)

def apply_transform(service, t, x):
    """Apply a transformation using `self` as object reference."""
    if t is None:
        return x
    else:
        return t(service, x)

def transform_map(f):
    def t(self, xs):
        return (f(self, x) for x in xs)

    return t

def grpc_call(service, method, request):
    f = getattr(service, to_lower_camel_case(method.name))
    try:
        result = f(request)
    except grpc.RpcError as e:
        raise make_exception(method, e) from None

    return result

def method_wrapper(m):
    """Generates a method from a `GrpcMethod` definition."""

    if m.is_simple:
        def simple_method(self):
            """TODO: no docstring!"""
            return apply_transform(
                self.__service__, m.output_transform,
                grpc_call(self.__service__, m, unwrap(self)))

        return simple_method

    elif m.input_transform is not None:
        def transform_method(self, *args, **kwargs):
            """TODO: no docstring!"""
            request = m.input_transform(self, *args, **kwargs)
            return apply_transform(
                self.__service__, m.output_transform,
                grpc_call(self.__service__, m, request))

        return transform_method

    elif m.static:
        def static_method(cls, *args, **kwargs):
            """TODO: no docstring!"""
            request = make_static_request(m, *args, **kwargs)
            return apply_transform(
                cls.__stub__(__server__), m.output_transform,
                grpc_call(cls.__stub__(__server__), m, request))

        return static_method

    else:
        def request_method(self, *args, **kwargs):
            """TODO: no docstring!"""
            request = make_request(self, m, *args, **kwargs)
            return apply_transform(
                self.__service__, m.output_transform,
                grpc_call(self.__service__, m, request))

        return request_method

class OopMeta(type):
    """Meta class for Grpc Object wrappers."""
    def __new__(cls, name, parents, dct):
        return super(OopMeta, cls).__new__(cls, name, parents, dct)

    def __init__(cls, name, parents, dct):
        super(OopMeta, cls).__init__(name, parents, dct)

        for m in cls.__methods__():
            if m.uses_request and not m.field_name:
                m.field_name = cls.__field_name__

            f = method_wrapper(m)
            if use_signature:
                f.__signature__ = m.signature

            if cls.__servicer__:
                f.__doc__ = m.docstring(cls.__servicer__)

            f.__name__ = m.name

            if m.static:
                setattr(cls, m.name, classmethod(f))
            else:
                setattr(cls, m.name, f)

        try:
            grpc_cls = getattr(xenon_pb2, name)
            if cls.__doc__ is None:
                cls.__doc__ = "Wrapped proto message."
            cls.__doc__ += "\n\n"
            for attr in list_attributes(grpc_cls.DESCRIPTOR):
                cls.__doc__ += \
                    "    :ivar {0}: {0}\n    :vartype {0}: {1}\n".format(*attr)

        except AttributeError:
            pass

class OopProxy(metaclass=OopMeta):
    """Base class for Grpc Object wrappers. Ensures basic object sanity,
    namely the existence of `__service__` and `__wrapped__` members and
    the using of `OopMeta` meta-class. Also manages retrieving attributes
    from the wrapped instance.

    :ivar __is_proxy__: If True, this value represents a wrapped value,
        from which the GRPC message can be extraced by getting the
        `__wrapped__` attribute.
    :ivar __servicer__: If applicable, this gives the GRPC servicer class
        associated with the proxy object; this is used to retrieve doc-strings.
    :ivar __field_name__: The default name to which an object of this class
        should be bound in a request. This can be overridden by specifying
        the `field_name` property in the `GRPCMethod` definition. For a
        well-designed API this should not be necessary though.
    """

    __is_proxy__ = True
    __servicer__ = None
    __field_name__ = None

    @classmethod
    def __methods__(cls):
        """This method should return a list of GRPCMethod objects."""
        return []

    def __init__(self, service, wrapped):
        self.__service__ = service
        self.__wrapped__ = wrapped

    @staticmethod
    def __stub__(server):
        """Return the GRPC stub class to which this object interfaces."""
        raise NotImplementedError()

    def __getattr__(self, attr):
        """Accesses fields of the corresponding GRPC message."""
        return getattr(self.__wrapped__, attr)

    def _repr_html_(self):
        members = [f.name for f in self.__wrapped__.DESCRIPTOR.fields]
        s = type(self).__name__ + ":
"
        for m in members:
            s += "
{}: {}".format(m, getattr(self, m))
            if m not in dir(self.__wrapped__):
                s += " (default)"
            s += "
"
        s += "
"
        return s

    def __str__(self):
        members = [f.name for f in self.__wrapped__.DESCRIPTOR.fields]
        s = type(self).__name__ + ":\n"
        for m in members:
            s += "    {}: {}".format(m, getattr(self, m))
            if m not in dir(self.__wrapped__):
                s += " (default)"
            s += "\n"
        return s

    def __dir__(self):
        members = [f.name for f in self.__wrapped__.DESCRIPTOR.fields]
        return dir(super(OopProxy, self)) + members
<commit_msg>fix flake8 error<commit_after>from .proto import xenon_pb2
from .server import __server__
from .exceptions import make_exception
import grpc

try:
    from enum import Enum
except ImportError as e:
    use_enum = False

    class Enum(object):
        """Minimal Enum replacement."""
        def __init__(self, name, items):
            for k, v in items:
                setattr(self, k, v)

else:
    use_enum = True

try:
    from inspect import (Signature, Parameter, signature)
except ImportError as e:
    use_signature = False
else:
    use_signature = True

def mirror_enum(parent, name):
    grpc_enum = getattr(parent, name)
    return Enum(name, grpc_enum.items())

def to_camel_case(name):
    return ''.join(w.title() for w in name.split('_'))

def to_lower_camel_case(name):
    words = name.split('_')
    return words[0] + ''.join(w.title() for w in words[1:])

def get_fields(msg_type):
    """Get a list of field names for Grpc message."""
    return list(f.name for f in msg_type.DESCRIPTOR.fields)

def get_field_type(f):
    """Obtain the type name of a GRPC Message field."""
    types = (t[5:] for t in dir(f) if t[:4] == 'TYPE' and
             getattr(f, t) == f.type)
    return next(types)

def get_field_description(f):
    """Get the type description of a GRPC Message field."""
    type_name = get_field_type(f)
    if type_name == 'MESSAGE' and \
            {sf.name for sf in f.message_type.fields} == {'key', 'value'}:
        return 'map<string, string>'
    elif type_name == 'MESSAGE':
        return f.message_type.full_name
    elif type_name == 'ENUM':
        return f.enum_type.full_name
    else:
        return type_name.lower()

def list_attributes(msg_type):
    """List all attributes with type description of a GRPC Message class."""
    return [(f.name, get_field_description(f))
            for f in msg_type.fields]

class GrpcMethod:
    """Data container for a GRPC method.

    :ivar name: underscore style of method name
    :ivar uses_request: wether this method has a request, if this value is
        `True`, the name is generated from `name`, if it is a string the
        contents of this string are used.
    :ivar field_name: name of `self` within the request.
    :ivar input_transform: custom method to generate a request from the
        method's arguments.
    :ivar output_transform: custom method to extract the return value from
        the return value.
    """
    def __init__(self, name, uses_request=False, field_name=None,
                 input_transform=None, output_transform=None,
                 static=False):
        self.name = name
        self.uses_request = uses_request
        self.field_name = field_name
        self.input_transform = input_transform
        self.output_transform = output_transform
        self.static = static

    @property
    def is_simple(self):
        return not self.uses_request and not self.input_transform \
            and not self.static

    @property
    def request_name(self):
        """Generate the name of the request."""
        if self.static and not self.uses_request:
            return 'Empty'

        if not self.uses_request:
            return None

        if isinstance(self.uses_request, str):
            return self.uses_request

        return to_camel_case(self.name) + "Request"

    @property
    def request_type(self):
        """Retrieve the type of the request, by fetching it from
        `xenon.proto.xenon_pb2`."""
        if self.static and not self.uses_request:
            return getattr(xenon_pb2, 'Empty')

        if not self.uses_request:
            return None

        return getattr(xenon_pb2, self.request_name)

    # python 3 only
    @property
    def signature(self):
        """Create a signature for this method, only in Python > 3.4"""
        if not use_signature:
            raise NotImplementedError("Python 3 only.")

        if self.static:
            parameters = \
                (Parameter(name='cls',
                           kind=Parameter.POSITIONAL_ONLY),)

        else:
            parameters = \
                (Parameter(name='self',
                           kind=Parameter.POSITIONAL_ONLY),)

        if self.input_transform:
            return signature(self.input_transform)

        if self.uses_request:
            fields = get_fields(self.request_type)
            if not self.static:
                if self.field_name not in fields:
                    raise NameError("field '{}' not found in {}".format(
                        self.field_name, self.request_name))
                fields.remove(self.field_name)
            parameters += tuple(
                Parameter(name=name, kind=Parameter.POSITIONAL_OR_KEYWORD,
                          default=None)
                for name in fields)

        return Signature(parameters)

    # TODO extend documentation rendered from proto
    def docstring(self, servicer):
        """Generate a doc-string."""
        s = getattr(servicer, to_lower_camel_case(self.name)).__doc__ \
            or "TODO: no docstring in .proto file"

        if self.uses_request:
            s += "\n"
            for field in get_fields(self.request_type):
                if field != self.field_name:
                    type_info = get_field_description(
                        self.request_type.DESCRIPTOR.fields_by_name[field])
                    s += "    :param {}: {}\n".format(field, field)
                    s += "    :type {0}: {1}\n".format(field, type_info)

        return s

def unwrap(arg):
    if hasattr(arg, '__is_proxy__'):
        return arg.__wrapped__
    else:
        return arg

def make_static_request(method, *args, **kwargs):
    """Creates a request from a static method function call."""
    if args and not use_signature:
        raise NotImplementedError("Only keyword arguments allowed in Python2")

    if use_signature:
        new_kwargs = {kw: unwrap(value) for kw, value in kwargs.items()}
        new_args = tuple(unwrap(value) for value in args)
        bound_args = method.signature.bind(
            None, *new_args, **new_kwargs).arguments

        # if we encounter any Enum arguments, replace them with their value
        for k in bound_args:
            if isinstance(bound_args[k], Enum):
                bound_args[k] = bound_args[k].value

        new_kwargs = {kw: v for kw, v in bound_args.items() if kw != 'cls'}

    else:
        new_kwargs = {kw: unwrap(value) for kw, value in kwargs.items()}

    return method.request_type(**new_kwargs)

def make_request(self, method, *args, **kwargs):
    """Creates a request from a method function call."""
    if args and not use_signature:
        raise NotImplementedError("Only keyword arguments allowed in Python2")

    new_kwargs = {kw: unwrap(value) for kw, value in kwargs.items()}

    if use_signature:
        new_args = tuple(unwrap(value) for value in args)
        bound_args = method.signature.bind(
            unwrap(self), *new_args, **new_kwargs).arguments

        # if we encounter any Enum arguments, replace them with their value
        def translate_enum(arg):
            return arg.value if isinstance(arg, Enum) else arg

        for k in bound_args:
            if isinstance(bound_args[k], str):
                continue
            if isinstance(bound_args[k], dict):
                continue

            try:
                x = [translate_enum(arg) for arg in bound_args[k]]
                bound_args[k] = x
            except TypeError:
                bound_args[k] = translate_enum(bound_args[k])

        # replace `self` with the correct keyword
        new_kwargs = {(kw if kw != 'self' else method.field_name): v
                      for kw, v in bound_args.items()}
        # args = tuple(x.value if isinstance(x, Enum) else x for x in args)

    else:
        new_kwargs[self.field_name] = unwrap(self)

    return method.request_type(**new_kwargs)

def apply_transform(service, t, x):
    """Apply a transformation using `self` as object reference."""
    if t is None:
        return x
    else:
        return t(service, x)

def transform_map(f):
    def t(self, xs):
        return (f(self, x) for x in xs)

    return t

def grpc_call(service, method, request):
    f = getattr(service, to_lower_camel_case(method.name))
    try:
        result = f(request)
    except grpc.RpcError as e:
        raise make_exception(method, e) from None

    return result

def method_wrapper(m):
    """Generates a method from a `GrpcMethod` definition."""

    if m.is_simple:
        def simple_method(self):
            """TODO: no docstring!"""
            return apply_transform(
                self.__service__, m.output_transform,
                grpc_call(self.__service__, m, unwrap(self)))

        return simple_method

    elif m.input_transform is not None:
        def transform_method(self, *args, **kwargs):
            """TODO: no docstring!"""
            request = m.input_transform(self, *args, **kwargs)
            return apply_transform(
                self.__service__, m.output_transform,
                grpc_call(self.__service__, m, request))

        return transform_method

    elif m.static:
        def static_method(cls, *args, **kwargs):
            """TODO: no docstring!"""
            request = make_static_request(m, *args, **kwargs)
            return apply_transform(
                cls.__stub__(__server__), m.output_transform,
                grpc_call(cls.__stub__(__server__), m, request))

        return static_method

    else:
        def request_method(self, *args, **kwargs):
            """TODO: no docstring!"""
            request = make_request(self, m, *args, **kwargs)
            return apply_transform(
                self.__service__, m.output_transform,
                grpc_call(self.__service__, m, request))

        return request_method

class OopMeta(type):
    """Meta class for Grpc Object wrappers."""
    def __new__(cls, name, parents, dct):
        return super(OopMeta, cls).__new__(cls, name, parents, dct)

    def __init__(cls, name, parents, dct):
        super(OopMeta, cls).__init__(name, parents, dct)

        for m in cls.__methods__():
            if m.uses_request and not m.field_name:
                m.field_name = cls.__field_name__

            f = method_wrapper(m)
            if use_signature:
                f.__signature__ = m.signature

            if cls.__servicer__:
                f.__doc__ = m.docstring(cls.__servicer__)

            f.__name__ = m.name

            if m.static:
                setattr(cls, m.name, classmethod(f))
            else:
                setattr(cls, m.name, f)

        try:
            grpc_cls = getattr(xenon_pb2, name)
            if cls.__doc__ is None:
                cls.__doc__ = "Wrapped proto message."
            cls.__doc__ += "\n\n"
            for attr in list_attributes(grpc_cls.DESCRIPTOR):
                cls.__doc__ += \
                    "    :ivar {0}: {0}\n    :vartype {0}: {1}\n".format(*attr)

        except AttributeError:
            pass

class OopProxy(metaclass=OopMeta):
    """Base class for Grpc Object wrappers. Ensures basic object sanity,
    namely the existence of `__service__` and `__wrapped__` members and
    the using of `OopMeta` meta-class. Also manages retrieving attributes
    from the wrapped instance.

    :ivar __is_proxy__: If True, this value represents a wrapped value,
        from which the GRPC message can be extraced by getting the
        `__wrapped__` attribute.
    :ivar __servicer__: If applicable, this gives the GRPC servicer class
        associated with the proxy object; this is used to retrieve doc-strings.
    :ivar __field_name__: The default name to which an object of this class
        should be bound in a request. This can be overridden by specifying
        the `field_name` property in the `GRPCMethod` definition. For a
        well-designed API this should not be necessary though.
    """

    __is_proxy__ = True
    __servicer__ = None
    __field_name__ = None

    @classmethod
    def __methods__(cls):
        """This method should return a list of GRPCMethod objects."""
        return []

    def __init__(self, service, wrapped):
        self.__service__ = service
        self.__wrapped__ = wrapped

    @staticmethod
    def __stub__(server):
        """Return the GRPC stub class to which this object interfaces."""
        raise NotImplementedError()

    def __getattr__(self, attr):
        """Accesses fields of the corresponding GRPC message."""
        return getattr(self.__wrapped__, attr)

    def _repr_html_(self):
        members = [f.name for f in self.__wrapped__.DESCRIPTOR.fields]
        s = type(self).__name__ + ":
"
        for m in members:
            s += "
{}: {}".format(m, getattr(self, m))
            if m not in dir(self.__wrapped__):
                s += " (default)"
            s += "
"
        s += "
"
        return s

    def __str__(self):
        members = [f.name for f in self.__wrapped__.DESCRIPTOR.fields]
        s = type(self).__name__ + ":\n"
        for m in members:
            s += "    {}: {}".format(m, getattr(self, m))
            if m not in dir(self.__wrapped__):
                s += " (default)"
            s += "\n"
        return s

    def __dir__(self):
        members = [f.name for f in self.__wrapped__.DESCRIPTOR.fields]
        return dir(super(OopProxy, self)) + members
<|endoftext|>


<commit_before>#
# Copyright (c) 2014 ThoughtWorks, Inc.
#
# Pixelated is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# Pixelated is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with Pixelated. If not, see <http://www.gnu.org/licenses/>.
import leap.keymanager

def patched_fetch_keys_from_server(self, address):
        """
        Fetch keys bound to C{address} from nickserver and insert them in
        local database.

        Instead of raising a KeyNotFound only for 404 responses, this implementation
        raises a KeyNotFound exception for all problems.

        For original see: https://github.com/leapcode/keymanager/blob/develop/src/leap/keymanager/__init__.py

        :param address: The address bound to the keys.
        :type address: str

        :raise KeyNotFound: If the key was not found on nickserver.
        """
        # request keys from the nickserver
        res = None
        try:
            res = self._get(self._nickserver_uri, {'address': address})
            res.raise_for_status()
            server_keys = res.json()
            # insert keys in local database
            if self.OPENPGP_KEY in server_keys:
                self._wrapper_map[OpenPGPKey].put_ascii_key(
                    server_keys['openpgp'])
        except requests.exceptions.HTTPError as e:
            logger.warning("HTTP error retrieving key: %r" % (e,))
            logger.warning("%s" % (res.content,))
            raise KeyNotFound(address)
        except Exception as e:
            logger.warning("Error retrieving key: %r" % (e,))
            raise KeyNotFound(address)

leap.keymanager.KeyManager._fetch_keys_from_server = patched_fetch_keys_from_server
<commit_msg>Added missing imports to monkey patch.<commit_after>#
# Copyright (c) 2014 ThoughtWorks, Inc.
#
# Pixelated is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# Pixelated is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with Pixelated. If not, see <http://www.gnu.org/licenses/>.
import leap.keymanager
import requests
import logging
from leap.keymanager.errors import KeyNotFound

logger = logging.getLogger(__name__)

def patched_fetch_keys_from_server(self, address):
        """
        Fetch keys bound to C{address} from nickserver and insert them in
        local database.

        Instead of raising a KeyNotFound only for 404 responses, this implementation
        raises a KeyNotFound exception for all problems.

        For original see: https://github.com/leapcode/keymanager/blob/develop/src/leap/keymanager/__init__.py

        :param address: The address bound to the keys.
        :type address: str

        :raise KeyNotFound: If the key was not found on nickserver.
        """
        # request keys from the nickserver
        res = None
        try:
            res = self._get(self._nickserver_uri, {'address': address})
            res.raise_for_status()
            server_keys = res.json()
            # insert keys in local database
            if self.OPENPGP_KEY in server_keys:
                self._wrapper_map[OpenPGPKey].put_ascii_key(
                    server_keys['openpgp'])
        except requests.exceptions.HTTPError as e:
            logger.warning("HTTP error retrieving key: %r" % (e,))
            logger.warning("%s" % (res.content,))
            raise KeyNotFound(address)
        except Exception as e:
            logger.warning("Error retrieving key: %r" % (e,))
            raise KeyNotFound(address)

leap.keymanager.KeyManager._fetch_keys_from_server = patched_fetch_keys_from_server
<|endoftext|>


<commit_before>def pct(value):
    return value / 100

def tax_calc(gross, deduction, residence):
    gross = float(gross)
    deduction = float(deduction)

    tax_total = 0

    city_tax_bracket = [
        # http://www.tax.ny.gov/pdf/current_forms/it/nyc_tax_rate_schedule.pdf
        { 'over': 500000, 'base': 18122, 'rate': pct(3.876) },
        { 'over':  50000, 'base':  1706, 'rate': pct(3.648) },
        { 'over':  25000, 'base':   808, 'rate': pct(3.591) },
        { 'over':  12000, 'base':   349, 'rate': pct(3.534) },
        { 'over':      0, 'base':   349, 'rate': pct(3.534) },
    ]

    if residence == 'NYC':
        for i in city_tax_bracket:
            if gross > i['over']:
                tax_total += i['base']
                tax_total += i['rate'] * ( gross - i['over'] )
                break
    else:
        # We don't yet support other areas.
        raise Exception

    state_tax_bracket = [
        # http://www.tax.ny.gov/pdf/current_forms/it/nyc_tax_rate_schedule.pdf
        { 'over': 1029250, 'base': 69612, 'rate': pct(8.82) },
        { 'over':  205850, 'base': 13209, 'rate': pct(6.85) },
        { 'over':   77150, 'base':  4651, 'rate': pct(6.65) },
        { 'over':   20550, 'base':  1000, 'rate': pct(5.90) },
        { 'over':   11300, 'base':   468, 'rate': pct(5.25) },
        { 'over':    8200, 'base':   328, 'rate': pct(4.50) },
        { 'over':       0, 'base':     0, 'rate': pct(4.00) },
    ]

    for i in state_tax_bracket:
        if gross > i['over']:
            tax_total += i['base']
            tax_total += i['rate'] * ( gross - i['over'] )
            break

    federal_tax_bracket = [
        # where AGI = income - state tax - city tax - standard deduction (5,800) - all other deductions (0)
        # http://www.irs.gov/pub/irs-pdf/i1040.pdf
        # http://www.irs.gov/publications/p501/index.html
        { 'over': 400000, 'base': 116163.75, 'rate': pct(39.6) },
        { 'over': 398350, 'base': 115586.25, 'rate': pct(35) },
        { 'over': 183250, 'base':  44603.25, 'rate': pct(33) },
        { 'over':  87850, 'base':  17891.25, 'rate': pct(28) },
        { 'over':  36250, 'base':   4991.25, 'rate': pct(25) },
        { 'over':   8925, 'base':    892.50, 'rate': pct(15) },
        { 'over':      0, 'base':         0, 'rate': pct(10) },
    ]

    agi = gross - tax_total - deduction

    for i in federal_tax_bracket:
        if agi > i['over']:
            tax_total += i['base']
            tax_total += i['rate'] * ( agi - i['over'] )
            break

    # FICA
    # 6.2% SSA (up to the wage base of 113700)
    # 1.45% Medicare
    # http://www.irs.gov/taxtopics/tc751.html
    tax_total += (agi * pct(1.45))
    wage_base = 113700
    if agi >= wage_base:
        tax_total += (wage_base * pct(6.2))
    else:
        tax_total += (agi * pct(6.2))

    return round(tax_total, 2)
<commit_msg>Instead of reusing agi, use new variable federal_agi, to indicate that is logically distinct<commit_after>def pct(value):
    return value / 100

def tax_calc(gross, deduction, residence):
    gross = float(gross)
    deduction = float(deduction)

    tax_total = 0

    city_tax_bracket = [
        # http://www.tax.ny.gov/pdf/current_forms/it/nyc_tax_rate_schedule.pdf
        { 'over': 500000, 'base': 18122, 'rate': pct(3.876) },
        { 'over':  50000, 'base':  1706, 'rate': pct(3.648) },
        { 'over':  25000, 'base':   808, 'rate': pct(3.591) },
        { 'over':  12000, 'base':   349, 'rate': pct(3.534) },
        { 'over':      0, 'base':   349, 'rate': pct(3.534) },
    ]

    if residence == 'NYC':
        for i in city_tax_bracket:
            if gross > i['over']:
                tax_total += i['base']
                tax_total += i['rate'] * ( gross - i['over'] )
                break
    else:
        # We don't yet support other areas.
        raise Exception

    state_tax_bracket = [
        # http://www.tax.ny.gov/pdf/current_forms/it/nyc_tax_rate_schedule.pdf
        { 'over': 1029250, 'base': 69612, 'rate': pct(8.82) },
        { 'over':  205850, 'base': 13209, 'rate': pct(6.85) },
        { 'over':   77150, 'base':  4651, 'rate': pct(6.65) },
        { 'over':   20550, 'base':  1000, 'rate': pct(5.90) },
        { 'over':   11300, 'base':   468, 'rate': pct(5.25) },
        { 'over':    8200, 'base':   328, 'rate': pct(4.50) },
        { 'over':       0, 'base':     0, 'rate': pct(4.00) },
    ]

    for i in state_tax_bracket:
        if gross > i['over']:
            tax_total += i['base']
            tax_total += i['rate'] * ( gross - i['over'] )
            break

    federal_tax_bracket = [
        # where AGI = income - state tax - city tax - standard deduction (5,800) - all other deductions (0)
        # http://www.irs.gov/pub/irs-pdf/i1040.pdf
        # http://www.irs.gov/publications/p501/index.html
        { 'over': 400000, 'base': 116163.75, 'rate': pct(39.6) },
        { 'over': 398350, 'base': 115586.25, 'rate': pct(35) },
        { 'over': 183250, 'base':  44603.25, 'rate': pct(33) },
        { 'over':  87850, 'base':  17891.25, 'rate': pct(28) },
        { 'over':  36250, 'base':   4991.25, 'rate': pct(25) },
        { 'over':   8925, 'base':    892.50, 'rate': pct(15) },
        { 'over':      0, 'base':         0, 'rate': pct(10) },
    ]

    federal_agi = gross - tax_total - deduction

    for i in federal_tax_bracket:
        if federal_agi > i['over']:
            tax_total += i['base']
            tax_total += i['rate'] * ( federal_agi - i['over'] )
            break

    # FICA
    # 6.2% SSA (up to the wage base of 113700)
    # 1.45% Medicare
    # http://www.irs.gov/taxtopics/tc751.html
    tax_total += (federal_agi * pct(1.45))
    wage_base = 113700
    if federal_agi >= wage_base:
        tax_total += (wage_base * pct(6.2))
    else:
        tax_total += (federal_agi * pct(6.2))

    return round(tax_total, 2)
<|endoftext|>


<commit_before>#!/usr/bin/env python
# XDSIntegrate.py
#   Copyright (C) 2006 CCLRC, Graeme Winter
#
#   This code is distributed under the BSD license, a copy of which is
#   included in the root directory of this package.
#
# A wrapper to handle the JOB=INTEGRATE module in XDS.
#

import os
import sys
import shutil
import math
import copy

if not os.environ.has_key('XIA2CORE_ROOT'):
    raise RuntimeError, 'XIA2CORE_ROOT not defined'

if not os.environ.has_key('XIA2_ROOT'):
    raise RuntimeError, 'XIA2_ROOT not defined'

if not os.path.join(os.environ['XIA2CORE_ROOT'],
                    'Python') in sys.path:
    sys.path.append(os.path.join(os.environ['XIA2CORE_ROOT'],
                                 'Python'))

if not os.environ['XIA2_ROOT'] in sys.path:
    sys.path.append(os.environ['XIA2_ROOT'])

from Driver.DriverFactory import DriverFactory

# interfaces that this inherits from ...
from Schema.Interfaces.FrameProcessor import FrameProcessor

# generic helper stuff
from XDS import header_to_xds, xds_check_version_supported, xds_check_error

# specific helper stuff
from XDSIntegrateHelpers import _parse_integrate_lp, \
     _parse_integrate_lp_updates

# global flags etc.
from Handlers.Flags import Flags
from Handlers.Streams import Chatter, Debug

# exceptions

from Schema.Exceptions.BadLatticeError import BadLatticeError

def XDSIntegrate(DriverType = None):

    DriverInstance = DriverFactory.Driver(DriverType)

    class XDSIntegrateWrapper(DriverInstance.__class__,
                              FrameProcessor):
        '''A wrapper for wrapping XDS in integrate mode.'''

        def __init__(self):

            # set up the object ancestors...

            DriverInstance.__class__.__init__(self)
            FrameProcessor.__init__(self)

            # now set myself up...

            self._parallel = Flags.get_parallel()

            if self._parallel <= 1:
                self.set_executable('xds')
            else:
                self.set_executable('xds_par')

            # generic bits

            self._data_range = (0, 0)

            self._input_data_files = { }
            self._output_data_files = { }

            self._input_data_files_list = ['X-CORRECTIONS.pck',
                                           'Y-CORRECTIONS.pck',
                                           'BLANK.pck',
                                           'BKGPIX.pck',
                                           'GAIN.pck',
                                           'XPARM.XDS']

            self._output_data_files_list = ['FRAME.pck']

            self._refined_xparm = False

            self._updates = { }

            # note well - INTEGRATE.HKL is not included in this list
            # because it is likely to be very large - this is treated
            # separately...

            self._integrate_hkl = None

            return

        # getter and setter for input / output data

        def set_input_data_file(self, name, data):
            self._input_data_files[name] = data
            return

        def get_output_data_file(self, name):
            return self._output_data_files[name]

        def get_integrate_hkl(self):
            return self._integrate_hkl

        # this needs setting up from setup_from_image in FrameProcessor

        def set_refined_xparm(self):
            self._refined_xparm = True

        def set_data_range(self, start, end):
            self._data_range = (start, end)

        def set_updates(self, updates):
            self._updates = updates

        def get_updates(self):
            return copy.deepcopy(self._updates)

        def run(self):
            '''Run integrate.'''

            image_header = self.get_header()

            # crank through the header dictionary and replace incorrect
            # information with updated values through the indexer
            # interface if available...

            # need to add distance, wavelength - that should be enough...

            if self.get_distance():
                image_header['distance'] = self.get_distance()

            if self.get_wavelength():
                image_header['wavelength'] = self.get_wavelength()

            header = header_to_xds(image_header)

            xds_inp = open(os.path.join(self.get_working_directory(),
                                        'XDS.INP'), 'w')

            # what are we doing?
            xds_inp.write('JOB=INTEGRATE\n')
            xds_inp.write('MAXIMUM_NUMBER_OF_PROCESSORS=%d\n' % \
                          self._parallel)

            # write out lots of output
            xds_inp.write('TEST=2\n')

            fixed_2401 = True

            if self._refined_xparm and fixed_2401:
                # allow only for crystal movement
                if Flags.get_refine():
                    Debug.write('Refining ORIENTATION CELL')
                    xds_inp.write('REFINE(INTEGRATE)=ORIENTATION CELL\n')
                else:
                    Debug.write('Not refining ORIENTATION CELL')
                    xds_inp.write('REFINE(INTEGRATE)=!\n')
            else:
                # bug 2420 - have found for some examples that the
                # refinement is unstable - perhaps some of this is
                # best postrefined? was ALL
                xds_inp.write('REFINE(INTEGRATE)=BEAM ORIENTATION CELL\n')

            # check for updated input parameters
            if self._updates.has_key('BEAM_DIVERGENCE') and \
                   self._updates.has_key('BEAM_DIVERGENCE_E.S.D.'):
                xds_inp.write(
                    'BEAM_DIVERGENCE=%f BEAM_DIVERGENCE_E.S.D.=%f\n' % \
                    (self._updates['BEAM_DIVERGENCE'],
                     self._updates['BEAM_DIVERGENCE_E.S.D.']))
                Debug.write(
                    'BEAM_DIVERGENCE=%f BEAM_DIVERGENCE_E.S.D.=%f' % \
                    (self._updates['BEAM_DIVERGENCE'],
                     self._updates['BEAM_DIVERGENCE_E.S.D.']))

            if self._updates.has_key('REFLECTING_RANGE') and \
                   self._updates.has_key('REFLECTING_RANGE_E.S.D.'):
                xds_inp.write(
                    'REFLECTING_RANGE=%f REFLECTING_RANGE_E.S.D.=%f\n' % \
                    (self._updates['REFLECTING_RANGE'],
                     self._updates['REFLECTING_RANGE_E.S.D.']))
                Debug.write(
                    'REFLECTING_RANGE=%f REFLECTING_RANGE_E.S.D.=%f' % \
                    (self._updates['REFLECTING_RANGE'],
                     self._updates['REFLECTING_RANGE_E.S.D.']))

            for record in header:
                xds_inp.write('%s\n' % record)

            name_template = os.path.join(self.get_directory(),
                                         self.get_template().replace('#', '?'))

            record = 'NAME_TEMPLATE_OF_DATA_FRAMES=%s\n' % \
                     name_template

            if len(record) < 80:
                xds_inp.write(record)

            else:
                # else we need to make a softlink, then run, then remove
                # softlink....

                try:
                    Debug.write('Linking %s to %s' % (
                        self.get_directory(),
                        os.path.join(self.get_working_directory(),
                                     'xds-image-directory')))

                    os.symlink(self.get_directory(),
                               os.path.join(self.get_working_directory(),
                                            'xds-image-directory'))
                except OSError, e:
                    Debug.write('Error linking: %s' % str(e))

                name_template = os.path.join('xds-image-directory',
                                             self.get_template().replace(
                    '#', '?'))
                record = 'NAME_TEMPLATE_OF_DATA_FRAMES=%s\n' % \
                         name_template

                xds_inp.write(record)

            xds_inp.write('DATA_RANGE=%d %d\n' % self._data_range)
            # xds_inp.write('MINIMUM_ZETA=0.1\n')

            xds_inp.close()

            # copy the input file...
            shutil.copyfile(os.path.join(self.get_working_directory(),
                                         'XDS.INP'),
                            os.path.join(self.get_working_directory(),
                                        'INTEGRATE.INP'))

            # write the input data files...

            for file in self._input_data_files_list:
                open(os.path.join(
                    self.get_working_directory(), file), 'wb').write(
                    self._input_data_files[file])

            self.start()
            self.close_wait()

            xds_check_version_supported(self.get_all_output())
            xds_check_error(self.get_all_output())

            # look for errors
            # like this perhaps - what the hell does this mean?
            #   !!! ERROR !!! "STRONGHKL": ASSERT VIOLATION

            # tidy up...
            try:
                os.remove(os.path.join(self.get_working_directory(),
                                       'xds-image-directory'))
            except OSError, e:
                pass

            # gather the output files

            for file in self._output_data_files_list:
                self._output_data_files[file] = open(os.path.join(
                    self.get_working_directory(), file), 'rb').read()

            self._integrate_hkl = os.path.join(self.get_working_directory(),
                                               'INTEGRATE.HKL')

            # look through integrate.lp for some useful information
            # to help with the analysis

            space_group_number = 0

            for o in open(os.path.join(
                self.get_working_directory(),
                'INTEGRATE.LP')).readlines():
                if 'SPACE_GROUP_NUMBER' in o:
                    space_group_number = int(o.split()[-1])

            stats = _parse_integrate_lp(os.path.join(
                self.get_working_directory(),
                'INTEGRATE.LP'))

            self._updates = _parse_integrate_lp_updates(os.path.join(
                self.get_working_directory(),
                'INTEGRATE.LP'))

            # analyse stats here, perhaps raising an exception if we
            # are unhappy with something, so that the indexing solution
            # can be eliminated in the integrater.

            images = stats.keys()
            images.sort()

            # these may not be present if only a couple of the
            # images were integrated...

            try:

                stddev_pixel = [stats[i]['rmsd_pixel'] for i in images]

                # fix to bug # 2501 - remove the extreme values from this
                # list...

                stddev_pixel = list(set(stddev_pixel))
                stddev_pixel.sort()
                stddev_pixel = stddev_pixel[1:-1]

                low, high = min(stddev_pixel), \
                            max(stddev_pixel)

                Chatter.write('Standard Deviation in pixel range: %f %f' % \
                              (low, high))

                # print a one-spot-per-image rendition of this...
                stddev_pixel = [stats[i]['rmsd_pixel'] for i in images]

                # FIXME need to allow for blank images in here etc.

                status_record = ''
                for stddev in stddev_pixel:
                    if stddev > 2.5:
                        status_record += '!'
                    elif stddev > 1.0:
                        status_record += '%'
                    else:
                        status_record += 'o'

                if len(status_record) > 60:
                    Chatter.write('Integration status per image (60/record):')
                else:
                    Chatter.write('Integration status per image:')

                for chunk in [status_record[i:i + 60] \
                              for i in range(0, len(status_record), 60)]:
                    Chatter.write(chunk)
                Chatter.write(
                    '"o" => ok          "%" => iffy rmsd "!" => bad rmsd')
                Chatter.write(
                    '"O" => overloaded  "#" => many bad  "." => blank')

                # next look for variations in the unit cell parameters
                unit_cells = [stats[i]['unit_cell'] for i in images]

                # compute average
                uc_mean = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0]

                for uc in unit_cells:
                    for j in range(6):
                        uc_mean[j] += uc[j]

                for j in range(6):
                    uc_mean[j] /= len(unit_cells)

                max_rel_dev = 0.0

                for uc in unit_cells:
                    for j in range(6):
                        if (math.fabs(uc[j] - uc_mean[j]) / \
                            uc_mean[j]) > max_rel_dev:
                            max_rel_dev = math.fabs(uc[j] - uc_mean[j]) / \
                                          uc_mean[j]

                Chatter.write('Maximum relative deviation in cell: %.3f' % \
                              max_rel_dev)

                # change - make this only a problem if not triclinic.

                if (high - low) / (0.5 * (high + low)) > 0.5 and \
                       space_group_number > 1:
                    # there was a very large variation in deviation
                    # FIXME 08/JAN/07 this should raise a BadLatticeException

                    # FIXME 26/MAY/07 this should also look for large
                    # variations in the unit cell parameters...
                    # need to have both to distinguish TS01 NAT and TS02.
                    # This doesn't really help...

                    raise BadLatticeError, \
                          'very large variation in pixel deviation'

            except KeyError, e:
                Chatter.write('Refinement not performed...')

            return

    return XDSIntegrateWrapper()

if __name__ == '__main__':

    integrate = XDSIntegrate()
    directory = os.path.join(os.environ['XIA2_ROOT'],
                             'Data', 'Test', 'Images')

    integrate.setup_from_image(os.path.join(directory, '12287_1_E1_001.img'))

    for file in ['X-CORRECTIONS.pck',
                 'Y-CORRECTIONS.pck',
                 'BLANK.pck',
                 'BKGPIX.pck',
                 'GAIN.pck',
                 'XPARM.XDS']:
        integrate.set_input_data_file(file, open(file, 'rb').read())

    integrate.set_data_range(1, 1)

    integrate.run()

<commit_msg>Remove the elimination based on the range of rmsd values...<commit_after>#!/usr/bin/env python
# XDSIntegrate.py
#   Copyright (C) 2006 CCLRC, Graeme Winter
#
#   This code is distributed under the BSD license, a copy of which is
#   included in the root directory of this package.
#
# A wrapper to handle the JOB=INTEGRATE module in XDS.
#

import os
import sys
import shutil
import math
import copy

if not os.environ.has_key('XIA2CORE_ROOT'):
    raise RuntimeError, 'XIA2CORE_ROOT not defined'

if not os.environ.has_key('XIA2_ROOT'):
    raise RuntimeError, 'XIA2_ROOT not defined'

if not os.path.join(os.environ['XIA2CORE_ROOT'],
                    'Python') in sys.path:
    sys.path.append(os.path.join(os.environ['XIA2CORE_ROOT'],
                                 'Python'))

if not os.environ['XIA2_ROOT'] in sys.path:
    sys.path.append(os.environ['XIA2_ROOT'])

from Driver.DriverFactory import DriverFactory

# interfaces that this inherits from ...
from Schema.Interfaces.FrameProcessor import FrameProcessor

# generic helper stuff
from XDS import header_to_xds, xds_check_version_supported, xds_check_error

# specific helper stuff
from XDSIntegrateHelpers import _parse_integrate_lp, \
     _parse_integrate_lp_updates

# global flags etc.
from Handlers.Flags import Flags
from Handlers.Streams import Chatter, Debug

# exceptions

from Schema.Exceptions.BadLatticeError import BadLatticeError

def XDSIntegrate(DriverType = None):

    DriverInstance = DriverFactory.Driver(DriverType)

    class XDSIntegrateWrapper(DriverInstance.__class__,
                              FrameProcessor):
        '''A wrapper for wrapping XDS in integrate mode.'''

        def __init__(self):

            # set up the object ancestors...

            DriverInstance.__class__.__init__(self)
            FrameProcessor.__init__(self)

            # now set myself up...

            self._parallel = Flags.get_parallel()

            if self._parallel <= 1:
                self.set_executable('xds')
            else:
                self.set_executable('xds_par')

            # generic bits

            self._data_range = (0, 0)

            self._input_data_files = { }
            self._output_data_files = { }

            self._input_data_files_list = ['X-CORRECTIONS.pck',
                                           'Y-CORRECTIONS.pck',
                                           'BLANK.pck',
                                           'BKGPIX.pck',
                                           'GAIN.pck',
                                           'XPARM.XDS']

            self._output_data_files_list = ['FRAME.pck']

            self._refined_xparm = False

            self._updates = { }

            # note well - INTEGRATE.HKL is not included in this list
            # because it is likely to be very large - this is treated
            # separately...

            self._integrate_hkl = None

            return

        # getter and setter for input / output data

        def set_input_data_file(self, name, data):
            self._input_data_files[name] = data
            return

        def get_output_data_file(self, name):
            return self._output_data_files[name]

        def get_integrate_hkl(self):
            return self._integrate_hkl

        # this needs setting up from setup_from_image in FrameProcessor

        def set_refined_xparm(self):
            self._refined_xparm = True

        def set_data_range(self, start, end):
            self._data_range = (start, end)

        def set_updates(self, updates):
            self._updates = updates

        def get_updates(self):
            return copy.deepcopy(self._updates)

        def run(self):
            '''Run integrate.'''

            image_header = self.get_header()

            # crank through the header dictionary and replace incorrect
            # information with updated values through the indexer
            # interface if available...

            # need to add distance, wavelength - that should be enough...

            if self.get_distance():
                image_header['distance'] = self.get_distance()

            if self.get_wavelength():
                image_header['wavelength'] = self.get_wavelength()

            header = header_to_xds(image_header)

            xds_inp = open(os.path.join(self.get_working_directory(),
                                        'XDS.INP'), 'w')

            # what are we doing?
            xds_inp.write('JOB=INTEGRATE\n')
            xds_inp.write('MAXIMUM_NUMBER_OF_PROCESSORS=%d\n' % \
                          self._parallel)

            # write out lots of output
            xds_inp.write('TEST=2\n')

            fixed_2401 = True

            if self._refined_xparm and fixed_2401:
                # allow only for crystal movement
                if Flags.get_refine():
                    Debug.write('Refining ORIENTATION CELL')
                    xds_inp.write('REFINE(INTEGRATE)=ORIENTATION CELL\n')
                else:
                    Debug.write('Not refining ORIENTATION CELL')
                    xds_inp.write('REFINE(INTEGRATE)=!\n')
            else:
                # bug 2420 - have found for some examples that the
                # refinement is unstable - perhaps some of this is
                # best postrefined? was ALL
                xds_inp.write('REFINE(INTEGRATE)=BEAM ORIENTATION CELL\n')

            # check for updated input parameters
            if self._updates.has_key('BEAM_DIVERGENCE') and \
                   self._updates.has_key('BEAM_DIVERGENCE_E.S.D.'):
                xds_inp.write(
                    'BEAM_DIVERGENCE=%f BEAM_DIVERGENCE_E.S.D.=%f\n' % \
                    (self._updates['BEAM_DIVERGENCE'],
                     self._updates['BEAM_DIVERGENCE_E.S.D.']))
                Debug.write(
                    'BEAM_DIVERGENCE=%f BEAM_DIVERGENCE_E.S.D.=%f' % \
                    (self._updates['BEAM_DIVERGENCE'],
                     self._updates['BEAM_DIVERGENCE_E.S.D.']))

            if self._updates.has_key('REFLECTING_RANGE') and \
                   self._updates.has_key('REFLECTING_RANGE_E.S.D.'):
                xds_inp.write(
                    'REFLECTING_RANGE=%f REFLECTING_RANGE_E.S.D.=%f\n' % \
                    (self._updates['REFLECTING_RANGE'],
                     self._updates['REFLECTING_RANGE_E.S.D.']))
                Debug.write(
                    'REFLECTING_RANGE=%f REFLECTING_RANGE_E.S.D.=%f' % \
                    (self._updates['REFLECTING_RANGE'],
                     self._updates['REFLECTING_RANGE_E.S.D.']))

            for record in header:
                xds_inp.write('%s\n' % record)

            name_template = os.path.join(self.get_directory(),
                                         self.get_template().replace('#', '?'))

            record = 'NAME_TEMPLATE_OF_DATA_FRAMES=%s\n' % \
                     name_template

            if len(record) < 80:
                xds_inp.write(record)

            else:
                # else we need to make a softlink, then run, then remove
                # softlink....

                try:
                    Debug.write('Linking %s to %s' % (
                        self.get_directory(),
                        os.path.join(self.get_working_directory(),
                                     'xds-image-directory')))

                    os.symlink(self.get_directory(),
                               os.path.join(self.get_working_directory(),
                                            'xds-image-directory'))
                except OSError, e:
                    Debug.write('Error linking: %s' % str(e))

                name_template = os.path.join('xds-image-directory',
                                             self.get_template().replace(
                    '#', '?'))
                record = 'NAME_TEMPLATE_OF_DATA_FRAMES=%s\n' % \
                         name_template

                xds_inp.write(record)

            xds_inp.write('DATA_RANGE=%d %d\n' % self._data_range)
            # xds_inp.write('MINIMUM_ZETA=0.1\n')

            xds_inp.close()

            # copy the input file...
            shutil.copyfile(os.path.join(self.get_working_directory(),
                                         'XDS.INP'),
                            os.path.join(self.get_working_directory(),
                                        'INTEGRATE.INP'))

            # write the input data files...

            for file in self._input_data_files_list:
                open(os.path.join(
                    self.get_working_directory(), file), 'wb').write(
                    self._input_data_files[file])

            self.start()
            self.close_wait()

            xds_check_version_supported(self.get_all_output())
            xds_check_error(self.get_all_output())

            # look for errors
            # like this perhaps - what the hell does this mean?
            #   !!! ERROR !!! "STRONGHKL": ASSERT VIOLATION

            # tidy up...
            try:
                os.remove(os.path.join(self.get_working_directory(),
                                       'xds-image-directory'))
            except OSError, e:
                pass

            # gather the output files

            for file in self._output_data_files_list:
                self._output_data_files[file] = open(os.path.join(
                    self.get_working_directory(), file), 'rb').read()

            self._integrate_hkl = os.path.join(self.get_working_directory(),
                                               'INTEGRATE.HKL')

            # look through integrate.lp for some useful information
            # to help with the analysis

            space_group_number = 0

            for o in open(os.path.join(
                self.get_working_directory(),
                'INTEGRATE.LP')).readlines():
                if 'SPACE_GROUP_NUMBER' in o:
                    space_group_number = int(o.split()[-1])

            stats = _parse_integrate_lp(os.path.join(
                self.get_working_directory(),
                'INTEGRATE.LP'))

            self._updates = _parse_integrate_lp_updates(os.path.join(
                self.get_working_directory(),
                'INTEGRATE.LP'))

            # analyse stats here, perhaps raising an exception if we
            # are unhappy with something, so that the indexing solution
            # can be eliminated in the integrater.

            images = stats.keys()
            images.sort()

            # these may not be present if only a couple of the
            # images were integrated...

            try:

                stddev_pixel = [stats[i]['rmsd_pixel'] for i in images]

                # fix to bug # 2501 - remove the extreme values from this
                # list...

                stddev_pixel = list(set(stddev_pixel))
                stddev_pixel.sort()
                stddev_pixel = stddev_pixel[1:-1]

                low, high = min(stddev_pixel), \
                            max(stddev_pixel)

                Chatter.write('Standard Deviation in pixel range: %f %f' % \
                              (low, high))

                # print a one-spot-per-image rendition of this...
                stddev_pixel = [stats[i]['rmsd_pixel'] for i in images]

                # FIXME need to allow for blank images in here etc.

                status_record = ''
                for stddev in stddev_pixel:
                    if stddev > 2.5:
                        status_record += '!'
                    elif stddev > 1.0:
                        status_record += '%'
                    else:
                        status_record += 'o'

                if len(status_record) > 60:
                    Chatter.write('Integration status per image (60/record):')
                else:
                    Chatter.write('Integration status per image:')

                for chunk in [status_record[i:i + 60] \
                              for i in range(0, len(status_record), 60)]:
                    Chatter.write(chunk)
                Chatter.write(
                    '"o" => ok          "%" => iffy rmsd "!" => bad rmsd')
                Chatter.write(
                    '"O" => overloaded  "#" => many bad  "." => blank')

                # next look for variations in the unit cell parameters
                unit_cells = [stats[i]['unit_cell'] for i in images]

                # compute average
                uc_mean = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0]

                for uc in unit_cells:
                    for j in range(6):
                        uc_mean[j] += uc[j]

                for j in range(6):
                    uc_mean[j] /= len(unit_cells)

                max_rel_dev = 0.0

                for uc in unit_cells:
                    for j in range(6):
                        if (math.fabs(uc[j] - uc_mean[j]) / \
                            uc_mean[j]) > max_rel_dev:
                            max_rel_dev = math.fabs(uc[j] - uc_mean[j]) / \
                                          uc_mean[j]

                Chatter.write('Maximum relative deviation in cell: %.3f' % \
                              max_rel_dev)

                # change - make this only a problem if not triclinic.

                if (high - low) / (0.5 * (high + low)) > 0.5 and \
                       space_group_number > 1 and False:
                    # there was a very large variation in deviation
                    # FIXME 08/JAN/07 this should raise a BadLatticeException

                    # FIXME 26/MAY/07 this should also look for large
                    # variations in the unit cell parameters...
                    # need to have both to distinguish TS01 NAT and TS02.
                    # This doesn't really help...

                    # 23/JUL/07 no longer use this as a metric -
                    # use the results of postrefinement...

                    raise BadLatticeError, \
                          'very large variation in pixel deviation'

            except KeyError, e:
                Chatter.write('Refinement not performed...')

            return

    return XDSIntegrateWrapper()

if __name__ == '__main__':

    integrate = XDSIntegrate()
    directory = os.path.join(os.environ['XIA2_ROOT'],
                             'Data', 'Test', 'Images')

    integrate.setup_from_image(os.path.join(directory, '12287_1_E1_001.img'))

    for file in ['X-CORRECTIONS.pck',
                 'Y-CORRECTIONS.pck',
                 'BLANK.pck',
                 'BKGPIX.pck',
                 'GAIN.pck',
                 'XPARM.XDS']:
        integrate.set_input_data_file(file, open(file, 'rb').read())

    integrate.set_data_range(1, 1)

    integrate.run()

<|endoftext|>


<commit_before>### import ####################################################################

import numpy as np
import matplotlib.pyplot as plt

### define ####################################################################

### subplot ###################################################################

class Subplot:

    def __init__(self, ax, energies, number_of_interactions=4, title='',
                 title_font_size=16, state_names=None, virtual=[None],
                 state_font_size=8, state_text_buffer=0.5, label_side='left'):
        '''
        Subplot.

        Parameters
        ----------
        ax : matplotlib axis
            The axis.
        energies : 1D array-like
            Energies (scaled between 0 and 1)
        number_of_interactions : integer
            Number of interactions in diagram.
        title : string (optional)
            Title of subplot. Default is empty string.
        state_names: list of str (optional)
            list of the names of the states
        virtual: list of ints (optional)
            list of indexes of any vitual energy states
        state_font_size: numtype (optional)
            font size for the state lables
        state_text_buffer: numtype (optional)
            space between the energy level bars and the state labels
        '''
        self.ax = ax
        self.energies = energies
        self.interactions = number_of_interactions
        self.state_names = state_names

        # Plot Energy Levels
        for i in range(len(self.energies)):
            if i in virtual:
                linestyle = '--'
            else:
                linestyle = '-'
            self.ax.axhline(self.energies[i], color='k', linewidth=2,
                            ls=linestyle, zorder=5)

        # add state names
        if isinstance(state_names, list):
            for i in range(len(self.energies)):
                if label_side == 'left':
                    ax.text(-state_text_buffer, energies[i], state_names[i],
                            fontsize=state_font_size, verticalalignment='center',
                            horizontalalignment='center')
                elif label_side == 'right':
                    ax.text(1 + state_text_buffer, energies[i], state_names[i],
                            fontsize=state_font_size, verticalalignment='center',
                            horizontalalignment='center')
        # calculate interaction_positons
        self.x_pos = np.linspace(0, 1, number_of_interactions)
        # set limits
        self.ax.set_xlim(-0.1, 1.1)
        self.ax.set_ylim(-0.01, 1.01)
        # remove guff
        self.ax.axis('off')
        # title
        self.ax.set_title(title, fontsize=title_font_size)

    def add_arrow(self, index, between, kind, label='', head_length=0.075,
                  head_aspect=2, font_size=7, color='k'):
        '''
        Add an arrow to the WMEL diagram.

        Parameters
        ----------
        index : integer
            The interaction, or start and stop interaction for the arrow.
        between : 2-element iterable of integers
            The inital and final state of the arrow
        kind : {'ket', 'bra'}
            The kind of interaction.
        label : string (optional)
            Interaction label. Default is empty string.
        head_length: number (optional)
            size of arrow head
        font_size : number (optional)
            Label font size. Default is 14.
        color : matplotlib color (optional)
            Arrow color. Default is black.

        Returns
        -------
        [line,arrow_head,text]
        '''
        if hasattr(index, 'index'):
            x_pos = list(index)
        else:
            x_pos = [index] * 2
        x_pos = [np.linspace(0, 1, self.interactions)[i] for i in x_pos]

        # calculate arrow length
        arrow_length = self.energies[between[1]] - self.energies[between[0]]
        arrow_end = self.energies[between[1]]
        if arrow_length > 0:
            direction = 1
            y_poss = [self.energies[between[0]], self.energies[between[1]] - head_length]
        elif arrow_length < 0:
            direction = -1
            y_poss = [self.energies[between[0]], self.energies[between[1]] + head_length]
        else:
            raise ValueError('between invalid!')

        length = abs(y_poss[0] - y_poss[1])
        if kind == 'ket':
            line = self.ax.plot(x_pos, y_poss, linestyle='-', color=color,
                                linewidth=2, zorder=9)
        elif kind == 'bra':
            line = self.ax.plot(x_pos, y_poss, linestyle='--', color=color,
                                linewidth=2, zorder=9)
        elif kind == 'out':
            yi = np.linspace(y_poss[0], y_poss[1], 100)
            xi = np.sin((yi - y_poss[0]) * int((1 / length) * 20)
                        * 2 * np.pi * length) / 40 + x_pos[0]
            line = self.ax.plot(xi, yi, linestyle='-', color=color,
                                linewidth=2, solid_capstyle='butt', zorder=9)
        else:
            raise ValueError('kind is not \'ket\', \'out\', or \'bra\'.')
        # add arrow head
        arrow_head = self.ax.arrow(x_pos[1], arrow_end - head_length * direction,
                                   0, 0.0001 * direction,
                                   head_width=head_length * head_aspect,
                                   head_length=head_length,
                                   fc=color, ec=color, linestyle='solid',
                                   linewidth=0, zorder=10)
        # add text
        text = self.ax.text(np.mean(x_pos), -0.1, label, fontsize=font_size,
                            horizontalalignment='center')
        return line, arrow_head, text

### artist ####################################################################

class Artist:

    def __init__(self, size, energies, state_names=None,
                 number_of_interactions=4, virtual=[None],
                 state_font_size=8, state_text_buffer=0.5):
        '''
        virtual a list of indicies
        '''
        # create figure
        figsize = [int(size[0] * ((number_of_interactions + 1.) / 6.)), size[1] * 2.5]
        fig, (subplots) = plt.subplots(size[1], size[0], figsize=figsize)
        self.fig = fig
        # wrap subplots if need be
        if size == [1, 1]:
            self.subplots = np.array([[subplots]])
            plt.subplots_adjust(left=0.3)
        elif size[1] == 1:
            self.subplots = np.array([subplots])
        else:
            self.subplots = subplots
        # add energy levels
        self.energies = energies
        for plot in self.subplots.flatten():
            for i in range(len(self.energies)):
                if i in virtual:
                    linestyle = '--'
                else:
                    linestyle = '-'
                plot.axhline(energies[i], color='k', linewidth=2, linestyle=linestyle)
        # add state names to leftmost plots
        if state_names:
            for i in range(size[1]):
                plot = self.subplots[i][0]
                for i in range(len(self.energies)):
                    plot.text(-state_text_buffer, energies[i], state_names[i], fontsize=state_font_size,
                              verticalalignment='center', horizontalalignment='center')
        # calculate interaction_positons
        self.x_pos = np.linspace(0, 1, number_of_interactions)
        # plot cleans up a bunch - call it now as well as later
        self.plot()

    def label_rows(self, labels, font_size=15, text_buffer=1.5):
        for i in range(len(self.subplots)):
            plot = self.subplots[i][-1]
            plot.text(text_buffer, 0.5, labels[i], fontsize=font_size,
                      verticalalignment='center', horizontalalignment='center')

    def label_columns(self, labels, font_size=15, text_buffer=1.15):
        for i in range(len(labels)):
            plot = self.subplots[0][i]
            plot.text(0.5, text_buffer, labels[i], fontsize=font_size,
                      verticalalignment='center', horizontalalignment='center')

    def clear_diagram(self, diagram):
        plot = self.subplots[diagram[1]][diagram[0]]
        plot.cla()

    def add_arrow(self, diagram, number, between, kind, label='',
                  head_length=0.075, font_size=7, color='k'):
        '''
        kind one in [ket, bra, out] \n
        returns [line, arrow_head, text]
        '''
        column, row = diagram
        x_pos = self.x_pos[number]
        # calculate arrow length
        arrow_length = self.energies[between[1]] - self.energies[between[0]]
        arrow_end = self.energies[between[1]]
        if arrow_length > 0:
            direction = 1
            y_poss = [self.energies[between[0]], self.energies[between[1]] - head_length]
        elif arrow_length < 0:
            direction = -1
            y_poss = [self.energies[between[0]], self.energies[between[1]] + head_length]
        else:
            print('between invalid!')
            return
        subplot = self.subplots[row][column]
        # add line
        length = abs(y_poss[0] - y_poss[1])
        if kind == 'ket':
            line = subplot.plot([x_pos, x_pos], y_poss, linestyle='-', color=color, linewidth=2)
        elif kind == 'bra':
            line = subplot.plot([x_pos, x_pos], y_poss, linestyle='--', color=color, linewidth=2)
        elif kind == 'out':
            yi = np.linspace(y_poss[0], y_poss[1], 100)
            xi = np.sin((yi - y_poss[0]) * int((1 / length) * 20)
                        * 2 * np.pi * length) / 40 + x_pos
            line = subplot.plot(xi, yi, linestyle='-', color=color,
                                linewidth=2, solid_capstyle='butt')
        # add arrow head
        arrow_head = subplot.arrow(self.x_pos[number], arrow_end - head_length * direction,
                                   0, 0.0001 * direction,
                                   head_width=head_length * 2,
                                   head_length=head_length,
                                   fc=color, ec=color, linestyle='solid', linewidth=0)
        # add text
        text = subplot.text(self.x_pos[number], -0.1, label,
                            fontsize=font_size, horizontalalignment='center')
        return line, arrow_head, text

    def plot(self, save_path=None, close=False, bbox_inches='tight', pad_inches=1):
        # final manipulations
        for plot in self.subplots.flatten():
            # set limits
            plot.set_xlim(-0.1, 1.1)
            plot.set_ylim(-0.1, 1.1)
            # remove guff
            plot.axis('off')
        # save
        if save_path:
            plt.savefig(save_path, transparent=True, dpi=300,
                        bbox_inches=bbox_inches, pad_inches=pad_inches)
        # close
        if close:
            plt.close()

### testing ###################################################################

if __name__ == '__main__':
    # testing code

    plt.close('all')

    diagram = Artist(size=[6, 3],
                     energies=[0., 0.4, 0.6, 1.],
                     state_names=['g', 'a', 'b', 'a+b'])

    diagram.label_rows([r'$\mathrm{\alpha}$', r'$\mathrm{\beta}$', r'$\mathrm{\gamma}$'])
    diagram.label_columns(['I', 'II', 'III', 'IV', 'V', 'VI'])

    # pw1 alpha
    diagram.add_arrow([0, 0], 0, [0, 2], 'ket', '1')
    diagram.add_arrow([0, 0], 1, [0, 1], 'bra', '-2')
    diagram.add_arrow([0, 0], 2, [1, 0], 'bra', '2\'')
    diagram.add_arrow([0, 0], 3, [3, 0], 'out')

    diagram.add_arrow([1, 0], 3, [2, 0], 'out')

    diagram.clear_diagram([2, 1])

    diagram.plot('WMEL_out.png')
<commit_msg>better WMEL defaults (resolves #119)<commit_after>### import ####################################################################

import numpy as np
import matplotlib.pyplot as plt

### define ####################################################################

### subplot ###################################################################

class Subplot:

    def __init__(self, ax, energies, number_of_interactions=4, title='',
                 title_font_size=16, state_names=None, virtual=[None],
                 state_font_size=14, state_text_buffer=0.5, label_side='left'):
        """ Subplot.

        Parameters
        ----------
        ax : matplotlib axis
            The axis.
        energies : 1D array-like
            Energies (scaled between 0 and 1)
        number_of_interactions : integer
            Number of interactions in diagram.
        title : string (optional)
            Title of subplot. Default is empty string.
        state_names: list of str (optional)
            list of the names of the states
        virtual: list of ints (optional)
            list of indexes of any vitual energy states
        state_font_size: numtype (optional)
            font size for the state lables
        state_text_buffer: numtype (optional)
            space between the energy level bars and the state labels
        """
        self.ax = ax
        self.energies = energies
        self.interactions = number_of_interactions
        self.state_names = state_names

        # Plot Energy Levels
        for i in range(len(self.energies)):
            if i in virtual:
                linestyle = '--'
            else:
                linestyle = '-'
            self.ax.axhline(self.energies[i], color='k', linewidth=2,
                            ls=linestyle, zorder=5)

        # add state names
        if isinstance(state_names, list):
            for i in range(len(self.energies)):
                if label_side == 'left':
                    ax.text(-state_text_buffer, energies[i], state_names[i],
                            fontsize=state_font_size, verticalalignment='center',
                            horizontalalignment='center')
                elif label_side == 'right':
                    ax.text(1 + state_text_buffer, energies[i], state_names[i],
                            fontsize=state_font_size, verticalalignment='center',
                            horizontalalignment='center')
        # calculate interaction_positons
        self.x_pos = np.linspace(0, 1, number_of_interactions)
        # set limits
        self.ax.set_xlim(-0.1, 1.1)
        self.ax.set_ylim(-0.01, 1.01)
        # remove guff
        self.ax.axis('off')
        # title
        self.ax.set_title(title, fontsize=title_font_size)

    def add_arrow(self, index, between, kind, label='', head_length=0.1,
                  head_aspect=2, font_size=14, color='k'):
        """ Add an arrow to the WMEL diagram.

        Parameters
        ----------
        index : integer
            The interaction, or start and stop interaction for the arrow.
        between : 2-element iterable of integers
            The inital and final state of the arrow
        kind : {'ket', 'bra'}
            The kind of interaction.
        label : string (optional)
            Interaction label. Default is empty string.
        head_length: number (optional)
            size of arrow head
        font_size : number (optional)
            Label font size. Default is 14.
        color : matplotlib color (optional)
            Arrow color. Default is black.

        Returns
        -------
        [line,arrow_head,text]
        """
        if hasattr(index, 'index'):
            x_pos = list(index)
        else:
            x_pos = [index] * 2
        x_pos = [np.linspace(0, 1, self.interactions)[i] for i in x_pos]

        # calculate arrow length
        arrow_length = self.energies[between[1]] - self.energies[between[0]]
        arrow_end = self.energies[between[1]]
        if arrow_length > 0:
            direction = 1
            y_poss = [self.energies[between[0]], self.energies[between[1]] - head_length]
        elif arrow_length < 0:
            direction = -1
            y_poss = [self.energies[between[0]], self.energies[between[1]] + head_length]
        else:
            raise ValueError('between invalid!')

        length = abs(y_poss[0] - y_poss[1])
        if kind == 'ket':
            line = self.ax.plot(x_pos, y_poss, linestyle='-', color=color,
                                linewidth=2, zorder=9)
        elif kind == 'bra':
            line = self.ax.plot(x_pos, y_poss, linestyle='--', color=color,
                                linewidth=2, zorder=9)
        elif kind == 'out':
            yi = np.linspace(y_poss[0], y_poss[1], 100)
            xi = np.sin((yi - y_poss[0]) * int((1 / length) * 20)
                        * 2 * np.pi * length) / 40 + x_pos[0]
            line = self.ax.plot(xi, yi, linestyle='-', color=color,
                                linewidth=2, solid_capstyle='butt', zorder=9)
        else:
            raise ValueError('kind is not \'ket\', \'out\', or \'bra\'.')
        # add arrow head
        arrow_head = self.ax.arrow(x_pos[1], arrow_end - head_length * direction,
                                   0, 0.0001 * direction,
                                   head_width=head_length * head_aspect,
                                   head_length=head_length,
                                   fc=color, ec=color, linestyle='solid',
                                   linewidth=0, zorder=10)
        # add text
        text = self.ax.text(np.mean(x_pos), -0.15, label, fontsize=font_size,
                            horizontalalignment='center')
        return line, arrow_head, text

### artist ####################################################################

class Artist:

    def __init__(self, size, energies, state_names=None,
                 number_of_interactions=4, virtual=[None],
                 state_font_size=8, state_text_buffer=0.5):
        """ virtual a list of indicies """
        # create figure
        figsize = [int(size[0] * ((number_of_interactions + 1.) / 6.)), size[1] * 2.5]
        fig, (subplots) = plt.subplots(size[1], size[0], figsize=figsize)
        self.fig = fig
        # wrap subplots if need be
        if size == [1, 1]:
            self.subplots = np.array([[subplots]])
            plt.subplots_adjust(left=0.3)
        elif size[1] == 1:
            self.subplots = np.array([subplots])
        else:
            self.subplots = subplots
        # add energy levels
        self.energies = energies
        for plot in self.subplots.flatten():
            for i in range(len(self.energies)):
                if i in virtual:
                    linestyle = '--'
                else:
                    linestyle = '-'
                plot.axhline(energies[i], color='k', linewidth=2, linestyle=linestyle)
        # add state names to leftmost plots
        if state_names:
            for i in range(size[1]):
                plot = self.subplots[i][0]
                for i in range(len(self.energies)):
                    plot.text(-state_text_buffer, energies[i], state_names[i], fontsize=state_font_size,
                              verticalalignment='center', horizontalalignment='center')
        # calculate interaction_positons
        self.x_pos = np.linspace(0, 1, number_of_interactions)
        # plot cleans up a bunch - call it now as well as later
        self.plot()

    def label_rows(self, labels, font_size=15, text_buffer=1.5):
        for i in range(len(self.subplots)):
            plot = self.subplots[i][-1]
            plot.text(text_buffer, 0.5, labels[i], fontsize=font_size,
                      verticalalignment='center', horizontalalignment='center')

    def label_columns(self, labels, font_size=15, text_buffer=1.15):
        for i in range(len(labels)):
            plot = self.subplots[0][i]
            plot.text(0.5, text_buffer, labels[i], fontsize=font_size,
                      verticalalignment='center', horizontalalignment='center')

    def clear_diagram(self, diagram):
        plot = self.subplots[diagram[1]][diagram[0]]
        plot.cla()

    def add_arrow(self, diagram, number, between, kind, label='',
                  head_length=0.075, font_size=7, color='k'):
        """
        kind one in [ket, bra, out] \n
        returns [line, arrow_head, text]
        """
        column, row = diagram
        x_pos = self.x_pos[number]
        # calculate arrow length
        arrow_length = self.energies[between[1]] - self.energies[between[0]]
        arrow_end = self.energies[between[1]]
        if arrow_length > 0:
            direction = 1
            y_poss = [self.energies[between[0]], self.energies[between[1]] - head_length]
        elif arrow_length < 0:
            direction = -1
            y_poss = [self.energies[between[0]], self.energies[between[1]] + head_length]
        else:
            print('between invalid!')
            return
        subplot = self.subplots[row][column]
        # add line
        length = abs(y_poss[0] - y_poss[1])
        if kind == 'ket':
            line = subplot.plot([x_pos, x_pos], y_poss, linestyle='-', color=color, linewidth=2)
        elif kind == 'bra':
            line = subplot.plot([x_pos, x_pos], y_poss, linestyle='--', color=color, linewidth=2)
        elif kind == 'out':
            yi = np.linspace(y_poss[0], y_poss[1], 100)
            xi = np.sin((yi - y_poss[0]) * int((1 / length) * 20)
                        * 2 * np.pi * length) / 40 + x_pos
            line = subplot.plot(xi, yi, linestyle='-', color=color,
                                linewidth=2, solid_capstyle='butt')
        # add arrow head
        arrow_head = subplot.arrow(self.x_pos[number], arrow_end - head_length * direction,
                                   0, 0.0001 * direction,
                                   head_width=head_length * 2,
                                   head_length=head_length,
                                   fc=color, ec=color, linestyle='solid', linewidth=0)
        # add text
        text = subplot.text(self.x_pos[number], -0.1, label,
                            fontsize=font_size, horizontalalignment='center')
        return line, arrow_head, text

    def plot(self, save_path=None, close=False, bbox_inches='tight', pad_inches=1):
        # final manipulations
        for plot in self.subplots.flatten():
            # set limits
            plot.set_xlim(-0.1, 1.1)
            plot.set_ylim(-0.1, 1.1)
            # remove guff
            plot.axis('off')
        # save
        if save_path:
            plt.savefig(save_path, transparent=True, dpi=300,
                        bbox_inches=bbox_inches, pad_inches=pad_inches)
        # close
        if close:
            plt.close()

### testing ###################################################################

if __name__ == '__main__':
    # testing code

    plt.close('all')

    diagram = Artist(size=[6, 3],
                     energies=[0., 0.4, 0.6, 1.],
                     state_names=['g', 'a', 'b', 'a+b'])

    diagram.label_rows([r'$\mathrm{\alpha}$', r'$\mathrm{\beta}$', r'$\mathrm{\gamma}$'])
    diagram.label_columns(['I', 'II', 'III', 'IV', 'V', 'VI'])

    # pw1 alpha
    diagram.add_arrow([0, 0], 0, [0, 2], 'ket', '1')
    diagram.add_arrow([0, 0], 1, [0, 1], 'bra', '-2')
    diagram.add_arrow([0, 0], 2, [1, 0], 'bra', '2\'')
    diagram.add_arrow([0, 0], 3, [3, 0], 'out')

    diagram.add_arrow([1, 0], 3, [2, 0], 'out')

    diagram.clear_diagram([2, 1])

    diagram.plot('WMEL_out.png')
<|endoftext|>