#!/usr/bin/env python
# This file is distributed under the terms of the 2-clause BSD License.
# Copyright (c) 2017-2018, Almar Klein

"""
Python implementation of the Binary Structured Data Format (BSDF).

BSDF is a binary format for serializing structured (scientific) data.
See http://bsdf.io for more information.

This is the reference implementation, which is relatively relatively
sophisticated, providing e.g. lazy loading of blobs and streamed
reading/writing. A simpler Python implementation is available as
``bsdf_lite.py``.

This module has no dependencies and works on Python 2.7 and 3.4+.

Note: on Legacy Python (Python 2.7), non-Unicode strings are encoded as bytes.
"""

# todo: in 2020, remove six stuff, __future__ and _isidentifier
# todo: in 2020, remove 'utf-8' args to encode/decode; it's faster

from __future__ import absolute_import, division, print_function

import bz2
import hashlib
import logging
import os
import re
import struct
import sys
import types
import zlib
from io import BytesIO

logger = logging.getLogger(__name__)

# Notes on versioning: the major and minor numbers correspond to the
# BSDF format version. The major number if increased when backward
# incompatible changes are introduced. An implementation must raise an
# exception when the file being read has a higher major version. The
# minor number is increased when new backward compatible features are
# introduced. An implementation must display a warning when the file
# being read has a higher minor version. The patch version is increased
# for subsequent releases of the implementation.
VERSION = 2, 1, 2
__version__ = ".".join(str(i) for i in VERSION)


# %% The encoder and decoder implementation

# From six.py
PY3 = sys.version_info[0] >= 3
if PY3:
    text_type = str
    string_types = str
    unicode_types = str
    integer_types = int
    classtypes = type
else:  # pragma: no cover
    logging.basicConfig()  # avoid "no handlers found" error
    text_type = unicode  # noqa
    string_types = basestring  # noqa
    unicode_types = unicode  # noqa
    integer_types = (int, long)  # noqa
    classtypes = type, types.ClassType

# Shorthands
spack = struct.pack
strunpack = struct.unpack


def lencode(x):
    """ Encode an unsigned integer into a variable sized blob of bytes.
    """
    # We could support 16 bit and 32 bit as well, but the gain is low, since
    # 9 bytes for collections with over 250 elements is marginal anyway.
    if x <= 250:
        return spack("<B", x)
    # elif x < 65536:
    #     return spack('<BH', 251, x)
    # elif x < 4294967296:
    #     return spack('<BI', 252, x)
    else:
        return spack("<BQ", 253, x)


# Include len decoder for completeness; we've inlined it for performance.
def lendecode(f):
    """ Decode an unsigned integer from a file.
    """
    n = strunpack("<B", f.read(1))[0]
    if n == 253:
        n = strunpack("<Q", f.read(8))[0]  # noqa
    return n


def encode_type_id(b, ext_id):
    """ Encode the type identifier, with or without extension id.
    """
    if ext_id is not None:
        bb = ext_id.encode("UTF-8")
        return b.upper() + lencode(len(bb)) + bb  # noqa
    else:
        return b  # noqa


def _isidentifier(s):  # pragma: no cover
    """ Use of str.isidentifier() for Legacy Python, but slower.
    """
    # http://stackoverflow.com/questions/2544972/
    return (
        isinstance(s, string_types)
        and re.match(r"^\w+$", s, re.UNICODE)
        and re.match(r"^[0-9]", s) is None
    )


