/mlabwrap_dev/mlabwrap_core.py

##############################################################################
################## mlabwrap: transparently wraps matlab(tm) ##################
##############################################################################
##
## o author: Alexander Schmolck <a.schmolck@gmx.net>
## o created: 2002-05-29 21:51:59+00:40
## o version: see `__version__`
## o keywords: matlab wrapper
## o license: MIT
## o FIXME:
##   - it seems proxies can somehow still 'disappear', maybe in connection
##     with exceptions in the matlab workspace?
##   - add test that defunct proxy-values are culled from matlab workspace
##     (for some reason ipython seems to keep them alive somehwere, even after
##     a zaphist, should find out what causes that!)
##   - add tests for exception handling!
##   - the proxy getitem/setitem only works quite properly for 1D arrays
##     (matlab's moronic syntax also means that 'foo(bar)(watz)' is not the
##     same as 'tmp = foo(bar); tmp(watz)' -- indeed chances are the former
##     will fail (not, apparently however with 'foo{bar}{watz}' (blech)). This
##     would make it quite hard to the proxy thing 'properly' for nested
##     proxies, so things may break down for complicated cases but can be
##     easily fixed manually e.g.: ``mlab._set('tmp', foo(bar));
##     mlab._get('tmp',remove=True)[watz]``
##   - Guess there should be some in principle avoidable problems with
##     assignments to sub-proxies (in addition to the more fundamental,
##     unavoidable problem that ``proxy[index].part = foo`` can't work as
##     expected if ``proxy[index]`` is a marshallable value that doesn't need
##     to be proxied itself; see below for workaround).
## o XXX:
##   - better support of string 'arrays'
##   - multi-dimensional arrays are unsupported
##   - treatment of lists, tuples and arrays with non-numerical values (these
##     should presumably be wrapped into wrapper classes MlabCell etc.)
##   - should test classes and further improve struct support?
##   - should we transform 1D vectors into row vectors when handing them to
##     matlab?
##   - what should be flattend? Should there be a scalarization opition?
##   - ``autosync_dirs`` is a bit of a hack (and maybe``handle_out``, too)...
##   - is ``global mlab`` in unpickling of proxies OK?
##   - hasattr fun for proxies (__deepcopy__ etc.)
##   - check pickling
## o TODO:
##   - delattr
##   - better error reporting: test for number of input args etc.
##   - add cloning of proxies.
##   - pickling for nested proxies (and session management for pickling)
##   - more tests
## o !!!:
##   - matlab complex arrays are intelligently of type 'double'
##   - ``class('func')`` but ``class(var)``


__version__ = '1.1'
__author__   = "Alexander Schmolck <a.schmolck@gmx.net>"
__all__ = ['matlab', 'MlabError']

import warnings
import os, sys, re
import weakref
import logging

import numpy as np

from mlabwrap_dev import mlabraw

MODULE_LOGGER = logging.getLogger(__name__)
MODULE_LOGGER.addHandler(logging.NullHandler())


#XXX: nested access
def _flush_write_stdout(s):
    """Writes `s` to stdout and flushes. Default value for ``handle_out``."""
    sys.stdout.write(s); sys.stdout.flush()


# XXX I changed this to no longer use weakrefs because it didn't seem 100%
# reliable on second thought; need to check if we need to do something to
# speed up proxy reclamation on the matlab side.
class CurlyIndexer(object):
    """A helper class to mimick ``foo{bar}``-style indexing in python."""
    def __init__(self, proxy):
        self.proxy = proxy

    def __getitem__(self, index):
        return self.proxy.__getitem__(index, '{}')

    def __setitem__(self, index, value):
        self.proxy.__setitem__(index, value, '{}')


