<commit_before>import gzip
import json
from django.db import models
from django.db import transaction

BULK_CREATE_BATCH_SIZE = 10000

class Annotation(object):
    """
    Can contain multiple Chunk_Feature objects merged together.
    """

    def __init__(self, base_first, base_last, chunk_feature, fragment=None):
        self.base_first = base_first
        self.base_last = base_last
        self.feature = chunk_feature.feature
        self.feature_base_first = chunk_feature.feature_base_first
        self.feature_base_last = chunk_feature.feature_base_last
        self.fragment = fragment

    @property
    def feature_name(self):
        if (
            self.feature_base_first != 1
            or self.feature_base_last != self.feature.length
        ):
            return "%s[%s:%s]" % (
                self.feature.name,
                self.feature_base_first,
                self.feature_base_last,
            )
        else:
            return self.feature.name

    def __str__(self):
        s = []
        s.append(self.feature_name)
        s.append(self.feature.type)
        if self.feature.strand == 1:
            s.append("+")
        else:
            s.append("-")
        return ", ".join(s)

    @staticmethod
    def from_chunk_feature_and_location_array(chunk_feature_locs):
        """
        chunk_feature_locs: an array of (Chunk_Feature, Fragment_Chunk_Location) tuples.
        """

        chunk_feature_locs = sorted(
            chunk_feature_locs,
            key=lambda t: (t[0].feature.id, t[1].base_first,
                           (t[0].feature.strand if t[0].feature.strand is not None else 1)
                           * t[0].feature_base_first)
        )

        annotations = []
        for cf, fcl in chunk_feature_locs:
            if (
                len(annotations) > 0
                and annotations[-1].feature.id == cf.feature_id
                and (((annotations[-1].feature.strand is None or annotations[-1].feature.strand > 0)
                      and annotations[-1].feature_base_last == cf.feature_base_first - 1)
                     or ((annotations[-1].feature.strand is not None
                          and annotations[-1].feature.strand < 0)
                         and annotations[-1].feature_base_first == cf.feature_base_last + 1))
                and annotations[-1].base_last == fcl.base_first - 1
            ):
                # merge annotation
                annotations[-1].base_last = fcl.base_last
                if (annotations[-1].feature.strand is None or annotations[-1].feature.strand > 0):
                    annotations[-1].feature_base_last = cf.feature_base_last
                else:
                    annotations[-1].feature_base_first = cf.feature_base_first
            else:
                annotations.append(
                    Annotation(
                        base_first=fcl.base_first,
                        base_last=fcl.base_last,
                        chunk_feature=cf,
                        fragment=fcl.fragment,
                    )
                )

        return sorted(annotations, key=lambda a: a.base_first)

class ChunkReference(object):
    """
    Parent class for storing chunk sequences in reference files
    1-based indexing example: (2, 4) of 5'-A(TCG)GCTA-3'
    """

    @staticmethod
    def generate_from_fragment(fragment, sequence=None, dirn='.'):
        return ChunkReference(None)

    @staticmethod
    def read_reference_sequence_at_position_opened_file(f, start, end):
        """
        Read sequence from 1-indexed positions from reference.
        """
        f.seek(start - 1)
        sequence = f.read(end - start + 1).decode('utf-8')
        return sequence

    def read_reference_sequence_at_position(self, start, end):
        """
        Read sequence from 1-indexed positions from reference.
        """
        with gzip.open(self.ref_fn, 'rb') as f:
            f.seek(start - 1)
            sequence = f.read(end - start + 1).decode('utf-8')
        return sequence

    def __init__(self, ref_fn):
        self.ref_fn = ref_fn

class LocalChunkReference(ChunkReference):
    """
    Class for locally storing chunk sequences in reference files
    """

    @staticmethod
    def generate_from_fragment(fragment, sequence=None, dirn='.'):
        if sequence is None:
            sequence = fragment.indexed_fragment().sequence

        ref_fn = f"{dirn}/{fragment.id}.fa.gz"
        with gzip.open(ref_fn, 'wb') as f:
            sequence = f.write(sequence.encode('utf-8'))

        return LocalChunkReference(ref_fn)

class NetAppChunkReference(ChunkReference):
    """
    Class for storing chunk sequences in reference files on NetApp
    """

    @staticmethod
    def generate_from_fragment(fragment, sequence=None, dirn='.'):
        # NOTE: dirn is ignored in this method, as we're storing on NetApp
        if sequence is None:
            sequence = fragment.indexed_fragment().sequence
        ref_fn = fragment.build_fragment_fasta_from_sequence(sequence)

        return NetAppChunkReference(ref_fn)

class BigIntPrimaryModel(models.Model):
    class Meta:
        app_label = "edge"
        abstract = True

    id = models.BigIntegerField(primary_key=True)

    @transaction.atomic()
    def save(self, *args, **kwargs):
        # mimic auto_increment
        if self.id is None:
            klass = type(self)
            try:
                self.id = (
                    klass.objects.select_for_update()
                    .order_by("-id")
                    .values("id")[0]["id"]
                    + 1
                )
            except IndexError:
                self.id = 1
        return super(BigIntPrimaryModel, self).save(*args, **kwargs)

class Chunk(BigIntPrimaryModel):
    CHUNK_REFERENCE_CLASS = NetAppChunkReference

    class Meta:
        app_label = "edge"

    initial_fragment = models.ForeignKey("Fragment", on_delete=models.PROTECT)
    sequence = models.TextField(null=True)
    ref_start_index = models.PositiveIntegerField(blank=True, null=True)
    ref_end_index = models.PositiveIntegerField(blank=True, null=True)

    @property
    def ref_fn(self):
        return self.initial_fragment.fragment_reference_fasta_gz_fn()

    @property
    def is_sequence_based(self):
        return self.sequence is not None

    @property
    def is_reference_based(self):
        return None not in [self.ref_start_index, self.ref_end_index]

    @property
    def length(self):
        if self.is_sequence_based:
            return len(self.sequence)
        elif self.is_reference_based:
            return self.ref_end_index - self.ref_start_index + 1
        raise Exception("invalid chunk data")

    def get_sequence(self, f=None):
        if self.is_sequence_based:
            return self.sequence
        elif self.is_reference_based:
            if f is None:
                cr = self.CHUNK_REFERENCE_CLASS(self.ref_fn)
                return cr.read_reference_sequence_at_position(
                    self.ref_start_index, self.ref_end_index
                )
            else:
                return self.CHUNK_REFERENCE_CLASS.read_reference_sequence_at_position_opened_file(
                    f, self.ref_start_index, self.ref_end_index
                )
        raise Exception("invalid chunk data")

    def reload(self):
        return Chunk.objects.get(pk=self.pk)

    @staticmethod
    @transaction.atomic()  # this method has to be in a transaction, see below
    def bulk_create(entries):
        # IMPORTANT: this entire method must be within a transaction, so we can
        # compute and assign unique BigIntPrimary IDs

        cur_id = 1
        try:
            cur_id = (
                Chunk.objects.select_for_update()
                .order_by("-id")
                .values("id")[0]["id"]
                + 1
            )
        except IndexError:
            pass

        for entry in entries:
            entry.id = cur_id
            cur_id += 1
        Chunk.objects.bulk_create(entries, batch_size=BULK_CREATE_BATCH_SIZE)

class Edge(BigIntPrimaryModel):
    class Meta:
        app_label = "edge"

    from_chunk = models.ForeignKey(
        Chunk, related_name="out_edges", on_delete=models.PROTECT
    )
    fragment = models.ForeignKey("Fragment", on_delete=models.PROTECT)
    # can be null, so we can supersede an edge from a child fragment
    to_chunk = models.ForeignKey(
        Chunk, null=True, related_name="in_edges", on_delete=models.PROTECT
    )

    @staticmethod
    @transaction.atomic()  # this method has to be in a transaction, see below
    def bulk_create(entries):
        # IMPORTANT: this entire method must be within a transaction, so we can
        # compute and assign unique BigIntPrimary IDs

        cur_id = 1
        try:
            cur_id = (
                Edge.objects.select_for_update()
                .order_by("-id")
                .values("id")[0]["id"]
                + 1
            )
        except IndexError:
            pass

        for entry in entries:
            entry.id = cur_id
            cur_id += 1
        Edge.objects.bulk_create(entries, batch_size=BULK_CREATE_BATCH_SIZE)

class Feature(models.Model):
    class Meta:
        app_label = "edge"

    name = models.CharField(max_length=100)
    type = models.CharField(max_length=100)
    strand = models.IntegerField(null=True)
    length = models.IntegerField()
    operation = models.ForeignKey("Operation", null=True, on_delete=models.CASCADE)
    _qualifiers = models.TextField(null=True, db_column="qualifiers")

    def set_qualifiers(self, qualifiers):
        self._qualifiers = json.dumps(qualifiers)

    @property
    def qualifiers(self):
        if self._qualifiers is not None:
            return json.loads(self._qualifiers)
        else:
            return None

class Chunk_Feature_Manager(models.Manager):
    def get_query_set(self):
        return (
            super(Chunk_Feature_Manager, self)
            .get_query_set()
            .select_related("chunk", "feature")
        )

