path: root/src/snakeoil/data_source.py

"""
A far more minimal form of file protocol encapsulation location and encoding into it

The primary use for data_source's is to encapsulate the following issues into a single object:

* is the data actually on disk (thus can I use more efficient ops against the file?).
* what is the preferred encoding?
* py3k compatibility concerns (bytes versus text file handles)

Note that all file like handles returned from `text_fileobj()` and `bytes_fileobj()`
have a required additional attribute- *exceptions*, either a single Exception class, or a
tuple of Exception classes that can be thrown by that file handle during usage.

This requirement exists purely to allow the consuming code to avoid having to know anything
about the backing of the file like object.

The proper way to use such a filehandle is as follows:

>>> from snakeoil.data_source import data_source
>>> source = data_source("It's a fez. I wear a fez now. Fezes are cool.", mutable=False)
>>> handle = source.text_fileobj()
>>> handle.write("You graffitied the oldest cliff face in the universe.")
Traceback (most recent call last):
TypeError:
>>> # if this where a normal file, it would be an IOError- it's impossible to guess the
>>> # correct exception to intercept, so instead we rely on the handle telling us what
>>> # we should catch;
>>> try:
...   handle.write("You wouldn't answer your phone.")
... except handle.exceptions as e:
...   print("we caught the exception.")
we caught the exception.
"""

__all__ = (
    "base", "bz2_source", "data_source", "local_source", "text_data_source",
    "bytes_data_source", "invokable_data_source",
)

import errno
from functools import partial
import io

from . import compression, fileutils, klass, stringio
from .currying import post_curry


def _mk_writable_cls(base, name):
    """
    inline mixin of writable overrides

    while a normal mixin is preferable, this is required due to
    differing slot layouts between py2k/py3k base classes of
    stringio.
    """

    class kls(base):
        __doc__ = f"""
        writable {name.split("_")[0]} StringIO instance suitable for usage as a data_source filehandle

        This adds a callback for updating the original data source, and appropriate
        exceptions attribute
        """


        base_cls = base
        exceptions = (MemoryError,)
        __slots__ = ('_callback',)

        def __init__(self, callback, data):
            """
            :param callback: functor invoked when this data source is modified;
                the functor takes a single value, the full content of the StringIO
            :param data: initial data for this instance
            """
            if not callable(callback):
                raise TypeError("callback must be callable")
            self.base_cls.__init__(self, data)
            self._callback = callback

        def close(self):
            self.flush()
            if self._callback is not None:
                self.seek(0)
                self._callback(self.read())
                self._callback = None
            self.base_cls.close(self)
    kls.__name__ = name
    return kls


text_wr_StringIO = _mk_writable_cls(io.StringIO, "text_wr_StringIO")
bytes_wr_StringIO = _mk_writable_cls(io.BytesIO, "bytes_wr_StringIO")


class text_ro_StringIO(stringio.text_readonly):
    """
    readonly text mode StringIO usable as a filehandle for a data_source

    Specifically this adds the necessary `exceptions` attribute; see
    :py:class:`snakeoil.stringio.text_readonly` for methods details.
    """
    __slots__ = ()
    exceptions = (MemoryError, TypeError)


class bytes_ro_StringIO(stringio.bytes_readonly):
    """
    readonly bytes mode StringIO usable as a filehandle for a data_source

    Specifically this adds the necessary `exceptions` attribute; see
    :py:class:`snakeoil.stringio.bytes_readonly` for methods details.
    """
    __slots__ = ()
    exceptions = (MemoryError, TypeError)


# derive our file classes- we derive *strictly* to append
# the exceptions class attribute for consumer usage.
def open_file(*args, **kwds):
    handle = open(*args, **kwds)
    handle.exceptions = (EnvironmentError,)
    return handle


