<commit_before>from __future__ import print_function
from functools import reduce
import boto3
import botocore
import logging
from io import BytesIO

from botocore.utils import fix_s3_host
from cottoncandy.utils import *
from .backend import CCBackEnd, CloudStream

try:
    from urllib import unquote
except ImportError:
    from urllib.parse import unquote
try:
    import cPickle as pickle
except ImportError:
    import pickle

class S3Client(CCBackEnd):
    """
    S3 client interface refactored out of cotton candy interfaces to allow for switching between
    S3 and Google Drive
    """

    @staticmethod
    def connect(ACCESS_KEY, SECRET_KEY, url, **kwargs):
        """Connect to S3 using boto

        Parameters
        ----------
        ACCESS_KEY
        SECRET_KEY
        url
        kwargs : dict
            Extra keyword arguments to `boto3.resource`

        Returns
        -------

        """
        s3 = boto3.resource('s3',
                            endpoint_url = url,
                            aws_access_key_id = ACCESS_KEY,
                            aws_secret_access_key = SECRET_KEY,
                            **kwargs)
        s3.meta.client.meta.events.unregister('before-sign.s3', fix_s3_host)
        return s3

    def __init__(self, bucket, access_key, secret_key, s3url, force_bucket_creation=False, **kwargs):
        """Constructor

        Parameters
        ----------
        bucket
        access_key
        secret_key
        s3url
        force_bucket_creation
        """
        super(S3Client, self).__init__()

        self.connection = S3Client.connect(access_key, secret_key, s3url, **kwargs)
        self.url = s3url
        if self.check_bucket_exists(bucket):
            self.set_current_bucket(bucket)
        else:
            print('* Bucket "%s" does not exist' % bucket)
            if force_bucket_creation:
                print('* Creating "%s" bucket...\n' % bucket)
                self.create_bucket(bucket)
            else:
                print('* cottoncandy instantiated without bucket...\n' \
                      '* Use with caution!\n' \
                      '* Many features will not work!!!\n')
                self.bucket_name = None

        if string2bool(ISBOTO_VERBOSE) is False:
            logging.getLogger('boto3').setLevel(logging.WARNING)
            logging.getLogger('botocore').setLevel(logging.WARNING)

    def get_bucket_name(self, bucket_name):
        """

        Parameters
        ----------
        bucket_name

        Returns
        -------

        """
        bucket_name = self.bucket_name \
            if bucket_name is None \
            else bucket_name
        return bucket_name

    @clean_object_name
    def check_file_exists(self, file_name, bucket_name=None):
        """Check whether object exists in bucket

        Parameters
        ----------
        file_name : str
            The object name
        bucket_name
        """
        bucket_name = self.get_bucket_name(bucket_name)
        ob = self.connection.Object(key = file_name, bucket_name = bucket_name)

        try:
            ob.load()
        except botocore.exceptions.ClientError as e:
            if e.response['Error']['Code'] == "404":
                exists = False
            else:
                raise e
        else:
            exists = True
        return exists

    def check_bucket_exists(self, bucket_name):
        """Check whether the bucket exists

        Parameters
        ----------
        bucket_name

        Returns
        -------

        """
        try:
            self.connection.meta.client.head_bucket(Bucket = bucket_name)
        except botocore.exceptions.ClientError as e:
            if e.response['Error']['Code'] == "404":
                exists = False
            else:
                raise e
        else:
            exists = True
        return exists

    def create_bucket(self, bucket_name, acl=DEFAULT_ACL):
        """Create a new bucket

        Parameters
        ----------
        bucket_name
        acl

        Returns
        -------

        """
        if MANDATORY_BUCKET_PREFIX:
            tt = len(MANDATORY_BUCKET_PREFIX)
            assert bucket_name[:tt] == MANDATORY_BUCKET_PREFIX

        self.connection.create_bucket(Bucket = bucket_name, ACL = acl)
        self.set_current_bucket(bucket_name)

    def set_current_bucket(self, bucket_name):
        """Sets which bucket to use

        Parameters
        ----------
        bucket_name

        Returns
        -------

        """
        if not self.check_bucket_exists(bucket_name):
            raise IOError('Bucket "%s" does not exist' % bucket_name)
        self.bucket_name = bucket_name

    def get_bucket(self):
        """Get bucket boto3 object

        Returns
        -------

        """
        s3_bucket = self.connection.Bucket(self.bucket_name)
        return s3_bucket

    def list_objects(self, **kwargs):
        """Get list of objects from the bucket
        This is a wrapper to ``self.get_bucket().bucket.objects``

        Parameters
        ----------
        limit : int, 1000
            Maximum number of items to return
        page_size : int, 1000
            The page size for pagination
        filter : dict
            A dictionary with key 'Prefix', specifying a prefix
            string.  Only return objects matching this string.
            Defaults to '/', all objects.
        kwargs : optional
            Dictionary of {method:value} for ``bucket.objects``

        Notes
        -----
        If you get a 'PaginationError', this means you have
        a lot of items on your bucket and should increase ``page_size``
        """
        defaults = dict(limit = 1000,
                        page_size = 1000,
                        filter = dict(Prefix = SEPARATOR),
                        )
        defaults.update(kwargs)
        bucket = self.get_bucket()
        prefix = defaults.pop('filter')

        if prefix['Prefix'] == SEPARATOR:
            request = bucket.objects
        else:
            request = bucket.objects.filter(**prefix)

        for method_name, value in defaults.items():
            if value is None:
                continue
            method = getattr(request, method_name)
            if isinstance(value, dict):
                request = method(**value)
            else:
                request = method(value)

        response = request.all()
        return response

    @property
    def size(self):
        return self.get_current_bucket_size()

    def get_current_bucket_size(self, limit=10 ** 6, page_size=10 ** 6):
        """Counts the size of all objects in the current bucket.

        Parameters
        ----------
        limit : int, 1000
            Maximum number of items to return
        page_size : int, 1000
            The page size for pagination

        Returns
        -------
        total_bytes : int
            The byte count of all objects in the bucket.

        Notes
        -----
        Because paging does not work properly, if there are more than
        limit,page_size number of objects in the bucket, this function will
        underestimate the total size. Check the printed number of objects for
        suspicious round numbers.
        TODO(anunez): Remove this note when the bug is fixed.
        """
        assert self.check_bucket_exists(self.bucket_name)
        obs = self.list_objects(limit = limit, page_size = page_size)
        object_sizes = [t.size for t in obs]
        total_bytes = sum(object_sizes)
        num_objects = len(object_sizes)
        del object_sizes

        txt = "%i bytes (%s) over %i objects"
        print(txt % (total_bytes, bytes2human(total_bytes), num_objects))
        return total_bytes

    def show_all_buckets(self):
        """

        Returns
        -------

        """
        all_buckets = list(self.connection.buckets.all())
        timeformat = lambda x: x.creation_date.astimezone(tzlocal()).strftime
        info = ['{0: <40} {1}'.format(unquote(t.name),
                                      timeformat(t)('%Y/%m/%d (%H:%M:%S)')) \
                for t in all_buckets]
        print('\n'.join(info))

    @clean_object_name
    def get_s3_object(self, file_name, bucket_name=None):
        """Get a boto3 object. Create it if it doesn't exist

        Parameters
        ----------
        file_name
        bucket_name

        Returns
        -------

        """
        bucket_name = self.get_bucket_name(bucket_name)
        return self.connection.Object(bucket_name = bucket_name, key = file_name)

    def upload_stream(self, body, file_name, metadata, permissions=DEFAULT_ACL):
        """Uploads a stream

        Parameters
        ----------
        file_name
        body : stream
        permissions
        metadata

        Returns
        -------

        """
        obj = self.get_s3_object(file_name)
        return obj.put(Body = body, ACL = permissions, Metadata = metadata)

    def download_stream(self, file_name):
        """Download object raw data.
        This simply calls the object body ``read()`` method.

        Parameters
        ---------
        file_name : str

        Returns
        -------
        stream
            file-like stream of object data
        """
        if not self.check_file_exists(file_name):
            raise IOError('Object "%s" does not exist' % file_name)
        s3_object = self.get_s3_object(file_name)
        return CloudStream(BytesIO(s3_object.get()['Body'].read()), s3_object.metadata)

    def upload_file(self, file_name, cloud_name=None, permissions=DEFAULT_ACL):
        """Upload a file to S3.

        Parameters
        ----------
        file_name : str
            Absolute path of file to upload
        cloud_name : str, None
            Name of uploaded object. If None, use
            the full file name as the object name.
        permissions: ???
            S3 permissions

        Returns
        -------
        response : boto3 response
        """
        assert os.path.exists(file_name)
        if file_name is None:
            file_name = os.path.abspath(file_name)
        s3_object = self.get_s3_object(cloud_name)
        return s3_object.upload_file(file_name, ExtraArgs = dict(ACL = permissions))

    def download_to_file(self, file_name, local_name):
        """Download S3 object to a file

        Parameters
        ----------
        file_name : str
        local_name : str
            Absolute path where the data will be downloaded on disk
        """
        assert self.check_file_exists(file_name)  # make sure object exists
        s3_object = self.get_s3_object(file_name)
        return s3_object.download_file(local_name)

    def upload_multipart(self, file_object, file_name, metadata, permissions=None,
                         buffersize=MPU_CHUNKSIZE, verbose=True, **kwargs):
        """Multi-part upload for a python file-object.

        This automatically creates a multipart upload of an object.
        Useful for large objects that are loaded in memory. This avoids
        having to write the file to disk and then using ``upload_from_file``.

        Parameters
        ----------
        file_name : str
        file_object :
            file-like python object (e.g. StringIO, file, etc)
        buffersize  : int
            Byte size of the individual parts to create.
            Defaults to 100MB
        verbose     : bool
            verbosity flag of whether to print mpu information to stdout
        metadata  : optional
            Metadata to store along with MPU object
        permissions : str?
            permissions for this file
        """
        client=self.connection.meta.client
        mpu = client.create_multipart_upload(Bucket = self.bucket_name,
                                             Key = file_name,
                                             Metadata = metadata)

        # get size
        nbytes_total = get_fileobject_size(file_object)
        file_object.seek(0)

        # check if our buffersize is sensible
        if nbytes_total < buffersize:
            npotential = nbytes_total / float(MIN_MPU_SIZE)
            if npotential > 10:
                # 10MB sensible minimum for 50MB < x < 100MB (default)
                buffersize = MIN_MPU_SIZE * 2
            else:
                # this file is smaller than chunksize by a little
                buffersize = MIN_MPU_SIZE

        # figure out parts split
        nparts = int(np.floor(nbytes_total / float(buffersize)))
        last_part_offset = int(nbytes_total - nparts * buffersize)

        # make sure we can upload
        assert nbytes_total < MAX_MPU_SIZE  # 5TB
        assert buffersize >= MIN_MPU_SIZE  # 5MB
        assert nparts < MAX_MPU_PARTS  # 10,000
        assert nbytes_total > buffersize

        mpu_info = dict(Parts = [])
        if verbose:
            print('MPU: %0.02fMB file in %i parts' % (nbytes_total / (2. ** 20), nparts))

        for chunk_idx in range(nparts):
            part_number = chunk_idx + 1
            if part_number == nparts:
                buffersize += last_part_offset
            txt = (part_number, nparts, buffersize / (2. ** 20), nbytes_total / (2. ** 20))
            if verbose:
                print('Uploading %i/%i: %0.02fMB of %0.02fMB' % txt)

            data_chunk = file_object.read(buffersize)
            response = client.upload_part(Bucket = self.bucket_name,
                                          Key = file_name,
                                          UploadId = mpu['UploadId'],
                                          PartNumber = part_number,
                                          Body = data_chunk)

            # store the part info
            part_info = dict(PartNumber = part_number,
                             ETag = response['ETag'])
            mpu_info['Parts'].append(part_info)

        # finalize
        mpu_response = client.complete_multipart_upload(Bucket = self.bucket_name,
                                                        Key = file_name,
                                                        UploadId = mpu['UploadId'],
                                                        MultipartUpload = mpu_info)
        return mpu_response

    def copy(self, source, destination, source_bucket, destination_bucket, overwrite):
        source_bucket = self.get_bucket_name(source_bucket)
        dest_bucket = source_bucket if (destination_bucket is None) else destination_bucket
        dest_bucket = self.get_bucket_name(dest_bucket)

        assert self.check_file_exists(source, bucket_name = source_bucket)
        ob_new = self.get_s3_object(destination, bucket_name = dest_bucket)

        if self.check_file_exists(ob_new.key, bucket_name = dest_bucket):
            assert overwrite is True

        fpath = pathjoin(source_bucket, source)
        ob_new.copy_from(CopySource = fpath)
        return ob_new

    def move(self, source, destination, source_bucket, destination_bucket, overwrite):
        new_ob = self.copy(source, destination, source_bucket, destination_bucket, overwrite)
        old_ob = self.get_s3_object(source, bucket_name = source_bucket)
        old_ob.delete()
        return new_ob

    def list_directory(self, path, limit):
        """List the contents of a "directory"

        Parameters
        ----------
        path : str (default: "/")

        Returns
        -------
        matches : list
            The children of the path.
        """
        if has_real_magic(path):
            raise ValueError('Use ``ls()`` when using search patterns: "%s"' % path)

        if (path != '') and (path != '/'):
            path = remove_root(path)
        path = remove_trivial_magic(path)
        path = mk_aws_path(path)

        response = self.get_bucket().meta.client.list_objects(Bucket = self.bucket_name,
                                                              Delimiter = SEPARATOR,
                                                              Prefix = path,
                                                              MaxKeys = limit)
        object_names = []
        if 'CommonPrefixes' in response:
            # we got common paths
            object_list = [t.values() for t in response['CommonPrefixes']]
            object_names += reduce(lambda x, y: x + y, object_list)
        if 'Contents' in response:
            # we got objects on the leaf nodes
            object_names += unquote_names([t['Key'] for t in response['Contents']])
        return map(os.path.normpath, object_names)

    def delete(self, file_name, recursive=False, delete=False):
        raise RuntimeError('Deleting on S3 backend is implemented by cottoncandy interface object')
