<commit_before># encoding: utf-8

"""MIME-encoded electronic mail message class."""

import imghdr
import os
import time
import base64

from datetime import datetime
from email.mime.text import MIMEText
from email.mime.multipart import MIMEMultipart
from email.mime.nonmultipart import MIMENonMultipart
from email.header import Header
from email.utils import make_msgid, formatdate
from mimetypes import guess_type

from marrow.mailer import release
from marrow.mailer.address import Address, AddressList, AutoConverter
from marrow.util.compat import basestring, unicode

__all__ = ['Message']

class Message(object):
    """Represents an e-mail message."""

    sender = AutoConverter('_sender', Address, False)
    author = AutoConverter('_author', AddressList)
    authors = author
    to = AutoConverter('_to', AddressList)
    cc = AutoConverter('_cc', AddressList)
    bcc = AutoConverter('_bcc', AddressList)
    reply = AutoConverter('_reply', AddressList)
    notify = AutoConverter('_notify', AddressList)

    def __init__(self, author=None, to=None, subject=None, **kw):
        """Instantiate a new Message object.

        No arguments are required, as everything can be set using class
        properties.  Alternatively, __everything__ can be set using the
        constructor, using named arguments.  The first three positional
        arguments can be used to quickly prepare a simple message.
        """

        # Internally used attributes
        self._id = None
        self._processed = False
        self._dirty = False
        self.mailer = None

        # Default values
        self.subject = None
        self.date = datetime.now()
        self.encoding = 'utf-8'
        self.organization = None
        self.priority = None
        self.plain = None
        self.rich = None
        self.attachments = []
        self.embedded = []
        self.headers = []
        self.retries = 3
        self.brand = True

        self._sender = None
        self._author = AddressList()
        self._to = AddressList()
        self._cc = AddressList()
        self._bcc = AddressList()
        self._reply = AddressList()
        self._notify = AddressList()

        # Overrides at initialization time
        if author is not None:
            self.author = author

        if to is not None:
            self.to = to

        if subject is not None:
            self.subject = subject

        for k in kw:
            if not hasattr(self, k):
                raise TypeError("Unexpected keyword argument: %s" % k)

            setattr(self, k, kw[k])

    def __setattr__(self, name, value):
        """Set the dirty flag as properties are updated."""
        object.__setattr__(self, name, value)
        if name not in ('bcc', '_id', '_dirty', '_processed'):
            object.__setattr__(self, '_dirty', True)

    def __str__(self):
        return self.mime.as_string()

    @property
    def id(self):
        if not self._id or (self._processed and self._dirty):
            self._id = make_msgid()
            self._processed = False
        return self._id

    @property
    def envelope(self):
        """Returns the address of the envelope sender address (SMTP from, if
        not set the sender, if this one isn't set too, the author)."""
        if not self.sender and not self.author:
            raise ValueError("Unable to determine message sender; no author or sender defined.")

        return self.sender or self.author[0]

    @property
    def recipients(self):
        return AddressList(self.to + self.cc + self.bcc)

    def _mime_document(self, plain, rich=None):
        if not rich:
            message = plain

        else:
            message = MIMEMultipart('alternative')
            message.attach(plain)

            if not self.embedded:
                message.attach(rich)

            else:
                embedded = MIMEMultipart('related')
                embedded.attach(rich)
                for attachment in self.embedded:
                    embedded.attach(attachment)
                message.attach(embedded)

        if self.attachments:
            attachments = MIMEMultipart()
            attachments.attach(message)
            for attachment in self.attachments:
                attachments.attach(attachment)
            message = attachments

        return message

    def _build_date_header_string(self, date_value):
        """Gets the date_value (may be None, basestring, float or
        datetime.datetime instance) and returns a valid date string as per
        RFC 2822."""
        if isinstance(date_value, datetime):
            date_value = time.mktime(date_value.timetuple())
        if not isinstance(date_value, basestring):
            date_value = formatdate(date_value, localtime=True)
        # Encode it here to avoid this:
        # Date: =?utf-8?q?Sat=2C_01_Sep_2012_13=3A08=3A29_-0300?=
        return date_value.encode('ascii')

    def _build_header_list(self, author, sender):
        date_value = self._build_date_header_string(self.date)
        headers = [
                ('Sender', sender),
                ('From', author),
                ('Reply-To', self.reply),
                ('Subject', self.subject),
                ('Date', date_value),
                ('To', self.to),
                ('Cc', self.cc),
                ('Disposition-Notification-To', self.notify),
                ('Organization', self.organization),
                ('X-Priority', self.priority),
            ]

        if self.brand:
            headers.extend([
                    ('X-Mailer', "marrow.mailer {0}".format(release.version))
                ])

        if isinstance(self.headers, dict):
            for key in self.headers:
                headers.append((key, self.headers[key]))

        else:
            headers.extend(self.headers)

        return headers

    def _add_headers_to_message(self, message, headers):
        for header in headers:
            if header[1] is None or (isinstance(header[1], list) and not header[1]):
                continue

            name, value = header

            if isinstance(value, Address):
                value = value.encode(self.encoding)
            elif isinstance(value, AddressList):
                value = value.encode(self.encoding)

            if isinstance(value, unicode):
                value = Header(value, self.encoding)
            else:
                value = Header(value)

            message[name] = value

    @property
    def mime(self):
        """Produce the final MIME message."""
        author = self.author
        sender = self.sender

        if not author:
            raise ValueError("You must specify an author.")

        if not self.subject:
            raise ValueError("You must specify a subject.")

        if len(self.recipients) == 0:
            raise ValueError("You must specify at least one recipient.")

        if not self.plain:
            raise ValueError("You must provide plain text content.")

        # DISCUSS: Take the first author, or raise this error?
        # if len(author) > 1 and len(sender) == 0:
        #     raise ValueError('If there are multiple authors of message, you must specify a sender!')

        # if len(sender) > 1:
        #     raise ValueError('You must not specify more than one sender!')

        if not self._dirty and self._processed:
            return self._mime

        self._processed = False

        plain = MIMEText(self._callable(self.plain), 'plain', self.encoding)

        rich = None
        if self.rich:
            rich = MIMEText(self._callable(self.rich), 'html', self.encoding)

        message = self._mime_document(plain, rich)
        headers = self._build_header_list(author, sender)
        self._add_headers_to_message(message, headers)

        self._mime = message
        self._processed = True
        self._dirty = False

        return message

    def attach(self, name, data=None, maintype=None, subtype=None,
        inline=False, filename=None, encoding=None):
        """Attach a file to this message.

        :param name: Path to the file to attach if data is None, or the name
                     of the file if the ``data`` argument is given
        :param data: Contents of the file to attach, or None if the data is to
                     be read from the file pointed to by the ``name`` argument
        :type data: bytes or a file-like object
        :param maintype: First part of the MIME type of the file -- will be
                         automatically guessed if not given
        :param subtype: Second part of the MIME type of the file -- will be
                        automatically guessed if not given
        :param inline: Whether to set the Content-Disposition for the file to
                       "inline" (True) or "attachment" (False)
        :param filename: The file name of the attached file as seen
                                    by the user in his/her mail client.
        :param encoding: Value of the Content-Encoding MIME header (e.g. "gzip"
                         in case of .tar.gz, but usually empty)
        """
        self._dirty = True

        if not maintype:
            maintype, guessed_encoding = guess_type(name)
            encoding = encoding or guessed_encoding
            if not maintype:
                maintype, subtype = 'application', 'octet-stream'
            else:
                maintype, _, subtype = maintype.partition('/')

        part = MIMENonMultipart(maintype, subtype)
        part.add_header('Content-Transfer-Encoding', 'base64')

        if encoding:
            part.add_header('Content-Encoding', encoding)

        if data is None:
            with open(name, 'rb') as fp:
                value = fp.read()
            name = os.path.basename(name)
        elif isinstance(data, bytes):
            value = data
        elif hasattr(data, 'read'):
            value = data.read()
        else:
            raise TypeError("Unable to read attachment contents")

        part.set_payload(base64.encodestring(value))

        if not filename:
            filename = name
        filename = os.path.basename(filename)

        if inline:
            part.add_header('Content-Disposition', 'inline', filename=filename)
            part.add_header('Content-ID', '<%s>' % filename)
            self.embedded.append(part)
        else:
            part.add_header('Content-Disposition', 'attachment', filename=filename)
            self.attachments.append(part)

    def embed(self, name, data=None):
        """Attach an image file and prepare for HTML embedding.

        This method should only be used to embed images.

        :param name: Path to the image to embed if data is None, or the name
                     of the file if the ``data`` argument is given
        :param data: Contents of the image to embed, or None if the data is to
                     be read from the file pointed to by the ``name`` argument
        """
        if data is None:
            with open(name, 'rb') as fp:
                data = fp.read()
            name = os.path.basename(name)
        elif isinstance(data, bytes):
            pass
        elif hasattr(data, 'read'):
            data = data.read()
        else:
            raise TypeError("Unable to read image contents")

        subtype = imghdr.what(None, data)
        self.attach(name, data, 'image', subtype, True)

    @staticmethod
    def _callable(var):
        if hasattr(var, '__call__'):
            return var()

        return var

    def send(self):
        if not self.mailer:
            raise NotImplementedError("Message instance is not bound to " \
                "a Mailer. Use mailer.send() instead.")
        return self.mailer.send(self)