class base:
    """
    base data_source class; implementations of the protocol are advised
    to derive from this.

    :ivar path: If None, no local path is available- else it's the ondisk path to
      the data
    """
    __slots__ = ("weakref",)

    path = None

    def text_fileobj(self, writable=False):
        """get a text level filehandle for for this data

        :param writable: whether or not we need to write to the handle
        :raise: TypeError if immutable and write is requested
        :return: file handle like object
        """
        raise NotImplementedError(self, "text_fileobj")

    def bytes_fileobj(self, writable=False):
        """get a bytes level filehandle for for this data

        :param writable: whether or not we need to write to the handle
        :raise: TypeError if immutable and write is requested
        :return: file handle like object
        """
        raise NotImplementedError(self, "bytes_fileobj")

    def transfer_to_path(self, path):
        return self.transfer_to_data_source(
            local_source(path, mutable=True, encoding=None))

    def transfer_to_data_source(self, write_source):
        read_f, m, write_f = None, None, None
        try:
            write_f = write_source.bytes_fileobj(True)
            if self.path is not None:
                m, read_f = fileutils.mmap_or_open_for_read(self.path)
            else:
                read_f = self.bytes_fileobj()

            if read_f is not None:
                transfer_between_files(read_f, write_f)
            else:
                write_f.write(m)
        finally:
            for x in (read_f, write_f, m):
                if x is None:
                    continue
                try:
                    x.close()
                except EnvironmentError:
                    pass


class local_source(base):

    """locally accessible data source

    Literally a file on disk.
    """

    __slots__ = ("path", "mutable", "encoding")

    buffering_window = 32768

    def __init__(self, path, mutable=False, encoding=None):
        """
        :param path: file path of the data source
        :param mutable: whether this data_source is considered modifiable or not
        :param encoding: the text encoding to force, if any
        """
        base.__init__(self)
        self.path = path
        self.mutable = mutable
        self.encoding = encoding

    @klass.steal_docs(base)
    def text_fileobj(self, writable=False):
        if writable and not self.mutable:
            raise TypeError("data source %s is immutable" % (self,))
        if self.encoding:
            opener = open_file
            opener = post_curry(opener, buffering=self.buffering_window,
                                encoding=self.encoding)
        else:
            opener = post_curry(open_file, self.buffering_window)
        if not writable:
            return opener(self.path, 'r')
        try:
            return opener(self.path, "r+")
        except IOError as ie:
            if ie.errno != errno.ENOENT:
                raise
            return opener(self.path, 'w+')

    @klass.steal_docs(base)
    def bytes_fileobj(self, writable=False):
        if not writable:
            return open_file(self.path, 'rb', self.buffering_window)
        if not self.mutable:
            raise TypeError("data source %s is immutable" % (self,))
        try:
            return open_file(self.path, 'rb+', self.buffering_window)
        except IOError as ie:
            if ie.errno != errno.ENOENT:
                raise
            return open_file(self.path, 'wb+', self.buffering_window)


class bz2_source(base):
    """
    locally accessible bz2 archive

    Literally a bz2 file on disk.
    """

    __slots__ = ("path", "mutable")

    def __init__(self, path, mutable=False):
        """
        :param path: file path of the data source
        :param mutable: whether this data source is considered modifiable or not
        """
        base.__init__(self)
        self.path = path
        self.mutable = mutable

    def text_fileobj(self, writable=False):
        data = compression.decompress_data(
            'bzip2', fileutils.readfile_bytes(self.path)).decode()
        if writable:
            if not self.mutable:
                raise TypeError(f"data source {self} is not mutable")
            return text_wr_StringIO(self._set_data, data)
        return text_ro_StringIO(data)

    def bytes_fileobj(self, writable=False):
        data = compression.decompress_data(
            'bzip2', fileutils.readfile_bytes(self.path))
        if writable:
            if not self.mutable:
                raise TypeError(f"data source {self} is not mutable")
            return bytes_wr_StringIO(self._set_data, data)
        return bytes_ro_StringIO(data)

    def _set_data(self, data):
        if isinstance(data, str):
            data = data.encode()
        with open(self.path, "wb") as f:
            f.write(compression.compress_data('bzip2', data))