<commit_msg>BF: use file name as object name if not given<commit_after>from __future__ import print_function
from functools import reduce
import boto3
import botocore
import logging
from io import BytesIO

from botocore.utils import fix_s3_host
from cottoncandy.utils import *
from .backend import CCBackEnd, CloudStream

try:
    from urllib import unquote
except ImportError:
    from urllib.parse import unquote
try:
    import cPickle as pickle
except ImportError:
    import pickle

class S3Client(CCBackEnd):
    """
    S3 client interface refactored out of cotton candy interfaces to allow for switching between
    S3 and Google Drive
    """

    @staticmethod
    def connect(ACCESS_KEY, SECRET_KEY, url, **kwargs):
        """Connect to S3 using boto

        Parameters
        ----------
        ACCESS_KEY
        SECRET_KEY
        url
        kwargs : dict
            Extra keyword arguments to `boto3.resource`

        Returns
        -------

        """
        s3 = boto3.resource('s3',
                            endpoint_url = url,
                            aws_access_key_id = ACCESS_KEY,
                            aws_secret_access_key = SECRET_KEY,
                            **kwargs)
        s3.meta.client.meta.events.unregister('before-sign.s3', fix_s3_host)
        return s3

    def __init__(self, bucket, access_key, secret_key, s3url, force_bucket_creation=False, **kwargs):
        """Constructor

        Parameters
        ----------
        bucket
        access_key
        secret_key
        s3url
        force_bucket_creation
        """
        super(S3Client, self).__init__()

        self.connection = S3Client.connect(access_key, secret_key, s3url, **kwargs)
        self.url = s3url
        if self.check_bucket_exists(bucket):
            self.set_current_bucket(bucket)
        else:
            print('* Bucket "%s" does not exist' % bucket)
            if force_bucket_creation:
                print('* Creating "%s" bucket...\n' % bucket)
                self.create_bucket(bucket)
            else:
                print('* cottoncandy instantiated without bucket...\n' \
                      '* Use with caution!\n' \
                      '* Many features will not work!!!\n')
                self.bucket_name = None

        if string2bool(ISBOTO_VERBOSE) is False:
            logging.getLogger('boto3').setLevel(logging.WARNING)
            logging.getLogger('botocore').setLevel(logging.WARNING)

    def get_bucket_name(self, bucket_name):
        """

        Parameters
        ----------
        bucket_name

        Returns
        -------

        """
        bucket_name = self.bucket_name \
            if bucket_name is None \
            else bucket_name
        return bucket_name

    @clean_object_name
    def check_file_exists(self, file_name, bucket_name=None):
        """Check whether object exists in bucket

        Parameters
        ----------
        file_name : str
            The object name
        bucket_name
        """
        bucket_name = self.get_bucket_name(bucket_name)
        ob = self.connection.Object(key = file_name, bucket_name = bucket_name)

        try:
            ob.load()
        except botocore.exceptions.ClientError as e:
            if e.response['Error']['Code'] == "404":
                exists = False
            else:
                raise e
        else:
            exists = True
        return exists

    def check_bucket_exists(self, bucket_name):
        """Check whether the bucket exists

        Parameters
        ----------
        bucket_name

        Returns
        -------

        """
        try:
            self.connection.meta.client.head_bucket(Bucket = bucket_name)
        except botocore.exceptions.ClientError as e:
            if e.response['Error']['Code'] == "404":
                exists = False
            else:
                raise e
        else:
            exists = True
        return exists

    def create_bucket(self, bucket_name, acl=DEFAULT_ACL):
        """Create a new bucket

        Parameters
        ----------
        bucket_name
        acl

        Returns
        -------

        """
        if MANDATORY_BUCKET_PREFIX:
            tt = len(MANDATORY_BUCKET_PREFIX)
            assert bucket_name[:tt] == MANDATORY_BUCKET_PREFIX

        self.connection.create_bucket(Bucket = bucket_name, ACL = acl)
        self.set_current_bucket(bucket_name)

    def set_current_bucket(self, bucket_name):
        """Sets which bucket to use

        Parameters
        ----------
        bucket_name

        Returns
        -------

        """
        if not self.check_bucket_exists(bucket_name):
            raise IOError('Bucket "%s" does not exist' % bucket_name)
        self.bucket_name = bucket_name

    def get_bucket(self):
        """Get bucket boto3 object

        Returns
        -------

        """
        s3_bucket = self.connection.Bucket(self.bucket_name)
        return s3_bucket

    def list_objects(self, **kwargs):
        """Get list of objects from the bucket
        This is a wrapper to ``self.get_bucket().bucket.objects``

        Parameters
        ----------
        limit : int, 1000
            Maximum number of items to return
        page_size : int, 1000
            The page size for pagination
        filter : dict
            A dictionary with key 'Prefix', specifying a prefix
            string.  Only return objects matching this string.
            Defaults to '/', all objects.
        kwargs : optional
            Dictionary of {method:value} for ``bucket.objects``

        Notes
        -----
        If you get a 'PaginationError', this means you have
        a lot of items on your bucket and should increase ``page_size``
        """
        defaults = dict(limit = 1000,
                        page_size = 1000,
                        filter = dict(Prefix = SEPARATOR),
                        )
        defaults.update(kwargs)
        bucket = self.get_bucket()
        prefix = defaults.pop('filter')

        if prefix['Prefix'] == SEPARATOR:
            request = bucket.objects
        else:
            request = bucket.objects.filter(**prefix)

        for method_name, value in defaults.items():
            if value is None:
                continue
            method = getattr(request, method_name)
            if isinstance(value, dict):
                request = method(**value)
            else:
                request = method(value)

        response = request.all()
        return response

    @property
    def size(self):
        return self.get_current_bucket_size()

    def get_current_bucket_size(self, limit=10 ** 6, page_size=10 ** 6):
        """Counts the size of all objects in the current bucket.

        Parameters
        ----------
        limit : int, 1000
            Maximum number of items to return
        page_size : int, 1000
            The page size for pagination

        Returns
        -------
        total_bytes : int
            The byte count of all objects in the bucket.

        Notes
        -----
        Because paging does not work properly, if there are more than
        limit,page_size number of objects in the bucket, this function will
        underestimate the total size. Check the printed number of objects for
        suspicious round numbers.
        TODO(anunez): Remove this note when the bug is fixed.
        """
        assert self.check_bucket_exists(self.bucket_name)
        obs = self.list_objects(limit = limit, page_size = page_size)
        object_sizes = [t.size for t in obs]
        total_bytes = sum(object_sizes)
        num_objects = len(object_sizes)
        del object_sizes

        txt = "%i bytes (%s) over %i objects"
        print(txt % (total_bytes, bytes2human(total_bytes), num_objects))
        return total_bytes

    def show_all_buckets(self):
        """

        Returns
        -------

        """
        all_buckets = list(self.connection.buckets.all())
        timeformat = lambda x: x.creation_date.astimezone(tzlocal()).strftime
        info = ['{0: <40} {1}'.format(unquote(t.name),
                                      timeformat(t)('%Y/%m/%d (%H:%M:%S)')) \
                for t in all_buckets]
        print('\n'.join(info))

    @clean_object_name
    def get_s3_object(self, file_name, bucket_name=None):
        """Get a boto3 object. Create it if it doesn't exist

        Parameters
        ----------
        file_name
        bucket_name

        Returns
        -------

        """
        bucket_name = self.get_bucket_name(bucket_name)
        return self.connection.Object(bucket_name = bucket_name, key = file_name)

    def upload_stream(self, body, file_name, metadata, permissions=DEFAULT_ACL):
        """Uploads a stream

        Parameters
        ----------
        file_name
        body : stream
        permissions
        metadata

        Returns
        -------

        """
        obj = self.get_s3_object(file_name)
        return obj.put(Body = body, ACL = permissions, Metadata = metadata)

    def download_stream(self, file_name):
        """Download object raw data.
        This simply calls the object body ``read()`` method.

        Parameters
        ---------
        file_name : str

        Returns
        -------
        stream
            file-like stream of object data
        """
        if not self.check_file_exists(file_name):
            raise IOError('Object "%s" does not exist' % file_name)
        s3_object = self.get_s3_object(file_name)
        return CloudStream(BytesIO(s3_object.get()['Body'].read()), s3_object.metadata)

    def upload_file(self, file_name, cloud_name=None, permissions=DEFAULT_ACL):
        """Upload a file to S3.

        Parameters
        ----------
        file_name : str
            Absolute path of file to upload
        cloud_name : str, None
            Name of uploaded object. If None, use
            the full file name as the object name.
        permissions: ???
            S3 permissions

        Returns
        -------
        response : boto3 response
        """
        assert os.path.exists(file_name)
        if cloud_name is None:
            cloud_name = file_name
        s3_object = self.get_s3_object(cloud_name)
        return s3_object.upload_file(file_name, ExtraArgs = dict(ACL = permissions))

    def download_to_file(self, file_name, local_name):
        """Download S3 object to a file

        Parameters
        ----------
        file_name : str
        local_name : str
            Absolute path where the data will be downloaded on disk
        """
        assert self.check_file_exists(file_name)  # make sure object exists
        s3_object = self.get_s3_object(file_name)
        return s3_object.download_file(local_name)

    def upload_multipart(self, file_object, file_name, metadata, permissions=None,
                         buffersize=MPU_CHUNKSIZE, verbose=True, **kwargs):
        """Multi-part upload for a python file-object.

        This automatically creates a multipart upload of an object.
        Useful for large objects that are loaded in memory. This avoids
        having to write the file to disk and then using ``upload_from_file``.

        Parameters
        ----------
        file_name : str
        file_object :
            file-like python object (e.g. StringIO, file, etc)
        buffersize  : int
            Byte size of the individual parts to create.
            Defaults to 100MB
        verbose     : bool
            verbosity flag of whether to print mpu information to stdout
        metadata  : optional
            Metadata to store along with MPU object
        permissions : str?
            permissions for this file
        """
        client=self.connection.meta.client
        mpu = client.create_multipart_upload(Bucket = self.bucket_name,
                                             Key = file_name,
                                             Metadata = metadata)

        # get size
        nbytes_total = get_fileobject_size(file_object)
        file_object.seek(0)

        # check if our buffersize is sensible
        if nbytes_total < buffersize:
            npotential = nbytes_total / float(MIN_MPU_SIZE)
            if npotential > 10:
                # 10MB sensible minimum for 50MB < x < 100MB (default)
                buffersize = MIN_MPU_SIZE * 2
            else:
                # this file is smaller than chunksize by a little
                buffersize = MIN_MPU_SIZE

        # figure out parts split
        nparts = int(np.floor(nbytes_total / float(buffersize)))
        last_part_offset = int(nbytes_total - nparts * buffersize)

        # make sure we can upload
        assert nbytes_total < MAX_MPU_SIZE  # 5TB
        assert buffersize >= MIN_MPU_SIZE  # 5MB
        assert nparts < MAX_MPU_PARTS  # 10,000
        assert nbytes_total > buffersize

        mpu_info = dict(Parts = [])
        if verbose:
            print('MPU: %0.02fMB file in %i parts' % (nbytes_total / (2. ** 20), nparts))

        for chunk_idx in range(nparts):
            part_number = chunk_idx + 1
            if part_number == nparts:
                buffersize += last_part_offset
            txt = (part_number, nparts, buffersize / (2. ** 20), nbytes_total / (2. ** 20))
            if verbose:
                print('Uploading %i/%i: %0.02fMB of %0.02fMB' % txt)

            data_chunk = file_object.read(buffersize)
            response = client.upload_part(Bucket = self.bucket_name,
                                          Key = file_name,
                                          UploadId = mpu['UploadId'],
                                          PartNumber = part_number,
                                          Body = data_chunk)

            # store the part info
            part_info = dict(PartNumber = part_number,
                             ETag = response['ETag'])
            mpu_info['Parts'].append(part_info)

        # finalize
        mpu_response = client.complete_multipart_upload(Bucket = self.bucket_name,
                                                        Key = file_name,
                                                        UploadId = mpu['UploadId'],
                                                        MultipartUpload = mpu_info)
        return mpu_response

    def copy(self, source, destination, source_bucket, destination_bucket, overwrite):
        source_bucket = self.get_bucket_name(source_bucket)
        dest_bucket = source_bucket if (destination_bucket is None) else destination_bucket
        dest_bucket = self.get_bucket_name(dest_bucket)

        assert self.check_file_exists(source, bucket_name = source_bucket)
        ob_new = self.get_s3_object(destination, bucket_name = dest_bucket)

        if self.check_file_exists(ob_new.key, bucket_name = dest_bucket):
            assert overwrite is True

        fpath = pathjoin(source_bucket, source)
        ob_new.copy_from(CopySource = fpath)
        return ob_new

    def move(self, source, destination, source_bucket, destination_bucket, overwrite):
        new_ob = self.copy(source, destination, source_bucket, destination_bucket, overwrite)
        old_ob = self.get_s3_object(source, bucket_name = source_bucket)
        old_ob.delete()
        return new_ob

    def list_directory(self, path, limit):
        """List the contents of a "directory"

        Parameters
        ----------
        path : str (default: "/")

        Returns
        -------
        matches : list
            The children of the path.
        """
        if has_real_magic(path):
            raise ValueError('Use ``ls()`` when using search patterns: "%s"' % path)

        if (path != '') and (path != '/'):
            path = remove_root(path)
        path = remove_trivial_magic(path)
        path = mk_aws_path(path)

        response = self.get_bucket().meta.client.list_objects(Bucket = self.bucket_name,
                                                              Delimiter = SEPARATOR,
                                                              Prefix = path,
                                                              MaxKeys = limit)
        object_names = []
        if 'CommonPrefixes' in response:
            # we got common paths
            object_list = [t.values() for t in response['CommonPrefixes']]
            object_names += reduce(lambda x, y: x + y, object_list)
        if 'Contents' in response:
            # we got objects on the leaf nodes
            object_names += unquote_names([t['Key'] for t in response['Contents']])
        return map(os.path.normpath, object_names)

    def delete(self, file_name, recursive=False, delete=False):
        raise RuntimeError('Deleting on S3 backend is implemented by cottoncandy interface object')
