<commit_before>from __future__ import (absolute_import, division, print_function,
                        unicode_literals)
import six

import logging
import h5py
import numpy as np
import os.path
import tifffile

from .handlers_base import HandlerBase
from .readers.spe import PrincetonSPEFile

logger = logging.getLogger(__name__)

from pims import FramesSequence, Frame

# The ImageCube class is used for a per event representation of
# a dataset

class ImageStack(FramesSequence):
    "One of these represents the data from an event: (num_images x w x h)"
    def __init__(self, dataset, start, stop):
        # `start` and `stop` are the limits of this cube
        # i indexes within the cube
        self._start = start
        self._stop = stop
        self._dataset = dataset

        # work around inconsistent naming choices in databroker's Image object
        self.dtype = self.pixel_type
        self.shape = self.frame_shape

    def get_frame(self, i):
        return Frame(self._dataset[self._start + i], frame_no=i)

    def __len__(self):
        return self._stop - self._start

    @property
    def pixel_type(self):
        return self._dataset.dtype

    @property
    def frame_shape(self):
        return self._dataset.shape[1:]

class IntegrityError(Exception):
    pass

class AreaDetectorSPEHandler(HandlerBase):
    specs = {'AD_SPE'} | HandlerBase.specs

    def __init__(self, fpath, template, filename,
                 frame_per_point=1):
        self._path = os.path.join(fpath, '')
        self._fpp = frame_per_point
        self._template = template
        self._filename = filename
        self._f_cache = dict()

    def __call__(self, point_number):
        if point_number not in self._f_cache:
            fname = self._template % (self._path,
                                      self._filename,
                                      point_number)
            spe_obj = PrincetonSPEFile(fname)
            self._f_cache[point_number] = spe_obj

        spe = self._f_cache[point_number]
        data = spe.getData()

        if data.shape[0] != self._fpp:
            raise IntegrityError("expected {} frames, found {} frames".format(
                                 self._fpp, data.shape[0]))
        return data.squeeze()

    def get_file_list(self, datum_kwarg_gen):
        return [self._template % (self._path,
                                  self._filename,
                                  d['point_number'])
                for d in datum_kwarg_gen]

class AreaDetectorTiffHandler(HandlerBase):
    specs = {'AD_TIFF'} | HandlerBase.specs

    def __init__(self, fpath, template, filename, frame_per_point=1):
        self._path = os.path.join(fpath, '')
        self._fpp = frame_per_point
        self._template = template
        self._filename = filename

    def _fnames_for_point(self, point_number):
        start, stop = point_number * self._fpp, (point_number + 1) * self._fpp
        for j in range(start, stop):
            yield self._template % (self._path, self._filename, j)

    def __call__(self, point_number):
        ret = []
        for fn in self._fnames_for_point(point_number):
            with tifffile.TiffFile(fn) as tif:
                ret.append(tif.asarray())
        return np.array(ret).squeeze()

    def get_file_list(self, datum_kwargs):
        ret = []
        for d_kw in datum_kwargs:
            ret.extend(self._fnames_for_point(**d_kw))
        return ret

class DummyAreaDetectorHandler(HandlerBase):
    def __init__(self, fpath, frame_per_point=1, **kwargs):

        self._fpp = frame_per_point

    def __call__(self, **kwargs):
        out_stack = np.ones((self._fpp, 10, 10)) * np.nan
        # return stacked and squeezed results
        return out_stack.squeeze()

class HDF5DatasetSliceHandler(HandlerBase):
    """
    Handler for data stored in one Dataset of an HDF5 file.

    Parameters
    ----------
    filename : string
        path to HDF5 file
    key : string
        key of the single HDF5 Dataset used by this Handler
    frame_per_point : integer, optional
        number of frames to return as one datum, default 1
    """
    def __init__(self, filename, key, frame_per_point=1):
        self._fpp = frame_per_point
        self._filename = filename
        self._key = key
        self._file = None
        self._dataset = None
        self._data_objects = {}
        self.open()

    def get_file_list(self, datum_kwarg_gen):
        return [self._filename]

    def __call__(self, point_number):
        # Don't read out the dataset until it is requested for the first time.
        if not self._dataset:
            self._dataset = self._file[self._key]

        if point_number not in self._data_objects:
            start = point_number * self._fpp
            stop = (point_number + 1) * self._fpp
            self._data_objects[point_number] = ImageStack(self._dataset,
                                                          start, stop)
        return self._data_objects[point_number]

    def open(self):
        if self._file:
            return

        self._file = h5py.File(self._filename, 'r')

    def close(self):
        super(HDF5DatasetSliceHandler, self).close()
        self._file.close()
        self._file = None

class AreaDetectorHDF5Handler(HDF5DatasetSliceHandler):
    """
    Handler for the 'AD_HDF5' spec used by Area Detectors.

    In this spec, the key (i.e., HDF5 dataset path) is always
    '/entry/data/data'.

    Parameters
    ----------
    filename : string
        path to HDF5 file
    frame_per_point : integer, optional
        number of frames to return as one datum, default 1
    """
    specs = {'AD_HDF5'} | HDF5DatasetSliceHandler.specs

    def __init__(self, filename, frame_per_point=1):
        hardcoded_key = '/entry/data/data'
        super(AreaDetectorHDF5Handler, self).__init__(
            filename=filename, key=hardcoded_key,
            frame_per_point=frame_per_point)

class AreaDetectorHDF5SWMRHandler(AreaDetectorHDF5Handler):
    """
    Handler for the 'AD_HDF5_SWMR' spec used by Area Detectors.

    In this spec, the key (i.e., HDF5 dataset path) is always
    '/entry/data/data'.

    Parameters
    ----------
    filename : string
        path to HDF5 file
    frame_per_point : integer, optional
        number of frames to return as one datum, default 1
    """
    specs = {'AD_HDF5_SWMR'} | HDF5DatasetSliceHandler.specs

    def open(self):
        if self._file:
            return

        self._file = h5py.File(self._filename, 'r', swmr=True)

    def __call__(self, point_number):
        if self._dataset is not None:
            self._dataset.id.refresh()
        rtn = super(AreaDetectorHDF5SWMRHandler, self).__call__(
            point_number)

        return rtn

class AreaDetectorHDF5TimestampHandler(HandlerBase):
    """ Handler to retrieve timestamps from Areadetector HDF5 File

    In this spec, the timestamps of the images are read.

    Parameters
    ----------
    filename : string
        path to HDF5 file
    frame_per_point : integer, optional
        number of frames to return as one datum, default 1
    """
    specs = {'AD_HDF5_TS'} | HandlerBase.specs

    def __init__(self, filename, frame_per_point=1):
        self._fpp = frame_per_point
        self._filename = filename
        self._key = ['/entry/instrument/NDAttributes/NDArrayEpicsTSSec',
                     '/entry/instrument/NDAttributes/NDArrayEpicsTSnSec']
        self._file = None
        self._dataset1 = None
        self._dataset2 = None
        self.open()

    def __call__(self, point_number):
        # Don't read out the dataset until it is requested for the first time.
        if not self._dataset1:
            self._dataset1 = self._file[self._key[0]]
        if not self._dataset2:
            self._dataset2 = self._file[self._key[1]]
        start, stop = point_number * self._fpp, (point_number + 1) * self._fpp
        rtn = self._dataset1[start:stop].squeeze()
        rtn = rtn + (self._dataset2[start:stop].squeeze() * 1e-9)
        return rtn

    def open(self):
        if self._file:
            return
        self._file = h5py.File(self._filename, 'r')

    def close(self):
        super(AreaDetectorHDF5TimestampHandler, self).close()
        self._file.close()
        self._file = None

class AreaDetectorHDF5SWMRTimestampHandler(AreaDetectorHDF5TimestampHandler):
    """ Handler to retrieve timestamps from Areadetector HDF5 File

    In this spec, the timestamps of the images are read. Reading
    is done using SWMR option to allow read during processing

    Parameters
    ----------
    filename : string
        path to HDF5 file
    frame_per_point : integer, optional
        number of frames to return as one datum, default 1
    """
    specs = {'AD_HDF5_SWMR_TS'} | HandlerBase.specs

    def open(self):
        if self._file:
            return
        self._file = h5py.File(self._filename, 'r', swmr=True)

    def __call__(self, point_number):
        if (self._dataset1 is not None) and (self._dataset2 is not None):
            self._dataset.id.refresh()
        rtn = super(AreaDetectorHDF5SWMRTimestampHandler, self).__call__(
            point_number)
        return rtn

class _HdfMapsHandlerBase(HDF5DatasetSliceHandler):
    """
    Reader for XRF data stored in hdf5 files.

    The data set is assumed to be in a group called MAPS and stored
    as a 3D array ordered [energy, x, y].

    Parameters
    ----------
    filename : str
        Path to physical location of file
    dset_path : str
        The path to the dataset inside of 'MAPS'
    """
    def __init__(self, filename, dset_path):
        self._filename = filename
        self._dset_path = dset_path
        self._file = None
        self._dset = None
        self._swmr = False
        self.open()

    def open(self):
        """
        Open the file for reading.

        Provided as a stand alone function to allow re-opening of the handler
        """
        super(_HdfMapsHandlerBase, self).open()
        self._dset = self._file['/'.join(['MAPS', self._dset_path])]

    def __call__(self):

        if not self._file:
            raise RuntimeError("File is not open")

        if self._swmr:
            self._dataset.id.refresh()

class HDFMapsSpectrumHandler(_HdfMapsHandlerBase):
    """
    Handler which selects energy spectrum from
    a MAPS XRF data product.
    """
    specs = {'MAPS_SPECTRUM'} | _HdfMapsHandlerBase.specs

    def __call__(self, x, y):
        """
        Return the spectrum at the x, y position

        Parameters
        ----------
        x : int
            raster index in the x direction

        y : int
            raster index in the y direction

        Returns
        -------
        spectrum : ndarray
            The MCA channels
        """
        super(HDFMapsSpectrumHandler, self).__call__()
        return self._dset[:, x, y]

class HDFMapsEnergyHandler(_HdfMapsHandlerBase):
    """
    Handler which select fixed-energy slices from
    a MAPS XRF data file.
    """
    specs = {'MAPS_PLANE'} | _HdfMapsHandlerBase.specs

    def __call__(self, e_index):
        """
        Return the raster plane at a fixed energy

        Parameters
        ----------
        e_index : int
            The index of the engery

        Returns
        -------
        plane : ndarray
            The raster image at a fixed energy.
        """
        super(HDFMapsEnergyHandler, self).__call__()
        return self._dset[e_index, :, :]