class MlabObjectProxy(object):
    """A proxy class for matlab objects that can't be converted to python
    types.

    WARNING: There are impedance-mismatch issues between python and matlab
    that make designing such a class difficult (e.g. dimensionality, indexing
    and ``length`` work fundamentally different in matlab than in python), so
    although this class currently tries to transparently support some stuff
    (notably (1D) indexing, slicing and attribute access), other operations
    (e.g. math operators and in particular __len__ and __iter__) are not yet
    supported. Don't depend on the indexing semantics not to change.

    Note:

    Assigning to parts of proxy objects (e.g. ``proxy[index].part =
    [[1,2,3]]``) should *largely* work as expected, the only exception
    would be if ``proxy.foo[index] = 3`` where ``proxy.foo[index]`` is some
    type that can be converted to python (i.e. an array or string, (or
    cell, if cell conversion has been enabled)), because then ``proxy.foo``
    returns a new python object. For these cases it's necessary to do::

      some_array[index] = 3; proxy.foo = some_array
    """
    def __init__(self, mlabwrap, name, parent=None):
        self.__dict__['_mlabwrap'] = mlabwrap
        self.__dict__['_name'] = name
        """The name is the name of the proxies representation in matlab."""
        self.__dict__['_parent'] = parent
        """To fake matlab's ``obj{foo}`` style indexing."""

    def __repr__(self):
        output = []
        mlab._do('disp(%s)' % self._name, nout=0, handle_out=output.append)
        rep = "".join(output)
        klass = self._mlabwrap._do("class(%s)" % self._name)
##         #XXX what about classes?
##         if klass == "struct":
##             rep = "\n" + self._mlabwrap._format_struct(self._name)
##         else:
##             rep = ""
        return "<%s of matlab-class: %r; internal name: %r; has parent: %s>\n%s" % (
            type(self).__name__, klass,
            self._name, ['yes', 'no'][self._parent is None],
            rep)
    def __del__(self):
        if self._parent is None:
            mlabraw.eval(self._mlabwrap._session, 'clear %s;' % self._name)
    def _get_part(self, to_get):
        if self._mlabwrap._var_type(to_get) in self._mlabwrap._mlabraw_can_convert:
            #!!! need assignment to TMP_VAL__ because `mlabraw.get` only works
            # with 'atomic' values like ``foo`` and not e.g. ``foo.bar``.
            mlabraw.eval(self._mlabwrap._session, "TMP_VAL__=%s" % to_get)
            return self._mlabwrap._get('TMP_VAL__', remove=True)
        return type(self)(self._mlabwrap, to_get, self)

    def _set_part(self, to_set, value):
        #FIXME s.a.
        if isinstance(value, MlabObjectProxy):
            mlabraw.eval(self._mlabwrap._session, "%s = %s;" % (to_set, value._name))
        else:
            self._mlabwrap._set("TMP_VAL__", value)
            mlabraw.eval(self._mlabwrap._session, "%s = TMP_VAL__;" % to_set)
            mlabraw.eval(self._mlabwrap._session, 'clear TMP_VAL__;')

    def __getattr__(self, attr):
        if attr == "_":
            return self.__dict__.setdefault('_', CurlyIndexer(self))
        else:
            return self._get_part("%s.%s" % (self._name, attr))

    def __setattr__(self, attr, value):
        self._set_part("%s.%s" % (self._name, attr), value)

    # FIXME still have to think properly about how to best translate Matlab semantics here...
    def __nonzero__(self):
        raise TypeError("%s does not yet implement truth testing" % type(self).__name__)

    def __len__(self):
        raise TypeError("%s does not yet implement __len__" % type(self).__name__)

    def __iter__(self):
        raise TypeError("%s does not yet implement iteration" % type(self).__name__)

    def _matlab_str_repr(s):
        if '\n' not in s:
            return "'%s'" % s.replace("'","''")
        else:
            # Matlab's string literals suck. They can't represent all
            # strings, so we need to use sprintf
            return "sprintf('%s')" % escape(s).replace("'","''").replace("%", "%%")

    _matlab_str_repr = staticmethod(_matlab_str_repr)

    #FIXME: those two only work ok for 1D indexing
    def _convert_index(self, index):
        if isinstance(index, int):
            return str(index + 1) # -> matlab 1-based indexing
        elif isinstance(index, basestring):
            return self._matlab_str_repr(index)
        elif isinstance(index, slice):
            if index == slice(None,None,None):
                return ":"
            elif index.step not in (None,1):
                raise ValueError("Illegal index for a proxy %r" % index)
            else:
                start = (index.start or 0) + 1
                if start == 0: start_s = 'end'
                elif start < 0: start_s = 'end%d' % start
                else: start_s = '%d' % start

                if index.stop is None: stop_s = 'end'
                elif index.stop < 0: stop_s = 'end%d' % index.stop
                else: stop_s = '%d' % index.stop

                return '%s:%s' % (start_s, stop_s)
        else:
            raise TypeError("Unsupported index type: %r." % type(index))

    def __getitem__(self, index, parens='()'):
        """WARNING: Semi-finished, semantics might change because it's not yet
           clear how to best bridge the matlab/python impedence match.
           HACK: Matlab decadently allows overloading *2* different indexing parens,
           ``()`` and ``{}``, hence the ``parens`` option."""
        index = self._convert_index(index)
        return self._get_part("".join([self._name,parens[0],index,parens[1]]))

    def __setitem__(self, index, value, parens='()'):
        """WARNING: see ``__getitem__``."""
        index = self._convert_index(index)
        return self._set_part("".join([self._name,parens[0],index,parens[1]]),
                                      value)