<|endoftext|>


<commit_before>#from flask import current_app as app
from flask import Blueprint, send_from_directory, abort, request
from werkzeug import secure_filename

from eve.methods.post import post_internal
#from eve.methods.patch import patch_internal
#from eve.methods.common import ratelimit
from eve.render import send_response
from eve.auth import requires_auth

import os
import errno

from pprint import pprint

UPLOAD_FOLDER = r'D:\Programmieren\amivapi\src\filedump'
STORAGE_URL = r'/storage'
ALLOWED_EXTENSIONS = ['.pdf', '.jpg', '.png']

""" Additional file functions """

def create(file):
    """Attempt to create the file on the drive
    First secures the filename
    Then, if necessary, changes it if a file with the same name exists
    Finally get size and type from file
    A dict will be returned that has all the information for the resource
    /files
    """
    filename = secure_filename(file.filename)

    # If needed, add number
    base, ext = os.path.splitext(filename)
    i = 1
    while exists(fullpath(filename)):
        filename = '%s_%s%s' % (base, str(i), ext)
        i += 1

    # Save file
    file.save(fullpath(filename))

    # Return new name
    return {
        'name': file.filename,
        'type': file.mimetype,
        'size': os.path.getsize(fullpath(filename)),
        'content_url': '%s/%s' % (STORAGE_URL, filename),
    }

def remove(filename):
    """Remove file from file system silently
    """
    try:
        os.remove(fullpath(filename))
    except OSError as e:
        if e.errno != errno.ENOENT:  # errno.ENOENT = no such file or directory
            raise  # re-raise exception if a different error occured

def exists(filename):
    """Check if file exists
    """
    return os.path.isfile(fullpath(filename))

def fullpath(filename):
    """Add specified path to file
    """
    return os.path.join(UPLOAD_FOLDER, filename)

def check_filename(filename):
    base, ext = os.path.splitext(filename)
    return ext in ALLOWED_EXTENSIONS

"""Here starts the routing"""

upload = Blueprint('upload ', __name__)

@requires_auth('files')
@upload.route('/files', methods=['POST'])
def upload_file():
    """ This will catch all POST requests to files
    If the request contains no files, it will abort
    Multiple files possible - parses all files and then posts them again
    inside the app as 'real' resources
    """
    if not request.files:
        abort(400, 'The request needs to contain a file.')

    payload = []
    for file in request.files:
        if not check_filename(request.files[file].filename):
            abort(400, 'Forbidden extension.')
        payload.append(create(request.files[file]))

    #app.test_client().post('file', response)
    response = post_internal('files', payload)

    return send_response('files', response)