class BsdfSerializer(object):
    """ Instances of this class represent a BSDF encoder/decoder.

    It acts as a placeholder for a set of extensions and encoding/decoding
    options. Use this to predefine extensions and options for high
    performance encoding/decoding. For general use, see the functions
    `save()`, `encode()`, `load()`, and `decode()`.

    This implementation of BSDF supports streaming lists (keep adding
    to a list after writing the main file), lazy loading of blobs, and
    in-place editing of blobs (for streams opened with a+).

    Options for encoding:

    * compression (int or str): ``0`` or "no" for no compression (default),
      ``1`` or "zlib" for Zlib compression (same as zip files and PNG), and
      ``2`` or "bz2" for Bz2 compression (more compact but slower writing).
      Note that some BSDF implementations (e.g. JavaScript) may not support
      compression.
    * use_checksum (bool): whether to include a checksum with binary blobs.
    * float64 (bool): Whether to write floats as 64 bit (default) or 32 bit.

    Options for decoding:

    * load_streaming (bool): if True, and the final object in the structure was
      a stream, will make it available as a stream in the decoded object.
    * lazy_blob (bool): if True, bytes are represented as Blob objects that can
      be used to lazily access the data, and also overwrite the data if the
      file is open in a+ mode.
    """

    def __init__(self, extensions=None, **options):
        self._extensions = {}  # name -> extension
        self._extensions_by_cls = {}  # cls -> (name, extension.encode)
        if extensions is None:
            extensions = standard_extensions
        for extension in extensions:
            self.add_extension(extension)
        self._parse_options(**options)

    def _parse_options(
        self,
        compression=0,
        use_checksum=False,
        float64=True,
        load_streaming=False,
        lazy_blob=False,
    ):

        # Validate compression
        if isinstance(compression, string_types):
            m = {"no": 0, "zlib": 1, "bz2": 2}
            compression = m.get(compression.lower(), compression)
        if compression not in (0, 1, 2):
            raise TypeError("Compression must be 0, 1, 2, " '"no", "zlib", or "bz2"')
        self._compression = compression

        # Other encoding args
        self._use_checksum = bool(use_checksum)
        self._float64 = bool(float64)

        # Decoding args
        self._load_streaming = bool(load_streaming)
        self._lazy_blob = bool(lazy_blob)

    def add_extension(self, extension_class):
        """ Add an extension to this serializer instance, which must be
        a subclass of Extension. Can be used as a decorator.
        """
        # Check class
        if not (
            isinstance(extension_class, type) and issubclass(extension_class, Extension)
        ):
            raise TypeError("add_extension() expects a Extension class.")
        extension = extension_class()

        # Get name
        name = extension.name
        if not isinstance(name, str):
            raise TypeError("Extension name must be str.")
        if len(name) == 0 or len(name) > 250:
            raise NameError(
                "Extension names must be nonempty and shorter " "than 251 chars."
            )
        if name in self._extensions:
            logger.warn(
                'BSDF warning: overwriting extension "%s", '
                "consider removing first" % name
            )

        # Get classes
        cls = extension.cls
        if not cls:
            clss = []
        elif isinstance(cls, (tuple, list)):
            clss = cls
        else:
            clss = [cls]
        for cls in clss:
            if not isinstance(cls, classtypes):
                raise TypeError("Extension classes must be types.")

        # Store
        for cls in clss:
            self._extensions_by_cls[cls] = name, extension.encode
        self._extensions[name] = extension
        return extension_class

    def remove_extension(self, name):
        """ Remove a converted by its unique name.
        """
        if not isinstance(name, str):
            raise TypeError("Extension name must be str.")
        if name in self._extensions:
            self._extensions.pop(name)
        for cls in list(self._extensions_by_cls.keys()):
            if self._extensions_by_cls[cls][0] == name:
                self._extensions_by_cls.pop(cls)

    def _encode(self, f, value, streams, ext_id):
        """ Main encoder function.
        """
        x = encode_type_id

        if value is None:
            f.write(x(b"v", ext_id))  # V for void
        elif value is True:
            f.write(x(b"y", ext_id))  # Y for yes
        elif value is False:
            f.write(x(b"n", ext_id))  # N for no
        elif isinstance(value, integer_types):
            if -32768 <= value <= 32767:
                f.write(x(b"h", ext_id) + spack("h", value))  # H for ...
            else:
                f.write(x(b"i", ext_id) + spack("<q", value))  # I for int
        elif isinstance(value, float):
            if self._float64:
                f.write(x(b"d", ext_id) + spack("<d", value))  # D for double
            else:
                f.write(x(b"f", ext_id) + spack("<f", value))  # f for float
        elif isinstance(value, unicode_types):
            bb = value.encode("UTF-8")
            f.write(x(b"s", ext_id) + lencode(len(bb)))  # S for str
            f.write(bb)
        elif isinstance(value, (list, tuple)):
            f.write(x(b"l", ext_id) + lencode(len(value)))  # L for list
            for v in value:
                self._encode(f, v, streams, None)
        elif isinstance(value, dict):
            f.write(x(b"m", ext_id) + lencode(len(value)))  # M for mapping
            for key, v in value.items():
                if PY3:
                    assert key.isidentifier()  # faster
                else:  # pragma: no cover
                    assert _isidentifier(key)
                # yield ' ' * indent + key
                name_b = key.encode("UTF-8")
                f.write(lencode(len(name_b)))
                f.write(name_b)
                self._encode(f, v, streams, None)
        elif isinstance(value, bytes):
            f.write(x(b"b", ext_id))  # B for blob
            blob = Blob(
                value, compression=self._compression, use_checksum=self._use_checksum
            )
            blob._to_file(f)  # noqa
        elif isinstance(value, Blob):
            f.write(x(b"b", ext_id))  # B for blob
            value._to_file(f)  # noqa
        elif isinstance(value, BaseStream):
            # Initialize the stream
            if value.mode != "w":
                raise ValueError("Cannot serialize a read-mode stream.")
            elif isinstance(value, ListStream):
                f.write(x(b"l", ext_id) + spack("<BQ", 255, 0))  # L for list
            else:
                raise TypeError("Only ListStream is supported")
            # Mark this as *the* stream, and activate the stream.
            # The save() function verifies this is the last written object.
            if len(streams) > 0:
                raise ValueError("Can only have one stream per file.")
            streams.append(value)
            value._activate(f, self._encode, self._decode)  # noqa
        else:
            if ext_id is not None:
                raise ValueError(
                    "Extension %s wronfully encodes object to another "
                    "extension object (though it may encode to a list/dict "
                    "that contains other extension objects)." % ext_id
                )
            # Try if the value is of a type we know
            ex = self._extensions_by_cls.get(value.__class__, None)
            # Maybe its a subclass of a type we know
            if ex is None:
                for name, c in self._extensions.items():
                    if c.match(self, value):
                        ex = name, c.encode
                        break
                else:
                    ex = None
            # Success or fail
            if ex is not None:
                ext_id2, extension_encode = ex
                self._encode(f, extension_encode(self, value), streams, ext_id2)
            else:
                t = (
                    "Class %r is not a valid base BSDF type, nor is it "
                    "handled by an extension."
                )
                raise TypeError(t % value.__class__.__name__)

    def _decode(self, f):
        """ Main decoder function.
        """

        # Get value
        char = f.read(1)
        c = char.lower()

        # Conversion (uppercase value identifiers signify converted values)
        if not char:
            raise EOFError()
        elif char != c:
            n = strunpack("<B", f.read(1))[0]
            # if n == 253: n = strunpack('<Q', f.read(8))[0]  # noqa - noneed
            ext_id = f.read(n).decode("UTF-8")
        else:
            ext_id = None

        if c == b"v":
            value = None
        elif c == b"y":
            value = True
        elif c == b"n":
            value = False
        elif c == b"h":
            value = strunpack("<h", f.read(2))[0]
        elif c == b"i":
            value = strunpack("<q", f.read(8))[0]
        elif c == b"f":
            value = strunpack("<f", f.read(4))[0]
        elif c == b"d":
            value = strunpack("<d", f.read(8))[0]
        elif c == b"s":
            n_s = strunpack("<B", f.read(1))[0]
            if n_s == 253:
                n_s = strunpack("<Q", f.read(8))[0]  # noqa
            value = f.read(n_s).decode("UTF-8")
        elif c == b"l":
            n = strunpack("<B", f.read(1))[0]
            if n >= 254:
                # Streaming
                closed = n == 254
                n = strunpack("<Q", f.read(8))[0]
                if self._load_streaming:
                    value = ListStream(n if closed else "r")
                    value._activate(f, self._encode, self._decode)  # noqa
                elif closed:
                    value = [self._decode(f) for i in range(n)]
                else:
                    value = []
                    try:
                        while True:
                            value.append(self._decode(f))
                    except EOFError:
                        pass
            else:
                # Normal
                if n == 253:
                    n = strunpack("<Q", f.read(8))[0]  # noqa
                value = [self._decode(f) for i in range(n)]
        elif c == b"m":
            value = dict()
            n = strunpack("<B", f.read(1))[0]
            if n == 253:
                n = strunpack("<Q", f.read(8))[0]  # noqa
            for i in range(n):
                n_name = strunpack("<B", f.read(1))[0]
                if n_name == 253:
                    n_name = strunpack("<Q", f.read(8))[0]  # noqa
                assert n_name > 0
                name = f.read(n_name).decode("UTF-8")
                value[name] = self._decode(f)
        elif c == b"b":
            if self._lazy_blob:
                value = Blob((f, True))
            else:
                blob = Blob((f, False))
                value = blob.get_bytes()
        else:
            raise RuntimeError("Parse error %r" % char)

        # Convert value if we have an extension for it
        if ext_id is not None:
            extension = self._extensions.get(ext_id, None)
            if extension is not None:
                value = extension.decode(self, value)
            else:
                logger.warn("BSDF warning: no extension found for %r" % ext_id)

        return value

    def encode(self, ob):
        """ Save the given object to bytes.
        """
        f = BytesIO()
        self.save(f, ob)
        return f.getvalue()

    def save(self, f, ob):
        """ Write the given object to the given file object.
        """
        f.write(b"BSDF")
        f.write(struct.pack("<B", VERSION[0]))
        f.write(struct.pack("<B", VERSION[1]))

        # Prepare streaming, this list will have 0 or 1 item at the end
        streams = []

        self._encode(f, ob, streams, None)

        # Verify that stream object was at the end, and add initial elements
        if len(streams) > 0:
            stream = streams[0]
            if stream._start_pos != f.tell():
                raise ValueError(
                    "The stream object must be " "the last object to be encoded."
                )

    def decode(self, bb):
        """ Load the data structure that is BSDF-encoded in the given bytes.
        """
        f = BytesIO(bb)
        return self.load(f)

    def load(self, f):
        """ Load a BSDF-encoded object from the given file object.
        """
        # Check magic string
        f4 = f.read(4)
        if f4 != b"BSDF":
            raise RuntimeError("This does not look like a BSDF file: %r" % f4)
        # Check version
        major_version = strunpack("<B", f.read(1))[0]
        minor_version = strunpack("<B", f.read(1))[0]
        file_version = "%i.%i" % (major_version, minor_version)
        if major_version != VERSION[0]:  # major version should be 2
            t = (
                "Reading file with different major version (%s) "
                "from the implementation (%s)."
            )
            raise RuntimeError(t % (__version__, file_version))
        if minor_version > VERSION[1]:  # minor should be < ours
            t = (
                "BSDF warning: reading file with higher minor version (%s) "
                "than the implementation (%s)."
            )
            logger.warn(t % (__version__, file_version))

        return self._decode(f)


