<commit_before>import json
import os
import tempfile
import uuid

from django.contrib.contenttypes.fields import GenericForeignKey, GenericRelation
from django.contrib.contenttypes.models import ContentType
from django.core.files.storage import FileSystemStorage
from django.core.serializers.json import DjangoJSONEncoder
from django.db import models
from django.db.models.fields.files import ImageFieldFile
from django.utils.translation import gettext_lazy as _

try:
    from PIL import Image
except ImportError:
    Image = None

class Foo(models.Model):
    a = models.CharField(max_length=10)
    d = models.DecimalField(max_digits=5, decimal_places=3)

def get_foo():
    return Foo.objects.get(id=1).pk

class Bar(models.Model):
    b = models.CharField(max_length=10)
    a = models.ForeignKey(Foo, models.CASCADE, default=get_foo, related_name="bars")

class Whiz(models.Model):
    CHOICES = (
        (
            "Group 1",
            (
                (1, "First"),
                (2, "Second"),
            ),
        ),
        (
            "Group 2",
            (
                (3, "Third"),
                (4, "Fourth"),
            ),
        ),
        (0, "Other"),
        (5, _("translated")),
    )
    c = models.IntegerField(choices=CHOICES, null=True)

class WhizDelayed(models.Model):
    c = models.IntegerField(choices=(), null=True)

# Contrived way of adding choices later.
WhizDelayed._meta.get_field("c").choices = Whiz.CHOICES

class WhizIter(models.Model):
    c = models.IntegerField(choices=iter(Whiz.CHOICES), null=True)

class WhizIterEmpty(models.Model):
    c = models.CharField(choices=iter(()), blank=True, max_length=1)

class Choiceful(models.Model):
    no_choices = models.IntegerField(null=True)
    empty_choices = models.IntegerField(choices=(), null=True)
    with_choices = models.IntegerField(choices=[(1, "A")], null=True)
    empty_choices_bool = models.BooleanField(choices=())
    empty_choices_text = models.TextField(choices=())

class BigD(models.Model):
    d = models.DecimalField(max_digits=32, decimal_places=30)

class FloatModel(models.Model):
    size = models.FloatField()

class BigS(models.Model):
    s = models.SlugField(max_length=255)

class UnicodeSlugField(models.Model):
    s = models.SlugField(max_length=255, allow_unicode=True)

class AutoModel(models.Model):
    value = models.AutoField(primary_key=True)

class BigAutoModel(models.Model):
    value = models.BigAutoField(primary_key=True)

class SmallAutoModel(models.Model):
    value = models.SmallAutoField(primary_key=True)

class SmallIntegerModel(models.Model):
    value = models.SmallIntegerField()

class IntegerModel(models.Model):
    value = models.IntegerField()

class BigIntegerModel(models.Model):
    value = models.BigIntegerField()
    null_value = models.BigIntegerField(null=True, blank=True)

class PositiveBigIntegerModel(models.Model):
    value = models.PositiveBigIntegerField()

class PositiveSmallIntegerModel(models.Model):
    value = models.PositiveSmallIntegerField()

class PositiveIntegerModel(models.Model):
    value = models.PositiveIntegerField()

class Post(models.Model):
    title = models.CharField(max_length=100)
    body = models.TextField()

class NullBooleanModel(models.Model):
    nbfield = models.BooleanField(null=True, blank=True)

class BooleanModel(models.Model):
    bfield = models.BooleanField()

class DateTimeModel(models.Model):
    d = models.DateField()
    dt = models.DateTimeField()
    t = models.TimeField()

class DurationModel(models.Model):
    field = models.DurationField()

class NullDurationModel(models.Model):
    field = models.DurationField(null=True)

class PrimaryKeyCharModel(models.Model):
    string = models.CharField(max_length=10, primary_key=True)

class FksToBooleans(models.Model):
    """Model with FKs to models with {Null,}BooleanField's, #15040"""

    bf = models.ForeignKey(BooleanModel, models.CASCADE)
    nbf = models.ForeignKey(NullBooleanModel, models.CASCADE)

class FkToChar(models.Model):
    """Model with FK to a model with a CharField primary key, #19299"""

    out = models.ForeignKey(PrimaryKeyCharModel, models.CASCADE)

class RenamedField(models.Model):
    modelname = models.IntegerField(name="fieldname", choices=((1, "One"),))

class VerboseNameField(models.Model):
    id = models.AutoField("verbose pk", primary_key=True)
    field1 = models.BigIntegerField("verbose field1")
    field2 = models.BooleanField("verbose field2", default=False)
    field3 = models.CharField("verbose field3", max_length=10)
    field4 = models.DateField("verbose field4")
    field5 = models.DateTimeField("verbose field5")
    field6 = models.DecimalField("verbose field6", max_digits=6, decimal_places=1)
    field7 = models.EmailField("verbose field7")
    field8 = models.FileField("verbose field8", upload_to="unused")
    field9 = models.FilePathField("verbose field9")
    field10 = models.FloatField("verbose field10")
    # Don't want to depend on Pillow in this test
    # field_image = models.ImageField("verbose field")
    field11 = models.IntegerField("verbose field11")
    field12 = models.GenericIPAddressField("verbose field12", protocol="ipv4")
    field13 = models.PositiveIntegerField("verbose field13")
    field14 = models.PositiveSmallIntegerField("verbose field14")
    field15 = models.SlugField("verbose field15")
    field16 = models.SmallIntegerField("verbose field16")
    field17 = models.TextField("verbose field17")
    field18 = models.TimeField("verbose field18")
    field19 = models.URLField("verbose field19")
    field20 = models.UUIDField("verbose field20")
    field21 = models.DurationField("verbose field21")

class GenericIPAddress(models.Model):
    ip = models.GenericIPAddressField(null=True, protocol="ipv4")

###############################################################################
# These models aren't used in any test, just here to ensure they validate
# successfully.

# See ticket #16570.
class DecimalLessThanOne(models.Model):
    d = models.DecimalField(max_digits=3, decimal_places=3)

# See ticket #18389.
class FieldClassAttributeModel(models.Model):
    field_class = models.CharField

###############################################################################

class DataModel(models.Model):
    short_data = models.BinaryField(max_length=10, default=b"\x08")
    data = models.BinaryField()

###############################################################################
# FileField

class Document(models.Model):
    myfile = models.FileField(upload_to="unused", unique=True)

###############################################################################
# ImageField

# If Pillow available, do these tests.
if Image:

    class TestImageFieldFile(ImageFieldFile):
        """
        Custom Field File class that records whether or not the underlying file
        was opened.
        """

        def __init__(self, *args, **kwargs):
            self.was_opened = False
            super().__init__(*args, **kwargs)

        def open(self):
            self.was_opened = True
            super().open()

    class TestImageField(models.ImageField):
        attr_class = TestImageFieldFile

    # Set up a temp directory for file storage.
    temp_storage_dir = tempfile.mkdtemp()
    temp_storage = FileSystemStorage(temp_storage_dir)
    temp_upload_to_dir = os.path.join(temp_storage.location, "tests")

    class Person(models.Model):
        """
        Model that defines an ImageField with no dimension fields.
        """

        name = models.CharField(max_length=50)
        mugshot = TestImageField(storage=temp_storage, upload_to="tests")

    class AbstractPersonWithHeight(models.Model):
        """
        Abstract model that defines an ImageField with only one dimension field
        to make sure the dimension update is correctly run on concrete subclass
        instance post-initialization.
        """

        mugshot = TestImageField(
            storage=temp_storage, upload_to="tests", height_field="mugshot_height"
        )
        mugshot_height = models.PositiveSmallIntegerField()

        class Meta:
            abstract = True

    class PersonWithHeight(AbstractPersonWithHeight):
        """
        Concrete model that subclass an abstract one with only on dimension
        field.
        """

        name = models.CharField(max_length=50)

    class PersonWithHeightAndWidth(models.Model):
        """
        Model that defines height and width fields after the ImageField.
        """

        name = models.CharField(max_length=50)
        mugshot = TestImageField(
            storage=temp_storage,
            upload_to="tests",
            height_field="mugshot_height",
            width_field="mugshot_width",
        )
        mugshot_height = models.PositiveSmallIntegerField()
        mugshot_width = models.PositiveSmallIntegerField()

    class PersonDimensionsFirst(models.Model):
        """
        Model that defines height and width fields before the ImageField.
        """

        name = models.CharField(max_length=50)
        mugshot_height = models.PositiveSmallIntegerField()
        mugshot_width = models.PositiveSmallIntegerField()
        mugshot = TestImageField(
            storage=temp_storage,
            upload_to="tests",
            height_field="mugshot_height",
            width_field="mugshot_width",
        )

    class PersonTwoImages(models.Model):
        """
        Model that:
        * Defines two ImageFields
        * Defines the height/width fields before the ImageFields
        * Has a nullable ImageField
        """

        name = models.CharField(max_length=50)
        mugshot_height = models.PositiveSmallIntegerField()
        mugshot_width = models.PositiveSmallIntegerField()
        mugshot = TestImageField(
            storage=temp_storage,
            upload_to="tests",
            height_field="mugshot_height",
            width_field="mugshot_width",
        )
        headshot_height = models.PositiveSmallIntegerField(blank=True, null=True)
        headshot_width = models.PositiveSmallIntegerField(blank=True, null=True)
        headshot = TestImageField(
            blank=True,
            null=True,
            storage=temp_storage,
            upload_to="tests",
            height_field="headshot_height",
            width_field="headshot_width",
        )

class CustomJSONDecoder(json.JSONDecoder):
    def __init__(self, object_hook=None, *args, **kwargs):
        return super().__init__(object_hook=self.as_uuid, *args, **kwargs)

    def as_uuid(self, dct):
        if "uuid" in dct:
            dct["uuid"] = uuid.UUID(dct["uuid"])
        return dct

class JSONModel(models.Model):
    value = models.JSONField()

    class Meta:
        required_db_features = {"supports_json_field"}

class NullableJSONModel(models.Model):
    value = models.JSONField(blank=True, null=True)
    value_custom = models.JSONField(
        encoder=DjangoJSONEncoder,
        decoder=CustomJSONDecoder,
        null=True,
    )

    class Meta:
        required_db_features = {"supports_json_field"}

class RelatedJSONModel(models.Model):
    value = models.JSONField()
    json_model = models.ForeignKey(NullableJSONModel, models.CASCADE)

    class Meta:
        required_db_features = {"supports_json_field"}

