<commit_before>"""
The classes in this module enable random access to a variety of file formats
(BAM, bigWig, bigBed, BED) using a uniform syntax, and allow you to compute
coverage across many features in parallel or just a single feature.

Using classes in the :mod:`metaseq.integration` and :mod:`metaseq.minibrowser`
modules, you can connect these objects to matplotlib figures that show a window
into the data, making exploration easy and interactive.

Generally, the :func:`genomic_signal` function is all you need -- just provide
a filename and the format and it will take care of the rest, returning
a genomic signal of the proper type.

Adding support for a new format is straightforward:

    * Write a new adapter for the format in :mod:`metaseq.filetype_adapters`
    * Subclass one of the existing classes below, setting the `adapter`
      attribute to be an instance of this new adapter
    * Add the new class to the `_registry` dictionary to enable support for the
      file format.

Note that to support parallel processing and to avoid repeating code, these
classes delegate their local_coverage methods to the
:func:`metaseq.array_helpers._local_coverage` function.
"""

import os
import sys
import subprocess

import numpy as np
from bx.bbi.bigwig_file import BigWigFile

import pybedtools

from array_helpers import _array, _array_parallel, _local_coverage, \
    _local_count
import filetype_adapters
import rebin
import helpers

def supported_formats():
    """
    Returns list of formats supported by metaseq's genomic signal objects.
    """
    return _registry.keys()

def genomic_signal(fn, kind):
    """
    Factory function that makes the right class for the file format.

    Typically you'll only need this function to create a new genomic signal
    object.

    :param fn: Filename
    :param kind:
        String.  Format of the file; see
        metaseq.genomic_signal._registry.keys()
    """
    try:
        klass = _registry[kind.lower()]
    except KeyError:
        raise ValueError(
            'No support for %s format, choices are %s'
            % (kind, _registry.keys()))
    m = klass(fn)
    m.kind = kind
    return m

class BaseSignal(object):
    """
    Base class to represent objects from which genomic signal can be
    calculated/extracted.

    `__getitem__` uses the underlying adapter the instance was created with
    (e.g., :class:`metaseq.filetype_adapters.BamAdapter` for
    a :class:`BamSignal` object).
    """
    def __init__(self, fn):
        self.fn = fn

    def array(self, features, processes=None, chunksize=1, ragged=False,
              **kwargs):
        """
        Creates an MxN NumPy array of genomic signal for the region defined by
        each feature in `features`, where M=len(features) and N=(bins or
        feature length)

        Parameters
        ----------
        features : iterable of interval-like objects
            An iterable of interval-like objects; see docstring for
            `local_coverage` method for more details.

        processes : int or None
            If not None, then create the array in parallel, giving each process
            chunks of length `chunksize` to work on.

        chunksize : int
            `features` will be split into `chunksize` pieces, and each piece
            will be given to a different process. The optimum value is
            dependent on the size of the features and the underlying data set,
            but `chunksize=100` is a good place to start.

        ragged : bool
            If False (default), then return a 2-D NumPy array.  This requires
            all rows to have the same number of columns, which you get when
            supplying `bins` or if all features are of uniform length.  If
            True, then return a list of 1-D NumPy arrays

        Additional keyword args are passed to local_coverage() which performs
        the work for each feature; see that method for more details.
        """
        if processes is not None:
            arrays = _array_parallel(
                self.fn, self.__class__, features, processes=processes,
                chunksize=chunksize, **kwargs)
        else:
            arrays = _array(self.fn, self.__class__, features, **kwargs)
        if not ragged:
            stacked_arrays = np.row_stack(arrays)
            del arrays
            return stacked_arrays
        else:
            return arrays

    def local_coverage(self, features, *args, **kwargs):
        processes = kwargs.pop('processes', None)
        if not processes:
            return _local_coverage(self.adapter, features, *args, **kwargs)

        if isinstance(features, (list, tuple)):
            raise ValueError(
                "only single features are supported for parallel "
                "local_coverage")

        # we don't want to have self.array do the binning
        bins = kwargs.pop('bins', None)

        # since if we got here processes is not None, then this will trigger
        # a parallel array creation
        features = helpers.tointerval(features)
        x = np.arange(features.start, features.stop)
        features = list(helpers.split_feature(features, processes))
        ys = self.array(
            features, *args, bins=None, processes=processes, ragged=True,
            **kwargs)
        # now we ravel() and re-bin
        y = np.column_stack(ys).ravel()
        if bins:
            xi, yi = rebin.rebin(x, y, bins)
            del x, y
            return xi, yi
        return x, y

class BigWigSignal(BaseSignal):
    def __init__(self, fn):
        """
        Class for operating on bigWig files
        """
        super(BigWigSignal, self).__init__(fn)
        self.adapter = filetype_adapters.BigWigAdapter(fn)

class IntervalSignal(BaseSignal):
    def __init__(self, fn):
        """
        Abstract class for bed, BAM and bigBed files.
        """
        BaseSignal.__init__(self, fn)

    def local_count(self, *args, **kwargs):
        return _local_count(self.adapter, *args, **kwargs)

    local_count.__doc__ = _local_count.__doc__