# %% Streaming and blob-files


class BaseStream(object):
    """ Base class for streams.
    """

    def __init__(self, mode="w"):
        self._i = 0
        self._count = -1
        if isinstance(mode, int):
            self._count = mode
            mode = "r"
        elif mode == "w":
            self._count = 0
        assert mode in ("r", "w")
        self._mode = mode
        self._f = None
        self._start_pos = 0

    def _activate(self, file, encode_func, decode_func):
        if self._f is not None:  # Associated with another write
            raise IOError("Stream object cannot be activated twice?")
        self._f = file
        self._start_pos = self._f.tell()
        self._encode = encode_func
        self._decode = decode_func

    @property
    def mode(self):
        """ The mode of this stream: 'r' or 'w'.
        """
        return self._mode


class ListStream(BaseStream):
    """ A streamable list object used for writing or reading.
    In read mode, it can also be iterated over.
    """

    @property
    def count(self):
        """ The number of elements in the stream (can be -1 for unclosed
        streams in read-mode).
        """
        return self._count

    @property
    def index(self):
        """ The current index of the element to read/write.
        """
        return self._i

    def append(self, item):
        """ Append an item to the streaming list. The object is immediately
        serialized and written to the underlying file.
        """
        # if self._mode != 'w':
        #     raise IOError('This ListStream is not in write mode.')
        if self._count != self._i:
            raise IOError("Can only append items to the end of the stream.")
        if self._f is None:
            raise IOError("List stream is not associated with a file yet.")
        if self._f.closed:
            raise IOError("Cannot stream to a close file.")
        self._encode(self._f, item, [self], None)
        self._i += 1
        self._count += 1

    def close(self, unstream=False):
        """ Close the stream, marking the number of written elements. New
        elements may still be appended, but they won't be read during decoding.
        If ``unstream`` is False, the stream is turned into a regular list
        (not streaming).
        """
        # if self._mode != 'w':
        #     raise IOError('This ListStream is not in write mode.')
        if self._count != self._i:
            raise IOError("Can only close when at the end of the stream.")
        if self._f is None:
            raise IOError("ListStream is not associated with a file yet.")
        if self._f.closed:
            raise IOError("Cannot close a stream on a close file.")
        i = self._f.tell()
        self._f.seek(self._start_pos - 8 - 1)
        self._f.write(spack("<B", 253 if unstream else 254))
        self._f.write(spack("<Q", self._count))
        self._f.seek(i)

    def next(self):
        """ Read and return the next element in the streaming list.
        Raises StopIteration if the stream is exhausted.
        """
        if self._mode != "r":
            raise IOError("This ListStream in not in read mode.")
        if self._f is None:
            raise IOError("ListStream is not associated with a file yet.")
        if getattr(self._f, "closed", None):  # not present on 2.7 http req :/
            raise IOError("Cannot read a stream from a close file.")
        if self._count >= 0:
            if self._i >= self._count:
                raise StopIteration()
            self._i += 1
            return self._decode(self._f)
        else:
            # This raises EOFError at some point.
            try:
                res = self._decode(self._f)
                self._i += 1
                return res
            except EOFError:
                self._count = self._i
                raise StopIteration()

    def __iter__(self):
        if self._mode != "r":
            raise IOError("Cannot iterate: ListStream in not in read mode.")
        return self

    def __next__(self):
        return self.next()