<commit_msg>Reindentation.<commit_after># encoding: utf-8

"""MIME-encoded electronic mail message class."""

import imghdr
import os
import time
import base64

from datetime import datetime
from email.mime.text import MIMEText
from email.mime.multipart import MIMEMultipart
from email.mime.nonmultipart import MIMENonMultipart
from email.header import Header
from email.utils import make_msgid, formatdate
from mimetypes import guess_type

from marrow.mailer import release
from marrow.mailer.address import Address, AddressList, AutoConverter
from marrow.util.compat import basestring, unicode

__all__ = ['Message']

class Message(object):
	"""Represents an e-mail message."""

	sender = AutoConverter('_sender', Address, False)
	author = AutoConverter('_author', AddressList)
	authors = author
	to = AutoConverter('_to', AddressList)
	cc = AutoConverter('_cc', AddressList)
	bcc = AutoConverter('_bcc', AddressList)
	reply = AutoConverter('_reply', AddressList)
	notify = AutoConverter('_notify', AddressList)

	def __init__(self, author=None, to=None, subject=None, **kw):
		"""Instantiate a new Message object.

		No arguments are required, as everything can be set using class
		properties.  Alternatively, __everything__ can be set using the
		constructor, using named arguments.  The first three positional
		arguments can be used to quickly prepare a simple message.
		"""

		# Internally used attributes
		self._id = None
		self._processed = False
		self._dirty = False
		self.mailer = None

		# Default values
		self.subject = None
		self.date = datetime.now()
		self.encoding = 'utf-8'
		self.organization = None
		self.priority = None
		self.plain = None
		self.rich = None
		self.attachments = []
		self.embedded = []
		self.headers = []
		self.retries = 3
		self.brand = True

		self._sender = None
		self._author = AddressList()
		self._to = AddressList()
		self._cc = AddressList()
		self._bcc = AddressList()
		self._reply = AddressList()
		self._notify = AddressList()

		# Overrides at initialization time
		if author is not None:
			self.author = author

		if to is not None:
			self.to = to

		if subject is not None:
			self.subject = subject

		for k in kw:
			if not hasattr(self, k):
				raise TypeError("Unexpected keyword argument: %s" % k)

			setattr(self, k, kw[k])

	def __setattr__(self, name, value):
		"""Set the dirty flag as properties are updated."""
		object.__setattr__(self, name, value)
		if name not in ('bcc', '_id', '_dirty', '_processed'):
			object.__setattr__(self, '_dirty', True)

	def __str__(self):
		return self.mime.as_string()

	@property
	def id(self):
		if not self._id or (self._processed and self._dirty):
			self._id = make_msgid()
			self._processed = False
		return self._id

	@property
	def envelope(self):
		"""Returns the address of the envelope sender address (SMTP from, if
		not set the sender, if this one isn't set too, the author)."""
		if not self.sender and not self.author:
			raise ValueError("Unable to determine message sender; no author or sender defined.")

		return self.sender or self.author[0]

	@property
	def recipients(self):
		return AddressList(self.to + self.cc + self.bcc)

	def _mime_document(self, plain, rich=None):
		if not rich:
			message = plain

		else:
			message = MIMEMultipart('alternative')
			message.attach(plain)

			if not self.embedded:
				message.attach(rich)

			else:
				embedded = MIMEMultipart('related')
				embedded.attach(rich)
				for attachment in self.embedded:
					embedded.attach(attachment)
				message.attach(embedded)

		if self.attachments:
			attachments = MIMEMultipart()
			attachments.attach(message)
			for attachment in self.attachments:
				attachments.attach(attachment)
			message = attachments

		return message

	def _build_date_header_string(self, date_value):
		"""Gets the date_value (may be None, basestring, float or
		datetime.datetime instance) and returns a valid date string as per
		RFC 2822."""
		if isinstance(date_value, datetime):
			date_value = time.mktime(date_value.timetuple())
		if not isinstance(date_value, basestring):
			date_value = formatdate(date_value, localtime=True)
		# Encode it here to avoid this:
		# Date: =?utf-8?q?Sat=2C_01_Sep_2012_13=3A08=3A29_-0300?=
		return date_value.encode('ascii')

	def _build_header_list(self, author, sender):
		date_value = self._build_date_header_string(self.date)
		headers = [
				('Sender', sender),
				('From', author),
				('Reply-To', self.reply),
				('Subject', self.subject),
				('Date', date_value),
				('To', self.to),
				('Cc', self.cc),
				('Disposition-Notification-To', self.notify),
				('Organization', self.organization),
				('X-Priority', self.priority),
			]

		if self.brand:
			headers.extend([
					('X-Mailer', "marrow.mailer {0}".format(release.version))
				])

		if isinstance(self.headers, dict):
			for key in self.headers:
				headers.append((key, self.headers[key]))

		else:
			headers.extend(self.headers)

		return headers

	def _add_headers_to_message(self, message, headers):
		for header in headers:
			if header[1] is None or (isinstance(header[1], list) and not header[1]):
				continue

			name, value = header

			if isinstance(value, Address):
				value = value.encode(self.encoding)
			elif isinstance(value, AddressList):
				value = value.encode(self.encoding)

			if isinstance(value, unicode):
				value = Header(value, self.encoding)
			else:
				value = Header(value)

			message[name] = value

	@property
	def mime(self):
		"""Produce the final MIME message."""
		author = self.author
		sender = self.sender

		if not author:
			raise ValueError("You must specify an author.")

		if not self.subject:
			raise ValueError("You must specify a subject.")

		if len(self.recipients) == 0:
			raise ValueError("You must specify at least one recipient.")

		if not self.plain:
			raise ValueError("You must provide plain text content.")

		# DISCUSS: Take the first author, or raise this error?
		# if len(author) > 1 and len(sender) == 0:
		#	 raise ValueError('If there are multiple authors of message, you must specify a sender!')

		# if len(sender) > 1:
		#	 raise ValueError('You must not specify more than one sender!')

		if not self._dirty and self._processed:
			return self._mime

		self._processed = False

		plain = MIMEText(self._callable(self.plain), 'plain', self.encoding)

		rich = None
		if self.rich:
			rich = MIMEText(self._callable(self.rich), 'html', self.encoding)

		message = self._mime_document(plain, rich)
		headers = self._build_header_list(author, sender)
		self._add_headers_to_message(message, headers)

		self._mime = message
		self._processed = True
		self._dirty = False

		return message

	def attach(self, name, data=None, maintype=None, subtype=None,
		inline=False, filename=None, encoding=None):
		"""Attach a file to this message.

		:param name: Path to the file to attach if data is None, or the name
					 of the file if the ``data`` argument is given
		:param data: Contents of the file to attach, or None if the data is to
					 be read from the file pointed to by the ``name`` argument
		:type data: bytes or a file-like object
		:param maintype: First part of the MIME type of the file -- will be
						 automatically guessed if not given
		:param subtype: Second part of the MIME type of the file -- will be
						automatically guessed if not given
		:param inline: Whether to set the Content-Disposition for the file to
					   "inline" (True) or "attachment" (False)
		:param filename: The file name of the attached file as seen
									by the user in his/her mail client.
		:param encoding: Value of the Content-Encoding MIME header (e.g. "gzip"
						 in case of .tar.gz, but usually empty)
		"""
		self._dirty = True

		if not maintype:
			maintype, guessed_encoding = guess_type(name)
			encoding = encoding or guessed_encoding
			if not maintype:
				maintype, subtype = 'application', 'octet-stream'
			else:
				maintype, _, subtype = maintype.partition('/')

		part = MIMENonMultipart(maintype, subtype)
		part.add_header('Content-Transfer-Encoding', 'base64')

		if encoding:
			part.add_header('Content-Encoding', encoding)

		if data is None:
			with open(name, 'rb') as fp:
				value = fp.read()
			name = os.path.basename(name)
		elif isinstance(data, bytes):
			value = data
		elif hasattr(data, 'read'):
			value = data.read()
		else:
			raise TypeError("Unable to read attachment contents")

		part.set_payload(base64.encodestring(value))

		if not filename:
			filename = name
		filename = os.path.basename(filename)

		if inline:
			part.add_header('Content-Disposition', 'inline', filename=filename)
			part.add_header('Content-ID', '<%s>' % filename)
			self.embedded.append(part)
		else:
			part.add_header('Content-Disposition', 'attachment', filename=filename)
			self.attachments.append(part)

	def embed(self, name, data=None):
		"""Attach an image file and prepare for HTML embedding.

		This method should only be used to embed images.

		:param name: Path to the image to embed if data is None, or the name
					 of the file if the ``data`` argument is given
		:param data: Contents of the image to embed, or None if the data is to
					 be read from the file pointed to by the ``name`` argument
		"""
		if data is None:
			with open(name, 'rb') as fp:
				data = fp.read()
			name = os.path.basename(name)
		elif isinstance(data, bytes):
			pass
		elif hasattr(data, 'read'):
			data = data.read()
		else:
			raise TypeError("Unable to read image contents")

		subtype = imghdr.what(None, data)
		self.attach(name, data, 'image', subtype, True)

	@staticmethod
	def _callable(var):
		if hasattr(var, '__call__'):
			return var()

		return var

	def send(self):
		if not self.mailer:
			raise NotImplementedError("Message instance is not bound to " \
				"a Mailer. Use mailer.send() instead.")
		return self.mailer.send(self)