class Chunk_Feature(BigIntPrimaryModel):
    class Meta:
        app_label = "edge"

    objects = Chunk_Feature_Manager()
    chunk = models.ForeignKey(Chunk, on_delete=models.PROTECT)
    feature = models.ForeignKey(Feature, on_delete=models.PROTECT)
    feature_base_first = models.IntegerField()
    feature_base_last = models.IntegerField()

    @staticmethod
    @transaction.atomic()  # this method has to be in a transaction, see below
    def bulk_create(entries):
        # IMPORTANT: this entire method must be within a transaction, so we can
        # compute and assign unique BigIntPrimary IDs

        cur_id = 1
        try:
            cur_id = (
                Chunk_Feature.objects.select_for_update()
                .order_by("-id")
                .values("id")[0]["id"]
                + 1
            )
        except IndexError:
            pass

        for entry in entries:
            entry.id = cur_id
            cur_id += 1
        Chunk_Feature.objects.bulk_create(entries, batch_size=BULK_CREATE_BATCH_SIZE)

class Fragment_Chunk_Location(BigIntPrimaryModel):
    class Meta:
        app_label = "edge"
        unique_together = (("fragment", "chunk"),)
        index_together = (("fragment", "base_last"), ("fragment", "base_first"))

    fragment = models.ForeignKey("Fragment", on_delete=models.PROTECT)
    chunk = models.ForeignKey(Chunk, on_delete=models.PROTECT)
    base_first = models.IntegerField()
    base_last = models.IntegerField()

    def __str__(self):
        return "%s: %s-%s" % (self.fragment.name, self.location[0], self.location[1])

    @property
    def next_chunk(self):
        for fcl in self.fragment.fragment_chunk_location_set.select_related(
            "chunk"
        ).filter(base_first=self.base_last + 1):
            return fcl.chunk
        return None

    @property
    def location(self):
        return (self.base_first, self.base_last)

    @property
    def prev_fragment_chunk(self):
        loc = self.location
        if loc[0] == 1:
            return None
        fc = Fragment_Chunk_Location.objects.get(
            fragment=self.fragment, base_last=loc[0] - 1
        )
        return fc

    def annotations(self):
        return [
            Annotation(
                base_first=self.base_first, base_last=self.base_last, chunk_feature=cf
            )
            for cf in self.chunk.chunk_feature_set.all()
        ]

    @staticmethod
    @transaction.atomic()  # this method has to be in a transaction, see below
    def bulk_create(entries):
        # IMPORTANT: this entire method must be within a transaction, so we can
        # compute and assign unique BigIntPrimary IDs

        cur_id = 1
        try:
            cur_id = (
                Fragment_Chunk_Location.objects.select_for_update()
                .order_by("-id")
                .values("id")[0]["id"]
                + 1
            )
        except IndexError:
            pass

        for entry in entries:
            entry.id = cur_id
            cur_id += 1
        Fragment_Chunk_Location.objects.bulk_create(entries, batch_size=BULK_CREATE_BATCH_SIZE)
<commit_msg>added bulk create to Feature<commit_after>import gzip
import json
from django.db import models
from django.db import transaction

BULK_CREATE_BATCH_SIZE = 10000

class Annotation(object):
    """
    Can contain multiple Chunk_Feature objects merged together.
    """

    def __init__(self, base_first, base_last, chunk_feature, fragment=None):
        self.base_first = base_first
        self.base_last = base_last
        self.feature = chunk_feature.feature
        self.feature_base_first = chunk_feature.feature_base_first
        self.feature_base_last = chunk_feature.feature_base_last
        self.fragment = fragment

    @property
    def feature_name(self):
        if (
            self.feature_base_first != 1
            or self.feature_base_last != self.feature.length
        ):
            return "%s[%s:%s]" % (
                self.feature.name,
                self.feature_base_first,
                self.feature_base_last,
            )
        else:
            return self.feature.name

    def __str__(self):
        s = []
        s.append(self.feature_name)
        s.append(self.feature.type)
        if self.feature.strand == 1:
            s.append("+")
        else:
            s.append("-")
        return ", ".join(s)

    @staticmethod
    def from_chunk_feature_and_location_array(chunk_feature_locs):
        """
        chunk_feature_locs: an array of (Chunk_Feature, Fragment_Chunk_Location) tuples.
        """

        chunk_feature_locs = sorted(
            chunk_feature_locs,
            key=lambda t: (t[0].feature.id, t[1].base_first,
                           (t[0].feature.strand if t[0].feature.strand is not None else 1)
                           * t[0].feature_base_first)
        )

        annotations = []
        for cf, fcl in chunk_feature_locs:
            if (
                len(annotations) > 0
                and annotations[-1].feature.id == cf.feature_id
                and (((annotations[-1].feature.strand is None or annotations[-1].feature.strand > 0)
                      and annotations[-1].feature_base_last == cf.feature_base_first - 1)
                     or ((annotations[-1].feature.strand is not None
                          and annotations[-1].feature.strand < 0)
                         and annotations[-1].feature_base_first == cf.feature_base_last + 1))
                and annotations[-1].base_last == fcl.base_first - 1
            ):
                # merge annotation
                annotations[-1].base_last = fcl.base_last
                if (annotations[-1].feature.strand is None or annotations[-1].feature.strand > 0):
                    annotations[-1].feature_base_last = cf.feature_base_last
                else:
                    annotations[-1].feature_base_first = cf.feature_base_first
            else:
                annotations.append(
                    Annotation(
                        base_first=fcl.base_first,
                        base_last=fcl.base_last,
                        chunk_feature=cf,
                        fragment=fcl.fragment,
                    )
                )

        return sorted(annotations, key=lambda a: a.base_first)

class ChunkReference(object):
    """
    Parent class for storing chunk sequences in reference files
    1-based indexing example: (2, 4) of 5'-A(TCG)GCTA-3'
    """

    @staticmethod
    def generate_from_fragment(fragment, sequence=None, dirn='.'):
        return ChunkReference(None)

    @staticmethod
    def read_reference_sequence_at_position_opened_file(f, start, end):
        """
        Read sequence from 1-indexed positions from reference.
        """
        f.seek(start - 1)
        sequence = f.read(end - start + 1).decode('utf-8')
        return sequence

    def read_reference_sequence_at_position(self, start, end):
        """
        Read sequence from 1-indexed positions from reference.
        """
        with gzip.open(self.ref_fn, 'rb') as f:
            f.seek(start - 1)
            sequence = f.read(end - start + 1).decode('utf-8')
        return sequence

    def __init__(self, ref_fn):
        self.ref_fn = ref_fn

class LocalChunkReference(ChunkReference):
    """
    Class for locally storing chunk sequences in reference files
    """

    @staticmethod
    def generate_from_fragment(fragment, sequence=None, dirn='.'):
        if sequence is None:
            sequence = fragment.indexed_fragment().sequence

        ref_fn = f"{dirn}/{fragment.id}.fa.gz"
        with gzip.open(ref_fn, 'wb') as f:
            sequence = f.write(sequence.encode('utf-8'))

        return LocalChunkReference(ref_fn)

class NetAppChunkReference(ChunkReference):
    """
    Class for storing chunk sequences in reference files on NetApp
    """

    @staticmethod
    def generate_from_fragment(fragment, sequence=None, dirn='.'):
        # NOTE: dirn is ignored in this method, as we're storing on NetApp
        if sequence is None:
            sequence = fragment.indexed_fragment().sequence
        ref_fn = fragment.build_fragment_fasta_from_sequence(sequence)

        return NetAppChunkReference(ref_fn)

class BigIntPrimaryModel(models.Model):
    class Meta:
        app_label = "edge"
        abstract = True

    id = models.BigIntegerField(primary_key=True)

    @transaction.atomic()
    def save(self, *args, **kwargs):
        # mimic auto_increment
        if self.id is None:
            klass = type(self)
            try:
                self.id = (
                    klass.objects.select_for_update()
                    .order_by("-id")
                    .values("id")[0]["id"]
                    + 1
                )
            except IndexError:
                self.id = 1
        return super(BigIntPrimaryModel, self).save(*args, **kwargs)

class Chunk(BigIntPrimaryModel):
    CHUNK_REFERENCE_CLASS = NetAppChunkReference

    class Meta:
        app_label = "edge"

    initial_fragment = models.ForeignKey("Fragment", on_delete=models.PROTECT)
    sequence = models.TextField(null=True)
    ref_start_index = models.PositiveIntegerField(blank=True, null=True)
    ref_end_index = models.PositiveIntegerField(blank=True, null=True)

    @property
    def ref_fn(self):
        return self.initial_fragment.fragment_reference_fasta_gz_fn()

    @property
    def is_sequence_based(self):
        return self.sequence is not None

    @property
    def is_reference_based(self):
        return None not in [self.ref_start_index, self.ref_end_index]

    @property
    def length(self):
        if self.is_sequence_based:
            return len(self.sequence)
        elif self.is_reference_based:
            return self.ref_end_index - self.ref_start_index + 1
        raise Exception("invalid chunk data")

    def get_sequence(self, f=None):
        if self.is_sequence_based:
            return self.sequence
        elif self.is_reference_based:
            if f is None:
                cr = self.CHUNK_REFERENCE_CLASS(self.ref_fn)
                return cr.read_reference_sequence_at_position(
                    self.ref_start_index, self.ref_end_index
                )
            else:
                return self.CHUNK_REFERENCE_CLASS.read_reference_sequence_at_position_opened_file(
                    f, self.ref_start_index, self.ref_end_index
                )
        raise Exception("invalid chunk data")

    def reload(self):
        return Chunk.objects.get(pk=self.pk)

    @staticmethod
    @transaction.atomic()  # this method has to be in a transaction, see below
    def bulk_create(entries):
        # IMPORTANT: this entire method must be within a transaction, so we can
        # compute and assign unique BigIntPrimary IDs

        cur_id = 1
        try:
            cur_id = (
                Chunk.objects.select_for_update()
                .order_by("-id")
                .values("id")[0]["id"]
                + 1
            )
        except IndexError:
            pass

        for entry in entries:
            entry.id = cur_id
            cur_id += 1
        Chunk.objects.bulk_create(entries, batch_size=BULK_CREATE_BATCH_SIZE)

