<commit_before>import logging
import cherrypy
import time
import json
import pickle
from cherrypy.process.plugins import Monitor
from sqlalchemy.orm import deferred
from sqlalchemy.exc import IntegrityError
from sqlalchemy import Column, Integer, String, BLOB, BigInteger, func
from sqlalchemy.dialects import mysql
from sqlalchemy.orm.exc import NoResultFound
from opmuse.database import Base, get_database, get_session

def debug(msg):
    cherrypy.log.error(msg, context='cache', severity=logging.DEBUG)

def log(msg):
    cherrypy.log(msg, context='cache')

class CacheObject(Base):
    __tablename__ = 'cache_objects'

    id = Column(Integer, primary_key=True)
    key = Column(String(255), index=True, unique=True)
    type = Column(String(128))
    updated = Column(BigInteger, index=True)
    value = deferred(Column(BLOB().with_variant(mysql.LONGBLOB(), 'mysql')))

class CacheStorage:
    def __init__(self):
        self._storage = {}

    def set(self, key, updated, value):
        self._storage[key] = {
            'value': value,
            'updated': updated
        }

    def get_updated(self, key):
        if key in self._storage:
            return self._storage[key]['updated']

    def get_value(self, key):
        if key in self._storage:
            return self._storage[key]['value']

    def has(self, key):
        return key in self._storage

    def values(self):
        return self._storage.items()

class Keep:
    pass

class Cache:
    KEEP_AGE = 30 * 60
    """
        For Keep values we ignore the provided age and use this instead if the
        provided age is larger than this value. The reason is that Keep values
        are only used to mark the key as used until the data is actually set
        (e.g. for slow bgtasks and such) so if that fails the Keep value is
        invalid and it's "safe" to try again.
    """

    def __init__(self):
        self.storage = CacheStorage()

    def needs_update(self, key, age):
        now = int(time.time())

        if self.storage.has(key):
            value = self.storage.get_value(key)
            updated = self.storage.get_updated(key)

            if value is Keep and age > Cache.KEEP_AGE:
                _age = Cache.KEEP_AGE
            else:
                _age = age

            if now - updated > _age:
                return True
        else:
            return True

        return False

    def get(self, key):
        value = self.storage.get_value(key)

        if value is Keep:
            return None
        else:
            return value

    def keep(self, key):
        """
            Updates the timestamp of the objects and creates it if it doesn't exist
            but keeps the value if there is one.
        """

        self.set(key, Keep)

    def set(self, key, value):
        if value is not None and value is not Keep and not isinstance(value, (str, bytes, dict, list, object)):
            raise ValueError("Unsupported value type.")

        updated = int(time.time())

        self.storage.set(key, updated, value)

class CachePlugin(Monitor):
    GC_AGE = 6 * 30 * 24 * 3600
    """
        How old an entry has to be considered stale and removed.
    """

    def __init__(self, bus):
        Monitor.__init__(self, bus, self.run, frequency = 20 * 60)

    def run(self):
        self._serialize()

    def start(self):
        self._gc(CachePlugin.GC_AGE)
        self._unserialize()

        Monitor.start(self)

    start.priority = 105

    def stop(self):
        self._serialize()

        Monitor.stop(self)

    stop.priority = 105

    def _gc(self, age):
        database = get_session()

        now = int(time.time())

        updated_filter = "(%d - updated) > %d" % (now, age)

        item_count = database.query(func.count(CacheObject.id)).filter(updated_filter).scalar()

        database.query(CacheObject).filter(updated_filter).delete(synchronize_session='fetch')

        database.commit()
        database.remove()

        log("Garbage collected %d cache objects" % item_count)

    def _serialize(self):
        # TODO actually use "updated" var to see if we even need to serialize/persist the changes...

        database = get_session()

        item_count = 0

        for key, item in cache.storage.values():
            item_count += 1

            value = item['value']
            updated = item['updated']

            # we don't serialize Keep values, in case of restart they'll just try again...
            if value is Keep:
                continue

            count = (database.query(func.count(CacheObject.id))
                     .filter(CacheObject.key == key).scalar())

            orig_value = value

            if isinstance(value, object):
                value_type = 'object'
            else:
                value_type = type(value).__name__

            if value_type == 'object':
                value = pickle.dumps(value)
            elif value_type == 'str':
                value = value.encode()
            elif value_type == 'dict' or value_type == 'list':
                value = json.dumps(value).encode()

            if count > 0:
                parameters = {'updated': updated}

                parameters['value'] = value
                parameters['type'] = value_type

                database.query(CacheObject).filter(CacheObject.key == key).update(parameters)
            else:
                parameters = {'key': key, 'value': value, 'updated': updated, 'type': value_type}
                database.execute(CacheObject.__table__.insert(), parameters)

        database.commit()
        database.remove()

        log("Serialized %d cache objects" % item_count)

    def _unserialize(self):
        database = get_session()

        item_count = 0

        for object in database.query(CacheObject).all():
            item_count += 1

            if object.type == 'object':
                value = pickle.loads(object.value)
            elif object.type == 'str':
                value = object.value.decode()
            elif object.type == 'dict' or object.type == 'list':
                value = json.loads(object.value.decode())
            else:
                value = object.value

            cache.storage.set(object.key, object.updated, value)

        database.remove()

        log("Unserialized %d cache objects" % item_count)

cache = Cache()
<commit_msg>persist Keep values because we want the updated timestamp<commit_after>import logging
import cherrypy
import time
import json
import pickle
from cherrypy.process.plugins import Monitor
from sqlalchemy.orm import deferred
from sqlalchemy.exc import IntegrityError
from sqlalchemy import Column, Integer, String, BLOB, BigInteger, func
from sqlalchemy.dialects import mysql
from sqlalchemy.orm.exc import NoResultFound
from opmuse.database import Base, get_database, get_session

def debug(msg):
    cherrypy.log.error(msg, context='cache', severity=logging.DEBUG)

def log(msg):
    cherrypy.log(msg, context='cache')

class CacheObject(Base):
    __tablename__ = 'cache_objects'

    id = Column(Integer, primary_key=True)
    key = Column(String(255), index=True, unique=True)
    type = Column(String(128))
    updated = Column(BigInteger, index=True)
    value = deferred(Column(BLOB().with_variant(mysql.LONGBLOB(), 'mysql')))

class CacheStorage:
    def __init__(self):
        self._storage = {}

    def set(self, key, updated, value):
        self._storage[key] = {
            'value': value,
            'updated': updated
        }

    def get_updated(self, key):
        if key in self._storage:
            return self._storage[key]['updated']

    def get_value(self, key):
        if key in self._storage:
            return self._storage[key]['value']

    def has(self, key):
        return key in self._storage

    def values(self):
        return self._storage.items()

class Keep:
    pass

class Cache:
    KEEP_AGE = 30 * 60
    """
        For Keep values we ignore the provided age and use this instead if the
        provided age is larger than this value. The reason is that Keep values
        are only used to mark the key as used until the data is actually set
        (e.g. for slow bgtasks and such) so if that fails the Keep value is
        invalid and it's "safe" to try again.
    """

    def __init__(self):
        self.storage = CacheStorage()

    def needs_update(self, key, age):
        now = int(time.time())

        if self.storage.has(key):
            value = self.storage.get_value(key)
            updated = self.storage.get_updated(key)

            if value is Keep and age > Cache.KEEP_AGE:
                _age = Cache.KEEP_AGE
            else:
                _age = age

            if now - updated > _age:
                return True
        else:
            return True

        return False

    def get(self, key):
        value = self.storage.get_value(key)

        if value is Keep:
            return None
        else:
            return value

    def keep(self, key):
        """
            Updates the timestamp of the objects and creates it if it doesn't exist
            but keeps the value if there is one.
        """

        self.set(key, Keep)

    def set(self, key, value):
        if value is not None and value is not Keep and not isinstance(value, (str, bytes, dict, list, object)):
            raise ValueError("Unsupported value type.")

        updated = int(time.time())

        self.storage.set(key, updated, value)

class CachePlugin(Monitor):
    GC_AGE = 6 * 30 * 24 * 3600
    """
        How old an entry has to be considered stale and removed.
    """

    def __init__(self, bus):
        Monitor.__init__(self, bus, self.run, frequency = 20 * 60)

    def run(self):
        self._serialize()

    def start(self):
        self._gc(CachePlugin.GC_AGE)
        self._unserialize()

        Monitor.start(self)

    start.priority = 105

    def stop(self):
        self._serialize()

        Monitor.stop(self)

    stop.priority = 105

    def _gc(self, age):
        database = get_session()

        now = int(time.time())

        updated_filter = "(%d - updated) > %d" % (now, age)

        item_count = database.query(func.count(CacheObject.id)).filter(updated_filter).scalar()

        database.query(CacheObject).filter(updated_filter).delete(synchronize_session='fetch')

        database.commit()
        database.remove()

        log("Garbage collected %d cache objects" % item_count)

    def _serialize(self):
        # TODO actually use "updated" var to see if we even need to serialize/persist the changes...

        database = get_session()

        item_count = 0

        for key, item in cache.storage.values():
            item_count += 1

            value = item['value']
            updated = item['updated']

            count = (database.query(func.count(CacheObject.id))
                     .filter(CacheObject.key == key).scalar())

            orig_value = value

            if isinstance(value, object):
                value_type = 'object'
            else:
                value_type = type(value).__name__

            if value_type == 'object':
                value = pickle.dumps(value)
            elif value_type == 'str':
                value = value.encode()
            elif value_type == 'dict' or value_type == 'list':
                value = json.dumps(value).encode()

            if count > 0:
                parameters = {'updated': updated}

                parameters['value'] = value
                parameters['type'] = value_type

                database.query(CacheObject).filter(CacheObject.key == key).update(parameters)
            else:
                parameters = {'key': key, 'value': value, 'updated': updated, 'type': value_type}
                database.execute(CacheObject.__table__.insert(), parameters)

        database.commit()
        database.remove()

        log("Serialized %d cache objects" % item_count)

    def _unserialize(self):
        database = get_session()

        item_count = 0

        for object in database.query(CacheObject).all():
            item_count += 1

            if object.type == 'object':
                value = pickle.loads(object.value)
            elif object.type == 'str':
                value = object.value.decode()
            elif object.type == 'dict' or object.type == 'list':
                value = json.loads(object.value.decode())
            else:
                value = object.value

            cache.storage.set(object.key, object.updated, value)

        database.remove()

        log("Unserialized %d cache objects" % item_count)

