<commit_before>#MIT License
#Copyright (c) 2017 Tim Wentzlau

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:

# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

""" Web socket ipc used by the web ui to communicate with Kervi """
import time
import inspect
import json
from kervi.spine import Spine
import kervi.utility.nethelper as nethelper
from kervi.core.authentication import Authorization
import kervi.utility.encryption as encryption
import logging

#from kervi.utility.kerviThread import KerviThread
from autobahn.asyncio.websocket import WebSocketServerProtocol

class _ObjectEncoder(json.JSONEncoder):
    def default(self, obj):
        if hasattr(obj, "to_json"):
            return self.default(obj.to_json())
        elif hasattr(obj, "__dict__"):
            data = dict(
                (key, value)
                for key, value in inspect.getmembers(obj)
                if not key.startswith("__")
                and not inspect.isabstract(value)
                and not inspect.isbuiltin(value)
                and not inspect.isfunction(value)
                and not inspect.isgenerator(value)
                and not inspect.isgeneratorfunction(value)
                and not inspect.ismethod(value)
                and not inspect.ismethoddescriptor(value)
                and not inspect.isroutine(value)
            )
            return self.default(data)
        return obj

class _WebCommandHandler(object):
    def __init__(self, command, protocol):
        self.protocol = protocol
        self.command = command
        spine = Spine()
        spine.register_command_handler(command, self.on_command, injected="socketSpine")

    def on_command(self, *args, **kwargs):
        injected = kwargs.get("injected", "")
        if not injected == "socketSpine":
            jsonres = json.dumps({"messageType":"command", "command":self.command, "args":args}, ensure_ascii=False).encode('utf8')
            self.protocol.sendMessage(jsonres, False)

class _WebQueryHandler(object):
    def __init__(self, query, protocol):
        self.protocol = protocol
        self.query = query
        spine = Spine()
        spine.register_query_handler(query, self.on_query, injected="socketSpine")

    def on_query(self, *args, **kwargs):
        injected = kwargs.get("injected", "")
        if not injected == "socketSpine":
            jsonres=json.dumps({"messageType":"query", "query":self.query, "args":args}, ensure_ascii=False).encode('utf8')
            self.protocol.sendMessage(jsonres, False)

class _WebEventHandler(object):
    def __init__(self, event, id_event, protocol):
        self.protocol = protocol
        self.event = event
        self.id_event = id_event
        self.spine = Spine()
        self.spine.register_event_handler(event, self.on_event, id_event, injected="socketSpine")

    def on_event(self, id_event, *args, **kwargs):
        injected = kwargs.get("injected", "")
        groups = kwargs.get("groups", None)
        process_id = kwargs.get("process_id", None)
        #self.spine.log.debug("WS relay event:{0} injected:{1}", self.event, injected)

        authorized = True

        if self.protocol.user != None and self.protocol.user["groups"] != None and groups != None and len(groups) > 0:
            for group in groups:
                if group in self.protocol.user["groups"]:
                    break
            else:
                authorized = False

        if authorized and self.protocol.authenticated and not injected == "socketSpine":

            cmd = {"messageType":"event", "event":self.event, "id":id_event, "args":args}
            jsonres = json.dumps(cmd, cls=_ObjectEncoder, ensure_ascii=False).encode('utf8')
            self.protocol.sendMessage(jsonres, False)

class _SpineProtocol(WebSocketServerProtocol):

    def __init__(self):
        self.spine = Spine()
        WebSocketServerProtocol.__init__(self)
        self.handlers = {"command":[], "query":[], "event":[]}
        self.authenticated = False
        self.session = None
        self.user = None
        self._authorization = Authorization()

    def add_command_handler(self, command):
        found = False
        for command_handler in self.handlers["command"]:
            if command_handler.command == command:
                found = True
        if not found:
            self.handlers["command"] += [_WebCommandHandler(command, self)]

    def add_query_handler(self, query):
        found = False
        for query_handler in self.handlers["query"]:
            if query_handler.query == query:
                found = True
        if not found:
            self.handlers["query"] += [_WebQueryHandler(query, self)]

    def add_event_handler(self, event, id_event):
        found = False
        for event_handler in self.handlers["event"]:
            if event_handler.event == event and event_handler.id_event == id_event:
                found = True
        if not found:
            self.handlers["event"] += [_WebEventHandler(event, id_event, self)]

    def send_response(self, id, response, state="ok", message=""):
        res = {
            "id":id,
            "messageType":"response",
            "state":state,
            "message":message,
            "response":response
        }
        jsonres = json.dumps(res, ensure_ascii=False).encode('utf8')
        self.sendMessage(jsonres, False)

    def onConnect(self, request):
        pass

    def onOpen(self):
        if self._authorization.active:
            res = {
                "messageType":"authenticate"
            }
        else:
            self.authenticated = True
            res = {
                "messageType":"session_authenticated",
                "session":"123456",
            }
        jsonres = json.dumps(res, ensure_ascii=False).encode('utf8')
        self.sendMessage(jsonres, False)

    def onMessage(self, payload, is_binary):
        try:
            obj = json.loads(payload.decode('utf8'))

            if obj["messageType"] == "authenticate":
                session, user = self._authorization.authorize(obj["userName"], obj["password"])

                if session is None:
                    self.spine.log.warning("authorization failed for:", obj["userName"])
                    res = {
                        "messageType":"authentication_failed",
                    }
                    #self.close()
                else:
                    self.session = session
                    self.user = user
                    self.authenticated = True
                    res = {
                        "messageType":"session_authenticated",
                        "session":session,
                    }
                jsonres = json.dumps(res, ensure_ascii=False).encode('utf8')
                self.sendMessage(jsonres, False)
            elif obj["messageType"] == "logoff":
                self.authenticated = False
                self.session = None
                self._authorization.remove_session(obj["session"])
                res = {
                    "messageType":"session_logoff"
                }
                jsonres = json.dumps(res, ensure_ascii=False).encode('utf8')
                self.sendMessage(jsonres, False)
            else:
                #self.spine.log.debug("WS onMessage:{0}", obj)
                if not self.authenticated:
                    pass
                elif obj["messageType"] == "query":
                    res = self.spine.send_query(obj["query"], *obj["args"], injected="socketSpine", session=self.user)
                    #self.spine.log.debug("query response:{0}", res)
                    self.send_response(obj["id"], res)
                elif obj["messageType"] == "registerQueryHandler":
                    self.add_query_handler(obj["query"])
                    self.send_response(None, None)
                elif obj["messageType"] == "command":
                    self.spine.send_command(obj["command"], *obj["args"], injected="socketSpine", session=self.user)
                    self.send_response(obj["id"], None)
                elif obj["messageType"] == "registerCommandHandler":
                    self.add_command_handler(obj["command"])
                    self.send_response(obj["id"], None)
                elif obj["messageType"] == "event":
                    self.spine.trigger_event(
                        obj["event"], obj["id"],
                        obj["args"],
                        injected="socketSpine"
                    )
                    self.send_response(obj["id"], None)
                elif obj["messageType"] == "registerEventHandler":
                    self.add_event_handler(obj["event"], obj["eventId"])
                    self.send_response(obj["id"], None)
        except:
            self.spine.log.exception("WS onMessage exception")
            #res={"execptionType":exc_type,"value":exc_value,"traceback":exc_traceback}
            #self.sendResponse(res,"exception")

class SocketSpine:
    def __init__(self, config):
        #coro = None
        self._started = False
        self._config = config
        self._spine = Spine()

        try:
            import asyncio
        except ImportError:
            ## Trollius >= 0.3 was renamed
            import trollius as asyncio

        from autobahn.asyncio.websocket import WebSocketServerFactory

        ssl_context = None

        if encryption.enabled():
            self._spine.log.debug("socket using ssl")
            cert_file, key_file = encryption.get_cert()
            try:
                import ssl
                ssl_context = ssl.SSLContext(ssl.PROTOCOL_SSLv23)
                ssl_context.load_cert_chain(cert_file, key_file)
                #self._spine.log.debug("socket ssl found")
            except:
                ssl_context = None
                self._spine.log.error("socket failed to use ssl")

        self._spine.log.debug(
            "start websocket on:{0}, port:{1}",
            self._config.network.ip,
            self._config.network.ws_port
        )
        self.factory = WebSocketServerFactory()
        self.factory.protocol = _SpineProtocol

        self.loop = asyncio.get_event_loop()
        self.coro = self.loop.create_server(
            self.factory,
            self._config.network.ip,
            self._config.network.ws_port,
            ssl=ssl_context
        )

    def start_socket(self):
        self._spine.log.verbose("start websocket: {0} {1} ", self._config.network.ip, self._config.network.ws_port)
        self.loop.run_until_complete(self.coro)
        self._started = True

    def step(self):
        if self._started:
            self.loop.run_until_complete(self.coro)
        #loop.run_until_complete(coro_local)
        time.sleep(.001)
<commit_msg>Fix issue with sending messages from threads<commit_after>#MIT License
#Copyright (c) 2017 Tim Wentzlau

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:

# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

""" Web socket ipc used by the web ui to communicate with Kervi """
import time
import inspect
import json
from kervi.spine import Spine
import kervi.utility.nethelper as nethelper
from kervi.core.authentication import Authorization
import kervi.utility.encryption as encryption
import logging

#from kervi.utility.kerviThread import KerviThread
from autobahn.asyncio.websocket import WebSocketServerProtocol
import asyncio