class MlabConversionError(Exception):
    """Raised when a Matlab type can't be converted to a python primitive."""
    pass


class MlabWrap(object):
    """An instance of this class manages a Matlab session.

    All attribute lookups on the class instance are translated into Matlab
    commands.  Python and Matlab objects are automatically converted as
    necessary.  The details of this handling can be controlled with a number of
    instance attributes, which are documented below.

    By creating multiple instances of this class, it is possible to have several
    independent Matlab sessions open simultaneously.
    """
    def __init__(self):
        """Create a new matlab(tm) wrapper object.
        """
        self._array_cast  = None
        """specifies a cast for arrays. If the result of an
        operation is a numpy array, ``return_type(res)`` will be returned
        instead."""
        self._autosync_dirs=True
        """`autosync_dirs` specifies whether the working directory of the
        matlab session should be kept in sync with that of python."""
        self._flatten_row_vecs = False
        """Automatically return 1xn matrices as flat numeric arrays."""
        self._flatten_col_vecs = False
        """Automatically return nx1 matrices as flat numeric arrays."""
        self._clear_call_args = True
        """Remove the function args from matlab workspace after each function
        call. Otherwise they are left to be (partly) overwritten by the next
        function call. This saves a function call in matlab but means that the
        memory used up by the arguments will remain unreclaimed till
        overwritten."""
        self._session = mlabraw.open(os.getenv("MLABRAW_CMD_STR", ""))
#        atexit.register(lambda handle=self._session: mlabraw.close(handle))
        self._proxies = weakref.WeakValueDictionary()
        """Use ``mlab._proxies.values()`` for a list of matlab object's that
        are currently proxied."""
        self._proxy_count = 0
        self._mlabraw_can_convert = ('double', 'char')
        """The matlab(tm) types that mlabraw will automatically convert for us."""
        self._dont_proxy = {'cell' : False}
        """The matlab(tm) types we can handle ourselves with a bit of
           effort. To turn on autoconversion for e.g. cell arrays do:
           ``mlab._dont_proxy["cell"] = True``."""

        self._logger = logging.getLogger(__name__ + '.MlabWrap')

    def sync_dirs(self):
        """Synchronize the Matlab current directory with the Python
        current directory.
        """
        # can't use self.cd, because calling it from self._do in order to
        # autosync directories will cause infinite recursion
        command = "cd('{0}')"
        self._mlabraw_eval(command.format(os.getcwd()))

    def close(self):
        """Close Matlab.
        """
        try:
            mlabraw.close(self._session)
        except AttributeError:
            pass

    def __del__(self):
        self.close()

    def __enter__(self):
        return self

    def __exit__(self, *exc_info):
        self.close()
        return False

    def _format_struct(self, varname):
        res = []
        fieldnames = self._do("fieldnames(%s)" % varname)
        size = np.ravel(self._do("size(%s)" % varname))
        return "%dx%d struct array with fields:\n%s" % (
            size[0], size[1], "\n   ".join([""] + fieldnames))