class BamSignal(IntervalSignal):
    def __init__(self, fn):
        """
        Class for operating on BAM files.
        """
        BaseSignal.__init__(self, fn)
        self._readcount = None
        self.adapter = filetype_adapters.BamAdapter(self.fn)

    def genome(self):
        """
        "genome" dictionary ready for pybedtools, based on the BAM header.
        """
        # This gets the underlying pysam Samfile object
        f = self.adapter.fileobj
        d = {}
        for ref, length in zip(f.references, f.lengths):
            d[ref] = (0, length)
        return d

    def million_mapped_reads(self, force=False):
        """
        Counts total reads in a BAM file and returns the result in millions of
        reads.

        If a file self.bam + '.scale' exists, then just read the first line of
        that file that doesn't start with a "#".  If such a file doesn't exist,
        then it will be created with the number of reads as the first and only
        line in the file.

        The result is also stored in self._readcount so that the time-consuming
        part only runs once; use force=True to force re-count.
        """
        # Already run?
        if self._readcount and not force:
            return self._readcount

        if os.path.exists(self.fn + '.scale') and not force:
            for line in open(self.fn + '.scale'):
                if line.startswith('#'):
                    continue
                self._readcount = float(line.strip())
                return self._readcount

        cmds = ['samtools',
                'view',
                '-c',
                '-F', '0x4',
                self.fn]
        p = subprocess.Popen(
            cmds, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        stdout, stderr = p.communicate()
        if stderr:
            sys.stderr.write('samtools says: %s' % stderr)
            return None
        mmr = int(stdout) / 1e6

        # write to file so the next time you need the lib size you can access
        # it quickly
        if not os.path.exists(self.fn + '.scale'):
            fout = open(self.fn + '.scale', 'w')
            fout.write(str(mmr) + '\n')
            fout.close()

        self._readcount = mmr
        return self._readcount

class BigBedSignal(IntervalSignal):
    def __init__(self, fn):
        """
        Class for operating on bigBed files.
        """
        IntervalSignal.__init__(self, fn)
        self.adapter = filetype_adapters.BigBedAdapter(fn)

class BedSignal(IntervalSignal):
    def __init__(self, fn):
        """
        Class for operating on BED files.
        """
        IntervalSignal.__init__(self, fn)
        self.adapter = filetype_adapters.BedAdapter(fn)

_registry = {
    'bam': BamSignal,
    'bed': BedSignal,
    'gff': BedSignal,
    'gtf': BedSignal,
    'vcf': BedSignal,
    'bigwig': BigWigSignal,
    'bigbed': BigBedSignal,
}
<commit_msg>use .mmr extension for cached million mapped reads files<commit_after>"""
The classes in this module enable random access to a variety of file formats
(BAM, bigWig, bigBed, BED) using a uniform syntax, and allow you to compute
coverage across many features in parallel or just a single feature.

Using classes in the :mod:`metaseq.integration` and :mod:`metaseq.minibrowser`
modules, you can connect these objects to matplotlib figures that show a window
into the data, making exploration easy and interactive.

Generally, the :func:`genomic_signal` function is all you need -- just provide
a filename and the format and it will take care of the rest, returning
a genomic signal of the proper type.

Adding support for a new format is straightforward:

    * Write a new adapter for the format in :mod:`metaseq.filetype_adapters`
    * Subclass one of the existing classes below, setting the `adapter`
      attribute to be an instance of this new adapter
    * Add the new class to the `_registry` dictionary to enable support for the
      file format.

Note that to support parallel processing and to avoid repeating code, these
classes delegate their local_coverage methods to the
:func:`metaseq.array_helpers._local_coverage` function.
"""

import os
import sys
import subprocess

import numpy as np
from bx.bbi.bigwig_file import BigWigFile

import pybedtools

from array_helpers import _array, _array_parallel, _local_coverage, \
    _local_count
import filetype_adapters
import rebin
import helpers

def supported_formats():
    """
    Returns list of formats supported by metaseq's genomic signal objects.
    """
    return _registry.keys()

def genomic_signal(fn, kind):
    """
    Factory function that makes the right class for the file format.

    Typically you'll only need this function to create a new genomic signal
    object.

    :param fn: Filename
    :param kind:
        String.  Format of the file; see
        metaseq.genomic_signal._registry.keys()
    """
    try:
        klass = _registry[kind.lower()]
    except KeyError:
        raise ValueError(
            'No support for %s format, choices are %s'
            % (kind, _registry.keys()))
    m = klass(fn)
    m.kind = kind
    return m

class BaseSignal(object):
    """
    Base class to represent objects from which genomic signal can be
    calculated/extracted.

    `__getitem__` uses the underlying adapter the instance was created with
    (e.g., :class:`metaseq.filetype_adapters.BamAdapter` for
    a :class:`BamSignal` object).
    """
    def __init__(self, fn):
        self.fn = fn

    def array(self, features, processes=None, chunksize=1, ragged=False,
              **kwargs):
        """
        Creates an MxN NumPy array of genomic signal for the region defined by
        each feature in `features`, where M=len(features) and N=(bins or
        feature length)

        Parameters
        ----------
        features : iterable of interval-like objects
            An iterable of interval-like objects; see docstring for
            `local_coverage` method for more details.

        processes : int or None
            If not None, then create the array in parallel, giving each process
            chunks of length `chunksize` to work on.

        chunksize : int
            `features` will be split into `chunksize` pieces, and each piece
            will be given to a different process. The optimum value is
            dependent on the size of the features and the underlying data set,
            but `chunksize=100` is a good place to start.

        ragged : bool
            If False (default), then return a 2-D NumPy array.  This requires
            all rows to have the same number of columns, which you get when
            supplying `bins` or if all features are of uniform length.  If
            True, then return a list of 1-D NumPy arrays

        Additional keyword args are passed to local_coverage() which performs
        the work for each feature; see that method for more details.
        """
        if processes is not None:
            arrays = _array_parallel(
                self.fn, self.__class__, features, processes=processes,
                chunksize=chunksize, **kwargs)
        else:
            arrays = _array(self.fn, self.__class__, features, **kwargs)
        if not ragged:
            stacked_arrays = np.row_stack(arrays)
            del arrays
            return stacked_arrays
        else:
            return arrays

    def local_coverage(self, features, *args, **kwargs):
        processes = kwargs.pop('processes', None)
        if not processes:
            return _local_coverage(self.adapter, features, *args, **kwargs)

        if isinstance(features, (list, tuple)):
            raise ValueError(
                "only single features are supported for parallel "
                "local_coverage")

        # we don't want to have self.array do the binning
        bins = kwargs.pop('bins', None)

        # since if we got here processes is not None, then this will trigger
        # a parallel array creation
        features = helpers.tointerval(features)
        x = np.arange(features.start, features.stop)
        features = list(helpers.split_feature(features, processes))
        ys = self.array(
            features, *args, bins=None, processes=processes, ragged=True,
            **kwargs)
        # now we ravel() and re-bin
        y = np.column_stack(ys).ravel()
        if bins:
            xi, yi = rebin.rebin(x, y, bins)
            del x, y
            return xi, yi
        return x, y

class BigWigSignal(BaseSignal):
    def __init__(self, fn):
        """
        Class for operating on bigWig files
        """
        super(BigWigSignal, self).__init__(fn)
        self.adapter = filetype_adapters.BigWigAdapter(fn)

class IntervalSignal(BaseSignal):
    def __init__(self, fn):
        """
        Abstract class for bed, BAM and bigBed files.
        """
        BaseSignal.__init__(self, fn)

    def local_count(self, *args, **kwargs):
        return _local_count(self.adapter, *args, **kwargs)

    local_count.__doc__ = _local_count.__doc__

class BamSignal(IntervalSignal):
    def __init__(self, fn):
        """
        Class for operating on BAM files.
        """
        BaseSignal.__init__(self, fn)
        self._readcount = None
        self.adapter = filetype_adapters.BamAdapter(self.fn)

    def genome(self):
        """
        "genome" dictionary ready for pybedtools, based on the BAM header.
        """
        # This gets the underlying pysam Samfile object
        f = self.adapter.fileobj
        d = {}
        for ref, length in zip(f.references, f.lengths):
            d[ref] = (0, length)
        return d

    def million_mapped_reads(self, force=False):
        """
        Counts total reads in a BAM file and returns the result in millions of
        reads.

        If a file self.bam + '.scale' exists, then just read the first line of
        that file that doesn't start with a "#".  If such a file doesn't exist,
        then it will be created with the number of reads as the first and only
        line in the file.

        The result is also stored in self._readcount so that the time-consuming
        part only runs once; use force=True to force re-count.
        """
        # Already run?
        if self._readcount and not force:
            return self._readcount

        if os.path.exists(self.fn + '.mmr') and not force:
            for line in open(self.fn + '.mmr'):
                if line.startswith('#'):
                    continue
                self._readcount = float(line.strip())
                return self._readcount

        cmds = ['samtools',
                'view',
                '-c',
                '-F', '0x4',
                self.fn]
        p = subprocess.Popen(
            cmds, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        stdout, stderr = p.communicate()
        if stderr:
            sys.stderr.write('samtools says: %s' % stderr)
            return None
        mmr = int(stdout) / 1e6

        # write to file so the next time you need the lib size you can access
        # it quickly
        if not os.path.exists(self.fn + '.mmr'):
            fout = open(self.fn + '.mmr', 'w')
            fout.write(str(mmr) + '\n')
            fout.close()

        self._readcount = mmr
        return self._readcount

class BigBedSignal(IntervalSignal):
    def __init__(self, fn):
        """
        Class for operating on bigBed files.
        """
        IntervalSignal.__init__(self, fn)
        self.adapter = filetype_adapters.BigBedAdapter(fn)

class BedSignal(IntervalSignal):
    def __init__(self, fn):
        """
        Class for operating on BED files.
        """
        IntervalSignal.__init__(self, fn)
        self.adapter = filetype_adapters.BedAdapter(fn)

_registry = {
    'bam': BamSignal,
    'bed': BedSignal,
    'gff': BedSignal,
    'gtf': BedSignal,
    'vcf': BedSignal,
    'bigwig': BigWigSignal,
    'bigbed': BigBedSignal,
}
<|endoftext|>


<commit_before>from struct import *
from GenFdsGlobalVariable import GenFdsGlobalVariable
import StringIO
from CommonDataClass.FdfClassObject import RegionClassObject
import os
class region(RegionClassObject):
    def __init__(self):
        RegionClassObject.__init__(self)

    """Add RegionData to Fd file"""

    def AddToBuffer(self, Buffer, BaseAddress, BlockSizeList, ErasePolarity, FvBinDict, vtfDict = None, MarcoDict = None):
        Size = self.Size
        GenFdsGlobalVariable.InfLogger('Generate Region')
        GenFdsGlobalVariable.InfLogger("   Region Size = %d" %Size)

        if self.RegionType == 'FV':
            #
            # Get Fv from FvDict
            #
            FvBuffer = StringIO.StringIO('')
            BlockSize = self.__BlockSizeOfRegion__(BlockSizeList)
            if len(self.RegionDataList) == 1:
                RegionData = self.RegionDataList[0]
                fv = GenFdsGlobalVariable.FdfParser.profile.FvDict.get(RegionData.upper())
            #
            # Create local Buffer            #

                if fv != None :
                    GenFdsGlobalVariable.InfLogger('   Region Name = FV')
                    #
                    # Call GenFv tool
                    #

                    self.FvAddress = int(BaseAddress, 16) + self.Offset
                    BlockSize = self.__BlockSizeOfRegion__(BlockSizeList)
                    BlockNum = self.__BlockNumOfRegion__(BlockSize)
                    FvBaseAddress = '0x%x' %self.FvAddress
                    FileName = fv.AddToBuffer(FvBuffer, FvBaseAddress, BlockSize, BlockNum, ErasePolarity, vtfDict)

                    #BinFile = open (FileName, 'r+b')
                    #FvBuffer.write(BinFile.read())
                    FvBinDict[self.RegionDataList[0].upper()] = FileName
            else:
                for RegionData in self.RegionDataList:
                    fv = GenFdsGlobalVariable.FdfParser.profile.FvDict.get(RegionData.upper())
                    if fv != None:
                        FileName = fv.AddToBuffer(FvBuffer, None, BlockSize, None, ErasePolarity, vtfDict)
                        FvBinDict[RegionData.upper()] = FileName

            if FvBuffer.len > Size:
                raise Exception ("Size of Region (%s) is large than Region Size ", self.Offset)
            elif FvBuffer.len < Size :
                raise Exception ("Size of Region (%s) is less than Region Size ", self.Offset)

            #BinFile = open (FileName, 'r+b')
            #Buffer.write(BinFile.read())
            Buffer.write(FvBuffer.getvalue())
            FvBuffer.close()

        if self.RegionType == 'FILE':
            FvBuffer = StringIO.StringIO('')
            for RegionData in self.RegionDataList:
                RegionData = GenFdsGlobalVariable.MarcoExend(RegionData, MarcoDict)
                GenFdsGlobalVariable.InfLogger('   Region File Name = FILE: %s'%RegionData)
                if RegionData[1] != ':' :
                    RegionData = os.path.join (GenFdsGlobalVariable.WorkSpaceDir, RegionData)
                if not os.path.exists(RegionData):
                    raise Excetpion ( 'File: %s dont exist !' %RegionData)

                BinFile = open (RegionData, 'r+b')
                FvBuffer.write(BinFile.read())
                if FvBuffer.len > Size :
                    raise Exception ("Size of File (%s) large than Region Size ", RegionData)

            #
            # If File contents less than region size, append "0xff" after it
            #
            if FvBuffer.len < Size:
                for index in range(0, (Size-FvBuffer.len)):
                    if (ErasePolarity == '1'):
                        FvBuffer.write(pack('B', int('0xFF', 16)))
                    else:
                        FvBuffer.write(pack('B', int('0x00', 16)))
            Buffer.write(FvBuffer.getvalue())
            FvBuffer.close()

        if self.RegionType == 'DATA' :
            GenFdsGlobalVariable.InfLogger('   Region Name = DATA')
            DataSize = 0
            for RegionData in self.RegionDataList:
                Data = RegionData.split(',')
                DataSize = DataSize + len(Data)
                if DataSize > Size:
                   raise Exception ("Size of DATA large than Region Size ")
                else:
                    for item in Data :
                        Buffer.write(pack('B', int(item, 16)))
            if DataSize < Size:
                if (ErasePolarity == '1'):
                    Buffer.write(pack(str(Size -DataSize)+'B', *(int('0xFF', 16) for i in range(Size - DataSize))))
                else:
                    Buffer.write(pack(str(Size - DataSize)+'B', *(int('0x00', 16) for i in range(Size - DataSize))))

        if self.RegionType == None:
            GenFdsGlobalVariable.InfLogger('   Region Name = None')
            if (ErasePolarity == '1') :
                Buffer.write(pack(str(Size)+'B', *(int('0xFF', 16) for i in range(0, Size))))
            else :
                Buffer.write(pack(str(Size)+'B', *(int('0x00', 16) for i in range(0, Size))))

    def __BlockSizeOfRegion__(self, BlockSizeList):
        Offset = 0x00
        BlockSize = 0
        for item in BlockSizeList:
            Offset = Offset + item[0]  * item[1]
            GenFdsGlobalVariable.VerboseLogger ("Offset = 0x%x" %Offset)
            GenFdsGlobalVariable.VerboseLogger ("self.Offset %x" %self.Offset)

            if self.Offset < Offset :
                BlockSize = item[0]
                GenFdsGlobalVariable.VerboseLogger ("BlockSize = %s" %BlockSize)
                return BlockSize
        return BlockSize

    def __BlockNumOfRegion__ (self, BlockSize):
        if BlockSize == 0 :
            raise Exception ("Region: %s doesn't in Fd address scope !" %self.Offset)
        BlockNum = self.Size / BlockSize
        GenFdsGlobalVariable.VerboseLogger ("BlockNum = %x" %BlockNum)
        return BlockNum

<commit_msg>Fix a typo.<commit_after>from struct import *
from GenFdsGlobalVariable import GenFdsGlobalVariable
import StringIO
from CommonDataClass.FdfClassObject import RegionClassObject
import os
class region(RegionClassObject):
    def __init__(self):
        RegionClassObject.__init__(self)

    """Add RegionData to Fd file"""

    def AddToBuffer(self, Buffer, BaseAddress, BlockSizeList, ErasePolarity, FvBinDict, vtfDict = None, MarcoDict = None):
        Size = self.Size
        GenFdsGlobalVariable.InfLogger('Generate Region')
        GenFdsGlobalVariable.InfLogger("   Region Size = %d" %Size)

        if self.RegionType == 'FV':
            #
            # Get Fv from FvDict
            #
            FvBuffer = StringIO.StringIO('')
            BlockSize = self.__BlockSizeOfRegion__(BlockSizeList)
            if len(self.RegionDataList) == 1:
                RegionData = self.RegionDataList[0]
                fv = GenFdsGlobalVariable.FdfParser.profile.FvDict.get(RegionData.upper())
            #
            # Create local Buffer            #

                if fv != None :
                    GenFdsGlobalVariable.InfLogger('   Region Name = FV')
                    #
                    # Call GenFv tool
                    #

                    self.FvAddress = int(BaseAddress, 16) + self.Offset
                    BlockSize = self.__BlockSizeOfRegion__(BlockSizeList)
                    BlockNum = self.__BlockNumOfRegion__(BlockSize)
                    FvBaseAddress = '0x%x' %self.FvAddress
                    FileName = fv.AddToBuffer(FvBuffer, FvBaseAddress, BlockSize, BlockNum, ErasePolarity, vtfDict)

                    #BinFile = open (FileName, 'r+b')
                    #FvBuffer.write(BinFile.read())
                    FvBinDict[self.RegionDataList[0].upper()] = FileName
            else:
                for RegionData in self.RegionDataList:
                    fv = GenFdsGlobalVariable.FdfParser.profile.FvDict.get(RegionData.upper())
                    if fv != None:
                        FileName = fv.AddToBuffer(FvBuffer, None, BlockSize, None, ErasePolarity, vtfDict)
                        FvBinDict[RegionData.upper()] = FileName

            if FvBuffer.len > Size:
                raise Exception ("Size of Region (%s) is large than Region Size ", self.Offset)
            elif FvBuffer.len < Size :
                raise Exception ("Size of Region (%s) is less than Region Size ", self.Offset)

            #BinFile = open (FileName, 'r+b')
            #Buffer.write(BinFile.read())
            Buffer.write(FvBuffer.getvalue())
            FvBuffer.close()

        if self.RegionType == 'FILE':
            FvBuffer = StringIO.StringIO('')
            for RegionData in self.RegionDataList:
                RegionData = GenFdsGlobalVariable.MarcoExend(RegionData, MarcoDict)
                GenFdsGlobalVariable.InfLogger('   Region File Name = FILE: %s'%RegionData)
                if RegionData[1] != ':' :
                    RegionData = os.path.join (GenFdsGlobalVariable.WorkSpaceDir, RegionData)
                if not os.path.exists(RegionData):
                    raise Exception ( 'File: %s dont exist !' %RegionData)

                BinFile = open (RegionData, 'r+b')
                FvBuffer.write(BinFile.read())
                if FvBuffer.len > Size :
                    raise Exception ("Size of File (%s) large than Region Size ", RegionData)

            #
            # If File contents less than region size, append "0xff" after it
            #
            if FvBuffer.len < Size:
                for index in range(0, (Size-FvBuffer.len)):
                    if (ErasePolarity == '1'):
                        FvBuffer.write(pack('B', int('0xFF', 16)))
                    else:
                        FvBuffer.write(pack('B', int('0x00', 16)))
            Buffer.write(FvBuffer.getvalue())
            FvBuffer.close()

        if self.RegionType == 'DATA' :
            GenFdsGlobalVariable.InfLogger('   Region Name = DATA')
            DataSize = 0
            for RegionData in self.RegionDataList:
                Data = RegionData.split(',')
                DataSize = DataSize + len(Data)
                if DataSize > Size:
                   raise Exception ("Size of DATA large than Region Size ")
                else:
                    for item in Data :
                        Buffer.write(pack('B', int(item, 16)))
            if DataSize < Size:
                if (ErasePolarity == '1'):
                    Buffer.write(pack(str(Size -DataSize)+'B', *(int('0xFF', 16) for i in range(Size - DataSize))))
                else:
                    Buffer.write(pack(str(Size - DataSize)+'B', *(int('0x00', 16) for i in range(Size - DataSize))))

        if self.RegionType == None:
            GenFdsGlobalVariable.InfLogger('   Region Name = None')
            if (ErasePolarity == '1') :
                Buffer.write(pack(str(Size)+'B', *(int('0xFF', 16) for i in range(0, Size))))
            else :
                Buffer.write(pack(str(Size)+'B', *(int('0x00', 16) for i in range(0, Size))))

    def __BlockSizeOfRegion__(self, BlockSizeList):
        Offset = 0x00
        BlockSize = 0
        for item in BlockSizeList:
            Offset = Offset + item[0]  * item[1]
            GenFdsGlobalVariable.VerboseLogger ("Offset = 0x%x" %Offset)
            GenFdsGlobalVariable.VerboseLogger ("self.Offset %x" %self.Offset)

            if self.Offset < Offset :
                BlockSize = item[0]
                GenFdsGlobalVariable.VerboseLogger ("BlockSize = %s" %BlockSize)
                return BlockSize
        return BlockSize

    def __BlockNumOfRegion__ (self, BlockSize):
        if BlockSize == 0 :
            raise Exception ("Region: %s doesn't in Fd address scope !" %self.Offset)
        BlockNum = self.Size / BlockSize
        GenFdsGlobalVariable.VerboseLogger ("BlockNum = %x" %BlockNum)
        return BlockNum

<|endoftext|>


<commit_before>"""Set of plots"""
# pylint: disable=invalid-name,too-many-arguments,too-few-public-methods

import datetime
import logging
import math

from abc import ABC
from typing import Optional, Sequence, Tuple

import matplotlib
import matplotlib.pyplot as plt
import numpy as np
import numpy.typing as npt

from matplotlib.backends.backend_pdf import PdfPages

from metatlas.io import write_utils
from metatlas.plots import utils

logger = logging.getLogger(__name__)

def restore_matplotlib_settings(func):
    """Stop function from permanently modifiying rcParams"""

    def wrapper(*args, **kwargs):
        original = matplotlib.rcParams
        out = func(*args, **kwargs)
        matplotlib.rcParams = original
        return out

    return wrapper

class Plot(ABC):
    """A single plot"""

    def __init__(self, title: str, group_name: str):
        self.title = title
        self.group_name = group_name

    def plot(self, ax: matplotlib.axes.Axes, back_color: utils.Color = "white") -> None:
        """
        Draw plot on ax
        back_color: background color for plot
        """
        ax.ticklabel_format(axis="y", scilimits=[0, 0])
        ax.set_facecolor(back_color)
        ax.set_title(f"{self.title}\n{self.group_name}")

class PlotSet(ABC):
    """A Set of related plots"""

    # pylint: disable=too-few-public-methods,too-many-locals,too-many-arguments

    @restore_matplotlib_settings
    def __init__(
        self,
        plots: Sequence[Plot],
        max_plots_per_fig: int = 30,
        x_min: Optional[float] = None,
        x_max: Optional[float] = None,
        y_min: Optional[float] = None,
        y_max: Optional[float] = None,
        sharey: bool = True,
        font_scale: float = 2,
    ):
        num_plots = len(plots)
        num_pages = math.ceil(num_plots / max_plots_per_fig)
        plots_per_page = math.ceil(num_plots / num_pages)
        self.figures = []
        color_generator = utils.colors()
        current_group = ""
        plot_idx = 0
        for _ in range(num_pages):
            plots_remaining = num_plots - plot_idx
            num_plots_this_page = min(plots_per_page, plots_remaining)
            _set_font_size(num_plots_this_page, font_scale)
            cur_fig, axs = utils.wrap_subplots(
                num_plots_this_page, sharey=sharey, sharex=True, constrained_layout=True
            )
            self.figures.append(cur_fig)
            for ax in axs:
                if plots[plot_idx].group_name != current_group:
                    current_color = next(color_generator)
                    current_group = plots[plot_idx].group_name
                plots[plot_idx].plot(ax, back_color=current_color)
                plot_idx += 1
        self.limit_axes(x_min, x_max, y_min, y_max)

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, exc_tb):
        self.close()

    def limit_axes(
        self,
        x_min: Optional[float],
        x_max: Optional[float],
        y_min: Optional[float],
        y_max: Optional[float],
    ) -> None:
        """Update all plots to have the desired axes limits"""
        if x_min is None and x_max is None and y_min is None and y_max is None:
            return
        sides = None
        if x_min is not None or x_max is not None:
            if y_min is None and y_max is None:
                sides = "both"
            elif y_min is None and y_max is not None:
                sides = "min"
            elif y_min is not None and y_max is None:
                sides = "max"
        for fig in self.figures:
            for ax in fig.axes:
                ax.set_xlim(left=x_min, right=x_max)
                ax.set_ylim(bottom=y_min, top=y_max)
                if sides is not None:
                    _autoscale(ax, axis="y", sides=sides)

    def save_pdf(self, file_name: str, title: str, overwrite: bool = False) -> None:
        """Create a PDF file containing one figure per page"""
        write_utils.check_existing_file(file_name, overwrite)
        plt.ioff()  # don't display the plots
        matplotlib.use("agg")  # mitigates a memory leak to not use backend_nbagg
        with PdfPages(file_name) as pdf:
            for fig in self.figures:
                pdf.savefig(fig)
            metadata = pdf.infodict()
            metadata["Title"] = title
            metadata["Author"] = "Joint Genome Institute"
            metadata["CreationDate"] = datetime.datetime.today()
        logger.debug("Exported PDF containing %s to %s.", title, file_name)

    def close(self):
        """Close all plots and free their memory"""
        for fig in self.figures:
            plt.close(fig)

# adapted from https://stackoverflow.com/questions/51323505/how-to-make-relim-and-autoscale-in-a-scatter-plot
def _get_xy(artist: matplotlib.artist.Artist) -> Tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]]:
    """Gets the xy coordinates of a given artist"""
    if "Collection" in str(artist):
        return artist.get_offsets().T
    if "Line" in str(artist):
        return artist.get_xdata(), artist.get_ydata()
    raise ValueError("This type of object isn't implemented yet")