class _ObjectEncoder(json.JSONEncoder):
    def default(self, obj):
        if hasattr(obj, "to_json"):
            return self.default(obj.to_json())
        elif hasattr(obj, "__dict__"):
            data = dict(
                (key, value)
                for key, value in inspect.getmembers(obj)
                if not key.startswith("__")
                and not inspect.isabstract(value)
                and not inspect.isbuiltin(value)
                and not inspect.isfunction(value)
                and not inspect.isgenerator(value)
                and not inspect.isgeneratorfunction(value)
                and not inspect.ismethod(value)
                and not inspect.ismethoddescriptor(value)
                and not inspect.isroutine(value)
            )
            return self.default(data)
        return obj

class _WebCommandHandler(object):
    def __init__(self, command, protocol):
        self.protocol = protocol
        self.command = command
        spine = Spine()
        spine.register_command_handler(command, self.on_command, injected="socketSpine")

    def on_command(self, *args, **kwargs):
        injected = kwargs.get("injected", "")
        if not injected == "socketSpine":
            jsonres = json.dumps({"messageType":"command", "command":self.command, "args":args}, ensure_ascii=False).encode('utf8')
            self.protocol.broadcast_message(self.protocol, jsonres)

class _WebQueryHandler(object):
    def __init__(self, query, protocol):
        self.protocol = protocol
        self.query = query
        spine = Spine()
        spine.register_query_handler(query, self.on_query, injected="socketSpine")

    def on_query(self, *args, **kwargs):
        injected = kwargs.get("injected", "")
        if not injected == "socketSpine":
            jsonres=json.dumps({"messageType":"query", "query":self.query, "args":args}, ensure_ascii=False).encode('utf8')
            self.protocol.broadcast_message(self.protocol, jsonres)

class _WebEventHandler(object):
    def __init__(self, event, id_event, protocol):
        self.protocol = protocol
        self.event = event
        self.id_event = id_event
        self.spine = Spine()
        self.spine.register_event_handler(event, self.on_event, id_event, injected="socketSpine")

    def on_event(self, id_event, *args, **kwargs):
        injected = kwargs.get("injected", "")
        groups = kwargs.get("groups", None)
        process_id = kwargs.get("process_id", None)
        #self.spine.log.debug("WS relay event:{0} injected:{1}", self.event, injected)

        authorized = True

        if self.protocol.user != None and self.protocol.user["groups"] != None and groups != None and len(groups) > 0:
            for group in groups:
                if group in self.protocol.user["groups"]:
                    break
            else:
                authorized = False

        if authorized and self.protocol.authenticated and not injected == "socketSpine":

            cmd = {"messageType":"event", "event":self.event, "id":id_event, "args":args}
            jsonres = json.dumps(cmd, cls=_ObjectEncoder, ensure_ascii=False).encode('utf8')
            self.protocol.broadcast_message(self.protocol, jsonres)

class _SpineProtocol(WebSocketServerProtocol):
    loop = None
    def __init__(self):
        self.spine = Spine()
        WebSocketServerProtocol.__init__(self)
        self.handlers = {"command":[], "query":[], "event":[]}
        self.authenticated = False
        self.session = None
        self.user = None
        self._authorization = Authorization()

    @classmethod
    def broadcast_message(cls, connection, json_data):
        #for c in set(cls.connections):
        cls.loop.call_soon_threadsafe(cls.sendMessage, connection, json_data)

    def add_command_handler(self, command):
        found = False
        for command_handler in self.handlers["command"]:
            if command_handler.command == command:
                found = True
        if not found:
            self.handlers["command"] += [_WebCommandHandler(command, self)]

    def add_query_handler(self, query):
        found = False
        for query_handler in self.handlers["query"]:
            if query_handler.query == query:
                found = True
        if not found:
            self.handlers["query"] += [_WebQueryHandler(query, self)]

    def add_event_handler(self, event, id_event):
        found = False
        for event_handler in self.handlers["event"]:
            if event_handler.event == event and event_handler.id_event == id_event:
                found = True
        if not found:
            self.handlers["event"] += [_WebEventHandler(event, id_event, self)]

    def send_response(self, id, response, state="ok", message=""):
        res = {
            "id":id,
            "messageType":"response",
            "state":state,
            "message":message,
            "response":response
        }
        jsonres = json.dumps(res, ensure_ascii=False).encode('utf8')
        self.sendMessage(jsonres, False)

    def onConnect(self, request):
        pass

    def onOpen(self):
        if self._authorization.active:
            res = {
                "messageType":"authenticate"
            }
        else:
            self.authenticated = True
            res = {
                "messageType":"session_authenticated",
                "session":"123456",
            }
        jsonres = json.dumps(res, ensure_ascii=False).encode('utf8')
        self.sendMessage(jsonres, False)

    # @asyncio.coroutine
    # def async_query(self, query, query_args, injected, session):
    #     def do_req():
    #         return self.spine.send_query(query, *query_args, injected=injected, session=session)
    #         #self.spine.log.debug("query response:{0}", res)
    #         #self.send_response(obj["id"], res)
    #         #return requests.get('https://api.github.com/user', auth=HTTPBasicAuth('user', 'pass'))
    #     req = self.loop.run_in_executor(None, do_req)
    #     resp = yield from req
    #     return resp
    def onMessage(self, payload, is_binary):
        try:
            obj = json.loads(payload.decode('utf8'))

            if obj["messageType"] == "authenticate":
                session, user = self._authorization.authorize(obj["userName"], obj["password"])

                if session is None:
                    self.spine.log.warning("authorization failed for:", obj["userName"])
                    res = {
                        "messageType":"authentication_failed",
                    }
                    #self.close()
                else:
                    self.session = session
                    self.user = user
                    self.authenticated = True
                    res = {
                        "messageType":"session_authenticated",
                        "session":session,
                    }
                jsonres = json.dumps(res, ensure_ascii=False).encode('utf8')
                self.sendMessage(jsonres, False)
            elif obj["messageType"] == "logoff":
                self.authenticated = False
                self.session = None
                self._authorization.remove_session(obj["session"])
                res = {
                    "messageType":"session_logoff"
                }
                jsonres = json.dumps(res, ensure_ascii=False).encode('utf8')
                self.sendMessage(jsonres, False)
            else:
                #self.spine.log.debug("WS onMessage:{0}", obj)
                if not self.authenticated:
                    pass
                elif obj["messageType"] == "query":
                    #res = yield from self.async_query(obj["query"], obj["args"], injected="socketSpine", session=self.user)
                    res = self.spine.send_query(obj["query"], *obj["args"], injected="socketSpine", session=self.user)
                    self.spine.log.debug("query response:{0}", res)
                    self.send_response(obj["id"], res)
                elif obj["messageType"] == "registerQueryHandler":
                    self.add_query_handler(obj["query"])
                    self.send_response(None, None)
                elif obj["messageType"] == "command":
                    self.spine.send_command(obj["command"], *obj["args"], injected="socketSpine", session=self.user)
                    self.send_response(obj["id"], None)
                elif obj["messageType"] == "registerCommandHandler":
                    self.add_command_handler(obj["command"])
                    self.send_response(obj["id"], None)
                elif obj["messageType"] == "event":
                    self.spine.trigger_event(
                        obj["event"], obj["id"],
                        obj["args"],
                        injected="socketSpine"
                    )
                    self.send_response(obj["id"], None)
                elif obj["messageType"] == "registerEventHandler":
                    self.add_event_handler(obj["event"], obj["eventId"])
                    self.send_response(obj["id"], None)
        except:
            self.spine.log.exception("WS onMessage exception")
            #res={"execptionType":exc_type,"value":exc_value,"traceback":exc_traceback}
            #self.sendResponse(res,"exception")

class SocketSpine:
    def __init__(self, config):
        #coro = None
        self._started = False
        self._config = config
        self._spine = Spine()

        try:
            import asyncio
        except ImportError:
            ## Trollius >= 0.3 was renamed
            import trollius as asyncio

        from autobahn.asyncio.websocket import WebSocketServerFactory

        ssl_context = None

        if encryption.enabled():
            self._spine.log.debug("socket using ssl")
            cert_file, key_file = encryption.get_cert()
            try:
                import ssl
                ssl_context = ssl.SSLContext(ssl.PROTOCOL_SSLv23)
                ssl_context.load_cert_chain(cert_file, key_file)
                #self._spine.log.debug("socket ssl found")
            except:
                ssl_context = None
                self._spine.log.error("socket failed to use ssl")

        self._spine.log.debug(
            "start websocket on:{0}, port:{1}",
            self._config.network.ip,
            self._config.network.ws_port
        )
        self.factory = WebSocketServerFactory()
        self.factory.protocol = _SpineProtocol

        self.loop = asyncio.get_event_loop()
        _SpineProtocol.loop = self.loop
        self.coro = self.loop.create_server(
            self.factory,
            self._config.network.ip,
            self._config.network.ws_port,
            ssl=ssl_context
        )

    def start_socket(self):
        self._spine.log.verbose("start websocket: {0} {1} ", self._config.network.ip, self._config.network.ws_port)
        self.loop.run_until_complete(self.coro)
        self._started = True

    def step(self):
        if self._started:
            self.loop.run_until_complete(self.coro)
        #loop.run_until_complete(coro_local)
        time.sleep(.001)
<|endoftext|>