##         fieldnames
##         fieldvalues = self._do(",".join(["%s.%s" % (varname, fn)
##                                          for fn in fieldnames]), nout=len(fieldnames))
##         maxlen = max(map(len, fieldnames))
##         return "\n".join(["%*s: %s" % (maxlen, (`fv`,`fv`[:20] + '...')[len(`fv`) > 23])
##                                        for fv in fieldvalues])

    def _var_type(self, varname):
        mlabraw.eval(self._session,
                     "TMP_CLS__ = class(%(x)s); if issparse(%(x)s),"
                     "TMP_CLS__ = [TMP_CLS__,'-sparse']; end;" % dict(x=varname))
        res_type = mlabraw.get(self._session, "TMP_CLS__")
        mlabraw.eval(self._session, "clear TMP_CLS__;") # unlikely to need try/finally to ensure clear
        return res_type

    def _make_proxy(self, varname, parent=None, constructor=MlabObjectProxy):
        """Creates a proxy for a variable.

        XXX create and cache nested proxies also here.
        """
        # FIXME why not just use gensym here?
        proxy_val_name = "PROXY_VAL%d__" % self._proxy_count
        self._proxy_count += 1
        mlabraw.eval(self._session, "%s = %s;" % (proxy_val_name, varname))
        res = constructor(self, proxy_val_name, parent)
        self._proxies[proxy_val_name] = res
        return res

    def _get_cell(self, varname):
        # XXX can currently only handle ``{}`` and 1D cells
        mlabraw.eval(self._session,
                   "TMP_SIZE_INFO__ = \
                   [all(size(%(vn)s) == 0), \
                    min(size(%(vn)s)) == 1 & ndims(%(vn)s) == 2, \
                    max(size(%(vn)s))];" % {'vn':varname})
        is_empty, is_rank1, cell_len = map(int,
                                           self._get("TMP_SIZE_INFO__", remove=True).flat)
        if is_empty:
            return []
        elif is_rank1:
            cell_bits = (["TMP%i%s__" % (i, gensym('_'))
                           for i in range(cell_len)])
            mlabraw.eval(self._session, '[%s] = deal(%s{:});' %
                       (",".join(cell_bits), varname))
            # !!! this recursive call means we have to take care with
            # overwriting temps!!!
            return self._get_values(cell_bits)
        else:
            raise MlabConversionError("Not a 1D cell array")

    def _manually_convert(self, varname, vartype):
        if vartype == 'cell':
            return self._get_cell(varname)

    def _get_values(self, varnames):
        """Retrieve the variable names from Matlab and then clear them.
        """
        if not varnames:
            raise ValueError("No varnames") #to prevent clear('')
        result = []
        for varname in varnames:
            result.append(self._get(varname))
        #FIXME wrap try/finally?
        command = "clear('{0}');".format("','".join(varnames))
        self._mlabraw_eval(command)
        return result

    def _do(self, command, *args, **kwargs):
        """Semi-raw execution of a matlab command.

        Smartly handle calls to matlab, figure out what to do with `args`,
        and when to use function call syntax and not.

        If no `args` are specified, the ``cmd`` not ``result = cmd()`` form is
        used in Matlab -- this also makes literal Matlab commands legal
        (eg. cmd=``get(gca, 'Children')``).

        If ``nout=0`` is specified, the Matlab command is executed as
        procedure, otherwise it is executed as function (default), nout
        specifying how many values should be returned (default 1).

        **Beware that if you use don't specify ``nout=0`` for a `cmd` that
        never returns a value will raise an error** (because assigning a
        variable to a call that doesn't return a value is illegal in matlab).

        ``cast`` specifies which typecast should be applied to the result
        (e.g. `int`), it defaults to none.

        XXX: should we add ``parens`` parameter?
        """
        handle_out = kwargs.get('handle_out', _flush_write_stdout)
        if self._autosync_dirs:
            self.sync_dirs()
        nout =  kwargs.get('nout', 1)
        #XXX what to do with matlab screen output
        arg_names = []
        try:
            for count, arg in enumerate(args):
                if isinstance(arg, MlabObjectProxy):
                    arg_names.append(arg._name)
                else:
                    next_name = 'ARGUMENT{0}__'.format(count)
                    arg_names.append(next_name)
                    self._mlabraw_put(arg_names[-1], arg)

            if args:
                command = "{0}({1})".format(command, ','.join(arg_names))
            # three cases for nout: 0, 1, more than 1
            # 0 -> None, 1 -> val, >1 -> [val1, val2, ...]
            if nout == 0:
                handle_out(self._mlabraw_eval(command + ';'))
                return
            # deal with matlab-style multiple value return
            result_vars = ['RESULT{0}__'.format(i) for i in xrange(nout)]
            command = '[{0}] = {1};'.format(','.join(result_vars), command)
            handle_out(self._mlabraw_eval(command))
            result = self._get_values(result_vars)

            if nout == 1:
                result = result[0]
            else:
                result = tuple(result)
            if kwargs.has_key('cast'):
                if nout == 0:
                    raise TypeError("Can't cast: 0 nout")
                return kwargs['cast'](result)
            else:
                return result
        finally:
            if len(arg_names) and self._clear_call_args:
                command = "clear('{0}');".format("','".join(arg_names))
                self._mlabraw_eval(command)

    def _mlabraw_eval(self, command):
        """Log the Matlab command and then pass it to mlabraw.eval.
        """
        self._logger.info('Sending command: ' + command)
        return mlabraw.eval(self._session, command)

    def _mlabraw_put(self, var_name, value):
        """Log the Matlab variable transmittal and then pass it to
        mlabraw.put.
        """
        self._logger.info('Putting variable: ' + var_name)
        return mlabraw.put(self._session, var_name, value)

    def _mlabraw_get(self, var_name):
        """Log the Matlab variable retrieval and then pass it to
        mlabraw.get.
        """
        self._logger.info('Getting variable: ' + var_name)
        return mlabraw.get(self._session, var_name)

    # this is really raw, no conversion of [[]] -> [], whatever
    def _get(self, name, remove=False):
        """Directly access a variable in matlab space.

        This should normally not be used by user code."""
        # FIXME should this really be needed in normal operation?
        if name in self._proxies:
            return self._proxies[name]
        varname = name
        vartype = self._var_type(varname)
        if vartype in self._mlabraw_can_convert:
            var = self._mlabraw_get(varname)
            if isinstance(var, np.ndarray):
                if self._flatten_row_vecs and np.shape(var)[0] == 1:
                    var.shape = var.shape[1:2]
                elif self._flatten_col_vecs and np.shape(var)[1] == 1:
                    var.shape = var.shape[0:1]
                if self._array_cast:
                    var = self._array_cast(var)
        else:
            var = None
            if self._dont_proxy.get(vartype):
                # manual conversions may fail (e.g. for multidimensional
                # cell arrays), in that case just fall back on proxying.
                try:
                    var = self._manually_convert(varname, vartype)
                except MlabConversionError: pass
            if var is None:
                # we can't convert this to a python object, so we just
                # create a proxy, and don't delete the real matlab
                # reference until the proxy is garbage collected
                var = self._make_proxy(varname)
        if remove:
            command = "clear('{0}');".format(varname)
            self._mlabraw_eval(command)
        return var

    def _set(self, var_name, value):
        """Directly set a variable `name` in matlab space to `value`.

        This should normally not be used in user code."""
        if isinstance(value, MlabObjectProxy):
            mlabraw.eval(self._session, "%s = %s;" % (var_name, value._name))
        else:
            self._mlabraw_put(var_name, value)

    def _make_mlab_command(self, name, nout, doc=None):

        def mlab_command(*args, **kwargs):
            if 'nout' not in kwargs:
                kwargs['nout'] = nout
            return self._do(name, *args, **kwargs)
        mlab_command.__doc__ = "\n" + doc

        return mlab_command

    # XXX this method needs some refactoring, but only after it is clear how
    # things should be done (e.g. what should be extracted from docstrings and
    # how)
    def __getattr__(self, attr):
        """Magically creates a wapper to a matlab function, procedure or
        object on-the-fly."""
        if re.search(r'\W', attr): # work around ipython <= 0.7.3 bug
            raise ValueError("Attributes don't look like this: %r" % attr)
        if attr.startswith('__'):
            raise AttributeError, attr
        assert not attr.startswith('_') # XXX
        # print_ -> print
        if attr[-1] == "_":
            name = attr[:-1]
        else:
            name = attr
        try:
            nout = int(self._do("nargout('{0}')".format(name), nout=1))
        except mlabraw.error, msg:
            # determine if "name" is a Matlab script and must be called
            # with nout=0
            msg = str(msg)
            if name in msg and 'is a script' in msg:
                nout = 0
            else:
                typ = np.ravel(self._do("exist('%s')" % name))[0]
                if typ == 0: # doesn't exist
                    raise AttributeError("No such matlab object: %s" % name)
                else:
                    msg = ("Couldn't determine number of output args "
                           "for {0}, assuming 1.")
                    warnings.warn(msg.format(name))
                    nout = 1
        # nargout returns -1 for functions with a variable number of output
        # arguments, if nout not specified, assume 1
        if nout == -1:
            nout = 1
        doc = self._do("help('%s')" % name)
        mlab_command = self._make_mlab_command(name, nout, doc)
        #!!! attr, *not* name, because we might have python keyword name!
        setattr(self, attr, mlab_command)
        return mlab_command


matlab = MlabWrap
MlabError = mlabraw.error