# adapted from https://stackoverflow.com/questions/51323505/how-to-make-relim-and-autoscale-in-a-scatter-plot
def _autoscale(ax: matplotlib.axes.Axes, axis: str = "y", sides: str = "both", margin: float = 0.1) -> None:
    """Autoscales the x or y axis of a given matplotlib ax object
    to fit the margins set by manually limits of the other axis,
    with margins in fraction of the width of the plot
    if sides is 'max' or 'min' then only adjust the limit on that side of axis"""
    assert axis in ["x", "y"]
    assert sides in ["both", "min", "max"]
    low, high = np.inf, -np.inf
    for artist in ax.collections + ax.lines:
        if axis == "y":
            set_lim = ax.set_ylim
            get_lim = ax.get_ylim
            cur_fixed_limit = ax.get_xlim()
            fixed, dependent = _get_xy(artist)
        else:
            set_lim = ax.set_xlim
            get_lim = ax.get_xlim
            cur_fixed_limit = ax.get_ylim()
            dependent, fixed = _get_xy(artist)
        low, high = _update_limts(low, high, fixed, dependent, cur_fixed_limit)
    margin = margin * (high - low)
    if low == np.inf and high == -np.inf:
        return
    assert low != np.inf and high != -np.inf
    new_min = (low - margin) if sides in ["both", "min"] else get_lim()[0]
    new_max = (high + margin) if sides in ["both", "max"] else get_lim()[1]
    set_lim(new_min, new_max)

def _update_limts(
    low: float,
    high: float,
    fixed: npt.NDArray[np.float64],
    dependent: npt.NDArray[np.float64],
    fixed_limits: Tuple[float, float],
) -> Tuple[float, float]:
    """Current limits low and high are updated to include data with ranges
    in the lists fixed and dependent subject to fixed_limits
    """
    local_low, local_high = _calculate_new_limit(fixed, dependent, fixed_limits)
    return min(local_low, low), max(local_high, high)