<commit_before>import matplotlib
matplotlib.use('Svg') # Change renderer so it doesn't use the GUI
matplotlib.rcParams['mathtext.fontset'] = 'cm' # Change font to Computer Modern (LaTeX font)
import matplotlib.pyplot as plt
plt.rcParams["figure.figsize"] = [10,10]
from shapes import Polygon, Circle, Ellipse

import numpy as np
import StringIO

class Figures:
	def __init__(self):
		self.fig, self.ax = plt.subplots()
		self.tick_interval = .25
		self.tick_label_interval = 1

	def __export__(self):
		export_str = StringIO.StringIO()
		self.fig.savefig(export_str, format='svg', bbox_inches='tight')
		export_str.seek(0)  # rewind the data
		return export_str.buf  # this is svg data

	def __writeFile__(self, file_location):
		plt.savefig(file_location)

	def __display__(self):
		plt.show()

	def format_axis(self, xyrange=None, grid=False, arrows=True, color='black'):
		# Modify the plot view to scale, remove axis, and center our shape

		color_dict = {
			"blue": 'b',
			"green": 'g',
			"red": 'r',
			"cyan": 'c',
			"magenta": 'm',
			"yellow": 'y',
			"black": 'k',
			"white": 'w'
		}

		def adjust_spines(ax, spines):
		    for loc, spine in ax.spines.items():
		        if loc in spines:
		            spine.set_position(('outward', 10))  # outward by 10 points
		            spine.set_smart_bounds(True)
		        else:
		            spine.set_color('none')  # don't draw spine

		    # turn off ticks where there is no spine
		    if 'left' in spines:
		        ax.yaxis.set_ticks_position('left')
		    else:
		        # no yaxis ticks
		        ax.yaxis.set_ticks([])

		    if 'bottom' in spines:
		        ax.xaxis.set_ticks_position('bottom')
		    else:
		        # no xaxis ticks
		        ax.xaxis.set_ticks([])

		if xyrange == None:
			plt.axis('off')
			plt.axis('scaled')
		else:
			plt.gca().set_aspect('equal', adjustable='box')
			self.ax.set_xlim(left=xyrange[0][0], right=xyrange[0][1])
			self.ax.set_ylim(bottom=xyrange[1][0], top=xyrange[1][1])

			self.ax.spines['right'].set_color('none')
			self.ax.spines['top'].set_color('none')

			self.ax.spines['left'].set_position(('data', 0))
			self.ax.spines['bottom'].set_position(('data', 0))

		if grid:
			self.ax.grid(color='k', linestyle='dashed', linewidth=.5, alpha=0.5)

		if arrows:
			xmin, xmax = self.ax.get_xlim()
			ymin, ymax = self.ax.get_ylim()

			self.ax.arrow(xmin, 0, xmax-xmin, 0., lw = 1,
			         head_width=self.tick_interval/3, head_length=self.tick_interval/3,
			         length_includes_head=True, clip_on=False,color=color_dict[color])

			self.ax.arrow(0, ymin, 0., ymax-ymin, lw = 1,
			         head_width=self.tick_interval/3, head_length=self.tick_interval/3,
					 length_includes_head=True, clip_on=False,color=color_dict[color])

		# Control color
		self.ax.spines['bottom'].set_color(color_dict[color])
		self.ax.spines['left'].set_color(color_dict[color])

		self.fig.set_tight_layout(True)

	def addPoint(self, xys, texts, pointsize=6, fontsize=12, colors='black'):
		color_dict = {
			"blue": 'b',
			"green": 'g',
			"red": 'r',
			"cyan": 'c',
			"magenta": 'm',
			"yellow": 'y',
			"black": 'k',
			"white": 'w'
		}
		if not isinstance(colors, list):
			colors = [colors]
			xys = [xys]
			texts = [texts]

		for xy, text, color in zip(xys, texts, colors):
			plt.plot(xy[0], xy[1], 'o{}'.format(color_dict[color]), ms=pointsize)
			self.ax.annotate(text, xytext=xy, xy=xy, fontsize=fontsize, horizontalalignment='center', textcoords='offset points')

	def addText(self, xy, text, color="black", fontsize=25, alignment='center'):
		color_dict = {
			"blue": 'b',
			"green": 'g',
			"red": 'r',
			"cyan": 'c',
			"magenta": 'm',
			"yellow": 'y',
			"black": 'k',
			"white": 'w'
		}
		if not isinstance(color, list):
			color = [color]
			xy = [xy]
			text = [text]
			alignment = [alignment]

		for xy, text, color, alignment in zip(xy, text, color, alignment):
			self.ax.annotate(text, xytext=xy, xy=xy, fontsize=fontsize, horizontalalignment=alignment)

	def addFunction(self, functions, xyranges, colors='black', linewidth=2):
		color_dict = {
			"blue": 'b',
			"green": 'g',
			"red": 'r',
			"cyan": 'c',
			"magenta": 'm',
			"yellow": 'y',
			"black": 'k',
			"white": 'w'
		}

		if not isinstance(functions, list):
			functions = [functions]
			xyranges = [xyranges]
			colors = [colors]

		for function, xyrange, color in zip(functions, xyranges, colors):
			x = np.linspace(xyrange[0][0], xyrange[0][1], 1000)
			y = function(x)
			self.ax.plot(x, y, color_dict[color])

	def addTicks(self, tick_label_interval=1, tick_interval=1):
		self.tick_interval = tick_interval
		self.tick_label_interval = tick_label_interval
		# Control ticks
		self.ax.xaxis.set_major_locator(matplotlib.ticker.MultipleLocator(tick_label_interval))
		self.ax.yaxis.set_major_locator(matplotlib.ticker.MultipleLocator(tick_label_interval))
		self.ax.xaxis.set_minor_locator(matplotlib.ticker.MultipleLocator(tick_interval))
		self.ax.yaxis.set_minor_locator(matplotlib.ticker.MultipleLocator(tick_interval))

	def addPolygon(self, vertices):
		polygon = Polygon.Polygon(vertices, self.fig, self.ax)
		return polygon

	def addCircle(self, xy=(0,0), diameter=None, radius=None, label=None):
		circle = Circle.Circle(self.fig, self.ax, xy, diameter, radius, label)
		return circle

	def addEllipse(self, xy=(0,0), width=None, height=None, wlabel=None, hlabel=None, is_radius=True):
		ellipse = Ellipse.Ellipse(self.fig, self.ax, xy, width, height, wlabel, hlabel, is_radius)
		return ellipse

	def addTriangle_angle(self, angle=(45*np.pi)/180, rotation=0):
		# Define the angles and sides
		alpha = angle
		beta = np.pi/2
		gamma = ((np.pi-beta-alpha))

		A = 1
		B = np.sin(beta)/np.sin(alpha)
		C = np.sin(gamma)/np.sin(alpha)

		# Define the vertices
		vertex_A = [0, A]
		vertex_B = [0, 0]
		vertex_C = [C, 0]

		polygon = plt.Polygon([vertex_A, vertex_B, vertex_C], fill=False, linewidth=3)

		# Perform the rotation if at all
		transformation = matplotlib.transforms.Affine2D().rotate(rotation) + self.ax.transData
		polygon.set_transform(transformation)

		# Create and add polygon
		self.ax.add_patch(polygon)

	def addTriangle_side(self):
		raise Exception('Not implemented yet!')
<commit_msg>Boxes around text & Latex optional<commit_after>import matplotlib
matplotlib.use('Svg') # Change renderer so it doesn't use the GUI
matplotlib.rcParams['mathtext.fontset'] = 'cm' # Change font to Computer Modern (LaTeX font)
import matplotlib.pyplot as plt
plt.rcParams["figure.figsize"] = [10,10]
from shapes import Polygon, Circle, Ellipse

import numpy as np
import StringIO