class AllFieldsModel(models.Model):
    big_integer = models.BigIntegerField()
    binary = models.BinaryField()
    boolean = models.BooleanField(default=False)
    char = models.CharField(max_length=10)
    date = models.DateField()
    datetime = models.DateTimeField()
    decimal = models.DecimalField(decimal_places=2, max_digits=2)
    duration = models.DurationField()
    email = models.EmailField()
    file_path = models.FilePathField()
    floatf = models.FloatField()
    integer = models.IntegerField()
    generic_ip = models.GenericIPAddressField()
    positive_integer = models.PositiveIntegerField()
    positive_small_integer = models.PositiveSmallIntegerField()
    slug = models.SlugField()
    small_integer = models.SmallIntegerField()
    text = models.TextField()
    time = models.TimeField()
    url = models.URLField()
    uuid = models.UUIDField()

    fo = models.ForeignObject(
        "self",
        on_delete=models.CASCADE,
        from_fields=["positive_integer"],
        to_fields=["id"],
        related_name="reverse",
    )
    fk = models.ForeignKey("self", models.CASCADE, related_name="reverse2")
    m2m = models.ManyToManyField("self")
    oto = models.OneToOneField("self", models.CASCADE)

    object_id = models.PositiveIntegerField()
    content_type = models.ForeignKey(ContentType, models.CASCADE)
    gfk = GenericForeignKey()
    gr = GenericRelation(DataModel)

class ManyToMany(models.Model):
    m2m = models.ManyToManyField("self")

###############################################################################

class UUIDModel(models.Model):
    field = models.UUIDField()

class NullableUUIDModel(models.Model):
    field = models.UUIDField(blank=True, null=True)

class PrimaryKeyUUIDModel(models.Model):
    id = models.UUIDField(primary_key=True, default=uuid.uuid4)

class RelatedToUUIDModel(models.Model):
    uuid_fk = models.ForeignKey("PrimaryKeyUUIDModel", models.CASCADE)

class UUIDChild(PrimaryKeyUUIDModel):
    pass

class UUIDGrandchild(UUIDChild):
    pass
<commit_msg>Removed unused path variable in model_fields tests.<commit_after>import json
import tempfile
import uuid

from django.contrib.contenttypes.fields import GenericForeignKey, GenericRelation
from django.contrib.contenttypes.models import ContentType
from django.core.files.storage import FileSystemStorage
from django.core.serializers.json import DjangoJSONEncoder
from django.db import models
from django.db.models.fields.files import ImageFieldFile
from django.utils.translation import gettext_lazy as _

try:
    from PIL import Image
except ImportError:
    Image = None

class Foo(models.Model):
    a = models.CharField(max_length=10)
    d = models.DecimalField(max_digits=5, decimal_places=3)

def get_foo():
    return Foo.objects.get(id=1).pk

class Bar(models.Model):
    b = models.CharField(max_length=10)
    a = models.ForeignKey(Foo, models.CASCADE, default=get_foo, related_name="bars")

class Whiz(models.Model):
    CHOICES = (
        (
            "Group 1",
            (
                (1, "First"),
                (2, "Second"),
            ),
        ),
        (
            "Group 2",
            (
                (3, "Third"),
                (4, "Fourth"),
            ),
        ),
        (0, "Other"),
        (5, _("translated")),
    )
    c = models.IntegerField(choices=CHOICES, null=True)

class WhizDelayed(models.Model):
    c = models.IntegerField(choices=(), null=True)

# Contrived way of adding choices later.
WhizDelayed._meta.get_field("c").choices = Whiz.CHOICES

class WhizIter(models.Model):
    c = models.IntegerField(choices=iter(Whiz.CHOICES), null=True)

class WhizIterEmpty(models.Model):
    c = models.CharField(choices=iter(()), blank=True, max_length=1)

class Choiceful(models.Model):
    no_choices = models.IntegerField(null=True)
    empty_choices = models.IntegerField(choices=(), null=True)
    with_choices = models.IntegerField(choices=[(1, "A")], null=True)
    empty_choices_bool = models.BooleanField(choices=())
    empty_choices_text = models.TextField(choices=())

class BigD(models.Model):
    d = models.DecimalField(max_digits=32, decimal_places=30)

class FloatModel(models.Model):
    size = models.FloatField()

class BigS(models.Model):
    s = models.SlugField(max_length=255)

class UnicodeSlugField(models.Model):
    s = models.SlugField(max_length=255, allow_unicode=True)

class AutoModel(models.Model):
    value = models.AutoField(primary_key=True)

class BigAutoModel(models.Model):
    value = models.BigAutoField(primary_key=True)

class SmallAutoModel(models.Model):
    value = models.SmallAutoField(primary_key=True)

class SmallIntegerModel(models.Model):
    value = models.SmallIntegerField()

class IntegerModel(models.Model):
    value = models.IntegerField()

class BigIntegerModel(models.Model):
    value = models.BigIntegerField()
    null_value = models.BigIntegerField(null=True, blank=True)

class PositiveBigIntegerModel(models.Model):
    value = models.PositiveBigIntegerField()

class PositiveSmallIntegerModel(models.Model):
    value = models.PositiveSmallIntegerField()

class PositiveIntegerModel(models.Model):
    value = models.PositiveIntegerField()

class Post(models.Model):
    title = models.CharField(max_length=100)
    body = models.TextField()

class NullBooleanModel(models.Model):
    nbfield = models.BooleanField(null=True, blank=True)

class BooleanModel(models.Model):
    bfield = models.BooleanField()

class DateTimeModel(models.Model):
    d = models.DateField()
    dt = models.DateTimeField()
    t = models.TimeField()

class DurationModel(models.Model):
    field = models.DurationField()

class NullDurationModel(models.Model):
    field = models.DurationField(null=True)

class PrimaryKeyCharModel(models.Model):
    string = models.CharField(max_length=10, primary_key=True)

class FksToBooleans(models.Model):
    """Model with FKs to models with {Null,}BooleanField's, #15040"""

    bf = models.ForeignKey(BooleanModel, models.CASCADE)
    nbf = models.ForeignKey(NullBooleanModel, models.CASCADE)

class FkToChar(models.Model):
    """Model with FK to a model with a CharField primary key, #19299"""

    out = models.ForeignKey(PrimaryKeyCharModel, models.CASCADE)

class RenamedField(models.Model):
    modelname = models.IntegerField(name="fieldname", choices=((1, "One"),))

class VerboseNameField(models.Model):
    id = models.AutoField("verbose pk", primary_key=True)
    field1 = models.BigIntegerField("verbose field1")
    field2 = models.BooleanField("verbose field2", default=False)
    field3 = models.CharField("verbose field3", max_length=10)
    field4 = models.DateField("verbose field4")
    field5 = models.DateTimeField("verbose field5")
    field6 = models.DecimalField("verbose field6", max_digits=6, decimal_places=1)
    field7 = models.EmailField("verbose field7")
    field8 = models.FileField("verbose field8", upload_to="unused")
    field9 = models.FilePathField("verbose field9")
    field10 = models.FloatField("verbose field10")
    # Don't want to depend on Pillow in this test
    # field_image = models.ImageField("verbose field")
    field11 = models.IntegerField("verbose field11")
    field12 = models.GenericIPAddressField("verbose field12", protocol="ipv4")
    field13 = models.PositiveIntegerField("verbose field13")
    field14 = models.PositiveSmallIntegerField("verbose field14")
    field15 = models.SlugField("verbose field15")
    field16 = models.SmallIntegerField("verbose field16")
    field17 = models.TextField("verbose field17")
    field18 = models.TimeField("verbose field18")
    field19 = models.URLField("verbose field19")
    field20 = models.UUIDField("verbose field20")
    field21 = models.DurationField("verbose field21")

class GenericIPAddress(models.Model):
    ip = models.GenericIPAddressField(null=True, protocol="ipv4")

###############################################################################
# These models aren't used in any test, just here to ensure they validate
# successfully.

# See ticket #16570.
class DecimalLessThanOne(models.Model):
    d = models.DecimalField(max_digits=3, decimal_places=3)

# See ticket #18389.
class FieldClassAttributeModel(models.Model):
    field_class = models.CharField

###############################################################################

class DataModel(models.Model):
    short_data = models.BinaryField(max_length=10, default=b"\x08")
    data = models.BinaryField()

###############################################################################
# FileField

class Document(models.Model):
    myfile = models.FileField(upload_to="unused", unique=True)

###############################################################################
# ImageField

# If Pillow available, do these tests.
if Image:

    class TestImageFieldFile(ImageFieldFile):
        """
        Custom Field File class that records whether or not the underlying file
        was opened.
        """

        def __init__(self, *args, **kwargs):
            self.was_opened = False
            super().__init__(*args, **kwargs)

        def open(self):
            self.was_opened = True
            super().open()

    class TestImageField(models.ImageField):
        attr_class = TestImageFieldFile

    # Set up a temp directory for file storage.
    temp_storage_dir = tempfile.mkdtemp()
    temp_storage = FileSystemStorage(temp_storage_dir)

    class Person(models.Model):
        """
        Model that defines an ImageField with no dimension fields.
        """

        name = models.CharField(max_length=50)
        mugshot = TestImageField(storage=temp_storage, upload_to="tests")

    class AbstractPersonWithHeight(models.Model):
        """
        Abstract model that defines an ImageField with only one dimension field
        to make sure the dimension update is correctly run on concrete subclass
        instance post-initialization.
        """

        mugshot = TestImageField(
            storage=temp_storage, upload_to="tests", height_field="mugshot_height"
        )
        mugshot_height = models.PositiveSmallIntegerField()

        class Meta:
            abstract = True

    class PersonWithHeight(AbstractPersonWithHeight):
        """
        Concrete model that subclass an abstract one with only on dimension
        field.
        """

        name = models.CharField(max_length=50)

    class PersonWithHeightAndWidth(models.Model):
        """
        Model that defines height and width fields after the ImageField.
        """

        name = models.CharField(max_length=50)
        mugshot = TestImageField(
            storage=temp_storage,
            upload_to="tests",
            height_field="mugshot_height",
            width_field="mugshot_width",
        )
        mugshot_height = models.PositiveSmallIntegerField()
        mugshot_width = models.PositiveSmallIntegerField()

    class PersonDimensionsFirst(models.Model):
        """
        Model that defines height and width fields before the ImageField.
        """

        name = models.CharField(max_length=50)
        mugshot_height = models.PositiveSmallIntegerField()
        mugshot_width = models.PositiveSmallIntegerField()
        mugshot = TestImageField(
            storage=temp_storage,
            upload_to="tests",
            height_field="mugshot_height",
            width_field="mugshot_width",
        )

    class PersonTwoImages(models.Model):
        """
        Model that:
        * Defines two ImageFields
        * Defines the height/width fields before the ImageFields
        * Has a nullable ImageField
        """

        name = models.CharField(max_length=50)
        mugshot_height = models.PositiveSmallIntegerField()
        mugshot_width = models.PositiveSmallIntegerField()
        mugshot = TestImageField(
            storage=temp_storage,
            upload_to="tests",
            height_field="mugshot_height",
            width_field="mugshot_width",
        )
        headshot_height = models.PositiveSmallIntegerField(blank=True, null=True)
        headshot_width = models.PositiveSmallIntegerField(blank=True, null=True)
        headshot = TestImageField(
            blank=True,
            null=True,
            storage=temp_storage,
            upload_to="tests",
            height_field="headshot_height",
            width_field="headshot_width",
        )