class NpyHandler(HandlerBase):
    """
    Class to deal with reading npy files

    Parameters
    ----------
    fpath : str
        Path to file

    mmap_mode : {'r', 'r+', c}, optional
        memmap mode to use to open file
    """
    specs = {'npy'} | HandlerBase.specs

    def __init__(self, filename, mmap_mode=None):
        self._mmap_mode = mmap_mode
        if not os.path.exists(filename):
            raise IOError("the requested file {fpath} does not exist")
        self._fpath = filename

    def __call__(self):
        return np.load(self._fpath, self._mmap_mode)

    def get_file_list(self, datum_kwarg_gen):
        return [self._fpath]

class NpyFrameWise(HandlerBase):
    specs = {'npy_FRAMEWISE'} | HandlerBase.specs

    def __init__(self, filename, mmap_mode=None):
        self._mmap_mode = mmap_mode
        if not os.path.exists(filename):
            raise IOError("the requested file {fpath} does not exist")
        self._fpath = filename
        self._data = np.load(self._fpath, self._mmap_mode)

    def __call__(self, frame_no):
        return self._data[frame_no]

    def get_file_list(self, datum_kwarg_gen):
        return [self._fpath]

class SingleTiffHandler(HandlerBase):
    specs = {'SINGLE_TIFF'} | HandlerBase.specs

    def __init__(self, filename):
        self._name = filename

    def __call__(self):
        return tifffile.imread(self._name)

class DATHandler(HandlerBase):
    ''' This handles vague text files through numpy's loadtxt module.
            Useful for the ubiquitous, vague and well loved ".DAT" file.
            Uses defaults from np.loadtxt.

        See Also
        --------
        np.loadtxt
    '''
    specs = {'DAT'} | HandlerBase.specs

    def __init__(self, fpath, **kwargs):
        self._path = fpath
        self.kwargs = kwargs

    def __call__(self):
        return np.loadtxt(self._path, **self.kwargs)
<commit_msg>ENH: lift CBF handler from LIX beamline config<commit_after>from __future__ import (absolute_import, division, print_function,
                        unicode_literals)
import six

import logging
import h5py
import numpy as np
import os.path
import tifffile

from .handlers_base import HandlerBase
from .readers.spe import PrincetonSPEFile

logger = logging.getLogger(__name__)

from pims import FramesSequence, Frame

# The ImageCube class is used for a per event representation of
# a dataset

class ImageStack(FramesSequence):
    "One of these represents the data from an event: (num_images x w x h)"
    def __init__(self, dataset, start, stop):
        # `start` and `stop` are the limits of this cube
        # i indexes within the cube
        self._start = start
        self._stop = stop
        self._dataset = dataset

        # work around inconsistent naming choices in databroker's Image object
        self.dtype = self.pixel_type
        self.shape = self.frame_shape

    def get_frame(self, i):
        return Frame(self._dataset[self._start + i], frame_no=i)

    def __len__(self):
        return self._stop - self._start

    @property
    def pixel_type(self):
        return self._dataset.dtype

    @property
    def frame_shape(self):
        return self._dataset.shape[1:]

class IntegrityError(Exception):
    pass

class AreaDetectorSPEHandler(HandlerBase):
    specs = {'AD_SPE'} | HandlerBase.specs

    def __init__(self, fpath, template, filename,
                 frame_per_point=1):
        self._path = os.path.join(fpath, '')
        self._fpp = frame_per_point
        self._template = template
        self._filename = filename
        self._f_cache = dict()

    def __call__(self, point_number):
        if point_number not in self._f_cache:
            fname = self._template % (self._path,
                                      self._filename,
                                      point_number)
            spe_obj = PrincetonSPEFile(fname)
            self._f_cache[point_number] = spe_obj

        spe = self._f_cache[point_number]
        data = spe.getData()

        if data.shape[0] != self._fpp:
            raise IntegrityError("expected {} frames, found {} frames".format(
                                 self._fpp, data.shape[0]))
        return data.squeeze()

    def get_file_list(self, datum_kwarg_gen):
        return [self._template % (self._path,
                                  self._filename,
                                  d['point_number'])
                for d in datum_kwarg_gen]

class AreaDetectorTiffHandler(HandlerBase):
    specs = {'AD_TIFF'} | HandlerBase.specs

    def __init__(self, fpath, template, filename, frame_per_point=1):
        self._path = os.path.join(fpath, '')
        self._fpp = frame_per_point
        self._template = template
        self._filename = filename

    def _fnames_for_point(self, point_number):
        start, stop = point_number * self._fpp, (point_number + 1) * self._fpp
        for j in range(start, stop):
            yield self._template % (self._path, self._filename, j)

    def __call__(self, point_number):
        ret = []
        for fn in self._fnames_for_point(point_number):
            with tifffile.TiffFile(fn) as tif:
                ret.append(tif.asarray())
        return np.array(ret).squeeze()

    def get_file_list(self, datum_kwargs):
        ret = []
        for d_kw in datum_kwargs:
            ret.extend(self._fnames_for_point(**d_kw))
        return ret

class DummyAreaDetectorHandler(HandlerBase):
    def __init__(self, fpath, frame_per_point=1, **kwargs):

        self._fpp = frame_per_point

    def __call__(self, **kwargs):
        out_stack = np.ones((self._fpp, 10, 10)) * np.nan
        # return stacked and squeezed results
        return out_stack.squeeze()

class HDF5DatasetSliceHandler(HandlerBase):
    """
    Handler for data stored in one Dataset of an HDF5 file.

    Parameters
    ----------
    filename : string
        path to HDF5 file
    key : string
        key of the single HDF5 Dataset used by this Handler
    frame_per_point : integer, optional
        number of frames to return as one datum, default 1
    """
    def __init__(self, filename, key, frame_per_point=1):
        self._fpp = frame_per_point
        self._filename = filename
        self._key = key
        self._file = None
        self._dataset = None
        self._data_objects = {}
        self.open()

    def get_file_list(self, datum_kwarg_gen):
        return [self._filename]

    def __call__(self, point_number):
        # Don't read out the dataset until it is requested for the first time.
        if not self._dataset:
            self._dataset = self._file[self._key]

        if point_number not in self._data_objects:
            start = point_number * self._fpp
            stop = (point_number + 1) * self._fpp
            self._data_objects[point_number] = ImageStack(self._dataset,
                                                          start, stop)
        return self._data_objects[point_number]

    def open(self):
        if self._file:
            return

        self._file = h5py.File(self._filename, 'r')

    def close(self):
        super(HDF5DatasetSliceHandler, self).close()
        self._file.close()
        self._file = None

class AreaDetectorHDF5Handler(HDF5DatasetSliceHandler):
    """
    Handler for the 'AD_HDF5' spec used by Area Detectors.

    In this spec, the key (i.e., HDF5 dataset path) is always
    '/entry/data/data'.

    Parameters
    ----------
    filename : string
        path to HDF5 file
    frame_per_point : integer, optional
        number of frames to return as one datum, default 1
    """
    specs = {'AD_HDF5'} | HDF5DatasetSliceHandler.specs

    def __init__(self, filename, frame_per_point=1):
        hardcoded_key = '/entry/data/data'
        super(AreaDetectorHDF5Handler, self).__init__(
            filename=filename, key=hardcoded_key,
            frame_per_point=frame_per_point)

class AreaDetectorHDF5SWMRHandler(AreaDetectorHDF5Handler):
    """
    Handler for the 'AD_HDF5_SWMR' spec used by Area Detectors.

    In this spec, the key (i.e., HDF5 dataset path) is always
    '/entry/data/data'.

    Parameters
    ----------
    filename : string
        path to HDF5 file
    frame_per_point : integer, optional
        number of frames to return as one datum, default 1
    """
    specs = {'AD_HDF5_SWMR'} | HDF5DatasetSliceHandler.specs

    def open(self):
        if self._file:
            return

        self._file = h5py.File(self._filename, 'r', swmr=True)

    def __call__(self, point_number):
        if self._dataset is not None:
            self._dataset.id.refresh()
        rtn = super(AreaDetectorHDF5SWMRHandler, self).__call__(
            point_number)

        return rtn

class AreaDetectorHDF5TimestampHandler(HandlerBase):
    """ Handler to retrieve timestamps from Areadetector HDF5 File

    In this spec, the timestamps of the images are read.

    Parameters
    ----------
    filename : string
        path to HDF5 file
    frame_per_point : integer, optional
        number of frames to return as one datum, default 1
    """
    specs = {'AD_HDF5_TS'} | HandlerBase.specs

    def __init__(self, filename, frame_per_point=1):
        self._fpp = frame_per_point
        self._filename = filename
        self._key = ['/entry/instrument/NDAttributes/NDArrayEpicsTSSec',
                     '/entry/instrument/NDAttributes/NDArrayEpicsTSnSec']
        self._file = None
        self._dataset1 = None
        self._dataset2 = None
        self.open()

    def __call__(self, point_number):
        # Don't read out the dataset until it is requested for the first time.
        if not self._dataset1:
            self._dataset1 = self._file[self._key[0]]
        if not self._dataset2:
            self._dataset2 = self._file[self._key[1]]
        start, stop = point_number * self._fpp, (point_number + 1) * self._fpp
        rtn = self._dataset1[start:stop].squeeze()
        rtn = rtn + (self._dataset2[start:stop].squeeze() * 1e-9)
        return rtn

    def open(self):
        if self._file:
            return
        self._file = h5py.File(self._filename, 'r')

    def close(self):
        super(AreaDetectorHDF5TimestampHandler, self).close()
        self._file.close()
        self._file = None

class AreaDetectorHDF5SWMRTimestampHandler(AreaDetectorHDF5TimestampHandler):
    """ Handler to retrieve timestamps from Areadetector HDF5 File

    In this spec, the timestamps of the images are read. Reading
    is done using SWMR option to allow read during processing

    Parameters
    ----------
    filename : string
        path to HDF5 file
    frame_per_point : integer, optional
        number of frames to return as one datum, default 1
    """
    specs = {'AD_HDF5_SWMR_TS'} | HandlerBase.specs

    def open(self):
        if self._file:
            return
        self._file = h5py.File(self._filename, 'r', swmr=True)

    def __call__(self, point_number):
        if (self._dataset1 is not None) and (self._dataset2 is not None):
            self._dataset.id.refresh()
        rtn = super(AreaDetectorHDF5SWMRTimestampHandler, self).__call__(
            point_number)
        return rtn

