<commit_before>
from django.contrib.contenttypes.generic import GenericForeignKey
from django.core.exceptions import ObjectDoesNotExist
from django.db import models
from django.db.models.base import ModelBase
from django.template.defaultfilters import truncatewords_html
from django.utils.html import strip_tags
from django.utils.timesince import timesince
from django.utils.translation import ugettext, ugettext_lazy as _

from mezzanine.core.fields import RichTextField
from mezzanine.core.managers import DisplayableManager, CurrentSiteManager
from mezzanine.generic.fields import KeywordsField
from mezzanine.utils.html import TagCloser
from mezzanine.utils.models import base_concrete_model
from mezzanine.utils.sites import current_site_id
from mezzanine.utils.timezone import now
from mezzanine.utils.urls import slugify

class SiteRelated(models.Model):
    """
    Abstract model for all things site-related. Adds a foreignkey to
    Django's ``Site`` model, and filters by site with all querysets.
    See ``mezzanine.utils.sites.current_site_id`` for implementation
    details.
    """

    objects = CurrentSiteManager()

    class Meta:
        abstract = True

    site = models.ForeignKey("sites.Site", editable=False)

    def save(self, update_site=False, *args, **kwargs):
        """
        Set the site to the current site when the record is first
        created, or the ``update_site`` argument is explicitly set
        to ``True``.
        """
        if update_site or not self.id:
            self.site_id = current_site_id()
        super(SiteRelated, self).save(*args, **kwargs)

class Slugged(SiteRelated):
    """
    Abstract model that handles auto-generating slugs. Each slugged
    object is also affiliated with a specific site object.
    """

    title = models.CharField(_("Title"), max_length=500)
    slug = models.CharField(_("URL"), max_length=2000, blank=True, null=True)

    class Meta:
        abstract = True
        ordering = ("title",)

    def __unicode__(self):
        return self.title

    def save(self, *args, **kwargs):
        """
        Create a unique slug by appending an index.
        """
        if not self.slug:
            # For custom content types, use the ``Page`` instance for
            # slug lookup.
            concrete_model = base_concrete_model(Slugged, self)
            self.slug = self.get_slug()
            i = 0
            while True:
                if i > 0:
                    if i > 1:
                        self.slug = self.slug.rsplit("-", 1)[0]
                    self.slug = "%s-%s" % (self.slug, i)
                qs = concrete_model.objects.all()
                if self.id is not None:
                    qs = qs.exclude(id=self.id)
                try:
                    qs.get(slug=self.slug)
                except ObjectDoesNotExist:
                    break
                i += 1
        super(Slugged, self).save(*args, **kwargs)

    def get_slug(self):
        """
        Allows subclasses to implement their own slug creation logic.
        """
        return slugify(self)

    def admin_link(self):
        return "%s" % (self.get_absolute_url(),
                                        ugettext("View on site"))
    admin_link.allow_tags = True
    admin_link.short_description = ""

class MetaData(models.Model):
    """
    Abstract model that provides meta data for content.
    """

    description = models.TextField(_("Description"), blank=True)
    gen_description = models.BooleanField(_("Generate description"),
        help_text=_("If checked, the description will be automatically "
                    "generated from content. Uncheck if you want to manually "
                    "set a custom description."), default=True)
    keywords = KeywordsField(verbose_name=_("Keywords"))

    class Meta:
        abstract = True

    def save(self, *args, **kwargs):
        """
        Set the description field on save.
        """
        if self.gen_description:
            self.description = strip_tags(self.description_from_content())
        super(MetaData, self).save(*args, **kwargs)

    def description_from_content(self):
        """
        Returns the first block or sentence of the first content-like
        field.
        """
        description = ""
        # Use the first RichTextField, or TextField if none found.
        for field_type in (RichTextField, models.TextField):
            if not description:
                for field in self._meta.fields:
                    if isinstance(field, field_type) and \
                        field.name != "description":
                        description = getattr(self, field.name)
                        if description:
                            break
        # Fall back to the title if description couldn't be determined.
        if not description:
            description = unicode(self)
        # Strip everything after the first block or sentence.
        ends = ("
", "
", "
", "
", "
",
                "\n", ". ", "! ", "? ")
        for end in ends:
            pos = description.lower().find(end)
            if pos > -1:
                description = TagCloser(description[:pos]).html
                break
        else:
            description = truncatewords_html(description, 100)
        return description

CONTENT_STATUS_DRAFT = 1
CONTENT_STATUS_PUBLISHED = 2
CONTENT_STATUS_CHOICES = (
    (CONTENT_STATUS_DRAFT, _("Draft")),
    (CONTENT_STATUS_PUBLISHED, _("Published")),
)

class Displayable(Slugged, MetaData):
    """
    Abstract model that provides features of a visible page on the
    website such as publishing fields. Basis of Mezzanine pages,
    blog posts, and Cartridge products.
    """

    status = models.IntegerField(_("Status"),
        choices=CONTENT_STATUS_CHOICES, default=CONTENT_STATUS_PUBLISHED)
    publish_date = models.DateTimeField(_("Published from"),
        help_text=_("With published checked, won't be shown until this time"),
        blank=True, null=True)
    expiry_date = models.DateTimeField(_("Expires on"),
        help_text=_("With published checked, won't be shown after this time"),
        blank=True, null=True)
    short_url = models.URLField(blank=True, null=True)

    objects = DisplayableManager()
    search_fields = {"keywords": 10, "title": 5}

    class Meta:
        abstract = True

    def save(self, *args, **kwargs):
        """
        Set default for ``publish_date``. We can't use ``auto_add`` on
        the field as it will be blank when a blog post is created from
        the quick blog form in the admin dashboard.
        """
        if self.publish_date is None:
            self.publish_date = now()
        super(Displayable, self).save(*args, **kwargs)

    def publish_date_since(self):
        """
        Returns the time since ``publish_date``.
        """
        return timesince(self.publish_date)
    publish_date_since.short_description = _("Published from")

    def get_absolute_url(self):
        """
        Raise an error if called on a subclass without
        ``get_absolute_url`` defined, to ensure all search results
        contains a URL.
        """
        name = self.__class__.__name__
        raise NotImplementedError("The model %s does not have "
                                  "get_absolute_url defined" % name)

class RichText(models.Model):
    """
    Provides a Rich Text field for managing general content and making
    it searchable.
    """

    content = RichTextField(_("Content"))

    search_fields = ("content",)

    class Meta:
        abstract = True

class OrderableBase(ModelBase):
    """
    Checks for ``order_with_respect_to`` on the model's inner ``Meta``
    class and if found, copies it to a custom attribute and deletes it
    since it will cause errors when used with ``ForeignKey("self")``.
    Also creates the ``ordering`` attribute on the ``Meta`` class if
    not yet provided.
    """

    def __new__(cls, name, bases, attrs):
        if "Meta" not in attrs:
            class Meta:
                pass
            attrs["Meta"] = Meta
        if hasattr(attrs["Meta"], "order_with_respect_to"):
            order_field = attrs["Meta"].order_with_respect_to
            attrs["order_with_respect_to"] = order_field
            del attrs["Meta"].order_with_respect_to
        if not hasattr(attrs["Meta"], "ordering"):
            setattr(attrs["Meta"], "ordering", ("_order",))
        return super(OrderableBase, cls).__new__(cls, name, bases, attrs)

class Orderable(models.Model):
    """
    Abstract model that provides a custom ordering integer field
    similar to using Meta's ``order_with_respect_to``, since to
    date (Django 1.2) this doesn't work with ``ForeignKey("self")``,
    or with Generic Relations. We may also want this feature for
    models that aren't ordered with respect to a particular field.
    """

    __metaclass__ = OrderableBase

    _order = models.IntegerField(_("Order"), null=True)

    class Meta:
        abstract = True

    def with_respect_to(self):
        """
        Returns a dict to use as a filter for ordering operations
        containing the original ``Meta.order_with_respect_to`` value
        if provided. If the field is a Generic Relation, the dict
        returned contains names and values for looking up the
        relation's ``ct_field`` and ``fk_field`` attributes.
        """
        try:
            name = self.order_with_respect_to
            value = getattr(self, name)
        except AttributeError:
            # No ``order_with_respect_to`` specified on the model.
            return {}
        # Support for generic relations.
        field = getattr(self.__class__, name)
        if isinstance(field, GenericForeignKey):
            names = (field.ct_field, field.fk_field)
            return dict([(name, getattr(self, name)) for name in names])
        return {name: value}

    def save(self, *args, **kwargs):
        """
        Set the initial ordering value.
        """
        if self._order is None:
            lookup = self.with_respect_to()
            lookup["_order__isnull"] = False
            concrete_model = base_concrete_model(Orderable, self)
            self._order = concrete_model.objects.filter(**lookup).count()
        super(Orderable, self).save(*args, **kwargs)

    def delete(self, *args, **kwargs):
        """
        Update the ordering values for siblings.
        """
        lookup = self.with_respect_to()
        lookup["_order__gte"] = self._order
        concrete_model = base_concrete_model(Orderable, self)
        after = concrete_model.objects.filter(**lookup)
        after.update(_order=models.F("_order") - 1)
        super(Orderable, self).delete(*args, **kwargs)

class Ownable(models.Model):
    """
    Abstract model that provides ownership of an object for a user.
    """

    user = models.ForeignKey("auth.User", verbose_name=_("Author"),
        related_name="%(class)ss")

    class Meta:
        abstract = True

    def is_editable(self, request):
        """
        Restrict in-line editing to the objects's owner and superusers.
        """
        return request.user.is_superuser or request.user.id == self.user_id
<commit_msg>Ensure unique slugs even when slug is provided. Closes #290.<commit_after>
from django.contrib.contenttypes.generic import GenericForeignKey
from django.core.exceptions import ObjectDoesNotExist
from django.db import models
from django.db.models.base import ModelBase
from django.template.defaultfilters import truncatewords_html
from django.utils.html import strip_tags
from django.utils.timesince import timesince
from django.utils.translation import ugettext, ugettext_lazy as _

from mezzanine.core.fields import RichTextField
from mezzanine.core.managers import DisplayableManager, CurrentSiteManager
from mezzanine.generic.fields import KeywordsField
from mezzanine.utils.html import TagCloser
from mezzanine.utils.models import base_concrete_model
from mezzanine.utils.sites import current_site_id
from mezzanine.utils.timezone import now
from mezzanine.utils.urls import slugify

class SiteRelated(models.Model):
    """
    Abstract model for all things site-related. Adds a foreignkey to
    Django's ``Site`` model, and filters by site with all querysets.
    See ``mezzanine.utils.sites.current_site_id`` for implementation
    details.
    """

    objects = CurrentSiteManager()

    class Meta:
        abstract = True

    site = models.ForeignKey("sites.Site", editable=False)

    def save(self, update_site=False, *args, **kwargs):
        """
        Set the site to the current site when the record is first
        created, or the ``update_site`` argument is explicitly set
        to ``True``.
        """
        if update_site or not self.id:
            self.site_id = current_site_id()
        super(SiteRelated, self).save(*args, **kwargs)