class CustomJSONDecoder(json.JSONDecoder):
    def __init__(self, object_hook=None, *args, **kwargs):
        return super().__init__(object_hook=self.as_uuid, *args, **kwargs)

    def as_uuid(self, dct):
        if "uuid" in dct:
            dct["uuid"] = uuid.UUID(dct["uuid"])
        return dct

class JSONModel(models.Model):
    value = models.JSONField()

    class Meta:
        required_db_features = {"supports_json_field"}

class NullableJSONModel(models.Model):
    value = models.JSONField(blank=True, null=True)
    value_custom = models.JSONField(
        encoder=DjangoJSONEncoder,
        decoder=CustomJSONDecoder,
        null=True,
    )

    class Meta:
        required_db_features = {"supports_json_field"}

class RelatedJSONModel(models.Model):
    value = models.JSONField()
    json_model = models.ForeignKey(NullableJSONModel, models.CASCADE)

    class Meta:
        required_db_features = {"supports_json_field"}

class AllFieldsModel(models.Model):
    big_integer = models.BigIntegerField()
    binary = models.BinaryField()
    boolean = models.BooleanField(default=False)
    char = models.CharField(max_length=10)
    date = models.DateField()
    datetime = models.DateTimeField()
    decimal = models.DecimalField(decimal_places=2, max_digits=2)
    duration = models.DurationField()
    email = models.EmailField()
    file_path = models.FilePathField()
    floatf = models.FloatField()
    integer = models.IntegerField()
    generic_ip = models.GenericIPAddressField()
    positive_integer = models.PositiveIntegerField()
    positive_small_integer = models.PositiveSmallIntegerField()
    slug = models.SlugField()
    small_integer = models.SmallIntegerField()
    text = models.TextField()
    time = models.TimeField()
    url = models.URLField()
    uuid = models.UUIDField()

    fo = models.ForeignObject(
        "self",
        on_delete=models.CASCADE,
        from_fields=["positive_integer"],
        to_fields=["id"],
        related_name="reverse",
    )
    fk = models.ForeignKey("self", models.CASCADE, related_name="reverse2")
    m2m = models.ManyToManyField("self")
    oto = models.OneToOneField("self", models.CASCADE)

    object_id = models.PositiveIntegerField()
    content_type = models.ForeignKey(ContentType, models.CASCADE)
    gfk = GenericForeignKey()
    gr = GenericRelation(DataModel)

class ManyToMany(models.Model):
    m2m = models.ManyToManyField("self")

###############################################################################

class UUIDModel(models.Model):
    field = models.UUIDField()

class NullableUUIDModel(models.Model):
    field = models.UUIDField(blank=True, null=True)

class PrimaryKeyUUIDModel(models.Model):
    id = models.UUIDField(primary_key=True, default=uuid.uuid4)

class RelatedToUUIDModel(models.Model):
    uuid_fk = models.ForeignKey("PrimaryKeyUUIDModel", models.CASCADE)

class UUIDChild(PrimaryKeyUUIDModel):
    pass

class UUIDGrandchild(UUIDChild):
    pass
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
#
#  This file is part of Sequana software
#
#  Copyright (c) 2016 - Sequana Development Team
#
#  File author(s):
#      Thomas Cokelaer <thomas.cokelaer@pasteur.fr>
#      Dimitri Desvillechabrol <dimitri.desvillechabrol@pasteur.fr>,
#          <d.desvillechabrol@gmail.com>
#
#  Distributed under the terms of the 3-clause BSD license.
#  The full license is in the LICENSE file, distributed with this software.
#
#  website: https://github.com/sequana/sequana
#  documentation: http://sequana.readthedocs.io
#
##############################################################################
""".. rubric:: Standalone application dedicated to coverage"""
import os
import shutil
import glob
import sys
from optparse import OptionParser
import argparse
from argparse import RawTextHelpFormatter

from sequana import bedtools
from sequana.reporting import report_mapping
from sequana.reporting import report_chromosome
from sequana.reporting import report_main
from sequana import logger

from easydev import shellcmd

from pylab import show, figure, savefig

from easydev.console import purple

# http://stackoverflow.com/questions/18462610/argumentparser-epilog-and-description-formatting-in-conjunction-with-argumentdef
class CustomFormatter(argparse.ArgumentDefaultsHelpFormatter, argparse.RawDescriptionHelpFormatter):
    pass

class Options(argparse.ArgumentParser):
    def  __init__(self, prog="sequana_coverage"):
        usage = purple("""\nWelcome to SEQUANA -- Coverage standalone

    Extract and plot coverage of one or more chromosomes/contigs in a BED or BAM
    file. In addition, running median used in conjunction with double thresholds
    extract regions of interests (low or high coverage). A reference may be
    provided to plot the coverage versus GC content.

    The input file should be one of the following:

    - a BED file that is a tabulated file at least 3 columns.
      The first column being the reference, the second is the position
      and the third column contains the coverage itself.
    - or a BAM file that is converted automatically
      into a BED file using the following command:

        samtools depth -aa input.bam > output.bed

    If the reference is provided, an additional plot showing the coverage versus
    GC content is also shown.

    Here are some examples

        sequana_coverage --input file.bed --window-median 1001
        sequana_coverage --input file.bam --window-median 1001 -r <REFERENCE.fa>

    An other interesting option is to provide a BED file with 4 columns. The
    fourth column being another coverage data created with a filter. One can
    create such a file only from the BAM file using samtools as follows:

        samtools view -q 35  -o data.filtered.bam input.fa.sorted.bam
        samtools depth data.filtered.bam input.fa.sorted.bam  -aa > test.bed
        sequana_coverage --input test.bed --show-html

    Note that the first file is the filtered one, and the second file is the
    unfiltered one.

        """)

        epilog = purple("""
----

AUTHORS: Thomas Cokelaer, Dimitri Desvillechabrol
Documentation: http://sequana.readthedocs.io
Issues: http://github.com/sequana/sequana
        """)

        description = """DESCRIPTION:
        """

        super(Options, self).__init__(usage=usage, prog=prog,
                description=description, epilog=epilog,
                formatter_class=CustomFormatter)

        # options to fill the config file
        group = self.add_argument_group("Required argument")
        group.add_argument("-i", "--input", dest="input", type=str,
            help=("Input file in BED or BAM format. If a BAM file is "
                 "provided, it will be converted locally to a BED file "
                 "using genomecov, which must be installed."))

        group = self.add_argument_group("Optional biological arguments")
        group.add_argument('-c', "--chromosome", dest="chromosome", type=int,
            default=1,
            help=(  "Chromosome number (if only one, no need to use: the first"
                    " and only chromosome is chosen automatically). Default is"
                    " first chromosome found in the BED file. You may want to"
                    " analyse all chromosomes at the same time. If so, set this"
                    " parameter to -1"))
        group.add_argument('-o', "--circular", dest="circular",
            default=False, action="store_true",
            help="""If the DNA of the organism is circular (typically
            viruses or bacteria), set to True""")

        group = self.add_argument_group("General")
        group.add_argument("--output-directory", dest="output_directory",
            default="report", help="name of the output (report) directory.")
        group.add_argument('--show', dest="show", default=False,
            action='store_true', help="""Show the pictures (matplotlib)""")
        group.add_argument('--show-html', dest="show_html", default=False,
            action='store_true',
            help="""When report is created, you can open
            the main page automatically with this option (default is False)""")
        group.add_argument("-q", "--quiet", dest="verbose",
            default=True, action="store_false")
        group.add_argument('--no-report', dest="create_report",
            default=True, action='store_false',
            help="""Do not create any HTML report""")

        group = self.add_argument_group('Annotation')
        group.add_argument("-b", "--genbank", dest="genbank",
            type=str, default=None, help='a valida genbank annotation')

        # Group related to GC content
        group = self.add_argument_group("GC content related")
        group.add_argument('-r', "--reference", dest="reference", type=str,
            default=None,
            help="""If available, you can provide a reference (ENA/NCBI). It
                 must have the same length as the one used to create the
                 BAM or BED file. If provided, it is used to create the
                 coverage versus GC content image""")
        group.add_argument("-g", "--window-gc", dest="w_gc", type=int,
            help="""Length of the running window to compute the GC content""", default=200)
        group.add_argument('-n', "--nlevels", dest="levels", type=int,
            default=3, help="""Number of levels in the contour""")

        #group running median
        group = self.add_argument_group("Running Median related")
        group.add_argument("-w", "--window-median", dest="w_median", type=int,
            help="""Length of the running median window (default 4001,
                 recommended for viruses).  For long genome, 20001
                 or 30001 is recommended but larger windows may be
                 useful in the presence of long deleted regions.""",
            default=4001)

        group.add_argument("-k", "--mixture-models", dest="k", type=int,
            help="""Number of mixture models to use (default 2, although if sequencing
        depth is below 8, k is set to 1 automatically). To ignore that behaviour
        set k to the required value""",
            default=None)

        group.add_argument("-L", "--low-threshold", dest="low_threshold",
            default=None, type=float,
            help=("lower threshold (zscore) of the confidence interval. "
                "Overwrite value given by --threshold/-T"))
        group.add_argument("-H", "--high-threshold", dest="high_threshold",
            default=None, type=float,
            help=("higher threshold (zscore) of the confidence interval. "
                "Overwrite value given by --threshold/-T"))
        group.add_argument("-T", "--threshold", dest="threshold",
            default=4, type=float,
            help="""set lower and higher thresholds of the confidence interval.""")

        group = self.add_argument_group("Download reference")
        group.add_argument("--download-reference", dest="accession",
            default=None, type=str)
        group.add_argument("--download-genbank", dest="download_genbank",
            default=None, type=str)
        group.add_argument("--database", dest="database",
            default="ENA", type=str,
            choices=["ENA", "EUtils"],
            help="Download the reference from one of these database (default ENA)")