class Blob(object):
    """ Object to represent a blob of bytes. When used to write a BSDF file,
    it's a wrapper for bytes plus properties such as what compression to apply.
    When used to read a BSDF file, it can be used to read the data lazily, and
    also modify the data if reading in 'r+' mode and the blob isn't compressed.
    """

    # For now, this does not allow re-sizing blobs (within the allocated size)
    # but this can be added later.

    def __init__(self, bb, compression=0, extra_size=0, use_checksum=False):
        if isinstance(bb, bytes):
            self._f = None
            self.compressed = self._from_bytes(bb, compression)
            self.compression = compression
            self.allocated_size = self.used_size + extra_size
            self.use_checksum = use_checksum
        elif isinstance(bb, tuple) and len(bb) == 2 and hasattr(bb[0], "read"):
            self._f, allow_seek = bb
            self.compressed = None
            self._from_file(self._f, allow_seek)
            self._modified = False
        else:
            raise TypeError("Wrong argument to create Blob.")

    def _from_bytes(self, value, compression):
        """ When used to wrap bytes in a blob.
        """
        if compression == 0:
            compressed = value
        elif compression == 1:
            compressed = zlib.compress(value, 9)
        elif compression == 2:
            compressed = bz2.compress(value, 9)
        else:  # pragma: no cover
            assert False, "Unknown compression identifier"

        self.data_size = len(value)
        self.used_size = len(compressed)
        return compressed

    def _to_file(self, f):
        """ Private friend method called by encoder to write a blob to a file.
        """
        # Write sizes - write at least in a size that allows resizing
        if self.allocated_size <= 250 and self.compression == 0:
            f.write(spack("<B", self.allocated_size))
            f.write(spack("<B", self.used_size))
            f.write(lencode(self.data_size))
        else:
            f.write(spack("<BQ", 253, self.allocated_size))
            f.write(spack("<BQ", 253, self.used_size))
            f.write(spack("<BQ", 253, self.data_size))
        # Compression and checksum
        f.write(spack("B", self.compression))
        if self.use_checksum:
            f.write(b"\xff" + hashlib.md5(self.compressed).digest())
        else:
            f.write(b"\x00")
        # Byte alignment (only necessary for uncompressed data)
        if self.compression == 0:
            alignment = 8 - (f.tell() + 1) % 8  # +1 for the byte to write
            f.write(spack("<B", alignment))  # padding for byte alignment
            f.write(b"\x00" * alignment)
        else:
            f.write(spack("<B", 0))
        # The actual data and extra space
        f.write(self.compressed)
        f.write(b"\x00" * (self.allocated_size - self.used_size))

    def _from_file(self, f, allow_seek):
        """ Used when a blob is read by the decoder.
        """
        # Read blob header data (5 to 42 bytes)
        # Size
        allocated_size = strunpack("<B", f.read(1))[0]
        if allocated_size == 253:
            allocated_size = strunpack("<Q", f.read(8))[0]  # noqa
        used_size = strunpack("<B", f.read(1))[0]
        if used_size == 253:
            used_size = strunpack("<Q", f.read(8))[0]  # noqa
        data_size = strunpack("<B", f.read(1))[0]
        if data_size == 253:
            data_size = strunpack("<Q", f.read(8))[0]  # noqa
        # Compression and checksum
        compression = strunpack("<B", f.read(1))[0]
        has_checksum = strunpack("<B", f.read(1))[0]
        if has_checksum:
            checksum = f.read(16)
        # Skip alignment
        alignment = strunpack("<B", f.read(1))[0]
        f.read(alignment)
        # Get or skip data + extra space
        if allow_seek:
            self.start_pos = f.tell()
            self.end_pos = self.start_pos + used_size
            f.seek(self.start_pos + allocated_size)
        else:
            self.start_pos = None
            self.end_pos = None
            self.compressed = f.read(used_size)
            f.read(allocated_size - used_size)
        # Store info
        self.alignment = alignment
        self.compression = compression
        self.use_checksum = checksum if has_checksum else None
        self.used_size = used_size
        self.allocated_size = allocated_size
        self.data_size = data_size

    def seek(self, p):
        """ Seek to the given position (relative to the blob start).
        """
        if self._f is None:
            raise RuntimeError(
                "Cannot seek in a blob " "that is not created by the BSDF decoder."
            )
        if p < 0:
            p = self.allocated_size + p
        if p < 0 or p > self.allocated_size:
            raise IOError("Seek beyond blob boundaries.")
        self._f.seek(self.start_pos + p)

    def tell(self):
        """ Get the current file pointer position (relative to the blob start).
        """
        if self._f is None:
            raise RuntimeError(
                "Cannot tell in a blob " "that is not created by the BSDF decoder."
            )
        return self._f.tell() - self.start_pos

    def write(self, bb):
        """ Write bytes to the blob.
        """
        if self._f is None:
            raise RuntimeError(
                "Cannot write in a blob " "that is not created by the BSDF decoder."
            )
        if self.compression:
            raise IOError("Cannot arbitrarily write in compressed blob.")
        if self._f.tell() + len(bb) > self.end_pos:
            raise IOError("Write beyond blob boundaries.")
        self._modified = True
        return self._f.write(bb)

    def read(self, n):
        """ Read n bytes from the blob.
        """
        if self._f is None:
            raise RuntimeError(
                "Cannot read in a blob " "that is not created by the BSDF decoder."
            )
        if self.compression:
            raise IOError("Cannot arbitrarily read in compressed blob.")
        if self._f.tell() + n > self.end_pos:
            raise IOError("Read beyond blob boundaries.")
        return self._f.read(n)

    def get_bytes(self):
        """ Get the contents of the blob as bytes.
        """
        if self.compressed is not None:
            compressed = self.compressed
        else:
            i = self._f.tell()
            self.seek(0)
            compressed = self._f.read(self.used_size)
            self._f.seek(i)
        if self.compression == 0:
            value = compressed
        elif self.compression == 1:
            value = zlib.decompress(compressed)
        elif self.compression == 2:
            value = bz2.decompress(compressed)
        else:  # pragma: no cover
            raise RuntimeError("Invalid compression %i" % self.compression)
        return value

    def update_checksum(self):
        """ Reset the blob's checksum if present. Call this after modifying
        the data.
        """
        # or ... should the presence of a checksum mean that data is proteced?
        if self.use_checksum and self._modified:
            self.seek(0)
            compressed = self._f.read(self.used_size)
            self._f.seek(self.start_pos - self.alignment - 1 - 16)
            self._f.write(hashlib.md5(compressed).digest())