class Slugged(SiteRelated):
    """
    Abstract model that handles auto-generating slugs. Each slugged
    object is also affiliated with a specific site object.
    """

    title = models.CharField(_("Title"), max_length=500)
    slug = models.CharField(_("URL"), max_length=2000, blank=True, null=True)

    class Meta:
        abstract = True
        ordering = ("title",)

    def __unicode__(self):
        return self.title

    def save(self, *args, **kwargs):
        """
        Create a unique slug by appending an index.
        """
        if not self.slug:
            self.slug = self.get_slug()
        # For custom content types, use the ``Page`` instance for
        # slug lookup.
        concrete_model = base_concrete_model(Slugged, self)
        i = 0
        while True:
            if i > 0:
                if i > 1:
                    self.slug = self.slug.rsplit("-", 1)[0]
                self.slug = "%s-%s" % (self.slug, i)
            qs = concrete_model.objects.all()
            if self.id is not None:
                qs = qs.exclude(id=self.id)
            try:
                qs.get(slug=self.slug)
            except ObjectDoesNotExist:
                break
            i += 1
        super(Slugged, self).save(*args, **kwargs)

    def get_slug(self):
        """
        Allows subclasses to implement their own slug creation logic.
        """
        return slugify(self)

    def admin_link(self):
        return "%s" % (self.get_absolute_url(),
                                        ugettext("View on site"))
    admin_link.allow_tags = True
    admin_link.short_description = ""

class MetaData(models.Model):
    """
    Abstract model that provides meta data for content.
    """

    description = models.TextField(_("Description"), blank=True)
    gen_description = models.BooleanField(_("Generate description"),
        help_text=_("If checked, the description will be automatically "
                    "generated from content. Uncheck if you want to manually "
                    "set a custom description."), default=True)
    keywords = KeywordsField(verbose_name=_("Keywords"))

    class Meta:
        abstract = True

    def save(self, *args, **kwargs):
        """
        Set the description field on save.
        """
        if self.gen_description:
            self.description = strip_tags(self.description_from_content())
        super(MetaData, self).save(*args, **kwargs)

    def description_from_content(self):
        """
        Returns the first block or sentence of the first content-like
        field.
        """
        description = ""
        # Use the first RichTextField, or TextField if none found.
        for field_type in (RichTextField, models.TextField):
            if not description:
                for field in self._meta.fields:
                    if isinstance(field, field_type) and \
                        field.name != "description":
                        description = getattr(self, field.name)
                        if description:
                            break
        # Fall back to the title if description couldn't be determined.
        if not description:
            description = unicode(self)
        # Strip everything after the first block or sentence.
        ends = ("
", "
", "
", "
", "
",
                "\n", ". ", "! ", "? ")
        for end in ends:
            pos = description.lower().find(end)
            if pos > -1:
                description = TagCloser(description[:pos]).html
                break
        else:
            description = truncatewords_html(description, 100)
        return description

CONTENT_STATUS_DRAFT = 1
CONTENT_STATUS_PUBLISHED = 2
CONTENT_STATUS_CHOICES = (
    (CONTENT_STATUS_DRAFT, _("Draft")),
    (CONTENT_STATUS_PUBLISHED, _("Published")),
)

class Displayable(Slugged, MetaData):
    """
    Abstract model that provides features of a visible page on the
    website such as publishing fields. Basis of Mezzanine pages,
    blog posts, and Cartridge products.
    """

    status = models.IntegerField(_("Status"),
        choices=CONTENT_STATUS_CHOICES, default=CONTENT_STATUS_PUBLISHED)
    publish_date = models.DateTimeField(_("Published from"),
        help_text=_("With published checked, won't be shown until this time"),
        blank=True, null=True)
    expiry_date = models.DateTimeField(_("Expires on"),
        help_text=_("With published checked, won't be shown after this time"),
        blank=True, null=True)
    short_url = models.URLField(blank=True, null=True)

    objects = DisplayableManager()
    search_fields = {"keywords": 10, "title": 5}

    class Meta:
        abstract = True

    def save(self, *args, **kwargs):
        """
        Set default for ``publish_date``. We can't use ``auto_add`` on
        the field as it will be blank when a blog post is created from
        the quick blog form in the admin dashboard.
        """
        if self.publish_date is None:
            self.publish_date = now()
        super(Displayable, self).save(*args, **kwargs)

    def publish_date_since(self):
        """
        Returns the time since ``publish_date``.
        """
        return timesince(self.publish_date)
    publish_date_since.short_description = _("Published from")

    def get_absolute_url(self):
        """
        Raise an error if called on a subclass without
        ``get_absolute_url`` defined, to ensure all search results
        contains a URL.
        """
        name = self.__class__.__name__
        raise NotImplementedError("The model %s does not have "
                                  "get_absolute_url defined" % name)

class RichText(models.Model):
    """
    Provides a Rich Text field for managing general content and making
    it searchable.
    """

    content = RichTextField(_("Content"))

    search_fields = ("content",)

    class Meta:
        abstract = True

class OrderableBase(ModelBase):
    """
    Checks for ``order_with_respect_to`` on the model's inner ``Meta``
    class and if found, copies it to a custom attribute and deletes it
    since it will cause errors when used with ``ForeignKey("self")``.
    Also creates the ``ordering`` attribute on the ``Meta`` class if
    not yet provided.
    """

    def __new__(cls, name, bases, attrs):
        if "Meta" not in attrs:
            class Meta:
                pass
            attrs["Meta"] = Meta
        if hasattr(attrs["Meta"], "order_with_respect_to"):
            order_field = attrs["Meta"].order_with_respect_to
            attrs["order_with_respect_to"] = order_field
            del attrs["Meta"].order_with_respect_to
        if not hasattr(attrs["Meta"], "ordering"):
            setattr(attrs["Meta"], "ordering", ("_order",))
        return super(OrderableBase, cls).__new__(cls, name, bases, attrs)

class Orderable(models.Model):
    """
    Abstract model that provides a custom ordering integer field
    similar to using Meta's ``order_with_respect_to``, since to
    date (Django 1.2) this doesn't work with ``ForeignKey("self")``,
    or with Generic Relations. We may also want this feature for
    models that aren't ordered with respect to a particular field.
    """

    __metaclass__ = OrderableBase

    _order = models.IntegerField(_("Order"), null=True)

    class Meta:
        abstract = True

    def with_respect_to(self):
        """
        Returns a dict to use as a filter for ordering operations
        containing the original ``Meta.order_with_respect_to`` value
        if provided. If the field is a Generic Relation, the dict
        returned contains names and values for looking up the
        relation's ``ct_field`` and ``fk_field`` attributes.
        """
        try:
            name = self.order_with_respect_to
            value = getattr(self, name)
        except AttributeError:
            # No ``order_with_respect_to`` specified on the model.
            return {}
        # Support for generic relations.
        field = getattr(self.__class__, name)
        if isinstance(field, GenericForeignKey):
            names = (field.ct_field, field.fk_field)
            return dict([(name, getattr(self, name)) for name in names])
        return {name: value}

    def save(self, *args, **kwargs):
        """
        Set the initial ordering value.
        """
        if self._order is None:
            lookup = self.with_respect_to()
            lookup["_order__isnull"] = False
            concrete_model = base_concrete_model(Orderable, self)
            self._order = concrete_model.objects.filter(**lookup).count()
        super(Orderable, self).save(*args, **kwargs)

    def delete(self, *args, **kwargs):
        """
        Update the ordering values for siblings.
        """
        lookup = self.with_respect_to()
        lookup["_order__gte"] = self._order
        concrete_model = base_concrete_model(Orderable, self)
        after = concrete_model.objects.filter(**lookup)
        after.update(_order=models.F("_order") - 1)
        super(Orderable, self).delete(*args, **kwargs)

class Ownable(models.Model):
    """
    Abstract model that provides ownership of an object for a user.
    """

    user = models.ForeignKey("auth.User", verbose_name=_("Author"),
        related_name="%(class)ss")

    class Meta:
        abstract = True

    def is_editable(self, request):
        """
        Restrict in-line editing to the objects's owner and superusers.
        """
        return request.user.is_superuser or request.user.id == self.user_id
<|endoftext|>


<commit_before>from __future__ import unicode_literals

from django.contrib.auth.tokens import default_token_generator
from django.core.mail import EmailMultiAlternatives
from django.core.urlresolvers import reverse
from django.template import loader, Context
from django.utils.http import int_to_base36

from mezzanine.conf import settings
from mezzanine.utils.urls import admin_url, next_url
from mezzanine.conf.context_processors import settings as context_settings

def split_addresses(email_string_list):
    """
    Converts a string containing comma separated email addresses
    into a list of email addresses.
    """
    return [f for f in [s.strip() for s in email_string_list.split(",")] if f]

def subject_template(template, context):
    """
    Loads and renders an email subject template, returning the
    subject string.
    """
    subject = loader.get_template(template).render(Context(context))
    return " ".join(subject.splitlines()).strip()

def send_mail_template(subject, template, addr_from, addr_to, context=None,
                       attachments=None, fail_silently=False, addr_bcc=None,
                       headers=None):
    """
    Send email rendering text and html versions for the specified
    template name using the context dictionary passed in.
    """
    if context is None:
        context = {}
    if attachments is None:
        attachments = []
    # Add template accessible settings from Mezzanine to the context
    # (normally added by a context processor for HTTP requests)
    context.update(context_settings())
    # Allow for a single address to be passed in.
    if not hasattr(addr_to, "__iter__"):
        addr_to = [addr_to]
    if addr_bcc is not None and not hasattr(addr_bcc, "__iter__"):
        addr_bcc = [addr_bcc]
    # Loads a template passing in vars as context.
    render = lambda type: loader.get_template("%s.%s" %
                          (template, type)).render(Context(context))
    # Create and send email.
    msg = EmailMultiAlternatives(subject, render("txt"),
                                 addr_from, addr_to, addr_bcc,
                                 headers=headers)
    msg.attach_alternative(render("html"), "text/html")
    for attachment in attachments:
        msg.attach(*attachment)
    msg.send(fail_silently=fail_silently)

def send_verification_mail(request, user, verification_type):
    """
    Sends an email with a verification link to users when
    ``ACCOUNTS_VERIFICATION_REQUIRED`` is ```True`` and they're signing
    up, or when they reset a lost password.
    The ``verification_type`` arg is both the name of the urlpattern for
    the verification link, as well as the names of the email templates
    to use.
    """
    verify_url = reverse(verification_type, kwargs={
        "uidb36": int_to_base36(user.id),
        "token": default_token_generator.make_token(user),
    }) + "?next=" + (next_url(request) or "/")
    context = {
        "request": request,
        "user": user,
        "verify_url": verify_url,
    }
    subject_template_name = "email/%s_subject.txt" % verification_type
    subject = subject_template(subject_template_name, context)
    send_mail_template(subject, "email/%s" % verification_type,
                       settings.DEFAULT_FROM_EMAIL, user.email,
                       context=context, fail_silently=settings.DEBUG)

def send_approve_mail(request, user):
    """
    Sends an email to staff in listed in the setting
    ``ACCOUNTS_APPROVAL_EMAILS``, when a new user signs up and the
    ``ACCOUNTS_APPROVAL_REQUIRED`` setting is ``True``.
    """
    settings.use_editable()
    approval_emails = split_addresses(settings.ACCOUNTS_APPROVAL_EMAILS)
    if not approval_emails:
        return
    context = {
        "request": request,
        "user": user,
        "change_url": admin_url(user.__class__, "change", user.id),
    }
    subject = subject_template("email/account_approve_subject.txt", context)
    send_mail_template(subject, "email/account_approve",
                       settings.DEFAULT_FROM_EMAIL, approval_emails,
                       context=context, fail_silently=settings.DEBUG)

def send_approved_mail(request, user):
    """
    Sends an email to a user once their ``is_active`` status goes from
    ``False`` to ``True`` when the ``ACCOUNTS_APPROVAL_REQUIRED``
    setting is ``True``.
    """
    context = {"request": request, "user": user}
    subject = subject_template("email/account_approved_subject.txt", context)
    send_mail_template(subject, "email/account_approved",
                       settings.DEFAULT_FROM_EMAIL, user.email,
                       context=context, fail_silently=settings.DEBUG)