def main(args=None):

    if args is None:
        args = sys.argv[:]

    user_options = Options(prog="sequana")

    # If --help or no options provided, show the help
    if len(args) == 1:
        user_options.parse_args(["prog", "--help"])
    else:
        options = user_options.parse_args(args[1:])

    if options.accession:
        logger.info("Download accession %s from %s\n" %
            (options.accession, options.database))

        from bioservices.apps import download_fasta as df
        df.download_fasta(options.accession, method=options.database)
        return

    if options.download_genbank:
        from sequana.snpeff import download_fasta_and_genbank
        download_fasta_and_genbank(options.download_genbank,
                                   options.download_genbank,
                                   genbank=True, fasta=False)
        return

    # We put the import here to make the --help faster
    from sequana import GenomeCov

    if options.genbank:
        assert os.path.exists(options.genbank), \
            "%s does not exists" % options.genbank

    if options.verbose:
        logger.info("Reading %s" % options.input)

    if options.input.endswith(".bam"):
        bedfile = options.input.replace(".bam", ".bed")
        if options.verbose:
            print("Converting BAM into BED file")
        shellcmd("bedtools genomecov -d -ibam %s > %s" % (options.input, bedfile))
    elif options.input.endswith(".bed"):
        bedfile = options.input
    else:
        raise ValueError("Input file must be a BAM or BED file")

    if options.low_threshold is None:
        options.low_threshold = -options.threshold

    if options.high_threshold is None:
        options.high_threshold = options.threshold

    gc = GenomeCov(bedfile, options.low_threshold, options.high_threshold,
            0.5, 0.5)

    if options.reference:
        print('Computing GC content')
        gc.compute_gc_content(options.reference, options.w_gc)

    if len(gc.chr_list) == 1:
        if options.verbose:
            print("There is only one chromosome. Selected automatically.")
        chrom = gc.chr_list[0]
        run_analysis(gc, chrom, 1, options)
    elif options.chromosome <-1 or options.chromosome > len(gc.chr_list) or\
            options.chromosome == 0:
        raise ValueError("invalid --chromosome value ; must be in [1-%s]" % len(gc.chr_list)+1)
    else:
        # For uses, we start at position 1 but in python, we start at zero
        if options.chromosome == -1:
            chromosomes = [this for this in range(len(gc.chr_list)) ]
        else:
            chromosomes = [options.chromosome-1]

        N = len(chromosomes)
        for i,chrom_index in enumerate(chromosomes):
            chrom = gc.chr_list[chrom_index]
            chrom_name = gc.chr_list[chrom_index].chrom_name
            if options.verbose:
                print("There are %s chromosomes/contigs." % len(gc.chr_list))
                print("==================== analysing chrom/contig %s/%s (%s)" % (i+1,N,chrom_name))
            run_analysis(gc, chrom, chrom_index, options)

def run_analysis(gc, chrom, chrom_index, options):

    if options.verbose:
        print(chrom)

    if options.verbose:
        print('Computing running median (w=%s)' % options.w_median)

    chrom.running_median(n=options.w_median, circular=options.circular)

    stats = chrom.get_stats(output="dataframe")
    stats.set_index("name", inplace=True)

    DOC = stats.ix['DOC'].Value
    if options.k is None and DOC < 8:
        options.k = 1
    elif options.k is None:
        options.k = 2

    if options.verbose:
        print("Number of mixture model %s " % options.k)
        print('Computing zscore')
    chrom.compute_zscore(k=options.k, verbose=options.verbose)

    if options.verbose:
        print("Computing centralness")

    # Let us save the thresholds first and then change it to compute centralness
    thresholds = chrom.thresholds.copy()

    chrom.thresholds.low = -3
    chrom.thresholds.high = 3
    c3 = chrom.get_centralness()

    chrom.thresholds.low = -4
    chrom.thresholds.high = 4
    c4 = chrom.get_centralness()
    chrom.thresholds = thresholds.copy()   # Get back to the original values

    if options.verbose and chrom.thresholds:
        print(chrom.thresholds)

    if options.verbose:
        res = chrom._get_best_gaussian()
        print("sigma and mu of the central distribution: mu=%s, sigma=%s" %
(round(res["mu"],3), round(res['sigma'],3)))
        print("Evenness: %8.3f" % chrom.get_evenness())
        print("Centralness (3 sigma): %f" % round(c3,3))
        print("Centralness (4 sigma): %f" % round(c4,4))

    if options.verbose:
        print("\n\n")

    figure(1)
    chrom.plot_coverage()

    # With the reference, we can plot others plots such as GC versus coverage
    if options.reference:
        figure(2)
        chrom.plot_gc_vs_coverage(Nlevels=options.levels, fontsize=20)
        filename_gc_cov = "coverage_vs_gc.chrom{0}.png".format(chrom_index)
        savefig(filename_gc_cov)

    if options.show:
        show()

    # Report chromosomes
    if options.verbose:
        print("Creating report in %s" % options.output_directory)

    if options.genbank:
        if options.verbose:
            print('Genbank: %s' % options.genbank)
        from sequana.tools import genbank_features_parser
        features = genbank_features_parser(options.genbank)
    else:
        features = None

    if options.create_report:
        sample_name = "coverage"
        r = report_chromosome.ChromosomeMappingReport(chrom,
                directory=options.output_directory, project="coverage",
                sample=sample_name, features=features, verbose=options.verbose)

        if options.reference:
            from snakemake import shell
            shell("cp %s report/images" %  filename_gc_cov)
            r.jinja['coverage_vs_gc'] = """<img src="images/%s">""" % filename_gc_cov

        r.jinja['standalone_command'] = " ".join(sys.argv)
        r.create_report()
        if options.verbose:
            print("Report created. See ./report directory content and look for the HTML file. You can also use --show-html option")
        if options.show_html:
            from easydev import onweb
            onweb("report/%s_mapping.chrom%s.html" % (sample_name,
                  chrom_index))

if __name__ == "__main__":
   import sys
   main()#sys.argv)

<commit_msg>add correlation gc versus coverage in stdout and report<commit_after># -*- coding: utf-8 -*-
#
#  This file is part of Sequana software
#
#  Copyright (c) 2016 - Sequana Development Team
#
#  File author(s):
#      Thomas Cokelaer <thomas.cokelaer@pasteur.fr>
#      Dimitri Desvillechabrol <dimitri.desvillechabrol@pasteur.fr>,
#          <d.desvillechabrol@gmail.com>
#
#  Distributed under the terms of the 3-clause BSD license.
#  The full license is in the LICENSE file, distributed with this software.
#
#  website: https://github.com/sequana/sequana
#  documentation: http://sequana.readthedocs.io
#
##############################################################################
""".. rubric:: Standalone application dedicated to coverage"""
import os
import shutil
import glob
import sys
from optparse import OptionParser
import argparse
from argparse import RawTextHelpFormatter

from sequana import bedtools
from sequana.reporting import report_mapping
from sequana.reporting import report_chromosome
from sequana.reporting import report_main
from sequana import logger

from easydev import shellcmd

from pylab import show, figure, savefig

from easydev.console import purple

# http://stackoverflow.com/questions/18462610/argumentparser-epilog-and-description-formatting-in-conjunction-with-argumentdef
class CustomFormatter(argparse.ArgumentDefaultsHelpFormatter, argparse.RawDescriptionHelpFormatter):
    pass

class Options(argparse.ArgumentParser):
    def  __init__(self, prog="sequana_coverage"):
        usage = purple("""\nWelcome to SEQUANA -- Coverage standalone

    Extract and plot coverage of one or more chromosomes/contigs in a BED or BAM
    file. In addition, running median used in conjunction with double thresholds
    extract regions of interests (low or high coverage). A reference may be
    provided to plot the coverage versus GC content.

    The input file should be one of the following:

    - a BED file that is a tabulated file at least 3 columns.
      The first column being the reference, the second is the position
      and the third column contains the coverage itself.
    - or a BAM file that is converted automatically
      into a BED file using the following command:

        samtools depth -aa input.bam > output.bed

    If the reference is provided, an additional plot showing the coverage versus
    GC content is also shown.

    Here are some examples

        sequana_coverage --input file.bed --window-median 1001
        sequana_coverage --input file.bam --window-median 1001 -r <REFERENCE.fa>

    An other interesting option is to provide a BED file with 4 columns. The
    fourth column being another coverage data created with a filter. One can
    create such a file only from the BAM file using samtools as follows:

        samtools view -q 35  -o data.filtered.bam input.fa.sorted.bam
        samtools depth data.filtered.bam input.fa.sorted.bam  -aa > test.bed
        sequana_coverage --input test.bed --show-html

    Note that the first file is the filtered one, and the second file is the
    unfiltered one.

        """)

        epilog = purple("""
----

AUTHORS: Thomas Cokelaer, Dimitri Desvillechabrol
Documentation: http://sequana.readthedocs.io
Issues: http://github.com/sequana/sequana
        """)

        description = """DESCRIPTION:
        """

        super(Options, self).__init__(usage=usage, prog=prog,
                description=description, epilog=epilog,
                formatter_class=CustomFormatter)

        # options to fill the config file
        group = self.add_argument_group("Required argument")
        group.add_argument("-i", "--input", dest="input", type=str,
            help=("Input file in BED or BAM format. If a BAM file is "
                 "provided, it will be converted locally to a BED file "
                 "using genomecov, which must be installed."))

        group = self.add_argument_group("Optional biological arguments")
        group.add_argument('-c', "--chromosome", dest="chromosome", type=int,
            default=1,
            help=(  "Chromosome number (if only one, no need to use: the first"
                    " and only chromosome is chosen automatically). Default is"
                    " first chromosome found in the BED file. You may want to"
                    " analyse all chromosomes at the same time. If so, set this"
                    " parameter to -1"))
        group.add_argument('-o', "--circular", dest="circular",
            default=False, action="store_true",
            help="""If the DNA of the organism is circular (typically
            viruses or bacteria), set to True""")

        group = self.add_argument_group("General")
        group.add_argument("--output-directory", dest="output_directory",
            default="report", help="name of the output (report) directory.")
        group.add_argument('--show', dest="show", default=False,
            action='store_true', help="""Show the pictures (matplotlib)""")
        group.add_argument('--show-html', dest="show_html", default=False,
            action='store_true',
            help="""When report is created, you can open
            the main page automatically with this option (default is False)""")
        group.add_argument("-q", "--quiet", dest="verbose",
            default=True, action="store_false")
        group.add_argument('--no-report', dest="create_report",
            default=True, action='store_false',
            help="""Do not create any HTML report""")

        group = self.add_argument_group('Annotation')
        group.add_argument("-b", "--genbank", dest="genbank",
            type=str, default=None, help='a valida genbank annotation')

        # Group related to GC content
        group = self.add_argument_group("GC content related")
        group.add_argument('-r', "--reference", dest="reference", type=str,
            default=None,
            help="""If available, you can provide a reference (ENA/NCBI). It
                 must have the same length as the one used to create the
                 BAM or BED file. If provided, it is used to create the
                 coverage versus GC content image""")
        group.add_argument("-g", "--window-gc", dest="w_gc", type=int,
            help="""Length of the running window to compute the GC content""", default=200)
        group.add_argument('-n', "--nlevels", dest="levels", type=int,
            default=3, help="""Number of levels in the contour""")

        #group running median
        group = self.add_argument_group("Running Median related")
        group.add_argument("-w", "--window-median", dest="w_median", type=int,
            help="""Length of the running median window (default 4001,
                 recommended for viruses).  For long genome, 20001
                 or 30001 is recommended but larger windows may be
                 useful in the presence of long deleted regions.""",
            default=4001)

        group.add_argument("-k", "--mixture-models", dest="k", type=int,
            help="""Number of mixture models to use (default 2, although if sequencing
        depth is below 8, k is set to 1 automatically). To ignore that behaviour
        set k to the required value""",
            default=None)

        group.add_argument("-L", "--low-threshold", dest="low_threshold",
            default=None, type=float,
            help=("lower threshold (zscore) of the confidence interval. "
                "Overwrite value given by --threshold/-T"))
        group.add_argument("-H", "--high-threshold", dest="high_threshold",
            default=None, type=float,
            help=("higher threshold (zscore) of the confidence interval. "
                "Overwrite value given by --threshold/-T"))
        group.add_argument("-T", "--threshold", dest="threshold",
            default=4, type=float,
            help="""set lower and higher thresholds of the confidence interval.""")

        group = self.add_argument_group("Download reference")
        group.add_argument("--download-reference", dest="accession",
            default=None, type=str)
        group.add_argument("--download-genbank", dest="download_genbank",
            default=None, type=str)
        group.add_argument("--database", dest="database",
            default="ENA", type=str,
            choices=["ENA", "EUtils"],
            help="Download the reference from one of these database (default ENA)")