# adapted from https://stackoverflow.com/questions/51323505/how-to-make-relim-and-autoscale-in-a-scatter-plot
def _calculate_new_limit(
    fixed: npt.NDArray[np.float64], dependent: npt.NDArray[np.float64], fixed_limits: Tuple[float, float]
) -> Tuple[float, float]:
    """Calculates the min/max of the dependent axis given a fixed axis with limits"""
    if len(fixed) > 2:
        mask = (fixed > fixed_limits[0]) & (fixed < fixed_limits[1])
        window = dependent[mask]
        if len(window) == 0:
            return np.inf, -np.inf
        return window.min(), window.max()
    low = dependent[0]
    high = dependent[-1]
    if low == 0.0 and high == 1.0:
        # This is a axhline in the autoscale direction
        return np.inf, -np.inf
    return low, high

def _set_font_size(num_plots: int, font_scale: float) -> None:
    """Scales the font size down based on the number of plots"""
    nrows, ncols = utils.subplot_dimensions(num_plots)
    scale_factor = font_scale / (nrows * ncols)**0.5
    matplotlib.rcParams.update({"font.size": 10 * scale_factor})
<commit_msg>PlotSet sharey works over multiple pages<commit_after>"""Set of plots"""
# pylint: disable=invalid-name,too-many-arguments,too-few-public-methods

import datetime
import logging
import math

from abc import ABC
from typing import Optional, Sequence, Tuple

import matplotlib
import matplotlib.pyplot as plt
import numpy as np
import numpy.typing as npt

from matplotlib.backends.backend_pdf import PdfPages

from metatlas.io import write_utils
from metatlas.plots import utils

logger = logging.getLogger(__name__)

def restore_matplotlib_settings(func):
    """Stop function from permanently modifiying rcParams"""

    def wrapper(*args, **kwargs):
        original = matplotlib.rcParams
        out = func(*args, **kwargs)
        matplotlib.rcParams = original
        return out

    return wrapper

class Plot(ABC):
    """A single plot"""

    def __init__(self, title: str, group_name: str):
        self.title = title
        self.group_name = group_name

    def plot(self, ax: matplotlib.axes.Axes, back_color: utils.Color = "white") -> None:
        """
        Draw plot on ax
        back_color: background color for plot
        """
        ax.ticklabel_format(axis="y", scilimits=[0, 0])
        ax.set_facecolor(back_color)
        ax.set_title(f"{self.title}\n{self.group_name}")

class PlotSet(ABC):
    """A Set of related plots"""

    # pylint: disable=too-few-public-methods,too-many-locals,too-many-arguments

    @restore_matplotlib_settings
    def __init__(
        self,
        plots: Sequence[Plot],
        max_plots_per_fig: int = 30,
        x_min: Optional[float] = None,
        x_max: Optional[float] = None,
        y_min: Optional[float] = None,
        y_max: Optional[float] = None,
        sharey: bool = True,
        font_scale: float = 2,
    ):
        num_plots = len(plots)
        num_pages = math.ceil(num_plots / max_plots_per_fig)
        plots_per_page = math.ceil(num_plots / num_pages)
        self.figures = []
        color_generator = utils.colors()
        current_group = ""
        plot_idx = 0
        for _ in range(num_pages):
            plots_remaining = num_plots - plot_idx
            num_plots_this_page = min(plots_per_page, plots_remaining)
            _set_font_size(num_plots_this_page, font_scale)
            cur_fig, axs = utils.wrap_subplots(
                num_plots_this_page, sharey=sharey, sharex=True, constrained_layout=True
            )
            self.figures.append(cur_fig)
            for ax in axs:
                if plots[plot_idx].group_name != current_group:
                    current_color = next(color_generator)
                    current_group = plots[plot_idx].group_name
                plots[plot_idx].plot(ax, back_color=current_color)
                plot_idx += 1
        if sharey:
            self.sharey_between_pages()
        self.limit_axes(x_min, x_max, y_min, y_max)

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, exc_tb):
        self.close()

    def sharey_between_pages(self):
        min_y = max_y = None
        for fig in self.figures:
            for ax in fig.axes:
                ax_min, ax_max = ax.get_ylim()
                if min_y is None or ax_min < min_y:
                    min_y = ax_min
                if max_y is None or ax_max > max_y:
                    max_y = ax_max
        if min_y is not None and max_y is not None:
            for fig in self.figures:
                for ax in fig.axes:
                    ax.set_ylim(bottom=min_y, top=max_y)

    def limit_axes(
        self,
        x_min: Optional[float],
        x_max: Optional[float],
        y_min: Optional[float],
        y_max: Optional[float],
    ) -> None:
        """Update all plots to have the desired axes limits"""
        if x_min is None and x_max is None and y_min is None and y_max is None:
            return
        sides = None
        if x_min is not None or x_max is not None:
            if y_min is None and y_max is None:
                sides = "both"
            elif y_min is None and y_max is not None:
                sides = "min"
            elif y_min is not None and y_max is None:
                sides = "max"
        for fig in self.figures:
            for ax in fig.axes:
                ax.set_xlim(left=x_min, right=x_max)
                ax.set_ylim(bottom=y_min, top=y_max)
                if sides is not None:
                    _autoscale(ax, axis="y", sides=sides)

    def save_pdf(self, file_name: str, title: str, overwrite: bool = False) -> None:
        """Create a PDF file containing one figure per page"""
        write_utils.check_existing_file(file_name, overwrite)
        plt.ioff()  # don't display the plots
        matplotlib.use("agg")  # mitigates a memory leak to not use backend_nbagg
        with PdfPages(file_name) as pdf:
            for fig in self.figures:
                pdf.savefig(fig)
            metadata = pdf.infodict()
            metadata["Title"] = title
            metadata["Author"] = "Joint Genome Institute"
            metadata["CreationDate"] = datetime.datetime.today()
        logger.debug("Exported PDF containing %s to %s.", title, file_name)

    def close(self):
        """Close all plots and free their memory"""
        for fig in self.figures:
            plt.close(fig)

# adapted from https://stackoverflow.com/questions/51323505/how-to-make-relim-and-autoscale-in-a-scatter-plot
def _get_xy(artist: matplotlib.artist.Artist) -> Tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]]:
    """Gets the xy coordinates of a given artist"""
    if "Collection" in str(artist):
        return artist.get_offsets().T
    if "Line" in str(artist):
        return artist.get_xdata(), artist.get_ydata()
    raise ValueError("This type of object isn't implemented yet")

# adapted from https://stackoverflow.com/questions/51323505/how-to-make-relim-and-autoscale-in-a-scatter-plot
def _autoscale(ax: matplotlib.axes.Axes, axis: str = "y", sides: str = "both", margin: float = 0.1) -> None:
    """Autoscales the x or y axis of a given matplotlib ax object
    to fit the margins set by manually limits of the other axis,
    with margins in fraction of the width of the plot
    if sides is 'max' or 'min' then only adjust the limit on that side of axis"""
    assert axis in ["x", "y"]
    assert sides in ["both", "min", "max"]
    low, high = np.inf, -np.inf
    for artist in ax.collections + ax.lines:
        if axis == "y":
            set_lim = ax.set_ylim
            get_lim = ax.get_ylim
            cur_fixed_limit = ax.get_xlim()
            fixed, dependent = _get_xy(artist)
        else:
            set_lim = ax.set_xlim
            get_lim = ax.get_xlim
            cur_fixed_limit = ax.get_ylim()
            dependent, fixed = _get_xy(artist)
        low, high = _update_limts(low, high, fixed, dependent, cur_fixed_limit)
    margin = margin * (high - low)
    if low == np.inf and high == -np.inf:
        return
    assert low != np.inf and high != -np.inf
    new_min = (low - margin) if sides in ["both", "min"] else get_lim()[0]
    new_max = (high + margin) if sides in ["both", "max"] else get_lim()[1]
    set_lim(new_min, new_max)

def _update_limts(
    low: float,
    high: float,
    fixed: npt.NDArray[np.float64],
    dependent: npt.NDArray[np.float64],
    fixed_limits: Tuple[float, float],
) -> Tuple[float, float]:
    """Current limits low and high are updated to include data with ranges
    in the lists fixed and dependent subject to fixed_limits
    """
    local_low, local_high = _calculate_new_limit(fixed, dependent, fixed_limits)
    return min(local_low, low), max(local_high, high)

# adapted from https://stackoverflow.com/questions/51323505/how-to-make-relim-and-autoscale-in-a-scatter-plot
def _calculate_new_limit(
    fixed: npt.NDArray[np.float64], dependent: npt.NDArray[np.float64], fixed_limits: Tuple[float, float]
) -> Tuple[float, float]:
    """Calculates the min/max of the dependent axis given a fixed axis with limits"""
    if len(fixed) > 2:
        mask = (fixed > fixed_limits[0]) & (fixed < fixed_limits[1])
        window = dependent[mask]
        if len(window) == 0:
            return np.inf, -np.inf
        return window.min(), window.max()
    low = dependent[0]
    high = dependent[-1]
    if low == 0.0 and high == 1.0:
        # This is a axhline in the autoscale direction
        return np.inf, -np.inf
    return low, high

def _set_font_size(num_plots: int, font_scale: float) -> None:
    """Scales the font size down based on the number of plots"""
    nrows, ncols = utils.subplot_dimensions(num_plots)
    scale_factor = font_scale / (nrows * ncols)**0.5
    matplotlib.rcParams.update({"font.size": 10 * scale_factor})
<|endoftext|>


<commit_before>import function_type
import math
import random
import string_utilities

BUILTIN_FUNCTIONS = {
    '!': function_type.make_type([1], handler=lambda x: not x),
    '&&': function_type.make_type([2], handler=lambda x, y: x and y),
    '||': function_type.make_type([2], handler=lambda x, y: x or y),
    '==': function_type.make_type([2], handler=lambda x, y: x == y),
    '!=': function_type.make_type([2], handler=lambda x, y: x != y),
    '<': function_type.make_type([2], handler=lambda x, y: x < y),
    '<=': function_type.make_type([2], handler=lambda x, y: x <= y),
    '>': function_type.make_type([2], handler=lambda x, y: x > y),
    '>=': function_type.make_type([2], handler=lambda x, y: x >= y),
    '~': function_type.make_type([1], handler=lambda x: -x),
    '+': function_type.make_type([2], handler=lambda x, y: x + y),
    '-': function_type.make_type([2], handler=lambda x, y: x - y),
    '*': function_type.make_type([2], handler=lambda x, y: x * y),
    '/': function_type.make_type([2], handler=lambda x, y: x / y),
    '%': function_type.make_type([2], handler=lambda x, y: x % y),
    'floor': function_type.make_type([1], handler=lambda x: math.floor(x)),
    'ceil': function_type.make_type([1], handler=lambda x: math.ceil(x)),
    'random': function_type.make_type([], handler=lambda: random.random()),
    '$': function_type.make_type([], handler=lambda: ()),
    ',': function_type.make_type([2], handler=lambda x, y: (x, y)),
    'head': function_type.make_type([1], handler=lambda x: x[0]),
    'tail': function_type.make_type([1], handler=lambda x: x[1]),
    'if': function_type.make_type([3], handler=lambda condition, true, false: true if condition else false),
    'out': function_type.make_type([1], handler=lambda x: print(x) or x),
    'out_str': function_type.make_type([1], handler=lambda x: print(string_utilities.string_from_list(x)) or x)
}