class Edge(BigIntPrimaryModel):
    class Meta:
        app_label = "edge"

    from_chunk = models.ForeignKey(
        Chunk, related_name="out_edges", on_delete=models.PROTECT
    )
    fragment = models.ForeignKey("Fragment", on_delete=models.PROTECT)
    # can be null, so we can supersede an edge from a child fragment
    to_chunk = models.ForeignKey(
        Chunk, null=True, related_name="in_edges", on_delete=models.PROTECT
    )

    @staticmethod
    @transaction.atomic()  # this method has to be in a transaction, see below
    def bulk_create(entries):
        # IMPORTANT: this entire method must be within a transaction, so we can
        # compute and assign unique BigIntPrimary IDs

        cur_id = 1
        try:
            cur_id = (
                Edge.objects.select_for_update()
                .order_by("-id")
                .values("id")[0]["id"]
                + 1
            )
        except IndexError:
            pass

        for entry in entries:
            entry.id = cur_id
            cur_id += 1
        Edge.objects.bulk_create(entries, batch_size=BULK_CREATE_BATCH_SIZE)

class Feature(models.Model):
    class Meta:
        app_label = "edge"

    name = models.CharField(max_length=100)
    type = models.CharField(max_length=100)
    strand = models.IntegerField(null=True)
    length = models.IntegerField()
    operation = models.ForeignKey("Operation", null=True, on_delete=models.CASCADE)
    _qualifiers = models.TextField(null=True, db_column="qualifiers")

    def set_qualifiers(self, qualifiers):
        self._qualifiers = json.dumps(qualifiers)

    @property
    def qualifiers(self):
        if self._qualifiers is not None:
            return json.loads(self._qualifiers)
        else:
            return None

    @staticmethod
    @transaction.atomic()  # this method has to be in a transaction, see below
    def bulk_create(entries):
        # IMPORTANT: this entire method must be within a transaction, so we can
        # compute and assign unique BigIntPrimary IDs

        cur_id = 1
        try:
            cur_id = (
                Feature.objects.select_for_update()
                .order_by("-id")
                .values("id")[0]["id"]
                + 1
            )
        except IndexError:
            pass

        for entry in entries:
            entry.id = cur_id
            cur_id += 1
        Feature.objects.bulk_create(entries, batch_size=BULK_CREATE_BATCH_SIZE)

class Chunk_Feature_Manager(models.Manager):
    def get_query_set(self):
        return (
            super(Chunk_Feature_Manager, self)
            .get_query_set()
            .select_related("chunk", "feature")
        )

class Chunk_Feature(BigIntPrimaryModel):
    class Meta:
        app_label = "edge"

    objects = Chunk_Feature_Manager()
    chunk = models.ForeignKey(Chunk, on_delete=models.PROTECT)
    feature = models.ForeignKey(Feature, on_delete=models.PROTECT)
    feature_base_first = models.IntegerField()
    feature_base_last = models.IntegerField()

    @staticmethod
    @transaction.atomic()  # this method has to be in a transaction, see below
    def bulk_create(entries):
        # IMPORTANT: this entire method must be within a transaction, so we can
        # compute and assign unique BigIntPrimary IDs

        cur_id = 1
        try:
            cur_id = (
                Chunk_Feature.objects.select_for_update()
                .order_by("-id")
                .values("id")[0]["id"]
                + 1
            )
        except IndexError:
            pass

        for entry in entries:
            entry.id = cur_id
            cur_id += 1
        Chunk_Feature.objects.bulk_create(entries, batch_size=BULK_CREATE_BATCH_SIZE)

class Fragment_Chunk_Location(BigIntPrimaryModel):
    class Meta:
        app_label = "edge"
        unique_together = (("fragment", "chunk"),)
        index_together = (("fragment", "base_last"), ("fragment", "base_first"))

    fragment = models.ForeignKey("Fragment", on_delete=models.PROTECT)
    chunk = models.ForeignKey(Chunk, on_delete=models.PROTECT)
    base_first = models.IntegerField()
    base_last = models.IntegerField()

    def __str__(self):
        return "%s: %s-%s" % (self.fragment.name, self.location[0], self.location[1])

    @property
    def next_chunk(self):
        for fcl in self.fragment.fragment_chunk_location_set.select_related(
            "chunk"
        ).filter(base_first=self.base_last + 1):
            return fcl.chunk
        return None

    @property
    def location(self):
        return (self.base_first, self.base_last)

    @property
    def prev_fragment_chunk(self):
        loc = self.location
        if loc[0] == 1:
            return None
        fc = Fragment_Chunk_Location.objects.get(
            fragment=self.fragment, base_last=loc[0] - 1
        )
        return fc

    def annotations(self):
        return [
            Annotation(
                base_first=self.base_first, base_last=self.base_last, chunk_feature=cf
            )
            for cf in self.chunk.chunk_feature_set.all()
        ]

    @staticmethod
    @transaction.atomic()  # this method has to be in a transaction, see below
    def bulk_create(entries):
        # IMPORTANT: this entire method must be within a transaction, so we can
        # compute and assign unique BigIntPrimary IDs

        cur_id = 1
        try:
            cur_id = (
                Fragment_Chunk_Location.objects.select_for_update()
                .order_by("-id")
                .values("id")[0]["id"]
                + 1
            )
        except IndexError:
            pass

        for entry in entries:
            entry.id = cur_id
            cur_id += 1
        Fragment_Chunk_Location.objects.bulk_create(entries, batch_size=BULK_CREATE_BATCH_SIZE)
<|endoftext|>


<commit_before>import os
import logging
import errno
import subprocess
import which

LOGGER = logging.getLogger(__name__)

COMPRESSION_METHODS = {
    'gzip'  : ('gzip', '.gz'),
    'pigz'  : ('pigz', '.gz'),
    'bzip2' : ('bzip2', '.bz2'),
    'pbzip2': ('pzip2', '.bz2'),
    'lzop'  : ('lzop', '.lzo'),
    'lzma'  : ('xz', '.xz'),
}

def lookup_compression(method):
    """
    Looks up the passed compression method in supported COMPRESSION_METHODS
    and returns a tuple in the form of ('command_name', 'file_extension').

    Arguments:

    method -- A string identifier of the compression method (i.e. 'gzip').
    """
    try:
        cmd, ext = COMPRESSION_METHODS[method]
        try:
            return which.which(cmd), ext
        except which.WhichError, e:
            raise OSError("No command found for compression method '%s'" %
                    method)
    except KeyError:
        raise OSError("Unsupported compression method '%s'" % method)

class CompressionInput(object):
    """
    Class to create a compressed file descriptor for reading.  Functions like
    a standard file descriptor such as from open().
    """
    def __init__(self, path, mode, cmd, bufsize=1024*1024):
        self.fileobj = open(path, 'r')
        self.pid = subprocess.Popen([cmd, '--decompress'],
                                    stdin=self.fileobj.fileno(),
                                    stdout=subprocess.PIPE,
                                    bufsize=bufsize)
        self.fd = self.pid.stdout.fileno()
        self.name = path
        self.closed = False

    def fileno(self):
        return self.fd

    def read(self, size):
        return os.read(self.fd, size)

    def next(self):
        return self.pid.stdout.next()

    def __iter__(self):
        return iter(self.pid.stdout)

    def close(self):
        import signal
        os.kill(self.pid.pid, signal.SIGTERM)
        self.fileobj.close()
        self.pid.stdout.close()
        self.pid.wait()
        self.closed = True

class CompressionOutput(object):
    """
    Class to create a compressed file descriptor for writing.  Functions like
    a standard file descriptor such as from open().
    """
    def __init__(self, path, mode, cmd, level, inline):
        self.cmd = cmd
        self.level = level
        self.inline = inline
        if not inline:
            self.fileobj = open(os.path.splitext(path)[0], mode)
            self.fd = self.fileobj.fileno()
        else:
            self.fileobj = open(path, 'w')
            args = [cmd]
            if level:
                args += ['-%d' % level]
            self.pid = subprocess.Popen(args,
                                        stdin=subprocess.PIPE,
                                        stdout=self.fileobj.fileno(),
                                        stderr=subprocess.PIPE)
            self.fd = self.pid.stdin.fileno()
        self.name = path
        self.closed = False

    def fileno(self):
        return self.fd

    def write(self, data):
        return os.write(self.fd, data)

    def close(self):
        self.closed = True
        if not self.inline:
            args = [self.cmd]
            if self.level:
                args += ['-%d' % self.level, '-']
            self.fileobj.close()
            self.fileobj = open(self.fileobj.name, 'r')
            cmp_f = open(self.name, 'w')
            LOGGER.debug("Running %r < %r[%d] > %r[%d]",
                         args, self.fileobj.name, self.fileobj.fileno(),
                         cmp_f.name, cmp_f.fileno())
            pid = subprocess.Popen(args,
                                   stdin=self.fileobj.fileno(),
                                   stdout=cmp_f.fileno())
            status = pid.wait()
            os.unlink(self.fileobj.name)
        else:
            self.pid.stdin.close()
            # Check for anything on stderr
            for line in self.pid.stderr:
                errmsg = line.strip()
                if not errmsg:
                    # gzip, among others, output a spurious blank line
                    continue
                LOGGER.error("Compression Error: %s", errmsg)
            self.fileobj.close()
            status = self.pid.wait()
            if status != 0:
                raise IOError(errno.EPIPE,
                              "Compression program %r exited with status %d" %
                                (self.cmd, status))