class Figures:
	def __init__(self):
		self.fig, self.ax = plt.subplots()
		self.tick_interval = .25
		self.tick_label_interval = 1

	def __export__(self):
		export_str = StringIO.StringIO()
		self.fig.savefig(export_str, format='svg', bbox_inches='tight')
		export_str.seek(0)  # rewind the data
		return export_str.buf  # this is svg data

	def __writeFile__(self, file_location):
		plt.savefig(file_location)

	def __display__(self):
		plt.show()

	def format_axis(self, xyrange=None, grid=False, arrows=True, color='black'):
		# Modify the plot view to scale, remove axis, and center our shape

		color_dict = {
			"blue": 'b',
			"green": 'g',
			"red": 'r',
			"cyan": 'c',
			"magenta": 'm',
			"yellow": 'y',
			"black": 'k',
			"white": 'w'
		}

		def adjust_spines(ax, spines):
		    for loc, spine in ax.spines.items():
		        if loc in spines:
		            spine.set_position(('outward', 10))  # outward by 10 points
		            spine.set_smart_bounds(True)
		        else:
		            spine.set_color('none')  # don't draw spine

		    # turn off ticks where there is no spine
		    if 'left' in spines:
		        ax.yaxis.set_ticks_position('left')
		    else:
		        # no yaxis ticks
		        ax.yaxis.set_ticks([])

		    if 'bottom' in spines:
		        ax.xaxis.set_ticks_position('bottom')
		    else:
		        # no xaxis ticks
		        ax.xaxis.set_ticks([])

		if xyrange == None:
			plt.axis('off')
			plt.axis('scaled')
		else:
			plt.gca().set_aspect('equal', adjustable='box')
			self.ax.set_xlim(left=xyrange[0][0], right=xyrange[0][1])
			self.ax.set_ylim(bottom=xyrange[1][0], top=xyrange[1][1])

			self.ax.spines['right'].set_color('none')
			self.ax.spines['top'].set_color('none')

			self.ax.spines['left'].set_position(('data', 0))
			self.ax.spines['bottom'].set_position(('data', 0))

		if grid:
			self.ax.grid(color='k', linestyle='dashed', linewidth=.5, alpha=0.5)

		if arrows:
			xmin, xmax = self.ax.get_xlim()
			ymin, ymax = self.ax.get_ylim()

			self.ax.arrow(xmin, 0, xmax-xmin, 0., lw = 1,
			         head_width=self.tick_interval/3, head_length=self.tick_interval/3,
			         length_includes_head=True, clip_on=False,color=color_dict[color])

			self.ax.arrow(0, ymin, 0., ymax-ymin, lw = 1,
			         head_width=self.tick_interval/3, head_length=self.tick_interval/3,
					 length_includes_head=True, clip_on=False,color=color_dict[color])

		# Control color
		self.ax.spines['bottom'].set_color(color_dict[color])
		self.ax.spines['left'].set_color(color_dict[color])

		self.fig.set_tight_layout(True)

	def addPoint(self, xys, texts, pointsize=6, fontsize=12, colors='black'):
		color_dict = {
			"blue": 'b',
			"green": 'g',
			"red": 'r',
			"cyan": 'c',
			"magenta": 'm',
			"yellow": 'y',
			"black": 'k',
			"white": 'w'
		}
		if not isinstance(colors, list):
			colors = [colors]
			xys = [xys]
			texts = [texts]

		for xy, text, color in zip(xys, texts, colors):
			plt.plot(xy[0], xy[1], 'o{}'.format(color_dict[color]), ms=pointsize)
			self.ax.annotate(text, xytext=xy, xy=xy, fontsize=fontsize, horizontalalignment='center', textcoords='offset points')

	def addText(self, xy, text, color="black", fontsize=12, halignment='center', valignment='top', bbox={}, latex=True):
		color_dict = {
			"blue": 'b',
			"green": 'g',
			"red": 'r',
			"cyan": 'c',
			"magenta": 'm',
			"yellow": 'y',
			"black": 'k',
			"white": 'w'
		}
		if not isinstance(color, list):
			color = [color]
			xy = [xy]
			text = [text]
			halignment = [halignment]
			valignment = [valignment]
			bbox = [bbox]
			latex = [latex]

		for xy, text, color, valignment, halignment, bbox, latex in zip(xy, text, color, valignment, halignment, bbox, latex):
			self.ax.annotate("$"+text+"$" if latex else text, xytext=xy, xy=xy, fontsize=fontsize, horizontalalignment=halignment, verticalalignment=valignment, bbox=bbox)

	def addFunction(self, functions, xyranges, colors='black', linewidth=2):
		color_dict = {
			"blue": 'b',
			"green": 'g',
			"red": 'r',
			"cyan": 'c',
			"magenta": 'm',
			"yellow": 'y',
			"black": 'k',
			"white": 'w'
		}

		if not isinstance(functions, list):
			functions = [functions]
			xyranges = [xyranges]
			colors = [colors]

		for function, xyrange, color in zip(functions, xyranges, colors):
			x = np.linspace(xyrange[0][0], xyrange[0][1], 1000)
			y = function(x)
			self.ax.plot(x, y, color_dict[color])

	def addTicks(self, tick_label_interval=1, tick_interval=1):
		self.tick_interval = tick_interval
		self.tick_label_interval = tick_label_interval
		# Control ticks
		self.ax.xaxis.set_major_locator(matplotlib.ticker.MultipleLocator(tick_label_interval))
		self.ax.yaxis.set_major_locator(matplotlib.ticker.MultipleLocator(tick_label_interval))
		self.ax.xaxis.set_minor_locator(matplotlib.ticker.MultipleLocator(tick_interval))
		self.ax.yaxis.set_minor_locator(matplotlib.ticker.MultipleLocator(tick_interval))

	def addPolygon(self, vertices):
		polygon = Polygon.Polygon(vertices, self.fig, self.ax)
		return polygon

	def addCircle(self, xy=(0,0), diameter=None, radius=None, label=None):
		circle = Circle.Circle(self.fig, self.ax, xy, diameter, radius, label)
		return circle

	def addEllipse(self, xy=(0,0), width=None, height=None, wlabel=None, hlabel=None, is_radius=True):
		ellipse = Ellipse.Ellipse(self.fig, self.ax, xy, width, height, wlabel, hlabel, is_radius)
		return ellipse

	def addTriangle_angle(self, angle=(45*np.pi)/180, rotation=0):
		# Define the angles and sides
		alpha = angle
		beta = np.pi/2
		gamma = ((np.pi-beta-alpha))

		A = 1
		B = np.sin(beta)/np.sin(alpha)
		C = np.sin(gamma)/np.sin(alpha)

		# Define the vertices
		vertex_A = [0, A]
		vertex_B = [0, 0]
		vertex_C = [C, 0]

		polygon = plt.Polygon([vertex_A, vertex_B, vertex_C], fill=False, linewidth=3)

		# Perform the rotation if at all
		transformation = matplotlib.transforms.Affine2D().rotate(rotation) + self.ax.transData
		polygon.set_transform(transformation)

		# Create and add polygon
		self.ax.add_patch(polygon)

	def addTriangle_side(self):
		raise Exception('Not implemented yet!')
<|endoftext|>


<commit_before>"""Utility object for sending emails reports via shutit

Example code:

        e = shutit.get_emailer('shutit.tk.mysql.mysql',shutit)
        for line in ['your message line 1', 'your message line 2']:
                e.add_line(line)
        for attach in ['/tmp/filetoattach1','/tmp/filetoattach2']:
                e.attach(attach)
        e.send()

Example cfg:

        [shutit.tk.mysql.mysql]
        mailto:recipient@example.com
        mailfrom:sender@example.com
        smtp_server:localhost
        subject:Shutit Report
        signature:--Angry Shutit
        compress:yes
        username:
        password:
        safe_mode: True
        mailto_maintainer: True

"""

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

from email.mime.text import MIMEText
from email.mime.application import MIMEApplication
from email.mime.multipart import MIMEMultipart
from smtplib import SMTP, SMTPSenderRefused
import os, gzip

class Emailer():
    """ Emailer class definition
    """

    def __init__(
        self,
        cfg_section,
        shutit
    ):
        """Initialise the emailer object
        cfg_section - section in shutit config to look for email configuration items
                      allowing easier config according to shutit_module.
                      e.g. 'com.my_module','subject': My Module Build Failed!
                      Config Items:
                      mailto      - address to send the mail to (no default)
                      mailfrom    - address to send the mail from (angry@shutit.tk)
                      smtp_server - server to send the mail (localhost)
                      smtp_port   - port to contact the smtp server on (587)
                      subject     - subject of the email (Shutit Report)
                      signature   - --Angry Shutit
                      compress    - gzip attachments? (True)
                      username    - mail username
                      password    - mail password
                      safe_mode   - don't fail the build if we get an exception
                      mailto_maintainer - email the maintainer of the module as
                                          well as the mailto address
        """
        self.shutit    = shutit
        self.config    = {}
        self.__set_config(cfg_section)
        self.lines     = []
        self.attaches  = []

    def __set_config(self, cfg_section):
        """Set a local config array up according to
        defaults and main shutit configuration

        cfg_section - see __init__
        """
        defaults = [
            'mailto', None,
            'mailfrom', 'angry@shutit.tk',
            'smtp_server', 'localhost',
            'smtp_port', 25,
            'send_mail', True,
            'subject', 'Shutit Report',
            'signature', '--Angry Shutit',
            'compress', True,
            'username', '',
            'password', '',
            'safe_mode', True,
            'maintainer','',
            'mailto_maintainer', True
        ]

        for cfg_name, cfg_default in zip(defaults[0::2], defaults[1::2]):
            try:
                self.config[cfg_name] = self.shutit.cfg[cfg_section][cfg_name]
            except KeyError:
                if cfg_default is None:
                    raise Exception(cfg_default + ' must be set')
                else:
                    self.config[cfg_name] = cfg_default

        # only send a mail to the module's maintainer if configured correctly
        if self.config['mailto_maintainer'] and \
            (self.config['maintainer'] == "" or \
            self.config['maintainer'] == self.config['mailto']):
            self.config['mailto_maintainer'] = False
            self.config['maintainer'] = ""

    @staticmethod
    def __gzip(filename):
        """ Compress a file returning the new filename (.gz)
        """
        zipname = filename + '.gz'
        file_pointer = open(filename,'rb')
        zip_pointer = gzip.open(zipname,'wb')
        zip_pointer.writelines(file_pointer)
        file_pointer.close()
        zip_pointer.close()
        return zipname

    def add_line(self, line):
        """Add a single line to the email body
        """
        self.lines.append(line)

    def add_body(self, msg):
        """Add an entire email body as a string, will be split on newlines
           and overwrite anything currently in the body (e.g added by add_lines)
        """
        self.lines = msg.rsplit('\n')

    def attach(self, filename, filetype="txt"):
        """Attach a file - currently needs to be entered as root (shutit)

        Filename - absolute path, relative to the docker container!
        filetype - MIMEApplication._subtype
        """
        shutit = self.shutit
        host_path = '/tmp'
        host_fn = os.path.join(host_path, os.path.basename(filename))
        shutit.get_file(filename, host_path)
        if self.config['compress']:
            filetype = 'x-gzip-compressed'
            filename = self.__gzip(host_fn)
        file_pointer = open(host_fn, 'rb')
        attach = MIMEApplication(file_pointer.read(), _subtype=filetype)
        file_pointer.close()
        attach.add_header(
            'Content-Disposition',
            'attachment',
            filename=os.path.basename(filename)
        )
        self.attaches.append(attach)

    def __compose(self):
        """ Compose the message, pulling together body, attachments etc
        """
        msg  = MIMEMultipart()
        msg['Subject'] = self.config['subject']
        msg['To']      = self.config['mailto']
        msg['From']    = self.config['mailfrom']
        # add the module's maintainer as a CC if configured
        if self.config['mailto_maintainer']:
            msg['Cc'] = self.config['maintainer']
        if self.config['signature'] != '':
            signature = '\n\n' + self.config['signature']
        else:
            signature = self.config['signature']
        body = MIMEText('\n'.join(self.lines) + signature)
        msg.attach(body)
        for attach in self.attaches:
            msg.attach(attach)
        return msg

    def send(self, attachment_failure=False):
        """Send the email according to the configured setup

           attachment_failure - used to indicate a recursive call after the
                                smtp server has refused based on file size.
                                Should not be used externally
        """
        if not self.config['send_mail']:
            self.shutit.log('emailer.send: Not configured to send mail!')
            return True
        msg = self.__compose()
        mailto = [self.config['mailto']]
        smtp = SMTP(self.config['smtp_server'], self.config['smtp_port'])
        smtp.starttls()
        if self.config['username'] != '':
            smtp.login(self.config['username'], self.config['password'])
        if self.config['mailto_maintainer']:
            mailto.append(self.config['maintainer'])
        try:
            smtp.sendmail(
                self.config['mailfrom'],
                mailto,
                msg.as_string()
            )
        except SMTPSenderRefused as refused:
            code = refused.args[0]
            if code == 552 and not attachment_failure:
                self.shutit.log("Mailserver rejected message due to " + \
                    "oversize attachments, attempting to resend without")
                self.send(attachment_failure=True)
                self.attaches = []
            else:
                self.shutit.log("Unhandled SMTP error:" + str(refused))
                if not self.config['safe_mode']:
                    raise refused
        except Exception as error:
            self.shutit.log('Unhandled exception: ' + str(error))
            if not self.config['safe_mode']:
                raise error
        finally:
            smtp.quit()