# @upload.route('/files/<id>', methods=['PATCH'])
# def modify_file(id):
#     """ This will catch all PATCH and PUT requests to a file
#     If the request contains no file or more than one file, it will abort

#     """
#     if len(request.files) != 1:
#         abort(400, 'The request needs to contain a single file.')

#     for file in request.files:
#         if not check_filename(request.files[file].filename):
#             abort(400, 'Forbidden extension.')
#         payload = create(request.files[file])

#     response = patch_internal('files/%s' % str(id), payload)

#     return send_response('files/%s' % str(id), response)

@upload.route('%s/<filename>' % STORAGE_URL, methods=['GET'])
def download_file(filename):
    return send_from_directory(UPLOAD_FOLDER, filename)

#hook for delete
def post_database_delete(item):
    """After the file item has been removed from the database, the file will
    be removed from the server
    """
    # Take the url (has actual filename) and remove path
    remove(os.path.basename(dict(item)['content_url']))

def post_studydoc_insert(item):
    """After a studydocument is inserted into the database, the relation is saved in a table
    """
    pprint((item))

    # if dict(item)['files'a]
    # db = app.data.driver.session
    #             db.add(thisconfirm)
    #             db.commit()
    return
<commit_msg>Removed the now obsolete filestorage.py<commit_after><|endoftext|>


<commit_before># coding=utf-8
from __future__ import absolute_import
from __future__ import unicode_literals

import logging
import sqlite3
import time as _time
import threading
import warnings

from . import sqlbase
from .exceptions import Empty

sqlite3.enable_callback_tracebacks(True)

log = logging.getLogger(__name__)

# 10 seconds internal for `wait` of event
TICK_FOR_WAIT = 10

class AckStatus(object):
    inited = '0'
    ready = '1'
    unack = '2'
    acked = '5'
    ack_failed = '9'

class SQLiteAckQueue(sqlbase.SQLiteBase):
    """SQLite3 based FIFO queue with ack support."""

    _TABLE_NAME = 'ack_queue'
    _KEY_COLUMN = '_id'  # the name of the key column, used in DB CRUD
    _MAX_ACKED_LENGTH = 1000  # depreciated
    # SQL to create a table
    _SQL_CREATE = (
        'CREATE TABLE IF NOT EXISTS {table_name} ('
        '{key_column} INTEGER PRIMARY KEY AUTOINCREMENT, '
        'data BLOB, timestamp FLOAT, status INTEGER)'
    )
    # SQL to insert a record
    _SQL_INSERT = (
        'INSERT INTO {table_name} (data, timestamp, status)'
        ' VALUES (?, ?, %s)' % AckStatus.inited
    )
    # SQL to select a record
    _SQL_SELECT_ID = (
        'SELECT {key_column}, data, timestamp, status FROM {table_name} WHERE'
        ' {key_column} = {rowid}'
    )
    _SQL_SELECT = (
        'SELECT {key_column}, data, timestamp, status FROM {table_name} '
        'WHERE {key_column} > {rowid} AND status < %s '
        'ORDER BY {key_column} ASC LIMIT 1' % AckStatus.unack
    )
    _SQL_MARK_ACK_UPDATE = (
        'UPDATE {table_name} SET status = ? WHERE {key_column} = ?'
    )
    _SQL_SELECT_WHERE = (
        'SELECT {key_column}, data, timestamp FROM {table_name}'
        ' WHERE {key_column} > {rowid} AND status < %s AND'
        ' {column} {op} ? ORDER BY {key_column} ASC'
        ' LIMIT 1 ' % AckStatus.unack
    )
    _SQL_UPDATE = 'UPDATE {table_name} SET data = ? WHERE {key_column} = ?'

    def __init__(self, path, auto_resume=True, **kwargs):
        super(SQLiteAckQueue, self).__init__(path, **kwargs)
        if not self.auto_commit:
            warnings.warn("disable auto commit is not support in ack queue")
            self.auto_commit = True
        self._unack_cache = {}
        if auto_resume:
            self.resume_unack_tasks()

    def resume_unack_tasks(self):
        unack_count = self.unack_count()
        if unack_count:
            log.warning("resume %d unack tasks", unack_count)
        sql = 'UPDATE {} set status = ?' \
              ' WHERE status = ?'.format(self._table_name)
        with self.tran_lock:
            with self._putter as tran:
                tran.execute(sql, (AckStatus.ready, AckStatus.unack,))
            self.total = self._count()

    def put(self, item):
        obj = self._serializer.dumps(item)
        _id = self._insert_into(obj, _time.time())
        self.total += 1
        self.put_event.set()
        return _id

    def _init(self):
        super(SQLiteAckQueue, self)._init()
        # Action lock to assure multiple action to be *atomic*
        self.action_lock = threading.Lock()
        self.total = self._count()

    def _count(self):
        sql = 'SELECT COUNT({}) FROM {} WHERE status < ?'.format(
            self._key_column, self._table_name
        )
        row = self._getter.execute(sql, (AckStatus.unack,)).fetchone()
        return row[0] if row else 0

    def _ack_count_via_status(self, status):
        sql = 'SELECT COUNT({}) FROM {} WHERE status = ?'.format(
            self._key_column, self._table_name
        )
        row = self._getter.execute(sql, (status,)).fetchone()
        return row[0] if row else 0

    def unack_count(self):
        return self._ack_count_via_status(AckStatus.unack)

    def acked_count(self):
        return self._ack_count_via_status(AckStatus.acked)

    def ready_count(self):
        return self._ack_count_via_status(AckStatus.ready)

    def ack_failed_count(self):
        return self._ack_count_via_status(AckStatus.ack_failed)

    @sqlbase.with_conditional_transaction
    def _mark_ack_status(self, key, status):
        return self._sql_mark_ack_status, (
            status,
            key,
        )

    @sqlbase.with_conditional_transaction
    def clear_acked_data(
        self, max_delete=1000, keep_latest=1000, clear_ack_failed=False
    ):
        acked_clear_all = ''
        acked_to_delete = ''
        acked_to_keep = ''
        if self._MAX_ACKED_LENGTH != 1000 and not max_delete:
            # Added for backward compatibility for
            # those that set the _MAX_ACKED_LENGTH
            print(
                "_MAX_ACKED_LENGTH has been depreciated.  \
                    Use clear_acked_data(keep_latest=1000, max_delete=1000)"
            )
            max_delete = self._MAX_ACKED_LENGTH
        if clear_ack_failed:
            acked_clear_all = 'OR status = %s' % AckStatus.ack_failed
        if max_delete and max_delete > 0:
            acked_to_delete = 'LIMIT %d' % max_delete
        if keep_latest and keep_latest > 0:
            acked_to_keep = 'OFFSET %d' % keep_latest
        sql = """DELETE FROM {table_name}
            WHERE {key_column} IN (
                SELECT _id FROM {table_name}
                WHERE status = ? {clear_ack_failed}
                ORDER BY {key_column}
                DESC {acked_to_delete} {acked_to_keep}
            )""".format(
            table_name=self._table_name,
            key_column=self._key_column,
            acked_to_delete=acked_to_delete,
            acked_to_keep=acked_to_keep,
            clear_ack_failed=acked_clear_all,
        )
        return sql, AckStatus.acked

    @sqlbase.with_conditional_transaction
    def shrink_disk_usage(self):
        sql = """VACUUM"""
        return sql, ()

    @property
    def _sql_mark_ack_status(self):
        return self._SQL_MARK_ACK_UPDATE.format(
            table_name=self._table_name, key_column=self._key_column
        )

    def _pop(self, rowid=None, next_in_order=False, raw=False):
        with self.action_lock:
            row = self._select(next_in_order=next_in_order, rowid=rowid)
            # Perhaps a sqlite3 bug, sometimes (None, None) is returned
            # by select, below can avoid these invalid records.
            if row and row[0] is not None:
                self._mark_ack_status(row[0], AckStatus.unack)
                serialized_data = row[1]
                item = self._serializer.loads(serialized_data)
                self._unack_cache[row[0]] = item
                self.total -= 1
                if raw:
                    return {'pqid': row[0], 'data': item, 'tiemstamp': row[2]}
                else:
                    return item
            return None

    def _find_item_id(self, item, search=True):
        if item is None:
            return None
        elif isinstance(item, dict) and "pqid" in item:
            return item.get("pqid")
        elif search:
            for key, value in self._unack_cache.items():
                if value is item:
                    return key
            log.warning("Can't find item in unack cache.")
        elif isinstance(item, int) or (
            isinstance(item, str) and item.isnumeric()
        ):
            return int(item)
        return None

    def _check_id(self, item, id):
        if id is not None and item is not None:
            raise ValueError("Specify an id or an item, not both.")
        elif id is None and item is None:
            raise ValueError("Specify an id or an item.")
        elif id is not None:
            search = False
            item = id
        else:
            search = True
        return item, search

    def ack(self, item=None, id=None):
        item, search = self._check_id(item, id)
        with self.action_lock:
            _id = self._find_item_id(item, search)
            if _id is None:
                return None
            self._mark_ack_status(_id, AckStatus.acked)
            if _id in self._unack_cache:
                self._unack_cache.pop(_id)
        return _id

    def ack_failed(self, item=None, id=None):
        item, search = self._check_id(item, id)
        with self.action_lock:
            _id = self._find_item_id(item, search)
            if _id is None:
                return None
            self._mark_ack_status(_id, AckStatus.ack_failed)
            if _id in self._unack_cache:
                self._unack_cache.pop(_id)
        return _id

    def nack(self, item=None, id=None):
        item, search = self._check_id(item, id)
        with self.action_lock:
            _id = self._find_item_id(item, search)
            if _id is None:
                return None
            self._mark_ack_status(_id, AckStatus.ready)
            if _id in self._unack_cache:
                self._unack_cache.pop(_id)
            self.total += 1
        return _id

    def update(self, item, id=None):
        if isinstance(item, dict) and "pqid" in item:
            _id = item.get("pqid")
            item = item.get("data")
        if id is not None:
            _id = id
        if _id is None:
            raise ValueError("Provide an id or raw item")
        obj = self._serializer.dumps(item)
        self._update(_id, obj)
        return _id

    def get(
        self, block=True, timeout=None, id=None, next_in_order=False, raw=False
    ):
        rowid = self._find_item_id(id, search=False)
        if rowid is None and next_in_order:
            raise ValueError(
                "'next_in_order' requires the preceding 'id' be specified."
            )
        if next_in_order and not isinstance(next_in_order, bool):
            raise ValueError("'next_in_order' must be a boolean (True/False)")
        if not block:
            serialized = self._pop(
                next_in_order=next_in_order, raw=raw, rowid=rowid
            )
            if serialized is None:
                raise Empty
        elif timeout is None:
            # block until a put event.
            serialized = self._pop(
                next_in_order=next_in_order, raw=raw, rowid=rowid
            )
            while serialized is None:
                self.put_event.clear()
                self.put_event.wait(TICK_FOR_WAIT)
                serialized = self._pop(
                    next_in_order=next_in_order, raw=raw, rowid=rowid
                )
        elif timeout < 0:
            raise ValueError("'timeout' must be a non-negative number")
        else:
            # block until the timeout reached
            endtime = _time.time() + timeout
            serialized = self._pop(
                next_in_order=next_in_order, raw=raw, rowid=rowid
            )
            while serialized is None:
                self.put_event.clear()
                remaining = endtime - _time.time()
                if remaining <= 0.0:
                    raise Empty
                self.put_event.wait(
                    TICK_FOR_WAIT if TICK_FOR_WAIT < remaining else remaining
                )
                serialized = self._pop(
                    next_in_order=next_in_order, raw=raw, rowid=rowid
                )
        return serialized

    def task_done(self):
        """Persist the current state if auto_commit=False."""
        if not self.auto_commit:
            self._task_done()

    def queue(self):
        rows = self._sql_queue()
        datarows = []
        for row in rows:
            item = {
                'id': row[0],
                'data': self._serializer.loads(row[1]),
                'timestamp': row[2],
                'status': row[3],
            }
            datarows.append(item)
        return datarows

    @property
    def size(self):
        return self.total

    def qsize(self):
        return max(0, self.size)

    def active_size(self):
        return max(0, self.size + len(self._unack_cache))

    def empty(self):
        return self.size == 0

    def __len__(self):
        return self.size