class _HdfMapsHandlerBase(HDF5DatasetSliceHandler):
    """
    Reader for XRF data stored in hdf5 files.

    The data set is assumed to be in a group called MAPS and stored
    as a 3D array ordered [energy, x, y].

    Parameters
    ----------
    filename : str
        Path to physical location of file
    dset_path : str
        The path to the dataset inside of 'MAPS'
    """
    def __init__(self, filename, dset_path):
        self._filename = filename
        self._dset_path = dset_path
        self._file = None
        self._dset = None
        self._swmr = False
        self.open()

    def open(self):
        """
        Open the file for reading.

        Provided as a stand alone function to allow re-opening of the handler
        """
        super(_HdfMapsHandlerBase, self).open()
        self._dset = self._file['/'.join(['MAPS', self._dset_path])]

    def __call__(self):

        if not self._file:
            raise RuntimeError("File is not open")

        if self._swmr:
            self._dataset.id.refresh()

class HDFMapsSpectrumHandler(_HdfMapsHandlerBase):
    """
    Handler which selects energy spectrum from
    a MAPS XRF data product.
    """
    specs = {'MAPS_SPECTRUM'} | _HdfMapsHandlerBase.specs

    def __call__(self, x, y):
        """
        Return the spectrum at the x, y position

        Parameters
        ----------
        x : int
            raster index in the x direction

        y : int
            raster index in the y direction

        Returns
        -------
        spectrum : ndarray
            The MCA channels
        """
        super(HDFMapsSpectrumHandler, self).__call__()
        return self._dset[:, x, y]

class HDFMapsEnergyHandler(_HdfMapsHandlerBase):
    """
    Handler which select fixed-energy slices from
    a MAPS XRF data file.
    """
    specs = {'MAPS_PLANE'} | _HdfMapsHandlerBase.specs

    def __call__(self, e_index):
        """
        Return the raster plane at a fixed energy

        Parameters
        ----------
        e_index : int
            The index of the engery

        Returns
        -------
        plane : ndarray
            The raster image at a fixed energy.
        """
        super(HDFMapsEnergyHandler, self).__call__()
        return self._dset[e_index, :, :]

class NpyHandler(HandlerBase):
    """
    Class to deal with reading npy files

    Parameters
    ----------
    fpath : str
        Path to file

    mmap_mode : {'r', 'r+', c}, optional
        memmap mode to use to open file
    """
    specs = {'npy'} | HandlerBase.specs

    def __init__(self, filename, mmap_mode=None):
        self._mmap_mode = mmap_mode
        if not os.path.exists(filename):
            raise IOError("the requested file {fpath} does not exist")
        self._fpath = filename

    def __call__(self):
        return np.load(self._fpath, self._mmap_mode)

    def get_file_list(self, datum_kwarg_gen):
        return [self._fpath]

class NpyFrameWise(HandlerBase):
    specs = {'npy_FRAMEWISE'} | HandlerBase.specs

    def __init__(self, filename, mmap_mode=None):
        self._mmap_mode = mmap_mode
        if not os.path.exists(filename):
            raise IOError("the requested file {fpath} does not exist")
        self._fpath = filename
        self._data = np.load(self._fpath, self._mmap_mode)

    def __call__(self, frame_no):
        return self._data[frame_no]

    def get_file_list(self, datum_kwarg_gen):
        return [self._fpath]

class SingleTiffHandler(HandlerBase):
    specs = {'SINGLE_TIFF'} | HandlerBase.specs

    def __init__(self, filename):
        self._name = filename

    def __call__(self):
        return tifffile.imread(self._name)

class DATHandler(HandlerBase):
    ''' This handles vague text files through numpy's loadtxt module.
            Useful for the ubiquitous, vague and well loved ".DAT" file.
            Uses defaults from np.loadtxt.

        See Also
        --------
        np.loadtxt
    '''
    specs = {'DAT'} | HandlerBase.specs

    def __init__(self, fpath, **kwargs):
        self._path = fpath
        self.kwargs = kwargs

    def __call__(self):
        return np.loadtxt(self._path, **self.kwargs)

class PilatusCBFHandler(HandlerBase):
    specs = {'AD_CBF'} | HandlerBase.specs

    def __init__(self, rpath, template, filename, frame_per_point=1,
                 initial_number=1):
        self._path = os.path.join(rpath, '')
        self._fpp = frame_per_point
        self._template = template
        self._filename = filename
        self._initial_number = initial_number

    def __call__(self, point_number):
        import fabio
        start, stop = (self._initial_number + point_number *
                       self._fpp, (point_number + 2) * self._fpp)
        ret = []
        # commented out by LY to test scan speed imperovement, 2017-01-24
        for j in range(start, stop):
            fn = self._template % (self._path, self._filename, j)
            img = fabio.open(fn)
            ret.append(img.data)
        return np.array(ret).squeeze()

    def get_file_list(self, datum_kwargs_gen):
        file_list = []
        for dk in datum_kwargs_gen:
            point_number = dk['point_number']
            start, stop = (self._initial_number + point_number *
                           self._fpp, (point_number + 2) * self._fpp)
            for j in range(start, stop):
                fn = self._template % (self._path, self._filename, j)
                file_list.append(fn)
        return file_list
<|endoftext|>


<commit_before>"""
A special gateway client for Voice.
"""
import enum
import json
import logging
import socket
import threading
import time
import typing
import zlib

import curio
from cuiows.client import WSClient
from cuiows.exc import WebsocketClosedError
from curio.thread import AWAIT, async_thread

from curious.core.gateway import Gateway, ReconnectWebsocket

logger = logging.getLogger("curious.voice")

class VGatewayOp(enum.IntEnum):
    IDENTIFY = 0
    SELECT_PROTOCOL = 1
    READY = 2
    HEARTBEAT = 3
    SESSION_DESCRIPTION = 4
    SPEAKING = 5
    HEARTBEAT_ACK = 6
    HELLO = 8

@async_thread(daemon=True)
def _heartbeat_loop(gw: 'VoiceGateway', heartbeat_interval: float):
    """
    Heartbeat looper that loops and sends heartbeats to the gateway.

    :param gw: The gateway to handle.
    """
    # async threads!
    logger.debug("Sending initial heartbeat.")
    AWAIT(gw.send_heartbeat())
    while True:
        # this is similar to the normal threaded event waiter
        # it will time out after heartbeat_interval seconds.
        try:
            AWAIT(curio.timeout_after(heartbeat_interval, gw._stop_heartbeating.wait()))
        except curio.TaskTimeout:
            pass
        else:
            break

        try:
            AWAIT(gw.send_heartbeat())
        except ReconnectWebsocket:
            break

class VoiceGateway(object):
    """
    A special type of websocket gateway that operates on the Voice connection.
    """

    GATEWAY_VERSION = 3

    def __init__(self, session_id: str, token: str, endpoint: str, user_id: str, guild_id: str):
        #: The current cuiows websocket object.
        self.websocket = None  # type: WSClient

        self._open = False

        self.session_id = session_id
        self.token = token
        self.user_id = user_id
        self.guild_id = guild_id
        self.sequence = 0
        self.ssrc_mapping = {}

        #: The main gateway object.
        self.main_gateway = None  # type: Gateway

        #: The current threading event used to signal if we need to stop heartbeating.
        self._stop_heartbeating = curio.Event()

        #: Voice server stuff
        self.endpoint = endpoint
        self.port = None  # type: int
        self.ssrc = None  # type: int
        self.secret_key = b""

        self._ready = curio.Event()
        self._got_secret_key = curio.Event()

    async def connect(self, url: str):
        """
        Connects the client to the websocket.

        This does not do ANY handshake steps.
        :param url: The URL to connect to.
        """
        logger.info("Opening connection to {}".format(url))
        self.websocket = await WSClient.connect(url)
        logger.info("Connected to gateway!")
        self._open = True
        return self.websocket

    async def _send(self, data):
        """
        Sends some data.

        This will actually send the data down the websocket, unlike `send` which only pretends to.
        """
        try:
            await self.websocket.send(data)
        except WebsocketClosedError:
            await self._close()

    def get_heartbeat(self):
        """
        :return: A heartbeat packet.
        """
        return {
            "op": int(VGatewayOp.HEARTBEAT),
            # ms time for v2/v3
            "d": int(round(time.time() * 1000))
        }

    async def send_heartbeat(self):
        """
        Sends a heartbeat.
        """
        hb = self.get_heartbeat()
        return await self._send_json(hb)

    def _send_json(self, payload: dict):
        """
        Sends a JSON payload down the websocket.

        This is a private method - use :meth:`send` instead.

        :param payload: The payload to send.
        """
        data = json.dumps(payload)
        return self.send(data)

    async def send(self, data: typing.Any):
        """
        Enqueues some data to be sent down the websocket.
        """
        await self.websocket.send(data)

    # Sending methods
    async def send_identify(self):
        """
        Sends an IDENTIFY payload.
        """
        payload = {
            "op": VGatewayOp.IDENTIFY,
            "d": {
                "server_id": str(self.guild_id),
                "user_id": str(self.user_id),
                "session_id": str(self.session_id),
                "token": str(self.token)
            }
        }
        await self._send_json(payload)

    async def send_select_protocol(self, ip: str, port: int):
        """
        Sends a SELECT PROTOCOL packet.
        """
        payload = {
            "op": VGatewayOp.SELECT_PROTOCOL,
            "d": {
                "protocol": "udp",
                "data": {
                    "address": ip,
                    "port": port,
                    "mode": "xsalsa20_poly1305"
                }
            }
        }

        await self._send_json(payload)

    async def send_speaking(self, speaking: bool = True):
        """
        Sends a SPEAKING packet.
        """
        payload = {
            "op": VGatewayOp.SPEAKING,
            "d": {
                "speaking": speaking,
                "delay": 0
            }
        }

        await self._send_json(payload)

    async def _start_heartbeating(self, heartbeat_interval: float) -> threading.Thread:
        """
        Starts the heartbeat thread.
        :param heartbeat_interval: The number of seconds between each heartbeat.
        """
        if not self._stop_heartbeating.is_set():
            # cancel some poor thread elsewhere
            await self._stop_heartbeating.set()
            del self._stop_heartbeating

        self._stop_heartbeating = curio.Event()

        # dont reference the task - it'll die by itself
        task = await curio.spawn(_heartbeat_loop(self, heartbeat_interval), daemon=True)

        return task

    async def _close(self):
        if not self.websocket.closed:
            await self.websocket.close_now(code=1000, reason="Client disconnected")

        self._open = False
        await self._stop_heartbeating.set()

    @classmethod
    async def from_gateway(cls, gw: Gateway, guild_id: int, channel_id: int) -> 'VoiceGateway':
        """
        Opens a new voice gateway connection from an existing Gateway connection.

        :param gw: The existing :class:`.Gateway` to create from.
        :param guild_id: The guild ID we're opening to.
        :param channel_id: The channel ID we're connecting to.
        """
        # Send our VOICE_STATE_UPDATE to tell Discord to start connecting us to voice.
        await gw.send_voice_state_update(guild_id, channel_id)
        state = await gw.state.wait_for_voice_data(guild_id)
        user_id = gw.state._user.id

        # Look up the gateway's IP address.
        endpoint = state["endpoint"]
        endpoint, port = endpoint.split(":")
        endpoint = await curio.abide(socket.gethostbyname, endpoint)

        obb = cls(token=state["token"], session_id=state["session_id"],
                  endpoint=endpoint,
                  user_id=user_id, guild_id=guild_id)
        # Open our connection to the voice websocket.
        if port == "80":
            await obb.connect(f"ws://{endpoint}/?v={cls.GATEWAY_VERSION}")
        elif port == "443":
            await obb.connect(f"wss://{endpoint}/?v={cls.GATEWAY_VERSION}")
        # Send our IDENTIFY.
        logger.info("Sending IDENTIFY...")
        await obb.send_identify()

        return obb

    async def next_event(self) -> None:
        """
        Gets the next event, in decoded form.
        """
        if not self._open:
            raise WebsocketClosedError(1006, reason="Connection lost")

        try:
            event = await self.websocket.poll()
        except WebsocketClosedError:
            # Close ourselves.
            await self._close()
            raise

        if isinstance(event, (bytes, bytearray)):
            # decompress the message
            event = zlib.decompress(event, 15, 10490000)
            event = event.decode("utf-8")

        if event is None:
            return

        event = json.loads(event)

        op = event.get("op")
        data = event.get("d")
        seq = event.get("seq")

        # Heartbeat logic, same as normal websocket
        if seq is not None:
            self.sequence = int(seq)

        # Switch based on operator.
        if op == VGatewayOp.HELLO:
            # Ignore these, they're useless for now.
            return

        elif op == VGatewayOp.READY:
            # Start heartbeating.
            heartbeat_interval = data.get("heartbeat_interval", 45000) / 1000.0
            logger.debug("Heartbeating every {} seconds.".format(heartbeat_interval))
            await self._start_heartbeating(heartbeat_interval)
            # Set our `ssrc`, `port` and `modes`.
            self.ssrc = data.get("ssrc")
            self.port = data.get("port")

            await self._ready.set()

        elif op == VGatewayOp.SESSION_DESCRIPTION:
            # Extract the secret key.
            self.secret_key = data["secret_key"]
            await self.send_speaking()
            await self._got_secret_key.set()

        elif op == VGatewayOp.SPEAKING:
            # build a cache of user_id -> ssrc
            user_id = int(data.get("user_id"))
            ssrc = data.get("ssrc")
            self.ssrc_mapping[ssrc] = user_id

        elif op == VGatewayOp.HEARTBEAT:
            # silence
            pass

        elif op == VGatewayOp.HEARTBEAT_ACK:
            # suppress
            pass

        else:
            logger.warning("Unhandled event: {}".format(op))