# %% High-level functions


def encode(ob, extensions=None, **options):
    """ Save (BSDF-encode) the given object to bytes.
    See `BSDFSerializer` for details on extensions and options.
    """
    s = BsdfSerializer(extensions, **options)
    return s.encode(ob)


def save(f, ob, extensions=None, **options):
    """ Save (BSDF-encode) the given object to the given filename or
    file object. See` BSDFSerializer` for details on extensions and options.
    """
    s = BsdfSerializer(extensions, **options)
    if isinstance(f, string_types):
        with open(f, "wb") as fp:
            return s.save(fp, ob)
    else:
        return s.save(f, ob)


def decode(bb, extensions=None, **options):
    """ Load a (BSDF-encoded) structure from bytes.
    See `BSDFSerializer` for details on extensions and options.
    """
    s = BsdfSerializer(extensions, **options)
    return s.decode(bb)


def load(f, extensions=None, **options):
    """ Load a (BSDF-encoded) structure from the given filename or file object.
    See `BSDFSerializer` for details on extensions and options.
    """
    s = BsdfSerializer(extensions, **options)
    if isinstance(f, string_types):
        if f.startswith(("~/", "~\\")):  # pragma: no cover
            f = os.path.expanduser(f)
        with open(f, "rb") as fp:
            return s.load(fp)
    else:
        return s.load(f)


