"""$Id: memmap.py,v 1.2 2002/03/01 22:36:15 jaytmiller Exp $

This module implements a class wrapper (Memmap) around the
mmap type which adds some additional properties.

The intended use of the class(es) (Memmap, MemmapSlice) is to create a
single tier of slices from a file and to use these slices as if they
were buffers within the Numarray framework.  The slices have these
properties:

1. MemmapSlices are non-overlapping.
2. MemmapSlices are resizable.
3. MemmapSlices from the same Memmap remain "related" and affect one
anothers buffer offsets.
4. Changing the size of a MemmapSlice changes the size of the parent Memmap.

For example:

Open a memory map windowed on file 'memmap.tst'.  In practice, the
file might be (a) much bigger and (b) already existing.   Other file modes
"r" and "r+" handle existing files for readonly and read-write cases.

>>> m = open("memmap.tst","w+",len=48)
>>> m
<Memmap on file 'memmap.tst' with mode='w+', length=48, 0 slices>

In this case, file is initialized to 0.  As in C stdio, when a file is
opened for 'w+', it is truncated (effectively deleted) if it already
exists.  The initial contents of a w+ Memmap are undefined.

Slice m into the buffers "n" and "p" which will correspond to arrays:

>>> n = m[0:16]
>>> n
<MemmapSlice of length: 16>

>>> p = m[24:48]
>>> p
<MemmapSlice of length: 24>

NOTE: You *can not* make overlapping slices of a Memmap:

>>> q = m[20:28]
Traceback (most recent call last):
    File "/home/jmiller/lib/python2.2/doctest.py", line 502, in _run_examples_inner
        exec compile(source, "<string>", "single") in globs
    File "<string>", line 1, in ?
    File "memmap.py", line 358, in __getitem__
        raise IndexError("Slice overlaps previous slice of same file.")
IndexError: Slice overlaps previous slice of same file.

NOTE: you *can not* make an array directly from a Memmap:

>>> import numarray
>>> c = numarray.NumArray(buffer=m, shape=(len(m)/4,), type=numarray.Int32)
Traceback (most recent call last):
    File "<stdin>", line 1, in ?
    File "numarray.py", line 477, in __init__
        aligned=aligned)
    File "ndarray.py", line 160, in __init__
        raise TypeError("buffer won't work as NDArray buffer")
TypeError: buffer won't work as NDArray buffer

This fails because m, being the root Memmap and not a MemmapSlice, does not
define __buffer__() or resize().

Finally, the good part.  Make arrays from the MemmapSlice "buffers":

>>> a = numarray.NumArray(buffer=n, shape=(len(n)/4,), type=numarray.Int32)
>>> a[:] = 0  # Since the initial contents of 'n' are undefined.
>>> a += 1
>>> a
array([1, 1, 1, 1])

>>> b = numarray.NumArray(buffer=p, shape=(len(p)/8,), type=numarray.Float64)
>>> b[:] = 0 # Since the initial contents of 'p' are undefined.
>>> b += 2.0
>>> b
array([ 2.,  2.,  2.])

Here's the other good part about MemmapSlices:  they're resizable.

>>> a.resize( 6 )
>>> a
array([1, 1, 1, 1, 1, 1])
>>> b
array([ 2.,  2.,  2.])

What you should note is that "b" retains the correct contents (i.e., offset
within "m") even though "a" grew, effectively moving "b".  In reality, "b"
stayed where it always was and "a" has moved to a bigger RAM-based buffer.

After doing resizes, call m.flush() to synchronize the underlying file
of "m" with any RAM based slices.  This step is required to avoid
implicitly shuffling huge amounts of file space for every insert or
delete.  After calling m.flush(), all slices are once again memory
mapped rather than purely RAM based.

>>> m.flush()

NOTE: Since memory maps don't guarantee when the underlying file will be
updated with the values you have written to the map, call m.sync() when
you want to be sure your changes are on disk.

>>> m.sync()

Now "a" and "b" are both memory mapped on "memmap.tst" again.

It is also possible for "a" or "b" to shrink:

>>> a.resize(0)
>>> a
array([], type=Int32)
>>> b
array([ 2.,  2.,  2.])

Note that even if "a" or "b" goes out of scope, "m" still has references
to "n" and "p" so they will persist until "m" is deleted.

When you're done with the memory map and arrays, call m.close().  m.close()
calls m.flush() which will do consolidation if any is needed.

>>> m.close()

It's considered an error to use "m" (or slices of m) any further after
closing it.

>>> m._buffer()
Traceback (most recent call last):
  File "<stdin>", line 1, in ?
  File "memmap.py", line 239, in _buffer
    raise RuntimeError("Memmap no longer valid.  (closed?)")
RuntimeError: Memmap no longer valid.  (closed?)

And that's basically it for the top level usage of Memmap.
"""