class data_source(base):

    """
    base class encapsulating a purely virtual data source lacking an on disk location.

    Whether this be due to transformation steps necessary (pulling the data out of
    an archive for example), or the data being generated on the fly, this classes's
    derivatives :py:class:`text_data_source` and :py:class:`bytes_data_source` are
    likely what you should be using for direct creation.

    :ivar data: the raw data- should either be a string or bytes depending on your
      derivative
    :ivar path: note that path is None for this class- no on disk location available.
    """

    __slots__ = ('data', 'mutable')

    def __init__(self, data, mutable=False):
        """
        :param data: data to wrap
        :param mutable: should this data_source be updatable?
        """
        base.__init__(self)
        self.data = data
        self.mutable = mutable

    def _convert_data(self, mode):
        if mode == 'bytes':
            if isinstance(self.data, bytes):
                return self.data
            return self.data.encode()
        if isinstance(self.data, str):
            return self.data
        return self.data.decode()

    @klass.steal_docs(base)
    def text_fileobj(self, writable=False):
        if writable:
            if not self.mutable:
                raise TypeError(f"data source {self} is not mutable")
            return text_wr_StringIO(self._reset_data,
                                    self._convert_data('text'))
        return text_ro_StringIO(self._convert_data('text'))

    def _reset_data(self, data):
        if isinstance(self.data, bytes):
            if not isinstance(data, bytes):
                data = data.encode()
        elif not isinstance(data, str):
            data = data.decode()
        self.data = data

    @klass.steal_docs(base)
    def bytes_fileobj(self, writable=False):
        if writable:
            if not self.mutable:
                raise TypeError(f"data source {self} is not mutable")
            return bytes_wr_StringIO(self._reset_data,
                                     self._convert_data('bytes'))
        return bytes_ro_StringIO(self._convert_data('bytes'))


class text_data_source(data_source):
    """Text data source.

    This does auto-conversion between bytes/text as needed.
    """

    __slots__ = ()

    @klass.steal_docs(data_source)
    def __init__(self, data, mutable=False):
        if not isinstance(data, str):
            raise TypeError("data must be a str")
        data_source.__init__(self, data, mutable=mutable)

    def _convert_data(self, mode):
        if mode != 'bytes':
            return self.data
        return self.data.encode()


class bytes_data_source(data_source):
    """Bytes data source.

    This does auto-conversion between bytes/text as needed.
    """

    __slots__ = ()

    @klass.steal_docs(data_source)
    def __init__(self, data, mutable=False):
        if not isinstance(data, bytes):
            raise TypeError("data must be bytes")
        data_source.__init__(self, data, mutable=mutable)

    def _convert_data(self, mode):
        if mode == 'bytes':
            return self.data
        return self.data.decode()


class invokable_data_source(data_source):

    """
    data source that takes a callable instead of the actual data item

    The callable takes a single argument- a boolean, True if a text fileobj
    is requested, False if None

    Note that this instance is explicitly readonly.
    """
    __slots__ = ()

    def __init__(self, data):
        """
        :param data: callable that accepts one argument- True if a text
          file obj was requested, False if a bytes file obj is requested.
        """
        data_source.__init__(self, data, mutable=False)

    @klass.steal_docs(data_source)
    def text_fileobj(self, writable=False):
        if writable:
            raise TypeError(f"data source {self} data is immutable")
        return self.data(True)

    @klass.steal_docs(data_source)
    def bytes_fileobj(self, writable=False):
        if writable:
            raise TypeError(f"data source {self} data is immutable")
        return self.data(False)

    @classmethod
    def wrap_function(cls, invokable, returns_text=True, returns_handle=False, encoding_hint=None):
        """
        Helper function to automatically convert a function that returns text or bytes into appropriate
        callable

        :param invokable: a callable that returns either text, or bytes, taking no args
        :param returns_text: True if the data returned is text/basestring, False if Not
        :param returns_handle: True if the object returned is a handle, False if not.  Note that returns_text
            still has meaning here- returns_text indicates what sort of data the handle returns from read
            invocations.
        :param encoding_hint: the preferred encoding to use for encoding
        :return: invokable_data_source instance
        """
        return cls(partial(cls._simple_wrapper, invokable, encoding_hint, returns_text, returns_handle))

    @staticmethod
    def _simple_wrapper(invokable, encoding_hint, returns_text, returns_handle, text_wanted):
        data = invokable()
        if returns_text != text_wanted:
            if text_wanted:
                if returns_handle:
                    data = data.read()
                if encoding_hint:
                    # we have an encoding, its bytes data, and text is wanted
                    data = data.decode(encoding_hint)
                else:
                    data = data.decode()
            else:
                # bytes were wanted...
                if returns_handle:
                    # pull in the data...
                    data = data.read()
                if encoding_hint is None:
                    # fallback to utf8
                    encoding_hint = 'utf8'
                data = data.encode(encoding_hint)
        elif returns_handle:
            return data
        if text_wanted:
            return text_ro_StringIO(data)
        return bytes_ro_StringIO(data)


def transfer_between_files(read_file, write_file, bufsize=(32 * 1024)):
    while data := read_file.read(bufsize):
        write_file.write(data)