cache = Cache()
<|endoftext|>


<commit_before># vim:sw=4:et:
"""This module contains some helpers that can be used to execute generator
functions in the GObject main loop.

This module provided the following classes:
GIdleThread - Thread like behavior for generators in a main loop
Queue - A simple queue implementation suitable for use with GIdleThread

Exceptions:
QueueEmpty - raised when one tried to get a value of an empty queue
QueueFull - raised when the queue reaches it's max size and the oldest item
            may not be disposed.
"""

import sys
import gobject
import time
import traceback

class GIdleThread(object):
    """This is a pseudo-"thread" for use with the GTK+ main loop.

    This class does act a bit like a thread, all code is executed in
    the callers thread though. The provided function should be a generator
    (or iterator).

    It can be started with start(). While the "thread" is running is_alive()
    can be called to see if it's alive. wait([timeout]) will wait till the
    generator is finished, or timeout seconds.

    If an exception is raised from within the generator, it is stored in
    the exc_info property. Execution of the generator is finished. The
    exc_info property contains a tuple (exc_type, exc_value, exc_traceback),
    see sys.exc_info() for details.

    Note that this routine runs in the current thread, so there is no need
    for nasty locking schemes.

    Example (runs a counter through the GLib main loop routine):
    >>> def counter(max):
    ...     for x in xrange(max):
    ...         yield x
    >>> t = GIdleThread(counter(123))
    >>> t.start()
    >>> while gen.is_alive():
    ...     main.iteration(False)
    """

    def __init__(self, generator, queue=None):
        assert hasattr(generator, 'next'), 'The generator should be an iterator'
        self._generator = generator
        self._queue = queue
        self._idle_id = 0
        self._error = None
        self._exc_info = (None, None, None)

    def start(self, priority=gobject.PRIORITY_LOW):
        """Start the generator. Default priority is low, so screen updates
        will be allowed to happen.
        """
        idle_id = gobject.idle_add(self.__generator_executer,
                                   priority=priority)
        self._idle_id = idle_id
        return idle_id

    def wait(self, timeout=0):
        """Wait until the corouine is finished or return after timeout seconds.
        This is achieved by running the GTK+ main loop.
        """
        clock = time.clock
        start_time = clock()
        main = gobject.main_context_default()
        while self.is_alive():
            main.iteration(False)
            if timeout and (clock() - start_time >= timeout):
                return

    def interrupt(self):
        """Force the generator to stop running.
        """
        if self.is_alive():
            gobject.source_remove(self._idle_id)
            self._idle_id = 0

    def is_alive(self):
        """Returns True if the generator is still running.
        """
        return self._idle_id != 0

    error = property(lambda self: self._exc_info[0],
                     doc="Return a possible exception that had occured "\
                         "during execution of the generator")

    exc_info = property(lambda self: self._exc_info,
                     doc="Return a exception information as provided by "\
                         "sys.exc_info()")

    def __generator_executer(self):
        try:
            result = self._generator.next()
            if self._queue:
                try:
                    self._queue.put(result)
                except QueueFull:
                    self.wait(0.5)
                    # If this doesn't work...
                    self._queue.put(result)
            return True
        except StopIteration:
            self._idle_id = 0
            return False
        except:
            self._exc_info = sys.exc_info()
            #traceback.print_exc()
            self._idle_id = 0
            return False

class QueueEmpty(Exception):
    """Exception raised whenever the queue is empty and someone tries to fetch
    a value.
    """
    pass

class QueueFull(Exception):
    """Exception raised when the queue is full and the oldest item may not be
    disposed.
    """
    pass

class Queue(object):
    """A FIFO queue. If the queue has a max size, the oldest item on the
    queue is dropped if that size id exceeded.
    """

    def __init__(self, size=0, dispose_oldest=True):
        self._queue = []
        self._size = size
        self._dispose_oldest = dispose_oldest

    def put(self, item):
        """Put item on the queue. If the queue size is limited ...
        """
        if self._size > 0 and len(self._queue) >= self._size:
            if self._dispose_oldest:
                self.get()
            else:
                raise QueueFull

        self._queue.insert(0, item)

    def get(self):
        """Get the oldest item off the queue.
        QueueEmpty is raised if no items are left on the queue.
        """
        try:
            return self._queue.pop()
        except IndexError:
            raise QueueEmpty

if __name__ == '__main__':
    def counter(max):
        for i in range(max):
            yield i

    def shower(queue):
        # Never stop reading the queue:
        while True:
            try:
                cnt = queue.get()
                print 'cnt =', cnt
            except QueueEmpty:
                pass
            yield None

    print 'Test 1: (should print range 0..22)'
    queue = Queue()
    c = GIdleThread(counter(23), queue)
    s = GIdleThread(shower(queue))

    main = gobject.main_context_default()
    c.start()
    s.start()
    s.wait(2)

    print 'Test 2: (should only print 22)'
    queue = Queue(size=1)
    c = GIdleThread(counter(23), queue)
    s = GIdleThread(shower(queue))

    main = gobject.main_context_default()
    c.start(priority=gobject.PRIORITY_DEFAULT)
    s.start()
    s.wait(3)
<commit_msg>- typo fixed<commit_after># vim:sw=4:et:
"""This module contains some helpers that can be used to execute generator
functions in the GObject main loop.

This module provided the following classes:
GIdleThread - Thread like behavior for generators in a main loop
Queue - A simple queue implementation suitable for use with GIdleThread

Exceptions:
QueueEmpty - raised when one tried to get a value of an empty queue
QueueFull - raised when the queue reaches it's max size and the oldest item
            may not be disposed.
"""

import sys
import gobject
import time
import traceback

class GIdleThread(object):
    """This is a pseudo-"thread" for use with the GTK+ main loop.

    This class does act a bit like a thread, all code is executed in
    the callers thread though. The provided function should be a generator
    (or iterator).

    It can be started with start(). While the "thread" is running is_alive()
    can be called to see if it's alive. wait([timeout]) will wait till the
    generator is finished, or timeout seconds.

    If an exception is raised from within the generator, it is stored in
    the exc_info property. Execution of the generator is finished. The
    exc_info property contains a tuple (exc_type, exc_value, exc_traceback),
    see sys.exc_info() for details.

    Note that this routine runs in the current thread, so there is no need
    for nasty locking schemes.

    Example (runs a counter through the GLib main loop routine):
    >>> def counter(max):
    ...     for x in xrange(max):
    ...         yield x
    >>> t = GIdleThread(counter(123))
    >>> t.start()
    >>> while t.is_alive():
    ...     main.iteration(False)
    """

    def __init__(self, generator, queue=None):
        assert hasattr(generator, 'next'), 'The generator should be an iterator'
        self._generator = generator
        self._queue = queue
        self._idle_id = 0
        self._error = None
        self._exc_info = (None, None, None)

    def start(self, priority=gobject.PRIORITY_LOW):
        """Start the generator. Default priority is low, so screen updates
        will be allowed to happen.
        """
        idle_id = gobject.idle_add(self.__generator_executer,
                                   priority=priority)
        self._idle_id = idle_id
        return idle_id

    def wait(self, timeout=0):
        """Wait until the corouine is finished or return after timeout seconds.
        This is achieved by running the GTK+ main loop.
        """
        clock = time.clock
        start_time = clock()
        main = gobject.main_context_default()
        while self.is_alive():
            main.iteration(False)
            if timeout and (clock() - start_time >= timeout):
                return

    def interrupt(self):
        """Force the generator to stop running.
        """
        if self.is_alive():
            gobject.source_remove(self._idle_id)
            self._idle_id = 0

    def is_alive(self):
        """Returns True if the generator is still running.
        """
        return self._idle_id != 0

    error = property(lambda self: self._exc_info[0],
                     doc="Return a possible exception that had occured "\
                         "during execution of the generator")

    exc_info = property(lambda self: self._exc_info,
                     doc="Return a exception information as provided by "\
                         "sys.exc_info()")

    def __generator_executer(self):
        try:
            result = self._generator.next()
            if self._queue:
                try:
                    self._queue.put(result)
                except QueueFull:
                    self.wait(0.5)
                    # If this doesn't work...
                    self._queue.put(result)
            return True
        except StopIteration:
            self._idle_id = 0
            return False
        except:
            self._exc_info = sys.exc_info()
            #traceback.print_exc()
            self._idle_id = 0
            return False

class QueueEmpty(Exception):
    """Exception raised whenever the queue is empty and someone tries to fetch
    a value.
    """
    pass

class QueueFull(Exception):
    """Exception raised when the queue is full and the oldest item may not be
    disposed.
    """
    pass

class Queue(object):
    """A FIFO queue. If the queue has a max size, the oldest item on the
    queue is dropped if that size id exceeded.
    """

    def __init__(self, size=0, dispose_oldest=True):
        self._queue = []
        self._size = size
        self._dispose_oldest = dispose_oldest

    def put(self, item):
        """Put item on the queue. If the queue size is limited ...
        """
        if self._size > 0 and len(self._queue) >= self._size:
            if self._dispose_oldest:
                self.get()
            else:
                raise QueueFull

        self._queue.insert(0, item)

    def get(self):
        """Get the oldest item off the queue.
        QueueEmpty is raised if no items are left on the queue.
        """
        try:
            return self._queue.pop()
        except IndexError:
            raise QueueEmpty