def stream_info(path, method=None, level=None):
    """
    Determine compression command, and compressed path based on original path
    and compression method.  If method is not passed, or level is 0 the
    original path is returned.

    Arguments:

    path    -- Path to file to compress/decompress
    method  -- Compression method (i.e. 'gzip', 'bzip2', 'pbzip2', 'lzop')
    level   -- Compression level (0-9)
    """
    if not method or level == 0:
        return path

    cmd, ext = lookup_compression(method)

    if not cmd:
        raise IOError("Unknown compression method '%s'" % cmd)

    if not path.endswith(ext):
        path += ext

    return cmd, path

def open_stream(path, mode, method=None, level=None, inline=True):
    """
    Opens a compressed data stream, and returns a file descriptor type object
    that acts much like os.open() does.  If no method is passed, or the
    compression level is 0, simply returns a file descriptor from open().

    Arguments:

    mode    -- File access mode (i.e. 'r' or 'w')
    method  -- Compression method (i.e. 'gzip', 'bzip2', 'pbzip2', 'lzop')
    level   -- Compression level
    inline  -- Boolean whether to compress inline, or after the file is written.
    """
    if not method or method == 'none' or level == 0:
        return open(path, mode)
    else:
        cmd, path = stream_info(path, method)
        if not cmd:
            raise IOError("Unknown compression method '%s'" % cmd)

        if mode == 'r':
            return CompressionInput(path, mode, cmd=cmd)
        elif mode == 'w':
            return CompressionOutput(path, mode, cmd=cmd, level=level,
                                     inline=inline)
        else:
            raise IOError("invalid mode: %s" % mode)
<commit_msg>Fix typo in pbzip2.<commit_after>import os
import logging
import errno
import subprocess
import which

LOGGER = logging.getLogger(__name__)

COMPRESSION_METHODS = {
    'gzip'  : ('gzip', '.gz'),
    'pigz'  : ('pigz', '.gz'),
    'bzip2' : ('bzip2', '.bz2'),
    'pbzip2': ('pbzip2', '.bz2'),
    'lzop'  : ('lzop', '.lzo'),
    'lzma'  : ('xz', '.xz'),
}

def lookup_compression(method):
    """
    Looks up the passed compression method in supported COMPRESSION_METHODS
    and returns a tuple in the form of ('command_name', 'file_extension').

    Arguments:

    method -- A string identifier of the compression method (i.e. 'gzip').
    """
    try:
        cmd, ext = COMPRESSION_METHODS[method]
        try:
            return which.which(cmd), ext
        except which.WhichError, e:
            raise OSError("No command found for compression method '%s'" %
                    method)
    except KeyError:
        raise OSError("Unsupported compression method '%s'" % method)

class CompressionInput(object):
    """
    Class to create a compressed file descriptor for reading.  Functions like
    a standard file descriptor such as from open().
    """
    def __init__(self, path, mode, cmd, bufsize=1024*1024):
        self.fileobj = open(path, 'r')
        self.pid = subprocess.Popen([cmd, '--decompress'],
                                    stdin=self.fileobj.fileno(),
                                    stdout=subprocess.PIPE,
                                    bufsize=bufsize)
        self.fd = self.pid.stdout.fileno()
        self.name = path
        self.closed = False

    def fileno(self):
        return self.fd

    def read(self, size):
        return os.read(self.fd, size)

    def next(self):
        return self.pid.stdout.next()

    def __iter__(self):
        return iter(self.pid.stdout)

    def close(self):
        import signal
        os.kill(self.pid.pid, signal.SIGTERM)
        self.fileobj.close()
        self.pid.stdout.close()
        self.pid.wait()
        self.closed = True

class CompressionOutput(object):
    """
    Class to create a compressed file descriptor for writing.  Functions like
    a standard file descriptor such as from open().
    """
    def __init__(self, path, mode, cmd, level, inline):
        self.cmd = cmd
        self.level = level
        self.inline = inline
        if not inline:
            self.fileobj = open(os.path.splitext(path)[0], mode)
            self.fd = self.fileobj.fileno()
        else:
            self.fileobj = open(path, 'w')
            args = [cmd]
            if level:
                args += ['-%d' % level]
            self.pid = subprocess.Popen(args,
                                        stdin=subprocess.PIPE,
                                        stdout=self.fileobj.fileno(),
                                        stderr=subprocess.PIPE)
            self.fd = self.pid.stdin.fileno()
        self.name = path
        self.closed = False

    def fileno(self):
        return self.fd

    def write(self, data):
        return os.write(self.fd, data)

    def close(self):
        self.closed = True
        if not self.inline:
            args = [self.cmd]
            if self.level:
                args += ['-%d' % self.level, '-']
            self.fileobj.close()
            self.fileobj = open(self.fileobj.name, 'r')
            cmp_f = open(self.name, 'w')
            LOGGER.debug("Running %r < %r[%d] > %r[%d]",
                         args, self.fileobj.name, self.fileobj.fileno(),
                         cmp_f.name, cmp_f.fileno())
            pid = subprocess.Popen(args,
                                   stdin=self.fileobj.fileno(),
                                   stdout=cmp_f.fileno())
            status = pid.wait()
            os.unlink(self.fileobj.name)
        else:
            self.pid.stdin.close()
            # Check for anything on stderr
            for line in self.pid.stderr:
                errmsg = line.strip()
                if not errmsg:
                    # gzip, among others, output a spurious blank line
                    continue
                LOGGER.error("Compression Error: %s", errmsg)
            self.fileobj.close()
            status = self.pid.wait()
            if status != 0:
                raise IOError(errno.EPIPE,
                              "Compression program %r exited with status %d" %
                                (self.cmd, status))

def stream_info(path, method=None, level=None):
    """
    Determine compression command, and compressed path based on original path
    and compression method.  If method is not passed, or level is 0 the
    original path is returned.

    Arguments:

    path    -- Path to file to compress/decompress
    method  -- Compression method (i.e. 'gzip', 'bzip2', 'pbzip2', 'lzop')
    level   -- Compression level (0-9)
    """
    if not method or level == 0:
        return path

    cmd, ext = lookup_compression(method)

    if not cmd:
        raise IOError("Unknown compression method '%s'" % cmd)

    if not path.endswith(ext):
        path += ext

    return cmd, path

def open_stream(path, mode, method=None, level=None, inline=True):
    """
    Opens a compressed data stream, and returns a file descriptor type object
    that acts much like os.open() does.  If no method is passed, or the
    compression level is 0, simply returns a file descriptor from open().

    Arguments:

    mode    -- File access mode (i.e. 'r' or 'w')
    method  -- Compression method (i.e. 'gzip', 'bzip2', 'pbzip2', 'lzop')
    level   -- Compression level
    inline  -- Boolean whether to compress inline, or after the file is written.
    """
    if not method or method == 'none' or level == 0:
        return open(path, mode)
    else:
        cmd, path = stream_info(path, method)
        if not cmd:
            raise IOError("Unknown compression method '%s'" % cmd)

        if mode == 'r':
            return CompressionInput(path, mode, cmd=cmd)
        elif mode == 'w':
            return CompressionOutput(path, mode, cmd=cmd, level=level,
                                     inline=inline)
        else:
            raise IOError("invalid mode: %s" % mode)
<|endoftext|>


<commit_before>import pandas as pd
import matplotlib.pyplot as plt
import networkx as nx
import argparse

def load_row(G, row, graphframe):
    src = row[1]
    dst = row[2]
    if src not in G:
        G.add_node(src, role=graphframe.ix[src][3])
    if dst not in G:
        G.add_node(dst, role=graphframe.ix[dst][4])
    G.add_edge(src, dst)

def mapcolors(nodelist):
    """Map each role to a color pass the values to nx.draw_networkx"""
    colors = dict()

    def color(role):
        """return the color string"""
        if 'PAT' in role:
            col = 'y'
        elif 'NUR' in role:
            col = 'r'
        elif 'MED' in role:
            col = 'g'
        elif 'ADM' in role:
            col = 'b'
        else:
            print(role)
            col = 'k'
        return col

    for vertex, properties in nodelist:
        colors[vertex] = color(properties["role"])
    return colors

def draw_graph(G,
               color_func=mapcolors,
               layout_func=nx.drawing.layout.spring_layout):
    """Draw the graph using the colors and layout that we picked"""
    colors = color_func(G.nodes(data=True))
    layout = layout_func(G)
    nx.draw_networkx(G, pos=layout, node_color=colors.values())