FIFOSQLiteAckQueue = SQLiteAckQueue

class FILOSQLiteAckQueue(SQLiteAckQueue):
    """SQLite3 based FILO queue with ack support."""

    _TABLE_NAME = 'ack_filo_queue'
    # SQL to select a record
    _SQL_SELECT = (
        'SELECT {key_column}, data FROM {table_name} '
        'WHERE {key_column} < {rowid} and status < %s '
        'ORDER BY {key_column} DESC LIMIT 1' % AckStatus.unack
    )

class UniqueAckQ(SQLiteAckQueue):
    _TABLE_NAME = 'ack_unique_queue'
    _SQL_CREATE = (
        'CREATE TABLE IF NOT EXISTS {table_name} ('
        '{key_column} INTEGER PRIMARY KEY AUTOINCREMENT, '
        'data BLOB, timestamp FLOAT, status INTEGER, UNIQUE (data))'
    )

    def put(self, item):
        obj = self._serializer.dumps(item, sort_keys=True)
        _id = None
        try:
            _id = self._insert_into(obj, _time.time())
        except sqlite3.IntegrityError:
            pass
        else:
            self.total += 1
            self.put_event.set()
        return _id
<commit_msg>update() - ref before assignment in raise<commit_after># coding=utf-8
from __future__ import absolute_import
from __future__ import unicode_literals

import logging
import sqlite3
import time as _time
import threading
import warnings

from . import sqlbase
from .exceptions import Empty

sqlite3.enable_callback_tracebacks(True)

log = logging.getLogger(__name__)

# 10 seconds internal for `wait` of event
TICK_FOR_WAIT = 10

class AckStatus(object):
    inited = '0'
    ready = '1'
    unack = '2'
    acked = '5'
    ack_failed = '9'

class SQLiteAckQueue(sqlbase.SQLiteBase):
    """SQLite3 based FIFO queue with ack support."""

    _TABLE_NAME = 'ack_queue'
    _KEY_COLUMN = '_id'  # the name of the key column, used in DB CRUD
    _MAX_ACKED_LENGTH = 1000  # depreciated
    # SQL to create a table
    _SQL_CREATE = (
        'CREATE TABLE IF NOT EXISTS {table_name} ('
        '{key_column} INTEGER PRIMARY KEY AUTOINCREMENT, '
        'data BLOB, timestamp FLOAT, status INTEGER)'
    )
    # SQL to insert a record
    _SQL_INSERT = (
        'INSERT INTO {table_name} (data, timestamp, status)'
        ' VALUES (?, ?, %s)' % AckStatus.inited
    )
    # SQL to select a record
    _SQL_SELECT_ID = (
        'SELECT {key_column}, data, timestamp, status FROM {table_name} WHERE'
        ' {key_column} = {rowid}'
    )
    _SQL_SELECT = (
        'SELECT {key_column}, data, timestamp, status FROM {table_name} '
        'WHERE {key_column} > {rowid} AND status < %s '
        'ORDER BY {key_column} ASC LIMIT 1' % AckStatus.unack
    )
    _SQL_MARK_ACK_UPDATE = (
        'UPDATE {table_name} SET status = ? WHERE {key_column} = ?'
    )
    _SQL_SELECT_WHERE = (
        'SELECT {key_column}, data, timestamp FROM {table_name}'
        ' WHERE {key_column} > {rowid} AND status < %s AND'
        ' {column} {op} ? ORDER BY {key_column} ASC'
        ' LIMIT 1 ' % AckStatus.unack
    )
    _SQL_UPDATE = 'UPDATE {table_name} SET data = ? WHERE {key_column} = ?'

    def __init__(self, path, auto_resume=True, **kwargs):
        super(SQLiteAckQueue, self).__init__(path, **kwargs)
        if not self.auto_commit:
            warnings.warn("disable auto commit is not support in ack queue")
            self.auto_commit = True
        self._unack_cache = {}
        if auto_resume:
            self.resume_unack_tasks()

    def resume_unack_tasks(self):
        unack_count = self.unack_count()
        if unack_count:
            log.warning("resume %d unack tasks", unack_count)
        sql = 'UPDATE {} set status = ?' \
              ' WHERE status = ?'.format(self._table_name)
        with self.tran_lock:
            with self._putter as tran:
                tran.execute(sql, (AckStatus.ready, AckStatus.unack,))
            self.total = self._count()

    def put(self, item):
        obj = self._serializer.dumps(item)
        _id = self._insert_into(obj, _time.time())
        self.total += 1
        self.put_event.set()
        return _id

    def _init(self):
        super(SQLiteAckQueue, self)._init()
        # Action lock to assure multiple action to be *atomic*
        self.action_lock = threading.Lock()
        self.total = self._count()

    def _count(self):
        sql = 'SELECT COUNT({}) FROM {} WHERE status < ?'.format(
            self._key_column, self._table_name
        )
        row = self._getter.execute(sql, (AckStatus.unack,)).fetchone()
        return row[0] if row else 0

    def _ack_count_via_status(self, status):
        sql = 'SELECT COUNT({}) FROM {} WHERE status = ?'.format(
            self._key_column, self._table_name
        )
        row = self._getter.execute(sql, (status,)).fetchone()
        return row[0] if row else 0

    def unack_count(self):
        return self._ack_count_via_status(AckStatus.unack)

    def acked_count(self):
        return self._ack_count_via_status(AckStatus.acked)

    def ready_count(self):
        return self._ack_count_via_status(AckStatus.ready)

    def ack_failed_count(self):
        return self._ack_count_via_status(AckStatus.ack_failed)

    @sqlbase.with_conditional_transaction
    def _mark_ack_status(self, key, status):
        return self._sql_mark_ack_status, (
            status,
            key,
        )

    @sqlbase.with_conditional_transaction
    def clear_acked_data(
        self, max_delete=1000, keep_latest=1000, clear_ack_failed=False
    ):
        acked_clear_all = ''
        acked_to_delete = ''
        acked_to_keep = ''
        if self._MAX_ACKED_LENGTH != 1000 and not max_delete:
            # Added for backward compatibility for
            # those that set the _MAX_ACKED_LENGTH
            print(
                "_MAX_ACKED_LENGTH has been depreciated.  \
                    Use clear_acked_data(keep_latest=1000, max_delete=1000)"
            )
            max_delete = self._MAX_ACKED_LENGTH
        if clear_ack_failed:
            acked_clear_all = 'OR status = %s' % AckStatus.ack_failed
        if max_delete and max_delete > 0:
            acked_to_delete = 'LIMIT %d' % max_delete
        if keep_latest and keep_latest > 0:
            acked_to_keep = 'OFFSET %d' % keep_latest
        sql = """DELETE FROM {table_name}
            WHERE {key_column} IN (
                SELECT _id FROM {table_name}
                WHERE status = ? {clear_ack_failed}
                ORDER BY {key_column}
                DESC {acked_to_delete} {acked_to_keep}
            )""".format(
            table_name=self._table_name,
            key_column=self._key_column,
            acked_to_delete=acked_to_delete,
            acked_to_keep=acked_to_keep,
            clear_ack_failed=acked_clear_all,
        )
        return sql, AckStatus.acked

    @sqlbase.with_conditional_transaction
    def shrink_disk_usage(self):
        sql = """VACUUM"""
        return sql, ()

    @property
    def _sql_mark_ack_status(self):
        return self._SQL_MARK_ACK_UPDATE.format(
            table_name=self._table_name, key_column=self._key_column
        )

    def _pop(self, rowid=None, next_in_order=False, raw=False):
        with self.action_lock:
            row = self._select(next_in_order=next_in_order, rowid=rowid)
            # Perhaps a sqlite3 bug, sometimes (None, None) is returned
            # by select, below can avoid these invalid records.
            if row and row[0] is not None:
                self._mark_ack_status(row[0], AckStatus.unack)
                serialized_data = row[1]
                item = self._serializer.loads(serialized_data)
                self._unack_cache[row[0]] = item
                self.total -= 1
                if raw:
                    return {'pqid': row[0], 'data': item, 'tiemstamp': row[2]}
                else:
                    return item
            return None

    def _find_item_id(self, item, search=True):
        if item is None:
            return None
        elif isinstance(item, dict) and "pqid" in item:
            return item.get("pqid")
        elif search:
            for key, value in self._unack_cache.items():
                if value is item:
                    return key
            log.warning("Can't find item in unack cache.")
        elif isinstance(item, int) or (
            isinstance(item, str) and item.isnumeric()
        ):
            return int(item)
        return None

    def _check_id(self, item, id):
        if id is not None and item is not None:
            raise ValueError("Specify an id or an item, not both.")
        elif id is None and item is None:
            raise ValueError("Specify an id or an item.")
        elif id is not None:
            search = False
            item = id
        else:
            search = True
        return item, search

    def ack(self, item=None, id=None):
        item, search = self._check_id(item, id)
        with self.action_lock:
            _id = self._find_item_id(item, search)
            if _id is None:
                return None
            self._mark_ack_status(_id, AckStatus.acked)
            if _id in self._unack_cache:
                self._unack_cache.pop(_id)
        return _id

    def ack_failed(self, item=None, id=None):
        item, search = self._check_id(item, id)
        with self.action_lock:
            _id = self._find_item_id(item, search)
            if _id is None:
                return None
            self._mark_ack_status(_id, AckStatus.ack_failed)
            if _id in self._unack_cache:
                self._unack_cache.pop(_id)
        return _id

    def nack(self, item=None, id=None):
        item, search = self._check_id(item, id)
        with self.action_lock:
            _id = self._find_item_id(item, search)
            if _id is None:
                return None
            self._mark_ack_status(_id, AckStatus.ready)
            if _id in self._unack_cache:
                self._unack_cache.pop(_id)
            self.total += 1
        return _id

    def update(self, item, id=None):
        _id = None
        if isinstance(item, dict) and "pqid" in item:
            _id = item.get("pqid")
            item = item.get("data")
        if id is not None:
            _id = id
        if _id is None:
            raise ValueError("Provide an id or raw item")
        obj = self._serializer.dumps(item)
        self._update(_id, obj)
        return _id

    def get(
        self, block=True, timeout=None, id=None, next_in_order=False, raw=False
    ):
        rowid = self._find_item_id(id, search=False)
        if rowid is None and next_in_order:
            raise ValueError(
                "'next_in_order' requires the preceding 'id' be specified."
            )
        if next_in_order and not isinstance(next_in_order, bool):
            raise ValueError("'next_in_order' must be a boolean (True/False)")
        if not block:
            serialized = self._pop(
                next_in_order=next_in_order, raw=raw, rowid=rowid
            )
            if serialized is None:
                raise Empty
        elif timeout is None:
            # block until a put event.
            serialized = self._pop(
                next_in_order=next_in_order, raw=raw, rowid=rowid
            )
            while serialized is None:
                self.put_event.clear()
                self.put_event.wait(TICK_FOR_WAIT)
                serialized = self._pop(
                    next_in_order=next_in_order, raw=raw, rowid=rowid
                )
        elif timeout < 0:
            raise ValueError("'timeout' must be a non-negative number")
        else:
            # block until the timeout reached
            endtime = _time.time() + timeout
            serialized = self._pop(
                next_in_order=next_in_order, raw=raw, rowid=rowid
            )
            while serialized is None:
                self.put_event.clear()
                remaining = endtime - _time.time()
                if remaining <= 0.0:
                    raise Empty
                self.put_event.wait(
                    TICK_FOR_WAIT if TICK_FOR_WAIT < remaining else remaining
                )
                serialized = self._pop(
                    next_in_order=next_in_order, raw=raw, rowid=rowid
                )
        return serialized

    def task_done(self):
        """Persist the current state if auto_commit=False."""
        if not self.auto_commit:
            self._task_done()

    def queue(self):
        rows = self._sql_queue()
        datarows = []
        for row in rows:
            item = {
                'id': row[0],
                'data': self._serializer.loads(row[1]),
                'timestamp': row[2],
                'status': row[3],
            }
            datarows.append(item)
        return datarows

    @property
    def size(self):
        return self.total

    def qsize(self):
        return max(0, self.size)

    def active_size(self):
        return max(0, self.size + len(self._unack_cache))

    def empty(self):
        return self.size == 0

    def __len__(self):
        return self.size