if __name__ == '__main__':
    def counter(max):
        for i in range(max):
            yield i

    def shower(queue):
        # Never stop reading the queue:
        while True:
            try:
                cnt = queue.get()
                print 'cnt =', cnt
            except QueueEmpty:
                pass
            yield None

    print 'Test 1: (should print range 0..22)'
    queue = Queue()
    c = GIdleThread(counter(23), queue)
    s = GIdleThread(shower(queue))

    main = gobject.main_context_default()
    c.start()
    s.start()
    s.wait(2)

    print 'Test 2: (should only print 22)'
    queue = Queue(size=1)
    c = GIdleThread(counter(23), queue)
    s = GIdleThread(shower(queue))

    main = gobject.main_context_default()
    c.start(priority=gobject.PRIORITY_DEFAULT)
    s.start()
    s.wait(3)
<|endoftext|>


<commit_before>"""Provides interfaces to various commands provided by FreeSurfer

   Change directory to provide relative paths for doctests
   >>> import os
   >>> filepath = os.path.dirname( os.path.realpath( __file__ ) )
   >>> datadir = os.path.realpath(os.path.join(filepath, '../testing/data'))
   >>> os.chdir(datadir)

"""

from __future__ import absolute_import
import os, string
from os import path
from glob import glob
from nipype.interfaces.base import (TraitedSpec,
                                    DynamicTraitedSpec,
                                    InputMultiPath,
                                    File,
                                    Directory,
                                    traits,
                                    BaseInterface,
                                   )
import nibabel as nb
from nipype.interfaces.traits_extension import isdefined, Undefined
import imghdr
from nipype.external import six

have_dcmstack = True
try:
    import dicom
    import dcmstack
    from dcmstack.dcmmeta import NiftiWrapper
except ImportError:
    have_dcmstack = False

def sanitize_path_comp(path_comp):
    result = []
    for char in path_comp:
        if not char in string.letters + string.digits + '-_.':
            result.append('_')
        else:
            result.append(char)
    return ''.join(result)

class NiftiGeneratorBaseInputSpec(TraitedSpec):
    out_format = traits.Str(desc="String which can be formatted with "
                            "meta data to create the output filename(s)")
    out_ext = traits.Str('.nii.gz', usedefault=True,
                         desc="Determines output file type")
    out_path = Directory(desc='output path, current working directory if not set')

class NiftiGeneratorBase(BaseInterface):
    '''Base class for interfaces that produce Nifti files, potentially with
    embeded meta data.'''
    def _get_out_path(self, meta, idx=None):
        '''Return the output path for the gernerated Nifti.'''
        if self.inputs.out_format:
            out_fmt = self.inputs.out_format
        else:
            #If no out_format is specified, use a sane default that will work
            #with the provided meta data.
            out_fmt = []
            if not idx is None:
                out_fmt.append('%03d' % idx)
            if 'SeriesNumber' in meta:
                out_fmt.append('%(SeriesNumber)03d')
            if 'ProtocolName' in meta:
                out_fmt.append('%(ProtocolName)s')
            elif 'SeriesDescription' in meta:
                out_fmt.append('%(SeriesDescription)s')
            else:
                out_fmt.append('sequence')
            out_fmt = '-'.join(out_fmt)
        out_fn = (out_fmt % meta) + self.inputs.out_ext
        out_fn = sanitize_path_comp(out_fn)

        out_path = os.getcwd()
        if isdefined(self.inputs.out_path):
            out_path = path.abspath(self.inputs.out_path)
        return path.join(out_path, out_fn)

class DcmStackInputSpec(NiftiGeneratorBaseInputSpec):
    dicom_files = traits.Either(InputMultiPath(File(exists=True)),
                                Directory(exists=True),
                                traits.Str(),
                                mandatory=True)
    embed_meta = traits.Bool(desc="Embed DICOM meta data into result")
    exclude_regexes = traits.List(desc="Meta data to exclude, suplementing "
                                  "any default exclude filters")
    include_regexes = traits.List(desc="Meta data to include, overriding any "
                                  "exclude filters")
    force_read = traits.Bool(True, usedefault=True,
                             desc=('Force reading files without DICM marker'))

class DcmStackOutputSpec(TraitedSpec):
    out_file = File(exists=True)

class DcmStack(NiftiGeneratorBase):
    '''Create one Nifti file from a set of DICOM files. Can optionally embed
    meta data.

    Example
    -------

    >>> from nipype.interfaces.dcmstack import DcmStack
    >>> stacker = DcmStack()
    >>> stacker.inputs.dicom_files = 'path/to/series/'
    >>> stacker.run() # doctest: +SKIP
    >>> result.outputs.out_file # doctest: +SKIP
    '/path/to/cwd/sequence.nii.gz'
    '''
    input_spec = DcmStackInputSpec
    output_spec = DcmStackOutputSpec

    def _get_filelist(self, trait_input):
        if isinstance(trait_input, six.string_types):
            if path.isdir(trait_input):
                return glob(path.join(trait_input, '*.dcm'))
            else:
                return glob(trait_input)

        return trait_input

    def _run_interface(self, runtime):
        src_paths = self._get_filelist(self.inputs.dicom_files)
        include_regexes = dcmstack.default_key_incl_res
        if isdefined(self.inputs.include_regexes):
            include_regexes += self.inputs.include_regexes
        exclude_regexes = dcmstack.default_key_excl_res
        if isdefined(self.inputs.exclude_regexes):
            exclude_regexes += self.inputs.exclude_regexes
        meta_filter = dcmstack.make_key_regex_filter(exclude_regexes,
                                                     include_regexes)
        stack = dcmstack.DicomStack(meta_filter=meta_filter)
        for src_path in src_paths:
            if not imghdr.what(src_path)=="gif":
                src_dcm = dicom.read_file(src_path, force=self.inputs.force_read)
                stack.add_dcm(src_dcm)
        nii = stack.to_nifti(embed_meta=True)
        nw = NiftiWrapper(nii)
        self.out_path = \
            self._get_out_path(nw.meta_ext.get_class_dict(('global', 'const')))
        if not self.inputs.embed_meta:
            nw.remove_extension()
        nb.save(nii, self.out_path)
        return runtime

    def _list_outputs(self):
        outputs = self._outputs().get()
        outputs["out_file"] = self.out_path
        return outputs

class GroupAndStackOutputSpec(TraitedSpec):
    out_list = traits.List(desc="List of output nifti files")

class GroupAndStack(DcmStack):
    '''Create (potentially) multiple Nifti files for a set of DICOM files.
    '''
    input_spec = DcmStackInputSpec
    output_spec = GroupAndStackOutputSpec

    def _run_interface(self, runtime):
        src_paths = self._get_filelist(self.inputs.dicom_files)
        stacks = dcmstack.parse_and_stack(src_paths)

        self.out_list = []
        for key, stack in stacks.iteritems():
            nw = NiftiWrapper(stack.to_nifti(embed_meta=True))
            const_meta = nw.meta_ext.get_class_dict(('global', 'const'))
            out_path =  self._get_out_path(const_meta)
            if not self.inputs.embed_meta:
                nw.remove_extension()
            nb.save(nw.nii_img, out_path)
            self.out_list.append(out_path)

        return runtime

    def _list_outputs(self):
        outputs = self._outputs().get()
        outputs["out_list"] = self.out_list
        return outputs

class LookupMetaInputSpec(TraitedSpec):
    in_file = File(mandatory=True,
                   exists=True,
                   desc='The input Nifti file')
    meta_keys = traits.Either(traits.List(),
                              traits.Dict(),
                              mandatory=True,
                              desc=("List of meta data keys to lookup, or a "
                              "dict where keys specify the meta data keys to "
                              "lookup and the values specify the output names")
                             )

class LookupMeta(BaseInterface):
    '''Lookup meta data values from a Nifti with embeded meta data.

    Example
    -------

    >>> from nipype.interfaces import dcmstack
    >>> lookup = dcmstack.LookupMeta()
    >>> lookup.inputs.in_file = 'functional.nii'
    >>> lookup.inputs.meta_keys = {'RepetitionTime' : 'TR', \
                                   'EchoTime' : 'TE'}
    >>> result = lookup.run() # doctest: +SKIP
    >>> result.outputs.TR # doctest: +SKIP
    9500.0
    >>> result.outputs.TE # doctest: +SKIP
    95.0
    '''
    input_spec = LookupMetaInputSpec
    output_spec = DynamicTraitedSpec

    def _make_name_map(self):
        if isinstance(self.inputs.meta_keys, list):
            self._meta_keys = {}
            for key in self.inputs.meta_keys:
                self._meta_keys[key] = key
        else:
            self._meta_keys = self.inputs.meta_keys

    def _outputs(self):
        self._make_name_map()
        outputs = super(LookupMeta, self)._outputs()
        undefined_traits = {}
        for out_name in self._meta_keys.values():
            outputs.add_trait(out_name, traits.Any)
            undefined_traits[out_name] = Undefined
        outputs.trait_set(trait_change_notify=False, **undefined_traits)
        #Not sure why this is needed
        for out_name in self._meta_keys.values():
            _ = getattr(outputs, out_name)
        return outputs

    def _run_interface(self, runtime):
        #If the 'meta_keys' input is a list, covert it to a dict
        self._make_name_map()
        nw = NiftiWrapper.from_filename(self.inputs.in_file)
        self.result = {}
        for meta_key, out_name in self._meta_keys.iteritems():
            self.result[out_name] = nw.meta_ext.get_values(meta_key)

        return runtime

    def _list_outputs(self):
        outputs = self._outputs().get()
        outputs.update(self.result)
        return outputs

class CopyMetaInputSpec(TraitedSpec):
    src_file = File(mandatory=True, exists=True)
    dest_file = File(mandatory=True, exists=True)
    include_classes = traits.List(desc="List of specific meta data "
                                  "classifications to include. If not "
                                  "specified include everything.")
    exclude_classes = traits.List(desc="List of meta data "
                                  "classifications to exclude")

class CopyMetaOutputSpec(TraitedSpec):
    dest_file = File(exists=True)