def main(args=None):

    if args is None:
        args = sys.argv[:]

    user_options = Options(prog="sequana")

    # If --help or no options provided, show the help
    if len(args) == 1:
        user_options.parse_args(["prog", "--help"])
    else:
        options = user_options.parse_args(args[1:])

    if options.accession:
        logger.info("Download accession %s from %s\n" %
            (options.accession, options.database))

        from bioservices.apps import download_fasta as df
        df.download_fasta(options.accession, method=options.database)
        return

    if options.download_genbank:
        from sequana.snpeff import download_fasta_and_genbank
        download_fasta_and_genbank(options.download_genbank,
                                   options.download_genbank,
                                   genbank=True, fasta=False)
        return

    # We put the import here to make the --help faster
    from sequana import GenomeCov

    if options.genbank:
        assert os.path.exists(options.genbank), \
            "%s does not exists" % options.genbank

    if options.verbose:
        logger.info("Reading %s" % options.input)

    if options.input.endswith(".bam"):
        bedfile = options.input.replace(".bam", ".bed")
        if options.verbose:
            print("Converting BAM into BED file")
        shellcmd("bedtools genomecov -d -ibam %s > %s" % (options.input, bedfile))
    elif options.input.endswith(".bed"):
        bedfile = options.input
    else:
        raise ValueError("Input file must be a BAM or BED file")

    if options.low_threshold is None:
        options.low_threshold = -options.threshold

    if options.high_threshold is None:
        options.high_threshold = options.threshold

    gc = GenomeCov(bedfile, options.low_threshold, options.high_threshold,
            0.5, 0.5)

    if options.reference:
        print('Computing GC content')
        gc.compute_gc_content(options.reference, options.w_gc)

    if len(gc.chr_list) == 1:
        if options.verbose:
            print("There is only one chromosome. Selected automatically.")
        chrom = gc.chr_list[0]
        run_analysis(gc, chrom, 1, options)
    elif options.chromosome <-1 or options.chromosome > len(gc.chr_list) or\
            options.chromosome == 0:
        raise ValueError("invalid --chromosome value ; must be in [1-%s]" % len(gc.chr_list)+1)
    else:
        # For uses, we start at position 1 but in python, we start at zero
        if options.chromosome == -1:
            chromosomes = [this for this in range(len(gc.chr_list)) ]
        else:
            chromosomes = [options.chromosome-1]

        N = len(chromosomes)
        for i,chrom_index in enumerate(chromosomes):
            chrom = gc.chr_list[chrom_index]
            chrom_name = gc.chr_list[chrom_index].chrom_name
            if options.verbose:
                print("There are %s chromosomes/contigs." % len(gc.chr_list))
                print("==================== analysing chrom/contig %s/%s (%s)" % (i+1,N,chrom_name))
            run_analysis(gc, chrom, chrom_index, options)

def run_analysis(gc, chrom, chrom_index, options):

    if options.verbose:
        print(chrom)

    if options.verbose:
        print('Computing running median (w=%s)' % options.w_median)

    chrom.running_median(n=options.w_median, circular=options.circular)

    stats = chrom.get_stats(output="dataframe")
    stats.set_index("name", inplace=True)

    DOC = stats.ix['DOC'].Value
    if options.k is None and DOC < 8:
        options.k = 1
    elif options.k is None:
        options.k = 2

    if options.verbose:
        print("Number of mixture model %s " % options.k)
        print('Computing zscore')
    chrom.compute_zscore(k=options.k, verbose=options.verbose)

    if options.verbose:
        print("Computing centralness")

    # Let us save the thresholds first and then change it to compute centralness
    thresholds = chrom.thresholds.copy()

    chrom.thresholds.low = -3
    chrom.thresholds.high = 3
    c3 = chrom.get_centralness()

    chrom.thresholds.low = -4
    chrom.thresholds.high = 4
    c4 = chrom.get_centralness()
    chrom.thresholds = thresholds.copy()   # Get back to the original values

    if options.verbose and chrom.thresholds:
        print(chrom.thresholds)

    if options.verbose:
        res = chrom._get_best_gaussian()
        print("sigma and mu of the central distribution: mu=%s, sigma=%s" %
(round(res["mu"],3), round(res['sigma'],3)))
        print("Evenness: %8.3f" % chrom.get_evenness())
        print("Centralness (3 sigma): %f" % round(c3,3))
        print("Centralness (4 sigma): %f" % round(c4,4))

    if options.verbose:
        print("\n\n")

    figure(1)
    chrom.plot_coverage()

    # With the reference, we can plot others plots such as GC versus coverage
    if options.reference:
        figure(2)
        chrom.plot_gc_vs_coverage(Nlevels=options.levels, fontsize=20)
        corr = chrom.get_gc_correlation()
        print("GC Correlation = %s" % corr)
        filename_gc_cov = "coverage_vs_gc.chrom{0}.png".format(chrom_index)
        savefig(filename_gc_cov)

    if options.show:
        show()

    # Report chromosomes
    if options.verbose:
        print("Creating report in %s" % options.output_directory)

    if options.genbank:
        if options.verbose:
            print('Genbank: %s' % options.genbank)
        from sequana.tools import genbank_features_parser
        features = genbank_features_parser(options.genbank)
    else:
        features = None

    if options.create_report:
        sample_name = "coverage"
        r = report_chromosome.ChromosomeMappingReport(chrom,
                directory=options.output_directory, project="coverage",
                sample=sample_name, features=features, verbose=options.verbose)

        if options.reference:
            from snakemake import shell
            shell("cp %s report/images" %  filename_gc_cov)
            r.jinja['coverage_vs_gc'] = """Correlation: %s </br> """ % corr
            r.jinja['coverage_vs_gc'] += """<img src="images/%s">""" % filename_gc_cov

        r.jinja['standalone_command'] = " ".join(sys.argv)
        r.create_report()
        if options.verbose:
            print("Report created. See ./report directory content and look for the HTML file. You can also use --show-html option")
        if options.show_html:
            from easydev import onweb
            onweb("report/%s_mapping.chrom%s.html" % (sample_name,
                  chrom_index))

if __name__ == "__main__":
   import sys
   main()#sys.argv)

<|endoftext|>


<commit_before>import copy
import matplotlib.patches as mpatches
from matplotlib.text import Text
from matplotlib.path import Path
from matplotlib.lines import Line2D
from matplotlib.transforms import Affine2D, Bbox, IdentityTransform

from matplotlib.text import Annotation

import numpy as np
from math import cos, sin, pi, atan2

def rotated_polygon(xy, ox, oy, angle):
    # angle in degree
    theta = angle/180.*pi

    st = sin(theta)
    ct = cos(theta)

    xy = np.asarray(xy, dtype="d")
    x, y = xy[:,0], xy[:,1]
    x1 = x - ox
    y1 = y - oy

    x2 =  ct*x1 + -st*y1
    y2 =  st*x1 +  ct*y1

    xp = x2 + ox
    yp = y2 + oy

    return np.hstack((xp.reshape((-1,1)), yp.reshape((-1,1))))

    # sss3 = [s1[0] for s1 in sss2 if isinstance(s1[0], parser_ds9.Shape)]

_point_type_dict=dict(circle="o",
                      box="s",
                      diamond="D",
                      x="x",
                      cross="+",
                      arrow="^",
                      boxcircle="*")

_ds9_to_mpl_colormap = dict(green="lime",
                            )

def properties_func_default(shape, saved_attrs):

    print "properties_func_default called"
    attr_list = copy.copy(shape.attr[0])
    attr_dict = copy.copy(shape.attr[1])

    attr_list.extend(saved_attrs[0])
    attr_dict.update(saved_attrs[1])

    color = attr_dict.get("color", None)
    color = _ds9_to_mpl_colormap.get(color, color)

    if shape.name ==  "text":
        kwargs = dict(color=color,
                      rotation=attr_dict.get("textangle", 0),
                      )
        font = attr_dict.get("font")
        if font:
            a = font.split()
            if len(a) == 3:
                fontsize=float(a[1])
            kwargs["fontsize"]=fontsize
    elif shape.name == "point":
        marker = _point_type_dict.get(attr_dict.get("point"), "o")
        kwargs = dict(markeredgecolor=color,
                      markerfacecolor="none",
                      marker=marker,
                      markersize=int(attr_dict.get("size", 1)),
                      markeredgewidth=int(attr_dict.get("width", 1)),
                      )
    elif shape.name in ["line", "vector"]:
        font = attr_dict.get("font")
        if font:
            fontsize = int(font.split()[1])
        else:
            fontsize = 10

        kwargs = dict(color=color,
                      linewidth=int(attr_dict.get("width", 1)),
                      mutation_scale=fontsize,
                      )
        if int(attr_dict.get("dash","0")):
            kwargs["linestyle"] = "dashed"

    else:
        kwargs = dict(edgecolor=color,
                      linewidth=int(attr_dict.get("width", 1)),
                      facecolor="none"
                      )

        if "background" in attr_list:
            kwargs["linestyle"] = "dashed"

        if int(attr_dict.get("dash","0")):
            kwargs["linestyle"] = "dashed"
        if shape.exclude:
            kwargs["hatch"] = "/"

    return kwargs

def _get_text(txt, x, y, dx, dy, ha="center", va="center", **kwargs):
    if "color" in kwargs:
        textcolor = kwargs["color"]
    elif "markeredgecolor" in kwargs:
        textcolor = kwargs["markeredgecolor"]
    ann = Annotation(txt, (x, y), xytext=(dx, dy),
                     xycoords='data',
                     textcoords="offset points",
                     ha=ha, va=va,
                     **kwargs)
    ann.set_transform(IdentityTransform())

    return ann

def as_mpl_artists(shape_list,
                   properties_func=None,
                   text_offset=5.0):
    """
    Converts a region list to a list of patches and a list of artists.

    Optional Keywords:
    [ text_offset ] - If there is text associated with the regions, add
    some vertical offset (in pixels) to the text so that it doesn't overlap
    with the regions.
    """

    patch_list = []
    artist_list = []

    if properties_func is None:
        properties_func = properties_func_default

    # properties for continued(? multiline?) regions
    saved_attrs = None

    for shape in shape_list:

        print "as_mpl_artists shape: ",shape

        patches = []

        if saved_attrs is None:
            _attrs = [], {}
        else:
            _attrs = copy.copy(saved_attrs[0]), copy.copy(saved_attrs[1])

        kwargs = properties_func(shape, _attrs)

        if shape.name == "composite":
            saved_attrs = shape.attr
            continue

        if saved_attrs is None and shape.continued:
            saved_attrs = shape.attr