<commit_msg>Change Py3-incompatible is-this-a-string type-check hack in send_mail_template<commit_after>from __future__ import unicode_literals
from future import bytes, isinstance, str

from django.contrib.auth.tokens import default_token_generator
from django.core.mail import EmailMultiAlternatives
from django.core.urlresolvers import reverse
from django.template import loader, Context
from django.utils.http import int_to_base36

from mezzanine.conf import settings
from mezzanine.utils.urls import admin_url, next_url
from mezzanine.conf.context_processors import settings as context_settings

def split_addresses(email_string_list):
    """
    Converts a string containing comma separated email addresses
    into a list of email addresses.
    """
    return [f for f in [s.strip() for s in email_string_list.split(",")] if f]

def subject_template(template, context):
    """
    Loads and renders an email subject template, returning the
    subject string.
    """
    subject = loader.get_template(template).render(Context(context))
    return " ".join(subject.splitlines()).strip()

def send_mail_template(subject, template, addr_from, addr_to, context=None,
                       attachments=None, fail_silently=False, addr_bcc=None,
                       headers=None):
    """
    Send email rendering text and html versions for the specified
    template name using the context dictionary passed in.
    """
    if context is None:
        context = {}
    if attachments is None:
        attachments = []
    # Add template accessible settings from Mezzanine to the context
    # (normally added by a context processor for HTTP requests)
    context.update(context_settings())
    # Allow for a single address to be passed in.
    # Python 3 strings have an __iter__ method, so the following hack
    # doesn't work: if not hasattr(addr_to, "__iter__"):
    if isinstance(addr_to, str) or isinstance(addr_to, bytes):
        addr_to = [addr_to]
    if addr_bcc is not None and (isinstance(addr_bcc, str) or
                                 isinstance(addr_bcc, bytes)):
        addr_bcc = [addr_bcc]
    # Loads a template passing in vars as context.
    render = lambda type: loader.get_template("%s.%s" %
                          (template, type)).render(Context(context))
    # Create and send email.
    msg = EmailMultiAlternatives(subject, render("txt"),
                                 addr_from, addr_to, addr_bcc,
                                 headers=headers)
    msg.attach_alternative(render("html"), "text/html")
    for attachment in attachments:
        msg.attach(*attachment)
    msg.send(fail_silently=fail_silently)

def send_verification_mail(request, user, verification_type):
    """
    Sends an email with a verification link to users when
    ``ACCOUNTS_VERIFICATION_REQUIRED`` is ```True`` and they're signing
    up, or when they reset a lost password.
    The ``verification_type`` arg is both the name of the urlpattern for
    the verification link, as well as the names of the email templates
    to use.
    """
    verify_url = reverse(verification_type, kwargs={
        "uidb36": int_to_base36(user.id),
        "token": default_token_generator.make_token(user),
    }) + "?next=" + (next_url(request) or "/")
    context = {
        "request": request,
        "user": user,
        "verify_url": verify_url,
    }
    subject_template_name = "email/%s_subject.txt" % verification_type
    subject = subject_template(subject_template_name, context)
    send_mail_template(subject, "email/%s" % verification_type,
                       settings.DEFAULT_FROM_EMAIL, user.email,
                       context=context, fail_silently=settings.DEBUG)

def send_approve_mail(request, user):
    """
    Sends an email to staff in listed in the setting
    ``ACCOUNTS_APPROVAL_EMAILS``, when a new user signs up and the
    ``ACCOUNTS_APPROVAL_REQUIRED`` setting is ``True``.
    """
    settings.use_editable()
    approval_emails = split_addresses(settings.ACCOUNTS_APPROVAL_EMAILS)
    if not approval_emails:
        return
    context = {
        "request": request,
        "user": user,
        "change_url": admin_url(user.__class__, "change", user.id),
    }
    subject = subject_template("email/account_approve_subject.txt", context)
    send_mail_template(subject, "email/account_approve",
                       settings.DEFAULT_FROM_EMAIL, approval_emails,
                       context=context, fail_silently=settings.DEBUG)

def send_approved_mail(request, user):
    """
    Sends an email to a user once their ``is_active`` status goes from
    ``False`` to ``True`` when the ``ACCOUNTS_APPROVAL_REQUIRED``
    setting is ``True``.
    """
    context = {"request": request, "user": user}
    subject = subject_template("email/account_approved_subject.txt", context)
    send_mail_template(subject, "email/account_approved",
                       settings.DEFAULT_FROM_EMAIL, user.email,
                       context=context, fail_silently=settings.DEBUG)
<|endoftext|>


<commit_before>"""
Ebook text

"""
from ui import gtk
import pango

import os
import gzip
import bz2
import zipfile
import textwrap
import re

from HTMLParser import HTMLParser
import htmlentitydefs

import plucker

class EbookText(gtk.TextBuffer):
    def __init__(self, filename):
        gtk.TextBuffer.__init__(self)

        self._text = None
        self._error = None
        self._load(filename)

        if self._text:
            self.set_text(self._text)

    @property
    def loaded(self):
        return self._text is not None

    @property
    def error(self):
        if not self._error:
            return None
        return str(self._error)

    def _load(self, filename):
        try:
            if filename.endswith('.gz'):
                f = gzip.open(filename, 'rb')
                filename = filename[:-4]
            elif filename.endswith('.bz2'):
                f = bz2.BZ2File(filename, 'rb')
                filename = filename[:-4]
            elif filename.endswith('.zip'):
                zf = zipfile.ZipFile(filename, 'rb')
                names = zf.namelist()
                if len(names) != 1:
                    raise IOError("Zip file does not contain a single file")
                f = zf.open(names[0], 'rb')
                filename = names[0]
            elif filename.endswith('.pdb'):
                f = plucker.PluckerFile(filename)
                # XXX: not implemented
                raise IOError("Plucker supprot not implemented yet.")
            elif filename.endswith('.txt'):
                f = open(filename, 'rb')
            else:
                raise IOError("Don't know how to open this type of files")

            self._load_stream(filename, f)
            f.close()
        except IOError, e:
            self._text = None
            self._error = e

    def _load_stream(self, filename, f):
        if filename.endswith('.html') or filename.endswith('.htm'):
            self._load_html(f)
        else:
            self._load_plain_text(f)

    def _load_html(self, f):
        raw_text = f.read()
        f.close()
        for encoding in detect_encoding(raw_text):
            try:
                raw_text = unicode(raw_text, encoding)
                break
            except UnicodeError:
                pass

        tag_bold = self.create_tag("bold", weight=pango.WEIGHT_BOLD)
        tag_emph = self.create_tag("emph", style=pango.STYLE_ITALIC)
        tag_big = self.create_tag("big", scale=1.5)

        parent = self

        class HandleHTML(HTMLParser):
            tags = []
            in_body = False
            encoding = 'latin1'
            para = u""
            omit = 0
            slurp_space = True

            def handle_starttag(self, tag, attrs):
                self.flush()
                if tag in ('h1', 'h2', 'h3', 'h4'):
                    self._append(u'\n\n')
                    self.tags.append(tag_big)
                elif tag == 'p' or tag == 'br' or tag == 'div':
                    self.tags = []
                    self._append(u'\n')
                    self.slurp_space = True
                elif tag == 'tr':
                    self.tags = []
                    self._append(u'\n')
                    self.slurp_space = True
                elif tag == 'td':
                    self._append(u'\t')
                elif tag == 'i' or tag == 'em':
                    self.tags.append(tag_emph)
                elif tag == 'b' or tag == 'strong' or tag == 'bold':
                    self.tags.append(tag_bold)
                elif tag == 'hr':
                    self._append(u'\n')
                elif tag == 'body':
                    self.in_body = True
                elif tag == 'style':
                    self.omit += 1
                elif tag == 'meta' and not self.in_body:
                    self.handle_meta(attrs)
                elif tag == 'img':
                    attrs = dict(attrs)
                    if 'alt' in attrs and attrs['alt'].strip():
                        self.handle_data('[IMAGE: %s]' % attrs['alt'].strip())
                    else:
                        self.handle_data('[IMAGE]')

            def handle_meta(self, attrs):
                attrs = dict(attrs)
                if not attrs.get('http-equiv', "").lower() == 'content-type':
                    return
                content = attrs.get('content', "")

                m = re.search(r'charset\s*=\s*([a-zA-Z0-9-]+)', content)
                if m:
                    encoding = m.group(1).lower()
                    try:
                        unicode('', encoding)
                        self.encoding = encoding
                    except UnicodeError:
                        pass

            def handle_endtag(self, tag):
                self.flush()
                if tag == 'style':
                    self.omit -= 1
                elif tag in ('h1', 'h2', 'h3', 'h4'):
                    if self.tags:
                        self.tags.pop()
                    self._append('\n')
                elif tag in ('i', 'em', 'b', 'strong', 'bold'):
                    if self.tags:
                        self.tags.pop()

            def handle_data(self, data):
                if self.omit:
                    return
                data = data.replace('\r', '')
                data = data.replace('\n', ' ')
                if self.slurp_space:
                    data = data.lstrip()
                self.para += data
                self.slurp_space = False

            def handle_charref(self, name):
                try:
                    self.para += unicode(chr(int(name)), 'latin1')
                except (UnicodeError, ValueError):
                    self.para += '?'

            def handle_entityref(self, name):
                self.para += u'%c' % htmlentitydefs.name2codepoint.get(name, 63)

            def _append(self, text):
                parent.insert_with_tags(parent.get_end_iter(),
                                        text.encode('utf-8'), *self.tags)

            def flush(self):
                if self.para:
                    self._append(self.para)
                    self.para = u""

        html = HandleHTML()
        html.feed(raw_text)
        html.flush()

    def _load_plain_text(self, f):
        raw_text = f.read()

        for encoding in detect_encoding(raw_text):
            try:
                self._text = unicode(raw_text, encoding)
                break
            except UnicodeError:
                pass

        self._text = rewrap(self._text)

def detect_encoding(text):
    encodings = ['utf-8']
    if '\x92' in text:
        encodings.append('windows-1250')
    else:
        encodings.append('latin1')
    return encodings

def rewrap(text):
    if not text:
        return

    # remove lines
    text = re.sub(r'-{10,}', '', text)
    text = re.sub(r'={10,}', '', text)

    text = text.replace(u'\r', u'')
    text = textwrap.dedent(text)

    sample = text[:80*1000]
    if len([x for x in sample.split('\n') if x.startswith(' ')]) > 20:
        # Paragraphs separated by indent
        text = re.sub(u'\n(?!\\s)', u' ', text)
    elif max(map(len, sample.split("\n"))) < 100 and sample.count('\n\n') > 5:
        # Paragraphs separated by empty lines
        text = re.sub(u'\n(?!\n)', u' ', text)
    else:
        # Paragraphs on a single line -- or couldn't determine formatting
        pass

    text = re.sub(u'\s{10,}', u'\n', text)
    text = re.sub(u'\n[ \t]+', u'\n', text)
    return text
<commit_msg>Fix file format detection<commit_after>"""
Ebook text

"""
from ui import gtk
import pango

import os
import gzip
import bz2
import zipfile
import textwrap
import re

from HTMLParser import HTMLParser
import htmlentitydefs

import plucker