class CopyMeta(BaseInterface):
    '''Copy meta data from one Nifti file to another. Useful for preserving
    meta data after some processing steps.'''
    input_spec = CopyMetaInputSpec
    output_spec = CopyMetaOutputSpec

    def _run_interface(self, runtime):
        src_nii = nb.load(self.inputs.src_file)
        src = NiftiWrapper(src_nii, make_empty=True)
        dest_nii = nb.load(self.inputs.dest_file)
        dest = NiftiWrapper(dest_nii, make_empty=True)
        classes = src.meta_ext.get_valid_classes()
        if self.inputs.include_classes:
            classes = [cls
                       for cls in classes
                       if cls in self.inputs.include_classes
                      ]
        if self.inputs.exclude_classes:
            classes = [cls
                       for cls in classes
                       if not cls in self.inputs.exclude_classes
                      ]

        for cls in classes:
            src_dict = src.meta_ext.get_class_dict(cls)
            dest_dict = dest.meta_ext.get_class_dict(cls)
            dest_dict.update(src_dict)
        # Update the shape and slice dimension to reflect the meta extension update.
        dest.meta_ext.slice_dim = src.meta_ext.slice_dim
        dest.meta_ext.shape = src.meta_ext.shape

        self.out_path = path.join(os.getcwd(),
                                  path.basename(self.inputs.dest_file))
        dest.to_filename(self.out_path)

        return runtime

    def _list_outputs(self):
        outputs = self._outputs().get()
        outputs['dest_file'] = self.out_path
        return outputs

class MergeNiftiInputSpec(NiftiGeneratorBaseInputSpec):
    in_files = traits.List(mandatory=True,
                           desc="List of Nifti files to merge")
    sort_order = traits.Either(traits.Str(),
                               traits.List(),
                               desc="One or more meta data keys to "
                               "sort files by.")
    merge_dim = traits.Int(desc="Dimension to merge along. If not "
                           "specified, the last singular or "
                           "non-existant dimension is used.")

class MergeNiftiOutputSpec(TraitedSpec):
    out_file = File(exists=True, desc="Merged Nifti file")

def make_key_func(meta_keys, index=None):
    def key_func(src_nii):
        result = [src_nii.get_meta(key, index) for key in meta_keys]
        return result

    return key_func

class MergeNifti(NiftiGeneratorBase):
    '''Merge multiple Nifti files into one. Merges together meta data
    extensions as well.'''
    input_spec = MergeNiftiInputSpec
    output_spec = MergeNiftiOutputSpec

    def _run_interface(self, runtime):
        niis = [nb.load(fn)
                for fn in self.inputs.in_files
               ]
        nws = [NiftiWrapper(nii, make_empty=True)
               for nii in niis
              ]
        if self.inputs.sort_order:
            sort_order = self.inputs.sort_order
            if isinstance(sort_order, six.string_types):
                sort_order = [sort_order]
            nws.sort(key=make_key_func(sort_order))
        if self.inputs.merge_dim == traits.Undefined:
            merge_dim = None
        else:
            merge_dim = self.inputs.merge_dim
        merged = NiftiWrapper.from_sequence(nws, merge_dim)
        const_meta = merged.meta_ext.get_class_dict(('global', 'const'))
        self.out_path = self._get_out_path(const_meta)
        nb.save(merged.nii_img, self.out_path)
        return runtime

    def _list_outputs(self):
        outputs = self._outputs().get()
        outputs['out_file'] = self.out_path
        return outputs

class SplitNiftiInputSpec(NiftiGeneratorBaseInputSpec):
    in_file = File(exists=True, mandatory=True, desc="Nifti file to split")
    split_dim = traits.Int(desc="Dimension to split along. If not "
                           "specified, the last dimension is used.")

class SplitNiftiOutputSpec(TraitedSpec):
    out_list = traits.List(File(exists=True),
                           desc="Split Nifti files")

class SplitNifti(NiftiGeneratorBase):
    '''Split one Nifti file into many along the specified dimension. Each
    result has an updated meta data extension as well.'''
    input_spec = SplitNiftiInputSpec
    output_spec = SplitNiftiOutputSpec

    def _run_interface(self, runtime):
        self.out_list = []
        nii = nb.load(self.inputs.in_file)
        nw = NiftiWrapper(nii, make_empty=True)
        split_dim = None
        if self.inputs.split_dim == traits.Undefined:
            split_dim = None
        else:
            split_dim = self.inputs.split_dim
        for split_idx, split_nw in enumerate(nw.split(split_dim)):
            const_meta = split_nw.meta_ext.get_class_dict(('global', 'const'))
            out_path = self._get_out_path(const_meta, idx=split_idx)
            nb.save(split_nw.nii_img, out_path)
            self.out_list.append(out_path)

        return runtime

    def _list_outputs(self):
        outputs = self._outputs().get()
        outputs['out_list'] = self.out_list
        return outputs

<commit_msg>dcmstack create directory if does not exist. pep8 fixes<commit_after>"""Provides interfaces to various commands provided by FreeSurfer

   Change directory to provide relative paths for doctests
   >>> import os
   >>> filepath = os.path.dirname( os.path.realpath( __file__ ) )
   >>> datadir = os.path.realpath(os.path.join(filepath, '../testing/data'))
   >>> os.chdir(datadir)

"""

from __future__ import absolute_import
import os
import string
import errno
from os import path as op
from glob import glob
from nipype.interfaces.base import (TraitedSpec,
                                    DynamicTraitedSpec,
                                    InputMultiPath,
                                    File,
                                    Directory,
                                    traits,
                                    BaseInterface,
                                   )
import nibabel as nb
from nipype.interfaces.traits_extension import isdefined, Undefined
import imghdr
from nipype.external import six

have_dcmstack = True
try:
    import dicom
    import dcmstack
    from dcmstack.dcmmeta import NiftiWrapper
except ImportError:
    have_dcmstack = False

def sanitize_path_comp(path_comp):
    result = []
    for char in path_comp:
        if not char in string.letters + string.digits + '-_.':
            result.append('_')
        else:
            result.append(char)
    return ''.join(result)

class NiftiGeneratorBaseInputSpec(TraitedSpec):
    out_format = traits.Str(desc="String which can be formatted with "
                            "meta data to create the output filename(s)")
    out_ext = traits.Str('.nii.gz', usedefault=True,
                         desc="Determines output file type")
    out_path = Directory(desc='output path, current working directory if not set')

class NiftiGeneratorBase(BaseInterface):
    '''Base class for interfaces that produce Nifti files, potentially with
    embeded meta data.'''
    def _get_out_path(self, meta, idx=None):
        '''Return the output path for the gernerated Nifti.'''
        if self.inputs.out_format:
            out_fmt = self.inputs.out_format
        else:
            #If no out_format is specified, use a sane default that will work
            #with the provided meta data.
            out_fmt = []
            if not idx is None:
                out_fmt.append('%03d' % idx)
            if 'SeriesNumber' in meta:
                out_fmt.append('%(SeriesNumber)03d')
            if 'ProtocolName' in meta:
                out_fmt.append('%(ProtocolName)s')
            elif 'SeriesDescription' in meta:
                out_fmt.append('%(SeriesDescription)s')
            else:
                out_fmt.append('sequence')
            out_fmt = '-'.join(out_fmt)
        out_fn = (out_fmt % meta) + self.inputs.out_ext
        out_fn = sanitize_path_comp(out_fn)

        out_path = os.getcwd()
        if isdefined(self.inputs.out_path):
            out_path = op.abspath(self.inputs.out_path)

            # now, mkdir -p $out_path
            try:
                op.makedirs(out_path)
            except OSError as exc:  # Python >2.5
                if exc.errno == errno.EEXIST and op.isdir(out_path):
                    pass
                else:
                    raise

        return op.join(out_path, out_fn)

class DcmStackInputSpec(NiftiGeneratorBaseInputSpec):
    dicom_files = traits.Either(InputMultiPath(File(exists=True)),
                                Directory(exists=True),
                                traits.Str(),
                                mandatory=True)
    embed_meta = traits.Bool(desc="Embed DICOM meta data into result")
    exclude_regexes = traits.List(desc="Meta data to exclude, suplementing "
                                  "any default exclude filters")
    include_regexes = traits.List(desc="Meta data to include, overriding any "
                                  "exclude filters")
    force_read = traits.Bool(True, usedefault=True,
                             desc=('Force reading files without DICM marker'))

class DcmStackOutputSpec(TraitedSpec):
    out_file = File(exists=True)

class DcmStack(NiftiGeneratorBase):
    '''Create one Nifti file from a set of DICOM files. Can optionally embed
    meta data.

    Example
    -------

    >>> from nipype.interfaces.dcmstack import DcmStack
    >>> stacker = DcmStack()
    >>> stacker.inputs.dicom_files = 'path/to/series/'
    >>> stacker.run() # doctest: +SKIP
    >>> result.outputs.out_file # doctest: +SKIP
    '/path/to/cwd/sequence.nii.gz'
    '''
    input_spec = DcmStackInputSpec
    output_spec = DcmStackOutputSpec

    def _get_filelist(self, trait_input):
        if isinstance(trait_input, six.string_types):
            if op.isdir(trait_input):
                return glob(op.join(trait_input, '*.dcm'))
            else:
                return glob(trait_input)

        return trait_input

    def _run_interface(self, runtime):
        src_paths = self._get_filelist(self.inputs.dicom_files)
        include_regexes = dcmstack.default_key_incl_res
        if isdefined(self.inputs.include_regexes):
            include_regexes += self.inputs.include_regexes
        exclude_regexes = dcmstack.default_key_excl_res
        if isdefined(self.inputs.exclude_regexes):
            exclude_regexes += self.inputs.exclude_regexes
        meta_filter = dcmstack.make_key_regex_filter(exclude_regexes,
                                                     include_regexes)
        stack = dcmstack.DicomStack(meta_filter=meta_filter)
        for src_path in src_paths:
            if not imghdr.what(src_path)=="gif":
                src_dcm = dicom.read_file(src_path, force=self.inputs.force_read)
                stack.add_dcm(src_dcm)
        nii = stack.to_nifti(embed_meta=True)
        nw = NiftiWrapper(nii)
        self.out_path = \
            self._get_out_path(nw.meta_ext.get_class_dict(('global', 'const')))
        if not self.inputs.embed_meta:
            nw.remove_extension()
        nb.save(nii, self.out_path)
        return runtime

    def _list_outputs(self):
        outputs = self._outputs().get()
        outputs["out_file"] = self.out_path
        return outputs