<|endoftext|>


<commit_before>import pytest
import numpy as np
import mbuild as mb
from mbuild.tests.base_test import BaseTest

class TestBox(BaseTest):

    def test_init_lengths(self):
        box = mb.Box(lengths=np.ones(3))
        assert np.array_equal(box.lengths, np.ones(3))
        assert np.array_equal(box.mins, np.zeros(3))
        assert np.array_equal(box.maxs, np.ones(3))

    def test_init_bounds(self):
        box = mb.Box(mins=np.zeros(3), maxs=np.ones(3))
        assert np.array_equal(box.lengths, np.ones(3))
        assert np.array_equal(box.mins, np.zeros(3))
        assert np.array_equal(box.maxs, np.ones(3))

    def test_scale(self):
        box = mb.Box(lengths=np.ones(3))
        scaling_factors = np.array([3, 4, 5])
        box.scale(scaling_factors)
        assert np.array_equal(box.lengths, scaling_factors)
        assert np.array_equal(box.mins, (np.ones(3) / 2) - (scaling_factors / 2))
        assert np.array_equal(box.maxs, (scaling_factors / 2) + (np.ones(3) / 2))

    def test_center(self):
        box = mb.Box(lengths=np.ones(3))
        box.center()
        assert np.array_equal(box.lengths, np.ones(3))
        assert np.array_equal(box.mins, np.ones(3) * -0.5)
        assert np.array_equal(box.maxs, np.ones(3) * 0.5)
<commit_msg>Remove box tests for non-existant functionality<commit_after>import pytest
import numpy as np
import mbuild as mb
from mbuild.tests.base_test import BaseTest

class TestBox(BaseTest):

    def test_init_lengths(self):
        box = mb.Box(lengths=np.ones(3))
        assert np.array_equal(box.lengths, np.ones(3))
        assert np.array_equal(box.mins, np.zeros(3))
        assert np.array_equal(box.maxs, np.ones(3))

    def test_init_bounds(self):
        box = mb.Box(mins=np.zeros(3), maxs=np.ones(3))
        assert np.array_equal(box.lengths, np.ones(3))
        assert np.array_equal(box.mins, np.zeros(3))
        assert np.array_equal(box.maxs, np.ones(3))
<|endoftext|>


<commit_before>__author__ = 'Richard Lincoln, r.w.lincoln@gmail.com'

""" This example demonstrates how to produce a publication plot of generator
costs using matplotlib. """

import matplotlib
matplotlib.use('WXAgg')#'TkAgg')

matplotlib.rc('font', **{'family': 'sans-serif',
                         'sans-serif': ['Computer Modern Sans serif']})
matplotlib.rc('font', **{'family': 'serif', 'serif': ['Computer Modern Roman']})
matplotlib.rc('text', usetex=True)

from os.path import join, dirname

import scipy
import pylab
import pylon

DATA_DIR = join(dirname(pylon.case.__file__), "test", "data")
CASE = join(DATA_DIR, "case24_ieee_rts", "case24_ieee_rts.pkl")

# Set up publication quality graphs.
#fig_width_pt = 246.0  # Get this from LaTeX using \showthe\columnwidth
#inches_per_pt = 1.0 / 72.27               # Convert pt to inch
golden_mean = (pylab.sqrt(5) - 1.0) / 2.0 # Aesthetic ratio
fig_width = 5.0#fig_width_pt * inches_per_pt  # width in inches
fig_height = fig_width * golden_mean      # height in inches
fig_size = [fig_width, fig_height]
params = {'backend': 'ps',
          'axes.labelsize': 10,
          'text.fontsize': 10,
          'legend.fontsize': 10,
          'xtick.labelsize': 8,
          'ytick.labelsize': 8,
          'text.usetex': True,
          'figure.figsize': fig_size}
pylab.rcParams.update(params)

case = pylon.Case.load(CASE)

g = [0,2,8,11,15,20,22,24,32]

style = ['-', '--', '-.', ':']
ns = len(style)

pylab.figure()
plots = []
for i in g:
    generator = case.generators[i]
    if generator.pcost_model == pylon.PW_LINEAR:
        x = [x for x, _ in generator.p_cost]
        y = [y for _, y in generator.p_cost]
    elif generator.pcost_model == pylon.POLYNOMIAL:
        x = scipy.arange(generator.p_min, generator.p_max, 0.1)
        y = scipy.polyval(scipy.array(generator.p_cost), x)
    else:
        raise
    plots.append(pylab.plot(x, y, style[i % ns]))
    pylab.xlabel("P (MW)")
    pylab.ylabel("Cost")
pylab.legend(plots, ["U%.0f" % case.generators[i].p_max for i in g])
#pylab.show()

pylab.savefig('/tmp/fig1.pdf')
<commit_msg>Fixing cost plot styles.<commit_after>__author__ = 'Richard Lincoln, r.w.lincoln@gmail.com'

""" This example demonstrates how to produce a publication plot of generator
costs using matplotlib. """

import matplotlib
#matplotlib.use('WXAgg')#'TkAgg')

#matplotlib.rc('font', **{'family': 'sans-serif',
#                         'sans-serif': ['Computer Modern Sans serif']})
matplotlib.rc('font', **{'family': 'serif', 'serif': ['Computer Modern Roman']})
matplotlib.rc('text', usetex=True)