class EbookText(gtk.TextBuffer):
    def __init__(self, filename):
        gtk.TextBuffer.__init__(self)

        self._text = None
        self._error = None
        self._load(filename)

        if self._text:
            self.set_text(self._text)

    @property
    def loaded(self):
        return self._text is not None

    @property
    def error(self):
        if not self._error:
            return None
        return str(self._error)

    def _load(self, filename):
        try:
            basefn, ext = os.path.splitext(filename)
            if ext == '.gz':
                f = gzip.open(filename, 'rb')
                filename = basefn
            elif ext == '.bz2':
                f = bz2.BZ2File(filename, 'rb')
                filename = basefn
            elif ext == '.zip':
                zf = zipfile.ZipFile(filename, 'rb')
                names = zf.namelist()
                if len(names) != 1:
                    raise IOError("Zip file does not contain a single file")
                f = zf.open(names[0], 'rb')
                filename = names[0]
            elif ext == '.pdb':
                f = plucker.PluckerFile(filename)
                # XXX: not implemented
                raise IOError("Plucker supprot not implemented yet.")
            else:
                f = open(filename, 'rb')

            self._load_stream(filename, f)
            f.close()
        except IOError, e:
            self._text = None
            self._error = e

    def _load_stream(self, filename, f):
        basefn, ext = os.path.splitext(filename)
        if ext in ('.html', '.htm'):
            self._load_html(f)
        elif ext in ('.txt', '.rst'):
            self._load_plain_text(f)
        else:
            raise IOError("Don't know how to open this type of files")

    def _load_html(self, f):
        raw_text = f.read()
        f.close()
        for encoding in detect_encoding(raw_text):
            try:
                raw_text = unicode(raw_text, encoding)
                break
            except UnicodeError:
                pass

        tag_bold = self.create_tag("bold", weight=pango.WEIGHT_BOLD)
        tag_emph = self.create_tag("emph", style=pango.STYLE_ITALIC)
        tag_big = self.create_tag("big", scale=1.5)

        parent = self

        class HandleHTML(HTMLParser):
            tags = []
            in_body = False
            encoding = 'latin1'
            para = u""
            omit = 0
            slurp_space = True

            def handle_starttag(self, tag, attrs):
                self.flush()
                if tag in ('h1', 'h2', 'h3', 'h4'):
                    self._append(u'\n\n')
                    self.tags.append(tag_big)
                elif tag == 'p' or tag == 'br' or tag == 'div':
                    self.tags = []
                    self._append(u'\n')
                    self.slurp_space = True
                elif tag == 'tr':
                    self.tags = []
                    self._append(u'\n')
                    self.slurp_space = True
                elif tag == 'td':
                    self._append(u'\t')
                elif tag == 'i' or tag == 'em':
                    self.tags.append(tag_emph)
                elif tag == 'b' or tag == 'strong' or tag == 'bold':
                    self.tags.append(tag_bold)
                elif tag == 'hr':
                    self._append(u'\n')
                elif tag == 'body':
                    self.in_body = True
                elif tag == 'style':
                    self.omit += 1
                elif tag == 'meta' and not self.in_body:
                    self.handle_meta(attrs)
                elif tag == 'img':
                    attrs = dict(attrs)
                    if 'alt' in attrs and attrs['alt'].strip():
                        self.handle_data('[IMAGE: %s]' % attrs['alt'].strip())
                    else:
                        self.handle_data('[IMAGE]')

            def handle_meta(self, attrs):
                attrs = dict(attrs)
                if not attrs.get('http-equiv', "").lower() == 'content-type':
                    return
                content = attrs.get('content', "")

                m = re.search(r'charset\s*=\s*([a-zA-Z0-9-]+)', content)
                if m:
                    encoding = m.group(1).lower()
                    try:
                        unicode('', encoding)
                        self.encoding = encoding
                    except UnicodeError:
                        pass

            def handle_endtag(self, tag):
                self.flush()
                if tag == 'style':
                    self.omit -= 1
                elif tag in ('h1', 'h2', 'h3', 'h4'):
                    if self.tags:
                        self.tags.pop()
                    self._append('\n')
                elif tag in ('i', 'em', 'b', 'strong', 'bold'):
                    if self.tags:
                        self.tags.pop()

            def handle_data(self, data):
                if self.omit:
                    return
                data = data.replace('\r', '')
                data = data.replace('\n', ' ')
                if self.slurp_space:
                    data = data.lstrip()
                self.para += data
                self.slurp_space = False

            def handle_charref(self, name):
                try:
                    self.para += unicode(chr(int(name)), 'latin1')
                except (UnicodeError, ValueError):
                    self.para += '?'

            def handle_entityref(self, name):
                self.para += u'%c' % htmlentitydefs.name2codepoint.get(name, 63)

            def _append(self, text):
                parent.insert_with_tags(parent.get_end_iter(),
                                        text.encode('utf-8'), *self.tags)

            def flush(self):
                if self.para:
                    self._append(self.para)
                    self.para = u""

        html = HandleHTML()
        html.feed(raw_text)
        html.flush()

    def _load_plain_text(self, f):
        raw_text = f.read()

        for encoding in detect_encoding(raw_text):
            try:
                self._text = unicode(raw_text, encoding)
                break
            except UnicodeError:
                pass

        self._text = rewrap(self._text)

def detect_encoding(text):
    encodings = ['utf-8']
    if '\x92' in text:
        encodings.append('windows-1250')
    else:
        encodings.append('latin1')
    return encodings

def rewrap(text):
    if not text:
        return

    # remove lines
    text = re.sub(r'-{10,}', '', text)
    text = re.sub(r'={10,}', '', text)

    text = text.replace(u'\r', u'')
    text = textwrap.dedent(text)

    sample = text[:80*1000]
    if len([x for x in sample.split('\n') if x.startswith(' ')]) > 20:
        # Paragraphs separated by indent
        text = re.sub(u'\n(?!\\s)', u' ', text)
    elif max(map(len, sample.split("\n"))) < 100 and sample.count('\n\n') > 5:
        # Paragraphs separated by empty lines
        text = re.sub(u'\n(?!\n)', u' ', text)
    else:
        # Paragraphs on a single line -- or couldn't determine formatting
        pass

    text = re.sub(u'\s{10,}', u'\n', text)
    text = re.sub(u'\n[ \t]+', u'\n', text)
    return text
<|endoftext|>


<commit_before>MAJOR = 0
MINOR = 18
PATCH = 3

def get_version(suffix=''):
    return '.'.join([str(v) for v in (MAJOR, MINOR, PATCH)]) + suffix

if __name__ == '__main__':
    print(get_version())
<commit_msg>version 0.18.4<commit_after>MAJOR = 0
MINOR = 18
PATCH = 4

def get_version(suffix=''):
    return '.'.join([str(v) for v in (MAJOR, MINOR, PATCH)]) + suffix

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


<commit_before># $Id$
#
# Copyright (C) 2006-2011 Greg Landrum
#  All Rights Reserved
#
import os
from rdkit.Chem.Draw.MolDrawing import MolDrawing,DrawingOptions

def _getCanvas():
  useAGG=False
  useCairo=False
  Canvas=None
  if not os.environ.get('RDKIT_CANVAS',''):
    try:
      from rdkit.Chem.Draw.cairoCanvas import Canvas
      useCairo=True
    except ImportError:
      try:
        from rdkit.Chem.Draw.aggCanvas import Canvas
        useAGG=True
      except ImportError:
        from rdkit.Chem.Draw.spingCanvas import Canvas
  else:
    canv=os.environ['RDKIT_CANVAS'].lower()
    if canv =='cairo':
      from rdkit.Chem.Draw.cairoCanvas import Canvas
      useCairo=True
    elif canv =='agg':
      from rdkit.Chem.Draw.aggCanvas import Canvas
      useAGG=True
    else:
      DrawingOptions.radicalSymbol='.' #<- the sping canvas doesn't support unicode well
      from rdkit.Chem.Draw.spingCanvas import Canvas
  return useAGG,useCairo,Canvas

def _createCanvas(size):
  useAGG,useCairo,Canvas=_getCanvas()
  if useAGG or useCairo:
    try:
      import Image
    except ImportError:
      from PIL import Image
    img = Image.new("RGBA",size,(0,0,0,0))
    canvas = Canvas(img)
  else:
    from rdkit.Chem.Draw.spingCanvas import Canvas
    canvas = Canvas(size=size,name='MolToImageFile')
    img = canvas._image
  return img,canvas

def MolToImage(mol, size=(300,300), kekulize=True, wedgeBonds=True,
               fitImage=False, options=None, canvas=None, **kwargs):
  """ returns a PIL image containing a drawing of the molecule

    Keyword arguments:
    kekulize -- run kekulization routine on input `mol` (default True)
    size -- final image size, in pixel (default (300,300))
    wedgeBonds -- draw wedge (stereo) bonds (default True)
    highlightAtoms -- list of atoms to highlight (default [])
    highlightMap -- dictionary of (atom, color) pairs (default None)
    highlightBonds -- list of bonds to highlight (default [])
  """
  if not mol:
    raise ValueError('Null molecule provided')
  if canvas is None:
    img,canvas=_createCanvas(size)
  else:
    img=None

  if options is None:
    options = DrawingOptions()
  if fitImage:
      options.dotsPerAngstrom = int(min(size) / 10)
  options.wedgeDashedBonds = wedgeBonds
  drawer = MolDrawing(canvas=canvas,drawingOptions=options)

  if kekulize:
    from rdkit import Chem
    mol = Chem.Mol(mol.ToBinary())
    Chem.Kekulize(mol)

  if not mol.GetNumConformers():
    from rdkit.Chem import AllChem
    AllChem.Compute2DCoords(mol)

  if 'legend' in kwargs:
    legend = kwargs['legend']
    del kwargs['legend']
  else:
    legend=''

  drawer.AddMol(mol,**kwargs)

  if legend:
    from rdkit.Chem.Draw.MolDrawing import Font
    bbox = drawer.boundingBoxes[mol]
    pos = size[0]/2,int(.94*size[1]),0 # the 0.94 is extremely empirical
    # canvas.addCanvasPolygon(((bbox[0],bbox[1]),(bbox[2],bbox[1]),(bbox[2],bbox[3]),(bbox[0],bbox[3])),
    #                         color=(1,0,0),fill=False,stroke=True)
    # canvas.addCanvasPolygon(((0,0),(0,size[1]),(size[0],size[1]),(size[0],0)   ),
    #                         color=(0,0,1),fill=False,stroke=True)
    font=Font(face='sans',size=12)
    canvas.addCanvasText(legend,pos,font)

  if kwargs.get('returnCanvas',False):
    return img,canvas,drawer
  else:
    canvas.flush()
    return img

def MolToFile(mol,fileName,size=(300,300),kekulize=True, wedgeBonds=True,
              imageType=None, fitImage=False, options=None, **kwargs):
  """ Generates a drawing of a molecule and writes it to a file
  """
  # original contribution from Uwe Hoffmann
  if not fileName:
    raise ValueError('no fileName provided')
  if not mol:
    raise ValueError('Null molecule provided')

  if imageType is None:
    imageType=os.path.splitext(fileName)[1][1:]

  if options is None:
    options = DrawingOptions()
  useAGG,useCairo,Canvas = _getCanvas()
  if fitImage:
      options.dotsPerAngstrom = int(min(size) / 10)
  options.wedgeDashedBonds = wedgeBonds
  if useCairo or useAGG:
    canvas = Canvas(size=size,imageType=imageType,
                              fileName=fileName)
  else:
    options.radicalSymbol = '.' #<- the sping canvas doesn't support unicode well
    canvas = Canvas(size=size,name=fileName,imageType=imageType)
  drawer = MolDrawing(canvas=canvas,drawingOptions=options)
  if kekulize:
    from rdkit import Chem
    mol = Chem.Mol(mol.ToBinary())
    Chem.Kekulize(mol)

  if not mol.GetNumConformers():
    from rdkit.Chem import AllChem
    AllChem.Compute2DCoords(mol)

  drawer.AddMol(mol,**kwargs)
  if useCairo or useAGG:
    canvas.flush()
  else:
    canvas.save()