class GroupAndStackOutputSpec(TraitedSpec):
    out_list = traits.List(desc="List of output nifti files")

class GroupAndStack(DcmStack):
    '''Create (potentially) multiple Nifti files for a set of DICOM files.
    '''
    input_spec = DcmStackInputSpec
    output_spec = GroupAndStackOutputSpec

    def _run_interface(self, runtime):
        src_paths = self._get_filelist(self.inputs.dicom_files)
        stacks = dcmstack.parse_and_stack(src_paths)

        self.out_list = []
        for key, stack in stacks.iteritems():
            nw = NiftiWrapper(stack.to_nifti(embed_meta=True))
            const_meta = nw.meta_ext.get_class_dict(('global', 'const'))
            out_path =  self._get_out_path(const_meta)
            if not self.inputs.embed_meta:
                nw.remove_extension()
            nb.save(nw.nii_img, out_path)
            self.out_list.append(out_path)

        return runtime

    def _list_outputs(self):
        outputs = self._outputs().get()
        outputs["out_list"] = self.out_list
        return outputs

class LookupMetaInputSpec(TraitedSpec):
    in_file = File(mandatory=True,
                   exists=True,
                   desc='The input Nifti file')
    meta_keys = traits.Either(traits.List(),
                              traits.Dict(),
                              mandatory=True,
                              desc=("List of meta data keys to lookup, or a "
                              "dict where keys specify the meta data keys to "
                              "lookup and the values specify the output names")
                             )

class LookupMeta(BaseInterface):
    '''Lookup meta data values from a Nifti with embeded meta data.

    Example
    -------

    >>> from nipype.interfaces import dcmstack
    >>> lookup = dcmstack.LookupMeta()
    >>> lookup.inputs.in_file = 'functional.nii'
    >>> lookup.inputs.meta_keys = {'RepetitionTime' : 'TR', \
                                   'EchoTime' : 'TE'}
    >>> result = lookup.run() # doctest: +SKIP
    >>> result.outputs.TR # doctest: +SKIP
    9500.0
    >>> result.outputs.TE # doctest: +SKIP
    95.0
    '''
    input_spec = LookupMetaInputSpec
    output_spec = DynamicTraitedSpec

    def _make_name_map(self):
        if isinstance(self.inputs.meta_keys, list):
            self._meta_keys = {}
            for key in self.inputs.meta_keys:
                self._meta_keys[key] = key
        else:
            self._meta_keys = self.inputs.meta_keys

    def _outputs(self):
        self._make_name_map()
        outputs = super(LookupMeta, self)._outputs()
        undefined_traits = {}
        for out_name in self._meta_keys.values():
            outputs.add_trait(out_name, traits.Any)
            undefined_traits[out_name] = Undefined
        outputs.trait_set(trait_change_notify=False, **undefined_traits)
        #Not sure why this is needed
        for out_name in self._meta_keys.values():
            _ = getattr(outputs, out_name)
        return outputs

    def _run_interface(self, runtime):
        #If the 'meta_keys' input is a list, covert it to a dict
        self._make_name_map()
        nw = NiftiWrapper.from_filename(self.inputs.in_file)
        self.result = {}
        for meta_key, out_name in self._meta_keys.iteritems():
            self.result[out_name] = nw.meta_ext.get_values(meta_key)

        return runtime

    def _list_outputs(self):
        outputs = self._outputs().get()
        outputs.update(self.result)
        return outputs

class CopyMetaInputSpec(TraitedSpec):
    src_file = File(mandatory=True, exists=True)
    dest_file = File(mandatory=True, exists=True)
    include_classes = traits.List(desc="List of specific meta data "
                                  "classifications to include. If not "
                                  "specified include everything.")
    exclude_classes = traits.List(desc="List of meta data "
                                  "classifications to exclude")

class CopyMetaOutputSpec(TraitedSpec):
    dest_file = File(exists=True)

class CopyMeta(BaseInterface):
    '''Copy meta data from one Nifti file to another. Useful for preserving
    meta data after some processing steps.'''
    input_spec = CopyMetaInputSpec
    output_spec = CopyMetaOutputSpec

    def _run_interface(self, runtime):
        src_nii = nb.load(self.inputs.src_file)
        src = NiftiWrapper(src_nii, make_empty=True)
        dest_nii = nb.load(self.inputs.dest_file)
        dest = NiftiWrapper(dest_nii, make_empty=True)
        classes = src.meta_ext.get_valid_classes()
        if self.inputs.include_classes:
            classes = [cls
                       for cls in classes
                       if cls in self.inputs.include_classes
                      ]
        if self.inputs.exclude_classes:
            classes = [cls
                       for cls in classes
                       if not cls in self.inputs.exclude_classes
                      ]

        for cls in classes:
            src_dict = src.meta_ext.get_class_dict(cls)
            dest_dict = dest.meta_ext.get_class_dict(cls)
            dest_dict.update(src_dict)
        # Update the shape and slice dimension to reflect the meta extension update.
        dest.meta_ext.slice_dim = src.meta_ext.slice_dim
        dest.meta_ext.shape = src.meta_ext.shape

        self.out_path = op.join(os.getcwd(),
                                  op.basename(self.inputs.dest_file))
        dest.to_filename(self.out_path)

        return runtime

    def _list_outputs(self):
        outputs = self._outputs().get()
        outputs['dest_file'] = self.out_path
        return outputs

class MergeNiftiInputSpec(NiftiGeneratorBaseInputSpec):
    in_files = traits.List(mandatory=True,
                           desc="List of Nifti files to merge")
    sort_order = traits.Either(traits.Str(),
                               traits.List(),
                               desc="One or more meta data keys to "
                               "sort files by.")
    merge_dim = traits.Int(desc="Dimension to merge along. If not "
                           "specified, the last singular or "
                           "non-existant dimension is used.")

class MergeNiftiOutputSpec(TraitedSpec):
    out_file = File(exists=True, desc="Merged Nifti file")

def make_key_func(meta_keys, index=None):
    def key_func(src_nii):
        result = [src_nii.get_meta(key, index) for key in meta_keys]
        return result

    return key_func

class MergeNifti(NiftiGeneratorBase):
    '''Merge multiple Nifti files into one. Merges together meta data
    extensions as well.'''
    input_spec = MergeNiftiInputSpec
    output_spec = MergeNiftiOutputSpec

    def _run_interface(self, runtime):
        niis = [nb.load(fn)
                for fn in self.inputs.in_files
               ]
        nws = [NiftiWrapper(nii, make_empty=True)
               for nii in niis
              ]
        if self.inputs.sort_order:
            sort_order = self.inputs.sort_order
            if isinstance(sort_order, six.string_types):
                sort_order = [sort_order]
            nws.sort(key=make_key_func(sort_order))
        if self.inputs.merge_dim == traits.Undefined:
            merge_dim = None
        else:
            merge_dim = self.inputs.merge_dim
        merged = NiftiWrapper.from_sequence(nws, merge_dim)
        const_meta = merged.meta_ext.get_class_dict(('global', 'const'))
        self.out_path = self._get_out_path(const_meta)
        nb.save(merged.nii_img, self.out_path)
        return runtime

    def _list_outputs(self):
        outputs = self._outputs().get()
        outputs['out_file'] = self.out_path
        return outputs

class SplitNiftiInputSpec(NiftiGeneratorBaseInputSpec):
    in_file = File(exists=True, mandatory=True, desc="Nifti file to split")
    split_dim = traits.Int(desc="Dimension to split along. If not "
                           "specified, the last dimension is used.")

class SplitNiftiOutputSpec(TraitedSpec):
    out_list = traits.List(File(exists=True),
                           desc="Split Nifti files")

class SplitNifti(NiftiGeneratorBase):
    '''
    Split one Nifti file into many along the specified dimension. Each
    result has an updated meta data extension as well.
    '''
    input_spec = SplitNiftiInputSpec
    output_spec = SplitNiftiOutputSpec

    def _run_interface(self, runtime):
        self.out_list = []
        nii = nb.load(self.inputs.in_file)
        nw = NiftiWrapper(nii, make_empty=True)
        split_dim = None
        if self.inputs.split_dim == traits.Undefined:
            split_dim = None
        else:
            split_dim = self.inputs.split_dim
        for split_idx, split_nw in enumerate(nw.split(split_dim)):
            const_meta = split_nw.meta_ext.get_class_dict(('global', 'const'))
            out_path = self._get_out_path(const_meta, idx=split_idx)
            nb.save(split_nw.nii_img, out_path)
            self.out_list.append(out_path)

        return runtime

    def _list_outputs(self):
        outputs = self._outputs().get()
        outputs['out_list'] = self.out_list
        return outputs

<|endoftext|>


<commit_before>from mahjong.tile import TilesConverter
from mahjong.utils import plus_dora, count_tiles_by_suits