from os.path import join, dirname

import scipy
import pylab
import pylon

DATA_DIR = join(dirname(pylon.case.__file__), "test", "data")
CASE = join(DATA_DIR, "case24_ieee_rts", "case24_ieee_rts.pkl")

# Set up publication quality graphs.
#fig_width_pt = 246.0  # Get this from LaTeX using \showthe\columnwidth
#inches_per_pt = 1.0 / 72.27               # Convert pt to inch
golden_mean = (pylab.sqrt(5) - 1.0) / 2.0 # Aesthetic ratio
fig_width = 5.5#fig_width_pt * inches_per_pt  # width in inches
fig_height = fig_width * golden_mean      # height in inches
fig_size = [fig_width, fig_height]
params = {'backend': 'ps',
          'axes.labelsize': 10,
          'text.fontsize': 10,
          'legend.fontsize': 8,
          'xtick.labelsize': 8,
          'ytick.labelsize': 8,
          'text.usetex': True,
#          'markup': 'tex',
#          'text.latex.unicode': True,
          'figure.figsize': fig_size}
pylab.rcParams.update(params)

case = pylon.Case.load(CASE)

#g = [0,2,8,11,15,20,22,24,32]
g = [15,0,24,2,8,20,11,32,22]

style = [('black', '-'), ('0.5', '-'), ('black', ':'), ('0.5', ':'),
    ('black', '--'), ('0.5', '--'), ('black', '-.'), ('0.5', '-.'), ('0.8', '-')]
ns = len(style)

pylab.figure()
plots = []
for i, gi in enumerate(g):
    generator = case.generators[gi]
    if generator.pcost_model == pylon.PW_LINEAR:
        x = [x for x, _ in generator.p_cost]
        y = [y for _, y in generator.p_cost]
    elif generator.pcost_model == pylon.POLYNOMIAL:
        x = scipy.arange(generator.p_min, generator.p_max, 0.1)
        y = scipy.polyval(scipy.array(generator.p_cost), x)
    else:
        raise
    clr, ls = style[i]
    plots.append(pylab.plot(x, y, linestyle=ls, color=clr))
    pylab.xlabel("P (MW)")
    pylab.ylabel("Cost (USD/h)")
pylab.legend(plots, ["U%.0f" % case.generators[i].p_max for i in g])

pylab.annotate("Oil", (90, 7000))
pylab.annotate("Coal", (170, 3600))
pylab.annotate("Nuclear", (210, 1800))

#pylab.show()

pylab.savefig('/tmp/fig1.pdf')
<|endoftext|>


<commit_before>from packtools.sps.utils.file_utils import (
    get_files_list_filtered,
    get_file_content_from_zip,
)
from packtools.sps.utils.xml_utils import get_xml_tree
from packtools.sps.validation import erratum
from packtools.sps.models.article_and_subarticles import ArticleAndSubArticles

class PackageErratumHasUnexpectedQuantityOfXMLFilesError(Exception):
    ...

class PackageErratumHasNoErrataXMLFileError(Exception):
    ...

class PackageErratumHasNoArticleXMLFileError(Exception):
    ...

class Package:
    def __init__(self, zip_path):
        self.zip_path = zip_path
        self.xmltree_article = None

class PackageWithErrata(Package):
    def __init__(self, zip_path, errata_types=['correction']):
        super().__init__(zip_path)

        self.xmltree_errata = None
        self.errata_types = errata_types
        self.discover_errata_and_article_xmls()

    def discover_errata_and_article_xmls(self):
        xmls = get_files_list_filtered(self.zip_path, ['.xml'])
        if len(xmls) != 2:
            raise PackageErratumHasUnexpectedQuantityOfXMLFilesError()

        for file_name in xmls:
            x_file_content = get_file_content_from_zip(file_name, self.zip_path)
            x_tree = get_xml_tree(x_file_content)
            x_article_and_subarticles = ArticleAndSubArticles(x_tree)

            if x_article_and_subarticles.main_article_type in self.errata_types:
                self.xmltree_errata = x_tree
            else:
                self.xmltree_article = x_tree

        if self.xmltree_article is None:
            raise PackageErratumHasNoArticleXMLFileError()

        if self.xmltree_errata is None:
            raise PackageErratumHasNoErrataXMLFileError()

    def is_valid(self):
        return erratum.has_compatible_errata_and_document(self.xmltree_errata, self.xmltree_article)

class PackageArticle(Package):
    def __init__(self, zip_path):
        super().__init__(zip_path)
        self.discover_article_xml()

    def discover_article_xml(self):
        xmls = get_files_list_filtered(self.zip_path, ['.xml'])
        if len(xmls) != 1:
            raise PackageErratumHasUnexpectedQuantityOfXMLFilesError()

        x_file_content = get_file_content_from_zip(xmls.pop(), self.zip_path)
        self.xmltree_article = get_xml_tree(x_file_content)
<commit_msg>Cria modelo para obter nome canônico de pacote associado a XMLTree<commit_after>from packtools.sps.utils.file_utils import (
    get_files_list_filtered,
    get_file_content_from_zip,
)
from packtools.sps.utils.xml_utils import get_xml_tree
from packtools.sps.validation import erratum
from packtools.sps.models.article_and_subarticles import ArticleAndSubArticles
from packtools.sps.models.article_doi_with_lang import DoiWithLang
from packtools.sps.models.front_articlemeta_issue import ArticleMetaIssue
from packtools.sps.models.front_journal_meta import ISSN, Acronym

class PackageErratumHasUnexpectedQuantityOfXMLFilesError(Exception):
    ...

class PackageErratumHasNoErrataXMLFileError(Exception):
    ...

class PackageErratumHasNoArticleXMLFileError(Exception):
    ...

class Package:
    def __init__(self, zip_path):
        self.zip_path = zip_path
        self.xmltree_article = None

class PackageWithErrata(Package):
    def __init__(self, zip_path, errata_types=['correction']):
        super().__init__(zip_path)

        self.xmltree_errata = None
        self.errata_types = errata_types
        self.discover_errata_and_article_xmls()

    def discover_errata_and_article_xmls(self):
        xmls = get_files_list_filtered(self.zip_path, ['.xml'])
        if len(xmls) != 2:
            raise PackageErratumHasUnexpectedQuantityOfXMLFilesError()

        for file_name in xmls:
            x_file_content = get_file_content_from_zip(file_name, self.zip_path)
            x_tree = get_xml_tree(x_file_content)
            x_article_and_subarticles = ArticleAndSubArticles(x_tree)

            if x_article_and_subarticles.main_article_type in self.errata_types:
                self.xmltree_errata = x_tree
            else:
                self.xmltree_article = x_tree

        if self.xmltree_article is None:
            raise PackageErratumHasNoArticleXMLFileError()

        if self.xmltree_errata is None:
            raise PackageErratumHasNoErrataXMLFileError()

    def is_valid(self):
        return erratum.has_compatible_errata_and_document(self.xmltree_errata, self.xmltree_article)

class PackageArticle(Package):
    def __init__(self, zip_path):
        super().__init__(zip_path)
        self.discover_article_xml()

    def discover_article_xml(self):
        xmls = get_files_list_filtered(self.zip_path, ['.xml'])
        if len(xmls) != 1:
            raise PackageErratumHasUnexpectedQuantityOfXMLFilesError()

        x_file_content = get_file_content_from_zip(xmls.pop(), self.zip_path)
        self.xmltree_article = get_xml_tree(x_file_content)

class PackageName:
    def __init__(self, xmltree):
        self.xmltree = xmltree

    @property
    def name(self):
        dwl = DoiWithLang(self.xmltree)
        _doi = dwl.main_doi and dwl.main_doi.split("/")[-1]

        ami = ArticleMetaIssue(self.xmltree)
        _fpage = ami.fpage
        if ami.fpage_seq:
            _fpage += ami.fpage_seq
        last_item = str(_fpage or ami.elocation_id or ami.order or _doi).zfill(5)

        issn = ISSN(self.xmltree)
        acron = Acronym(self.xmltree)
        data = (
            issn.epub or issn.ppub,
            acron.text,
            ami.volume,
            ami.number,
            ami.suppl and f"s{ami.suppl}",
            last_item,
        )
        return "-".join([str(item).zfill(2) for item in data if item])