def MolToImageFile(mol,filename,size=(300,300),kekulize=True, wedgeBonds=True,
                   **kwargs):
  """  DEPRECATED:  please use MolToFile instead

  """
  img = MolToImage(mol,size=size,kekulize=kekulize,wedgeBonds=wedgeBonds,**kwargs)
  img.save(filename)

tkRoot=None
tkLabel=None
tkPI=None
def ShowMol(mol,size=(300,300),kekulize=True,wedgeBonds=True,
            title='RDKit Molecule',**kwargs):
  """ Generates a picture of a molecule and displays it in a Tkinter window
  """
  global tkRoot,tkLabel,tkPI
  import Tkinter
  try:
    import ImageTk
  except ImportError:
    from PIL import ImageTk

  img = MolToImage(mol,size,kekulize,wedgeBonds,**kwargs)

  if not tkRoot:
    tkRoot = Tkinter.Tk()
    tkRoot.title(title)
    tkPI = ImageTk.PhotoImage(img)
    tkLabel = Tkinter.Label(tkRoot,image=tkPI)
    tkLabel.place(x=0,y=0,width=img.size[0],height=img.size[1])
  else:
    tkPI.paste(img)
  tkRoot.geometry('%dx%d'%(img.size))

def MolToMPL(mol,size=(300,300),kekulize=True, wedgeBonds=True,
             imageType=None, fitImage=False, options=None, **kwargs):
  """ Generates a drawing of a molecule on a matplotlib canvas
  """
  if not mol:
    raise ValueError('Null molecule provided')
  from rdkit.Chem.Draw.mplCanvas import Canvas
  canvas = Canvas(size)
  if options is None:
    options = DrawingOptions()
    options.bgColor=None
  if fitImage:
      drawingOptions.dotsPerAngstrom = int(min(size) / 10)
  options.wedgeDashedBonds=wedgeBonds
  drawer = MolDrawing(canvas=canvas, drawingOptions=options)
  omol=mol
  if kekulize:
    from rdkit import Chem
    mol = Chem.Mol(mol.ToBinary())
    Chem.Kekulize(mol)

  if not mol.GetNumConformers():
    from rdkit.Chem import AllChem
    AllChem.Compute2DCoords(mol)

  drawer.AddMol(mol,**kwargs)
  omol._atomPs=drawer.atomPs[mol]
  for k,v in omol._atomPs.iteritems():
    omol._atomPs[k]=canvas.rescalePt(v)
  canvas._figure.set_size_inches(float(size[0])/100,float(size[1])/100)
  return canvas._figure

def calcAtomGaussians(mol,a=0.03,step=0.02,weights=None):
  """
useful things to do with these:
fig.axes[0].imshow(z,cmap=cm.gray,interpolation='bilinear',origin='lower',extent=(0,1,0,1))
fig.axes[0].contour(x,y,z,20,colors='k')

fig=Draw.MolToMPL(m);
contribs=Crippen.rdMolDescriptors._CalcCrippenContribs(m)
logps,mrs=zip(*contribs)
x,y,z=Draw.calcAtomGaussians(m,0.03,step=0.01,weights=logps)
fig.axes[0].imshow(z,cmap=cm.jet,interpolation='bilinear',origin='lower',extent=(0,1,0,1))
fig.axes[0].contour(x,y,z,20,colors='k',alpha=0.5)
fig.savefig('coumlogps.colored.png',bbox_inches='tight')

  """
  import numpy
  from matplotlib import mlab
  x = numpy.arange(0,1,step)
  y = numpy.arange(0,1,step)
  X,Y = numpy.meshgrid(x,y)
  if weights is None:
    weights=[1.]*mol.GetNumAtoms()
  Z = mlab.bivariate_normal(X,Y,a,a,mol._atomPs[0][0], mol._atomPs[0][1])*weights[0]
  for i in range(1,mol.GetNumAtoms()):
    Zp = mlab.bivariate_normal(X,Y,a,a,mol._atomPs[i][0], mol._atomPs[i][1])
    Z += Zp*weights[i]
  return X,Y,Z

def MolsToImage(mols, subImgSize=(200,200),legends=None,**kwargs):
  """
  """
  try:
    import Image
  except ImportError:
    from PIL import Image
  if legends is None: legends = [None]*len(mols)
  res = Image.new("RGBA",(subImgSize[0]*len(mols),subImgSize[1]))
  for i,mol in enumerate(mols):
    res.paste(MolToImage(mol,subImgSize,legend=legends[i],**kwargs),(i*subImgSize[0],0))
  return res

def MolsToGridImage(mols,molsPerRow=3,subImgSize=(200,200),legends=None,
                    highlightAtomLists=None,**kwargs):
  """
  """
  try:
    import Image
  except ImportError:
    from PIL import Image
  if legends is None: legends = [None]*len(mols)

  nRows = len(mols)//molsPerRow
  if len(mols)%molsPerRow : nRows+=1

  res = Image.new("RGBA",(molsPerRow*subImgSize[0],nRows*subImgSize[1]),(255,255,255,0))
  for i,mol in enumerate(mols):
    row = i//molsPerRow
    col = i%molsPerRow
    highlights=None
    if highlightAtomLists and highlightAtomLists[i]:
      highlights=highlightAtomLists[i]
    res.paste(MolToImage(mol,subImgSize,legend=legends[i],highlightAtoms=highlights,
                         **kwargs),(col*subImgSize[0],row*subImgSize[1]))
  return res