FIFOSQLiteAckQueue = SQLiteAckQueue

class FILOSQLiteAckQueue(SQLiteAckQueue):
    """SQLite3 based FILO queue with ack support."""

    _TABLE_NAME = 'ack_filo_queue'
    # SQL to select a record
    _SQL_SELECT = (
        'SELECT {key_column}, data FROM {table_name} '
        'WHERE {key_column} < {rowid} and status < %s '
        'ORDER BY {key_column} DESC LIMIT 1' % AckStatus.unack
    )

class UniqueAckQ(SQLiteAckQueue):
    _TABLE_NAME = 'ack_unique_queue'
    _SQL_CREATE = (
        'CREATE TABLE IF NOT EXISTS {table_name} ('
        '{key_column} INTEGER PRIMARY KEY AUTOINCREMENT, '
        'data BLOB, timestamp FLOAT, status INTEGER, UNIQUE (data))'
    )

    def put(self, item):
        obj = self._serializer.dumps(item, sort_keys=True)
        _id = None
        try:
            _id = self._insert_into(obj, _time.time())
        except sqlite3.IntegrityError:
            pass
        else:
            self.total += 1
            self.put_event.set()
        return _id
<|endoftext|>


<commit_before>#!/usr/bin/env python
# $Id$

#  pyftpdlib is released under the MIT license, reproduced below:
#  ======================================================================
#  Copyright (C) 2007-2012 Giampaolo Rodola' <g.rodola@gmail.com>
#
#                         All Rights Reserved
#
# Permission is hereby granted, free of charge, to any person
# obtaining a copy of this software and associated documentation
# files (the "Software"), to deal in the Software without
# restriction, including without limitation the rights to use,
# copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the
# Software is furnished to do so, subject to the following
# conditions:
#
# The above copyright notice and this permission notice shall be
# included in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES
# OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
# NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
# HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
# WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
# OTHER DEALINGS IN THE SOFTWARE.
#
#  ======================================================================

"""
A FTP server banning clients in case of commands flood.

If client sends more than 300 requests per-second it will be
disconnected and won't be able to re-connect for 1 hour.
"""

from pyftpdlib.ftpserver import FTPHandler, FTPServer, DummyAuthorizer, CallEvery

class AntiFloodHandler(FTPHandler):

    cmds_per_second = 300  # max number of cmds per second
    ban_for = 60 * 60      # 1 hour
    banned_ips = []

    def __init__(self, *args, **kwargs):
        super(AntiFloodHandler, self).__init__(*args, **kwargs)
        self.processed_cmds = 0
        self.pcmds_callback = CallEvery(1, self.check_processed_cmds)

    def handle(self):
        # called when client connects.
        if self.remote_ip in self.banned_ips:
            self.respond('550 you are banned')
            self.close()
        else:
            super(AntiFloodHandler, self).handle()

    def check_processed_cmds(self):
        # called every second; checks for the number of commands
        # sent in the last second.
        if self.processed_cmds > self.cmds_per_second:
            self.ban(self.remote_ip)
        else:
            self.processed_cmds = 0

    def process_command(self, *args, **kwargs):
        # increase counter for every received command
        self.processed_cmds += 1
        super(AntiFloodHandler, self).process_command(*args, **kwargs)

    def ban(self, ip):
        # ban ip and schedule next un-ban
        if ip not in self.banned_ips:
            self.log('banned IP %s for command flooding' % ip)
        self.respond('550 you are banned for %s seconds' % self.ban_for)
        self.close()
        self.banned_ips.append(ip)

    def unban(self, ip):
        # unban ip
        try:
            self.banned_ips.remove(ip)
        except ValueError:
            pass
        else:
            self.log('unbanning IP %s' % ip)

    def close(self):
        super(AntiFloodHandler, self).close()
        if not self.pcmds_callback.cancelled:
            self.pcmds_callback.cancel()

def main():
    authorizer = DummyAuthorizer()
    authorizer.add_user('user', '12345', '.', perm='elradfmw')
    authorizer.add_anonymous('.')
    handler = AntiFloodHandler
    handler.authorizer = authorizer
    ftpd = FTPServer(('', 21), handler)
    ftpd.serve_forever(timeout=1)

if __name__ == '__main__':
    main()
<commit_msg>anti flood demo script: response strings should use Upper cased letters as per RFC-959<commit_after>#!/usr/bin/env python
# $Id$

#  pyftpdlib is released under the MIT license, reproduced below:
#  ======================================================================
#  Copyright (C) 2007-2012 Giampaolo Rodola' <g.rodola@gmail.com>
#
#                         All Rights Reserved
#
# Permission is hereby granted, free of charge, to any person
# obtaining a copy of this software and associated documentation
# files (the "Software"), to deal in the Software without
# restriction, including without limitation the rights to use,
# copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the
# Software is furnished to do so, subject to the following
# conditions:
#
# The above copyright notice and this permission notice shall be
# included in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES
# OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
# NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
# HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
# WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
# OTHER DEALINGS IN THE SOFTWARE.
#
#  ======================================================================

"""
A FTP server banning clients in case of commands flood.

If client sends more than 300 requests per-second it will be
disconnected and won't be able to re-connect for 1 hour.
"""

from pyftpdlib.ftpserver import FTPHandler, FTPServer, DummyAuthorizer, CallEvery

class AntiFloodHandler(FTPHandler):

    cmds_per_second = 300  # max number of cmds per second
    ban_for = 60 * 60      # 1 hour
    banned_ips = []

    def __init__(self, *args, **kwargs):
        super(AntiFloodHandler, self).__init__(*args, **kwargs)
        self.processed_cmds = 0
        self.pcmds_callback = CallEvery(1, self.check_processed_cmds)

    def handle(self):
        # called when client connects.
        if self.remote_ip in self.banned_ips:
            self.respond('550 You are banned.')
            self.close()
        else:
            super(AntiFloodHandler, self).handle()

    def check_processed_cmds(self):
        # called every second; checks for the number of commands
        # sent in the last second.
        if self.processed_cmds > self.cmds_per_second:
            self.ban(self.remote_ip)
        else:
            self.processed_cmds = 0

    def process_command(self, *args, **kwargs):
        # increase counter for every received command
        self.processed_cmds += 1
        super(AntiFloodHandler, self).process_command(*args, **kwargs)

    def ban(self, ip):
        # ban ip and schedule next un-ban
        if ip not in self.banned_ips:
            self.log('banned IP %s for command flooding' % ip)
        self.respond('550 You are banned for %s seconds.' % self.ban_for)
        self.close()
        self.banned_ips.append(ip)

    def unban(self, ip):
        # unban ip
        try:
            self.banned_ips.remove(ip)
        except ValueError:
            pass
        else:
            self.log('unbanning IP %s' % ip)

    def close(self):
        super(AntiFloodHandler, self).close()
        if not self.pcmds_callback.cancelled:
            self.pcmds_callback.cancel()