<|endoftext|>


<commit_before>import datetime
import multiprocessing

import json
import time
import serial
from ..utils.logger import get_logger
from ..utils.config import load_config
from ..utils.messaging import PanMessaging

class EnvironmentalMonitor(object):

    """
    Sets and logs an environmental monitor

    Creates a serial data reader that is connected to a serial port and puts the
    reading of that serial line into a loop in a separate process. Also sets up a
    loop that will monitor the value on that serial line and send out a zmq message.

    Args:
        name (str):                 Name for the process that runs the serial reader, defaults to
                                    'Generic Sensor'
        connect_on_startup(bool):   Whether monitor should start on creation, defaults to True
    """

    def __init__(self, config=None, messaging=None, name="Generic Sensor", connect_on_startup=False):
        self.logger = get_logger(self)
        self.config = load_config()
        assert self.config is not None, self.logger.warning("Config not set for environmental monitor")

        self._sleep_interval = 1
        self._is_running = False

        self.serial_readers = dict()

        config = self.config.get('environment', {})

        for sensor in config.keys():

            self.serial_port = config[sensor].get('serial_port')
            self.name = sensor
            self.logger.debug("Creating sensor on port: {} {}".format(self.name, self.serial_port))

            try:
                # Create the actual reader
                serial_reader = serial.SerialData(port=self.serial_port, name=sensor)
                self.serial_readers[sensor] = serial_reader

            except:
                self.logger.warning("Cannot connect to environmental sensor")

        if connect_on_startup:
            self.start_monitoring()

        # Set up ZMQ publisher
        if messaging is None:
            messaging = PanMessaging()

        self.messaging = messaging
        self.publisher = multiprocessing.Process(target=self.get_reading)
        self.publisher.name = "PANOPTES_environment"
        self.publisher.daemon = True
        self.publisher.start()

    def start_monitoring(self):
        """
        Connects over the serial port and starts the thread listening
        """
        for sensor, serial_reader in self.serial_readers.items():
            self.logger.info("Starting {} monitoring".format(sensor))

            try:
                serial_reader.connect()
                self.is_running = True
                serial_reader.start()

            except:
                self.logger.warning("Cannot connect to monitor via serial port")

    def stop_monitoring(self):
        """ Stops the monitor """
        self.logger.info("Stopping environmental monitoring")
        self.is_running = False

    @property
    def is_running(self):
        return self._is_running

    @is_running.setter
    def is_running(self, value):
        self._is_running = value

    def get_reading(self):
        """ Gets a reading from the sensor

        This is run as a loop, which gets the current serial reading, loads the
        value up as a JSON message, and then inserts into the mongo db. The `current`
        reading for the sensor is also updated and a message is sent out via zmq
        """
        self.logger.debug("Starting get_reading loop for {}".format(self.name))
        while True and self.is_running:

            for name, serial_reader in self.serial_readers.items():
                sensor_value = serial_reader.get_reading()

                sensor_data = dict()

                # Parse the value as JSON
                if len(sensor_value) > 0:
                    try:
                        sensor_data = json.loads(sensor_value.replace('nan', 'null'))
                    except ValueError:
                        print("Bad JSON: {0}".format(sensor_value))

                # Create a message object
                message = {
                    "date": datetime.datetime.utcnow(),
                    "type": self.name,
                    "data": sensor_data
                }

                # Insert message in mongo
                self.logger.debug("Inserting data to mongo")
                self.sensors.insert(message)

                # Update the 'current' reading
                self.logger.debug("Updating the 'current' value in mongo")
                self.sensors.update(
                    {"status": "current", 'type': self.name},
                    {"$set": {
                        "date": datetime.datetime.utcnow(),
                        "type": self.name,
                        "data": sensor_data
                    }},
                    True
                )

                # Send out on ZMQ
                self.messaging.send_message(channel=self.name, message=message)

            self.logger.debug("Sleeping for {} seconds".format(self._sleep_interval))
            time.sleep(self._sleep_interval)
<commit_msg>Import module properly<commit_after>import datetime
import multiprocessing

import json
import time
from ..utils.logger import get_logger
from ..utils.config import load_config
from ..utils.messaging import PanMessaging
from ..utils.rs232 import SerialData

class EnvironmentalMonitor(object):

    """
    Sets and logs an environmental monitor

    Creates a serial data reader that is connected to a serial port and puts the
    reading of that serial line into a loop in a separate process. Also sets up a
    loop that will monitor the value on that serial line and send out a zmq message.

    Args:
        name (str):                 Name for the process that runs the serial reader, defaults to
                                    'Generic Sensor'
        connect_on_startup(bool):   Whether monitor should start on creation, defaults to True
    """

    def __init__(self, config=None, messaging=None, name="Generic Sensor", connect_on_startup=False):
        self.logger = get_logger(self)
        self.config = load_config()
        assert self.config is not None, self.logger.warning("Config not set for environmental monitor")

        self._sleep_interval = 1
        self._is_running = False

        self.serial_readers = dict()

        config = self.config.get('environment', {})

        for sensor in config.keys():

            self.serial_port = config[sensor].get('serial_port')
            self.name = sensor
            self.logger.debug("Creating sensor on port: {} {}".format(self.name, self.serial_port))

            try:
                # Create the actual reader
                serial_reader = SerialData(port=self.serial_port, name=sensor)
                self.serial_readers[sensor] = serial_reader

            except:
                self.logger.warning("Cannot connect to environmental sensor")

        if connect_on_startup:
            self.start_monitoring()

        # Set up ZMQ publisher
        if messaging is None:
            messaging = PanMessaging()

        self.messaging = messaging
        self.publisher = multiprocessing.Process(target=self.get_reading)
        self.publisher.name = "PANOPTES_environment"
        self.publisher.daemon = True
        self.publisher.start()

    def start_monitoring(self):
        """
        Connects over the serial port and starts the thread listening
        """
        for sensor, serial_reader in self.serial_readers.items():
            self.logger.info("Starting {} monitoring".format(sensor))

            try:
                serial_reader.connect()
                self.is_running = True
                serial_reader.start()

            except:
                self.logger.warning("Cannot connect to monitor via serial port")

    def stop_monitoring(self):
        """ Stops the monitor """
        self.logger.info("Stopping environmental monitoring")
        self.is_running = False

    @property
    def is_running(self):
        return self._is_running

    @is_running.setter
    def is_running(self, value):
        self._is_running = value

    def get_reading(self):
        """ Gets a reading from the sensor

        This is run as a loop, which gets the current serial reading, loads the
        value up as a JSON message, and then inserts into the mongo db. The `current`
        reading for the sensor is also updated and a message is sent out via zmq
        """
        self.logger.debug("Starting get_reading loop for {}".format(self.name))
        while True and self.is_running:

            for name, serial_reader in self.serial_readers.items():
                sensor_value = serial_reader.get_reading()

                sensor_data = dict()

                # Parse the value as JSON
                if len(sensor_value) > 0:
                    try:
                        sensor_data = json.loads(sensor_value.replace('nan', 'null'))
                    except ValueError:
                        print("Bad JSON: {0}".format(sensor_value))

                # Create a message object
                message = {
                    "date": datetime.datetime.utcnow(),
                    "type": self.name,
                    "data": sensor_data
                }

                # Insert message in mongo
                self.logger.debug("Inserting data to mongo")
                self.sensors.insert(message)

                # Update the 'current' reading
                self.logger.debug("Updating the 'current' value in mongo")
                self.sensors.update(
                    {"status": "current", 'type': self.name},
                    {"$set": {
                        "date": datetime.datetime.utcnow(),
                        "type": self.name,
                        "data": sensor_data
                    }},
                    True
                )

                # Send out on ZMQ
                self.messaging.send_message(channel=self.name, message=message)

            self.logger.debug("Sleeping for {} seconds".format(self._sleep_interval))
            time.sleep(self._sleep_interval)