<commit_msg>Bring voice code to asyncwebsockets.<commit_after>"""
A special gateway client for Voice.
"""
import enum
import json
import logging
import socket
import threading
import time
import typing
import zlib

import curio
from asyncwebsockets import Websocket, WebsocketBytesMessage, WebsocketClosed, connect_websocket
from curio.thread import AWAIT, async_thread

from curious.core.gateway import Gateway, ReconnectWebsocket

logger = logging.getLogger("curious.voice")

class VGatewayOp(enum.IntEnum):
    IDENTIFY = 0
    SELECT_PROTOCOL = 1
    READY = 2
    HEARTBEAT = 3
    SESSION_DESCRIPTION = 4
    SPEAKING = 5
    HEARTBEAT_ACK = 6
    HELLO = 8

@async_thread(daemon=True)
def _heartbeat_loop(gw: 'VoiceGateway', heartbeat_interval: float):
    """
    Heartbeat looper that loops and sends heartbeats to the gateway.

    :param gw: The gateway to handle.
    """
    # async threads!
    logger.debug("Sending initial heartbeat.")
    AWAIT(gw.send_heartbeat())
    while True:
        # this is similar to the normal threaded event waiter
        # it will time out after heartbeat_interval seconds.
        try:
            AWAIT(curio.timeout_after(heartbeat_interval, gw._stop_heartbeating.wait()))
        except curio.TaskTimeout:
            pass
        else:
            break

        try:
            AWAIT(gw.send_heartbeat())
        except ReconnectWebsocket:
            break

class VoiceGateway(object):
    """
    A special type of websocket gateway that operates on the Voice connection.
    """

    GATEWAY_VERSION = 3

    def __init__(self, session_id: str, token: str, endpoint: str, user_id: str, guild_id: str):
        #: The current websocket object.
        self.websocket = None  # type: Websocket

        self._open = False

        self.session_id = session_id
        self.token = token
        self.user_id = user_id
        self.guild_id = guild_id
        self.sequence = 0
        self.ssrc_mapping = {}

        #: The main gateway object.
        self.main_gateway = None  # type: Gateway

        #: The current threading event used to signal if we need to stop heartbeating.
        self._stop_heartbeating = curio.Event()

        #: Voice server stuff
        self.endpoint = endpoint
        self.port = None  # type: int
        self.ssrc = None  # type: int
        self.secret_key = b""

        self._ready = curio.Event()
        self._got_secret_key = curio.Event()

    async def connect(self, url: str):
        """
        Connects the client to the websocket.

        This does not do ANY handshake steps.
        :param url: The URL to connect to.
        """
        logger.info("Opening connection to {}".format(url))
        self.websocket = await connect_websocket(url)
        logger.info("Connected to gateway!")
        self._open = True
        return self.websocket

    async def _send(self, data):
        """
        Sends some data.

        This will actually send the data down the websocket, unlike `send` which only pretends to.
        """
        try:
            await self.websocket.send_message(data)
        except WebsocketClosed:
            await self._close()

    def get_heartbeat(self):
        """
        :return: A heartbeat packet.
        """
        return {
            "op": int(VGatewayOp.HEARTBEAT),
            # ms time for v2/v3
            "d": int(round(time.time() * 1000))
        }

    async def send_heartbeat(self):
        """
        Sends a heartbeat.
        """
        hb = self.get_heartbeat()
        return await self._send_json(hb)

    def _send_json(self, payload: dict):
        """
        Sends a JSON payload down the websocket.

        This is a private method - use :meth:`send` instead.

        :param payload: The payload to send.
        """
        data = json.dumps(payload)
        return self.send(data)

    async def send(self, data: typing.Any):
        """
        Enqueues some data to be sent down the websocket.
        """
        await self.websocket.send_message(data)

    # Sending methods
    async def send_identify(self):
        """
        Sends an IDENTIFY payload.
        """
        payload = {
            "op": VGatewayOp.IDENTIFY,
            "d": {
                "server_id": str(self.guild_id),
                "user_id": str(self.user_id),
                "session_id": str(self.session_id),
                "token": str(self.token)
            }
        }
        await self._send_json(payload)

    async def send_select_protocol(self, ip: str, port: int):
        """
        Sends a SELECT PROTOCOL packet.
        """
        payload = {
            "op": VGatewayOp.SELECT_PROTOCOL,
            "d": {
                "protocol": "udp",
                "data": {
                    "address": ip,
                    "port": port,
                    "mode": "xsalsa20_poly1305"
                }
            }
        }

        await self._send_json(payload)

    async def send_speaking(self, speaking: bool = True):
        """
        Sends a SPEAKING packet.
        """
        payload = {
            "op": VGatewayOp.SPEAKING,
            "d": {
                "speaking": speaking,
                "delay": 0
            }
        }

        await self._send_json(payload)

    async def _start_heartbeating(self, heartbeat_interval: float) -> threading.Thread:
        """
        Starts the heartbeat thread.
        :param heartbeat_interval: The number of seconds between each heartbeat.
        """
        if not self._stop_heartbeating.is_set():
            # cancel some poor thread elsewhere
            await self._stop_heartbeating.set()
            del self._stop_heartbeating

        self._stop_heartbeating = curio.Event()

        # dont reference the task - it'll die by itself
        task = await curio.spawn(_heartbeat_loop(self, heartbeat_interval), daemon=True)

        return task

    async def _close(self):
        if not self.websocket.closed:
            await self.websocket.close(code=1000, reason="Client disconnected")

        self._open = False
        await self._stop_heartbeating.set()

    @classmethod
    async def from_gateway(cls, gw: Gateway, guild_id: int, channel_id: int) -> 'VoiceGateway':
        """
        Opens a new voice gateway connection from an existing Gateway connection.

        :param gw: The existing :class:`.Gateway` to create from.
        :param guild_id: The guild ID we're opening to.
        :param channel_id: The channel ID we're connecting to.
        """
        # Send our VOICE_STATE_UPDATE to tell Discord to start connecting us to voice.
        await gw.send_voice_state_update(guild_id, channel_id)
        state = await gw.state.wait_for_voice_data(guild_id)
        user_id = gw.state._user.id

        # Look up the gateway's IP address.
        endpoint = state["endpoint"]
        endpoint, port = endpoint.split(":")
        endpoint = await curio.abide(socket.gethostbyname, endpoint)

        obb = cls(token=state["token"], session_id=state["session_id"],
                  endpoint=endpoint,
                  user_id=user_id, guild_id=guild_id)
        # Open our connection to the voice websocket.
        if port == "80":
            await obb.connect(f"ws://{endpoint}/?v={cls.GATEWAY_VERSION}")
        elif port == "443":
            await obb.connect(f"wss://{endpoint}/?v={cls.GATEWAY_VERSION}")
        # Send our IDENTIFY.
        logger.info("Sending IDENTIFY...")
        await obb.send_identify()

        return obb

    async def next_event(self) -> None:
        """
        Gets the next event, in decoded form.
        """
        if not self._open:
            raise WebsocketClosed(1006, reason="Connection lost")

        try:
            event = await self.websocket.next_message()
        except WebsocketClosed:
            # Close ourselves.
            await self._close()
            raise

        if isinstance(event, WebsocketBytesMessage):
            # decompress the message
            data = zlib.decompress(event.data, 15, 10490000)
            data = data.decode("utf-8")
        else:
            data = event.data

        if data is None:
            return

        event = json.loads(data)

        op = event.get("op")
        data = event.get("d")
        seq = event.get("seq")

        # Heartbeat logic, same as normal websocket
        if seq is not None:
            self.sequence = int(seq)

        # Switch based on operator.
        if op == VGatewayOp.HELLO:
            # Ignore these, they're useless for now.
            return

        elif op == VGatewayOp.READY:
            # Start heartbeating.
            heartbeat_interval = data.get("heartbeat_interval", 45000) / 1000.0
            logger.debug("Heartbeating every {} seconds.".format(heartbeat_interval))
            await self._start_heartbeating(heartbeat_interval)
            # Set our `ssrc`, `port` and `modes`.
            self.ssrc = data.get("ssrc")
            self.port = data.get("port")

            await self._ready.set()

        elif op == VGatewayOp.SESSION_DESCRIPTION:
            # Extract the secret key.
            self.secret_key = data["secret_key"]
            await self.send_speaking()
            await self._got_secret_key.set()

        elif op == VGatewayOp.SPEAKING:
            # build a cache of user_id -> ssrc
            user_id = int(data.get("user_id"))
            ssrc = data.get("ssrc")
            self.ssrc_mapping[ssrc] = user_id

        elif op == VGatewayOp.HEARTBEAT:
            # silence
            pass

        elif op == VGatewayOp.HEARTBEAT_ACK:
            # suppress
            pass

        else:
            logger.warning("Unhandled event: {}".format(op))