<commit_msg>Always log, fix gzipping, clear attachments before recursing<commit_after>"""Utility object for sending emails reports via shutit

Example code:

        e = shutit.get_emailer('shutit.tk.mysql.mysql',shutit)
        for line in ['your message line 1', 'your message line 2']:
                e.add_line(line)
        for attach in ['/tmp/filetoattach1','/tmp/filetoattach2']:
                e.attach(attach)
        e.send()

Example cfg:

        [shutit.tk.mysql.mysql]
        mailto:recipient@example.com
        mailfrom:sender@example.com
        smtp_server:localhost
        subject:Shutit Report
        signature:--Angry Shutit
        compress:yes
        username:
        password:
        safe_mode: True
        mailto_maintainer: True

"""

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

from email.mime.text import MIMEText
from email.mime.application import MIMEApplication
from email.mime.multipart import MIMEMultipart
from smtplib import SMTP, SMTPSenderRefused
import os, gzip

class Emailer():
    """ Emailer class definition
    """

    def __init__(
        self,
        cfg_section,
        shutit
    ):
        """Initialise the emailer object
        cfg_section - section in shutit config to look for email configuration items
                      allowing easier config according to shutit_module.
                      e.g. 'com.my_module','subject': My Module Build Failed!
                      Config Items:
                      mailto      - address to send the mail to (no default)
                      mailfrom    - address to send the mail from (angry@shutit.tk)
                      smtp_server - server to send the mail (localhost)
                      smtp_port   - port to contact the smtp server on (587)
                      subject     - subject of the email (Shutit Report)
                      signature   - --Angry Shutit
                      compress    - gzip attachments? (True)
                      username    - mail username
                      password    - mail password
                      safe_mode   - don't fail the build if we get an exception
                      mailto_maintainer - email the maintainer of the module as
                                          well as the mailto address
        """
        self.shutit    = shutit
        self.config    = {}
        self.__set_config(cfg_section)
        self.lines     = []
        self.attaches  = []

    def __set_config(self, cfg_section):
        """Set a local config array up according to
        defaults and main shutit configuration

        cfg_section - see __init__
        """
        defaults = [
            'mailto', None,
            'mailfrom', 'angry@shutit.tk',
            'smtp_server', 'localhost',
            'smtp_port', 25,
            'send_mail', True,
            'subject', 'Shutit Report',
            'signature', '--Angry Shutit',
            'compress', True,
            'username', '',
            'password', '',
            'safe_mode', True,
            'maintainer','',
            'mailto_maintainer', True
        ]

        for cfg_name, cfg_default in zip(defaults[0::2], defaults[1::2]):
            try:
                self.config[cfg_name] = self.shutit.cfg[cfg_section][cfg_name]
            except KeyError:
                if cfg_default is None:
                    raise Exception(cfg_default + ' must be set')
                else:
                    self.config[cfg_name] = cfg_default

        # only send a mail to the module's maintainer if configured correctly
        if self.config['mailto_maintainer'] and \
            (self.config['maintainer'] == "" or \
            self.config['maintainer'] == self.config['mailto']):
            self.config['mailto_maintainer'] = False
            self.config['maintainer'] = ""

    @staticmethod
    def __gzip(filename):
        """ Compress a file returning the new filename (.gz)
        """
        zipname = filename + '.gz'
        file_pointer = open(filename,'rb')
        zip_pointer = gzip.open(zipname,'wb')
        zip_pointer.writelines(file_pointer)
        file_pointer.close()
        zip_pointer.close()
        return zipname

    def add_line(self, line):
        """Add a single line to the email body
        """
        self.lines.append(line)

    def add_body(self, msg):
        """Add an entire email body as a string, will be split on newlines
           and overwrite anything currently in the body (e.g added by add_lines)
        """
        self.lines = msg.rsplit('\n')

    def attach(self, filename, filetype="txt"):
        """Attach a file - currently needs to be entered as root (shutit)

        Filename - absolute path, relative to the docker container!
        filetype - MIMEApplication._subtype
        """
        shutit = self.shutit
        host_path = '/tmp'
        host_fn = os.path.join(host_path, os.path.basename(filename))
        shutit.get_file(filename, host_path)
        if self.config['compress']:
            filetype = 'x-gzip-compressed'
            filename = self.__gzip(host_fn)
            host_fn = os.path.join(host_path, os.path.basename(filename))
        file_pointer = open(host_fn, 'rb')
        attach = MIMEApplication(file_pointer.read(), _subtype=filetype)
        file_pointer.close()
        attach.add_header(
            'Content-Disposition',
            'attachment',
            filename=os.path.basename(filename)
        )
        self.attaches.append(attach)

    def __compose(self):
        """ Compose the message, pulling together body, attachments etc
        """
        msg  = MIMEMultipart()
        msg['Subject'] = self.config['subject']
        msg['To']      = self.config['mailto']
        msg['From']    = self.config['mailfrom']
        # add the module's maintainer as a CC if configured
        if self.config['mailto_maintainer']:
            msg['Cc'] = self.config['maintainer']
        if self.config['signature'] != '':
            signature = '\n\n' + self.config['signature']
        else:
            signature = self.config['signature']
        body = MIMEText('\n'.join(self.lines) + signature)
        msg.attach(body)
        for attach in self.attaches:
            msg.attach(attach)
        return msg

    def send(self, attachment_failure=False):
        """Send the email according to the configured setup

           attachment_failure - used to indicate a recursive call after the
                                smtp server has refused based on file size.
                                Should not be used externally
        """
        if not self.config['send_mail']:
            self.shutit.log(
                'emailer.send: Not configured to send mail!',
                force_stdout=True
            )
            return True
        msg = self.__compose()
        mailto = [self.config['mailto']]
        smtp = SMTP(self.config['smtp_server'], self.config['smtp_port'])
        smtp.starttls()
        if self.config['username'] != '':
            smtp.login(self.config['username'], self.config['password'])
        if self.config['mailto_maintainer']:
            mailto.append(self.config['maintainer'])
        try:
            self.shutit.log('Attempting to send email', force_stdout=True)
            smtp.sendmail(
                self.config['mailfrom'],
                mailto,
                msg.as_string()
            )
        except SMTPSenderRefused as refused:
            code = refused.args[0]
            if code == 552 and not attachment_failure:
                self.shutit.log(
                    "Mailserver rejected message due to " + \
                    "oversize attachments, attempting to resend without",
                    force_stdout=True
                )
                self.attaches = []
                self.lines.append("Oversized attachments not sent")
                self.send(attachment_failure=True)
            else:
                self.shutit.log(
                    "Unhandled SMTP error:" + str(refused),
                    force_stdout=True
                )
                if not self.config['safe_mode']:
                    raise refused
        except Exception as error:
            self.shutit.log(
                'Unhandled exception: ' + str(error),
                force_stdout=True
            )
            if not self.config['safe_mode']:
                raise error
        finally:
            smtp.quit()

<|endoftext|>


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

"""
Copyright (C) 2013
Author(s): Matthew Loper

See LICENCE.txt for licensing and contact information.
"""