<|endoftext|>


<commit_before>from __future__ import unicode_literals, division, absolute_import
import glob
import logging
import os
import urllib
import urlparse
import flexget.bencode

from sqlalchemy import Column, Integer, String, ForeignKey, or_, and_, select, update, DateTime, Boolean, Unicode
from sqlalchemy.orm.exc import NoResultFound

from datetime import datetime, date, time
from flexget import db_schema, plugin
from flexget.entry import Entry
from flexget.event import event
from flexget.task import TaskAbort
from flexget.utils import requests
from flexget.utils.database import quality_requirement_property, with_session
from flexget.utils.sqlalchemy_utils import table_exists, table_schema
from flexget.utils.tools import parse_timedelta
from sqlalchemy.ext.hybrid import Comparator, hybrid_property
from sqlalchemy.orm import relation, backref, relationship
from sqlalchemy.schema import Table

try:
    from flexget.plugins.filter import queue_base
except ImportError:
    raise plugin.DependencyError(issued_by='subtitle_queue', missing='queue_base',
                                 message='subtitle_queue requires the queue_base plugin')
try:
    from flexget.plugins.metainfo.subtitles_check import MetainfoSubs
except ImportError:
    raise plugin.DependencyError(issued_by='subtitle_queue', missing='subtitles_check',
                                 message='subtitle_queue requires way of checking for subtitles')
try:
    from babelfish import Language
except ImportError:
    raise plugin.DependencyError(issued_by='subtitle_queue', missing='babelfish',
                                 message='subtitle_queue requires the babelfish plugin')

log = logging.getLogger('subtitle_queue')
Base = db_schema.versioned_base('subtitle_queue', 0)

class NormalizedComparator(Comparator):
    def operate(self, op, other):
        return op(self.__clause_element__(), os.path.normcase(os.path.normpath(other)))

class LangComparator(Comparator):
    def operate(self, op, other):
        return op(self.__clause_element__(), other)

association_table = Table('association', Base.metadata,
                          Column('sub_queue_id', Integer, ForeignKey('subtitle_queue.id')),
                          Column('lang_id', Integer, ForeignKey('subtitle_language.id'))
                          )

class SubtitleLanguages(Base):
    __tablename__ = 'subtitle_language'

    id = Column(Integer, primary_key=True)
    _language = Column('language', Unicode, unique=True, index=True)
    #queued_sub_id = Column(Integer, ForeignKey('subtitle_queue.id'))

    def name_setter(self, value):
        self._language = value

    def name_getter(self):
        return Language.fromietf(self._language)

    def lang_comparator(self):
        return LangComparator(self._language)

    language = hybrid_property(name_getter, name_setter)
    language.comparator(lang_comparator)

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

    def __str__(self):
        return '<SubtitleLanguage(%s)>' % (self.language)

class QueuedSubtitle(Base):

    __tablename__ = 'subtitle_queue'

    id = Column(Integer, primary_key=True)
    title = Column(String)  # Not completely necessary
    _path = Column('path', Unicode)  # Absolute path of file to fetch subtitles for
    _path_normalized = Column('path_normalized', Unicode, index=True, unique=True)
    _alternate_path = Column('alternate_path', Unicode)  # Absolute path of file to fetch subtitles for
    _alternate_path_normalized = Column('alternate_path_normalized', Unicode, index=True, unique=True)
    #imdb_id = Column(String)  # not sure this is strictly necessary
    added = Column(DateTime)  # Used to determine age
    stop_after = Column(String)
    downloaded = Column(Boolean)
    languages = relationship(SubtitleLanguages, secondary=association_table, backref="primary")

    def name_setter(self, value):
        self._path = value
        self._path_normalized = os.path.normcase(os.path.normpath(value))

    def name_getter(self):
        return self._path

    def name_comparator(self):
        return NormalizedComparator(self._path_normalized)

    path = hybrid_property(name_getter, name_setter)
    path.comparator(name_comparator)

    def alt_setter(self, value):
        self._alternate_path = value
        self._alternate_path_normalized = os.path.normcase(os.path.normpath(value))

    def alt_getter(self):
        return self._alternate_path

    def alt_comparator(self):
        return NormalizedComparator(self._alternate_path_normalized)

    alternate_path = hybrid_property(alt_getter, alt_setter)
    alternate_path.comparator(alt_comparator)

    def __init__(self, path, alternate_path, title, stop_after="7 days"):
        self.path = path
        self.alternate_path = alternate_path
        self.added = datetime.now()
        self.stop_after = stop_after
        self.title = title
        self.downloaded = False

    def __str__(self):
        return '<SubtitleQueue(%s, %s, %s)>' % (self.path, self.added, self.languages[0])

class SubtitleQueue(object):
    schema = {
        "oneOf": [
            {
                'type': 'object',
                'properties': {
                    'action': {'type': 'string', 'enum': ['add', 'remove', 'emit']},
                    'stop_after': {'type': 'string', 'format': 'interval'},
                    'languages': {"oneOf": [
                        {'type': 'array', 'items': {'type': 'string'}, 'minItems': 1},
                        {'type': 'string'},
                    ]},
                },
                'required': ['action'],
                'additionalProperties': False
            },
        ]
    }

    def emit(self, task, config):
        if not config:
            return
        subtitles = queue_get(session=task.session)
        entries = []
        for sub_item in subtitles:
            entry = Entry()
            if os.path.exists(sub_item.path):
                path = sub_item.path
            elif sub_item.alternate_path and os.path.exists(sub_item.alternate_path):
                path = sub_item.alternate_path
            elif sub_item.path:
                log.debug('%s: File not found.' % sub_item.path)
                continue
            elif sub_item.alternate_path:
                log.debug('%s: File not found.' % sub_item.alternate_path)
                continue
            else:
                log.error('Bad entry. No paths specified.')
                continue
            entry['url'] = urlparse.urljoin('file:', urllib.pathname2url(path))
            entry['location'] = path
            entry['title'] = sub_item.title

            primary = []
            for language in sub_item.languages:
                primary.append(language.language)
            entry['subtitle_languages'] = primary

            index = path.rfind('.')
            for lang in primary:
                if not glob.glob(path[:index] + "." + str(lang) + ".srt"):
                    break
            else:
                log.debug('All primary subtitles already fetched for %s.' % entry['title'])
                continue
            entries.append(entry)
            log.debug('Emitted subtitle entry for %s.' % entry['title'])
        return entries

    def on_task_filter(self, task, config):
        if config:
            for entry in task.entries:
                entry.accept()

    def on_task_input(self, task, config):
        if not config:
            return
        return self.emit(task, config)

    def on_task_output(self, task, config):
        if not config:
            return
        action = config.get('action')
        if action == 'emit':
            return
        for entry in task.accepted:
            try:
                if action == "add":
                    # is it a local file?
                    if os.path.exists(entry.get('location', '')):
                        src = entry.get('location', '')
                        dest = entry.get('output', '')
                        if dest:
                            # file has been moved and original location may already be in db
                            queue_edit(src, dest, entry.get('title', ''), config)
                        else:
                            queue_add(src, entry.get('title', ''), config)
                    else:
                        content_filename = entry.get('content_filename', '')
                        for url in entry.get('urls', ''):
                            try:
                                info = flexget.bencode.bdecode(requests.get(url).text)['info']
                                files = info['files']
                                break
                            except Exception:
                                continue
                        if not info or not files:
                            raise TaskAbort('Invalid entry. Cannot queue.')
                        if len(files) == 1:
                            title = files[0]['path'][0]  # extract file name
                            if content_filename:
                                # need to queue the content_filename
                                ext = title.rfind('.')
                                title = content_filename + title[ext:]
                        else:
                            # TODO: queue dir or specific files?
                            # paths = []
                            # for file in files:
                            #     p = file['path'][0]
                            #     # this is bad
                            #     if p.endswith(('.avi', '.mkv', '.mp4', '.mov', '.mpg', '.wmv')):
                            #         paths.append(p)
                            #

                            # set the title to be torrent name, which hopefully is download dir
                            title = info['name']
                        if entry.get('movedone', ''):
                            src = entry.get('path', '') + title
                            dest = entry.get('movedone', '') + title
                            queue_add(src, entry.get('title', ''), config, alternate_path=dest)
                        elif entry.get('path', ''):
                            queue_add(src, entry.get('title', ''), config)
                        else:
                            raise TaskAbort("Cannot queue invalid entries.")
                elif action == "remove":
                    queue_del(entry['path'])
                return
            except QueueError as e:
                # ignore already in queue
                if e.errno != 1:
                    entry.fail("Error adding file to subtitle queue: %s" % e.message)