#         elif (shape.name in shape.attr[1]):
#             if (shape.attr[1][shape.name] != "ignore"):
#                 saved_attrs = shape.attr

        if not shape.continued:
            saved_attrs = None

        # text associated with the shape
        txt = shape.attr[1].get("text","")

        if shape.name == "polygon":
            xy = np.array(shape.coord_list)
            xy.shape = -1,2

            # -1 for change origin to 0,0
            patches=[mpatches.Polygon(xy-1, closed=True, **kwargs)]

        elif shape.name == "rotbox" or shape.name == "box":
            xc, yc, w, h, rot = shape.coord_list
            # -1 for change origin to 0,0
            xc, yc = xc-1, yc-1
            _box = np.array([[-w/2., -h/2.],
                             [-w/2., h/2.],
                             [w/2., h/2.],
                             [w/2., -h/2.]])
            box = _box + [xc, yc]
            rotbox = rotated_polygon(box, xc, yc, rot)
            patches = [mpatches.Polygon(rotbox, closed=True, **kwargs)]

        elif shape.name == "ellipse":
            xc, yc  = shape.coord_list[:2]
            # -1 for change origin to 0,0
            xc, yc = xc-1, yc-1
            angle = shape.coord_list[-1]

            maj_list, min_list = shape.coord_list[2:-1:2], shape.coord_list[3:-1:2]

            patches = [mpatches.Ellipse((xc, yc), 2*maj, 2*min,
                                        angle=angle, **kwargs) \
                       for maj, min in zip(maj_list, min_list)]

        elif shape.name == "annulus":
            xc, yc  = shape.coord_list[:2]
            # -1 for change origin to 0,0
            xc, yc = xc-1, yc-1
            r_list = shape.coord_list[2:]

            patches = [mpatches.Ellipse((xc, yc), 2*r, 2*r,
                                        **kwargs) \
                       for r in r_list]

        elif shape.name == "circle":
            xc, yc, major = shape.coord_list
            # -1 for change origin to 0,0
            xc, yc = xc-1, yc-1
            patches = [mpatches.Ellipse((xc, yc), 2*major, 2*major,
                                        angle=0, **kwargs)]

        elif shape.name == "panda":
            xc, yc, a1, a2, an, r1, r2, rn = shape.coord_list
            # -1 for change origin to 0,0
            xc, yc = xc-1, yc-1
            patches = [mpatches.Arc((xc, yc), rr*2, rr*2, angle=0,
                                    theta1=a1, theta2=a2, **kwargs) \
                       for rr in np.linspace(r1, r2, rn+1)]

            for aa in np.linspace(a1, a2, an+1):
                xx = np.array([r1, r2]) * np.cos(aa/180.*np.pi) + xc
                yy = np.array([r1, r2]) * np.sin(aa/180.*np.pi) + yc
                p = Path(np.transpose([xx,yy]))
                patches.append(mpatches.PathPatch(p, **kwargs))

        elif shape.name == "pie":
            xc, yc, r1, r2, a1, a2 = shape.coord_list
            # -1 for change origin to 0,0
            xc, yc = xc-1, yc-1

            patches = [mpatches.Arc((xc, yc), rr*2, rr*2, angle=0,
                                   theta1=a1, theta2=a2, **kwargs) \
                       for rr in [r1, r2]]

            for aa in [a1, a2]:
                xx = np.array([r1, r2]) * np.cos(aa/180.*np.pi) + xc
                yy = np.array([r1, r2]) * np.sin(aa/180.*np.pi) + yc
                p = Path(np.transpose([xx,yy]))
                patches.append(mpatches.PathPatch(p, **kwargs))

        elif shape.name == "epanda":
            xc, yc, a1, a2, an, r11, r12,r21, r22, rn, angle = shape.coord_list
            # -1 for change origin to 0,0
            xc, yc = xc-1, yc-1

            # mpl takes angle a1, a2 as angle as in circle before
            # transformation to ellipse.

            x1, y1 = cos(a1/180.*pi), sin(a1/180.*pi)*r11/r12
            x2, y2 = cos(a2/180.*pi), sin(a2/180.*pi)*r11/r12

            a1, a2 = atan2(y1, x1)/pi*180., atan2(y2, x2)/pi*180.

            patches = [mpatches.Arc((xc, yc), rr1*2, rr2*2,
                                    angle=angle, theta1=a1, theta2=a2,
                                  **kwargs) \
                       for rr1, rr2 in zip(np.linspace(r11, r21, rn+1),
                                           np.linspace(r12, r22, rn+1))]

            for aa in np.linspace(a1, a2, an+1):
                xx = np.array([r11, r21]) * np.cos(aa/180.*np.pi)
                yy = np.array([r11, r21]) * np.sin(aa/180.*np.pi)
                p = Path(np.transpose([xx,yy]))
                tr = Affine2D().scale(1, r12/r11).rotate_deg(angle).translate(xc, yc)
                p2 = tr.transform_path(p)
                patches.append(mpatches.PathPatch(p2, **kwargs))

        elif shape.name == "text":
            xc, yc  = shape.coord_list[:2]
            # -1 for change origin to 0,0
            xc, yc = xc-1, yc-1

            if txt:
                _t = _get_text(txt, xc, yc, 0, 0, **kwargs)
                artist_list.append(_t)

        elif shape.name == "point":
            xc, yc  = shape.coord_list[:2]
            # -1 for change origin to 0,0
            xc, yc = xc-1, yc-1
            artist_list.append(Line2D([xc], [yc],
                                      **kwargs))

            if txt:
                textshape = copy.copy(shape)
                textshape.name = "text"
                textkwargs = properties_func(textshape, _attrs)
                _t = _get_text(txt, xc, yc, 0, text_offset,
                               va="bottom",
                               **textkwargs)
                artist_list.append(_t)

        elif shape.name in ["line", "vector"]:
            if shape.name == "line":
                x1, y1, x2, y2  = shape.coord_list[:4]
                # -1 for change origin to 0,0
                x1, y1, x2, y2 = x1-1, y1-1, x2-1, y2-1

                a1, a2 = shape.attr[1].get("line", "0 0").strip().split()[:2]

                arrowstyle = "-"
                if int(a1):
                    arrowstyle = "<" + arrowstyle
                if int(a2):
                    arrowstyle = arrowstyle + ">"

            else: # shape.name == "vector"
                x1, y1, l, a  = shape.coord_list[:4]
                # -1 for change origin to 0,0
                x1, y1 = x1-1, y1-1
                x2, y2 = x1 + l * np.cos(a/180.*np.pi), y1 + l * np.sin(a/180.*np.pi)
                v1 = int(shape.attr[1].get("vector", "0").strip())

                if v1:
                    arrowstyle="->"
                else:
                    arrowstyle="-"

            patches = [mpatches.FancyArrowPatch(posA=(x1, y1),
                                                posB=(x2, y2),
                                                arrowstyle=arrowstyle,
                                                arrow_transmuter=None,
                                                connectionstyle="arc3",
                                                patchA=None, patchB=None,
                                                shrinkA=0, shrinkB=0,
                                                connector=None,
                                                **kwargs)]

        else:
            print "'as_mpl_artists' does not know how to convert '%s' to mpl artist" % (shape.name,)

        patch_list.extend(patches)

        if txt and patches:
            # the text associated with a shape uses different
            # matplotlib keywords than the shape itself for, e.g.,
            # color
            textshape = copy.copy(shape)
            textshape.name = "text"
            textkwargs = properties_func(textshape, _attrs)

            # calculate the text position
            _bb = [p.get_window_extent() for p in patches]
            _bbox = Bbox.union(_bb)
            x0, y0, x1, y1 = _bbox.extents
            xc = .5*(x0+x1)

            _t = _get_text(txt, xc, y1, 0, text_offset,
                           va="bottom",
                           **textkwargs)
            artist_list.append(_t)

    return patch_list, artist_list

if 0:
    import parser_ds9
    regname = "az_region.reg"
    shape_list = region_read_as_imagecoord(regname, wcs)
    p = as_mpl_patches(shape_list)

    clf()
    ax = gca()
    [ax.add_patch(p1) for p1 in p]
    xlim(0,1500)
    ylim(0,1500)

    draw()

<commit_msg>Removed leftover debug statements from last commit (whoops)<commit_after>import copy
import matplotlib.patches as mpatches
from matplotlib.text import Text
from matplotlib.path import Path
from matplotlib.lines import Line2D
from matplotlib.transforms import Affine2D, Bbox, IdentityTransform

from matplotlib.text import Annotation

import numpy as np
from math import cos, sin, pi, atan2

def rotated_polygon(xy, ox, oy, angle):
    # angle in degree
    theta = angle/180.*pi

    st = sin(theta)
    ct = cos(theta)

    xy = np.asarray(xy, dtype="d")
    x, y = xy[:,0], xy[:,1]
    x1 = x - ox
    y1 = y - oy

    x2 =  ct*x1 + -st*y1
    y2 =  st*x1 +  ct*y1

    xp = x2 + ox
    yp = y2 + oy

    return np.hstack((xp.reshape((-1,1)), yp.reshape((-1,1))))

    # sss3 = [s1[0] for s1 in sss2 if isinstance(s1[0], parser_ds9.Shape)]

_point_type_dict=dict(circle="o",
                      box="s",
                      diamond="D",
                      x="x",
                      cross="+",
                      arrow="^",
                      boxcircle="*")

_ds9_to_mpl_colormap = dict(green="lime",
                            )

def properties_func_default(shape, saved_attrs):

    attr_list = copy.copy(shape.attr[0])
    attr_dict = copy.copy(shape.attr[1])

    attr_list.extend(saved_attrs[0])
    attr_dict.update(saved_attrs[1])

    color = attr_dict.get("color", None)
    color = _ds9_to_mpl_colormap.get(color, color)

    if shape.name ==  "text":
        kwargs = dict(color=color,
                      rotation=attr_dict.get("textangle", 0),
                      )
        font = attr_dict.get("font")
        if font:
            a = font.split()
            if len(a) == 3:
                fontsize=float(a[1])
            kwargs["fontsize"]=fontsize
    elif shape.name == "point":
        marker = _point_type_dict.get(attr_dict.get("point"), "o")
        kwargs = dict(markeredgecolor=color,
                      markerfacecolor="none",
                      marker=marker,
                      markersize=int(attr_dict.get("size", 1)),
                      markeredgewidth=int(attr_dict.get("width", 1)),
                      )
    elif shape.name in ["line", "vector"]:
        font = attr_dict.get("font")
        if font:
            fontsize = int(font.split()[1])
        else:
            fontsize = 10

        kwargs = dict(color=color,
                      linewidth=int(attr_dict.get("width", 1)),
                      mutation_scale=fontsize,
                      )
        if int(attr_dict.get("dash","0")):
            kwargs["linestyle"] = "dashed"

    else:
        kwargs = dict(edgecolor=color,
                      linewidth=int(attr_dict.get("width", 1)),
                      facecolor="none"
                      )

        if "background" in attr_list:
            kwargs["linestyle"] = "dashed"

        if int(attr_dict.get("dash","0")):
            kwargs["linestyle"] = "dashed"
        if shape.exclude:
            kwargs["hatch"] = "/"

    return kwargs