import cv2
import numpy as np
from copy import deepcopy
import scipy.sparse as sp
from chumpy.ch import MatVecMult, Ch, depends_on

def laplacian_pyramid(input_objective, imshape, normalization, n_levels, as_list):

    if normalization is None:
        norm2 = lambda x : x
    elif normalization is 'SSE':
        norm2 = lambda x : x / np.sqrt(np.sum(x.r**2.))
    elif normalization is 'size':
        norm2 = lambda x : x / np.sqrt(x.r.size)
    else:
        raise Exception('Normalization must be None, SSE, or size.')

    output_objs = []
    for level in range(n_levels):

        blur_mtx = filter_for(imshape[0], imshape[1], imshape[2], kernel = GaussianKernel2D(3, 1))
        blurred = MatVecMult(blur_mtx, input_objective).reshape(imshape)
        output_objs.append(norm2(input_objective - blurred))

        halfsampler_mtx, imshape = halfsampler_for(imshape)
        input_objective = MatVecMult(halfsampler_mtx, blurred).reshape(imshape)

    output_objs.append(norm2(input_objective).reshape(imshape))

    return output_objs if as_list else reduce(lambda x, y : x & y, output_objs)

def gaussian_pyramid(input_objective, imshape, normalization, n_levels, as_list=True, label=None):

    if normalization is None:
        norm2 = lambda x : x
    elif normalization is 'SSE':
        norm2 = lambda x : x / np.sqrt(np.sum(x.r**2.))
    elif normalization is 'size':
        norm2 = lambda x : x / np.sqrt(x.r.size)
    else:
        raise Exception('Normalization must be None, SSE, or size.')

    cur_imshape = deepcopy(imshape)
    cur_obj = input_objective

    input_objective = norm2(input_objective)
    output_objectives = [input_objective]

    for ik in range(n_levels):
        cur_obj = GaussPyrDownOne(px=cur_obj, im_shape = cur_imshape)
        cur_imshape = cur_obj.output_shape
        output_objectives.append(norm2(cur_obj) if label is None else norm2(cur_obj) >> '%s%d' % (label,ik))

    if not as_list:
        andit = lambda a : reduce(lambda x, y : x & y, a)
        output_objectives = andit(output_objectives)

    return output_objectives

def GaussianKernel2D(ksize, sigma):
    if ksize % 2 != 1:
        raise Exception('ksize should be an even number')
    if sigma <= 0:
        raise Exception('sigma should be positive')
    oneway = np.tile(cv2.getGaussianKernel(ksize,sigma), (1, ksize))
    return oneway * oneway.T

class GaussPyrDownOne(Ch):
    terms =  'im_shape', 'want_downsampling', 'kernel'
    dterms = 'px'

    # Approximation to a 3x3 Gaussian kernel
    default_kernel = GaussianKernel2D(3, 1)

    def on_changed(self, which):
        if not hasattr(self, 'kernel'):
            self.kernel = self.default_kernel.copy()

        if 'im_shape' in which:
            sh = self.im_shape
            self.transf_mtx, self._output_shape = filter_for_nopadding(sh, self.kernel)
            if True: # self.want_downsampling: <-- setting want_downsampling to "False" is broken
                halfsampler, self._output_shape = halfsampler_for(self._output_shape)
                self.transf_mtx = halfsampler.dot(self.transf_mtx)
            self.add_dterm('transform', MatVecMult(mtx=self.transf_mtx))

        if 'px' in which:
            self.transform.vec = self.px

    @property
    def output_shape(self):
        self._call_on_changed() # trigger changes from im_shape
        return self._output_shape

    def compute_r(self):
        return self.transf_mtx.dot(self.px.r.ravel()).reshape(self.output_shape)

    def compute_dr_wrt(self, wrt):
        if wrt is self.transform:
            return 1

def halfsampler_for(shape):
    h = shape[0]
    w = shape[1]
    d = shape[2]

    JS = np.arange(h*w*d).reshape((h,w,d))
    JS = JS[::2,::2,:]
    JS = JS.flatten()

    IS = np.arange(len(JS))
    data = np.ones(len(JS))

    shape = (int(np.ceil(h/2.)), int(np.ceil(w/2.)), int(d))
    return sp.csc_matrix((data, (IS, JS)), shape=(len(IS), h*w*d)), shape

def filter_for_nopadding(shape, kernel):
    new_shape = (shape - np.array([kernel.shape[0] - 1, kernel.shape[1] - 1, 0])).astype(np.uint32)

    IS = []
    JS = []
    data = []

    new_h = int(new_shape[0])
    new_w = int(new_shape[1])
    old_h = shape[0]
    old_w = shape[1]
    xs = np.tile(np.arange(old_w), (old_h, 1))
    ys = np.tile(np.arange(old_h), (old_w, 1)).T.copy()

    for ky in range(kernel.shape[0]):
        for kx in range(kernel.shape[1]):

            xs2 = xs[ky:new_h+ky, kx:new_w+kx]
            ys2 = ys[ky:new_h+ky, kx:new_w+kx]

            JS.append(xs2.ravel() + ys2.ravel()*old_w)
            IS.append(np.arange(new_shape[0]*new_shape[1]))
            data.append(np.ones(IS[-1].size) * kernel[ky, kx])

    IS = np.concatenate(IS)
    JS = np.concatenate(JS)
    data = np.concatenate(data)

    d = int(shape[2])
    if d > 1:
        IS = [IS*d+k for k in range(d)]
        JS = [JS*d+k for k in range(d)]
        data = [data for k in range(d)]
        IS = np.concatenate(IS)
        JS = np.concatenate(JS)
        data = np.concatenate(data)

    return sp.csc_matrix((data, (IS, JS))), new_shape

def filter_for(h, w, d, kernel):
    kernel = np.atleast_3d(kernel)
    if kernel.shape[2] != d:
        kernel = np.tile(kernel, (1, 1, d))

    kxm = (kernel.shape[1]+1) / 2
    kym = (kernel.shape[0]+1) / 2

    xs = np.tile(np.arange(w), (h, 1))
    ys = np.tile(np.arange(h), (w, 1)).T.copy()

    IS = []
    JS = []
    data = []
    for ky in range(kernel.shape[0]):
        for kx in range(kernel.shape[1]):
            for channel in range(d):
                cky = ky - kym
                ckx = kx - kxm

                xs2 = np.clip(xs + ckx, 0, w-1)
                ys2 = np.clip(ys + cky, 0, h-1)

                IS.append(np.arange(w*h)*d+channel)
                JS.append((xs2.ravel() + ys2.ravel()*w)*d+channel)
                data.append(np.ones(IS[-1].size) * kernel[ky, kx, channel])

    IS = np.concatenate(IS)
    JS = np.concatenate(JS)
    data = np.concatenate(data)

    #if d > 1:
    #    IS = [IS*d+k for k in range(d)]
    #    JS = [JS*d+k for k in range(d)]
    #    data = [data for k in range(d)]
    #    IS = np.concatenate(IS)
    #    JS = np.concatenate(JS)
    #    data = np.concatenate(data)
    #

    return sp.csc_matrix((data, (IS, JS)), shape=(h*w*d, h*w*d))

def main():
    pass

if __name__ == '__main__':
    main()
<commit_msg>filters.py: check input parms<commit_after>#!/usr/bin/env python

"""
Copyright (C) 2013
Author(s): Matthew Loper

See LICENCE.txt for licensing and contact information.
"""

import cv2
import numpy as np
from copy import deepcopy
import scipy.sparse as sp
from chumpy.ch import MatVecMult, Ch, depends_on

def laplacian_pyramid(input_objective, imshape, normalization, n_levels, as_list):

    if normalization is None:
        norm2 = lambda x : x
    elif normalization is 'SSE':
        norm2 = lambda x : x / np.sqrt(np.sum(x.r**2.))
    elif normalization is 'size':
        norm2 = lambda x : x / np.sqrt(x.r.size)
    else:
        raise Exception('Normalization must be None, SSE, or size.')

    output_objs = []
    for level in range(n_levels):

        blur_mtx = filter_for(imshape[0], imshape[1], imshape[2], kernel = GaussianKernel2D(3, 1))
        blurred = MatVecMult(blur_mtx, input_objective).reshape(imshape)
        output_objs.append(norm2(input_objective - blurred))

        halfsampler_mtx, imshape = halfsampler_for(imshape)
        input_objective = MatVecMult(halfsampler_mtx, blurred).reshape(imshape)

    output_objs.append(norm2(input_objective).reshape(imshape))

    return output_objs if as_list else reduce(lambda x, y : x & y, output_objs)

def gaussian_pyramid(input_objective, imshape, normalization, n_levels, as_list=True, label=None):

    if normalization is None:
        norm2 = lambda x : x
    elif normalization is 'SSE':
        norm2 = lambda x : x / np.sqrt(np.sum(x.r**2.))
    elif normalization is 'size':
        norm2 = lambda x : x / np.sqrt(x.r.size)
    else:
        raise Exception('Normalization must be None, SSE, or size.')

    cur_imshape = deepcopy(imshape)
    cur_obj = input_objective

    input_objective = norm2(input_objective)
    output_objectives = [input_objective]

    for ik in range(n_levels):
        cur_obj = GaussPyrDownOne(px=cur_obj, im_shape = cur_imshape)
        cur_imshape = cur_obj.output_shape
        output_objectives.append(norm2(cur_obj) if label is None else norm2(cur_obj) >> '%s%d' % (label,ik))

    if not as_list:
        andit = lambda a : reduce(lambda x, y : x & y, a)
        output_objectives = andit(output_objectives)

    return output_objectives