def main():
    authorizer = DummyAuthorizer()
    authorizer.add_user('user', '12345', '.', perm='elradfmw')
    authorizer.add_anonymous('.')
    handler = AntiFloodHandler
    handler.authorizer = authorizer
    ftpd = FTPServer(('', 21), handler)
    ftpd.serve_forever(timeout=1)

if __name__ == '__main__':
    main()
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
'''
Genshi Renderer for Salt
'''

from __future__ import absolute_import

# Import 3rd party libs
try:
    from genshi.template import MarkupTemplate
    from genshi.template import NewTextTemplate
    from genshi.template import OldTextTemplate
    HAS_LIBS = True
except ImportError:
    HAS_LIBS = False

# Import salt libs
from six import string_types

def render(genshi_data, saltenv='base', sls='', method='xml', **kws):
    '''
    Render a Genshi template. A method should be passed in as part of the
    kwargs. If no method is passed in, xml is assumed. Valid methods are:

    .. code-block:

        - xml
        - xhtml
        - html
        - text
        - newtext
        - oldtext

    Note that the ``text`` method will call ``NewTextTemplate``. If ``oldtext``
    is desired, it must be called explicitly

    :rtype: A Python data structure
    '''
    if not HAS_LIBS:
        return {}

    if not isinstance(genshi_data, string_types):
        genshi_data = genshi_data.read()

    if genshi_data.startswith('#!'):
        genshi_data = genshi_data[(genshi_data.find('\n') + 1):]
    if not genshi_data.strip():
        return {}

    if method == 'text' or method == 'newtext':
        tmpl = NewTextTemplate(genshi_data)
    elif method == 'oldtext':
        tmpl = OldTextTemplate(genshi_data)
    else:
        tmpl = MarkupTemplate(genshi_data)

    return tmpl.generate(**kws).render(method)
<commit_msg>Replaced module six in file /salt/renderers/genshi.py<commit_after># -*- coding: utf-8 -*-
'''
Genshi Renderer for Salt
'''

from __future__ import absolute_import

# Import 3rd party libs
try:
    from genshi.template import MarkupTemplate
    from genshi.template import NewTextTemplate
    from genshi.template import OldTextTemplate
    HAS_LIBS = True
except ImportError:
    HAS_LIBS = False

# Import salt libs
from salt.utils.six import string_types

def render(genshi_data, saltenv='base', sls='', method='xml', **kws):
    '''
    Render a Genshi template. A method should be passed in as part of the
    kwargs. If no method is passed in, xml is assumed. Valid methods are:

    .. code-block:

        - xml
        - xhtml
        - html
        - text
        - newtext
        - oldtext

    Note that the ``text`` method will call ``NewTextTemplate``. If ``oldtext``
    is desired, it must be called explicitly

    :rtype: A Python data structure
    '''
    if not HAS_LIBS:
        return {}

    if not isinstance(genshi_data, string_types):
        genshi_data = genshi_data.read()

    if genshi_data.startswith('#!'):
        genshi_data = genshi_data[(genshi_data.find('\n') + 1):]
    if not genshi_data.strip():
        return {}

    if method == 'text' or method == 'newtext':
        tmpl = NewTextTemplate(genshi_data)
    elif method == 'oldtext':
        tmpl = OldTextTemplate(genshi_data)
    else:
        tmpl = MarkupTemplate(genshi_data)

    return tmpl.generate(**kws).render(method)
<|endoftext|>


<commit_before>"""
VolumeVisualizationSimple

:Authors:
	Berend Klein Haneveld
"""
from VolumeVisualization import VolumeVisualization
from VolumeVisualization import VisualizationTypeSimple
from vtk import vtkVolumeProperty
from vtk import vtkColorTransferFunction
from vtk import vtkPiecewiseFunction
from ui.widgets.SliderFloatWidget import SliderFloatWidget
from ui.widgets.ColorWidget import ColorChoiceWidget
from core.decorators import overrides
from PySide.QtGui import QWidget
from PySide.QtGui import QGridLayout
from PySide.QtGui import QGroupBox
from PySide.QtCore import Qt

class VolumeVisualizationSimple(VolumeVisualization):
	"""
	VolumeVisualization subclass for a simple visualization.
	"""
	def __init__(self):
		super(VolumeVisualizationSimple, self).__init__()

		self.visualizationType = VisualizationTypeSimple

		# Create the volume property
		self.volProp = vtkVolumeProperty()
		self.volProp.SetIndependentComponents(True)
		self.volProp.SetInterpolationTypeToLinear()
		self.volProp.ShadeOn()
		self.volProp.SetAmbient(0.1)
		self.volProp.SetDiffuse(0.9)
		self.volProp.SetSpecular(0.2)
		self.volProp.SetSpecularPower(10.0)
		self.volProp.SetScalarOpacityUnitDistance(0.8919)

		self.minimum = 0
		self.maximum = 1
		self.lowerBound = 0
		self.upperBound = 1
		colors = [[255, 139, 0], [0, 147, 255], [0, 255, 147], [213, 100, 255], [255, 75, 75]]
		self.colors = map(lambda x: [x[0] / 255.0, x[1] / 255.0, x[2] / 255.0], colors)
		self.color = self.colors[0]
		self.opacity = 1.0
		self.colorFunction = None
		self.opacityFunction = None

	@overrides(VolumeVisualization)
	def getParameterWidget(self):
		"""
		Returns a widget with sliders / fields with which properties of this
		volume property can be adjusted.
		:rtype: QWidget
		"""
		self.lowerBoundSlider = SliderFloatWidget()
		self.lowerBoundSlider.setName("Lower:")
		self.lowerBoundSlider.setRange([self.minimum, self.maximum])
		self.lowerBoundSlider.setValue(self.lowerBound)
		self.lowerBoundSlider.label.setAlignment(Qt.AlignRight | Qt.AlignVCenter)
		self.lowerBoundSlider.valueChanged.connect(self.valueChanged)

		self.upperBoundSlider = SliderFloatWidget()
		self.upperBoundSlider.setName("Upper:")
		self.upperBoundSlider.setRange([self.minimum, self.maximum])
		self.upperBoundSlider.setValue(self.upperBound)
		self.upperBoundSlider.label.setAlignment(Qt.AlignRight | Qt.AlignVCenter)
		self.upperBoundSlider.valueChanged.connect(self.valueChanged)

		layoutGroup = QGridLayout()
		layoutGroup.setAlignment(Qt.AlignTop)
		layoutGroup.setContentsMargins(5, 0, 0, 0)
		layoutGroup.setSpacing(0)
		layoutGroup.addWidget(self.lowerBoundSlider)
		layoutGroup.addWidget(self.upperBoundSlider)

		groupbox = QGroupBox("Thresholds:")
		groupbox.setLayout(layoutGroup)

		self.opacitySlider = SliderFloatWidget()
		self.opacitySlider.setName("Opacity:")
		self.opacitySlider.setRange([0.0, 1.0])
		self.opacitySlider.setValue(self.opacity)
		self.opacitySlider.label.setAlignment(Qt.AlignRight | Qt.AlignVCenter)
		self.opacitySlider.valueChanged.connect(self.valueChanged)

		self.colorChooser = ColorChoiceWidget()
		self.colorChooser.setName("Color:")
		self.colorChooser.setColors(self.colors)
		self.colorChooser.setColor(self.color)
		self.colorChooser.setMinimumHeight(self.upperBoundSlider.sizeHint().height())
		self.colorChooser.label.setAlignment(Qt.AlignRight | Qt.AlignVCenter)
		self.colorChooser.valueChanged.connect(self.valueChanged)

		otherLayoutGroup = QGridLayout()
		otherLayoutGroup.setAlignment(Qt.AlignTop)
		otherLayoutGroup.setContentsMargins(5, 0, 0, 0)
		otherLayoutGroup.setSpacing(0)
		otherLayoutGroup.addWidget(self.opacitySlider)
		# otherLayoutGroup.addWidget(self.colorChooser)

		# otherBox = QGroupBox("Color and opacity:")
		otherBox = QGroupBox()
		otherBox.setLayout(otherLayoutGroup)

		layout = QGridLayout()
		layout.setAlignment(Qt.AlignTop)
		layout.setContentsMargins(0, 0, 0, 0)
		layout.setHorizontalSpacing(0)
		layout.addWidget(groupbox)
		layout.addWidget(otherBox)

		widget = QWidget()
		widget.setLayout(layout)

		try:
			from ColumnResizer import ColumnResizer
			self.columnResizer = ColumnResizer()
			self.columnResizer.addWidgetsFromLayout(self.lowerBoundSlider.layout(), 0)
			self.columnResizer.addWidgetsFromLayout(self.upperBoundSlider.layout(), 0)
			self.columnResizer.addWidgetsFromLayout(self.colorChooser.layout(), 0)
			self.columnResizer.addWidgetsFromLayout(self.opacitySlider.layout(), 0)

			self.otherColRes = ColumnResizer()
			self.otherColRes.addWidgetsFromLayout(self.lowerBoundSlider.layout(), 2)
			self.otherColRes.addWidgetsFromLayout(self.upperBoundSlider.layout(), 2)
			self.otherColRes.addWidgetsFromLayout(self.opacitySlider.layout(), 2)
		except Exception, e:
			print e

		return widget

	@overrides(VolumeVisualization)
	def setImageData(self, imageData):
		if imageData is None:
			self.minimum = 0.0
			self.maximum = 1.0
			self.lowerBound = self.minimum
			self.upperBound = self.maximum
			self.opacity = 1.0
			return

		self.minimum, self.maximum = imageData.GetScalarRange()
		self.lowerBound = self.minimum
		self.upperBound = self.maximum
		self.opacity = 1.0

	@overrides(VolumeVisualization)
	def setMapper(self, mapper):
		pass

	@overrides(VolumeVisualization)
	def shaderType(self):
		return 0

	@overrides(VolumeVisualization)
	def updateTransferFunction(self):
		r, g, b = self.color

		# Transfer functions and properties
		if not self.colorFunction:
			self.colorFunction = vtkColorTransferFunction()
		else:
			self.colorFunction.RemoveAllPoints()
		self.colorFunction.AddRGBPoint(self.minimum, r*0.7, g*0.7, b*0.7)
		self.colorFunction.AddRGBPoint(self.maximum, r, g, b)

		if not self.opacityFunction:
			self.opacityFunction = vtkPiecewiseFunction()
		else:
			self.opacityFunction.RemoveAllPoints()
		self.opacityFunction.AddPoint(self.minimum, 0)
		self.opacityFunction.AddPoint(self.lowerBound, 0)
		self.opacityFunction.AddPoint(self.lowerBound+0.0001, self.opacity)
		self.opacityFunction.AddPoint(self.upperBound-0.0001, self.opacity)
		self.opacityFunction.AddPoint(self.upperBound, 0)
		self.opacityFunction.AddPoint(self.maximum+0.0001, 0)

		self.volProp.SetColor(self.colorFunction)
		self.volProp.SetScalarOpacity(self.opacityFunction)

		self.updatedTransferFunction.emit()

	@overrides(VolumeVisualization)
	def valueChanged(self, value):
		"""
		This method is called when the value of one of the sliders / fields is
		adjusted. Argument value is unused. It is just there so that it can be
		connected to the signals of the interface elements.

		:type value: int
		"""
		self.lowerBound = min(self.lowerBoundSlider.value(), self.upperBoundSlider.value())
		self.upperBound = max(self.lowerBoundSlider.value(), self.upperBoundSlider.value())
		self.color = self.colorChooser.color
		self.opacity = self.opacitySlider.value()
		self.updateTransferFunction()