if __name__ == '__main__':
    import read_code
    import lexer
    import preparser
    import parser
    import evaluate

    code = read_code.read_code()
    specific_lexer = lexer.Lexer()
    specific_preparser = preparser.Preparser(specific_lexer)
    preast = specific_preparser.preparse(code)
    specific_parser = parser.Parser()
    ast = specific_parser.parse(preast, BUILTIN_FUNCTIONS)
    result = evaluate.evaluate(ast, BUILTIN_FUNCTIONS)
    print(result)

    for some_error in specific_lexer.get_errors() + specific_preparser.get_errors() + specific_parser.get_errors():
        some_error.detect_position(code)
        print(some_error)
<commit_msg>Correct built-in functions<commit_after>import function_type
import math
import random
import string_utilities

BUILTIN_FUNCTIONS = {
    '!': function_type.make_type([1], handler=lambda x: not x),
    '&&': function_type.make_type([2], handler=lambda x, y: x and y),
    '||': function_type.make_type([2], handler=lambda x, y: x or y),
    '==': function_type.make_type([2], handler=lambda x, y: x == y),
    '!=': function_type.make_type([2], handler=lambda x, y: x != y),
    '<': function_type.make_type([2], handler=lambda x, y: x < y),
    '<=': function_type.make_type([2], handler=lambda x, y: x <= y),
    '>': function_type.make_type([2], handler=lambda x, y: x > y),
    '>=': function_type.make_type([2], handler=lambda x, y: x >= y),
    '~': function_type.make_type([1], handler=lambda x: -x),
    '+': function_type.make_type([2], handler=lambda x, y: x + y),
    '-': function_type.make_type([2], handler=lambda x, y: x - y),
    '*': function_type.make_type([2], handler=lambda x, y: x * y),
    '/': function_type.make_type([2], handler=lambda x, y: x / y),
    '%': function_type.make_type([2], handler=lambda x, y: x % y),
    'floor': function_type.make_type([1], handler=lambda x: math.floor(x)),
    'ceil': function_type.make_type([1], handler=lambda x: math.ceil(x)),
    'random': function_type.make_type([], handler=lambda: random.random()),
    '$': function_type.make_type([], handler=lambda: ()),
    ',': function_type.make_type([2], handler=lambda x, y: (x, y)),
    'head': function_type.make_type([1], handler=lambda x: x[0]),
    'tail': function_type.make_type([1], handler=lambda x: x[1]),
    'if': function_type.make_type([3], handler=lambda condition, true, false: true if condition else false),
    'str': function_type.make_type([1], handler=lambda x: string_utilities.string_to_list(str(x))),
    'out': function_type.make_type([1], handler=lambda x: print(string_utilities.string_from_list(x)))
}

if __name__ == '__main__':
    import read_code
    import lexer
    import preparser
    import parser
    import evaluate

    code = read_code.read_code()
    specific_lexer = lexer.Lexer()
    specific_preparser = preparser.Preparser(specific_lexer)
    preast = specific_preparser.preparse(code)
    specific_parser = parser.Parser()
    ast = specific_parser.parse(preast, BUILTIN_FUNCTIONS)
    result = evaluate.evaluate(ast, BUILTIN_FUNCTIONS)
    print(result)

    for some_error in specific_lexer.get_errors() + specific_preparser.get_errors() + specific_parser.get_errors():
        some_error.detect_position(code)
        print(some_error)
<|endoftext|>


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

"""
The POX Messenger system.

The Messenger is a way to build services in POX that can be consumed by
external clients.

Sometimes a controller might need to interact with the outside world.
Sometimes you need to integrate with an existing piece of software and
maybe you don't get to choose how you communicate with it.  Other times,
you have the opportunity and burden of rolling your own.  The Messenger
is meant to help you with the latter case.

In short, the POX Messenger is a system for communicating between NOX and
external programs by exchanging messages encoded in JSON.  It is intended
to be quite general, and supports multiple communication models (for
example, you can use it for pub/sub, virtual circuits, broadcasts...).
It is also transport-independent (as of this writing, it supports a
straightforward TCP socket transport and an HTTP transport).

*Connections* are somehow established when a client connects via some
*Transport*, and this causes a ConnectionStarted event on MessengerHub.
Messages sent by the client at this point raise a MessageReceived event
on MessengerHub, and a disconnection will raise ConnectionClosed.  A
listener of ConnectionStarted or MessageReceived messages may use the
Connection's send() method to send a message back to the client, or may
store the Connection to send() a reply later (e.g., to implement a pubsub
communication model).

Incoming messages are often intended to be received only be specific
services.  To make this happen, a service can create a Target object
specifying name for the target, and then listen to events on this
specific Target object (e.g., MessageReceived, ConnectionClosed).
Then clients can include a TARGET key in their messages with the
value set to be the target name (or a list of target names, or null to
broadcast through MessengerHub).  Messages which have specified an
available TARGET will only appear as messages on the corresponding Target
object -- not on MessengerHub itself.  Messages with a TARGET that is
not available will continue to show up on MessengerHub.

There are also times when a service would like to do a lot communication
with a client over a specific connection.  In this case, the service can
"claim" a connection (via the .claim() method of the event) and listen
to it directly.  After being claimed, the connection will no longer
raise events on the MessengerHub.  This allows for connection-oriented
communication patterns.

Generally, listeners will listen to MessageReceived events and respond
to those.  However, it is also possible to set a connection to buffer
the messages it receives, and it can be polled later.
"""

from pox.lib.revent import *
from pox.lib.recoco import *
import pox.core
from pox.core import core as core
import weakref
import json

log = pox.core.getLogger()

# JSON decoder used by default
defaultDecoder = json.JSONDecoder()

class ConnectionClosed (Event):
  def __init__ (self, connection):
    Event.__init__(self)
    self.con = connection

class ConnectionStarted (Event):
  def __init__ (self, connection):
    Event.__init__(self)
    self.con = connection
    self._claimed = False

  def claim (self):
    #assert self._claimed == False
    self._claimed = True
    self.halt = True
    return self.con

class MessageReceived (Event):
  def __init__ (self, connection, msg):
    Event.__init__(self)
    self.con = connection
    self._claimed = False
    self.msg = msg

  def claim (self):
    assert self._claimed == False
    self._claimed = True
    self.halt = True
    return self.con

  def _invoke (self, handler, *args, **kw):
    # Special handling.  If the handler doesn't return any disposition, then
    # we look to see if they read the message.  If so, we stop processing
    # this event now.
    l = len(self.con._msgs) # Sleazy
    r = handler(self, self.msg, *args, **kw)
    if r is not None:
      return r
    if len(self.con._msgs) < l:
      # Looks like they read it.  Stop processing.
      return EventHalt

class Target (object):
  pass
  #TODO

class MessengerConnection (EventMixin):
  _eventMixin_events = set([
    MessageReceived,
    ConnectionClosed,
  ])

  ID = "INVALID ID"

  def __init__ (self, source, ID = None):
    EventMixin.__init__(self)
    self._isConnected = True
    self._buf = bytes()
    self._msgs = []
    self._source = source
    self._newlines = False
    self.buffered = False
    if ID is not None:
      self.ID = ID

    claimed = False
    e = core.messenger.raiseEventNoErrors(ConnectionStarted, self)
    if e is not None:
      claimed = e._claimed

    if not claimed:
      # Unclaimed events get forwarded to here too
      self.addListener(MessageReceived, self._defaultMessageReceived, priority=-1) # Low priority

  def _close (self):
    # Called internally
    if self._isConnected is False: return
    self._source._forget(self)
    self._isConnected = False
    self.raiseEventNoErrors(ConnectionClosed, self)
    self.raiseEventNoErrors(MessageReceived, self, None)
    core.messenger.raiseEventNoErrors(ConnectionClosed, self)

  def send (self, whatever, **kw):
    if self._isConnected is False: return False
    s = json.dumps(whatever, **kw)
    if self._newlines: s += "\n"
    self.sendRaw(s)
    return True

  def sendRaw (self, data):
    raise RuntimeError("Not implemented")

  def isConnected (self):
    return self._isConnected

  def isReadable (self):
    return len(self._msgs) > 0

  def peek (self, default = None):
    return self.read(default = default, peek=True)

  def read (self, default = None, peek = False):
    if len(self._msgs) == 0: return default
    if peek:
      return self._msgs[0]
    else:
      return self._msgs.pop(0)

  def _recv_msg (self, msg):
    #print self,"recv:",msg
    self._msgs.append(msg)
    self.raiseEventNoErrors(MessageReceived, self, msg)
    if not self.buffered:
      del self._msgs[:]

  def _recv_raw (self, data):
    if len(data) == 0: return
    if len(self._buf) == 0:
      if data[0].isspace():
        self._buf = data.lstrip()
      else:
        self._buf = data
    else:
      self._buf += data

    while len(self._buf) > 0:
      try:
        msg, l = defaultDecoder.raw_decode(self._buf)
      except:
        # Need more data before it's a valid message
        # (.. or the stream is corrupt and things will never be okay ever again)
        return

      self._buf = self._buf[l:]
      if len(self._buf) != 0 and self._buf[0].isspace():
        self._buf = self._buf.lstrip()
      self._recv_msg(msg)

  def __str__ (self):
    # Subclasses probably want to change this
    return "<" + self.__class__.__name__ + "/" + self.ID + ">"

  def _defaultMessageReceived (self, event, msg):
    #print self,"default recv:",msg
    #TODO: make sure this actually works. I have never tried re-raising an event.
    core.messenger.raiseEventNoErrors(event)
    if event._claimed:
      # Someone claimed this connection -- stop forwarding it globally
      return EventRemove

  def close (self):
    self._close()

class TCPMessengerConnection (MessengerConnection, Task):
  def __init__ (self, source, socket):
    self._socket = socket
    MessengerConnection.__init__(self, source, ID=str(id(self))) #TODO: better ID
    Task.__init__(self)

    #self.start()

  def _close (self):
    super(TCPMessengerConnection, self)._close()
    try:
      self._socket.shutdown(socket.SHUT_RDWR)
    except:
      pass

  def sendRaw (self, data):
    try:
      l = self._socket.send(data)
      if l == len(data): return
    except:
      pass
    #TODO: do something more graceful!
    self._close()

  def run (self):
    log.debug("%s started" % (self,))
    while self.isConnected():
      d = yield Recv(self._socket)
      if d is None or len(d) == 0:
        break
      self._recv_raw(d)
    self._close()
    log.debug("%s stopped" % (self,))

  #TODO: __str__ with port numbers, etc

#TODO: a queued listener base class (like for HTTP)

class MessengerHub (EventMixin):
  _eventMixin_events = set([
    ConnectionStarted,  # Immediately when a connection goes up
    ConnectionClosed,   # When a connection goes down
    MessageReceived,    # For unclaimed messages
  ])
  def __init__ (self):
    EventMixin.__init__(self)
    self.connections = weakref.WeakSet()