@with_session
def queue_add(path, title, config, alternate_path='', session=None):
    if session.query(QueuedSubtitle).filter(or_(QueuedSubtitle.path == path,
                                                QueuedSubtitle.alternate_path == path)).first():
        raise QueueError("ERROR: %s is already in the queue." % path, errno=1)
    else:
        primary = make_lang_list(config.get('languages', []))
        if config.get('stop_after', ''):
            item = QueuedSubtitle(path, alternate_path, title, stop_after=config.get('stop_after'))
        else:
            item = QueuedSubtitle(path, alternate_path, title)
        session.add(item)
        item.languages = primary if primary else [SubtitleLanguages('eng')]
        log.debug('Added %s to queue with %s primary languages.' % (item.path, len(item.languages)))
        return True

@with_session
def queue_del(path, session=None):
    item = session.query(QueuedSubtitle).filter(QueuedSubtitle.path == path)
    item = item.first()
    if not item:
        # Not necessarily an error?
        raise QueueError("DEBUG: %s is not in the queue." % path)
    else:
        session.delete(item)
        return item.path

@with_session
def queue_edit(src, dest, title, config, session=None):
    item = session.query(QueuedSubtitle).filter(QueuedSubtitle.path == src).first()
    if not item:
        log.debug("DEBUG: %s was moved but is not in the queue." % src)
        queue_add(dest, title, config, alternate_path=src, session=session)
        #raise QueueError("ERROR: %s is not in the queue." % src)
    else:
        item.alternate_path = src
        item.path = dest
        stop_after = config.get('stop_after', '')
        if stop_after and item.stop_after is not stop_after:
            item.stop_after = stop_after
        primary = make_lang_list(config.get('languages', []))
        item.languages = primary

@with_session
def queue_get(session=None):
    subs = session.query(QueuedSubtitle).all()
    for sub_item in subs:
        if sub_item.added + parse_timedelta(sub_item.stop_after) < datetime.combine(date.today(), time()):
            subs.remove(sub_item)
            session.delete(sub_item)
    return subs

# must always pass the session
def get_lang(lang, session):
    l = session.query(SubtitleLanguages).filter(SubtitleLanguages.language == lang).first()
    return l

# TODO: prettify? ugly shit code fuck me
def make_lang_list(languages):
    primary = []
    if not isinstance(languages, list):
        languages = [languages]
    for language in languages:
        l = get_lang(language)
        if l:
            primary.append(l)
        else:
            primary.append(SubtitleLanguages(language))
    return primary

class QueueError(Exception):
    """Exception raised if there is an error with a queue operation"""

    # TODO: taken from movie_queue
    def __init__(self, message, errno=0):
        self.message = message
        self.errno = errno

@event('plugin.register')
def register_plugin():
    plugin.register(SubtitleQueue, 'subtitle_queue', api_ver=2)<commit_msg>fixed bugs and changed some logic<commit_after>from __future__ import unicode_literals, division, absolute_import
import glob
import logging
import os
import urllib
import urlparse
import flexget.bencode

from sqlalchemy import Column, Integer, String, ForeignKey, or_, and_, select, update, DateTime, Boolean, Unicode
from sqlalchemy.orm.exc import NoResultFound

from datetime import datetime, date, time
from flexget import db_schema, plugin
from flexget.entry import Entry
from flexget.event import event
from flexget.task import TaskAbort
from flexget.utils import requests
from flexget.utils.database import quality_requirement_property, with_session
from flexget.utils.sqlalchemy_utils import table_exists, table_schema
from flexget.utils.tools import parse_timedelta
from sqlalchemy.ext.hybrid import Comparator, hybrid_property
from sqlalchemy.orm import relation, backref, relationship
from sqlalchemy.schema import Table

try:
    from flexget.plugins.filter import queue_base
except ImportError:
    raise plugin.DependencyError(issued_by='subtitle_queue', missing='queue_base',
                                 message='subtitle_queue requires the queue_base plugin')
try:
    from flexget.plugins.metainfo.subtitles_check import MetainfoSubs
except ImportError:
    raise plugin.DependencyError(issued_by='subtitle_queue', missing='subtitles_check',
                                 message='subtitle_queue requires way of checking for subtitles')
try:
    from babelfish import Language
except ImportError:
    raise plugin.DependencyError(issued_by='subtitle_queue', missing='babelfish',
                                 message='subtitle_queue requires the babelfish plugin')

log = logging.getLogger('subtitle_queue')
Base = db_schema.versioned_base('subtitle_queue', 0)

class NormalizedComparator(Comparator):
    def operate(self, op, other):
        return op(self.__clause_element__(), os.path.normcase(os.path.normpath(other)))

class LangComparator(Comparator):
    def operate(self, op, other):
        return op(self.__clause_element__(), other)

association_table = Table('association', Base.metadata,
                          Column('sub_queue_id', Integer, ForeignKey('subtitle_queue.id')),
                          Column('lang_id', Integer, ForeignKey('subtitle_language.id'))
                          )

class SubtitleLanguages(Base):
    __tablename__ = 'subtitle_language'

    id = Column(Integer, primary_key=True)
    _language = Column('language', Unicode, unique=True, index=True)
    #queued_sub_id = Column(Integer, ForeignKey('subtitle_queue.id'))

    def name_setter(self, value):
        self._language = value

    def name_getter(self):
        return Language.fromietf(self._language)

    def lang_comparator(self):
        return LangComparator(self._language)

    language = hybrid_property(name_getter, name_setter)
    language.comparator(lang_comparator)

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

    def __str__(self):
        return '<SubtitleLanguage(%s)>' % (self.language)

class QueuedSubtitle(Base):

    __tablename__ = 'subtitle_queue'

    id = Column(Integer, primary_key=True)
    title = Column(String)  # Not completely necessary
    _path = Column('path', Unicode)  # Absolute path of file to fetch subtitles for
    _path_normalized = Column('path_normalized', Unicode, index=True, unique=True)
    _alternate_path = Column('alternate_path', Unicode)  # Absolute path of file to fetch subtitles for
    _alternate_path_normalized = Column('alternate_path_normalized', Unicode, index=True, unique=True)
    #imdb_id = Column(String)  # not sure this is strictly necessary
    added = Column(DateTime)  # Used to determine age
    stop_after = Column(String)
    downloaded = Column(Boolean)
    languages = relationship(SubtitleLanguages, secondary=association_table, backref="primary")

    def name_setter(self, value):
        self._path = value
        self._path_normalized = os.path.normcase(os.path.normpath(value))

    def name_getter(self):
        return self._path

    def name_comparator(self):
        return NormalizedComparator(self._path_normalized)

    path = hybrid_property(name_getter, name_setter)
    path.comparator(name_comparator)

    def alt_setter(self, value):
        self._alternate_path = value
        self._alternate_path_normalized = os.path.normcase(os.path.normpath(value)) if value is not None else None

    def alt_getter(self):
        return self._alternate_path

    def alt_comparator(self):
        return NormalizedComparator(self._alternate_path_normalized)

    alternate_path = hybrid_property(alt_getter, alt_setter)
    alternate_path.comparator(alt_comparator)

    def __init__(self, path, alternate_path, title, stop_after="7 days"):
        self.path = path
        self.alternate_path = alternate_path
        self.added = datetime.now()
        self.stop_after = stop_after
        self.title = title
        self.downloaded = False

    def __str__(self):
        return '<SubtitleQueue(%s, %s, %s)>' % (self.path, self.added, self.languages[0])