<commit_msg>Adjust ambient and diffuse parameters of visualization<commit_after>"""
VolumeVisualizationSimple

:Authors:
	Berend Klein Haneveld
"""
from VolumeVisualization import VolumeVisualization
from VolumeVisualization import VisualizationTypeSimple
from vtk import vtkVolumeProperty
from vtk import vtkColorTransferFunction
from vtk import vtkPiecewiseFunction
from ui.widgets.SliderFloatWidget import SliderFloatWidget
from ui.widgets.ColorWidget import ColorChoiceWidget
from core.decorators import overrides
from PySide.QtGui import QWidget
from PySide.QtGui import QGridLayout
from PySide.QtGui import QGroupBox
from PySide.QtCore import Qt

class VolumeVisualizationSimple(VolumeVisualization):
	"""
	VolumeVisualization subclass for a simple visualization.
	"""
	def __init__(self):
		super(VolumeVisualizationSimple, self).__init__()

		self.visualizationType = VisualizationTypeSimple

		# Create the volume property
		self.volProp = vtkVolumeProperty()
		self.volProp.SetIndependentComponents(True)
		self.volProp.SetInterpolationTypeToLinear()
		self.volProp.ShadeOn()
		self.volProp.SetAmbient(0.3)
		self.volProp.SetDiffuse(0.8)
		self.volProp.SetSpecular(0.2)
		self.volProp.SetSpecularPower(10.0)
		self.volProp.SetScalarOpacityUnitDistance(0.8919)

		self.minimum = 0
		self.maximum = 1
		self.lowerBound = 0
		self.upperBound = 1
		colors = [[255, 139, 0], [0, 147, 255], [0, 255, 147], [213, 100, 255], [255, 75, 75]]
		self.colors = map(lambda x: [x[0] / 255.0, x[1] / 255.0, x[2] / 255.0], colors)
		self.color = self.colors[0]
		self.opacity = 1.0
		self.colorFunction = None
		self.opacityFunction = None

	@overrides(VolumeVisualization)
	def getParameterWidget(self):
		"""
		Returns a widget with sliders / fields with which properties of this
		volume property can be adjusted.
		:rtype: QWidget
		"""
		self.lowerBoundSlider = SliderFloatWidget()
		self.lowerBoundSlider.setName("Lower:")
		self.lowerBoundSlider.setRange([self.minimum, self.maximum])
		self.lowerBoundSlider.setValue(self.lowerBound)
		self.lowerBoundSlider.label.setAlignment(Qt.AlignRight | Qt.AlignVCenter)
		self.lowerBoundSlider.valueChanged.connect(self.valueChanged)

		self.upperBoundSlider = SliderFloatWidget()
		self.upperBoundSlider.setName("Upper:")
		self.upperBoundSlider.setRange([self.minimum, self.maximum])
		self.upperBoundSlider.setValue(self.upperBound)
		self.upperBoundSlider.label.setAlignment(Qt.AlignRight | Qt.AlignVCenter)
		self.upperBoundSlider.valueChanged.connect(self.valueChanged)

		layoutGroup = QGridLayout()
		layoutGroup.setAlignment(Qt.AlignTop)
		layoutGroup.setContentsMargins(5, 0, 0, 0)
		layoutGroup.setSpacing(0)
		layoutGroup.addWidget(self.lowerBoundSlider)
		layoutGroup.addWidget(self.upperBoundSlider)

		groupbox = QGroupBox("Thresholds:")
		groupbox.setLayout(layoutGroup)

		self.opacitySlider = SliderFloatWidget()
		self.opacitySlider.setName("Opacity:")
		self.opacitySlider.setRange([0.0, 1.0])
		self.opacitySlider.setValue(self.opacity)
		self.opacitySlider.label.setAlignment(Qt.AlignRight | Qt.AlignVCenter)
		self.opacitySlider.valueChanged.connect(self.valueChanged)

		self.colorChooser = ColorChoiceWidget()
		self.colorChooser.setName("Color:")
		self.colorChooser.setColors(self.colors)
		self.colorChooser.setColor(self.color)
		self.colorChooser.setMinimumHeight(self.upperBoundSlider.sizeHint().height())
		self.colorChooser.label.setAlignment(Qt.AlignRight | Qt.AlignVCenter)
		self.colorChooser.valueChanged.connect(self.valueChanged)

		otherLayoutGroup = QGridLayout()
		otherLayoutGroup.setAlignment(Qt.AlignTop)
		otherLayoutGroup.setContentsMargins(5, 0, 0, 0)
		otherLayoutGroup.setSpacing(0)
		otherLayoutGroup.addWidget(self.opacitySlider)
		# otherLayoutGroup.addWidget(self.colorChooser)

		# otherBox = QGroupBox("Color and opacity:")
		otherBox = QGroupBox()
		otherBox.setLayout(otherLayoutGroup)

		layout = QGridLayout()
		layout.setAlignment(Qt.AlignTop)
		layout.setContentsMargins(0, 0, 0, 0)
		layout.setHorizontalSpacing(0)
		layout.addWidget(groupbox)
		layout.addWidget(otherBox)

		widget = QWidget()
		widget.setLayout(layout)

		try:
			from ColumnResizer import ColumnResizer
			self.columnResizer = ColumnResizer()
			self.columnResizer.addWidgetsFromLayout(self.lowerBoundSlider.layout(), 0)
			self.columnResizer.addWidgetsFromLayout(self.upperBoundSlider.layout(), 0)
			self.columnResizer.addWidgetsFromLayout(self.colorChooser.layout(), 0)
			self.columnResizer.addWidgetsFromLayout(self.opacitySlider.layout(), 0)

			self.otherColRes = ColumnResizer()
			self.otherColRes.addWidgetsFromLayout(self.lowerBoundSlider.layout(), 2)
			self.otherColRes.addWidgetsFromLayout(self.upperBoundSlider.layout(), 2)
			self.otherColRes.addWidgetsFromLayout(self.opacitySlider.layout(), 2)
		except Exception, e:
			print e

		return widget

	@overrides(VolumeVisualization)
	def setImageData(self, imageData):
		if imageData is None:
			self.minimum = 0.0
			self.maximum = 1.0
			self.lowerBound = self.minimum
			self.upperBound = self.maximum
			self.opacity = 1.0
			return

		self.minimum, self.maximum = imageData.GetScalarRange()
		self.lowerBound = self.minimum
		self.upperBound = self.maximum
		self.opacity = 1.0

	@overrides(VolumeVisualization)
	def setMapper(self, mapper):
		pass

	@overrides(VolumeVisualization)
	def shaderType(self):
		return 0

	@overrides(VolumeVisualization)
	def updateTransferFunction(self):
		r, g, b = self.color

		# Transfer functions and properties
		if not self.colorFunction:
			self.colorFunction = vtkColorTransferFunction()
		else:
			self.colorFunction.RemoveAllPoints()
		self.colorFunction.AddRGBPoint(self.minimum, r*0.7, g*0.7, b*0.7)
		self.colorFunction.AddRGBPoint(self.maximum, r, g, b)

		if not self.opacityFunction:
			self.opacityFunction = vtkPiecewiseFunction()
		else:
			self.opacityFunction.RemoveAllPoints()
		self.opacityFunction.AddPoint(self.minimum, 0)
		self.opacityFunction.AddPoint(self.lowerBound, 0)
		self.opacityFunction.AddPoint(self.lowerBound+0.0001, self.opacity)
		self.opacityFunction.AddPoint(self.upperBound-0.0001, self.opacity)
		self.opacityFunction.AddPoint(self.upperBound, 0)
		self.opacityFunction.AddPoint(self.maximum+0.0001, 0)

		self.volProp.SetColor(self.colorFunction)
		self.volProp.SetScalarOpacity(self.opacityFunction)

		self.updatedTransferFunction.emit()

	@overrides(VolumeVisualization)
	def valueChanged(self, value):
		"""
		This method is called when the value of one of the sliders / fields is
		adjusted. Argument value is unused. It is just there so that it can be
		connected to the signals of the interface elements.

		:type value: int
		"""
		self.lowerBound = min(self.lowerBoundSlider.value(), self.upperBoundSlider.value())
		self.upperBound = max(self.lowerBoundSlider.value(), self.upperBoundSlider.value())
		self.color = self.colorChooser.color
		self.opacity = self.opacitySlider.value()
		self.updateTransferFunction()
<|endoftext|>