#TODO: make a superclass source
class TCPMessengerSource (Task):
  def __init__ (self, address = "0.0.0.0", port = 7790):
    port = int(port)
    Task.__init__(self)
    self._addr = (address,port)
    self._connections = set()

  def _forget (self, connection):
    """ Forget about a connection (because it has closed) """
    if connection in self._connections:
      #print "forget about",connection
      self._connections.remove(connection)

  def run (self):
    listener = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    listener.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    listener.bind(self._addr)
    listener.listen(0)

    log.debug("Listening for connections on %s:%i" % (self._addr))

    con = None
    while core.running:
      try:
        rlist, wlist, elist = yield Select([listener])
        if len(rlist) == 0:
          # Must have been interrupted
          break

        rc = TCPMessengerConnection(self, listener.accept()[0])
        self._connections.add(rc)
        rc.start()
      except:
        traceback.print_exc()
        break

    try:
      listener.close()
    except:
      pass
    log.debug("No longer listening for connections")

def launch (tcp_address = "0.0.0.0", tcp_port = 7790):
  core.register("messenger", MessengerHub())
  t = TCPMessengerSource(tcp_address, tcp_port)
  core.addListener(pox.core.GoingUpEvent, lambda event: t.start())
<commit_msg>pox.messenger: don't assume a source for TCPMessageConnection + elimate race condition<commit_after># Copyright 2011 James McCauley
#
# This file is part of POX.
#
# POX is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# POX is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with POX.  If not, see <http://www.gnu.org/licenses/>.

"""
The POX Messenger system.

The Messenger is a way to build services in POX that can be consumed by
external clients.

Sometimes a controller might need to interact with the outside world.
Sometimes you need to integrate with an existing piece of software and
maybe you don't get to choose how you communicate with it.  Other times,
you have the opportunity and burden of rolling your own.  The Messenger
is meant to help you with the latter case.

In short, the POX Messenger is a system for communicating between NOX and
external programs by exchanging messages encoded in JSON.  It is intended
to be quite general, and supports multiple communication models (for
example, you can use it for pub/sub, virtual circuits, broadcasts...).
It is also transport-independent (as of this writing, it supports a
straightforward TCP socket transport and an HTTP transport).

*Connections* are somehow established when a client connects via some
*Transport*, and this causes a ConnectionStarted event on MessengerHub.
Messages sent by the client at this point raise a MessageReceived event
on MessengerHub, and a disconnection will raise ConnectionClosed.  A
listener of ConnectionStarted or MessageReceived messages may use the
Connection's send() method to send a message back to the client, or may
store the Connection to send() a reply later (e.g., to implement a pubsub
communication model).

Incoming messages are often intended to be received only be specific
services.  To make this happen, a service can create a Target object
specifying name for the target, and then listen to events on this
specific Target object (e.g., MessageReceived, ConnectionClosed).
Then clients can include a TARGET key in their messages with the
value set to be the target name (or a list of target names, or null to
broadcast through MessengerHub).  Messages which have specified an
available TARGET will only appear as messages on the corresponding Target
object -- not on MessengerHub itself.  Messages with a TARGET that is
not available will continue to show up on MessengerHub.

There are also times when a service would like to do a lot communication
with a client over a specific connection.  In this case, the service can
"claim" a connection (via the .claim() method of the event) and listen
to it directly.  After being claimed, the connection will no longer
raise events on the MessengerHub.  This allows for connection-oriented
communication patterns.

Generally, listeners will listen to MessageReceived events and respond
to those.  However, it is also possible to set a connection to buffer
the messages it receives, and it can be polled later.
"""

from pox.lib.revent import *
from pox.lib.recoco import *
import pox.core
from pox.core import core as core
import weakref
import json

log = pox.core.getLogger()

# JSON decoder used by default
defaultDecoder = json.JSONDecoder()

class MessengerListening(Event):
  pass

class ConnectionClosed (Event):
  def __init__ (self, connection):
    Event.__init__(self)
    self.con = connection

class ConnectionStarted (Event):
  def __init__ (self, connection):
    Event.__init__(self)
    self.con = connection
    self._claimed = False

  def claim (self):
    #assert self._claimed == False
    self._claimed = True
    self.halt = True
    return self.con

class MessageReceived (Event):
  def __init__ (self, connection, msg):
    Event.__init__(self)
    self.con = connection
    self._claimed = False
    self.msg = msg

  def claim (self):
    assert self._claimed == False
    self._claimed = True
    self.halt = True
    return self.con

  def _invoke (self, handler, *args, **kw):
    # Special handling.  If the handler doesn't return any disposition, then
    # we look to see if they read the message.  If so, we stop processing
    # this event now.
    l = len(self.con._msgs) # Sleazy
    r = handler(self, self.msg, *args, **kw)
    if r is not None:
      return r
    if len(self.con._msgs) < l:
      # Looks like they read it.  Stop processing.
      return EventHalt

class Target (object):
  pass
  #TODO

class MessengerConnection (EventMixin):
  _eventMixin_events = set([
    MessageReceived,
    ConnectionClosed,
  ])

  ID = "INVALID ID"

  def __init__ (self, source=None, ID = None):
    EventMixin.__init__(self)
    self._isConnected = True
    self._buf = bytes()
    self._msgs = []
    self._source = source
    self._newlines = False
    self.buffered = False
    if ID is not None:
      self.ID = ID

    claimed = False
    e = core.messenger.raiseEventNoErrors(ConnectionStarted, self)
    if e is not None:
      claimed = e._claimed

    if not claimed:
      # Unclaimed events get forwarded to here too
      self.addListener(MessageReceived, self._defaultMessageReceived, priority=-1) # Low priority

  def _close (self):
    # Called internally
    if self._isConnected is False: return
    if self._source:
      self._source._forget(self)
    self._isConnected = False
    self.raiseEventNoErrors(ConnectionClosed, self)
    self.raiseEventNoErrors(MessageReceived, self, None)
    core.messenger.raiseEventNoErrors(ConnectionClosed, self)

  def send (self, whatever, **kw):
    if self._isConnected is False: return False
    s = json.dumps(whatever, **kw)
    if self._newlines: s += "\n"
    self.sendRaw(s)
    return True

  def sendRaw (self, data):
    raise RuntimeError("Not implemented")

  def isConnected (self):
    return self._isConnected

  def isReadable (self):
    return len(self._msgs) > 0

  def peek (self, default = None):
    return self.read(default = default, peek=True)

  def read (self, default = None, peek = False):
    if len(self._msgs) == 0: return default
    if peek:
      return self._msgs[0]
    else:
      return self._msgs.pop(0)

  def _recv_msg (self, msg):
    #print self,"recv:",msg
    self._msgs.append(msg)
    self.raiseEventNoErrors(MessageReceived, self, msg)
    if not self.buffered:
      del self._msgs[:]

  def _recv_raw (self, data):
    if len(data) == 0: return
    if len(self._buf) == 0:
      if data[0].isspace():
        self._buf = data.lstrip()
      else:
        self._buf = data
    else:
      self._buf += data

    while len(self._buf) > 0:
      try:
        msg, l = defaultDecoder.raw_decode(self._buf)
      except:
        # Need more data before it's a valid message
        # (.. or the stream is corrupt and things will never be okay ever again)
        return

      self._buf = self._buf[l:]
      if len(self._buf) != 0 and self._buf[0].isspace():
        self._buf = self._buf.lstrip()
      self._recv_msg(msg)

  def __str__ (self):
    # Subclasses probably want to change this
    return "<" + self.__class__.__name__ + "/" + self.ID + ">"

  def _defaultMessageReceived (self, event, msg):
    #print self,"default recv:",msg
    #TODO: make sure this actually works. I have never tried re-raising an event.
    core.messenger.raiseEventNoErrors(event)
    if event._claimed:
      # Someone claimed this connection -- stop forwarding it globally
      return EventRemove

  def close (self):
    self._close()

class TCPMessengerConnection (MessengerConnection, Task):
  def __init__ (self, source=None, socket=None):
    self._socket = socket
    MessengerConnection.__init__(self, source, ID=str(id(self))) #TODO: better ID
    Task.__init__(self)

    #self.start()

  def _close (self):
    super(TCPMessengerConnection, self)._close()
    try:
      self._socket.shutdown(socket.SHUT_RDWR)
    except:
      pass

  def sendRaw (self, data):
    try:
      l = self._socket.send(data)
      if l == len(data): return
    except:
      pass
    #TODO: do something more graceful!
    self._close()

  def run (self):
    log.debug("%s started" % (self,))
    while self.isConnected():
      d = yield Recv(self._socket)
      if d is None or len(d) == 0:
        break
      self._recv_raw(d)
    self._close()
    log.debug("%s stopped" % (self,))

  #TODO: __str__ with port numbers, etc

#TODO: a queued listener base class (like for HTTP)

class MessengerHub (EventMixin):
  _eventMixin_events = set([
    ConnectionStarted,  # Immediately when a connection goes up
    ConnectionClosed,   # When a connection goes down
    MessageReceived,    # For unclaimed messages
    MessengerListening, # The TCP listening port is up
  ])
  def __init__ (self):
    EventMixin.__init__(self)
    self.connections = weakref.WeakSet()

#TODO: make a superclass source
class TCPMessengerSource (Task):
  def __init__ (self, address = "0.0.0.0", port = 7790):
    port = int(port)
    Task.__init__(self)
    self._addr = (address,port)
    self._connections = set()

  def _forget (self, connection):
    """ Forget about a connection (because it has closed) """
    if connection in self._connections:
      #print "forget about",connection
      self._connections.remove(connection)

  def run (self):
    listener = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    listener.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    listener.bind(self._addr)
    listener.listen(0)

    log.debug("Listening for connections on %s:%i" % (self._addr))
    core.messenger.raiseEventNoErrors(MessengerListening())

    con = None
    while core.running:
      try:
        rlist, wlist, elist = yield Select([listener])
        if len(rlist) == 0:
          # Must have been interrupted
          break

        rc = TCPMessengerConnection(self, listener.accept()[0])
        self._connections.add(rc)
        rc.start()
      except:
        traceback.print_exc()
        break

    try:
      listener.close()
    except:
      pass
    log.debug("No longer listening for connections")

def launch (tcp_address = "0.0.0.0", tcp_port = 7790):
  core.register("messenger", MessengerHub())
  t = TCPMessengerSource(tcp_address, tcp_port)
  core.addListener(pox.core.GoingUpEvent, lambda event: t.start())
<|endoftext|>


<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-
import unittest
from SpiffyWorld import Dungeon, UnitGenerator, Unit