import sys
from __builtin__ import open as __open
import os
import mmap
import types
import operator
import memory

valid_filemodes = ["r", "r+", "w+"]
writeable_filemodes = ["r+","w+"]

mode_equivalents = {
    "readonly":"r",
    "readwrite":"r+",
    "write":"w+"
    }

def _open(file, mode):
    return __open(file, mode+"b")

class Memmap:
    def __init__(self, filename, mode="r+", len=None):
        """
        Open "memmap.tst" as "copy-on-write" by specifying mode="r".

        NOTE: Currently copy-on-write is unsupported on Win-32 platforms
        due to limitations of python's mmap module.

        Here create a small test file with 100 zeroes.

        >>> _open("memmap.tst","w+").write(chr(0)*100)

        ### m = Memmap("memmap.tst","r")
        ### n = m[50:60]

        Insertions and deletions are not permitted on slices of a readonly
        Memmap, since these would be reflected onto the underlying file.

        Open "memmap.tst" for update as a memory map.  Any length
        specified should be >= the length of "memmap.tst"; the length of
        "memmap.tst" is the default.  As in C stdio, it is considered an
        error if "memmap.tst" does not already exist when opening for
        mode="r+" or "r".

        >>> m = Memmap("memmap.tst","r+")
        """
        self._filename = filename
        self._slices = []
        self._mmap = None

        if mode in mode_equivalents.keys():
            mode = mode_equivalents[mode]
        elif mode not in valid_filemodes:
            raise ValueError("mode must be one of %s" % \
                             (valid_filemodes + mode_equivalents.keys()))

        if mode == "w+" and len is None:
            raise ValueError("Must specify 'len' if mode is 'w+'")

        self._readonly  = (mode == "r")
        self._mode = mode

        if (self._readonly and os.name != "posix" and sys.version_info < (2,2,0,'beta',2)):
            raise RuntimeError(
                "Readonly Memmaps only work on POSIX platforms.")

        file = _open(filename, mode)
        file.seek(0, 2)
        flen = file.tell()

        if len is None:
            len = flen

        if mode == "w+" or (mode == "r+" and flen < len):
            if len:
                file.seek(len-1, 0) # seek to the last byte and write it.
                file.write(chr(0))
            flen = len

        if len:
            if sys.version_info < (2,2,0,'beta',2):
                if os.name == "posix":
                    if mode == "r":
                        flags = mmap.MAP_PRIVATE
                    else:
                        flags = mmap.MAP_SHARED
                        prot = mmap.PROT_READ | mmap.PROT_WRITE
                        self._mmap = mmap.mmap(file.fileno(), len, flags, prot)
                else:
                    self._mmap = mmap.mmap(file.fileno(), len)
            else:
                if mode == "r":
                    acc = mmap.ACCESS_COPY
                else:
                    acc = mmap.ACCESS_WRITE
                self._mmap = mmap.mmap(file.fileno(), len, access=acc)
        else:
            self._mmap = None

        file.close()

    def __repr__(self):
        return "<Memmap on file '%s' with mode='%s', length=%d, %d slices>" % \
            (self._filename, self._mode, len(self), len(self._slices))

    def close(self):
        """
        close(self) unites the memory map and any RAM based slices
        with its underlying file and removes the mapping and all
        references to its slices.  Don't call this till you're done
        using the Memmap!
        """
        if not self._readonly:
            self.flush()
        self._lose_map()

    def _buffer(self, begin=0, end=None):
        """_buffer(self) returns a buffer object for Memmap 'self'.
        """
        if self._mmap is None:
            raise RuntimeError("Memmap no longer valid.  (closed?)")
        if end is None:
            end = len(self)
        return memory.writeable_buffer(self._mmap, begin, end-begin)

    def _chkOverlaps(self, begin, end):
        """_chkOverlaps(self, begin, end) is called to raise an exception
        if the requested slice will overlap any slices which have already
        been taken.
        """
        for b,e,obj in self._slices:
            if b < begin < e or \
               b < end < e  or \
               begin < b < end or \
               begin < e < end or \
               b == begin and e == end and b != e:
                raise IndexError("Slice overlaps previous slice of same file.")

    def __len__(self):
        """len(self) is the number of bytes in Memmap 'self'.
        """
        if self._mmap:
            maplen = len(self._mmap)
        else:
            maplen = 0
        orig_len = reduce(operator.add, [ e-b for b,e,o in self._slices ], 0)
        obj_len = reduce(operator.add, [ len(o) for b,e,o in self._slices], 0)
        return int(maplen - orig_len + obj_len)

    def __str__(self):
        self.flush()
        if self._mmap is not None:
            return self._mmap[:]
        else:
            return ""

    def _fix_slice(self, i):
        """_fix_slice(self, i) converts a __getitem__ 'key' into slice
        parameters, and returns a tuple (beg, end).
        """
        if isinstance(i, types.SliceType):
            if i.step is not None:
                raise IndexError("Memmap does not support strides.")
            j, i = i.stop, i.start
        else:
            j = i+1
        i, j = self._chkIndex(i, 1), self._chkIndex(j, 1)
        return i, j

    def _chkIndex(self, i, isSlice=0):
        """_chkIndex(self, i) raises an IndexError if 'i' is not a valid
        index of 'self'
        """
        if i == sys.maxint:  # Assume i not maxint unless it's a slice stop
            return len(self)
        if i < 0:
            i += len(self)
        if not(0 <= i < len(self)+isSlice):
            raise IndexError("Invalid Memmap index: %d" % (i,))
        return i

    def __getitem__(self,i):
        """__getitem__(self,i) returns a MemmapSlice corresponding to the
        index 'i' of the Memmap 'self'.  The Memmap keeps a record of the
        slice so that it can coordinate it with other slices from the same
        file.  Slices of a Memmap are not permitted to overlap.
        """
        i, j = self._fix_slice(i)
        self._chkOverlaps(i, j)
        obj = MemmapSlice(self._buffer(i, j), self._readonly)
        self._slices.append((i, j, obj))
        return obj

    def __delitem__(self,i):
        """__delitem__(self,i) deletes a slice from a Memmap, removing the
        record of the "slice", but not the underlying data footprint.
        """
        i, j = self._fix_slice(i)
        for k in range(len(self._slices)):
            b,e,o = self._slices[k]
            if b==i and e==j:
                del self._slices[k]
                return
        else:
            raise ValueError("Can't find slice (%d,%d)" % (i,j))

    def _writeCheck(self):
        if self._readonly:
            raise TypeError("Attempt to modify a readonly Memmap")

    def sync(self):
        """sync(self) ultimately calls msync, guaranteeing that updates
        to a MMap are already written to the underlying file system
        device when the call returns.
        """
        if self._mmap is not None:
            self._mmap.flush()

    def _dirty(self):
        """_dirty(self) is 1 if any slice of self is "dirty".  A slice is
        dirty if it has been resized in any way, or was not part of the
        original Memmap.  _dirty(self) specifically excludes in-place
        modification, since this can happen at the C-level and there's
        no way to know whether it has happened or not.
        """
        return reduce( operator.or_, [ o.dirty() for b,e,o in self._slices ],
                       0 )

    def _lose_map(self):
        """_lose_map(self) eliminates all references to the underlying mmap,
        so that it will be deleted.  This appeared necessary on Win-NT to be
        able to re-map the same file.
        """
        for b,e,o in self._slices:
            o._rebuffer(None)
        if self._mmap is not None:
            self._mmap.close()
            self._mmap = None

    def _consolidate(self):
        """_consolidate(self) re-writes the memory map file, interspersing
        RAM based slices with the content of the mmap which has been updated
        in-place.  The new memory mapped file is then mapped in place of the
        old one.
        """
        temp_map = "Memmap.tmp"
        f = _open(temp_map, "w+")
        mlen = len(self)
        l = self._slices
        self._slices = []
        l.sort()
        m = 0
        for b, e, obj in l:
            if b > m:       # copy original mmap between slices
                f.write(self._buffer(m,b))
            ob = f.tell()
            f.write(obj.__buffer__())
            oe = f.tell()
            self._slices.append((ob,oe,obj))
            m = e
        f.write(self._buffer(m, mlen))
        f.close()
        self._lose_map()
        os.remove(self._filename);
        os.rename(temp_map, self._filename)
        f = _open(self._filename,"r+")
        self._mmap = mmap.mmap(f.fileno(), mlen)
        for b,e,o in self._slices:
            o._rebuffer(self._buffer(b, e))
        f.close()

    def flush(self):
        """flush(self) flushes the memory map, consolidates it with any
        RAM based slices if required, and re-maps it.  Both slice
        offsets and buffers change.  If there are no RAM based
        buffers, no consolidation is required.
        """
        self._writeCheck()
        if self._mmap is not None:
            self._mmap.flush()
        if self._dirty():
            self._consolidate()

    def find(self, string, offset=0):
        """find(string, offset=0) returns the first index at which string
        is found, or -1 on failure.

        NOTE: find is known not to work on Python systems earlier than
        v2.2 due to a bug in mmapmodule.c

        ### _open("memmap.tst","w+").write("this is a test")
        ### Memmap("memmap.tst",len=14).find("is")
        2
        ### Memmap("memmap.tst",len=14).find("is", 3)
        5
        ### _open("memmap.tst","w+").write("x")
        ### Memmap("memmap.tst",len=1).find("is")
        -1
        """
        if self._mmap is None:
            raise RuntimeError("_mmap is None;  Memmap closed?")
        else:
            return self._mmap.find(string, offset)

    def move(self, dest, src, count):
        """move(dest, src, count) moves 'count' characters from 'src' to 'dest'
        within a Memmap.

        >>>
        """
        # self._writeCheck()  # copy-on-write
        self._buffer()[dest:dest+count] = self._buffer()[src:src+count]

    # Method for adding new slices in between existing slices.
    def _insert(self, offset, bufferable):
        """
        _insert(self, offset, buffer) inserts a new object meeting the buffer
        api at offset, and returns the resulting MemmapSlice.
        """
        # self._writeCheck()  # copy-on-write
        self._chkOverlaps(offset, offset)
        self._chkIndex(offset, isSlice=1)
        obj = MemmapSlice(memory.writeable_buffer(bufferable), dirty=1)
        self._slices.append((offset, offset, obj))
        return obj

    def insert(self, offset, size):
        """
        insert(self, offset, size) inserts a new slice of 'size'
        bytes, possibly between two adjacent slices, at byte 'offset'.
        It is not legal to insert into the middle of another slice, but
        pretty much everything else is legal.  The resulting
        MemmapSlice is returned.

        >>> m = open("memmap.tst",mode="w+",len=100)
        >>> n = m[0:50]
        >>> p = m[50:100]
        >>> q=m.insert(0, 200)
        >>> r=m.insert(50, 100)
        >>> s=m.insert(100, 300)
        >>> t=m.insert(45, 100)
        Traceback (most recent call last):
        File "/home/jmiller/lib/python2.2/doctest.py", line 502, in _run_examples_inner
        exec compile(source, "<string>", "single") in globs
        File "<string>", line 1, in ?
        File "memmap.py", line 440, in insert
        return self._insert(offset, memory.memory_buffer(size))
        File "memmap.py", line 413, in _insert
        self._chkOverlaps(offset, offset)
        File "memmap.py", line 255, in _chkOverlaps
        raise IndexError("Slice overlaps previous slice of same file.")
        IndexError: Slice overlaps previous slice of same file.
        >>> m.flush()
        >>> len(m)
        700
        >>> m.close()
        """
        return self._insert(offset, memory.memory_buffer(size))


    def __del__(self):
        self.close()