class SubtitleQueue(object):
    schema = {
        "oneOf": [
            {
                'type': 'object',
                'properties': {
                    'action': {'type': 'string', 'enum': ['add', 'remove', 'emit']},
                    'stop_after': {'type': 'string', 'format': 'interval'},
                    'languages': {"oneOf": [
                        {'type': 'array', 'items': {'type': 'string'}, 'minItems': 1},
                        {'type': 'string'},
                    ]},
                },
                'required': ['action'],
                'additionalProperties': False
            },
        ]
    }

    def emit(self, task, config):
        if not config:
            return
        subtitles = queue_get(session=task.session)
        entries = []
        for sub_item in subtitles:
            entry = Entry()
            if os.path.exists(sub_item.path):
                path = sub_item.path
            elif sub_item.alternate_path and os.path.exists(sub_item.alternate_path):
                path = sub_item.alternate_path
            elif sub_item.path:
                log.debug('%s: File not found.' % sub_item.path)
                continue
            elif sub_item.alternate_path:
                log.debug('%s: File not found.' % sub_item.alternate_path)
                continue
            else:
                log.error('Bad entry. No paths specified.')
                continue
            entry['url'] = urlparse.urljoin('file:', urllib.pathname2url(path))
            entry['location'] = path
            entry['title'] = sub_item.title

            primary = []
            for language in sub_item.languages:
                primary.append(language.language)
            entry['subtitle_languages'] = primary

            index = path.rfind('.')
            if len(primary) > 1:
                for lang in primary:
                    if not glob.glob(path[:index] + "." + str(lang) + ".srt"):
                        break
                else:
                    log.debug('All primary subtitles already fetched for %s.' % entry['title'])
                    continue
            else:
                if glob.glob(path[:index] + ".srt"):
                    log.debug('All primary subtitles already fetched for %s.' % entry['title'])
                    continue
            entries.append(entry)
            log.debug('Emitted subtitle entry for %s.' % entry['title'])
        return entries

    def on_task_filter(self, task, config):
        if config:
            for entry in task.entries:
                entry.accept()

    def on_task_input(self, task, config):
        if not config:
            return
        return self.emit(task, config)

    def on_task_output(self, task, config):
        if not config:
            return
        action = config.get('action')
        if action == 'emit':
            return
        for entry in task.accepted:
            try:
                if action == "add":
                    # is it a local file?
                    if os.path.isfile(entry.get('location', '')):
                        src = entry.get('location', '')
                        dest = entry.get('output', '')
                        if dest:
                            # file has been moved and original location may already be in db
                            queue_edit(src, dest, entry.get('title', ''), config)
                        else:
                            queue_add(src, entry.get('title', ''), config)
                    elif entry.get('url', '').endswith('.torrent'):
                        content_filename = entry.get('content_filename', '')
                        info = None
                        files = None
                        for url in entry.get('urls', ''):
                            try:
                                info = flexget.bencode.bdecode(requests.get(url).text)['info']
                                files = info['files']
                                break
                            except Exception:
                                continue
                        if not info or not files:
                            raise TaskAbort('Invalid entry. Cannot queue.')
                        if len(files) == 1:
                            title = files[0]['path'][0]  # extract file name
                            if content_filename:
                                # need to queue the content_filename
                                ext = title.rfind('.')
                                title = content_filename + title[ext:]
                        else:
                            # TODO: queue dir or specific files?
                            # paths = []
                            # for file in files:
                            #     p = file['path'][0]
                            #     # this is bad
                            #     if p.endswith(('.avi', '.mkv', '.mp4', '.mov', '.mpg', '.wmv')):
                            #         paths.append(p)
                            #

                            # set the title to be torrent name, which hopefully is download dir
                            title = info['name']
                        if entry.get('movedone', ''):
                            src = entry.get('path', '') + title
                            dest = entry.get('movedone', '') + title
                            queue_add(src, entry.get('title', ''), config, alternate_path=dest)
                        elif entry.get('path', ''):
                            queue_add(entry.get('path'), entry.get('title', ''), config)
                        else:
                            raise TaskAbort("Cannot queue invalid entries.")
                    else:
                        raise TaskAbort("Invalid entry for subtitle queue.")
                elif action == "remove":
                    if entry.get('location', ''):
                        queue_del(entry['location'])
                    else:
                        raise TaskAbort("Cannot delete non-local files.")
                return
            except QueueError as e:
                # ignore already in queue
                if e.errno != 1:
                    entry.fail("Error adding file to subtitle queue: %s" % e.message)

@with_session
def queue_add(path, title, config, alternate_path=None, session=None):
    path = unicode(path)
    item = session.query(QueuedSubtitle).filter(or_(QueuedSubtitle.path == path,
                                                    QueuedSubtitle.alternate_path == path)).first()
    primary = make_lang_list(config.get('languages', []), session=session)
    default = [SubtitleLanguages('eng')]
    if item:
        log.debug('Already queued. Updating values.')
        if item.path == path and alternate_path:
            item.alternate_path = alternate_path
        elif item.alternate_path == path and alternate_path:
            item.path = alternate_path
            #raise QueueError("ERROR: %s is already in the queue." % path, errno=1)
        item.languages = primary if primary else default
    else:
        if config.get('stop_after', ''):
            item = QueuedSubtitle(path, alternate_path, title, stop_after=config.get('stop_after'))
        else:
            item = QueuedSubtitle(path, alternate_path, title)
        session.add(item)
        item.languages = primary if primary else default
        log.debug('Added %s to queue with %s primary languages.' % (item.path, len(item.languages)))
        return True

@with_session
def queue_del(path, session=None):
    item = session.query(QueuedSubtitle).filter(or_(QueuedSubtitle.path == path,
                                                    QueuedSubtitle.alternate_path == path)).first()
    if not item:
        # Not necessarily an error?
        raise QueueError("DEBUG: %s is not in the queue." % path)
    else:
        session.delete(item)
        return item.path

@with_session
def queue_edit(src, dest, title, config, session=None):
    item = session.query(QueuedSubtitle).filter(QueuedSubtitle.path == src).first()
    if not item:
        log.debug("DEBUG: %s was moved but is not in the queue." % src)
        queue_add(dest, title, config, alternate_path=src, session=session)
        #raise QueueError("ERROR: %s is not in the queue." % src)
    else:
        item.alternate_path = src
        item.path = dest
        stop_after = config.get('stop_after', '')
        if stop_after and item.stop_after is not stop_after:
            item.stop_after = stop_after
        primary = make_lang_list(config.get('languages', []), session=session)
        item.languages = primary

@with_session
def queue_get(session=None):
    subs = session.query(QueuedSubtitle).all()
    for sub_item in subs:
        if sub_item.added + parse_timedelta(sub_item.stop_after) < datetime.combine(date.today(), time()):
            subs.remove(sub_item)
            session.delete(sub_item)
    return subs

# must always pass the session
@with_session
def get_lang(lang, session=None):
    l = session.query(SubtitleLanguages).filter(SubtitleLanguages.language == unicode(lang)).first()
    return l

# TODO: prettify? ugly shit code fuck me
@with_session
def make_lang_list(languages, session=None):
    primary = []
    if not isinstance(languages, list):
        languages = [languages]
    for language in languages:
        l = get_lang(language, session=session)
        if l:
            primary.append(l)
        else:
            primary.append(SubtitleLanguages(unicode(language)))
    return primary

class QueueError(Exception):
    """Exception raised if there is an error with a queue operation"""

    # TODO: taken from movie_queue
    def __init__(self, message, errno=0):
        self.message = message
        self.errno = errno

@event('plugin.register')
def register_plugin():
    plugin.register(SubtitleQueue, 'subtitle_queue', api_ver=2)<|endoftext|>