def GaussianKernel2D(ksize, sigma):
    if ksize % 2 != 1:
        raise Exception('ksize should be an even number')
    if sigma <= 0:
        raise Exception('sigma should be positive')
    oneway = np.tile(cv2.getGaussianKernel(ksize,sigma), (1, ksize))
    return oneway * oneway.T

class GaussPyrDownOne(Ch):
    terms =  'im_shape', 'want_downsampling', 'kernel'
    dterms = 'px'

    # Approximation to a 3x3 Gaussian kernel
    default_kernel = GaussianKernel2D(3, 1)

    def on_changed(self, which):
        if not hasattr(self, 'kernel'):
            self.kernel = self.default_kernel.copy()

        if 'im_shape' in which:
            sh = self.im_shape
            self.transf_mtx, self._output_shape = filter_for_nopadding(sh, self.kernel)
            if True: # self.want_downsampling: <-- setting want_downsampling to "False" is broken
                halfsampler, self._output_shape = halfsampler_for(self._output_shape)
                self.transf_mtx = halfsampler.dot(self.transf_mtx)
            self.add_dterm('transform', MatVecMult(mtx=self.transf_mtx))

        if 'px' in which:
            self.transform.vec = self.px

    @property
    def output_shape(self):
        self._call_on_changed() # trigger changes from im_shape
        return self._output_shape

    def compute_r(self):
        return self.transf_mtx.dot(self.px.r.ravel()).reshape(self.output_shape)

    def compute_dr_wrt(self, wrt):
        if wrt is self.transform:
            return 1

def halfsampler_for(shape):
    h = shape[0]
    w = shape[1]
    d = shape[2]

    JS = np.arange(h*w*d).reshape((h,w,d))
    JS = JS[::2,::2,:]
    JS = JS.flatten()

    IS = np.arange(len(JS))
    data = np.ones(len(JS))

    shape = (int(np.ceil(h/2.)), int(np.ceil(w/2.)), int(d))
    return sp.csc_matrix((data, (IS, JS)), shape=(len(IS), h*w*d)), shape

def filter_for_nopadding(shape, kernel):
    assert(len(shape)==3)
    new_shape = (shape - np.array([kernel.shape[0] - 1, kernel.shape[1] - 1, 0])).astype(np.uint32)

    IS = []
    JS = []
    data = []

    new_h = int(new_shape[0])
    new_w = int(new_shape[1])
    old_h = shape[0]
    old_w = shape[1]
    xs = np.tile(np.arange(old_w), (old_h, 1))
    ys = np.tile(np.arange(old_h), (old_w, 1)).T.copy()

    for ky in range(kernel.shape[0]):
        for kx in range(kernel.shape[1]):

            xs2 = xs[ky:new_h+ky, kx:new_w+kx]
            ys2 = ys[ky:new_h+ky, kx:new_w+kx]

            JS.append(xs2.ravel() + ys2.ravel()*old_w)
            IS.append(np.arange(new_shape[0]*new_shape[1]))
            data.append(np.ones(IS[-1].size) * kernel[ky, kx])

    IS = np.concatenate(IS)
    JS = np.concatenate(JS)
    data = np.concatenate(data)

    d = int(shape[2])
    if d > 1:
        IS = [IS*d+k for k in range(d)]
        JS = [JS*d+k for k in range(d)]
        data = [data for k in range(d)]
        IS = np.concatenate(IS)
        JS = np.concatenate(JS)
        data = np.concatenate(data)

    return sp.csc_matrix((data, (IS, JS))), new_shape

def filter_for(h, w, d, kernel):
    kernel = np.atleast_3d(kernel)
    if kernel.shape[2] != d:
        kernel = np.tile(kernel, (1, 1, d))

    kxm = (kernel.shape[1]+1) / 2
    kym = (kernel.shape[0]+1) / 2

    xs = np.tile(np.arange(w), (h, 1))
    ys = np.tile(np.arange(h), (w, 1)).T.copy()

    IS = []
    JS = []
    data = []
    for ky in range(kernel.shape[0]):
        for kx in range(kernel.shape[1]):
            for channel in range(d):
                cky = ky - kym
                ckx = kx - kxm

                xs2 = np.clip(xs + ckx, 0, w-1)
                ys2 = np.clip(ys + cky, 0, h-1)

                IS.append(np.arange(w*h)*d+channel)
                JS.append((xs2.ravel() + ys2.ravel()*w)*d+channel)
                data.append(np.ones(IS[-1].size) * kernel[ky, kx, channel])

    IS = np.concatenate(IS)
    JS = np.concatenate(JS)
    data = np.concatenate(data)

    #if d > 1:
    #    IS = [IS*d+k for k in range(d)]
    #    JS = [JS*d+k for k in range(d)]
    #    data = [data for k in range(d)]
    #    IS = np.concatenate(IS)
    #    JS = np.concatenate(JS)
    #    data = np.concatenate(data)
    #

    return sp.csc_matrix((data, (IS, JS)), shape=(h*w*d, h*w*d))

def main():
    pass

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


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

'''
Query the ADS and INSPIRE databases to collect the citations used in a LaTeX file into a BibTeX file

Idea and main implementation is taken by Michele Vallisneri, see http://www.vallis.org/salon/summary-2.html
I changed the script to avoid the requests module, and added the INSPIRE database part.

Usage:
python fillbib.py <tex_file> <bib_file>. If <bib_file> is absent, it will try to guess it from the aux file"

'''
from __future__ import absolute_import, print_function
import sys, os, re, html

if sys.version_info.major>=3:
    import urllib.request as urllib
else:
    import urllib

def ads_citation(c): # download single ADS citation
    #f= urllib.urlopen("http://adsabs.harvard.edu/cgi-bin/nph-bib_query?bibcode="+c+"&data_type=BIBTEX&db_key=AST&nocookieset=1")
    #print("http://adsabs.harvard.edu/cgi-bin/nph-bib_query?bibcode="+c+"&data_type=BIBTEX&db_key=AST&nocookieset=1")
    f= urllib.urlopen("https://ui.adsabs.harvard.edu/abs/"+c+"/exportcitation")

    bib = f.read()
    f.close()
    if sys.version_info.major>=3:
        bib=bib.decode()
    bib = "@"+list(filter(lambda x:'adsnote' in x, bib.split("@")))[0].split("</textarea>")[0]
    bib=html.unescape(bib)

    if 'arXiv' in c: # Take care of preprint on ADS
        bib = bib.split("{")[0]+"{"+c+","+",".join(bib.split(",")[1:])
        return bib

    elif bib.split("{")[1].split(',')[0] == c: # Check you got what you where looking for
        return bib
    else:
        return None