def get_args():
    """Produce the args namespace based on command line arguments"""
    parser = argparse.ArgumentParser()
    parser.add_argument('filename')
    parser.add_argument('--draw-graph', action='store_true')
    parser.add_argument('--drawfile', type=str)
    arg = parser.parse_args()
    return arg

if __name__ == '__main__':
    arg = get_args()
    graphframe = pd.read_table(arg.filename, header=None)
    G = nx.Graph()
    insertion_count = 0
    batchsize = 100
    for row_tuple in graphframe.iterrows():
        load_row(G, row_tuple[1], graphframe)
        insertion_count += 1
        if insertion_count % batchsize == 0:
            #do something for each batch
            pass
    if arg.draw_graph:
        draw_graph(G)
        if arg.drawfile:
            plt.savefig(arg.drawfile)
        else:
            plt.show()
<commit_msg>add code to produce graph matrices in matlab format<commit_after>import pandas as pd
import matplotlib.pyplot as plt
import networkx as nx
import argparse
import sys

def load_row(G, row, graphframe):
    src = row[1]
    dst = row[2]
    if src not in G:
        G.add_node(src, role=graphframe.ix[src][3])
    if dst not in G:
        G.add_node(dst, role=graphframe.ix[dst][4])
    G.add_edge(src, dst)

def mapcolors(nodelist):
    """Map each role to a color pass the values to nx.draw_networkx"""
    colors = dict()

    def color(role):
        """return the color string"""
        if 'PAT' in role:
            col = 'y'
        elif 'NUR' in role:
            col = 'r'
        elif 'MED' in role:
            col = 'g'
        elif 'ADM' in role:
            col = 'b'
        else:
            print(role)
            col = 'k'
        return col

    for vertex, properties in nodelist:
        colors[vertex] = color(properties["role"])
    return colors

def draw_graph(G,
               color_func=mapcolors,
               layout_func=nx.drawing.layout.spring_layout):
    """Draw the graph using the colors and layout that we picked"""
    colors = color_func(G.nodes(data=True))
    layout = layout_func(G)
    nx.draw_networkx(G, pos=layout, node_color=colors.values())

def get_args():
    """Produce the args namespace based on command line arguments"""
    parser = argparse.ArgumentParser()
    parser.add_argument('filename')
    parser.add_argument('--draw-graph', action='store_true')
    parser.add_argument('--drawfile', type=str)
    parser.add_argument('--matfile', type=str,
                        help='Store the graph matrices into this file in matlab format')
    arg = parser.parse_args()
    return arg

def write_adj_mat(G, fileobj=sys.stdout):
    """Write G to a sparse matrix format that Julia and Matlab can read."""
    lapmatrix = nx.laplacian_matrix(G)
    norm_lapl = nx.normalized_laplacian_matrix(G)
    adjmatrix = nx.adjacency_matrix(G)
    mdict = {'laplacian': lapmatrix,
             'norm_lapl': norm_lapl,
             'adjacency': adjmatrix}
    sio.savemat(fileobj, mdict)
    return mdict

if __name__ == '__main__':
    arg = get_args()
    graphframe = pd.read_table(arg.filename, header=None)
    G = nx.Graph()
    insertion_count = 0
    batchsize = 100
    for row_tuple in graphframe.iterrows():
        load_row(G, row_tuple[1], graphframe)
        insertion_count += 1
        if insertion_count % batchsize == 0:
            #do something for each batch
            pass
    if arg.draw_graph:
        draw_graph(G)
        if arg.drawfile:
            plt.savefig(arg.drawfile)
        else:
            plt.show()
    if arg.matfile:
        import scipy.io as sio
        write_adj_mat(G, arg.matfile)
<|endoftext|>


<commit_before># Copyright ClusterHQ Inc.  See LICENSE file for details.

"""
Functionality for creating ``IBlockDeviceAPI`` providers suitable for use in
the current execution environment.

This depends on a ``CLOUD_CONFIG_FILE`` environment variable being set.

See `acceptance testing <acceptance-testing>`_ for details.

.. code-block:: python

    from .blockdevicefactory import ProviderType, get_blockdeviceapi

    api = get_blockdeviceapi(ProviderType.openstack)
    volume = api.create_volume(...)

"""

from os import environ
from uuid import uuid4

from yaml import safe_load

from twisted.trial.unittest import SkipTest
from twisted.python.constants import Names, NamedConstant

from keystoneclient_rackspace.v2_0 import RackspaceAuth
from keystoneclient.session import Session

from cinderclient.client import Client as CinderClient
from novaclient.client import Client as NovaClient

from ..cinder import CinderBlockDeviceAPI
from ..ebs import EBSBlockDeviceAPI, ec2_client
from ..test.test_blockdevice import detach_destroy_volumes

# The Rackspace authentication endpoint
# See http://docs.rackspace.com/cbs/api/v1.0/cbs-devguide/content/Authentication-d1e647.html # noqa
RACKSPACE_AUTH_URL = "https://identity.api.rackspacecloud.com/v2.0"

class ConfigMissing(Exception):
    """
    The cloud configuration could not be found.
    """

class ProviderType(Names):
    """
    Kinds of compute/storage cloud providers for which this module is able to
    build ``IBlockDeviceAPI`` providers.
    """
    openstack = NamedConstant()
    aws = NamedConstant()

def get_blockdeviceapi(provider):
    """
    Validate and load cloud provider's yml config file.
    Default to ``~/acceptance.yml`` in the current user home directory, since
    that's where buildbot puts its acceptance test credentials file.
    """
    cls, args = get_blockdeviceapi_args(provider)
    return cls(**args)

def get_blockdeviceapi_args(provider):
    """
    Get initializer arguments suitable for use in the instantiation of an
    ``IBlockDeviceAPI`` implementation compatible with the given provider.

    :param provider: A provider type the ``IBlockDeviceAPI`` is to be
        compatible with.  A value from ``ProviderType``.

    :raises: ``ConfigMissing`` if a ``CLOUD_CONFIG_FILE`` was not set and the
        default config file could not be read.

    :return: A ``dict`` that can initialize the matching implementation.
    """
    config_file_path = environ.get('CLOUD_CONFIG_FILE')
    if config_file_path is None:
        raise ConfigMissing(
            'Supply the path to a cloud credentials file '
            'using the CLOUD_CONFIG_FILE environment variable. '
            'See: '
            'https://docs.clusterhq.com/en/latest/gettinginvolved/acceptance-testing.html '  # noqa
            'for details of the expected format.'
        )

    with open(config_file_path) as config_file:
        config = safe_load(config_file.read())

    section = config[provider.name]
    cls, get_kwargs = _BLOCKDEVICE_TYPES[provider]
    kwargs = dict(cluster_id=uuid4())
    kwargs.update(get_kwargs(**section))
    return cls, kwargs

def _rackspace_session(username, key, **kwargs):
    """
    Create a Keystone session capable of authenticating with Rackspace.

    :param unicode keyname: A RackSpace API username.
    :param unicode key: A RackSpace API key.

    :return: A ``keystoneclient.session.Session``.
    """
    auth = RackspaceAuth(
        auth_url=RACKSPACE_AUTH_URL,
        username=username,
        api_key=key
    )
    return Session(auth=auth)

def _openstack(region, **config):
    """
    Create Cinder and Nova volume managers suitable for use in the creation of
    a ``CinderBlockDeviceAPI``.

    :param bytes region: The name of the region to which to connect.
    :param config: Any additional configuration (possibly provider-specific)
        necessary to authenticate a session for use with the CinderClient and
        NovaClient.

    :return: A ``dict`` giving initializer arguments for
        ``CinderBlockDeviceAPI``.
    """
    # TODO: Look up the right session factory in the config and use it here
    # instead of assuming Rackspace.
    session = _rackspace_session(**config)
    cinder_client = CinderClient(
        session=session, region_name=region, version=1
    )
    nova_client = NovaClient(
        session=session, region_name=region, version=2
    )
    return dict(
        cinder_volume_manager=cinder_client.volumes,
        nova_volume_manager=nova_client.volumes,
    )

def _aws(**config):
    """
    Create an EC2 client suitable for use in the creation of an
    ``EBSBlockDeviceAPI``.

    See ``flocker.node.agents.ebs.ec2_client`` for parameter documentation.
    """
    return dict(
        ec2_client=ec2_client(**config),
    )

_BLOCKDEVICE_TYPES = {
    ProviderType.openstack: (CinderBlockDeviceAPI, _openstack),
    ProviderType.aws: (_aws, EBSBlockDeviceAPI),
}

# ^^^^^^^^^^^^^^^^^^^^ generally useful implementation code, put it somewhere
# nice and use it
#
#
# vvvvvvvvvvvvvvvvvvvv testing helper that actually belongs in this module

def get_blockdeviceapi_with_cleanup(test_case, provider):
    """
    Instantiate an ``IBlockDeviceAPI`` implementation appropriate to the given
    provider and configured to work in the current environment.  Arrange for
    all volumes created by it to be cleaned up at the end of the current test
    run.

    :param TestCase test_case: The running test.
    :param provider: A provider type the ``IBlockDeviceAPI`` is to be
        compatible with.  A value from ``ProviderType``.

    :raises: ``SkipTest`` if a ``CLOUD_CONFIG_FILE`` was not set and the
        default config file could not be read.

    :return: The new ``IBlockDeviceAPI`` provider.
    """
    try:
        api = get_blockdeviceapi(provider)
    except ConfigMissing as e:
        raise SkipTest(str(e))
    test_case.addCleanup(detach_destroy_volumes, api)
    return api