class EnemyAnalyzer(object):
    player = None
    chosen_suit = None

    def __init__(self, player):
        """
        :param player: instance of EnemyPlayer
        """
        self.player = player
        self.table = player.table
        self.chosen_suit = None

    @property
    def in_tempai(self):
        """
        Try to detect is user in tempai or not
        :return: boolean
        """
        # simplest case, user in riichi
        if self.player.in_riichi:
            return True

        if len(self.player.melds) == 4:
            return True

        return False

    @property
    def is_threatening(self):
        """
        Should we fold against this player or not
        :return: boolean
        """
        if self.player.in_riichi:
            return True

        discards = self.player.discards
        discards_34 = TilesConverter.to_34_array([x.value for x in discards])

        is_honitsu_open_sets, open_hand_suit = False, None
        is_honitsu_discards, discard_suit = self._is_honitsu_discards(discards_34)

        meld_tiles = self.player.meld_tiles
        meld_tiles_34 = TilesConverter.to_34_array(meld_tiles)
        if meld_tiles:
            dora_count = sum([plus_dora(x, self.table.dora_indicators) for x in meld_tiles])
            # enemy has a lot of dora tiles in his opened sets
            # so better to fold against him
            if dora_count >= 3:
                return True

            # check that user has a discard and melds that looks like honitsu
            is_honitsu_open_sets, open_hand_suit = self._is_honitsu_open_sets(meld_tiles_34)

        if is_honitsu_discards:
            # user's discard looks like honitsu
            # but he opened sets from different suits
            if meld_tiles and not is_honitsu_open_sets:
                return False

            # users open sets and his discards tell to us
            # that he is collecting specific suit
            self.chosen_suit = open_hand_suit
            return True

        return False

    def _is_honitsu_open_sets(self, meld_tiles_34):
        """
        Check that user opened all sets with same suit
        :param meld_tiles_34:
        :return:
        """
        total_sets = sum(meld_tiles_34) // 3
        if total_sets < 2:
            return False, None

        result = count_tiles_by_suits(meld_tiles_34)

        suits = [x for x in result if x['name'] != 'honor']
        suits = sorted(suits, key=lambda x: x['count'], reverse=True)

        # for honitsu we can have tiles only in one suit
        if suits[1]['count'] == 0 or suits[2]['count'] == 0:
            return True, suits[0]['function']

        return False, None

    def _is_honitsu_discards(self, discards_34):
        """
        Check that user opened all sets with same suit
        :param discards_34:
        :return:
        """
        total_discards = sum(discards_34)

        # there is no sense to analyze earlier discards
        if total_discards < 6:
            return False, None

        result = count_tiles_by_suits(discards_34)

        honor = [x for x in result if x['name'] == 'honor'][0]
        suits = [x for x in result if x['name'] != 'honor']
        suits = sorted(suits, key=lambda x: x['count'], reverse=False)

        less_suit = suits[0]['count'] + honor['count']
        percentage_of_less_suit = (less_suit / total_discards) * 100

        # there is not too much one suit + honor tiles in the discard
        # so we can tell that user trying to collect honitsu
        if percentage_of_less_suit <= 20:
            return True, suits[0]['function']

        return False, None
<commit_msg>Improve the way to detect honitsu hands<commit_after>from mahjong.tile import TilesConverter
from mahjong.utils import plus_dora, count_tiles_by_suits

class EnemyAnalyzer(object):
    player = None
    chosen_suit = None

    def __init__(self, player):
        """
        :param player: instance of EnemyPlayer
        """
        self.player = player
        self.table = player.table
        self.chosen_suit = None

    @property
    def in_tempai(self):
        """
        Try to detect is user in tempai or not
        :return: boolean
        """
        # simplest case, user in riichi
        if self.player.in_riichi:
            return True

        if len(self.player.melds) == 4:
            return True

        return False

    @property
    def is_threatening(self):
        """
        Should we fold against this player or not
        :return: boolean
        """
        if self.player.in_riichi:
            return True

        discards = self.player.discards
        discards_34 = TilesConverter.to_34_array([x.value for x in discards])

        is_honitsu_open_sets, open_hand_suit = False, None
        is_honitsu_discards, discard_suit = self._is_honitsu_discards(discards_34)

        meld_tiles = self.player.meld_tiles
        meld_tiles_34 = TilesConverter.to_34_array(meld_tiles)
        if meld_tiles:
            dora_count = sum([plus_dora(x, self.table.dora_indicators) for x in meld_tiles])
            # enemy has a lot of dora tiles in his opened sets
            # so better to fold against him
            if dora_count >= 3:
                return True

            # check that user has a discard and melds that looks like honitsu
            is_honitsu_open_sets, open_hand_suit = self._is_honitsu_open_sets(meld_tiles_34)

        if is_honitsu_open_sets and is_honitsu_discards:
            # users open sets and his discards tell to us
            # that he is collecting specific suit
            self.chosen_suit = open_hand_suit
            return True

        return False

    def _is_honitsu_open_sets(self, meld_tiles_34):
        """
        Check that user opened all sets with same suit
        :param meld_tiles_34:
        :return:
        """
        total_sets = sum(meld_tiles_34) // 3
        if total_sets < 2:
            return False, None

        result = count_tiles_by_suits(meld_tiles_34)

        suits = [x for x in result if x['name'] != 'honor']
        suits = sorted(suits, key=lambda x: x['count'], reverse=True)

        # for honitsu we can have tiles only in one suit
        if suits[1]['count'] == 0 or suits[2]['count'] == 0:
            return True, suits[0]['function']

        return False, None

    def _is_honitsu_discards(self, discards_34):
        """
        Check that user opened all sets with same suit
        :param discards_34:
        :return:
        """
        total_discards = sum(discards_34)

        # there is no sense to analyze earlier discards
        if total_discards < 6:
            return False, None

        result = count_tiles_by_suits(discards_34)

        honor = [x for x in result if x['name'] == 'honor'][0]
        suits = [x for x in result if x['name'] != 'honor']
        suits = sorted(suits, key=lambda x: x['count'], reverse=False)

        less_suit = suits[0]['count'] + honor['count']
        percentage_of_less_suit = (less_suit / total_discards) * 100

        # there is not too much one suit + honor tiles in the discard
        # so we can tell that user trying to collect honitsu
        if percentage_of_less_suit <= 20:
            return True, suits[0]['function']

        return False, None
<|endoftext|>


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

import os
import sys
import glob
from py4j.protocol import Py4JJavaError
from py4j.java_gateway import JavaObject
from py4j.java_collections import ListConverter, JavaArray, JavaList, JavaMap

from pyspark import RDD, SparkContext
from pyspark.serializers import PickleSerializer, AutoBatchedSerializer
from pyspark.sql import DataFrame, SQLContext
from pyspark.mllib.common import callJavaFunc
from pyspark import SparkConf
import numpy as np
import threading
from bigdl.util.engine import prepare_env

INTMAX = 2147483647
INTMIN = -2147483648
DOUBLEMAX = 1.7976931348623157E308

if sys.version >= '3':
    long = int
    unicode = str

class SingletonMixin(object):
    _lock = threading.RLock()
    _instance = None

    @classmethod
    def instance(cls,
                 bigdl_type="float"):
        if not cls._instance:
            with cls._lock:
                if not cls._instance:
                    cls._instance = cls(bigdl_type)
        return cls._instance

class JavaCreator(SingletonMixin):
    __creator_class="com.intel.analytics.bigdl.python.api.PythonBigDL"

    @classmethod
    def get_creator_class(cls):
        with JavaCreator._lock:
            return JavaCreator.__creator_class

    @classmethod
    def set_creator_class(cls, cclass):
        with JavaCreator._lock:
            JavaCreator.__creator_class = cclass
            JavaCreator._instance = None

    def __init__(self, bigdl_type):
        sc = get_spark_context()
        jclass = getattr(sc._jvm, JavaCreator.get_creator_class())
        if bigdl_type == "float":
            self.value = getattr(jclass, "ofFloat")()
        elif bigdl_type == "double":
            self.value = getattr(jclass, "ofDouble")()
        else:
            raise Exception("Not supported bigdl_type: %s" % bigdl_type)

class JavaValue(object):
    def jvm_class_constructor(self):
        name = "create" + self.__class__.__name__
        print("creating: " + name)
        return name

    def __init__(self, jvalue, bigdl_type, *args):
        self.value = jvalue if jvalue else callBigDlFunc(
            bigdl_type, self.jvm_class_constructor(), *args)
        self.bigdl_type = bigdl_type

    def __str__(self):
        return self.value.toString()

class EvaluatedResult():
    """
    A testing result used to benchmark the model quality.
    """
    def __init__(self, result, total_num, method):
        """

        :param result: the validation result. i.e: top1 accuracy percentage.
        :param total_num: the total processed records.
        :param method: the validation method. i.e: Top1Accuracy
        """
        self.result = result
        self.total_num = total_num
        self.method = method

    def __reduce__(self):
        return (EvaluatedResult, (self.result, self.total_num, self.method))

    def __str__(self):
        return "Evaluated result: %s, total_num: %s, method: %s" % (
            self.result, self.total_num, self.method)

class JTensor(object):
    """
    A wrapper to easy our work when need to pass or return Tensor to/from Scala.

    >>> import numpy as np
    >>> from bigdl.util.common import JTensor
    >>> np.random.seed(123)
    >>>
    """
    def __init__(self, storage, shape, bigdl_type="float"):
        self.storage = storage
        self.shape = shape
        self.bigdl_type = bigdl_type

    @classmethod
    def from_ndarray(cls, a_ndarray, bigdl_type="float"):
        """
        Convert a ndarray to Tensor which would be used in Java side.

        >>> import numpy as np
        >>> from bigdl.util.common import JTensor
        >>> from bigdl.util.common import callBigDlFunc
        >>> np.random.seed(123)
        >>> data = np.random.uniform(0, 1, (2, 3)).astype("float32")
        >>> result = JTensor.from_ndarray(data)
        >>> data_back = result.to_ndarray()
        >>> (data == data_back).all()
        True
        >>> tensor1 = callBigDlFunc("float", "testTensor", JTensor.from_ndarray(data))  # noqa
        >>> array_from_tensor = tensor1.to_ndarray()
        >>> (array_from_tensor == data).all()
        True
        """
        return cls(*JTensor.flatten_ndarray(a_ndarray),
                   bigdl_type= bigdl_type) if a_ndarray is not None else None  # noqa

    def to_ndarray(self):
        def get_dtype():
            if "float" == self.bigdl_type:
                return "float32"
            else:
                return "float64"
        return np.array(self.storage, dtype=get_dtype()).reshape(self.shape)  # noqa

    @classmethod
    def flatten_ndarray(cls, a_ndarray):
        """
        Utility method to flatten a ndarray

        :return: (storage, shape)

        >>> from bigdl.util.common import JTensor
        >>> np.random.seed(123)
        >>> data = np.random.uniform(0, 1, (2, 3))
        >>> (storage, shape) = JTensor.flatten_ndarray(data)
        >>> shape
        [2, 3]
        >>> (storage, shape) = JTensor.flatten_ndarray(np.array(2))
        >>> shape
        [1]
        """
        storage = [float(i) for i in a_ndarray.ravel()]
        shape = list(a_ndarray.shape) if a_ndarray.shape else [a_ndarray.size]
        return storage, shape

    def __reduce__(self):
        return (JTensor, (self.storage, self.shape, self.bigdl_type))

    def __str__(self):
        return "storage: %s, shape: %s," % (self.storage, self.storage)