def inspire_citation(c): # download single INSPIRE citation

    f= urllib.urlopen("https://old.inspirehep.net/search?p="+c+"&of=hx&em=B&sf=year&so=d&rg=1")
    bib = f.read()
    if sys.version_info.major>=3:
        bib=bib.decode()
    f.close()
    bib = "@"+bib.split("@")[1].split('
')[-2]

    if bib.split("{")[1].split(',')[0] == c: # Check you got what you where looking for
        return bib
    else:
        return None

def test_ads(): # test single ADS web scraping (both published articles and preprints)
    test_key = ["2016PhRvL.116f1102A","2016arXiv160203837T"]
    known_output= '@ARTICLE{2016PhRvL.116f1102A,\n   author = {{Abbott}, B.~P. and {Abbott}, R. and {Abbott}, T.~D. and {Abernathy}, M.~R. and \n\t{Acernese}, F. and {Ackley}, K. and {Adams}, C. and {Adams}, T. and \n\t{Addesso}, P. and {Adhikari}, R.~X. and et al.},\n    title = "{Observation of Gravitational Waves from a Binary Black Hole Merger}",\n  journal = {Physical Review Letters},\narchivePrefix = "arXiv",\n   eprint = {1602.03837},\n primaryClass = "gr-qc",\n     year = 2016,\n    month = feb,\n   volume = 116,\n   number = 6,\n      eid = {061102},\n    pages = {061102},\n      doi = {10.1103/PhysRevLett.116.061102},\n   adsurl = {http://adsabs.harvard.edu/abs/2016PhRvL.116f1102A},\n  adsnote = {Provided by the SAO/NASA Astrophysics Data System}\n}\n\n'
    assert [ads_citation(tk) == known_output for tk in test_key]

def test_inspire(): # test single INSPIRE web scraping
    test_key = "Abbott:2016blz"
    known_output= '@article{Abbott:2016blz,\n      author         = "Abbott, B. P. and others",\n      title          = "{Observation of Gravitational Waves from a Binary Black\n                        Hole Merger}",\n      collaboration  = "Virgo, LIGO Scientific",\n      journal        = "Phys. Rev. Lett.",\n      volume         = "116",\n      year           = "2016",\n      number         = "6",\n      pages          = "061102",\n      doi            = "10.1103/PhysRevLett.116.061102",\n      eprint         = "1602.03837",\n      archivePrefix  = "arXiv",\n      primaryClass   = "gr-qc",\n      reportNumber   = "LIGO-P150914",\n      SLACcitation   = "%%CITATION = ARXIV:1602.03837;%%"\n}\n'
    assert inspire_citation(test_key) == known_output

if __name__ == "__main__":

    if len(sys.argv)==2:     # Get the name of the bibfile from the aux file
        basename = sys.argv[1].split('.tex')[0]
        auxfile = basename + '.aux'
        for line in open(auxfile,'r'):
            m = re.search(r'\\bibdata\{(.*)\}',line)   # match \citation{...}, collect the ... note that we escape \, {, and }
            if m:
                bibfile = list(filter(lambda x:x!=basename+'Notes', m.group(1).split(',')))[0]  # Remove that annyoing feature of revtex which creates a *Notes.bib bibfile. Note this is not solid if you want to handle multiple bib files.
        bibfile = bibfile + '.bib'

    elif len(sys.argv)==3:    # Bibfile specified from argv
        basename = sys.argv[1].split('.tex')[0]
        auxfile = basename + '.aux'
        bibfile = sys.argv[2].split('.bib')[0] + '.bib'
    else:
        print("Usage: python fillbib.py <tex_file> <bib_file>. If <bib_file> is absent, assume the two are the same.")
        sys.exit()

    # Get all citations from aux file. Citations will look like \citation{2004PhRvD..69j4017P,2004PhRvD..69j4017P}
    cites = set() # use a set (no repetitions)
    for line in open(auxfile,'r'):
        m = re.search(r'\\citation\{(.*)\}',line)   # find \citation{...}
        if m:
            cites.update(m.group(1).split(','))     # split by commas

    cites= cites.difference(['REVTEX41Control','apsrev41Control']) # Remove annoying entries of revtex

    print("Seek:", cites)

    # Check what you already have in the bib file
    haves = []
    if os.path.isfile(bibfile):
        for line in open(bibfile,'r'):
            m = re.search(r'@.*?\{(.*),',line)  # .*\{ means "any # of any char followed by {"; .*?\{ means "the shortest string matching any # of any char followed by {"
            if m:
                haves.append(m.group(1))
    print("Have:", haves)

    # Query ADS and INSPIRE
    bibtex = open(bibfile,'a')      # open for appending

    for c in cites:
        if c and c not in haves: # c is something and you don't have it already

            if not c[0].isalpha(): # The first charachter is a number: could be on ADS

                try:
                    bib = ads_citation(c)
                    bibtex.write(bib)
                    print("ADS Found:", c)
                except:
                    print("ADS Not found:", c)

            else: # The first charachter is not a number: could be on INSPIRE

                try:
                    bib = inspire_citation(c)
                    bibtex.write(bib)
                    print("INSPIRE Found:", c)
                except:
                    print("INSPIRE Not found:", c)

    bibtex.close()
<commit_msg>Fix reference string in test_inspire<commit_after>#!/usr/bin/env python

'''
Query the ADS and INSPIRE databases to collect the citations used in a LaTeX file into a BibTeX file

Idea and main implementation is taken by Michele Vallisneri, see http://www.vallis.org/salon/summary-2.html
I changed the script to avoid the requests module, and added the INSPIRE database part.

Usage:
python fillbib.py <tex_file> <bib_file>. If <bib_file> is absent, it will try to guess it from the aux file"

'''
from __future__ import absolute_import, print_function
import sys, os, re, html

if sys.version_info.major>=3:
    import urllib.request as urllib
else:
    import urllib

def ads_citation(c): # download single ADS citation
    #f= urllib.urlopen("http://adsabs.harvard.edu/cgi-bin/nph-bib_query?bibcode="+c+"&data_type=BIBTEX&db_key=AST&nocookieset=1")
    #print("http://adsabs.harvard.edu/cgi-bin/nph-bib_query?bibcode="+c+"&data_type=BIBTEX&db_key=AST&nocookieset=1")
    f= urllib.urlopen("https://ui.adsabs.harvard.edu/abs/"+c+"/exportcitation")

    bib = f.read()
    f.close()
    if sys.version_info.major>=3:
        bib=bib.decode()
    bib = "@"+list(filter(lambda x:'adsnote' in x, bib.split("@")))[0].split("</textarea>")[0]
    bib=html.unescape(bib)

    if 'arXiv' in c: # Take care of preprint on ADS
        bib = bib.split("{")[0]+"{"+c+","+",".join(bib.split(",")[1:])
        return bib

    elif bib.split("{")[1].split(',')[0] == c: # Check you got what you where looking for
        return bib
    else:
        return None

def inspire_citation(c): # download single INSPIRE citation

    f= urllib.urlopen("https://old.inspirehep.net/search?p="+c+"&of=hx&em=B&sf=year&so=d&rg=1")
    bib = f.read()
    if sys.version_info.major>=3:
        bib=bib.decode()
    f.close()
    bib = "@"+bib.split("@")[1].split('
')[-2]

    if bib.split("{")[1].split(',')[0] == c: # Check you got what you where looking for
        return bib
    else:
        return None

def test_ads(): # test single ADS web scraping (both published articles and preprints)
    test_key = ["2016PhRvL.116f1102A","2016arXiv160203837T"]
    known_output= '@ARTICLE{2016PhRvL.116f1102A,\n   author = {{Abbott}, B.~P. and {Abbott}, R. and {Abbott}, T.~D. and {Abernathy}, M.~R. and \n\t{Acernese}, F. and {Ackley}, K. and {Adams}, C. and {Adams}, T. and \n\t{Addesso}, P. and {Adhikari}, R.~X. and et al.},\n    title = "{Observation of Gravitational Waves from a Binary Black Hole Merger}",\n  journal = {Physical Review Letters},\narchivePrefix = "arXiv",\n   eprint = {1602.03837},\n primaryClass = "gr-qc",\n     year = 2016,\n    month = feb,\n   volume = 116,\n   number = 6,\n      eid = {061102},\n    pages = {061102},\n      doi = {10.1103/PhysRevLett.116.061102},\n   adsurl = {http://adsabs.harvard.edu/abs/2016PhRvL.116f1102A},\n  adsnote = {Provided by the SAO/NASA Astrophysics Data System}\n}\n\n'
    assert [ads_citation(tk) == known_output for tk in test_key]

def test_inspire(): # test single INSPIRE web scraping
    test_key = "Abbott:2016blz"
    known_output= '@article{Abbott:2016blz,\n      author         = "Abbott, B. P. and others",\n      title          = "{Observation of Gravitational Waves from a Binary Black\n                        Hole Merger}",\n      collaboration  = "LIGO Scientific, Virgo",\n      journal        = "Phys. Rev. Lett.",\n      volume         = "116",\n      year           = "2016",\n      number         = "6",\n      pages          = "061102",\n      doi            = "10.1103/PhysRevLett.116.061102",\n      eprint         = "1602.03837",\n      archivePrefix  = "arXiv",\n      primaryClass   = "gr-qc",\n      reportNumber   = "LIGO-P150914",\n      SLACcitation   = "%%CITATION = ARXIV:1602.03837;%%"\n}\n'
    assert inspire_citation(test_key) == known_output

if __name__ == "__main__":

    if len(sys.argv)==2:     # Get the name of the bibfile from the aux file
        basename = sys.argv[1].split('.tex')[0]
        auxfile = basename + '.aux'
        for line in open(auxfile,'r'):
            m = re.search(r'\\bibdata\{(.*)\}',line)   # match \citation{...}, collect the ... note that we escape \, {, and }
            if m:
                bibfile = list(filter(lambda x:x!=basename+'Notes', m.group(1).split(',')))[0]  # Remove that annyoing feature of revtex which creates a *Notes.bib bibfile. Note this is not solid if you want to handle multiple bib files.
        bibfile = bibfile + '.bib'

    elif len(sys.argv)==3:    # Bibfile specified from argv
        basename = sys.argv[1].split('.tex')[0]
        auxfile = basename + '.aux'
        bibfile = sys.argv[2].split('.bib')[0] + '.bib'
    else:
        print("Usage: python fillbib.py <tex_file> <bib_file>. If <bib_file> is absent, assume the two are the same.")
        sys.exit()

    # Get all citations from aux file. Citations will look like \citation{2004PhRvD..69j4017P,2004PhRvD..69j4017P}
    cites = set() # use a set (no repetitions)
    for line in open(auxfile,'r'):
        m = re.search(r'\\citation\{(.*)\}',line)   # find \citation{...}
        if m:
            cites.update(m.group(1).split(','))     # split by commas

    cites= cites.difference(['REVTEX41Control','apsrev41Control']) # Remove annoying entries of revtex

    print("Seek:", cites)

    # Check what you already have in the bib file
    haves = []
    if os.path.isfile(bibfile):
        for line in open(bibfile,'r'):
            m = re.search(r'@.*?\{(.*),',line)  # .*\{ means "any # of any char followed by {"; .*?\{ means "the shortest string matching any # of any char followed by {"
            if m:
                haves.append(m.group(1))
    print("Have:", haves)

    # Query ADS and INSPIRE
    bibtex = open(bibfile,'a')      # open for appending

    for c in cites:
        if c and c not in haves: # c is something and you don't have it already

            if not c[0].isalpha(): # The first charachter is a number: could be on ADS

                try:
                    bib = ads_citation(c)
                    bibtex.write(bib)
                    print("ADS Found:", c)
                except:
                    print("ADS Not found:", c)

            else: # The first charachter is not a number: could be on INSPIRE

                try:
                    bib = inspire_citation(c)
                    bibtex.write(bib)
                    print("INSPIRE Found:", c)
                except:
                    print("INSPIRE Not found:", c)

    bibtex.close()
<|endoftext|>