<|endoftext|>


<commit_before># Copyright (c) 2012 Web Notes Technologies Pvt Ltd (http://erpnext.com)
#
# MIT License (MIT)
#
# Permission is hereby granted, free of charge, to any person obtaining a
# copy of this software and associated documentation files (the "Software"),
# to deal in the Software without restriction, including without limitation
# the rights to use, copy, modify, merge, publish, distribute, sublicense,
# and/or sell copies of the Software, and to permit persons to whom the
# Software is furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED,
# INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A
# PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
# HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF
# CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE
# OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
#

"""
	This module contains classes for managing incoming emails
"""

class IncomingMail:
	"""
		Single incoming email object. Extracts, text / html and attachments from the email
	"""
	def __init__(self, content):
		"""
			Parse the incoming mail content
		"""
		import email

		self.mail = email.message_from_string(content)
		self.text_content = ''
		self.html_content = ''
		self.attachments = []
		self.parse()

	def get_text_content(self):
		"""
			Returns the text parts of the email. If None, then HTML parts
		"""
		return self.text_content or self.html_content

	def get_charset(self, part):
		"""
			Guesses character set
		"""
		charset = part.get_content_charset()
		if not charset:
			import chardet
			charset = chardet.detect(str(part))['encoding']

		return charset

	def get_payload(self, part, charset):
		"""
			get utf-8 encoded part content
		"""
		try:
			return unicode(part.get_payload(decode=True),str(charset),"ignore").encode('utf8','replace')
		except LookupError, e:
			return part.get_payload()

	def get_attachment(self, part, charset):
		"""
			Extracts an attachment
		"""
		self.attachments.append({
			'content-type': part.get_content_type(),
			'filename': part.get_filename(),
			'content': part.get_payload(decode=True)
		})

	def parse(self):
		"""
			Extracts text, html and attachments from the mail
		"""
		for part in self.mail.walk():
			self.process_part(part)

	def get_thread_id(self):
		"""
			Extracts thread id of the message between first []
			from the subject
		"""
		import re
		subject = self.mail.get('Subject', '')

		return re.findall('(?<=\[)[\w/]+', subject)

	def process_part(self, part):
		"""
			Process a single part of an email
		"""
		charset = self.get_charset(part)
		content_type = part.get_content_type()

		if content_type == 'text/plain':
			self.text_content += self.get_payload(part, charset)

		if content_type == 'text/html':
			self.html_content += self.get_payload(part, charset)

		if part.get_filename():
			self.get_attachment(part, charset)

class POP3Mailbox:
	"""
		A simple pop3 mailbox, abstracts connection and mail extraction
		To use, subclass it and override method process_message(from, subject, text, thread_id)
	"""

	def __init__(self, settings_doc):
		"""
			settings_doc must contain
			use_ssl, host, username, password
			(by name or object)
		"""
		if isinstance(settings_doc, basestring):
			from webnotes.model.doc import Document
			self.settings = Document(settings_doc, settings_doc)
		else:
			self.settings = settings_doc

	def connect(self):
		"""
			Connects to the mailbox
		"""
		import poplib

		if self.settings.use_ssl:
			self.pop = poplib.POP3_SSL(self.settings.host)
		else:
			self.pop = poplib.POP3(self.settings.host)
		self.pop.user(self.settings.username)
		self.pop.pass_(self.settings.password)

	def get_messages(self):
		"""
			Loads messages from the mailbox and calls
			process_message for each message
		"""

		if not self.check_mails():
			return # nothing to do

		self.connect()
		num = len(self.pop.list()[1])
		for m in range(num):
			msg = self.pop.retr(m+1)
			try:
				self.process_message(IncomingMail('\n'.join(msg[1])))
			except:
				pass
			self.pop.dele(m+1)
		self.pop.quit()

	def check_mails(self):
		"""
			To be overridden
			If mailbox is to be scanned, returns true
		"""
		return True

	def process_message(self, mail):
		"""
			To be overriden
		"""
		pass
<commit_msg>restrict number of support ticket emails<commit_after># Copyright (c) 2012 Web Notes Technologies Pvt Ltd (http://erpnext.com)
#
# MIT License (MIT)
#
# Permission is hereby granted, free of charge, to any person obtaining a
# copy of this software and associated documentation files (the "Software"),
# to deal in the Software without restriction, including without limitation
# the rights to use, copy, modify, merge, publish, distribute, sublicense,
# and/or sell copies of the Software, and to permit persons to whom the
# Software is furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED,
# INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A
# PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
# HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF
# CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE
# OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
#

"""
	This module contains classes for managing incoming emails
"""

class IncomingMail:
	"""
		Single incoming email object. Extracts, text / html and attachments from the email
	"""
	def __init__(self, content):
		"""
			Parse the incoming mail content
		"""
		import email

		self.mail = email.message_from_string(content)
		self.text_content = ''
		self.html_content = ''
		self.attachments = []
		self.parse()

	def get_text_content(self):
		"""
			Returns the text parts of the email. If None, then HTML parts
		"""
		return self.text_content or self.html_content

	def get_charset(self, part):
		"""
			Guesses character set
		"""
		charset = part.get_content_charset()
		if not charset:
			import chardet
			charset = chardet.detect(str(part))['encoding']

		return charset

	def get_payload(self, part, charset):
		"""
			get utf-8 encoded part content
		"""
		try:
			return unicode(part.get_payload(decode=True),str(charset),"ignore").encode('utf8','replace')
		except LookupError, e:
			return part.get_payload()

	def get_attachment(self, part, charset):
		"""
			Extracts an attachment
		"""
		self.attachments.append({
			'content-type': part.get_content_type(),
			'filename': part.get_filename(),
			'content': part.get_payload(decode=True)
		})

	def parse(self):
		"""
			Extracts text, html and attachments from the mail
		"""
		for part in self.mail.walk():
			self.process_part(part)

	def get_thread_id(self):
		"""
			Extracts thread id of the message between first []
			from the subject
		"""
		import re
		subject = self.mail.get('Subject', '')

		return re.findall('(?<=\[)[\w/]+', subject)

	def process_part(self, part):
		"""
			Process a single part of an email
		"""
		charset = self.get_charset(part)
		content_type = part.get_content_type()

		if content_type == 'text/plain':
			self.text_content += self.get_payload(part, charset)

		if content_type == 'text/html':
			self.html_content += self.get_payload(part, charset)

		if part.get_filename():
			self.get_attachment(part, charset)

class POP3Mailbox:
	"""
		A simple pop3 mailbox, abstracts connection and mail extraction
		To use, subclass it and override method process_message(from, subject, text, thread_id)
	"""

	def __init__(self, settings_doc):
		"""
			settings_doc must contain
			use_ssl, host, username, password
			(by name or object)
		"""
		if isinstance(settings_doc, basestring):
			from webnotes.model.doc import Document
			self.settings = Document(settings_doc, settings_doc)
		else:
			self.settings = settings_doc

	def connect(self):
		"""
			Connects to the mailbox
		"""
		import poplib

		if self.settings.use_ssl:
			self.pop = poplib.POP3_SSL(self.settings.host)
		else:
			self.pop = poplib.POP3(self.settings.host)
		self.pop.user(self.settings.username)
		self.pop.pass_(self.settings.password)

	def get_messages(self):
		"""
			Loads messages from the mailbox and calls
			process_message for each message
		"""

		if not self.check_mails():
			return # nothing to do

		self.connect()
		num = len(self.pop.list()[1])
		if num > 25: num = 25
		for m in range(num):
			msg = self.pop.retr(m+1)
			try:
				self.process_message(IncomingMail('\n'.join(msg[1])))
			except:
				pass
			self.pop.dele(m+1)
		self.pop.quit()

	def check_mails(self):
		"""
			To be overridden
			If mailbox is to be scanned, returns true
		"""
		return True

	def process_message(self, mail):
		"""
			To be overriden
		"""
		pass
<|endoftext|>


<commit_before>"""
Tests for behaviour related to type annotations.
"""

from sys import version_info

from pyflakes import messages as m
from pyflakes.test.harness import TestCase, skipIf