def _get_text(txt, x, y, dx, dy, ha="center", va="center", **kwargs):
    if "color" in kwargs:
        textcolor = kwargs["color"]
    elif "markeredgecolor" in kwargs:
        textcolor = kwargs["markeredgecolor"]
    ann = Annotation(txt, (x, y), xytext=(dx, dy),
                     xycoords='data',
                     textcoords="offset points",
                     ha=ha, va=va,
                     **kwargs)
    ann.set_transform(IdentityTransform())

    return ann

def as_mpl_artists(shape_list,
                   properties_func=None,
                   text_offset=5.0):
    """
    Converts a region list to a list of patches and a list of artists.

    Optional Keywords:
    [ text_offset ] - If there is text associated with the regions, add
    some vertical offset (in pixels) to the text so that it doesn't overlap
    with the regions.
    """

    patch_list = []
    artist_list = []

    if properties_func is None:
        properties_func = properties_func_default

    # properties for continued(? multiline?) regions
    saved_attrs = None

    for shape in shape_list:

        patches = []

        if saved_attrs is None:
            _attrs = [], {}
        else:
            _attrs = copy.copy(saved_attrs[0]), copy.copy(saved_attrs[1])

        kwargs = properties_func(shape, _attrs)

        if shape.name == "composite":
            saved_attrs = shape.attr
            continue

        if saved_attrs is None and shape.continued:
            saved_attrs = shape.attr
#         elif (shape.name in shape.attr[1]):
#             if (shape.attr[1][shape.name] != "ignore"):
#                 saved_attrs = shape.attr

        if not shape.continued:
            saved_attrs = None

        # text associated with the shape
        txt = shape.attr[1].get("text","")

        if shape.name == "polygon":
            xy = np.array(shape.coord_list)
            xy.shape = -1,2

            # -1 for change origin to 0,0
            patches=[mpatches.Polygon(xy-1, closed=True, **kwargs)]

        elif shape.name == "rotbox" or shape.name == "box":
            xc, yc, w, h, rot = shape.coord_list
            # -1 for change origin to 0,0
            xc, yc = xc-1, yc-1
            _box = np.array([[-w/2., -h/2.],
                             [-w/2., h/2.],
                             [w/2., h/2.],
                             [w/2., -h/2.]])
            box = _box + [xc, yc]
            rotbox = rotated_polygon(box, xc, yc, rot)
            patches = [mpatches.Polygon(rotbox, closed=True, **kwargs)]

        elif shape.name == "ellipse":
            xc, yc  = shape.coord_list[:2]
            # -1 for change origin to 0,0
            xc, yc = xc-1, yc-1
            angle = shape.coord_list[-1]

            maj_list, min_list = shape.coord_list[2:-1:2], shape.coord_list[3:-1:2]

            patches = [mpatches.Ellipse((xc, yc), 2*maj, 2*min,
                                        angle=angle, **kwargs) \
                       for maj, min in zip(maj_list, min_list)]

        elif shape.name == "annulus":
            xc, yc  = shape.coord_list[:2]
            # -1 for change origin to 0,0
            xc, yc = xc-1, yc-1
            r_list = shape.coord_list[2:]

            patches = [mpatches.Ellipse((xc, yc), 2*r, 2*r,
                                        **kwargs) \
                       for r in r_list]

        elif shape.name == "circle":
            xc, yc, major = shape.coord_list
            # -1 for change origin to 0,0
            xc, yc = xc-1, yc-1
            patches = [mpatches.Ellipse((xc, yc), 2*major, 2*major,
                                        angle=0, **kwargs)]

        elif shape.name == "panda":
            xc, yc, a1, a2, an, r1, r2, rn = shape.coord_list
            # -1 for change origin to 0,0
            xc, yc = xc-1, yc-1
            patches = [mpatches.Arc((xc, yc), rr*2, rr*2, angle=0,
                                    theta1=a1, theta2=a2, **kwargs) \
                       for rr in np.linspace(r1, r2, rn+1)]

            for aa in np.linspace(a1, a2, an+1):
                xx = np.array([r1, r2]) * np.cos(aa/180.*np.pi) + xc
                yy = np.array([r1, r2]) * np.sin(aa/180.*np.pi) + yc
                p = Path(np.transpose([xx,yy]))
                patches.append(mpatches.PathPatch(p, **kwargs))

        elif shape.name == "pie":
            xc, yc, r1, r2, a1, a2 = shape.coord_list
            # -1 for change origin to 0,0
            xc, yc = xc-1, yc-1

            patches = [mpatches.Arc((xc, yc), rr*2, rr*2, angle=0,
                                   theta1=a1, theta2=a2, **kwargs) \
                       for rr in [r1, r2]]

            for aa in [a1, a2]:
                xx = np.array([r1, r2]) * np.cos(aa/180.*np.pi) + xc
                yy = np.array([r1, r2]) * np.sin(aa/180.*np.pi) + yc
                p = Path(np.transpose([xx,yy]))
                patches.append(mpatches.PathPatch(p, **kwargs))

        elif shape.name == "epanda":
            xc, yc, a1, a2, an, r11, r12,r21, r22, rn, angle = shape.coord_list
            # -1 for change origin to 0,0
            xc, yc = xc-1, yc-1

            # mpl takes angle a1, a2 as angle as in circle before
            # transformation to ellipse.

            x1, y1 = cos(a1/180.*pi), sin(a1/180.*pi)*r11/r12
            x2, y2 = cos(a2/180.*pi), sin(a2/180.*pi)*r11/r12

            a1, a2 = atan2(y1, x1)/pi*180., atan2(y2, x2)/pi*180.

            patches = [mpatches.Arc((xc, yc), rr1*2, rr2*2,
                                    angle=angle, theta1=a1, theta2=a2,
                                  **kwargs) \
                       for rr1, rr2 in zip(np.linspace(r11, r21, rn+1),
                                           np.linspace(r12, r22, rn+1))]

            for aa in np.linspace(a1, a2, an+1):
                xx = np.array([r11, r21]) * np.cos(aa/180.*np.pi)
                yy = np.array([r11, r21]) * np.sin(aa/180.*np.pi)
                p = Path(np.transpose([xx,yy]))
                tr = Affine2D().scale(1, r12/r11).rotate_deg(angle).translate(xc, yc)
                p2 = tr.transform_path(p)
                patches.append(mpatches.PathPatch(p2, **kwargs))

        elif shape.name == "text":
            xc, yc  = shape.coord_list[:2]
            # -1 for change origin to 0,0
            xc, yc = xc-1, yc-1

            if txt:
                _t = _get_text(txt, xc, yc, 0, 0, **kwargs)
                artist_list.append(_t)

        elif shape.name == "point":
            xc, yc  = shape.coord_list[:2]
            # -1 for change origin to 0,0
            xc, yc = xc-1, yc-1
            artist_list.append(Line2D([xc], [yc],
                                      **kwargs))

            if txt:
                textshape = copy.copy(shape)
                textshape.name = "text"
                textkwargs = properties_func(textshape, _attrs)
                _t = _get_text(txt, xc, yc, 0, text_offset,
                               va="bottom",
                               **textkwargs)
                artist_list.append(_t)

        elif shape.name in ["line", "vector"]:
            if shape.name == "line":
                x1, y1, x2, y2  = shape.coord_list[:4]
                # -1 for change origin to 0,0
                x1, y1, x2, y2 = x1-1, y1-1, x2-1, y2-1

                a1, a2 = shape.attr[1].get("line", "0 0").strip().split()[:2]

                arrowstyle = "-"
                if int(a1):
                    arrowstyle = "<" + arrowstyle
                if int(a2):
                    arrowstyle = arrowstyle + ">"

            else: # shape.name == "vector"
                x1, y1, l, a  = shape.coord_list[:4]
                # -1 for change origin to 0,0
                x1, y1 = x1-1, y1-1
                x2, y2 = x1 + l * np.cos(a/180.*np.pi), y1 + l * np.sin(a/180.*np.pi)
                v1 = int(shape.attr[1].get("vector", "0").strip())

                if v1:
                    arrowstyle="->"
                else:
                    arrowstyle="-"

            patches = [mpatches.FancyArrowPatch(posA=(x1, y1),
                                                posB=(x2, y2),
                                                arrowstyle=arrowstyle,
                                                arrow_transmuter=None,
                                                connectionstyle="arc3",
                                                patchA=None, patchB=None,
                                                shrinkA=0, shrinkB=0,
                                                connector=None,
                                                **kwargs)]

        else:
            print "'as_mpl_artists' does not know how to convert '%s' to mpl artist" % (shape.name,)

        patch_list.extend(patches)

        if txt and patches:
            # the text associated with a shape uses different
            # matplotlib keywords than the shape itself for, e.g.,
            # color
            textshape = copy.copy(shape)
            textshape.name = "text"
            textkwargs = properties_func(textshape, _attrs)

            # calculate the text position
            _bb = [p.get_window_extent() for p in patches]
            _bbox = Bbox.union(_bb)
            x0, y0, x1, y1 = _bbox.extents
            xc = .5*(x0+x1)

            _t = _get_text(txt, xc, y1, 0, text_offset,
                           va="bottom",
                           **textkwargs)
            artist_list.append(_t)

    return patch_list, artist_list

if 0:
    import parser_ds9
    regname = "az_region.reg"
    shape_list = region_read_as_imagecoord(regname, wcs)
    p = as_mpl_patches(shape_list)

    clf()
    ax = gca()
    [ax.add_patch(p1) for p1 in p]
    xlim(0,1500)
    ylim(0,1500)

    draw()

<|endoftext|>


<commit_before>#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
#

from __future__ import unicode_literals
from __future__ import division
from __future__ import absolute_import
from __future__ import print_function

import os
import time
import traceback
from threading import Event
from threading import Timer

from system_test import TestCase, Qdrouterd, main_module, TIMEOUT
from system_test import Logger
from system_test import QdManager
from system_test import unittest
from system_test import Process
from system_test import DIR
from qpid_dispatch.management.client import Node
from subprocess import PIPE, STDOUT
from TCP_echo_client import TcpEchoClient
from TCP_echo_server import TcpEchoServer