class Sample(object):
    def __init__(self, features, label, features_shape, label_shape,
                 bigdl_type="float"):
        def get_dtype():
            if "float" == bigdl_type:
                return "float32"
            else:
                return "float64"
        self.features = np.array(features, dtype=get_dtype()).reshape(features_shape)  # noqa
        self.label = np.array(label, dtype=get_dtype()).reshape(label_shape)
        self.bigdl_type = bigdl_type

    @classmethod
    def from_ndarray(cls, features, label, bigdl_type="float"):
        return cls(
            features=[float(i) for i in features.ravel()],
            label=[float(i) for i in label.ravel()],
            features_shape=list(features.shape),
            label_shape=list(label.shape) if label.shape else [label.size],
            bigdl_type=bigdl_type)

    def __reduce__(self):
        (features_storage, features_shape) = JTensor.flatten_ndarray(self.features)
        (label_storage, label_shape) = JTensor.flatten_ndarray(self.label)
        return (Sample, (
            features_storage, label_storage, features_shape, label_shape,
            self.bigdl_type))

    def __str__(self):
        return "features: %s, label: %s," % (self.features, self.label)

class RNG():
    """
    generate tensor data with seed
    """
    def __init__(self, bigdl_type="float"):
        self.bigdl_type = bigdl_type

    def set_seed(self, seed):
        callBigDlFunc(self.bigdl_type, "setModelSeed", seed)

    def uniform(self, a, b, size):
        return callBigDlFunc(self.bigdl_type, "uniform", a, b, size).to_ndarray() # noqa

_picklable_classes = [
    'LinkedList',
    'SparseVector',
    'DenseVector',
    'DenseMatrix',
    'Rating',
    'LabeledPoint',
    'Sample',
    'EvaluatedResult',
    'JTensor'
]

def init_engine(bigdl_type="float"):
    callBigDlFunc(bigdl_type, "initEngine")

def get_bigdl_conf():
    bigdl_conf_file = "spark-bigdl.conf"
    bigdl_python_wrapper = "python-api.zip"

    def load_conf(conf_str):
        return dict(line.split() for line in conf_str.split("\n") if
                    "#" not in line and line.strip())

    for p in sys.path:
        if bigdl_conf_file in p and os.path.isfile(p):
            with open(p) if sys.version_info < (3,) else open(p, encoding='latin-1') as conf_file: # noqa
                return load_conf(conf_file.read())
        if bigdl_python_wrapper in p and os.path.isfile(p):
            import zipfile
            with zipfile.ZipFile(p, 'r') as zip_conf:
                content = zip_conf.read(bigdl_conf_file)
                if sys.version_info >= (3,):
                    content = str(content, 'latin-1')
                return load_conf(content)
    raise Exception("Cannot find spark-bigdl.conf.Pls add it to PYTHONPATH.")

def to_list(a):
    if type(a) is list:
        return a
    return [a]

def create_spark_conf():
    bigdl_conf = get_bigdl_conf()
    sparkConf = SparkConf()
    sparkConf.setAll(bigdl_conf.items())
    return sparkConf

def get_spark_context(conf = None):
    """
    Get the current active spark context and create one if no active instance
    :param conf: combining bigdl configs into spark conf
    :return: SparkContext
    """
    with SparkContext._lock:  # Compatible with Spark1.5.1
        if SparkContext._active_spark_context is None:
            SparkContext(conf=conf or create_spark_conf())
        return SparkContext._active_spark_context

def get_spark_sql_context(sc):
    if "getOrCreate" in SQLContext.__dict__:
        return SQLContext.getOrCreate()
    else:
        return SQLContext(sc)  # Compatible with Spark1.5.1

def callBigDlFunc(bigdl_type, name, *args):
    """ Call API in PythonBigDL """
    jinstance = JavaCreator.instance(bigdl_type=bigdl_type).value
    sc = get_spark_context()
    api = getattr(jinstance, name)
    return callJavaFunc(sc, api, *args)

def _java2py(sc, r, encoding="bytes"):
    if isinstance(r, JavaObject):
        clsName = r.getClass().getSimpleName()
        # convert RDD into JavaRDD
        if clsName != 'JavaRDD' and clsName.endswith("RDD"):
            r = r.toJavaRDD()
            clsName = 'JavaRDD'

        if clsName == 'JavaRDD':
            jrdd = sc._jvm.SerDe.javaToPython(r)
            return RDD(jrdd, sc)

        if clsName == 'DataFrame':
            return DataFrame(r, get_spark_sql_context(sc))

        if clsName in _picklable_classes:
            r = sc._jvm.org.apache.spark.bigdl.api.python.BigDLSerDe.dumps(r)
        elif isinstance(r, (JavaArray, JavaList, JavaMap)):
            try:
                r = sc._jvm.org.apache.spark.bigdl.api.python.BigDLSerDe.dumps(
                    r)
            except Py4JJavaError:
                pass  # not pickable

    if isinstance(r, (bytearray, bytes)):
        r = PickleSerializer().loads(bytes(r), encoding=encoding)
    return r

def callJavaFunc(sc, func, *args):
    """ Call Java Function """
    args = [_py2java(sc, a) for a in args]
    result = func(*args)
    return _java2py(sc, result)

def _to_java_object_rdd(rdd):
    """ Return a JavaRDD of Object by unpickling

    It will convert each Python object into Java object by Pyrolite, whenever
    the RDD is serialized in batch or not.
    """
    rdd = rdd._reserialize(AutoBatchedSerializer(PickleSerializer()))
    return \
        rdd.ctx._jvm.org.apache.spark.bigdl.api.python.BigDLSerDe.pythonToJava(
            rdd._jrdd, True)

def _py2java(sc, obj):
    """ Convert Python object into Java """
    if isinstance(obj, RDD):
        obj = _to_java_object_rdd(obj)
    elif isinstance(obj, DataFrame):
        obj = obj._jdf
    elif isinstance(obj, SparkContext):
        obj = obj._jsc
    elif isinstance(obj, (list, tuple)):
        obj = ListConverter().convert([_py2java(sc, x) for x in obj],
                                      sc._gateway._gateway_client)
    elif isinstance(obj, dict):
        result = {}
        for (key, value) in obj.items():
            result[key] = _py2java(sc, value) if isinstance(value, JavaValue) else value  # noqa
        obj = result

    elif isinstance(obj, JavaValue):
        obj = obj.value
    elif isinstance(obj, JavaObject):
        pass
    elif isinstance(obj, (int, long, float, bool, bytes, unicode)):
        pass
    else:
        data = bytearray(PickleSerializer().dumps(obj))
        obj = sc._jvm.org.apache.spark.bigdl.api.python.BigDLSerDe.loads(data)
    return obj

def _test():
    import doctest
    from pyspark import SparkContext
    from bigdl.nn import layer
    globs = layer.__dict__.copy()
    sc = SparkContext(master="local[2]", appName="test common utility")
    globs['sc'] = sc
    (failure_count, test_count) = doctest.testmod(globs=globs,
                                                  optionflags=doctest.ELLIPSIS)
    if failure_count:
        exit(-1)

if __name__ == "__main__":
    _test()