<commit_msg>Doc fix<commit_after># Copyright ClusterHQ Inc.  See LICENSE file for details.

"""
Functionality for creating ``IBlockDeviceAPI`` providers suitable for use in
the current execution environment.

This depends on a ``CLOUD_CONFIG_FILE`` environment variable being set.

See `acceptance testing <acceptance-testing>`_ for details.

.. code-block:: python

    from .blockdevicefactory import ProviderType, get_blockdeviceapi

    api = get_blockdeviceapi(ProviderType.openstack)
    volume = api.create_volume(...)

"""

from os import environ
from uuid import uuid4

from yaml import safe_load

from twisted.trial.unittest import SkipTest
from twisted.python.constants import Names, NamedConstant

from keystoneclient_rackspace.v2_0 import RackspaceAuth
from keystoneclient.session import Session

from cinderclient.client import Client as CinderClient
from novaclient.client import Client as NovaClient

from ..cinder import CinderBlockDeviceAPI
from ..ebs import EBSBlockDeviceAPI, ec2_client
from ..test.test_blockdevice import detach_destroy_volumes

# The Rackspace authentication endpoint
# See http://docs.rackspace.com/cbs/api/v1.0/cbs-devguide/content/Authentication-d1e647.html # noqa
RACKSPACE_AUTH_URL = "https://identity.api.rackspacecloud.com/v2.0"

class ConfigMissing(Exception):
    """
    The cloud configuration could not be found.
    """

class ProviderType(Names):
    """
    Kinds of compute/storage cloud providers for which this module is able to
    build ``IBlockDeviceAPI`` providers.
    """
    openstack = NamedConstant()
    aws = NamedConstant()

def get_blockdeviceapi(provider):
    """
    Validate and load cloud provider's yml config file.
    Default to ``~/acceptance.yml`` in the current user home directory, since
    that's where buildbot puts its acceptance test credentials file.
    """
    cls, args = get_blockdeviceapi_args(provider)
    return cls(**args)

def get_blockdeviceapi_args(provider):
    """
    Get initializer arguments suitable for use in the instantiation of an
    ``IBlockDeviceAPI`` implementation compatible with the given provider.

    :param provider: A provider type the ``IBlockDeviceAPI`` is to be
        compatible with.  A value from ``ProviderType``.

    :raises: ``ConfigMissing`` if a ``CLOUD_CONFIG_FILE`` was not set and the
        default config file could not be read.

    :return: A two-tuple of an ``IBlockDeviceAPI`` implementation and a
        ``dict`` of keyword arguments that can be used instantiate that
        implementation.
    """
    config_file_path = environ.get('CLOUD_CONFIG_FILE')
    if config_file_path is None:
        raise ConfigMissing(
            'Supply the path to a cloud credentials file '
            'using the CLOUD_CONFIG_FILE environment variable. '
            'See: '
            'https://docs.clusterhq.com/en/latest/gettinginvolved/acceptance-testing.html '  # noqa
            'for details of the expected format.'
        )

    with open(config_file_path) as config_file:
        config = safe_load(config_file.read())

    section = config[provider.name]
    cls, get_kwargs = _BLOCKDEVICE_TYPES[provider]
    kwargs = dict(cluster_id=uuid4())
    kwargs.update(get_kwargs(**section))
    return cls, kwargs

def _rackspace_session(username, key, **kwargs):
    """
    Create a Keystone session capable of authenticating with Rackspace.

    :param unicode keyname: A RackSpace API username.
    :param unicode key: A RackSpace API key.

    :return: A ``keystoneclient.session.Session``.
    """
    auth = RackspaceAuth(
        auth_url=RACKSPACE_AUTH_URL,
        username=username,
        api_key=key
    )
    return Session(auth=auth)

def _openstack(region, **config):
    """
    Create Cinder and Nova volume managers suitable for use in the creation of
    a ``CinderBlockDeviceAPI``.

    :param bytes region: The name of the region to which to connect.
    :param config: Any additional configuration (possibly provider-specific)
        necessary to authenticate a session for use with the CinderClient and
        NovaClient.

    :return: A ``dict`` giving initializer arguments for
        ``CinderBlockDeviceAPI``.
    """
    # TODO: Look up the right session factory in the config and use it here
    # instead of assuming Rackspace.
    session = _rackspace_session(**config)
    cinder_client = CinderClient(
        session=session, region_name=region, version=1
    )
    nova_client = NovaClient(
        session=session, region_name=region, version=2
    )
    return dict(
        cinder_volume_manager=cinder_client.volumes,
        nova_volume_manager=nova_client.volumes,
    )

def _aws(**config):
    """
    Create an EC2 client suitable for use in the creation of an
    ``EBSBlockDeviceAPI``.

    See ``flocker.node.agents.ebs.ec2_client`` for parameter documentation.
    """
    return dict(
        ec2_client=ec2_client(**config),
    )

_BLOCKDEVICE_TYPES = {
    ProviderType.openstack: (CinderBlockDeviceAPI, _openstack),
    ProviderType.aws: (_aws, EBSBlockDeviceAPI),
}

# ^^^^^^^^^^^^^^^^^^^^ generally useful implementation code, put it somewhere
# nice and use it
#
#
# vvvvvvvvvvvvvvvvvvvv testing helper that actually belongs in this module

def get_blockdeviceapi_with_cleanup(test_case, provider):
    """
    Instantiate an ``IBlockDeviceAPI`` implementation appropriate to the given
    provider and configured to work in the current environment.  Arrange for
    all volumes created by it to be cleaned up at the end of the current test
    run.

    :param TestCase test_case: The running test.
    :param provider: A provider type the ``IBlockDeviceAPI`` is to be
        compatible with.  A value from ``ProviderType``.

    :raises: ``SkipTest`` if a ``CLOUD_CONFIG_FILE`` was not set and the
        default config file could not be read.

    :return: The new ``IBlockDeviceAPI`` provider.
    """
    try:
        api = get_blockdeviceapi(provider)
    except ConfigMissing as e:
        raise SkipTest(str(e))
    test_case.addCleanup(detach_destroy_volumes, api)
    return api
<|endoftext|>


<commit_before>import logging

import numpy as np

from sklearn.preprocessing import label_binarize
from sklearn.preprocessing import LabelEncoder
from sklearn.preprocessing import MultiLabelBinarizer

from .base import BaseClassifier

__all__ = ['LabelsClassifier', 'BinaryLabelsClassifier', 'MultiLabelsClassifier', 'MulticlassLabelsClassifier']

logger = logging.getLogger(__name__)

class LabelsClassifier(BaseClassifier):
    def __init__(self, include=[], exclude=[], **kwargs):
        super(LabelsClassifier, self).__init__(**kwargs)

        self.exclude = set(exclude)
        self.include = set(include)
    #end def

    def _fit(self, X, Y_labels, *, binarize_args={}, fit_args={}, **kwargs):
        Y_binarized = self.binarize_labels(Y_labels, **binarize_args)
        return self.fit_binarized(X, Y_binarized, **fit_args)
    #end def

    def fit_binarized(self, X_featurized, Y_binarized, **kwargs): raise NotImplementedError('fit_binarized is not implemented.')

    def predict(self, X_featurized, binarized=True, **kwargs):
        Y_predict_binarized = super(LabelsClassifier, self).predict(X_featurized, **kwargs)
        if binarized: return Y_predict_binarized

        return self.unbinarize_labels(Y_predict_binarized)
    #end def

    def predict_and_proba(self, X_featurized, *, binarized=True, **kwargs):
        Y_proba, Y_predict_binarized = super(LabelsClassifier, self).predict_and_proba(X_featurized, **kwargs)
        if Y_proba.ndim == 1:
            Y_proba_2d = np.zeros((Y_proba.shape[0], 2), dtype=Y_proba.dtype)
            Y_proba_2d[:, 0] = Y_proba
            Y_proba_2d[:, 1] = 1.0 - Y_proba
            Y_proba = Y_proba_2d
        #end if

        if binarized: return Y_proba, Y_predict_binarized

        return Y_proba, self.unbinarize_labels(Y_predict_binarized)
    #end def

    def binarize_labels(self, Y_labels, **kwargs): raise NotImplementedError('binarize_labels is not implemented.')

    def binarize_dicts(self, Y_dicts, *, default=0.0): raise NotImplementedError('binarize_dicts is not implemented.')

    def unbinarize_labels(self, Y_binarized, *, epsilon=0.0): raise NotImplementedError('unbinarize_labels is not implemented.')

    def _filter_labels(self, Y_labels):
        if self.exclude or self.include:
            Y_labels_filtered = np.empty(Y_labels.shape, dtype=np.object)
            removed_labels = 0
            for i in range(Y_labels.shape[0]):
                Y_labels_filtered[i] = [l for l in Y_labels[i] if (l in self.include or not self.include) and l not in self.exclude]
                removed_labels += len(Y_labels[i]) - len(Y_labels_filtered[i])
            #end for
            logger.info('{} label-instances removed from the training data.'.format(removed_labels))
        else: Y_labels_filtered = Y_labels

        return Y_labels_filtered
    #end def
#end def