class TcpAdaptorOneRouterEcho(TestCase, Process):
    """
    Run echo tests through a stand-alone router
    """
    amqp_listener_port       = None
    tcp_client_listener_port = None
    tcp_server_listener_port = None

    echo_timeout = 30 # local timeout to wait for one echo client to finish

    @classmethod
    def setUpClass(cls):
        """Start a router"""
        super(TcpAdaptorOneRouterEcho, cls).setUpClass()

        def router(name, mode, l_amqp, l_tcp_client, l_tcp_server, addr, site, extra=None):
            config = [
                ('router', {'mode': mode, 'id': name}),
                ('listener', {'port': l_amqp, 'stripAnnotations': 'no'}),
                ('tcpConnector', {"host": "127.0.0.1",
                                  "port": l_tcp_server,
                                  "address": addr,
                                  "siteId": site}),
                ('tcpListener', {"host": "0.0.0.0",
                                 "port": l_tcp_client,
                                 "address": addr,
                                 "siteId": site})
            ]

            if extra:
                config.append(extra)
            config = Qdrouterd.Config(config)
            cls.routers.append(cls.tester.qdrouterd(name, config, wait=True))

        cls.routers = []

        cls.amqp_listener_port       = cls.tester.get_port()
        cls.tcp_client_listener_port = cls.tester.get_port()
        cls.tcp_server_listener_port = cls.tester.get_port()

        router('A', 'interior', cls.amqp_listener_port, cls.tcp_client_listener_port,
               cls.tcp_server_listener_port, "some_address", "best_site")

        cls.logger = Logger(title="TcpAdaptorOneRouterEcho-testClass",
                            print_to_console=True,
                            save_for_dump=False)
        cls.logger.log("  amqp_listener_port       : %d" % cls.amqp_listener_port)
        cls.logger.log("  tcp_client_listener_port : %d" % cls.tcp_client_listener_port)
        cls.logger.log("  tcp_server_listener_port : %d" % cls.tcp_server_listener_port)

    def do_test_echo(self, test_name, logger, host, port, size, count, print_client_logs):
        # Run echo client. Return true if it works.
        name = "%s_%s_%s_%s" % (test_name, port, size, count)
        client_prefix = "ECHO_CLIENT %s" % name
        client_logger = Logger(title=client_prefix,
                               print_to_console=print_client_logs,
                               save_for_dump=False)
        result = True # assume it works
        start_time = time.time()
        try:
            # start client
            client = TcpEchoClient(prefix=client_prefix,
                                   host=host,
                                   port=port,
                                   size=size,
                                   count=count,
                                   timeout=TIMEOUT,
                                   logger=client_logger)
            #assert client.is_running

            # wait for client to finish
            keep_waiting = True
            while keep_waiting:
                time.sleep(0.1)
                elapsed = time.time() - start_time
                if elapsed > self.echo_timeout:
                    client.error = "TIMEOUT - local wait time exceeded"
                    logger.log("%s %s" % (name, client.error))
                    keep_waiting = False
                    result = False
                if client.error is not None:
                    logger.log("%s Client stopped with error: %s" % (name, client.error))
                    keep_waiting = False
                    result = False
                if client.exit_status is not None:
                    logger.log("%s Client stopped with status: %s" % (name, client.exit_status))
                    keep_waiting = False
                    result = False
                if keep_waiting and not client.is_running:
                    logger.log("%s Client stopped with no error or status" % (name))
                    keep_waiting = False

        except Exception as exc:
            logger.log("EchoClient %s failed. Exception: %s" %
                       (name, traceback.format_exc()))
            result = False

        # wait for client to exit
        client.wait()

        if not result:
            pass

        return result

    def test_01_tcp_echo_one_router(self):
        """
        Run one echo server.
        Run many echo clients.
        :return:
        """
        # define logging
        print_logs_server = True
        print_logs_client = True

        # start echo server
        test_name = "test_01_tcp_echo_one_router"
        server_prefix = "ECHO_SERVER %s" % test_name
        server_logger = Logger(title=test_name,
                               print_to_console=print_logs_server,
                               save_for_dump=False)
        server = TcpEchoServer(prefix=server_prefix,
                               port=self.tcp_server_listener_port,
                               timeout=TIMEOUT,
                               logger=server_logger)
        assert server.is_running

        # run series of clients to test
        result = True
        for size in [1, 50, 1000]:
            for count in [1, 10]:
                # make sure server is still running
                if server.error is not None:
                    self.logger.log("%s Server stopped with error: %s" % (self.name, server.error))
                    result = False
                if server.exit_status is not None:
                    self.logger.log("%s Server stopped with status: %s" % (self.name, server.exit_status))
                    result = False
                # run another test client
                if result:
                    test_info = "Starting echo client %s host:localhost, port:%d, size:%d, count:%d" % \
                               (test_name, self.tcp_client_listener_port, size, count)
                    self.logger.log(test_info)
                    result = self.do_test_echo(test_name, self.logger, "localhost",
                                               self.tcp_client_listener_port, size, count,
                                               print_logs_client)
                if not result:
                    break
            if not result:
                break
        # stop echo server
        server.wait()
        assert result, "Test case failed %s" % test_info
        self.logger.log("Test test_01_tcp_echo_one_router: SUCCESS")

if __name__== '__main__':
    unittest.main(main_module())
<commit_msg>DISPATCH-1825: Skip TCP self test if selectors module is absent<commit_after>#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
#

from __future__ import unicode_literals
from __future__ import division
from __future__ import absolute_import
from __future__ import print_function

import os
import time
import traceback
from threading import Event
from threading import Timer

from system_test import TestCase, Qdrouterd, main_module, TIMEOUT
from system_test import Logger
from system_test import QdManager
from system_test import unittest
from system_test import DIR
from qpid_dispatch.management.client import Node
from subprocess import PIPE, STDOUT

try:
    from TCP_echo_client import TcpEchoClient
    from TCP_echo_server import TcpEchoServer
except ImportError:
    class TCP_echo_client(object):
        pass
    class TCP_echo_server(object):
        pass

class TcpAdaptorOneRouterEcho(TestCase):
    """
    Run echo tests through a stand-alone router
    """
    amqp_listener_port       = None
    tcp_client_listener_port = None
    tcp_server_listener_port = None

    echo_timeout = 30 # local timeout to wait for one echo client to finish

    @classmethod
    def setUpClass(cls):
        """Start a router"""
        super(TcpAdaptorOneRouterEcho, cls).setUpClass()

        def router(name, mode, l_amqp, l_tcp_client, l_tcp_server, addr, site, extra=None):
            config = [
                ('router', {'mode': mode, 'id': name}),
                ('listener', {'port': l_amqp, 'stripAnnotations': 'no'}),
                ('tcpConnector', {"host": "127.0.0.1",
                                  "port": l_tcp_server,
                                  "address": addr,
                                  "siteId": site}),
                ('tcpListener', {"host": "0.0.0.0",
                                 "port": l_tcp_client,
                                 "address": addr,
                                 "siteId": site})
            ]

            if extra:
                config.append(extra)
            config = Qdrouterd.Config(config)
            cls.routers.append(cls.tester.qdrouterd(name, config, wait=True))

        cls.routers = []

        cls.amqp_listener_port       = cls.tester.get_port()
        cls.tcp_client_listener_port = cls.tester.get_port()
        cls.tcp_server_listener_port = cls.tester.get_port()

        router('A', 'interior', cls.amqp_listener_port, cls.tcp_client_listener_port,
               cls.tcp_server_listener_port, "some_address", "best_site")

        cls.logger = Logger(title="TcpAdaptorOneRouterEcho-testClass",
                            print_to_console=True,
                            save_for_dump=False)
        cls.logger.log("  amqp_listener_port       : %d" % cls.amqp_listener_port)
        cls.logger.log("  tcp_client_listener_port : %d" % cls.tcp_client_listener_port)
        cls.logger.log("  tcp_server_listener_port : %d" % cls.tcp_server_listener_port)

    def do_test_echo(self, test_name, logger, host, port, size, count, print_client_logs):
        # Run echo client. Return true if it works.
        name = "%s_%s_%s_%s" % (test_name, port, size, count)
        client_prefix = "ECHO_CLIENT %s" % name
        client_logger = Logger(title=client_prefix,
                               print_to_console=print_client_logs,
                               save_for_dump=False)
        result = True # assume it works
        start_time = time.time()
        try:
            # start client
            client = TcpEchoClient(prefix=client_prefix,
                                   host=host,
                                   port=port,
                                   size=size,
                                   count=count,
                                   timeout=TIMEOUT,
                                   logger=client_logger)
            #assert client.is_running

            # wait for client to finish
            keep_waiting = True
            while keep_waiting:
                time.sleep(0.1)
                elapsed = time.time() - start_time
                if elapsed > self.echo_timeout:
                    client.error = "TIMEOUT - local wait time exceeded"
                    logger.log("%s %s" % (name, client.error))
                    keep_waiting = False
                    result = False
                if client.error is not None:
                    logger.log("%s Client stopped with error: %s" % (name, client.error))
                    keep_waiting = False
                    result = False
                if client.exit_status is not None:
                    logger.log("%s Client stopped with status: %s" % (name, client.exit_status))
                    keep_waiting = False
                    result = False
                if keep_waiting and not client.is_running:
                    logger.log("%s Client stopped with no error or status" % (name))
                    keep_waiting = False

        except Exception as exc:
            logger.log("EchoClient %s failed. Exception: %s" %
                       (name, traceback.format_exc()))
            result = False

        # wait for client to exit
        client.wait()

        if not result:
            pass

        return result

    def test_01_tcp_echo_one_router(self):
        """
        Run one echo server.
        Run many echo clients.
        :return:
        """
        # run test only if selectors module is available
        try:
            import selectors
        except ImportError:
            # server and client modules require selectors
            self.skipTest("Python selectors module is not available on this platform.")

        # define logging
        print_logs_server = True
        print_logs_client = True

        # start echo server
        test_name = "test_01_tcp_echo_one_router"
        server_prefix = "ECHO_SERVER %s" % test_name
        server_logger = Logger(title=test_name,
                               print_to_console=print_logs_server,
                               save_for_dump=False)
        server = TcpEchoServer(prefix=server_prefix,
                               port=self.tcp_server_listener_port,
                               timeout=TIMEOUT,
                               logger=server_logger)
        assert server.is_running

        # run series of clients to test
        result = True
        for size in [1, 50, 1000]:
            for count in [1, 10]:
                # make sure server is still running
                if server.error is not None:
                    self.logger.log("%s Server stopped with error: %s" % (self.name, server.error))
                    result = False
                if server.exit_status is not None:
                    self.logger.log("%s Server stopped with status: %s" % (self.name, server.exit_status))
                    result = False
                # run another test client
                if result:
                    test_info = "Starting echo client %s host:localhost, port:%d, size:%d, count:%d" % \
                               (test_name, self.tcp_client_listener_port, size, count)
                    self.logger.log(test_info)
                    result = self.do_test_echo(test_name, self.logger, "localhost",
                                               self.tcp_client_listener_port, size, count,
                                               print_logs_client)
                if not result:
                    break
            if not result:
                break
        # stop echo server
        server.wait()
        assert result, "Test case failed %s" % test_info
        self.logger.log("Test test_01_tcp_echo_one_router: SUCCESS")

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