class TestDungeon(unittest.TestCase):

    def _get_dungeon_model(self):
        dungeon_model = {
            "id": 1,
            "name": "TestDungeon",
            "min_level": 1,
            "max_level": 100,
            "description": "foo",
            "channel": "#testing",
            "units": []
        }

        return dungeon_model

    def test_create_dungeon(self):
        dungeon_model = self._get_dungeon_model()
        dungeon = Dungeon(dungeon=dungeon_model, announcer={})

        self.assertEqual(len(dungeon.units), 0)

    def test_add_dungeon_units(self):
        dungeon_model = self._get_dungeon_model()
        dungeon = Dungeon(dungeon=dungeon_model, announcer={})

        unit_generator = UnitGenerator()
        unit = unit_generator.generate()

        dungeon.add_unit(unit)

        dungeon_unit = dungeon.units[0]

        self.assertIsInstance(dungeon_unit, Unit)

    def test_get_unit_by_user_id(self):
        dungeon_model = self._get_dungeon_model()
        dungeon = Dungeon(dungeon=dungeon_model, announcer={})
        user_id = 42
        unit_generator = UnitGenerator()
        unit = unit_generator.generate(user_id=user_id)

        dungeon.add_unit(unit)

        actual = dungeon.get_unit_by_user_id(user_id)

        self.assertIsInstance(actual, Unit)

    def test_get_player_by_user_id(self):
        dungeon_model = self._get_dungeon_model()
        dungeon = Dungeon(dungeon=dungeon_model, announcer={})

        user_id = 42
        unit_generator = UnitGenerator()
        unit = unit_generator.generate(user_id=user_id)

        dungeon.add_unit(unit)

        self.assertTrue(len(dungeon.units), 0)

        actual = dungeon.get_player_by_user_id(user_id)

        self.assertIsInstance(actual, Unit)

    def test_get_unit_by_name(self):
        dungeon_model = self._get_dungeon_model()
        dungeon = Dungeon(dungeon=dungeon_model, announcer={})

        unit_name = "quux"

        unit_generator = UnitGenerator()
        unit = unit_generator.generate(unit_name=unit_name)

        dungeon.add_unit(unit)

        self.assertTrue(len(dungeon.units), 0)

        # test exact match
        actual = dungeon.get_unit_by_name(unit_name)
        self.assertIsInstance(actual, Unit)

        # test invalid match
        actual = dungeon.get_unit_by_name("foo")
        self.assertIsNone(actual)

        # test case sensitivity
        actual = dungeon.get_unit_by_name(unit_name.upper())
        self.assertIsInstance(actual, Unit)

        # test startswith
        actual = dungeon.get_unit_by_name(unit_name[0:1])
        self.assertIsInstance(actual, Unit)

    def test_get_living_unit_by_name(self):
        dungeon_model = self._get_dungeon_model()
        dungeon = Dungeon(dungeon=dungeon_model, announcer={})

        unit_name = "baz"
        unit_generator = UnitGenerator()
        unit = unit_generator.generate(unit_name=unit_name)

        dungeon.add_unit(unit)

        # test typical use case
        actual = dungeon.get_living_unit_by_name(unit.get_name())
        self.assertIsInstance(actual, Unit)

        unit.kill()

        actual = dungeon.get_living_unit_by_name(unit.get_name())
        self.assertIsNone(actual)

    def test_get_dead_unit_by_name(self):
        dungeon_model = self._get_dungeon_model()
        dungeon = Dungeon(dungeon=dungeon_model, announcer={})

        unit_name = "Oliver Queen"
        unit_generator = UnitGenerator()
        unit = unit_generator.generate(unit_name=unit_name)

        dungeon.add_unit(unit)

        unit.kill()

        # test typical use case
        actual = dungeon.get_dead_unit_by_name(unit.get_name())
        self.assertIsInstance(actual, Unit)

        unit.hp = 1

        actual_alive = dungeon.get_living_unit_by_name(unit.get_name())
        self.assertIsInstance(actual_alive, Unit)

    def test_get_living_players(self):
        dungeon_model = self._get_dungeon_model()
        dungeon = Dungeon(dungeon=dungeon_model, announcer={})

        unit_generator = UnitGenerator()

        unit = unit_generator.generate(is_player=True)

        dungeon.add_unit(unit)

        dead_unit = unit_generator.generate()
        dead_unit.kill()

        dungeon.add_unit(dead_unit)

        # test typical use case
        actual = dungeon.get_living_players()
        self.assertEqual(len(actual), 1)

        dead_players = dungeon.get_dead_players()
        self.assertEqual(len(dead_players), 0)

    def test_get_unit_status_distribution(self):
        dungeon_model = self._get_dungeon_model()
        dungeon = Dungeon(dungeon=dungeon_model, announcer={})

        """ living hostile player """
        unit_generator = UnitGenerator()
        living_player = unit_generator.generate(is_player=True,
                                                combat_status="hostile")
        dungeon.add_unit(living_player)

        """ dead friendly NPC """
        dead_friendly_unit = unit_generator.generate(combat_status="friendly")
        dead_friendly_unit.kill()
        dungeon.add_unit(dead_friendly_unit)

        """ living friendly NPC """
        living_npc = unit_generator.generate(combat_status="friendly")
        dungeon.add_unit(living_npc)

        """ living hostile NPC """
        hostile_living_npc = unit_generator.generate(combat_status="hostile")
        dungeon.add_unit(hostile_living_npc)

        # TODO: add undead here
        actual = dungeon.get_unit_status_distribution()

        expected = {
            "players": {
                "living": 1,
                "dead": 0,
                "undead": 0
            },
            "npc": {
                "hostile": {
                    "living": 1,
                    "dead": 0,
                    "undead": 0
                },
                "friendly": {
                    "living": 1,
                    "dead": 1,
                    "undead": 0
                }
            }
        }

        self.assertEqual(expected, actual)

if __name__ == '__main__':
    unittest.main()
<commit_msg>Refactors Dungeon tests to accommodate new logger change<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-
import unittest
from SpiffyWorld import Dungeon, UnitGenerator, Unit
import logging
from testfixtures import LogCapture

class TestDungeon(unittest.TestCase):

    def _get_dungeon_model(self):
        dungeon_model = {
            "id": 1,
            "name": "TestDungeon",
            "min_level": 1,
            "max_level": 100,
            "description": "foo",
            "channel": "#testing",
            "units": []
        }

        return dungeon_model

    def test_create_dungeon(self):
        dungeon_model = self._get_dungeon_model()
        with LogCapture():
            logger = logging.getLogger()
            dungeon = Dungeon(dungeon=dungeon_model,
                              announcer={},
                              log=logger)

        self.assertEqual(len(dungeon.units), 0)

    def test_add_dungeon_units(self):
        dungeon_model = self._get_dungeon_model()
        with LogCapture():
            logger = logging.getLogger()
            dungeon = Dungeon(dungeon=dungeon_model,
                              announcer={},
                              log=logger)

        unit_generator = UnitGenerator()
        unit = unit_generator.generate()

        dungeon.add_unit(unit)

        dungeon_unit = dungeon.units[0]

        self.assertIsInstance(dungeon_unit, Unit)

    def test_get_unit_by_user_id(self):
        dungeon_model = self._get_dungeon_model()
        with LogCapture():
            logger = logging.getLogger()
            dungeon = Dungeon(dungeon=dungeon_model,
                              announcer={},
                              log=logger)
        user_id = 42
        unit_generator = UnitGenerator()
        unit = unit_generator.generate(user_id=user_id)

        dungeon.add_unit(unit)

        actual = dungeon.get_unit_by_user_id(user_id)

        self.assertIsInstance(actual, Unit)

    def test_get_player_by_user_id(self):
        dungeon_model = self._get_dungeon_model()
        with LogCapture():
            logger = logging.getLogger()
            dungeon = Dungeon(dungeon=dungeon_model,
                              announcer={},
                              log=logger)

        user_id = 42
        unit_generator = UnitGenerator()
        unit = unit_generator.generate(user_id=user_id)

        dungeon.add_unit(unit)

        self.assertTrue(len(dungeon.units), 0)

        actual = dungeon.get_player_by_user_id(user_id)

        self.assertIsInstance(actual, Unit)

    def test_get_unit_by_name(self):
        dungeon_model = self._get_dungeon_model()
        with LogCapture():
            logger = logging.getLogger()
            dungeon = Dungeon(dungeon=dungeon_model,
                              announcer={},
                              log=logger)

        unit_name = "quux"

        unit_generator = UnitGenerator()
        unit = unit_generator.generate(unit_name=unit_name)

        dungeon.add_unit(unit)

        self.assertTrue(len(dungeon.units), 0)

        # test exact match
        actual = dungeon.get_unit_by_name(unit_name)
        self.assertIsInstance(actual, Unit)

        # test invalid match
        actual = dungeon.get_unit_by_name("foo")
        self.assertIsNone(actual)

        # test case sensitivity
        actual = dungeon.get_unit_by_name(unit_name.upper())
        self.assertIsInstance(actual, Unit)

        # test startswith
        actual = dungeon.get_unit_by_name(unit_name[0:1])
        self.assertIsInstance(actual, Unit)

    def test_get_living_unit_by_name(self):
        dungeon_model = self._get_dungeon_model()
        with LogCapture():
            logger = logging.getLogger()
            dungeon = Dungeon(dungeon=dungeon_model,
                              announcer={},
                              log=logger)

        unit_name = "baz"
        unit_generator = UnitGenerator()
        unit = unit_generator.generate(unit_name=unit_name)

        dungeon.add_unit(unit)

        # test typical use case
        actual = dungeon.get_living_unit_by_name(unit.get_name())
        self.assertIsInstance(actual, Unit)

        unit.kill()

        actual = dungeon.get_living_unit_by_name(unit.get_name())
        self.assertIsNone(actual)

    def test_get_dead_unit_by_name(self):
        dungeon_model = self._get_dungeon_model()
        with LogCapture():
            logger = logging.getLogger()
            dungeon = Dungeon(dungeon=dungeon_model,
                              announcer={},
                              log=logger)

        unit_name = "Oliver Queen"
        unit_generator = UnitGenerator()
        unit = unit_generator.generate(unit_name=unit_name)

        dungeon.add_unit(unit)

        unit.kill()

        # test typical use case
        actual = dungeon.get_dead_unit_by_name(unit.get_name())
        self.assertIsInstance(actual, Unit)

        unit.hp = 1

        actual_alive = dungeon.get_living_unit_by_name(unit.get_name())
        self.assertIsInstance(actual_alive, Unit)

    def test_get_living_players(self):
        dungeon_model = self._get_dungeon_model()
        with LogCapture():
            logger = logging.getLogger()
            dungeon = Dungeon(dungeon=dungeon_model,
                              announcer={},
                              log=logger)

        unit_generator = UnitGenerator()

        unit = unit_generator.generate(is_player=True)

        dungeon.add_unit(unit)

        dead_unit = unit_generator.generate()
        dead_unit.kill()

        dungeon.add_unit(dead_unit)

        # test typical use case
        actual = dungeon.get_living_players()
        self.assertEqual(len(actual), 1)

        dead_players = dungeon.get_dead_players()
        self.assertEqual(len(dead_players), 0)

    def test_get_unit_status_distribution(self):
        dungeon_model = self._get_dungeon_model()
        with LogCapture():
            logger = logging.getLogger()
            dungeon = Dungeon(dungeon=dungeon_model,
                              announcer={},
                              log=logger)

        """ living hostile player """
        unit_generator = UnitGenerator()
        living_player = unit_generator.generate(is_player=True,
                                                combat_status="hostile")
        dungeon.add_unit(living_player)

        """ dead friendly NPC """
        dead_friendly_unit = unit_generator.generate(combat_status="friendly")
        dead_friendly_unit.kill()
        dungeon.add_unit(dead_friendly_unit)

        """ living friendly NPC """
        living_npc = unit_generator.generate(combat_status="friendly")
        dungeon.add_unit(living_npc)

        """ living hostile NPC """
        hostile_living_npc = unit_generator.generate(combat_status="hostile")
        dungeon.add_unit(hostile_living_npc)

        # TODO: add undead here
        actual = dungeon.get_unit_status_distribution()

        expected = {
            "players": {
                "living": 1,
                "dead": 0,
                "undead": 0
            },
            "npc": {
                "hostile": {
                    "living": 1,
                    "dead": 0,
                    "undead": 0
                },
                "friendly": {
                    "living": 1,
                    "dead": 1,
                    "undead": 0
                }
            }
        }

        self.assertEqual(expected, actual)

