import bz2
import StringIO

from django.db import models
from django.utils.six import with_metaclass


class CompressedBinaryField(with_metaclass(models.SubfieldBase, models.BinaryField)):

    description = "A binary field with bzip2 compression"

    def to_python(self, value):
        # Check value
        value = super(CompressedBinaryField, self).to_python(value)
        if value is None:
            return None
        if isinstance(value, bytes):
            return value

        # Decompress value if necessary
        compressed_content = StringIO.StringIO(value).getvalue()
        return bz2.decompress(compressed_content)

    def get_prep_value(self, value):
        # Check for empty values
        if value is None:
            return None

        # Compress
        return bz2.compress(value)