# Aliases for json compat
loads = decode
dumps = encode


# %% Standard extensions

# Defining extensions as a dict would be more compact and feel lighter, but
# that would only allow lambdas, which is too limiting, e.g. for ndarray
# extension.


class Extension(object):
    """ Base class to implement BSDF extensions for special data types.

    Extension classes are provided to the BSDF serializer, which
    instantiates the class. That way, the extension can be somewhat dynamic:
    e.g. the NDArrayExtension exposes the ndarray class only when numpy
    is imported.

    A extension instance must have two attributes. These can be attribiutes of
    the class, or of the instance set in ``__init__()``:

    * name (str): the name by which encoded values will be identified.
    * cls (type): the type (or list of types) to match values with.
      This is optional, but it makes the encoder select extensions faster.

    Further, it needs 3 methods:

    * `match(serializer, value) -> bool`: return whether the extension can
      convert the given value. The default is ``isinstance(value, self.cls)``.
    * `encode(serializer, value) -> encoded_value`: the function to encode a
      value to more basic data types.
    * `decode(serializer, encoded_value) -> value`: the function to decode an
      encoded value back to its intended representation.

    """

    name = ""
    cls = ()

    def __repr__(self):
        return "<BSDF extension %r at 0x%s>" % (self.name, hex(id(self)))

    def match(self, s, v):
        return isinstance(v, self.cls)

    def encode(self, s, v):
        raise NotImplementedError()

    def decode(self, s, v):
        raise NotImplementedError()


class ComplexExtension(Extension):

    name = "c"
    cls = complex

    def encode(self, s, v):
        return (v.real, v.imag)

    def decode(self, s, v):
        return complex(v[0], v[1])


class NDArrayExtension(Extension):

    name = "ndarray"

    def __init__(self):
        if "numpy" in sys.modules:
            import numpy as np

            self.cls = np.ndarray

    def match(self, s, v):  # pragma: no cover - e.g. work for nd arrays in JS
        return hasattr(v, "shape") and hasattr(v, "dtype") and hasattr(v, "tobytes")

    def encode(self, s, v):
        return dict(shape=v.shape, dtype=text_type(v.dtype), data=v.tobytes())

    def decode(self, s, v):
        try:
            import numpy as np
        except ImportError:  # pragma: no cover
            return v
        a = np.frombuffer(v["data"], dtype=v["dtype"])
        a.shape = v["shape"]
        return a


standard_extensions = [ComplexExtension, NDArrayExtension]


if __name__ == "__main__":
    # Invoke CLI
    import bsdf_cli

    bsdf_cli.main()