def ReactionToImage(rxn, subImgSize=(200,200),**kwargs):
  """
  """
  try:
    import Image
  except ImportError:
    from PIL import Image

  mols = []
  for i in range(rxn.GetNumReactantTemplates()):
    tmpl=rxn.GetReactantTemplate(i)
    tmpl.UpdatePropertyCache(False)
    mols.append(tmpl)
  mols.append(None)
  for i in range(rxn.GetNumProductTemplates()):
    tmpl = rxn.GetProductTemplate(i)
    tmpl.UpdatePropertyCache(False)
    mols.append(tmpl)

  res = Image.new("RGBA",(subImgSize[0]*len(mols),subImgSize[1]),(255,255,255,0))
  for i,mol in enumerate(mols):
    if mol is not None:
      nimg = MolToImage(mol,subImgSize,kekulize=False,**kwargs)
    else:
      nimg,canvas = _createCanvas(subImgSize)
      p0 = (10,subImgSize[1]//2)
      p1 = (subImgSize[0]-10,subImgSize[1]//2)
      p3 = (subImgSize[0]-20,subImgSize[1]//2-10)
      p4 = (subImgSize[0]-20,subImgSize[1]//2+10)
      canvas.addCanvasLine(p0,p1,lineWidth=2,color=(0,0,0))
      canvas.addCanvasLine(p3,p1,lineWidth=2,color=(0,0,0))
      canvas.addCanvasLine(p4,p1,lineWidth=2,color=(0,0,0))
      if hasattr(canvas,'flush'):
        canvas.flush()
      else:
        canvas.save()
    res.paste(nimg,(i*subImgSize[0],0))
  return res

<commit_msg>Custom color to highlight atoms in Mol2Image<commit_after># $Id$
#
# Copyright (C) 2006-2011 Greg Landrum
#  All Rights Reserved
#
import os
from rdkit.Chem.Draw.MolDrawing import MolDrawing,DrawingOptions

def _getCanvas():
  useAGG=False
  useCairo=False
  Canvas=None
  if not os.environ.get('RDKIT_CANVAS',''):
    try:
      from rdkit.Chem.Draw.cairoCanvas import Canvas
      useCairo=True
    except ImportError:
      try:
        from rdkit.Chem.Draw.aggCanvas import Canvas
        useAGG=True
      except ImportError:
        from rdkit.Chem.Draw.spingCanvas import Canvas
  else:
    canv=os.environ['RDKIT_CANVAS'].lower()
    if canv =='cairo':
      from rdkit.Chem.Draw.cairoCanvas import Canvas
      useCairo=True
    elif canv =='agg':
      from rdkit.Chem.Draw.aggCanvas import Canvas
      useAGG=True
    else:
      DrawingOptions.radicalSymbol='.' #<- the sping canvas doesn't support unicode well
      from rdkit.Chem.Draw.spingCanvas import Canvas
  return useAGG,useCairo,Canvas

def _createCanvas(size):
  useAGG,useCairo,Canvas=_getCanvas()
  if useAGG or useCairo:
    try:
      import Image
    except ImportError:
      from PIL import Image
    img = Image.new("RGBA",size,(0,0,0,0))
    canvas = Canvas(img)
  else:
    from rdkit.Chem.Draw.spingCanvas import Canvas
    canvas = Canvas(size=size,name='MolToImageFile')
    img = canvas._image
  return img,canvas

def MolToImage(mol, size=(300,300), kekulize=True, wedgeBonds=True,
               fitImage=False, options=None, canvas=None, **kwargs):
  """Returns a PIL image containing a drawing of the molecule

      ARGUMENTS:

        - kekulize: run kekulization routine on input `mol` (default True)

        - size: final image size, in pixel (default (300,300))

        - wedgeBonds: draw wedge (stereo) bonds (default True)

        - highlightAtoms: list of atoms to highlight (default [])

        - highlightMap: dictionary of (atom, color) pairs (default None)

        - highlightBonds: list of bonds to highlight (default [])

        - highlightColor: RGB color as tuple (default [1, 0, 0])

      NOTE:

            use 'matplotlib.colors.to_rgb()' to convert string and
            HTML color codes into the RGB tuple representation, eg.

              from matplotlib.colors import ColorConverter
              img = Draw.MolToImage(m, highlightAtoms=[1,2], highlightColor=ColorConverter().to_rgb('aqua'))
              img.save("molecule.png")

      RETURNS:

        a PIL Image object
  """

  if not mol:
    raise ValueError('Null molecule provided')
  if canvas is None:
    img,canvas=_createCanvas(size)
  else:
    img=None

  if options is None:
    options = DrawingOptions()
  if fitImage:
      options.dotsPerAngstrom = int(min(size) / 10)
  options.wedgeDashedBonds = wedgeBonds
  if 'highlightColor' in kwargs:
      color = kwargs.pop('highlightColor', (1, 0, 0))
      options.selectColor = color

  drawer = MolDrawing(canvas=canvas,drawingOptions=options)

  if kekulize:
    from rdkit import Chem
    mol = Chem.Mol(mol.ToBinary())
    Chem.Kekulize(mol)

  if not mol.GetNumConformers():
    from rdkit.Chem import AllChem
    AllChem.Compute2DCoords(mol)

  if 'legend' in kwargs:
    legend = kwargs['legend']
    del kwargs['legend']
  else:
    legend=''

  drawer.AddMol(mol,**kwargs)

  if legend:
    from rdkit.Chem.Draw.MolDrawing import Font
    bbox = drawer.boundingBoxes[mol]
    pos = size[0]/2,int(.94*size[1]),0 # the 0.94 is extremely empirical
    # canvas.addCanvasPolygon(((bbox[0],bbox[1]),(bbox[2],bbox[1]),(bbox[2],bbox[3]),(bbox[0],bbox[3])),
    #                         color=(1,0,0),fill=False,stroke=True)
    # canvas.addCanvasPolygon(((0,0),(0,size[1]),(size[0],size[1]),(size[0],0)   ),
    #                         color=(0,0,1),fill=False,stroke=True)
    font=Font(face='sans',size=12)
    canvas.addCanvasText(legend,pos,font)

  if kwargs.get('returnCanvas',False):
    return img,canvas,drawer
  else:
    canvas.flush()
    return img

def MolToFile(mol,fileName,size=(300,300),kekulize=True, wedgeBonds=True,
              imageType=None, fitImage=False, options=None, **kwargs):
  """ Generates a drawing of a molecule and writes it to a file
  """
  # original contribution from Uwe Hoffmann
  if not fileName:
    raise ValueError('no fileName provided')
  if not mol:
    raise ValueError('Null molecule provided')

  if imageType is None:
    imageType=os.path.splitext(fileName)[1][1:]

  if options is None:
    options = DrawingOptions()
  useAGG,useCairo,Canvas = _getCanvas()
  if fitImage:
      options.dotsPerAngstrom = int(min(size) / 10)
  options.wedgeDashedBonds = wedgeBonds
  if useCairo or useAGG:
    canvas = Canvas(size=size,imageType=imageType,
                              fileName=fileName)
  else:
    options.radicalSymbol = '.' #<- the sping canvas doesn't support unicode well
    canvas = Canvas(size=size,name=fileName,imageType=imageType)
  drawer = MolDrawing(canvas=canvas,drawingOptions=options)
  if kekulize:
    from rdkit import Chem
    mol = Chem.Mol(mol.ToBinary())
    Chem.Kekulize(mol)

  if not mol.GetNumConformers():
    from rdkit.Chem import AllChem
    AllChem.Compute2DCoords(mol)

  drawer.AddMol(mol,**kwargs)
  if useCairo or useAGG:
    canvas.flush()
  else:
    canvas.save()

def MolToImageFile(mol,filename,size=(300,300),kekulize=True, wedgeBonds=True,
                   **kwargs):
  """  DEPRECATED:  please use MolToFile instead

  """
  img = MolToImage(mol,size=size,kekulize=kekulize,wedgeBonds=wedgeBonds,**kwargs)
  img.save(filename)

tkRoot=None
tkLabel=None
tkPI=None
def ShowMol(mol,size=(300,300),kekulize=True,wedgeBonds=True,
            title='RDKit Molecule',**kwargs):
  """ Generates a picture of a molecule and displays it in a Tkinter window
  """
  global tkRoot,tkLabel,tkPI
  import Tkinter
  try:
    import ImageTk
  except ImportError:
    from PIL import ImageTk

  img = MolToImage(mol,size,kekulize,wedgeBonds,**kwargs)

  if not tkRoot:
    tkRoot = Tkinter.Tk()
    tkRoot.title(title)
    tkPI = ImageTk.PhotoImage(img)
    tkLabel = Tkinter.Label(tkRoot,image=tkPI)
    tkLabel.place(x=0,y=0,width=img.size[0],height=img.size[1])
  else:
    tkPI.paste(img)
  tkRoot.geometry('%dx%d'%(img.size))

def MolToMPL(mol,size=(300,300),kekulize=True, wedgeBonds=True,
             imageType=None, fitImage=False, options=None, **kwargs):
  """ Generates a drawing of a molecule on a matplotlib canvas
  """
  if not mol:
    raise ValueError('Null molecule provided')
  from rdkit.Chem.Draw.mplCanvas import Canvas
  canvas = Canvas(size)
  if options is None:
    options = DrawingOptions()
    options.bgColor=None
  if fitImage:
      drawingOptions.dotsPerAngstrom = int(min(size) / 10)
  options.wedgeDashedBonds=wedgeBonds
  drawer = MolDrawing(canvas=canvas, drawingOptions=options)
  omol=mol
  if kekulize:
    from rdkit import Chem
    mol = Chem.Mol(mol.ToBinary())
    Chem.Kekulize(mol)

  if not mol.GetNumConformers():
    from rdkit.Chem import AllChem
    AllChem.Compute2DCoords(mol)

  drawer.AddMol(mol,**kwargs)
  omol._atomPs=drawer.atomPs[mol]
  for k,v in omol._atomPs.iteritems():
    omol._atomPs[k]=canvas.rescalePt(v)
  canvas._figure.set_size_inches(float(size[0])/100,float(size[1])/100)
  return canvas._figure

def calcAtomGaussians(mol,a=0.03,step=0.02,weights=None):
  """
useful things to do with these:
fig.axes[0].imshow(z,cmap=cm.gray,interpolation='bilinear',origin='lower',extent=(0,1,0,1))
fig.axes[0].contour(x,y,z,20,colors='k')

fig=Draw.MolToMPL(m);
contribs=Crippen.rdMolDescriptors._CalcCrippenContribs(m)
logps,mrs=zip(*contribs)
x,y,z=Draw.calcAtomGaussians(m,0.03,step=0.01,weights=logps)
fig.axes[0].imshow(z,cmap=cm.jet,interpolation='bilinear',origin='lower',extent=(0,1,0,1))
fig.axes[0].contour(x,y,z,20,colors='k',alpha=0.5)
fig.savefig('coumlogps.colored.png',bbox_inches='tight')

  """
  import numpy
  from matplotlib import mlab
  x = numpy.arange(0,1,step)
  y = numpy.arange(0,1,step)
  X,Y = numpy.meshgrid(x,y)
  if weights is None:
    weights=[1.]*mol.GetNumAtoms()
  Z = mlab.bivariate_normal(X,Y,a,a,mol._atomPs[0][0], mol._atomPs[0][1])*weights[0]
  for i in range(1,mol.GetNumAtoms()):
    Zp = mlab.bivariate_normal(X,Y,a,a,mol._atomPs[i][0], mol._atomPs[i][1])
    Z += Zp*weights[i]
  return X,Y,Z

def MolsToImage(mols, subImgSize=(200,200),legends=None,**kwargs):
  """
  """
  try:
    import Image
  except ImportError:
    from PIL import Image
  if legends is None: legends = [None]*len(mols)
  res = Image.new("RGBA",(subImgSize[0]*len(mols),subImgSize[1]))
  for i,mol in enumerate(mols):
    res.paste(MolToImage(mol,subImgSize,legend=legends[i],**kwargs),(i*subImgSize[0],0))
  return res

def MolsToGridImage(mols,molsPerRow=3,subImgSize=(200,200),legends=None,
                    highlightAtomLists=None,**kwargs):
  """
  """
  try:
    import Image
  except ImportError:
    from PIL import Image
  if legends is None: legends = [None]*len(mols)

  nRows = len(mols)//molsPerRow
  if len(mols)%molsPerRow : nRows+=1

  res = Image.new("RGBA",(molsPerRow*subImgSize[0],nRows*subImgSize[1]),(255,255,255,0))
  for i,mol in enumerate(mols):
    row = i//molsPerRow
    col = i%molsPerRow
    highlights=None
    if highlightAtomLists and highlightAtomLists[i]:
      highlights=highlightAtomLists[i]
    res.paste(MolToImage(mol,subImgSize,legend=legends[i],highlightAtoms=highlights,
                         **kwargs),(col*subImgSize[0],row*subImgSize[1]))
  return res

def ReactionToImage(rxn, subImgSize=(200,200),**kwargs):
  """
  """
  try:
    import Image
  except ImportError:
    from PIL import Image

  mols = []
  for i in range(rxn.GetNumReactantTemplates()):
    tmpl=rxn.GetReactantTemplate(i)
    tmpl.UpdatePropertyCache(False)
    mols.append(tmpl)
  mols.append(None)
  for i in range(rxn.GetNumProductTemplates()):
    tmpl = rxn.GetProductTemplate(i)
    tmpl.UpdatePropertyCache(False)
    mols.append(tmpl)

  res = Image.new("RGBA",(subImgSize[0]*len(mols),subImgSize[1]),(255,255,255,0))
  for i,mol in enumerate(mols):
    if mol is not None:
      nimg = MolToImage(mol,subImgSize,kekulize=False,**kwargs)
    else:
      nimg,canvas = _createCanvas(subImgSize)
      p0 = (10,subImgSize[1]//2)
      p1 = (subImgSize[0]-10,subImgSize[1]//2)
      p3 = (subImgSize[0]-20,subImgSize[1]//2-10)
      p4 = (subImgSize[0]-20,subImgSize[1]//2+10)
      canvas.addCanvasLine(p0,p1,lineWidth=2,color=(0,0,0))
      canvas.addCanvasLine(p3,p1,lineWidth=2,color=(0,0,0))
      canvas.addCanvasLine(p4,p1,lineWidth=2,color=(0,0,0))
      if hasattr(canvas,'flush'):
        canvas.flush()
      else:
        canvas.save()
    res.paste(nimg,(i*subImgSize[0],0))
  return res

<|endoftext|>


<commit_before>from utils.exceptions import AuthenticationRequired, NoInternetConnection
import utils.structures as structures
from reader import reader
import writer
import urllib2

def internet_required(fn):
	def wrapper(*args, **kwargs):
		try:
			return fn(*args, **kwargs)
		except urllib2.URLError:
			raise NoInternetConnection
	return wrapper

def login_required(fn):
	def wrapper(self, *args, **kwargs):
		if not self.user.loggedin:
			raise AuthenticationRequired
		return fn(self, *args, **kwargs)
	return wrapper

class MinervaSession:

	def __init__(self, user):

		self.site = structures.minervaSite()
		self.user = structures.minervaCred(user[0],user[1])
		self.writer = writer.MinervaWriter(self.site)
		self.reader = reader.MinervaReader(self.site)

	@internet_required
	def login(self) :

		# Logs in, This page can confirm a succesful login.
		loginResponse = self.writer.login(self.user)

		loggedin, msg = self.reader.login(loginResponse)

		self.user.loggedin = loggedin

		return msg[0]

	@internet_required
	def logout(self):

		if not self.user.loggedin:
			print 'You are already logged out.'
		else:
			# Logs out.
			logoutPage = self.writer.logout()

		return logoutPage.read()

	@internet_required
	@login_required
	def drop(self, crn):

		print 'Dropping', crn, '...'

		dropPage = self.writer.drop(crn)

		return dropPage.read()

	@internet_required
	@login_required
	def transcript(self):

		transcriptPage = self.writer.transcript()

		curriculum = self.reader.transcript(transcriptPage)

		return curriculum

	@internet_required
	@login_required
	def list(self):

		print 'Working on this function now.'

	@internet_required
	def deal_with_request(self, req):

		response = ''

		if req['login']:
			response = self.login()

		elif req['drop']:
			response = self.drop(req['<class_crn>'])

		elif req['transcript']:
			response = self.transcript()

		else:
			response = 'Not implemented yet.'

		return response<commit_msg>Del: Removed unused list function.<commit_after>from utils.exceptions import AuthenticationRequired, NoInternetConnection
import utils.structures as structures
from reader import reader
import writer
import urllib2

def internet_required(fn):
	def wrapper(*args, **kwargs):
		try:
			return fn(*args, **kwargs)
		except urllib2.URLError:
			raise NoInternetConnection
	return wrapper

def login_required(fn):
	def wrapper(self, *args, **kwargs):
		if not self.user.loggedin:
			raise AuthenticationRequired
		return fn(self, *args, **kwargs)
	return wrapper

class MinervaSession:

	def __init__(self, user):

		self.site = structures.minervaSite()
		self.user = structures.minervaCred(user[0],user[1])
		self.writer = writer.MinervaWriter(self.site)
		self.reader = reader.MinervaReader(self.site)

	@internet_required
	def login(self) :

		# Logs in, This page can confirm a succesful login.
		loginResponse = self.writer.login(self.user)

		loggedin, msg = self.reader.login(loginResponse)

		self.user.loggedin = loggedin

		return msg[0]

	@internet_required
	def logout(self):

		if not self.user.loggedin:
			print 'You are already logged out.'
		else:
			# Logs out.
			logoutPage = self.writer.logout()

		return logoutPage.read()

	@internet_required
	@login_required
	def drop(self, crn):

		print 'Dropping', crn, '...'

		dropPage = self.writer.drop(crn)

		return dropPage.read()

	@internet_required
	@login_required
	def transcript(self):

		transcriptPage = self.writer.transcript()

		curriculum = self.reader.transcript(transcriptPage)

		return curriculum

	@internet_required
	def deal_with_request(self, req):

		response = ''

		if req['login']:
			response = self.login()

		elif req['drop']:
			response = self.drop(req['<class_crn>'])

		elif req['transcript']:
			response = self.transcript()

		else:
			response = 'Not implemented yet.'

		return response<|endoftext|>


<commit_before># coding=utf-8
# Copyright 2020 The TF-Agents Authors.
#
# Licensed 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
#
#     https://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.

r"""Train and Eval PPO.

To run:

```bash
tensorboard --logdir $HOME/tmp/ppo/gym/HalfCheetah-v2/ --port 2223 &

python tf_agents/agents/ppo/examples/v2/train_eval_clip_agent.py \
  --root_dir=$HOME/tmp/ppo/gym/HalfCheetah-v2/ \
  --logtostderr
```
"""

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

import functools
import os
import time

from absl import app
from absl import flags
from absl import logging

import gin
import tensorflow as tf  # pylint: disable=g-explicit-tensorflow-version-import

from tf_agents.agents.ppo import ppo_clip_agent
from tf_agents.drivers import dynamic_episode_driver
from tf_agents.environments import parallel_py_environment
from tf_agents.environments import suite_mujoco
from tf_agents.environments import tf_py_environment
from tf_agents.eval import metric_utils
from tf_agents.metrics import tf_metrics
from tf_agents.networks import actor_distribution_network
from tf_agents.networks import actor_distribution_rnn_network
from tf_agents.networks import value_network
from tf_agents.networks import value_rnn_network
from tf_agents.policies import policy_saver
from tf_agents.replay_buffers import tf_uniform_replay_buffer
from tf_agents.system import system_multiprocessing as multiprocessing
from tf_agents.utils import common

flags.DEFINE_string('root_dir', os.getenv('TEST_UNDECLARED_OUTPUTS_DIR'),
                    'Root directory for writing logs/summaries/checkpoints.')
flags.DEFINE_string('env_name', 'HalfCheetah-v2', 'Name of an environment')
flags.DEFINE_integer('replay_buffer_capacity', 1001,
                     'Replay buffer capacity per env.')
flags.DEFINE_integer('num_parallel_environments', 30,
                     'Number of environments to run in parallel')
flags.DEFINE_integer('num_environment_steps', 25000000,
                     'Number of environment steps to run before finishing.')
flags.DEFINE_integer('num_epochs', 25,
                     'Number of epochs for computing policy updates.')
flags.DEFINE_integer(
    'collect_episodes_per_iteration', 30,
    'The number of episodes to take in the environment before '
    'each update. This is the total across all parallel '
    'environments.')
flags.DEFINE_integer('num_eval_episodes', 30,
                     'The number of episodes to run eval on.')
flags.DEFINE_boolean('use_rnns', False,
                     'If true, use RNN for policy and value function.')
FLAGS = flags.FLAGS

@gin.configurable
def train_eval(
    root_dir,
    env_name='HalfCheetah-v2',
    env_load_fn=suite_mujoco.load,
    random_seed=None,
    # TODO(b/127576522): rename to policy_fc_layers.
    actor_fc_layers=(200, 100),
    value_fc_layers=(200, 100),
    use_rnns=False,
    # Params for collect
    num_environment_steps=25000000,
    collect_episodes_per_iteration=30,
    num_parallel_environments=30,
    replay_buffer_capacity=1001,  # Per-environment
    # Params for train
    num_epochs=25,
    learning_rate=1e-3,
    # Params for eval
    num_eval_episodes=30,
    eval_interval=500,
    # Params for summaries and logging
    train_checkpoint_interval=500,
    policy_checkpoint_interval=500,
    log_interval=50,
    summary_interval=50,
    summaries_flush_secs=1,
    use_tf_functions=True,
    debug_summaries=False,
    summarize_grads_and_vars=False):
  """A simple train and eval for PPO."""
  if root_dir is None:
    raise AttributeError('train_eval requires a root_dir.')

  root_dir = os.path.expanduser(root_dir)
  train_dir = os.path.join(root_dir, 'train')
  eval_dir = os.path.join(root_dir, 'eval')
  saved_model_dir = os.path.join(root_dir, 'policy_saved_model')

  train_summary_writer = tf.compat.v2.summary.create_file_writer(
      train_dir, flush_millis=summaries_flush_secs * 1000)
  train_summary_writer.set_as_default()

  eval_summary_writer = tf.compat.v2.summary.create_file_writer(
      eval_dir, flush_millis=summaries_flush_secs * 1000)
  eval_metrics = [
      tf_metrics.AverageReturnMetric(buffer_size=num_eval_episodes),
      tf_metrics.AverageEpisodeLengthMetric(buffer_size=num_eval_episodes)
  ]

  global_step = tf.compat.v1.train.get_or_create_global_step()
  with tf.compat.v2.summary.record_if(
      lambda: tf.math.equal(global_step % summary_interval, 0)):
    if random_seed is not None:
      tf.compat.v1.set_random_seed(random_seed)
    eval_tf_env = tf_py_environment.TFPyEnvironment(env_load_fn(env_name))
    tf_env = tf_py_environment.TFPyEnvironment(
        parallel_py_environment.ParallelPyEnvironment(
            [lambda: env_load_fn(env_name)] * num_parallel_environments))
    optimizer = tf.compat.v1.train.AdamOptimizer(learning_rate=learning_rate)

    if use_rnns:
      actor_net = actor_distribution_rnn_network.ActorDistributionRnnNetwork(
          tf_env.observation_spec(),
          tf_env.action_spec(),
          input_fc_layer_params=actor_fc_layers,
          output_fc_layer_params=None)
      value_net = value_rnn_network.ValueRnnNetwork(
          tf_env.observation_spec(),
          input_fc_layer_params=value_fc_layers,
          output_fc_layer_params=None)
    else:
      actor_net = actor_distribution_network.ActorDistributionNetwork(
          tf_env.observation_spec(),
          tf_env.action_spec(),
          fc_layer_params=actor_fc_layers,
          activation_fn=tf.keras.activations.tanh)
      value_net = value_network.ValueNetwork(
          tf_env.observation_spec(),
          fc_layer_params=value_fc_layers,
          activation_fn=tf.keras.activations.tanh)

    tf_agent = ppo_clip_agent.PPOClipAgent(
        tf_env.time_step_spec(),
        tf_env.action_spec(),
        optimizer,
        actor_net=actor_net,
        value_net=value_net,
        entropy_regularization=0.0,
        importance_ratio_clipping=0.2,
        normalize_observations=False,
        normalize_rewards=False,
        use_gae=True,
        num_epochs=num_epochs,
        debug_summaries=debug_summaries,
        summarize_grads_and_vars=summarize_grads_and_vars,
        train_step_counter=global_step)
    tf_agent.initialize()

    environment_steps_metric = tf_metrics.EnvironmentSteps()
    step_metrics = [
        tf_metrics.NumberOfEpisodes(),
        environment_steps_metric,
    ]

    train_metrics = step_metrics + [
        tf_metrics.AverageReturnMetric(
            batch_size=num_parallel_environments),
        tf_metrics.AverageEpisodeLengthMetric(
            batch_size=num_parallel_environments),
    ]

    eval_policy = tf_agent.policy
    collect_policy = tf_agent.collect_policy

    replay_buffer = tf_uniform_replay_buffer.TFUniformReplayBuffer(
        tf_agent.collect_data_spec,
        batch_size=num_parallel_environments,
        max_length=replay_buffer_capacity)

    train_checkpointer = common.Checkpointer(
        ckpt_dir=train_dir,
        agent=tf_agent,
        global_step=global_step,
        metrics=metric_utils.MetricsGroup(train_metrics, 'train_metrics'))
    policy_checkpointer = common.Checkpointer(
        ckpt_dir=os.path.join(train_dir, 'policy'),
        policy=eval_policy,
        global_step=global_step)
    saved_model = policy_saver.PolicySaver(
        eval_policy, train_step=global_step)

    train_checkpointer.initialize_or_restore()

    collect_driver = dynamic_episode_driver.DynamicEpisodeDriver(
        tf_env,
        collect_policy,
        observers=[replay_buffer.add_batch] + train_metrics,
        num_episodes=collect_episodes_per_iteration)

    def train_step():
      trajectories = replay_buffer.gather_all()
      return tf_agent.train(experience=trajectories)

    if use_tf_functions:
      # TODO(b/123828980): Enable once the cause for slowdown was identified.
      collect_driver.run = common.function(collect_driver.run, autograph=False)
      tf_agent.train = common.function(tf_agent.train, autograph=False)
      train_step = common.function(train_step)

    collect_time = 0
    train_time = 0
    timed_at_step = global_step.numpy()

    while environment_steps_metric.result() < num_environment_steps:
      global_step_val = global_step.numpy()
      if global_step_val % eval_interval == 0:
        metric_utils.eager_compute(
            eval_metrics,
            eval_tf_env,
            eval_policy,
            num_episodes=num_eval_episodes,
            train_step=global_step,
            summary_writer=eval_summary_writer,
            summary_prefix='Metrics',
        )

      start_time = time.time()
      collect_driver.run()
      collect_time += time.time() - start_time

      start_time = time.time()
      total_loss, _ = train_step()
      replay_buffer.clear()
      train_time += time.time() - start_time

      for train_metric in train_metrics:
        train_metric.tf_summaries(
            train_step=global_step, step_metrics=step_metrics)

      if global_step_val % log_interval == 0:
        logging.info('step = %d, loss = %f', global_step_val, total_loss)
        steps_per_sec = (
            (global_step_val - timed_at_step) / (collect_time + train_time))
        logging.info('%.3f steps/sec', steps_per_sec)
        logging.info('collect_time = %.3f, train_time = %.3f', collect_time,
                     train_time)
        with tf.compat.v2.summary.record_if(True):
          tf.compat.v2.summary.scalar(
              name='global_steps_per_sec', data=steps_per_sec, step=global_step)

        if global_step_val % train_checkpoint_interval == 0:
          train_checkpointer.save(global_step=global_step_val)

        if global_step_val % policy_checkpoint_interval == 0:
          policy_checkpointer.save(global_step=global_step_val)
          saved_model_path = os.path.join(
              saved_model_dir, 'policy_' + ('%d' % global_step_val).zfill(9))
          saved_model.save(saved_model_path)

        timed_at_step = global_step_val
        collect_time = 0
        train_time = 0

    # One final eval before exiting.
    metric_utils.eager_compute(
        eval_metrics,
        eval_tf_env,
        eval_policy,
        num_episodes=num_eval_episodes,
        train_step=global_step,
        summary_writer=eval_summary_writer,
        summary_prefix='Metrics',
    )

def main(_):
  logging.set_verbosity(logging.INFO)
  tf.compat.v1.enable_v2_behavior()
  train_eval(
      FLAGS.root_dir,
      env_name=FLAGS.env_name,
      use_rnns=FLAGS.use_rnns,
      num_environment_steps=FLAGS.num_environment_steps,
      collect_episodes_per_iteration=FLAGS.collect_episodes_per_iteration,
      num_parallel_environments=FLAGS.num_parallel_environments,
      replay_buffer_capacity=FLAGS.replay_buffer_capacity,
      num_epochs=FLAGS.num_epochs,
      num_eval_episodes=FLAGS.num_eval_episodes)

if __name__ == '__main__':
  flags.mark_flag_as_required('root_dir')
  multiprocessing.handle_main(functools.partial(app.run, main))
<commit_msg>Fix train_eval_clip_agent with RNNs test.<commit_after># coding=utf-8
# Copyright 2020 The TF-Agents Authors.
#
# Licensed 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
#
#     https://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.

r"""Train and Eval PPO.

To run:

```bash
tensorboard --logdir $HOME/tmp/ppo/gym/HalfCheetah-v2/ --port 2223 &

python tf_agents/agents/ppo/examples/v2/train_eval_clip_agent.py \
  --root_dir=$HOME/tmp/ppo/gym/HalfCheetah-v2/ \
  --logtostderr
```
"""

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

import functools
import os
import time

from absl import app
from absl import flags
from absl import logging

import gin
import tensorflow as tf  # pylint: disable=g-explicit-tensorflow-version-import

from tf_agents.agents.ppo import ppo_clip_agent
from tf_agents.drivers import dynamic_episode_driver
from tf_agents.environments import parallel_py_environment
from tf_agents.environments import suite_mujoco
from tf_agents.environments import tf_py_environment
from tf_agents.eval import metric_utils
from tf_agents.metrics import tf_metrics
from tf_agents.networks import actor_distribution_network
from tf_agents.networks import actor_distribution_rnn_network
from tf_agents.networks import value_network
from tf_agents.networks import value_rnn_network
from tf_agents.policies import policy_saver
from tf_agents.replay_buffers import tf_uniform_replay_buffer
from tf_agents.system import system_multiprocessing as multiprocessing
from tf_agents.utils import common

flags.DEFINE_string('root_dir', os.getenv('TEST_UNDECLARED_OUTPUTS_DIR'),
                    'Root directory for writing logs/summaries/checkpoints.')
flags.DEFINE_string('env_name', 'HalfCheetah-v2', 'Name of an environment')
flags.DEFINE_integer('replay_buffer_capacity', 1001,
                     'Replay buffer capacity per env.')
flags.DEFINE_integer('num_parallel_environments', 30,
                     'Number of environments to run in parallel')
flags.DEFINE_integer('num_environment_steps', 25000000,
                     'Number of environment steps to run before finishing.')
flags.DEFINE_integer('num_epochs', 25,
                     'Number of epochs for computing policy updates.')
flags.DEFINE_integer(
    'collect_episodes_per_iteration', 30,
    'The number of episodes to take in the environment before '
    'each update. This is the total across all parallel '
    'environments.')
flags.DEFINE_integer('num_eval_episodes', 30,
                     'The number of episodes to run eval on.')
flags.DEFINE_boolean('use_rnns', False,
                     'If true, use RNN for policy and value function.')
FLAGS = flags.FLAGS

@gin.configurable
def train_eval(
    root_dir,
    env_name='HalfCheetah-v2',
    env_load_fn=suite_mujoco.load,
    random_seed=None,
    # TODO(b/127576522): rename to policy_fc_layers.
    actor_fc_layers=(200, 100),
    value_fc_layers=(200, 100),
    use_rnns=False,
    lstm_size=(20,),
    # Params for collect
    num_environment_steps=25000000,
    collect_episodes_per_iteration=30,
    num_parallel_environments=30,
    replay_buffer_capacity=1001,  # Per-environment
    # Params for train
    num_epochs=25,
    learning_rate=1e-3,
    # Params for eval
    num_eval_episodes=30,
    eval_interval=500,
    # Params for summaries and logging
    train_checkpoint_interval=500,
    policy_checkpoint_interval=500,
    log_interval=50,
    summary_interval=50,
    summaries_flush_secs=1,
    use_tf_functions=True,
    debug_summaries=False,
    summarize_grads_and_vars=False):
  """A simple train and eval for PPO."""
  if root_dir is None:
    raise AttributeError('train_eval requires a root_dir.')

  root_dir = os.path.expanduser(root_dir)
  train_dir = os.path.join(root_dir, 'train')
  eval_dir = os.path.join(root_dir, 'eval')
  saved_model_dir = os.path.join(root_dir, 'policy_saved_model')

  train_summary_writer = tf.compat.v2.summary.create_file_writer(
      train_dir, flush_millis=summaries_flush_secs * 1000)
  train_summary_writer.set_as_default()

  eval_summary_writer = tf.compat.v2.summary.create_file_writer(
      eval_dir, flush_millis=summaries_flush_secs * 1000)
  eval_metrics = [
      tf_metrics.AverageReturnMetric(buffer_size=num_eval_episodes),
      tf_metrics.AverageEpisodeLengthMetric(buffer_size=num_eval_episodes)
  ]

  global_step = tf.compat.v1.train.get_or_create_global_step()
  with tf.compat.v2.summary.record_if(
      lambda: tf.math.equal(global_step % summary_interval, 0)):
    if random_seed is not None:
      tf.compat.v1.set_random_seed(random_seed)
    eval_tf_env = tf_py_environment.TFPyEnvironment(env_load_fn(env_name))
    tf_env = tf_py_environment.TFPyEnvironment(
        parallel_py_environment.ParallelPyEnvironment(
            [lambda: env_load_fn(env_name)] * num_parallel_environments))
    optimizer = tf.compat.v1.train.AdamOptimizer(learning_rate=learning_rate)

    if use_rnns:
      actor_net = actor_distribution_rnn_network.ActorDistributionRnnNetwork(
          tf_env.observation_spec(),
          tf_env.action_spec(),
          input_fc_layer_params=actor_fc_layers,
          output_fc_layer_params=None,
          lstm_size=lstm_size)
      value_net = value_rnn_network.ValueRnnNetwork(
          tf_env.observation_spec(),
          input_fc_layer_params=value_fc_layers,
          output_fc_layer_params=None)
    else:
      actor_net = actor_distribution_network.ActorDistributionNetwork(
          tf_env.observation_spec(),
          tf_env.action_spec(),
          fc_layer_params=actor_fc_layers,
          activation_fn=tf.keras.activations.tanh)
      value_net = value_network.ValueNetwork(
          tf_env.observation_spec(),
          fc_layer_params=value_fc_layers,
          activation_fn=tf.keras.activations.tanh)

    tf_agent = ppo_clip_agent.PPOClipAgent(
        tf_env.time_step_spec(),
        tf_env.action_spec(),
        optimizer,
        actor_net=actor_net,
        value_net=value_net,
        entropy_regularization=0.0,
        importance_ratio_clipping=0.2,
        normalize_observations=False,
        normalize_rewards=False,
        use_gae=True,
        num_epochs=num_epochs,
        debug_summaries=debug_summaries,
        summarize_grads_and_vars=summarize_grads_and_vars,
        train_step_counter=global_step)
    tf_agent.initialize()

    environment_steps_metric = tf_metrics.EnvironmentSteps()
    step_metrics = [
        tf_metrics.NumberOfEpisodes(),
        environment_steps_metric,
    ]

    train_metrics = step_metrics + [
        tf_metrics.AverageReturnMetric(
            batch_size=num_parallel_environments),
        tf_metrics.AverageEpisodeLengthMetric(
            batch_size=num_parallel_environments),
    ]

    eval_policy = tf_agent.policy
    collect_policy = tf_agent.collect_policy

    replay_buffer = tf_uniform_replay_buffer.TFUniformReplayBuffer(
        tf_agent.collect_data_spec,
        batch_size=num_parallel_environments,
        max_length=replay_buffer_capacity)

    train_checkpointer = common.Checkpointer(
        ckpt_dir=train_dir,
        agent=tf_agent,
        global_step=global_step,
        metrics=metric_utils.MetricsGroup(train_metrics, 'train_metrics'))
    policy_checkpointer = common.Checkpointer(
        ckpt_dir=os.path.join(train_dir, 'policy'),
        policy=eval_policy,
        global_step=global_step)
    saved_model = policy_saver.PolicySaver(
        eval_policy, train_step=global_step)

    train_checkpointer.initialize_or_restore()

    collect_driver = dynamic_episode_driver.DynamicEpisodeDriver(
        tf_env,
        collect_policy,
        observers=[replay_buffer.add_batch] + train_metrics,
        num_episodes=collect_episodes_per_iteration)

    def train_step():
      trajectories = replay_buffer.gather_all()
      return tf_agent.train(experience=trajectories)

    if use_tf_functions:
      # TODO(b/123828980): Enable once the cause for slowdown was identified.
      collect_driver.run = common.function(collect_driver.run, autograph=False)
      tf_agent.train = common.function(tf_agent.train, autograph=False)
      train_step = common.function(train_step)

    collect_time = 0
    train_time = 0
    timed_at_step = global_step.numpy()

    while environment_steps_metric.result() < num_environment_steps:
      global_step_val = global_step.numpy()
      if global_step_val % eval_interval == 0:
        metric_utils.eager_compute(
            eval_metrics,
            eval_tf_env,
            eval_policy,
            num_episodes=num_eval_episodes,
            train_step=global_step,
            summary_writer=eval_summary_writer,
            summary_prefix='Metrics',
        )

      start_time = time.time()
      collect_driver.run()
      collect_time += time.time() - start_time

      start_time = time.time()
      total_loss, _ = train_step()
      replay_buffer.clear()
      train_time += time.time() - start_time

      for train_metric in train_metrics:
        train_metric.tf_summaries(
            train_step=global_step, step_metrics=step_metrics)

      if global_step_val % log_interval == 0:
        logging.info('step = %d, loss = %f', global_step_val, total_loss)
        steps_per_sec = (
            (global_step_val - timed_at_step) / (collect_time + train_time))
        logging.info('%.3f steps/sec', steps_per_sec)
        logging.info('collect_time = %.3f, train_time = %.3f', collect_time,
                     train_time)
        with tf.compat.v2.summary.record_if(True):
          tf.compat.v2.summary.scalar(
              name='global_steps_per_sec', data=steps_per_sec, step=global_step)

        if global_step_val % train_checkpoint_interval == 0:
          train_checkpointer.save(global_step=global_step_val)

        if global_step_val % policy_checkpoint_interval == 0:
          policy_checkpointer.save(global_step=global_step_val)
          saved_model_path = os.path.join(
              saved_model_dir, 'policy_' + ('%d' % global_step_val).zfill(9))
          saved_model.save(saved_model_path)

        timed_at_step = global_step_val
        collect_time = 0
        train_time = 0

    # One final eval before exiting.
    metric_utils.eager_compute(
        eval_metrics,
        eval_tf_env,
        eval_policy,
        num_episodes=num_eval_episodes,
        train_step=global_step,
        summary_writer=eval_summary_writer,
        summary_prefix='Metrics',
    )

def main(_):
  logging.set_verbosity(logging.INFO)
  tf.compat.v1.enable_v2_behavior()
  train_eval(
      FLAGS.root_dir,
      env_name=FLAGS.env_name,
      use_rnns=FLAGS.use_rnns,
      num_environment_steps=FLAGS.num_environment_steps,
      collect_episodes_per_iteration=FLAGS.collect_episodes_per_iteration,
      num_parallel_environments=FLAGS.num_parallel_environments,
      replay_buffer_capacity=FLAGS.replay_buffer_capacity,
      num_epochs=FLAGS.num_epochs,
      num_eval_episodes=FLAGS.num_eval_episodes)

if __name__ == '__main__':
  flags.mark_flag_as_required('root_dir')
  multiprocessing.handle_main(functools.partial(app.run, main))
<|endoftext|>