class BinaryLabelsClassifier(LabelsClassifier):
    def __init__(self, pos_label, *, not_pos_label=None, **kwargs):
        super(BinaryLabelsClassifier, self).__init__(**kwargs)

        self.pos_label = pos_label
        self.not_pos_label = 'not ' + pos_label if not_pos_label is None else not_pos_label
    #end def

    def predict_and_proba(self, *args, **kwargs):
        Y_proba, Y_predict = super(BinaryLabelsClassifier, self).predict_and_proba(*args, **kwargs)
        binarized = kwargs.get('binarized', True)

        if binarized:
            Y_predict = Y_predict[:, 0]  # must be 0th one coz classes_ only has 1 thing

        return Y_proba, Y_predict
    #end def

    def binarize_labels(self, Y_labels):
        if Y_labels.shape[0] == 0:
            return np.zeros((0, 1))

        Y_labels_pos = [self.pos_label if self.pos_label in Y_labels[i] else self.not_pos_label for i in range(Y_labels.shape[0])]
        Y_binarized = label_binarize(Y_labels_pos, classes=[self.pos_label, self.not_pos_label]).reshape(Y_labels.shape[0])

        return Y_binarized
    #end def

    def binarize_dicts(self, Y_dicts, *, default=0.0):
        return np.array([Y_dicts[i].get(self.pos_label, default) for i in range(Y_dicts.shape[0])])
    #end def

    def unbinarize_labels(self, Y_binarized, *, epsilon=0.0):
        return np.array([[self.pos_label] if Y_binarized[i] > epsilon else [] for i in range(Y_binarized.shape[0])], dtype=np.object)
    #end def

    @property
    def classes_(self):
        return [self.pos_label, self.not_pos_label]
#end def

# TODO: use multilabel format and inherit multilabelclassifier
class MulticlassLabelsClassifier(LabelsClassifier):
    def __init__(self, **kwargs):
        super(MulticlassLabelsClassifier, self).__init__(**kwargs)
    #end def

    def _fit(self, X, Y_labels, *, binarize_args={}, fit_args={}, **kwargs):
        Y_labels_filtered = self._filter_labels(Y_labels)
        Y_labels_filtered = np.array([Y_labels_filtered[i][0] if Y_labels_filtered[i] else '<none>' for i in range(Y_labels_filtered.shape[0])])

        self.label_encoder_ = LabelEncoder().fit(Y_labels_filtered)
        logger.info('{} labels found in training instances.'.format(len(self.classes_)))

        if not len(self.classes_): raise ValueError('There are no labels available for fitting model.')

        return super(MulticlassLabelsClassifier, self)._fit(X, Y_labels, binarize_args=binarize_args, fit_args=fit_args, **kwargs)
    #end def

    def predict_and_proba(self, X_featurized, *, binarized=True, **kwargs):
        Y_proba = self.predict_proba(X_featurized, **kwargs)
        Y_predict_binarized = np.argmax(Y_proba, axis=1)

        if binarized: return Y_proba, Y_predict_binarized

        return Y_proba, self.unbinarize_labels(Y_predict_binarized)
    #end def

    def predict_proba(self, X_featurized, **kwargs):
        if 'thresholds' in kwargs: logger.warning('thresholds argument is ignored by <{}>.'.format(self))
        if 'denominators' in kwargs: logger.warning('denominators argument is ignored by <{}>.'.format(self))
        if 'rescale' in kwargs:
            del kwargs['rescale']
            logger.warning('rescale argument is ignored by <{}>.'.format(self))
        #end if

        return super(MulticlassLabelsClassifier, self).predict_proba(X_featurized, rescale=False, **kwargs)
    #end def

    def binarize_labels(self, Y_labels):
        classes_ = set(self.classes_)

        Y_labels_filtered = np.empty(Y_labels.shape[0], dtype=np.object)
        for i in range(Y_labels.shape[0]):
            labels = list(filter(classes_.__contains__, Y_labels[i]))
            Y_labels_filtered[i] = labels[0] if labels else '<none>'

        return self.label_encoder_.transform(Y_labels_filtered)
    #end def

    def binarize_dicts(self, Y_dicts, *, default=0.0):
        binarized = np.fill((Y_dicts.shape[0], len(self.classes_)), default, dtype=np.float)
        classes_map = dict((c, i) for i, c in enumerate(self.clsases_))

        for i in range(Y_dicts.shape[0]):
            d = Y_dicts[i]
            for k, p in d.items():
                try: binarized[i, classes_map[k]] = p
                except IndexError: pass
            #end for
        #end for

        return binarized
    #end def

    def unbinarize_labels(self, Y_proba, *, epsilon=0.0, to_dict=False):
        unbinarized = np.empty(Y_proba.shape[0], dtype=np.object)
        if len(Y_proba.shape) == 1:
            for i in range(Y_proba.shape[0]):
                if to_dict: unbinarized[i] = dict([(self.classes_[Y_proba[i]], Y_proba[i])])
                else: unbinarized[i] = [self.classes_[Y_proba[i]]]
            #end for
        else:
            Y_argmax = np.argmax(Y_proba, axis=1)
            for i in range(Y_proba.shape[0]):
                if to_dict: unbinarized[i] = dict((self.classes_[j], Y_proba[i, j]) for j in range(Y_proba.shape[1]) if Y_proba[i, j] > epsilon)
                else: unbinarized[i] = [self.classes_[Y_argmax[i]]]
            #end for
        #end if

        return unbinarized
    #end def

    @property
    def classes_(self):
        return self.label_encoder_.classes_
#end def

class MultiLabelsClassifier(LabelsClassifier):
    def __init__(self, n_jobs=1, include=[], exclude=[], **kwargs):
        super(MultiLabelsClassifier, self).__init__(**kwargs)

        self.exclude = set(exclude)
        self.include = set(include)
    #end def

    def _fit(self, X, Y_labels, *, binarize_args={}, fit_args={}, **kwargs):
        Y_labels_filtered = self._filter_labels(Y_labels)
        self.label_binarizer_ = MultiLabelBinarizer(sparse_output=False).fit(Y_labels_filtered)
        logger.info('{} labels found in training instances.'.format(len(self.classes_)))

        if not len(self.classes_): raise ValueError('There are no labels available for fitting model.')

        return super(MultiLabelsClassifier, self)._fit(X, Y_labels, binarize_args=binarize_args, fit_args=fit_args, **kwargs)
    #end def

    def binarize_labels(self, Y_labels):
        classes_ = set(self.classes_)

        return self.label_binarizer_.transform((filter(classes_.__contains__, labels) for labels in Y_labels))
    #end def

    def binarize_dicts(self, Y_dicts, *, default=0.0):
        binarized = np.fill((Y_dicts.shape[0], len(self.classes_)), default, dtype=np.float)
        classes_map = dict((c, i) for i, c in enumerate(self.clsases_))

        for i in range(Y_dicts.shape[0]):
            d = Y_dicts[i]
            for k, p in d.items():
                try: binarized[i, classes_map[k]] = p
                except IndexError: pass
            #end for
        #end for

        return binarized
    #end def

    def unbinarize_labels(self, Y_proba, *, epsilon=0.0, to_dict=False):
        assert len(self.classes_) == Y_proba.shape[1]
        unbinarized = np.empty(Y_proba.shape[0], dtype=np.object)
        for i in range(Y_proba.shape[0]):
            if to_dict: unbinarized[i] = dict((self.classes_[j], Y_proba[i, j]) for j in range(Y_proba.shape[1]) if Y_proba[i, j] > epsilon)
            else: unbinarized[i] = [self.classes_[j] for j in range(Y_proba.shape[1]) if Y_proba[i, j] > epsilon]
        #end for

        return unbinarized
    #end def

    @property
    def classes_(self):
        return self.label_binarizer_.classes_
    #end def
#end def
<commit_msg>Clean up some parts<commit_after>import logging

import numpy as np

from sklearn.preprocessing import label_binarize
from sklearn.preprocessing import LabelEncoder
from sklearn.preprocessing import MultiLabelBinarizer

from .base import BaseClassifier

__all__ = ['LabelsClassifier', 'BinaryLabelsClassifier', 'MultiLabelsClassifier', 'MulticlassLabelsClassifier']

logger = logging.getLogger(__name__)