class TestTypeAnnotations(TestCase):

    def test_typingOverload(self):
        """Allow intentional redefinitions via @typing.overload"""
        self.flakes("""
        import typing
        from typing import overload

        @overload
        def f(s):  # type: (None) -> None
            pass

        @overload
        def f(s):  # type: (int) -> int
            pass

        def f(s):
            return s

        @typing.overload
        def g(s):  # type: (None) -> None
            pass

        @typing.overload
        def g(s):  # type: (int) -> int
            pass

        def g(s):
            return s
        """)

    def test_overload_with_multiple_decorators(self):
        self.flakes("""
            from typing import overload
            dec = lambda f: f

            @dec
            @overload
            def f(x):  # type: (int) -> int
                pass

            @dec
            @overload
            def f(x):  # type: (str) -> str
                pass

            @dec
            def f(x): return x
       """)

    def test_overload_in_class(self):
        self.flakes("""
        from typing import overload

        class C:
            @overload
            def f(self, x):  # type: (int) -> int
                pass

            @overload
            def f(self, x):  # type: (str) -> str
                pass

            def f(self, x): return x
        """)

    def test_not_a_typing_overload(self):
        """regression test for @typing.overload detection bug in 2.1.0"""
        self.flakes("""
            x = lambda f: f

            @x
            def t():
                pass

            y = lambda f: f

            @x
            @y
            def t():
                pass

            @x
            @y
            def t():
                pass
        """, m.RedefinedWhileUnused, m.RedefinedWhileUnused)

    @skipIf(version_info < (3, 6), 'new in Python 3.6')
    def test_variable_annotations(self):
        self.flakes('''
        name: str
        age: int
        ''')
        self.flakes('''
        name: str = 'Bob'
        age: int = 18
        ''')
        self.flakes('''
        class C:
            name: str
            age: int
        ''')
        self.flakes('''
        class C:
            name: str = 'Bob'
            age: int = 18
        ''')
        self.flakes('''
        def f():
            name: str
            age: int
        ''')
        self.flakes('''
        def f():
            name: str = 'Bob'
            age: int = 18
            foo: not_a_real_type = None
        ''', m.UnusedVariable, m.UnusedVariable, m.UnusedVariable, m.UndefinedName)
        self.flakes('''
        def f():
            name: str
            print(name)
        ''', m.UndefinedName)
        self.flakes('''
        from typing import Any
        def f():
            a: Any
        ''')
        self.flakes('''
        foo: not_a_real_type
        ''', m.UndefinedName)
        self.flakes('''
        foo: not_a_real_type = None
        ''', m.UndefinedName)
        self.flakes('''
        class C:
            foo: not_a_real_type
        ''', m.UndefinedName)
        self.flakes('''
        class C:
            foo: not_a_real_type = None
        ''', m.UndefinedName)
        self.flakes('''
        def f():
            class C:
                foo: not_a_real_type
        ''', m.UndefinedName)
        self.flakes('''
        def f():
            class C:
                foo: not_a_real_type = None
        ''', m.UndefinedName)
        self.flakes('''
        from foo import Bar
        bar: Bar
        ''')
        self.flakes('''
        from foo import Bar
        bar: 'Bar'
        ''')
        self.flakes('''
        import foo
        bar: foo.Bar
        ''')
        self.flakes('''
        import foo
        bar: 'foo.Bar'
        ''')
        self.flakes('''
        from foo import Bar
        def f(bar: Bar): pass
        ''')
        self.flakes('''
        from foo import Bar
        def f(bar: 'Bar'): pass
        ''')
        self.flakes('''
        from foo import Bar
        def f(bar) -> Bar: return bar
        ''')
        self.flakes('''
        from foo import Bar
        def f(bar) -> 'Bar': return bar
        ''')
        self.flakes('''
        bar: 'Bar'
        ''', m.UndefinedName)
        self.flakes('''
        bar: 'foo.Bar'
        ''', m.UndefinedName)
        self.flakes('''
        from foo import Bar
        bar: str
        ''', m.UnusedImport)
        self.flakes('''
        from foo import Bar
        def f(bar: str): pass
        ''', m.UnusedImport)
        self.flakes('''
        def f(a: A) -> A: pass
        class A: pass
        ''', m.UndefinedName, m.UndefinedName)
        self.flakes('''
        def f(a: 'A') -> 'A': return a
        class A: pass
        ''')
        self.flakes('''
        a: A
        class A: pass
        ''', m.UndefinedName)
        self.flakes('''
        a: 'A'
        class A: pass
        ''')
        self.flakes('''
        a: 'A B'
        ''', m.ForwardAnnotationSyntaxError)
        self.flakes('''
        a: 'A; B'
        ''', m.ForwardAnnotationSyntaxError)
        self.flakes('''
        a: '1 + 2'
        ''')
        self.flakes('''
        a: 'a: "A"'
        ''', m.ForwardAnnotationSyntaxError)

    @skipIf(version_info < (3, 5), 'new in Python 3.5')
    def test_annotated_async_def(self):
        self.flakes('''
        class c: pass
        async def func(c: c) -> None: pass
        ''')

    @skipIf(version_info < (3, 7), 'new in Python 3.7')
    def test_postponed_annotations(self):
        self.flakes('''
        from __future__ import annotations
        def f(a: A) -> A: pass
        class A:
            b: B
        class B: pass
        ''')

        self.flakes('''
        from __future__ import annotations
        def f(a: A) -> A: pass
        class A:
            b: Undefined
        class B: pass
        ''', m.UndefinedName)

    def test_typeCommentsMarkImportsAsUsed(self):
        self.flakes("""
        from mod import A, B, C, D, E, F, G

        def f(
            a,  # type: A
        ):
            # type: (...) -> B
            for b in a:  # type: C
                with b as c:  # type: D
                    d = c.x  # type: E
                    return d

        def g(x):  # type: (F) -> G
            return x.y
        """)

    def test_typeCommentsFullSignature(self):
        self.flakes("""
        from mod import A, B, C, D
        def f(a, b):
            # type: (A, B[C]) -> D
            return a + b
        """)

    def test_typeCommentsStarArgs(self):
        self.flakes("""
        from mod import A, B, C, D
        def f(a, *b, **c):
            # type: (A, *B, **C) -> D
            return a + b
        """)

    def test_typeCommentsFullSignatureWithDocstring(self):
        self.flakes('''
        from mod import A, B, C, D
        def f(a, b):
            # type: (A, B[C]) -> D
            """do the thing!"""
            return a + b
        ''')

    def test_typeCommentsAdditionalComemnt(self):
        self.flakes("""
        from mod import F

        x = 1 # type: F  # noqa
        """)

    def test_typeCommentsNoWhitespaceAnnotation(self):
        self.flakes("""
        from mod import F

        x = 1  #type:F
        """)

    def test_typeCommentsInvalidDoesNotMarkAsUsed(self):
        self.flakes("""
        from mod import F

        # type: F
        """, m.UnusedImport)

    def test_typeCommentsSyntaxError(self):
        self.flakes("""
        def f(x):  # type: (F[) -> None
            pass
        """, m.CommentAnnotationSyntaxError)

    def test_typeCommentsSyntaxErrorCorrectLine(self):
        checker = self.flakes("""\
        x = 1
        # type: definitely not a PEP 484 comment
        """, m.CommentAnnotationSyntaxError)
        self.assertEqual(checker.messages[0].lineno, 2)

    def test_typeCommentsAssignedToPreviousNode(self):
        # This test demonstrates an issue in the implementation which
        # associates the type comment with a node above it, however the type
        # comment isn't valid according to mypy.  If an improved approach
        # which can detect these "invalid" type comments is implemented, this
        # test should be removed / improved to assert that new check.
        self.flakes("""
        from mod import F
        x = 1
        # type: F
        """)