class MemmapSlice:
    def __init__(self, buffer, dirty=0, readonly=0):
        self._buffer = buffer
        self._dirty = dirty
        self._readonly = readonly

    def __repr__(self):
        return "<MemmapSlice of length: %d>" % len(self)

    def __len__(self):
        if self._buffer:
            return len(self.__buffer__())
        else:
            return 0

    def dirty(self):
        """dirty(self, set=None) is 1 iff 'self' has changed its buffer since it
        was created.
        """
        return self._dirty

    def __getitem__(self,i):
        if type(i) is types.IntType:
            return self.__buffer__()[i]
        elif type(i) is types.SliceType:
            return self.__buffer__()[i.start:i.stop]

    def __buffer__(self):
        if self._buffer is not None:
            return self._buffer
        else:
            raise RuntimeError("MemmapSlice no longer valid...(Memmap closed?)")

    def _rebuffer(self, b):
        self._buffer = b
        self._dirty = 0

    def _modify_buffer(self, offset, size):
        """_modify_buffer(self, offset, size) replaces the slice's mmap
        buffer with a resized RAM buffer, and copies the contents.
        """
        self._dirty = 1
        # self._writeCheck()  # copy-on-write
        olen = len(self)
        nlen = olen+size
        newMaster = memory.memory_buffer(nlen)
        newMaster[0:offset] = self._buffer[0:offset]
        if size > 0:
            newMaster[offset+size:] = self._buffer[offset:]
        else:
            newMaster[offset:] = self._buffer[offset-size:]
        self._buffer = newMaster

    def _insert(self, offset, size):
        """_insert(self, offset, size) expands the MMap at 'offset'
        by 'size' bytes.

        'offset' refers to a position between two existing characters, the
        beginning, or the end.
        """
        # self._writeCheck()  # copy-on-write
        # Since insertion points aren't indices, tolerate the end
        self._chkIndex(offset, 1)
        if offset + size > sys.maxint:
            raise ValueError("Insert makes file too big for integer offsets")
        self._modify_buffer(offset, size)

    PADC = chr(0)
    def insert(self, offset, value, size=None, padc=PADC):
        """insert(self, offset, value, size=None) inserts string 'value' at
        'offset', possibly padding it with extra characters of value 'padc'.
        If size is None, the the size of the insert is len(value).
        """
        l = len(value)
        if size is None:
            size = l
        elif l < size:
            value += padc * (size-l)
        elif l > size:
            raise ValueError("'value' too long for 'size'")
        self._insert(offset, size)
        self.__buffer__()[offset:offset+size] = value

    def _append(self, size):
        """append(self, size) is similar to 'insert', but assumes the offset is
        the end of the current slice.
        """
        self._insert(len(self), size)

    def append(self, value):
        size = len(value)
        self._append(size)
        self.__buffer__()[-size:] = value

    def delete(self, offset, size):
        """delete(self, offset, size) removes 'size' bytes from the MMap,
        starting at 'offset'.

        'offset' refers to a position between two existing characters, the
        beginning, or the end.
        """
        self._chkIndex(offset+size, 1)
        self._insert(offset, -size)

    def truncate(self, size):
        """truncate(self, size) is similar to 'delete', but assumes the offset
        is the end of the current slice.
        """
        self.delete(len(self)-size, size)

    def resize(self, newsize):
        """resize(self, newsize) appends or truncates the MemmapSlice to
        'newsize'.  Any newly added region is uninitialized.
        """
        olen = len(self)
        if newsize > olen:
            self._append( newsize-olen )
        elif newsize < olen:
            if newsize < 0:
                raise ValueError("Negative resize value")
            self.truncate(olen - newsize)

    def flush(self):
        """flush(self)
        """
        raise ValueError("Only the 'root' Memmap should be flushed.")

    def _writeCheck(self):
        if self._readonly:
            raise TypeError("Attempt to modify a readonly MemmapSlice.")

    def _chkIndex(self, i, End=0):
        """_chkIndex(self, i) raises an IndexError if 'i' is not a valid
        index of 'self'
        """
        olen = len(self)
        if i == sys.maxint:  # Assume i not maxint unless it's a slice stop
            return olen
        if i < 0:
            i += olen
        if not(0 <= i < olen+End):
            raise IndexError("Invalid Memmap index: %d" % (i,))
        return i

    def __str__(self):
        return str(self.__buffer__())

def open(filename, mode="r+", len=None):
    """open(filename, mode="r+", len=None) creates a new Memmap object.
    """
    return Memmap(filename, mode, len)

def close(map):
    return map.close()

def test():
    """
    >>> import os
    >>> os.remove("memmap.tst")
    """
    import doctest, memmap
    return doctest.testmod(memmap)

if __name__ == "__main__":
    test()