class LabelsClassifier(BaseClassifier):
    def __init__(self, include=[], exclude=[], **kwargs):
        super(LabelsClassifier, self).__init__(**kwargs)

        self.exclude = set(exclude)
        self.include = set(include)
    #end def

    def _fit(self, X, Y_labels, *, binarize_args={}, fit_args={}, **kwargs):
        Y_binarized = self.binarize_labels(Y_labels, **binarize_args)
        return self.fit_binarized(X, Y_binarized, **fit_args)
    #end def

    def fit_binarized(self, X_featurized, Y_binarized, **kwargs): raise NotImplementedError('fit_binarized is not implemented.')

    def predict(self, X_featurized, binarized=True, **kwargs):
        Y_predict_binarized = super(LabelsClassifier, self).predict(X_featurized, **kwargs)
        if binarized: return Y_predict_binarized

        return self.unbinarize_labels(Y_predict_binarized)
    #end def

    def predict_and_proba(self, X_featurized, *, binarized=True, **kwargs):
        Y_proba, Y_predict_binarized = super(LabelsClassifier, self).predict_and_proba(X_featurized, **kwargs)
        if binarized: return Y_proba, Y_predict_binarized

        return Y_proba, self.unbinarize_labels(Y_predict_binarized)
    #end def

    def binarize_labels(self, Y_labels, **kwargs): raise NotImplementedError('binarize_labels is not implemented.')

    def binarize_dicts(self, Y_dicts, *, default=0.0): raise NotImplementedError('binarize_dicts is not implemented.')

    def unbinarize_labels(self, Y_binarized, *, epsilon=0.0): raise NotImplementedError('unbinarize_labels is not implemented.')

    def _filter_labels(self, Y_labels):
        if self.exclude or self.include:
            Y_labels_filtered = np.empty(Y_labels.shape, dtype=np.object)
            removed_labels = 0
            for i in range(Y_labels.shape[0]):
                Y_labels_filtered[i] = [l for l in Y_labels[i] if (l in self.include or not self.include) and l not in self.exclude]
                removed_labels += len(Y_labels[i]) - len(Y_labels_filtered[i])
            #end for
            logger.info('{} label-instances removed from the training data.'.format(removed_labels))
        else: Y_labels_filtered = Y_labels

        return Y_labels_filtered
    #end def
#end def

class BinaryLabelsClassifier(LabelsClassifier):
    def __init__(self, pos_label, *, not_pos_label=None, **kwargs):
        super(BinaryLabelsClassifier, self).__init__(**kwargs)

        self.pos_label = pos_label
        self.not_pos_label = 'not ' + pos_label if not_pos_label is None else not_pos_label
    #end def

    def predict_and_proba(self, X_featurized, *, binarized=True, **kwargs):
        Y_proba, Y_predict = super(BinaryLabelsClassifier, self).predict_and_proba(X_featurized, binarized=binarized, **kwargs)
        if Y_proba.ndim == 1:
            Y_proba_2d = np.zeros((Y_proba.shape[0], 2), dtype=Y_proba.dtype)
            Y_proba_2d[:, 0] = Y_proba
            Y_proba_2d[:, 1] = 1.0 - Y_proba
            Y_proba = Y_proba_2d
        #end if

        if binarized:
            Y_predict = Y_predict[:, 0]  # must be 0th one coz classes_ only has 1 thing

        return Y_proba, Y_predict
    #end def

    def binarize_labels(self, Y_labels):
        if Y_labels.shape[0] == 0:
            return np.zeros((0, 1))

        Y_labels_pos = [self.pos_label if self.pos_label in Y_labels[i] else self.not_pos_label for i in range(Y_labels.shape[0])]
        Y_binarized = label_binarize(Y_labels_pos, classes=[self.pos_label, self.not_pos_label]).reshape(Y_labels.shape[0])

        return Y_binarized
    #end def

    def binarize_dicts(self, Y_dicts, *, default=0.0):
        return np.array([Y_dicts[i].get(self.pos_label, default) for i in range(Y_dicts.shape[0])])
    #end def

    def unbinarize_labels(self, Y_binarized, *, epsilon=0.0):
        return np.array([[self.pos_label] if Y_binarized[i] > epsilon else [] for i in range(Y_binarized.shape[0])], dtype=np.object)
    #end def

    @property
    def classes_(self):
        return [self.pos_label, self.not_pos_label]
#end def

# TODO: use multilabel format and inherit multilabelclassifier
class MulticlassLabelsClassifier(LabelsClassifier):
    def __init__(self, **kwargs):
        super(MulticlassLabelsClassifier, self).__init__(**kwargs)
    #end def

    def _fit(self, X, Y_labels, *, binarize_args={}, fit_args={}, **kwargs):
        Y_labels_filtered = self._filter_labels(Y_labels)
        Y_labels_filtered = np.array([Y_labels_filtered[i][0] if Y_labels_filtered[i] else '<none>' for i in range(Y_labels_filtered.shape[0])])

        self.label_encoder_ = LabelEncoder().fit(Y_labels_filtered)
        logger.info('{} labels found in training instances.'.format(len(self.classes_)))

        if not len(self.classes_): raise ValueError('There are no labels available for fitting model.')

        return super(MulticlassLabelsClassifier, self)._fit(X, Y_labels, binarize_args=binarize_args, fit_args=fit_args, **kwargs)
    #end def

    def predict_and_proba(self, X_featurized, *, binarized=True, **kwargs):
        Y_proba = self.predict_proba(X_featurized, **kwargs)
        Y_predict_binarized = np.argmax(Y_proba, axis=1)

        if binarized: return Y_proba, Y_predict_binarized

        return Y_proba, self.unbinarize_labels(Y_predict_binarized)
    #end def

    def predict_proba(self, X_featurized, **kwargs):
        if 'thresholds' in kwargs: logger.warning('thresholds argument is ignored by <{}>.'.format(self))
        if 'denominators' in kwargs: logger.warning('denominators argument is ignored by <{}>.'.format(self))
        if 'rescale' in kwargs:
            del kwargs['rescale']
            logger.warning('rescale argument is ignored by <{}>.'.format(self))
        #end if

        return super(MulticlassLabelsClassifier, self).predict_proba(X_featurized, rescale=False, **kwargs)
    #end def

    def binarize_labels(self, Y_labels):
        classes_ = set(self.classes_)

        Y_labels_filtered = np.empty(Y_labels.shape[0], dtype=np.object)
        for i in range(Y_labels.shape[0]):
            labels = list(filter(classes_.__contains__, Y_labels[i]))
            Y_labels_filtered[i] = labels[0] if labels else '<none>'

        return self.label_encoder_.transform(Y_labels_filtered)
    #end def

    def binarize_dicts(self, Y_dicts, *, default=0.0):
        binarized = np.fill((Y_dicts.shape[0], len(self.classes_)), default, dtype=np.float)
        classes_map = dict((c, i) for i, c in enumerate(self.clsases_))

        for i in range(Y_dicts.shape[0]):
            d = Y_dicts[i]
            for k, p in d.items():
                try: binarized[i, classes_map[k]] = p
                except IndexError: pass
            #end for
        #end for

        return binarized
    #end def

    def unbinarize_labels(self, Y_proba, *, epsilon=0.0, to_dict=False):
        unbinarized = np.empty(Y_proba.shape[0], dtype=np.object)
        if len(Y_proba.shape) == 1:
            for i in range(Y_proba.shape[0]):
                if to_dict: unbinarized[i] = dict([(self.classes_[Y_proba[i]], Y_proba[i])])
                else: unbinarized[i] = [self.classes_[Y_proba[i]]]
            #end for
        else:
            Y_argmax = np.argmax(Y_proba, axis=1)
            for i in range(Y_proba.shape[0]):
                if to_dict: unbinarized[i] = dict((self.classes_[j], Y_proba[i, j]) for j in range(Y_proba.shape[1]) if Y_proba[i, j] > epsilon)
                else: unbinarized[i] = [self.classes_[Y_argmax[i]]]
            #end for
        #end if

        return unbinarized
    #end def

    @property
    def classes_(self):
        return self.label_encoder_.classes_
#end def

class MultiLabelsClassifier(LabelsClassifier):
    def __init__(self, n_jobs=1, include=[], exclude=[], **kwargs):
        super(MultiLabelsClassifier, self).__init__(**kwargs)

        self.exclude = set(exclude)
        self.include = set(include)
    #end def

    def _fit(self, X, Y_labels, *, binarize_args={}, fit_args={}, **kwargs):
        Y_labels_filtered = self._filter_labels(Y_labels)
        self.label_binarizer_ = MultiLabelBinarizer(sparse_output=False).fit(Y_labels_filtered)
        logger.info('{} labels found in training instances.'.format(len(self.classes_)))

        if not len(self.classes_): raise ValueError('There are no labels available for fitting model.')

        return super(MultiLabelsClassifier, self)._fit(X, Y_labels, binarize_args=binarize_args, fit_args=fit_args, **kwargs)
    #end def

    def binarize_labels(self, Y_labels):
        classes_ = set(self.classes_)

        return self.label_binarizer_.transform((filter(classes_.__contains__, labels) for labels in Y_labels))
    #end def

    def binarize_dicts(self, Y_dicts, *, default=0.0):
        binarized = np.fill((Y_dicts.shape[0], len(self.classes_)), default, dtype=np.float)
        classes_map = dict((c, i) for i, c in enumerate(self.clsases_))

        for i in range(Y_dicts.shape[0]):
            d = Y_dicts[i]
            for k, p in d.items():
                try: binarized[i, classes_map[k]] = p
                except IndexError: pass
            #end for
        #end for

        return binarized
    #end def

    def unbinarize_labels(self, Y_proba, *, epsilon=0.0, to_dict=False):
        assert len(self.classes_) == Y_proba.shape[1]
        unbinarized = np.empty(Y_proba.shape[0], dtype=np.object)
        for i in range(Y_proba.shape[0]):
            if to_dict: unbinarized[i] = dict((self.classes_[j], Y_proba[i, j]) for j in range(Y_proba.shape[1]) if Y_proba[i, j] > epsilon)
            else: unbinarized[i] = [self.classes_[j] for j in range(Y_proba.shape[1]) if Y_proba[i, j] > epsilon]
        #end for

        return unbinarized
    #end def

    @property
    def classes_(self):
        return self.label_binarizer_.classes_
    #end def
#end def
<|endoftext|>