<commit_msg>Add another test for #422 (#433)<commit_after>"""
Tests for behaviour related to type annotations.
"""

from sys import version_info

from pyflakes import messages as m
from pyflakes.test.harness import TestCase, skipIf

class TestTypeAnnotations(TestCase):

    def test_typingOverload(self):
        """Allow intentional redefinitions via @typing.overload"""
        self.flakes("""
        import typing
        from typing import overload

        @overload
        def f(s):  # type: (None) -> None
            pass

        @overload
        def f(s):  # type: (int) -> int
            pass

        def f(s):
            return s

        @typing.overload
        def g(s):  # type: (None) -> None
            pass

        @typing.overload
        def g(s):  # type: (int) -> int
            pass

        def g(s):
            return s
        """)

    def test_overload_with_multiple_decorators(self):
        self.flakes("""
            from typing import overload
            dec = lambda f: f

            @dec
            @overload
            def f(x):  # type: (int) -> int
                pass

            @dec
            @overload
            def f(x):  # type: (str) -> str
                pass

            @dec
            def f(x): return x
       """)

    def test_overload_in_class(self):
        self.flakes("""
        from typing import overload

        class C:
            @overload
            def f(self, x):  # type: (int) -> int
                pass

            @overload
            def f(self, x):  # type: (str) -> str
                pass

            def f(self, x): return x
        """)

    def test_not_a_typing_overload(self):
        """regression test for @typing.overload detection bug in 2.1.0"""
        self.flakes("""
            def foo(x):
                return x

            @foo
            def bar():
                pass

            def bar():
                pass
        """, m.RedefinedWhileUnused)

    @skipIf(version_info < (3, 6), 'new in Python 3.6')
    def test_variable_annotations(self):
        self.flakes('''
        name: str
        age: int
        ''')
        self.flakes('''
        name: str = 'Bob'
        age: int = 18
        ''')
        self.flakes('''
        class C:
            name: str
            age: int
        ''')
        self.flakes('''
        class C:
            name: str = 'Bob'
            age: int = 18
        ''')
        self.flakes('''
        def f():
            name: str
            age: int
        ''')
        self.flakes('''
        def f():
            name: str = 'Bob'
            age: int = 18
            foo: not_a_real_type = None
        ''', m.UnusedVariable, m.UnusedVariable, m.UnusedVariable, m.UndefinedName)
        self.flakes('''
        def f():
            name: str
            print(name)
        ''', m.UndefinedName)
        self.flakes('''
        from typing import Any
        def f():
            a: Any
        ''')
        self.flakes('''
        foo: not_a_real_type
        ''', m.UndefinedName)
        self.flakes('''
        foo: not_a_real_type = None
        ''', m.UndefinedName)
        self.flakes('''
        class C:
            foo: not_a_real_type
        ''', m.UndefinedName)
        self.flakes('''
        class C:
            foo: not_a_real_type = None
        ''', m.UndefinedName)
        self.flakes('''
        def f():
            class C:
                foo: not_a_real_type
        ''', m.UndefinedName)
        self.flakes('''
        def f():
            class C:
                foo: not_a_real_type = None
        ''', m.UndefinedName)
        self.flakes('''
        from foo import Bar
        bar: Bar
        ''')
        self.flakes('''
        from foo import Bar
        bar: 'Bar'
        ''')
        self.flakes('''
        import foo
        bar: foo.Bar
        ''')
        self.flakes('''
        import foo
        bar: 'foo.Bar'
        ''')
        self.flakes('''
        from foo import Bar
        def f(bar: Bar): pass
        ''')
        self.flakes('''
        from foo import Bar
        def f(bar: 'Bar'): pass
        ''')
        self.flakes('''
        from foo import Bar
        def f(bar) -> Bar: return bar
        ''')
        self.flakes('''
        from foo import Bar
        def f(bar) -> 'Bar': return bar
        ''')
        self.flakes('''
        bar: 'Bar'
        ''', m.UndefinedName)
        self.flakes('''
        bar: 'foo.Bar'
        ''', m.UndefinedName)
        self.flakes('''
        from foo import Bar
        bar: str
        ''', m.UnusedImport)
        self.flakes('''
        from foo import Bar
        def f(bar: str): pass
        ''', m.UnusedImport)
        self.flakes('''
        def f(a: A) -> A: pass
        class A: pass
        ''', m.UndefinedName, m.UndefinedName)
        self.flakes('''
        def f(a: 'A') -> 'A': return a
        class A: pass
        ''')
        self.flakes('''
        a: A
        class A: pass
        ''', m.UndefinedName)
        self.flakes('''
        a: 'A'
        class A: pass
        ''')
        self.flakes('''
        a: 'A B'
        ''', m.ForwardAnnotationSyntaxError)
        self.flakes('''
        a: 'A; B'
        ''', m.ForwardAnnotationSyntaxError)
        self.flakes('''
        a: '1 + 2'
        ''')
        self.flakes('''
        a: 'a: "A"'
        ''', m.ForwardAnnotationSyntaxError)

    @skipIf(version_info < (3, 5), 'new in Python 3.5')
    def test_annotated_async_def(self):
        self.flakes('''
        class c: pass
        async def func(c: c) -> None: pass
        ''')

    @skipIf(version_info < (3, 7), 'new in Python 3.7')
    def test_postponed_annotations(self):
        self.flakes('''
        from __future__ import annotations
        def f(a: A) -> A: pass
        class A:
            b: B
        class B: pass
        ''')

        self.flakes('''
        from __future__ import annotations
        def f(a: A) -> A: pass
        class A:
            b: Undefined
        class B: pass
        ''', m.UndefinedName)

    def test_typeCommentsMarkImportsAsUsed(self):
        self.flakes("""
        from mod import A, B, C, D, E, F, G

        def f(
            a,  # type: A
        ):
            # type: (...) -> B
            for b in a:  # type: C
                with b as c:  # type: D
                    d = c.x  # type: E
                    return d

        def g(x):  # type: (F) -> G
            return x.y
        """)

    def test_typeCommentsFullSignature(self):
        self.flakes("""
        from mod import A, B, C, D
        def f(a, b):
            # type: (A, B[C]) -> D
            return a + b
        """)

    def test_typeCommentsStarArgs(self):
        self.flakes("""
        from mod import A, B, C, D
        def f(a, *b, **c):
            # type: (A, *B, **C) -> D
            return a + b
        """)

    def test_typeCommentsFullSignatureWithDocstring(self):
        self.flakes('''
        from mod import A, B, C, D
        def f(a, b):
            # type: (A, B[C]) -> D
            """do the thing!"""
            return a + b
        ''')

    def test_typeCommentsAdditionalComemnt(self):
        self.flakes("""
        from mod import F

        x = 1 # type: F  # noqa
        """)

    def test_typeCommentsNoWhitespaceAnnotation(self):
        self.flakes("""
        from mod import F

        x = 1  #type:F
        """)

    def test_typeCommentsInvalidDoesNotMarkAsUsed(self):
        self.flakes("""
        from mod import F

        # type: F
        """, m.UnusedImport)

    def test_typeCommentsSyntaxError(self):
        self.flakes("""
        def f(x):  # type: (F[) -> None
            pass
        """, m.CommentAnnotationSyntaxError)

    def test_typeCommentsSyntaxErrorCorrectLine(self):
        checker = self.flakes("""\
        x = 1
        # type: definitely not a PEP 484 comment
        """, m.CommentAnnotationSyntaxError)
        self.assertEqual(checker.messages[0].lineno, 2)

    def test_typeCommentsAssignedToPreviousNode(self):
        # This test demonstrates an issue in the implementation which
        # associates the type comment with a node above it, however the type
        # comment isn't valid according to mypy.  If an improved approach
        # which can detect these "invalid" type comments is implemented, this
        # test should be removed / improved to assert that new check.
        self.flakes("""
        from mod import F
        x = 1
        # type: F
        """)
<|endoftext|>


<commit_before># coding: utf-8
# Copyright (c) Pymatgen Development Team.
# Distributed under the terms of the MIT License.

from __future__ import division, unicode_literals

import re
import csv
import collections
import itertools
from io import open
import math
from six.moves import zip
import logging

from monty.json import MSONable, MontyDecoder

import numpy as np

from pymatgen.core.structure import Molecule
from pymatgen.io.qchem.outputs import edges_from_babel, build_MoleculeGraph, is_isomorphic
import networkx as nx

# have_babel = True
# try:
#     from pymatgen.io.babel import BabelMolAdaptor
#     import openbabel as ob
# except ImportError:
#     print("Cannot find OpenBabel! Thus, bonds must be provided by the user.")
#     have_babel = False

"""
This module defines tools to calculate bond dissocation energies.
"""

__author__ = "Samuel Blau"
__copyright__ = "Copyright 2018, The Materials Project"
__version__ = "1.0"
__maintainer__ = "Samuel Blau"
__email__ = "samblau1@gmail.com"
__status__ = "Alpha"
__date__ = "7/26/18"

logger = logging.getLogger(__name__)

class BondDissociationEnergies(MSONable):

    def __init__(self, molecule_entry, fragment_entries):
        """
        Standard constructor for bond dissociation energies.

        Args:

        Get list of all bonds in target molecule
        Build molecule graph of target molecule
        Loop through bonds aka edges and attempt to split into disconnected subgraphs
            On success:
                Search through fragments: find entries with initial OR final molecules which are isomorphic
                Should be three entires per fragment: original charge, OC+1, OC-1
                If there are too few entries:
                    If one of the fragments is an H atom, you should only find two. H+ should be set (to what value???)
                    If fragment matches initial, but not final structure:
                        Just accept the rearrangement as part of the energetic cost of dissociation?
                        At the very least, flag this structure change in the eventual entry
                    If fragment matches final, but not initial structure:
                        Ignore unless no other fragment matches either?
                    If fragment is unstable:
                        Find its subfragments? Is that even sensical?
                        Just pass?
                    If fragment truly cannot be found, aka has not been calculated or calculate failed to complete:
                        Tell user to calculate it and pass
                If there are too many entries:
                    Get rid of duplicates
                        Going to be annoying given time stamps, rounding errors, etc
                    Remove those whose initial rems don't conform to current defaults - Happens before passing
                    If still more than three, for now, isolate differences, report to user, and grab most recent
                    Just take more negative?
                Grab energies of each fragment at all three charges
                If target molecule is neutral:
                    Calculate E_frag1(0)  + E_frag2(0)
                    Calculate E_frag1(+1) + E_frag2(-1)
                    Calculate E_frag1(-1) + E_frag2(+1)
                If target molecule is +1:
                    Calculate E_frag1(0)  + E_frag2(+1)
                    Calculate E_frag1(+1) + E_frag2(0)
                    Calculate E_frag1(+2) + E_frag2(-1)?
                    Calculate E_frag1(-1) + E_frag2(+2)?
                If target molecule is -1:
                    Calculate E_frag1(0)  + E_frag2(-1)
                    Calculate E_frag1(-1) + E_frag2(0)
                    Calculate E_frag1(-2) + E_frag2(+1)?
                    Calculate E_frag1(+1) + E_frag2(-2)?
                Don't limit - just make final charge add to original charge
                Save E_molecule - calculated values associated with bond indices & charges in some data structure
            On failure:
                Track bonds that do not break to yield subgraphs as ring bonds
        Make all possible pairs of ring bonds
        Loop through ring bond pairs and attempt to split into disconnected subgraphs
            Same success/failure procedure as above. Should be wrapped in a function.
        Resolve individual rings
        Attempt to solve for single ring bond dissociation energies, careful with charge

        """

        self.molecule_entry = molecule_entry
        self.fragment_entries = fragment_entries
        self.bond_dissociation_energies = []
        # self.all_bonds = edges_from_babel(target_molecule)
        self.ring_bonds = []
        self.bad_pairs = []
        self.mol_graph = build_MoleculeGraph(Molecule.from_dict(molecule_entry["output"]["optimized_molecule"]))#, self.all_bonds)
        for bond in self.mol_graph.graph.edges: #switch to all_bonds?
            bonds = [(bond[0],bond[1])]
            self.fragment_and_process(bonds)
        self.bond_pairs = []
        for ii,bond in enumerate(self.ring_bonds):
            for jj in range(ii+1,len(self.ring_bonds)):
                bond_pair = [bond, self.ring_bonds[jj]]
                self.bond_pairs += [bond_pair]
        if len(self.bond_pairs) > 0:
            print("Ring bonds detected!")
        for bond_pair in self.bond_pairs:
            self.fragment_and_process(bond_pair)
        self.solve_ring_bonds()
        # return self.bond_dissociation_energies

    def fragment_and_process(self, bonds):
        try:
            frags = self.mol_graph.split_molecule_subgraphs(bonds,allow_reverse=True)
            frag_success = True
        except RuntimeError:
            if len(bonds) == 1:
                self.ring_bonds += bonds
            elif len(bonds) == 2:
                self.bad_pairs += bonds
            else:
                print('No reason to try and break more than two bonds at once! Exiting...')
                raise ValueError
            frag_success = False
        if frag_success:
            frag1_entries = self.search_fragment_entries(frags[0])[0]
            frag2_entries = self.search_fragment_entries(frags[1])[0]
            for frag1 in frag1_entries:
                for frag2 in frag2_entries:
                    if frag1["calcs_reversed"][-1]["input"]["rem"]["method"] == frag2["calcs_reversed"][-1]["input"]["rem"]["method"] and frag1["calcs_reversed"][-1]["input"]["rem"]["basis"] == frag2["calcs_reversed"][-1]["input"]["rem"]["basis"]: # Add PCM check?
                        if frag1["output"]["initial_molecule"]["charge"] + frag2["output"]["initial_molecule"]["charge"] == self.molecule_entry["output"]["optimized_molecule"]["charge"] and self.molecule_entry["calcs_reversed"][-1]["input"]["rem"]["basis"] == frag1["calcs_reversed"][-1]["input"]["rem"]["basis"] and self.molecule_entry["calcs_reversed"][-1]["input"]["rem"]["method"] == frag1["calcs_reversed"][-1]["input"]["rem"]["method"]:
                            coords = nx.get_node_attributes(self.mol_graph.graph, "coords")
                            specie = nx.get_node_attributes(self.mol_graph.graph, "specie")
                            if frag1["output"]["final_energy"] == "unstable" or frag2["output"]["final_energy"] == "unstable":
                                new_entry = ["unstable", bonds, specie[bonds[0][0]], specie[bonds[0][1]], coords[bonds[0][0]], coords[bonds[0][1]], Molecule.from_dict(frag1["output"]["initial_molecule"]).composition.reduced_formula, frag1["output"]["initial_molecule"]["charge"], frag1["output"]["final_energy"], Molecule.from_dict(frag2["output"]["initial_molecule"]).composition.reduced_formula, frag2["output"]["initial_molecule"]["charge"], frag2["output"]["final_energy"]]
                            else:
                                new_entry = [self.molecule_entry["output"]["final_energy"] - (frag1["output"]["final_energy"] + frag2["output"]["final_energy"]), bonds, specie[bonds[0][0]], specie[bonds[0][1]], coords[bonds[0][0]], coords[bonds[0][1]], Molecule.from_dict(frag1["output"]["initial_molecule"]).composition.reduced_formula, frag1["output"]["initial_molecule"]["charge"], frag1["output"]["final_energy"], Molecule.from_dict(frag2["output"]["initial_molecule"]).composition.reduced_formula, frag2["output"]["initial_molecule"]["charge"], frag2["output"]["final_energy"]]
                            self.bond_dissociation_energies += [new_entry]

    def search_fragment_entries(self, frag):
        entries = []
        initial_entries = []
        final_entries = []
        for entry in self.fragment_entries:
            initial_graph = build_MoleculeGraph(Molecule.from_dict(entry["input"]["initial_molecule"])).graph
            final_graph = build_MoleculeGraph(Molecule.from_dict(entry["output"]["initial_molecule"])).graph
            if is_isomorphic(frag.graph, initial_graph) and is_isomorphic(frag.graph, final_graph):
                entries += [entry]
            elif is_isomorphic(frag.graph, initial_graph):
                initial_entries += [entry]
            elif is_isomorphic(frag.graph, final_graph):
                final_entries += [entry]
        return [entries, initial_entries, final_entries]

    def solve_ring_bonds(self):
        pass
<commit_msg>debugging<commit_after># coding: utf-8
# Copyright (c) Pymatgen Development Team.
# Distributed under the terms of the MIT License.

from __future__ import division, unicode_literals

import re
import csv
import collections
import itertools
from io import open
import math
from six.moves import zip
import logging

from monty.json import MSONable, MontyDecoder

import numpy as np

from pymatgen.core.structure import Molecule
from pymatgen.io.qchem.outputs import edges_from_babel, build_MoleculeGraph, is_isomorphic
import networkx as nx

# have_babel = True
# try:
#     from pymatgen.io.babel import BabelMolAdaptor
#     import openbabel as ob
# except ImportError:
#     print("Cannot find OpenBabel! Thus, bonds must be provided by the user.")
#     have_babel = False

"""
This module defines tools to calculate bond dissocation energies.
"""

__author__ = "Samuel Blau"
__copyright__ = "Copyright 2018, The Materials Project"
__version__ = "1.0"
__maintainer__ = "Samuel Blau"
__email__ = "samblau1@gmail.com"
__status__ = "Alpha"
__date__ = "7/26/18"

logger = logging.getLogger(__name__)

class BondDissociationEnergies(MSONable):

    def __init__(self, molecule_entry, fragment_entries):
        """
        Standard constructor for bond dissociation energies.

        Args:

        Get list of all bonds in target molecule
        Build molecule graph of target molecule
        Loop through bonds aka edges and attempt to split into disconnected subgraphs
            On success:
                Search through fragments: find entries with initial OR final molecules which are isomorphic
                Should be three entires per fragment: original charge, OC+1, OC-1
                If there are too few entries:
                    If one of the fragments is an H atom, you should only find two. H+ should be set (to what value???)
                    If fragment matches initial, but not final structure:
                        Just accept the rearrangement as part of the energetic cost of dissociation?
                        At the very least, flag this structure change in the eventual entry
                    If fragment matches final, but not initial structure:
                        Ignore unless no other fragment matches either?
                    If fragment is unstable:
                        Find its subfragments? Is that even sensical?
                        Just pass?
                    If fragment truly cannot be found, aka has not been calculated or calculate failed to complete:
                        Tell user to calculate it and pass
                If there are too many entries:
                    Get rid of duplicates
                        Going to be annoying given time stamps, rounding errors, etc
                    Remove those whose initial rems don't conform to current defaults - Happens before passing
                    If still more than three, for now, isolate differences, report to user, and grab most recent
                    Just take more negative?
                Grab energies of each fragment at all three charges
                If target molecule is neutral:
                    Calculate E_frag1(0)  + E_frag2(0)
                    Calculate E_frag1(+1) + E_frag2(-1)
                    Calculate E_frag1(-1) + E_frag2(+1)
                If target molecule is +1:
                    Calculate E_frag1(0)  + E_frag2(+1)
                    Calculate E_frag1(+1) + E_frag2(0)
                    Calculate E_frag1(+2) + E_frag2(-1)?
                    Calculate E_frag1(-1) + E_frag2(+2)?
                If target molecule is -1:
                    Calculate E_frag1(0)  + E_frag2(-1)
                    Calculate E_frag1(-1) + E_frag2(0)
                    Calculate E_frag1(-2) + E_frag2(+1)?
                    Calculate E_frag1(+1) + E_frag2(-2)?
                Don't limit - just make final charge add to original charge
                Save E_molecule - calculated values associated with bond indices & charges in some data structure
            On failure:
                Track bonds that do not break to yield subgraphs as ring bonds
        Make all possible pairs of ring bonds
        Loop through ring bond pairs and attempt to split into disconnected subgraphs
            Same success/failure procedure as above. Should be wrapped in a function.
        Resolve individual rings
        Attempt to solve for single ring bond dissociation energies, careful with charge

        """

        self.molecule_entry = molecule_entry
        self.fragment_entries = fragment_entries
        self.bond_dissociation_energies = []
        # self.all_bonds = edges_from_babel(target_molecule)
        self.ring_bonds = []
        self.bad_pairs = []
        self.mol_graph = build_MoleculeGraph(Molecule.from_dict(molecule_entry["output"]["optimized_molecule"]))#, self.all_bonds)
        for bond in self.mol_graph.graph.edges: #switch to all_bonds?
            bonds = [(bond[0],bond[1])]
            self.fragment_and_process(bonds)
        self.bond_pairs = []
        for ii,bond in enumerate(self.ring_bonds):
            for jj in range(ii+1,len(self.ring_bonds)):
                bond_pair = [bond, self.ring_bonds[jj]]
                self.bond_pairs += [bond_pair]
        if len(self.bond_pairs) > 0:
            print("Ring bonds detected!")
        for bond_pair in self.bond_pairs:
            self.fragment_and_process(bond_pair)
        self.solve_ring_bonds()
        # return self.bond_dissociation_energies

    def fragment_and_process(self, bonds):
        try:
            frags = self.mol_graph.split_molecule_subgraphs(bonds,allow_reverse=True)
            frag_success = True
        except RuntimeError:
            if len(bonds) == 1:
                self.ring_bonds += bonds
            elif len(bonds) == 2:
                self.bad_pairs += bonds
            else:
                print('No reason to try and break more than two bonds at once! Exiting...')
                raise ValueError
            frag_success = False
        if frag_success:
            frag1_entries = self.search_fragment_entries(frags[0])[0]
            frag2_entries = self.search_fragment_entries(frags[1])[0]
            if bonds == [(0, 14)] or bonds == [(6, 10)]:
                print()
                print(bonds)
                print()
                print(frag1_entries)
                print()
                print(frag2_entries)
                print()
            for frag1 in frag1_entries:
                for frag2 in frag2_entries:
                    if frag1["calcs_reversed"][-1]["input"]["rem"]["method"] == frag2["calcs_reversed"][-1]["input"]["rem"]["method"] and frag1["calcs_reversed"][-1]["input"]["rem"]["basis"] == frag2["calcs_reversed"][-1]["input"]["rem"]["basis"]: # Add PCM check?
                        if frag1["output"]["initial_molecule"]["charge"] + frag2["output"]["initial_molecule"]["charge"] == self.molecule_entry["output"]["optimized_molecule"]["charge"] and self.molecule_entry["calcs_reversed"][-1]["input"]["rem"]["basis"] == frag1["calcs_reversed"][-1]["input"]["rem"]["basis"] and self.molecule_entry["calcs_reversed"][-1]["input"]["rem"]["method"] == frag1["calcs_reversed"][-1]["input"]["rem"]["method"]:
                            coords = nx.get_node_attributes(self.mol_graph.graph, "coords")
                            specie = nx.get_node_attributes(self.mol_graph.graph, "specie")
                            if frag1["output"]["final_energy"] == "unstable" or frag2["output"]["final_energy"] == "unstable":
                                new_entry = ["unstable", bonds, specie[bonds[0][0]], specie[bonds[0][1]], coords[bonds[0][0]], coords[bonds[0][1]], Molecule.from_dict(frag1["output"]["initial_molecule"]).composition.reduced_formula, frag1["output"]["initial_molecule"]["charge"], frag1["output"]["final_energy"], Molecule.from_dict(frag2["output"]["initial_molecule"]).composition.reduced_formula, frag2["output"]["initial_molecule"]["charge"], frag2["output"]["final_energy"]]
                            else:
                                new_entry = [self.molecule_entry["output"]["final_energy"] - (frag1["output"]["final_energy"] + frag2["output"]["final_energy"]), bonds, specie[bonds[0][0]], specie[bonds[0][1]], coords[bonds[0][0]], coords[bonds[0][1]], Molecule.from_dict(frag1["output"]["initial_molecule"]).composition.reduced_formula, frag1["output"]["initial_molecule"]["charge"], frag1["output"]["final_energy"], Molecule.from_dict(frag2["output"]["initial_molecule"]).composition.reduced_formula, frag2["output"]["initial_molecule"]["charge"], frag2["output"]["final_energy"]]
                            self.bond_dissociation_energies += [new_entry]

    def search_fragment_entries(self, frag):
        entries = []
        initial_entries = []
        final_entries = []
        for entry in self.fragment_entries:
            initial_graph = build_MoleculeGraph(Molecule.from_dict(entry["input"]["initial_molecule"])).graph
            final_graph = build_MoleculeGraph(Molecule.from_dict(entry["output"]["initial_molecule"])).graph
            if is_isomorphic(frag.graph, initial_graph) and is_isomorphic(frag.graph, final_graph):
                entries += [entry]
            elif is_isomorphic(frag.graph, initial_graph):
                initial_entries += [entry]
            elif is_isomorphic(frag.graph, final_graph):
                final_entries += [entry]
        return [entries, initial_entries, final_entries]

    def solve_ring_bonds(self):
        pass
<|endoftext|>