<commit_msg>Ignore some folders from doctest scanning (#1465)<commit_after>#
# Copyright 2016 The BigDL Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

import os
import sys
import glob
from py4j.protocol import Py4JJavaError
from py4j.java_gateway import JavaObject
from py4j.java_collections import ListConverter, JavaArray, JavaList, JavaMap

from pyspark import RDD, SparkContext
from pyspark.serializers import PickleSerializer, AutoBatchedSerializer
from pyspark.sql import DataFrame, SQLContext
from pyspark.mllib.common import callJavaFunc
from pyspark import SparkConf
import numpy as np
import threading
from bigdl.util.engine import prepare_env

INTMAX = 2147483647
INTMIN = -2147483648
DOUBLEMAX = 1.7976931348623157E308

if sys.version >= '3':
    long = int
    unicode = str

class SingletonMixin(object):
    _lock = threading.RLock()
    _instance = None

    @classmethod
    def instance(cls,
                 bigdl_type="float"):
        if not cls._instance:
            with cls._lock:
                if not cls._instance:
                    cls._instance = cls(bigdl_type)
        return cls._instance

class JavaCreator(SingletonMixin):
    __creator_class="com.intel.analytics.bigdl.python.api.PythonBigDL"

    @classmethod
    def get_creator_class(cls):
        with JavaCreator._lock:
            return JavaCreator.__creator_class

    @classmethod
    def set_creator_class(cls, cclass):
        with JavaCreator._lock:
            JavaCreator.__creator_class = cclass
            JavaCreator._instance = None

    def __init__(self, bigdl_type):
        sc = get_spark_context()
        jclass = getattr(sc._jvm, JavaCreator.get_creator_class())
        if bigdl_type == "float":
            self.value = getattr(jclass, "ofFloat")()
        elif bigdl_type == "double":
            self.value = getattr(jclass, "ofDouble")()
        else:
            raise Exception("Not supported bigdl_type: %s" % bigdl_type)

class JavaValue(object):
    def jvm_class_constructor(self):
        name = "create" + self.__class__.__name__
        print("creating: " + name)
        return name

    def __init__(self, jvalue, bigdl_type, *args):
        self.value = jvalue if jvalue else callBigDlFunc(
            bigdl_type, self.jvm_class_constructor(), *args)
        self.bigdl_type = bigdl_type

    def __str__(self):
        return self.value.toString()

class EvaluatedResult():
    """
    A testing result used to benchmark the model quality.
    """
    def __init__(self, result, total_num, method):
        """

        :param result: the validation result. i.e: top1 accuracy percentage.
        :param total_num: the total processed records.
        :param method: the validation method. i.e: Top1Accuracy
        """
        self.result = result
        self.total_num = total_num
        self.method = method

    def __reduce__(self):
        return (EvaluatedResult, (self.result, self.total_num, self.method))

    def __str__(self):
        return "Evaluated result: %s, total_num: %s, method: %s" % (
            self.result, self.total_num, self.method)

class JTensor(object):
    """
    A wrapper to easy our work when need to pass or return Tensor to/from Scala.

    >>> import numpy as np
    >>> from bigdl.util.common import JTensor
    >>> np.random.seed(123)
    >>>
    """
    def __init__(self, storage, shape, bigdl_type="float"):
        self.storage = storage
        self.shape = shape
        self.bigdl_type = bigdl_type

    @classmethod
    def from_ndarray(cls, a_ndarray, bigdl_type="float"):
        """
        Convert a ndarray to Tensor which would be used in Java side.

        >>> import numpy as np
        >>> from bigdl.util.common import JTensor
        >>> from bigdl.util.common import callBigDlFunc
        >>> np.random.seed(123)
        >>> data = np.random.uniform(0, 1, (2, 3)).astype("float32")
        >>> result = JTensor.from_ndarray(data)
        >>> data_back = result.to_ndarray()
        >>> (data == data_back).all()
        True
        >>> tensor1 = callBigDlFunc("float", "testTensor", JTensor.from_ndarray(data))  # noqa
        >>> array_from_tensor = tensor1.to_ndarray()
        >>> (array_from_tensor == data).all()
        True
        """
        return cls(*JTensor.flatten_ndarray(a_ndarray),
                   bigdl_type= bigdl_type) if a_ndarray is not None else None  # noqa

    def to_ndarray(self):
        def get_dtype():
            if "float" == self.bigdl_type:
                return "float32"
            else:
                return "float64"
        return np.array(self.storage, dtype=get_dtype()).reshape(self.shape)  # noqa

    @classmethod
    def flatten_ndarray(cls, a_ndarray):
        """
        Utility method to flatten a ndarray

        :return: (storage, shape)

        >>> from bigdl.util.common import JTensor
        >>> np.random.seed(123)
        >>> data = np.random.uniform(0, 1, (2, 3))
        >>> (storage, shape) = JTensor.flatten_ndarray(data)
        >>> shape
        [2, 3]
        >>> (storage, shape) = JTensor.flatten_ndarray(np.array(2))
        >>> shape
        [1]
        """
        storage = [float(i) for i in a_ndarray.ravel()]
        shape = list(a_ndarray.shape) if a_ndarray.shape else [a_ndarray.size]
        return storage, shape

    def __reduce__(self):
        return (JTensor, (self.storage, self.shape, self.bigdl_type))

    def __str__(self):
        return "storage: %s, shape: %s," % (self.storage, self.storage)

class Sample(object):
    def __init__(self, features, label, features_shape, label_shape,
                 bigdl_type="float"):
        def get_dtype():
            if "float" == bigdl_type:
                return "float32"
            else:
                return "float64"
        self.features = np.array(features, dtype=get_dtype()).reshape(features_shape)  # noqa
        self.label = np.array(label, dtype=get_dtype()).reshape(label_shape)
        self.bigdl_type = bigdl_type

    @classmethod
    def from_ndarray(cls, features, label, bigdl_type="float"):
        return cls(
            features=[float(i) for i in features.ravel()],
            label=[float(i) for i in label.ravel()],
            features_shape=list(features.shape),
            label_shape=list(label.shape) if label.shape else [label.size],
            bigdl_type=bigdl_type)

    def __reduce__(self):
        (features_storage, features_shape) = JTensor.flatten_ndarray(self.features)
        (label_storage, label_shape) = JTensor.flatten_ndarray(self.label)
        return (Sample, (
            features_storage, label_storage, features_shape, label_shape,
            self.bigdl_type))

    def __str__(self):
        return "features: %s, label: %s," % (self.features, self.label)

class RNG():
    """
    generate tensor data with seed
    """
    def __init__(self, bigdl_type="float"):
        self.bigdl_type = bigdl_type

    def set_seed(self, seed):
        callBigDlFunc(self.bigdl_type, "setModelSeed", seed)

    def uniform(self, a, b, size):
        return callBigDlFunc(self.bigdl_type, "uniform", a, b, size).to_ndarray() # noqa

_picklable_classes = [
    'LinkedList',
    'SparseVector',
    'DenseVector',
    'DenseMatrix',
    'Rating',
    'LabeledPoint',
    'Sample',
    'EvaluatedResult',
    'JTensor'
]

def init_engine(bigdl_type="float"):
    callBigDlFunc(bigdl_type, "initEngine")

def get_bigdl_conf():
    bigdl_conf_file = "spark-bigdl.conf"
    bigdl_python_wrapper = "python-api.zip"

    def load_conf(conf_str):
        return dict(line.split() for line in conf_str.split("\n") if
                    "#" not in line and line.strip())

    for p in sys.path:
        if bigdl_conf_file in p and os.path.isfile(p):
            with open(p) if sys.version_info < (3,) else open(p, encoding='latin-1') as conf_file: # noqa
                return load_conf(conf_file.read())
        if bigdl_python_wrapper in p and os.path.isfile(p):
            import zipfile
            with zipfile.ZipFile(p, 'r') as zip_conf:
                content = zip_conf.read(bigdl_conf_file)
                if sys.version_info >= (3,):
                    content = str(content, 'latin-1')
                return load_conf(content)
    raise Exception("Cannot find spark-bigdl.conf.Pls add it to PYTHONPATH.")

def to_list(a):
    if type(a) is list:
        return a
    return [a]

def create_spark_conf():
    bigdl_conf = get_bigdl_conf()
    sparkConf = SparkConf()
    sparkConf.setAll(bigdl_conf.items())
    return sparkConf

def get_spark_context(conf = None):
    """
    Get the current active spark context and create one if no active instance
    :param conf: combining bigdl configs into spark conf
    :return: SparkContext
    """
    if hasattr(SparkContext, "getOrCreate"):
        return SparkContext.getOrCreate(conf=conf or create_spark_conf())
    else:
        # Might have threading issue but we cann't add _lock here
        # as it's not RLock in spark1.5
        if SparkContext._active_spark_context is None:
            SparkContext(conf=conf or create_spark_conf())
        return SparkContext._active_spark_context

def get_spark_sql_context(sc):
    if "getOrCreate" in SQLContext.__dict__:
        return SQLContext.getOrCreate()
    else:
        return SQLContext(sc)  # Compatible with Spark1.5.1

def callBigDlFunc(bigdl_type, name, *args):
    """ Call API in PythonBigDL """
    jinstance = JavaCreator.instance(bigdl_type=bigdl_type).value
    sc = get_spark_context()
    api = getattr(jinstance, name)
    return callJavaFunc(sc, api, *args)

def _java2py(sc, r, encoding="bytes"):
    if isinstance(r, JavaObject):
        clsName = r.getClass().getSimpleName()
        # convert RDD into JavaRDD
        if clsName != 'JavaRDD' and clsName.endswith("RDD"):
            r = r.toJavaRDD()
            clsName = 'JavaRDD'

        if clsName == 'JavaRDD':
            jrdd = sc._jvm.SerDe.javaToPython(r)
            return RDD(jrdd, sc)

        if clsName == 'DataFrame':
            return DataFrame(r, get_spark_sql_context(sc))

        if clsName in _picklable_classes:
            r = sc._jvm.org.apache.spark.bigdl.api.python.BigDLSerDe.dumps(r)
        elif isinstance(r, (JavaArray, JavaList, JavaMap)):
            try:
                r = sc._jvm.org.apache.spark.bigdl.api.python.BigDLSerDe.dumps(
                    r)
            except Py4JJavaError:
                pass  # not pickable

    if isinstance(r, (bytearray, bytes)):
        r = PickleSerializer().loads(bytes(r), encoding=encoding)
    return r

def callJavaFunc(sc, func, *args):
    """ Call Java Function """
    args = [_py2java(sc, a) for a in args]
    result = func(*args)
    return _java2py(sc, result)

def _to_java_object_rdd(rdd):
    """ Return a JavaRDD of Object by unpickling

    It will convert each Python object into Java object by Pyrolite, whenever
    the RDD is serialized in batch or not.
    """
    rdd = rdd._reserialize(AutoBatchedSerializer(PickleSerializer()))
    return \
        rdd.ctx._jvm.org.apache.spark.bigdl.api.python.BigDLSerDe.pythonToJava(
            rdd._jrdd, True)

def _py2java(sc, obj):
    """ Convert Python object into Java """
    if isinstance(obj, RDD):
        obj = _to_java_object_rdd(obj)
    elif isinstance(obj, DataFrame):
        obj = obj._jdf
    elif isinstance(obj, SparkContext):
        obj = obj._jsc
    elif isinstance(obj, (list, tuple)):
        obj = ListConverter().convert([_py2java(sc, x) for x in obj],
                                      sc._gateway._gateway_client)
    elif isinstance(obj, dict):
        result = {}
        for (key, value) in obj.items():
            result[key] = _py2java(sc, value) if isinstance(value, JavaValue) else value  # noqa
        obj = result

    elif isinstance(obj, JavaValue):
        obj = obj.value
    elif isinstance(obj, JavaObject):
        pass
    elif isinstance(obj, (int, long, float, bool, bytes, unicode)):
        pass
    else:
        data = bytearray(PickleSerializer().dumps(obj))
        obj = sc._jvm.org.apache.spark.bigdl.api.python.BigDLSerDe.loads(data)
    return obj

def _test():
    import doctest
    from pyspark import SparkContext
    from bigdl.nn import layer
    globs = layer.__dict__.copy()
    sc = SparkContext(master="local[2]", appName="test common utility")
    globs['sc'] = sc
    (failure_count, test_count) = doctest.testmod(globs=globs,
                                                  optionflags=doctest.ELLIPSIS)
    if failure_count:
        exit(-1)

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