if __name__ == '__main__':
    unittest.main()
<|endoftext|>


<commit_before>'''
==============
pg_bawler.core
==============

Base classes for LISTEN / NOTIFY.  Postgresql documentation for
`LISTEN<https://www.postgresql.org/docs/current/static/sql-listen.html>`_ /
`NOTIFY<https://www.postgresql.org/docs/current/static/sql-notify.html>`_.
'''
import asyncio
import logging

import aiopg

LOGGER = logging.getLogger(name='pg_bawler.core')

def cache_async_def(func):
    cache_attr_name = '_cache_async_def_{func.__name__}'.format(func=func)
    async def _cache_method(self, *args, **kwargs):
        if not hasattr(self, cache_attr_name):
            setattr(self, cache_attr_name, await func(self, *args, **kwargs))
        return getattr(self, cache_attr_name)
    # simulate functools.update_wrapper
    _cache_method.__name__ = func.__name__
    _cache_method.__doc__ = func.__doc__
    _cache_method.__module__ = func.__module__
    # save cache_attr_name on function
    # so delattr(self, func.cache_attr_name) will clear the cache
    _cache_method.cache_attr_name = cache_attr_name
    return _cache_method

class BawlerBase:
    '''
    Base ``pg_bawler`` class with convenience methods around ``aiopg``.
    '''

    def __init__(self, *, connection_params, loop=None):
        self.connection_params = connection_params
        self._connection = None
        self.loop = asyncio.get_event_loop() if loop is None else loop

    @cache_async_def
    async def pg_pool(self):
        return await aiopg.create_pool(**self.connection_params)

    @cache_async_def
    async def pg_connection(self):
        return await (await self.pg_pool()).acquire()

    async def drop_connection(self):
        '''
        Drops current connection

        Next call to the ``self.pg_connection`` will acquire new connection
        from pool. Use this method to drop dead connections on server restart.
        '''
        pg_conn = (await self.pg_connection())
        pg_conn.close()
        await (await self.pg_pool()).release(pg_conn)
        # clear cached connection property (cache_async_def)
        delattr(self, self.pg_connection.cache_attr_name)

class SenderMixin:

    NOTIFY_SEND_TPL = 'SELECT pg_notify(\'{channel}\', \'{payload}\')'

    def get_notify_statement(self, *, channel, payload):
        return self.NOTIFY_SEND_TPL.format(channel=channel, payload=payload)

    async def send(self, *, channel, payload):
        async with (await self.pg_connection()).cursor() as pg_cursor:
            await pg_cursor.execute(
                self.get_notify_statement(channel=channel, payload=payload))

class ListenerMixin:

    CHANNEL_REGISTRATION_TPL = 'LISTEN {channel}'
    listen_timeout = None
    _stopped = False

    async def stop(self):
        await self.drop_connection()
        self._stopped = True

    @property
    def is_stopped(self):
        return self._stopped

    def _get_listen_statement(self, channel):
        return self.CHANNEL_REGISTRATION_TPL.format(channel=channel)

    async def register_channel(self, channel):
        async with (await self.pg_connection()).cursor() as cursor:
            await cursor.execute(self._get_listen_statement(channel))

    async def get_notification(self):
        try:
            notification = await asyncio.wait_for(
                (await self.pg_connection()).notifies.get(),
                self.listen_timeout)
        except asyncio.TimeoutError:
            LOGGER.debug(
                'Timed out. No notification for last %s seconds.',
                self.listen_timeout)
            LOGGER.debug(
                'Checking health of connection.')
            async with (await self.pg_connection()).cursor() as pg_cursor:
                await pg_cursor.execute('SELECT 1')
                is_healthy = await pg_cursor.fetchone() == (1, )
            if is_healthy:
                LOGGER.debug('Connection seems to be OK.')
            else:
                LOGGER.error('Failed postgres connection!')
                LOGGER.info('Dropping this connection and creating new one.')
                await self.drop_connection()
            return None
        else:
            LOGGER.debug(
                'Received notification from channel %s: %s',
                notification.channel, notification.payload)
            return notification

    def register_handler(self, handler):
        '''
        Registers ``handler`` with this listener and returns handler's index in
        the list of handlers.
        '''
        if not hasattr(self, 'handlers'):
            self.handlers = [handler]
        else:
            self.handlers.append(handler)
        return len(self.handlers) - 1

    def unregister_handler(self, handler):
        '''
        Unregisters ``handler`` by removing it from handlers list.

        Returns ``True`` if handler was successfully removed, ``False``
        otherwise.
        '''
        # TODO: Should we raise exception here?
        #       Like [].remove, raises ValueError
        try:
            self.handlers.remove(handler)
        except AttributeError:
            LOGGER.debug('No handlers registered yet.')
            return False
        except ValueError:
            LOGGER.debug('Handler is not registered.')
            return False
        else:
            LOGGER.debug('Handler %s unregistered.')
            return True

    async def listen(self):
        while not self.is_stopped:
            notification = await self.get_notification()
            if notification is not None:
                for handler in getattr(self, 'handlers', []):
                    self.loop.create_task(handler(notification))
<commit_msg>Make possible to register more channels with separate handlers<commit_after>'''
==============
pg_bawler.core
==============

Base classes for LISTEN / NOTIFY.  Postgresql documentation for
`LISTEN<https://www.postgresql.org/docs/current/static/sql-listen.html>`_ /
`NOTIFY<https://www.postgresql.org/docs/current/static/sql-notify.html>`_.
'''
import asyncio
import logging

import aiopg

LOGGER = logging.getLogger(name='pg_bawler.core')

def cache_async_def(func):
    cache_attr_name = '_cache_async_def_{func.__name__}'.format(func=func)
    async def _cache_method(self, *args, **kwargs):
        if not hasattr(self, cache_attr_name):
            setattr(self, cache_attr_name, await func(self, *args, **kwargs))
        return getattr(self, cache_attr_name)
    # simulate functools.update_wrapper
    _cache_method.__name__ = func.__name__
    _cache_method.__doc__ = func.__doc__
    _cache_method.__module__ = func.__module__
    # save cache_attr_name on function
    # so delattr(self, func.cache_attr_name) will clear the cache
    _cache_method.cache_attr_name = cache_attr_name
    return _cache_method

class BawlerBase:
    '''
    Base ``pg_bawler`` class with convenience methods around ``aiopg``.
    '''

    def __init__(self, *, connection_params, loop=None):
        self.connection_params = connection_params
        self._connection = None
        self.loop = asyncio.get_event_loop() if loop is None else loop

    @cache_async_def
    async def pg_pool(self):
        return await aiopg.create_pool(**self.connection_params)

    @cache_async_def
    async def pg_connection(self):
        return await (await self.pg_pool()).acquire()

    async def drop_connection(self):
        '''
        Drops current connection

        Next call to the ``self.pg_connection`` will acquire new connection
        from pool. Use this method to drop dead connections on server restart.
        '''
        pg_conn = (await self.pg_connection())
        pg_conn.close()
        await (await self.pg_pool()).release(pg_conn)
        # clear cached connection property (cache_async_def)
        delattr(self, self.pg_connection.cache_attr_name)

class SenderMixin:

    NOTIFY_SEND_TPL = 'SELECT pg_notify(\'{channel}\', \'{payload}\')'

    def get_notify_statement(self, *, channel, payload):
        return self.NOTIFY_SEND_TPL.format(channel=channel, payload=payload)

    async def send(self, *, channel, payload):
        async with (await self.pg_connection()).cursor() as pg_cursor:
            await pg_cursor.execute(
                self.get_notify_statement(channel=channel, payload=payload))

class ListenerMixin:

    CHANNEL_REGISTRATION_TPL = 'LISTEN {channel}'
    listen_timeout = None
    stop_on_timeout = False
    _stopped = False

    async def stop(self):
        await self.drop_connection()
        self._stopped = True

    @property
    def is_stopped(self):
        return self._stopped

    @property
    def registered_channels(self):
        prop_name = '_registered_channels'
        if not hasattr(self, prop_name):
            setattr(self, prop_name, {})
        return getattr(self, prop_name)

    def _get_listen_statement(self, channel):
        return self.CHANNEL_REGISTRATION_TPL.format(channel=channel)

    async def register_channel(self, channel):
        async with (await self.pg_connection()).cursor() as cursor:
            await cursor.execute(self._get_listen_statement(channel))
        self.registered_channels.setdefault(channel, [])

    async def get_notification(self):
        try:
            notification = await asyncio.wait_for(
                (await self.pg_connection()).notifies.get(),
                self.listen_timeout)
        except asyncio.TimeoutError:
            LOGGER.debug(
                'Timed out. No notification for last %s seconds.',
                self.listen_timeout)
            if self.stop_on_timeout:
                await self.stop()
            else:
                LOGGER.debug(
                    'Checking health of connection.')
                async with (await self.pg_connection()).cursor() as pg_cursor:
                    await pg_cursor.execute('SELECT 1')
                    is_healthy = await pg_cursor.fetchone() == (1, )
                if is_healthy:
                    LOGGER.debug('Connection seems to be OK.')
                else:
                    LOGGER.error('Failed postgres connection!')
                    LOGGER.info(
                        'Dropping this connection and creating new one.')
                    await self.drop_connection()
            return None
        else:
            LOGGER.debug(
                'Received notification from channel %s: %s',
                notification.channel, notification.payload)
            return notification

    def register_handler(self, channel, handler):
        '''
        Registers ``handler`` with given ``channel``

        :param channel: Name of channel
        :param handler: Coroutine that will handle notifications from
            ``channel``
        :returns: None
        '''
        if channel in self.registered_channels:
            self.registered_channels.append(handler)
        else:
            self.registered_channels[channel] = [handler]

    def unregister_handler(self, channel, handler):
        '''
        Unregisters ``handler`` by removing it from handlers list of given
        ``channel``.

        :param channel: Name of channel
        :param handler: Coroutine to unregister from ``channel``
        :returns: None
        '''
        if channel in self.registered_channels:
            try:
                self.registered_channels[channel].remove(handler)
            except ValueError:
                LOGGER.debug('Handler is not registered.')
            else:
                LOGGER.debug('Handler %s unregistered.')
        return None

    async def listen(self):
        while not self.is_stopped:
            notification = await self.get_notification()
            if notification is not None:
                handlers = self.registered_channels[notification.channel]
                for handler in handlers:
                    self.loop.create_task(handler(notification, self))
<|endoftext|>