<commit_before># -*- coding: utf-8 -*-

import httplib
import pprint
import re
import urllib
import urllib2

import bunch
import simplejson as json
from functools import partial
from urlobject import URLObject

p = "^\(#(\d+)\)"
code_re = re.compile(p)

__all__ = ['Graph']

class Graph(object):

    """
    Proxy for accessing the Facebook Graph API.

    This class uses dynamic attribute handling to provide a flexible and
    future-proof interface to the Graph API.

    Tutorial
    ========

    To get started using the API, create a new `Graph` instance with an access
    token:

        >>> g = Graph(access_token)  # Access token is optional.
        >>> g
        <Graph('https://graph.facebook.com/') at 0x...>

    Addressing Nodes
    ----------------

    Each `Graph` contains an access token and a URL. The graph you just created
    will have its URL set to 'https://graph.facebook.com/' by default (this is
    defined as the class attribute `Graph.API_ROOT`). The URL is represented as
    a `URLObject`; see <http://github.com/zacharyvoase/urlobject> for more
    information. Remember that you can treat it exactly as you would a `unicode`
    string; it just supports a few more methods to enable easy URL manipulation.

        >>> g.url
        <URLObject(u'https://graph.facebook.com/') at 0x...>
        >>> print g.url
        https://graph.facebook.com/

    To address child nodes within the Graph API, `Graph` supports dynamic
    attribute and item lookups:

        >>> g.me
        <Graph('https://graph.facebook.com/me') at 0x...>
        >>> g.me.home
        <Graph('https://graph.facebook.com/me/home') at 0x...>
        >>> g['me']['home']
        <Graph('https://graph.facebook.com/me/home') at 0x...>
        >>> g[123456789]
        <Graph('https://graph.facebook.com/123456789') at 0x...>

    Note that a `Graph` instance is rarely modified; these methods return copies
    of the original object. In addition, the API is lazy: HTTP requests will
    never be made unless you explicitly make them.

    Retrieving Nodes
    ----------------

    You can fetch data by calling a `Graph` instance:

        >>> about_me = g.me()
        >>> about_me
        Node({'about': '...', 'id': '1503223370'})

    This returns a `Node` object, which contains the retrieved data. `Node` is
    a subclass of `bunch.Bunch`, so you can access keys using attribute syntax:

        >>> about_me.id
        '1503223370'
        >>> about_me.first_name
        'Zachary'
        >>> about_me.hometown.name
        'London, United Kingdom'

    Accessing non-existent attributes or items will return a `Graph` instance
    corresponding to a child node. This `Graph` can then be called normally, to
    retrieve the child node it represents:

        >>> about_me.home
        <Graph('https://graph.facebook.com/me/home') at 0x...>
        >>> about_me.home()
        Node({'data': [...]})

    See `Node`’s documentation for further examples.

    Creating, Updating and Deleting Nodes
    -------------------------------------

    With the Graph API, node manipulation is done via HTTP POST requests. The
    `post()` method on `Graph` instances will POST to the current URL, with
    varying semantics for each endpoint:

        >>> post = g.me.feed.post(message="Test.")  # Status update
        >>> post
        Node({'id': '...'})
        >>> g[post.id].comments.post(message="A comment.") # Comment on the post
        Node({'id': '...'})
        >>> g[post.id].likes.post()  # Like the post
        True

        >>> event = g[121481007877204]()
        >>> event.name
        'Facebook Developer Garage London May 2010'
        >>> event.rsvp_status is None
        True
        >>> event.attending.post()  # Attend the given event
        True

    Deletes are just POST requests with `?method=delete`; the `delete()` method
    is a helpful shortcut:

        >>> g[post.id].delete()
        True

    """

    API_ROOT = URLObject.parse('https://graph.facebook.com/')

    def __init__(self, access_token=None, err_handler=None, **state):
        self.access_token = access_token
        self.err_handler = err_handler
        self.url = self.API_ROOT
        self.__dict__.update(state)

    def __repr__(self):
        return '<Graph(%r) at 0x%x>' % (str(self.url), id(self))

    def copy(self, **update):
        """Copy this Graph, optionally overriding some attributes."""

        return type(self)(access_token=self.access_token, err_handler=self.err_handler, **update)

    def __getitem__(self, item):
        if isinstance(item, slice):
            params = {'offset': item.start,
                      'limit': item.stop - item.start}
            return self.copy(url=(self.url & params))()
        return self.copy(url=(self.url / unicode(item)))

    def __getattr__(self, attr):
        return self[attr]

    def __or__(self, params):
        return self.copy(url=(self.url | params))

    def __and__(self, params):
        return self.copy(url=(self.url & params))

    def __call__(self, **params):
        """Read the current URL, and JSON-decode the results."""

        if self.access_token:
            params['access_token'] = self.access_token
        data = json.loads(self.fetch(self.url | params))
        return self.node(data)

    def fields(self, *fields):
        """Shortcut for `?fields=x,y,z`."""

        return self | ('fields', ','.join(fields))

    def ids(self, *ids):
        """Shortcut for `?ids=1,2,3`."""

        return self | ('ids', ','.join(map(str, ids)))

    def node(self, data):
        return Node._new(self, data, err_handler=self.err_handler)

    def post(self, **params):

        """
        POST to this URL (with parameters); return the JSON-decoded result.

        Example:

            >>> Graph('ACCESS TOKEN').me.feed.post(message="Test.")
            Node({'id': '...'})

        Some methods allow file attachments so uses MIME request to send those through.
        Must pass in a file object as 'file'
        """

        if self.access_token:
            params['access_token'] = self.access_token

        if self.url.path.split('/')[-1] in ['photos']:
            fetch = partial(self.post_mime, self.url, **params)
        else:
            params = dict([(k, v.encode('UTF-8')) for (k,v) in params.iteritems() if v is not None])
            fetch = partial(self.fetch, self.url, data=urllib.urlencode(params))

        data = json.loads(fetch())
        return self.node(data)

    def post_file(self, file, **params):

        if self.access_token:
            params['access_token'] = self.access_token
        params['file'] = file

        data = json.loads(self.post_mime(self.url, **params))

        return self.node(data)

    @staticmethod
    def post_mime(url, **kwargs):

        body = []
        crlf = '\r\n'
        boundary = "graphBoundary"

        # UTF8 params
        utf8_kwargs = dict([(k, v.encode('UTF-8')) for (k,v) in kwargs.iteritems() if k != 'file' and v is not None])

        # Add args
        for (k,v) in utf8_kwargs.iteritems():
            body.append("--"+boundary)
            body.append('Content-Disposition: form-data; name="%s"' % k)
            body.append('')
            body.append(str(v))

        # Add raw data
        file = kwargs.get('file')
        if file:
            file.open()
            data = file.read()
            file.close()

            body.append("--"+boundary)
            body.append('Content-Disposition: form-data; filename="facegraphfile.png"')
            body.append('')
            body.append(data)

            body.append("--"+boundary+"--")
            body.append('')

        body = crlf.join(body)

        # Post to server
        r = httplib.HTTPSConnection(url.host)
        headers = {'Content-Type': 'multipart/form-data; boundary=%s' % boundary,
                   'Content-Length': str(len(body)),
                   'MIME-Version': '1.0'}

        r.request('POST', url.path, body, headers)

        try:
            return r.getresponse().read()
        finally:
            r.close()

    def delete(self):
        """Delete this resource. Sends a POST with `?method=delete`."""

        return self.post(method='delete')

    @staticmethod
    def fetch(url, data=None):

        """
        Fetch the specified URL, with optional form data; return a string.

        This method exists mainly for dependency injection purposes. By default
        it uses urllib2; you may override it and use an alternative library.
        """
        conn = None
        try:
            conn = urllib2.urlopen(url, data=data)
            return conn.read()
        except urllib2.HTTPError, e:
            return e.fp.read()
        finally:
            conn and conn.close()

    def __sentry__(self):
        '''Transform the graph object into something that sentry can
        understand'''
        return "Graph(%s, %s)" % (self.url, str(self.__dict__))

class Node(bunch.Bunch):

    """
    Represent a JSON dictionary result from the Facebook Graph API.

    Accessing Items
    ---------------

    You can access items using either attribute or item syntax:

        >>> n = Node._new(None, {'a': 1, 'b': {'c': 3}})
        >>> n.a
        1
        >>> n['a']
        1
        >>> n.b.c
        3
        >>> n['b'].c
        3
        >>> n.b['c']
        3
        >>> n['b']['c']

    The same applies to assignment, although this is not recommended. Under the
    hood, this uses `bunch.Bunch` (<http://pypi.python.org/pypi/bunch>).

    Requesting Child Nodes
    ----------------------

    Many types of node in the Graph API have child nodes accessible through
    further API calls. For example:

        >>> g = Graph('access token')
        >>> g.me
        <Graph('https://graph.facebook.com/me') at 0x...>
        >>> g.me()  # User
        Node(...)
        >>> g.me.feed
        <Graph('https://graph.facebook.com/me/feed) at 0x...>
        >>> g.me.feed()  # News Feed for current user
        Node(...)

    `Node` makes it easier to access these child nodes using dynamic attribute
    handling:

        >>> me = g.me()
        >>> me.feed
        <Graph('https://graph.facebook.com/me/feed') at 0x...>

    Accessing a non-existent attribute or item will return a new `Graph`
    pointing to the child node, which you can then retrieve (as a `Node`) by
    calling directly.

    Because `Node` uses dynamic attribute handling, accessing non-existent
    attributes will still return a `Graph`. This can sometimes seem like more of
    a bug than a feature. You’ll have to check for the presence of a key first,
    to make sure what you're trying to access actually exists:

        def user_location(user_id):
            user = graph[user_id]()
            if 'location' in user:
                return user.location
            return False

    """

    @classmethod
    def _new(cls, api, data, err_handler=None):

        """
        Create a new `Node` from a `Graph` and a JSON-decoded object.

        If the object is not a dictionary, it will be returned as-is.
        """

        if isinstance(data, dict):
            if data.get("error"):
                code = data["error"].get("code")
                msg = data["error"]["message"]
                if code is None:
                    try:
                        code = int(code_re.match(msg).group(1))
                    except AttributeError:
                        pass
                e = GraphException(code, msg)
                if err_handler:
                    err_handler(e=e)
                else:
                    raise e
            return cls(api, bunch.bunchify(data))
        return data

    def __init__(self, api, data):
        super(Node, self).__init__(data)
        object.__setattr__(self, '_api', api)

    def __repr__(self):
        return 'Node(%r,\n%s)' % (
            self._api,
            indent(pprint.pformat(bunch.unbunchify(self)), prefix='     '))

    def __getitem__(self, item):
        try:
            return bunch.Bunch.__getitem__(self, item)
        except KeyError:
            return self._api[item]

    def __getattr__(self, attr):
        try:
            return bunch.Bunch.__getattr__(self, attr)
        except AttributeError:
            return self._api[attr]

    @property
    def next_page(self):
        """Shortcut for a `Graph` pointing to the next page."""

        return self._api.copy(url=URLObject.parse(self.paging.next))

    @property
    def previous_page(self):
        """Shortcut for a `Graph` pointing to the previous page."""

        return self._api.copy(url=URLObject.parse(self.paging.next))

class GraphException(Exception):
    def __init__(self, code, message, args=None):
        Exception.__init__(self)
        if args is not None:
            self.args = args
        self.message = message
        self.code = code

    def __repr__(self):
        return str(self)

    def __str__(self):
        return "%s: %s\n%s" % (self.code, self.message, self.args)

def indent(string, prefix='    '):
    """Indent each line of a string with a prefix (default is 4 spaces)."""

    return ''.join(prefix + line for line in string.splitlines(True))
<commit_msg>Dont print args on new line<commit_after># -*- coding: utf-8 -*-

import httplib
import pprint
import re
import urllib
import urllib2

import bunch
import simplejson as json
from functools import partial
from urlobject import URLObject

p = "^\(#(\d+)\)"
code_re = re.compile(p)

__all__ = ['Graph']

class Graph(object):

    """
    Proxy for accessing the Facebook Graph API.

    This class uses dynamic attribute handling to provide a flexible and
    future-proof interface to the Graph API.

    Tutorial
    ========

    To get started using the API, create a new `Graph` instance with an access
    token:

        >>> g = Graph(access_token)  # Access token is optional.
        >>> g
        <Graph('https://graph.facebook.com/') at 0x...>

    Addressing Nodes
    ----------------

    Each `Graph` contains an access token and a URL. The graph you just created
    will have its URL set to 'https://graph.facebook.com/' by default (this is
    defined as the class attribute `Graph.API_ROOT`). The URL is represented as
    a `URLObject`; see <http://github.com/zacharyvoase/urlobject> for more
    information. Remember that you can treat it exactly as you would a `unicode`
    string; it just supports a few more methods to enable easy URL manipulation.

        >>> g.url
        <URLObject(u'https://graph.facebook.com/') at 0x...>
        >>> print g.url
        https://graph.facebook.com/

    To address child nodes within the Graph API, `Graph` supports dynamic
    attribute and item lookups:

        >>> g.me
        <Graph('https://graph.facebook.com/me') at 0x...>
        >>> g.me.home
        <Graph('https://graph.facebook.com/me/home') at 0x...>
        >>> g['me']['home']
        <Graph('https://graph.facebook.com/me/home') at 0x...>
        >>> g[123456789]
        <Graph('https://graph.facebook.com/123456789') at 0x...>

    Note that a `Graph` instance is rarely modified; these methods return copies
    of the original object. In addition, the API is lazy: HTTP requests will
    never be made unless you explicitly make them.

    Retrieving Nodes
    ----------------

    You can fetch data by calling a `Graph` instance:

        >>> about_me = g.me()
        >>> about_me
        Node({'about': '...', 'id': '1503223370'})

    This returns a `Node` object, which contains the retrieved data. `Node` is
    a subclass of `bunch.Bunch`, so you can access keys using attribute syntax:

        >>> about_me.id
        '1503223370'
        >>> about_me.first_name
        'Zachary'
        >>> about_me.hometown.name
        'London, United Kingdom'

    Accessing non-existent attributes or items will return a `Graph` instance
    corresponding to a child node. This `Graph` can then be called normally, to
    retrieve the child node it represents:

        >>> about_me.home
        <Graph('https://graph.facebook.com/me/home') at 0x...>
        >>> about_me.home()
        Node({'data': [...]})

    See `Node`’s documentation for further examples.

    Creating, Updating and Deleting Nodes
    -------------------------------------

    With the Graph API, node manipulation is done via HTTP POST requests. The
    `post()` method on `Graph` instances will POST to the current URL, with
    varying semantics for each endpoint:

        >>> post = g.me.feed.post(message="Test.")  # Status update
        >>> post
        Node({'id': '...'})
        >>> g[post.id].comments.post(message="A comment.") # Comment on the post
        Node({'id': '...'})
        >>> g[post.id].likes.post()  # Like the post
        True

        >>> event = g[121481007877204]()
        >>> event.name
        'Facebook Developer Garage London May 2010'
        >>> event.rsvp_status is None
        True
        >>> event.attending.post()  # Attend the given event
        True

    Deletes are just POST requests with `?method=delete`; the `delete()` method
    is a helpful shortcut:

        >>> g[post.id].delete()
        True

    """

    API_ROOT = URLObject.parse('https://graph.facebook.com/')

    def __init__(self, access_token=None, err_handler=None, **state):
        self.access_token = access_token
        self.err_handler = err_handler
        self.url = self.API_ROOT
        self.__dict__.update(state)

    def __repr__(self):
        return '<Graph(%r) at 0x%x>' % (str(self.url), id(self))

    def copy(self, **update):
        """Copy this Graph, optionally overriding some attributes."""

        return type(self)(access_token=self.access_token, err_handler=self.err_handler, **update)

    def __getitem__(self, item):
        if isinstance(item, slice):
            params = {'offset': item.start,
                      'limit': item.stop - item.start}
            return self.copy(url=(self.url & params))()
        return self.copy(url=(self.url / unicode(item)))

    def __getattr__(self, attr):
        return self[attr]

    def __or__(self, params):
        return self.copy(url=(self.url | params))

    def __and__(self, params):
        return self.copy(url=(self.url & params))

    def __call__(self, **params):
        """Read the current URL, and JSON-decode the results."""

        if self.access_token:
            params['access_token'] = self.access_token
        data = json.loads(self.fetch(self.url | params))
        return self.node(data)

    def fields(self, *fields):
        """Shortcut for `?fields=x,y,z`."""

        return self | ('fields', ','.join(fields))

    def ids(self, *ids):
        """Shortcut for `?ids=1,2,3`."""

        return self | ('ids', ','.join(map(str, ids)))

    def node(self, data):
        return Node._new(self, data, err_handler=self.err_handler)

    def post(self, **params):

        """
        POST to this URL (with parameters); return the JSON-decoded result.

        Example:

            >>> Graph('ACCESS TOKEN').me.feed.post(message="Test.")
            Node({'id': '...'})

        Some methods allow file attachments so uses MIME request to send those through.
        Must pass in a file object as 'file'
        """

        if self.access_token:
            params['access_token'] = self.access_token

        if self.url.path.split('/')[-1] in ['photos']:
            fetch = partial(self.post_mime, self.url, **params)
        else:
            params = dict([(k, v.encode('UTF-8')) for (k,v) in params.iteritems() if v is not None])
            fetch = partial(self.fetch, self.url, data=urllib.urlencode(params))

        data = json.loads(fetch())
        return self.node(data)

    def post_file(self, file, **params):

        if self.access_token:
            params['access_token'] = self.access_token
        params['file'] = file

        data = json.loads(self.post_mime(self.url, **params))

        return self.node(data)

    @staticmethod
    def post_mime(url, **kwargs):

        body = []
        crlf = '\r\n'
        boundary = "graphBoundary"

        # UTF8 params
        utf8_kwargs = dict([(k, v.encode('UTF-8')) for (k,v) in kwargs.iteritems() if k != 'file' and v is not None])

        # Add args
        for (k,v) in utf8_kwargs.iteritems():
            body.append("--"+boundary)
            body.append('Content-Disposition: form-data; name="%s"' % k)
            body.append('')
            body.append(str(v))

        # Add raw data
        file = kwargs.get('file')
        if file:
            file.open()
            data = file.read()
            file.close()

            body.append("--"+boundary)
            body.append('Content-Disposition: form-data; filename="facegraphfile.png"')
            body.append('')
            body.append(data)

            body.append("--"+boundary+"--")
            body.append('')

        body = crlf.join(body)

        # Post to server
        r = httplib.HTTPSConnection(url.host)
        headers = {'Content-Type': 'multipart/form-data; boundary=%s' % boundary,
                   'Content-Length': str(len(body)),
                   'MIME-Version': '1.0'}

        r.request('POST', url.path, body, headers)

        try:
            return r.getresponse().read()
        finally:
            r.close()

    def delete(self):
        """Delete this resource. Sends a POST with `?method=delete`."""

        return self.post(method='delete')

    @staticmethod
    def fetch(url, data=None):

        """
        Fetch the specified URL, with optional form data; return a string.

        This method exists mainly for dependency injection purposes. By default
        it uses urllib2; you may override it and use an alternative library.
        """
        conn = None
        try:
            conn = urllib2.urlopen(url, data=data)
            return conn.read()
        except urllib2.HTTPError, e:
            return e.fp.read()
        finally:
            conn and conn.close()

    def __sentry__(self):
        '''Transform the graph object into something that sentry can
        understand'''
        return "Graph(%s, %s)" % (self.url, str(self.__dict__))

class Node(bunch.Bunch):

    """
    Represent a JSON dictionary result from the Facebook Graph API.

    Accessing Items
    ---------------

    You can access items using either attribute or item syntax:

        >>> n = Node._new(None, {'a': 1, 'b': {'c': 3}})
        >>> n.a
        1
        >>> n['a']
        1
        >>> n.b.c
        3
        >>> n['b'].c
        3
        >>> n.b['c']
        3
        >>> n['b']['c']

    The same applies to assignment, although this is not recommended. Under the
    hood, this uses `bunch.Bunch` (<http://pypi.python.org/pypi/bunch>).

    Requesting Child Nodes
    ----------------------

    Many types of node in the Graph API have child nodes accessible through
    further API calls. For example:

        >>> g = Graph('access token')
        >>> g.me
        <Graph('https://graph.facebook.com/me') at 0x...>
        >>> g.me()  # User
        Node(...)
        >>> g.me.feed
        <Graph('https://graph.facebook.com/me/feed) at 0x...>
        >>> g.me.feed()  # News Feed for current user
        Node(...)

    `Node` makes it easier to access these child nodes using dynamic attribute
    handling:

        >>> me = g.me()
        >>> me.feed
        <Graph('https://graph.facebook.com/me/feed') at 0x...>

    Accessing a non-existent attribute or item will return a new `Graph`
    pointing to the child node, which you can then retrieve (as a `Node`) by
    calling directly.

    Because `Node` uses dynamic attribute handling, accessing non-existent
    attributes will still return a `Graph`. This can sometimes seem like more of
    a bug than a feature. You’ll have to check for the presence of a key first,
    to make sure what you're trying to access actually exists:

        def user_location(user_id):
            user = graph[user_id]()
            if 'location' in user:
                return user.location
            return False

    """

    @classmethod
    def _new(cls, api, data, err_handler=None):

        """
        Create a new `Node` from a `Graph` and a JSON-decoded object.

        If the object is not a dictionary, it will be returned as-is.
        """

        if isinstance(data, dict):
            if data.get("error"):
                code = data["error"].get("code")
                msg = data["error"]["message"]
                if code is None:
                    try:
                        code = int(code_re.match(msg).group(1))
                    except AttributeError:
                        pass
                e = GraphException(code, msg)
                if err_handler:
                    err_handler(e=e)
                else:
                    raise e
            return cls(api, bunch.bunchify(data))
        return data

    def __init__(self, api, data):
        super(Node, self).__init__(data)
        object.__setattr__(self, '_api', api)

    def __repr__(self):
        return 'Node(%r,\n%s)' % (
            self._api,
            indent(pprint.pformat(bunch.unbunchify(self)), prefix='     '))

    def __getitem__(self, item):
        try:
            return bunch.Bunch.__getitem__(self, item)
        except KeyError:
            return self._api[item]

    def __getattr__(self, attr):
        try:
            return bunch.Bunch.__getattr__(self, attr)
        except AttributeError:
            return self._api[attr]

    @property
    def next_page(self):
        """Shortcut for a `Graph` pointing to the next page."""

        return self._api.copy(url=URLObject.parse(self.paging.next))

    @property
    def previous_page(self):
        """Shortcut for a `Graph` pointing to the previous page."""

        return self._api.copy(url=URLObject.parse(self.paging.next))

class GraphException(Exception):
    def __init__(self, code, message, args=None):
        Exception.__init__(self)
        if args is not None:
            self.args = args
        self.message = message
        self.code = code

    def __repr__(self):
        return str(self)

    def __str__(self):
        return "%s: %s %s" % (self.code, self.message, self.args)

def indent(string, prefix='    '):
    """Indent each line of a string with a prefix (default is 4 spaces)."""

    return ''.join(prefix + line for line in string.splitlines(True))
<|endoftext|>


<commit_before>"""
Allow local testing of mailer and templates by replaying an SQS message.

MAILER_FILE input is a file containing the exact base64-encoded, gzipped
data that's enqueued to SQS via :py:meth:`c7n.actions.Notify.send_sqs`.

Alternatively, with -p|--plain specified, the file will be assumed to be
JSON data that can be loaded directly.
"""
from __future__ import absolute_import, division, print_function, unicode_literals

import argparse
import base64
import json
import logging
import os
import zlib

import boto3
import jsonschema
from c7n_mailer.cli import CONFIG_SCHEMA
from c7n_mailer.email_delivery import EmailDelivery
from c7n_mailer.utils import setup_defaults
from c7n_mailer.utils_email import get_mimetext_message
from ruamel import yaml

logger = logging.getLogger(__name__)

class MailerTester(object):

    def __init__(self, msg_file, config, msg_plain=False, json_dump_file=None):
        if not os.path.exists(msg_file):
            raise RuntimeError("File does not exist: %s" % msg_file)
        logger.debug('Reading message from: %s', msg_file)
        with open(msg_file, 'r') as fh:
            raw = fh.read()
        logger.debug('Read %d byte message', len(raw))
        if msg_plain:
            raw = raw.strip()
        else:
            logger.debug('base64-decoding and zlib decompressing message')
            raw = zlib.decompress(base64.b64decode(raw))
            if json_dump_file is not None:
                with open(json_dump_file, 'wb') as fh:  # pragma: no cover
                    fh.write(raw)
        self.data = json.loads(raw)
        logger.debug('Loaded message JSON')
        self.config = config
        self.session = boto3.Session()

    def run(self, dry_run=False, print_only=False):
        emd = EmailDelivery(self.config, self.session, logger)
        addrs_to_msgs = emd.get_to_addrs_email_messages_map(self.data)
        logger.info('Would send email to: %s', addrs_to_msgs.keys())
        if print_only:
            mime = get_mimetext_message(
                self.config,
                logger,
                self.data,
                self.data['resources'],
                ['foo@example.com']
            )
            logger.info('Send mail with subject: "%s"', mime['Subject'])
            print(mime.get_payload(None, True))
            return
        if dry_run:
            for to_addrs, mimetext_msg in addrs_to_msgs.items():
                print('-> SEND MESSAGE TO: %s' % '; '.join(to_addrs))
                print(mimetext_msg.get_payload(None, True))
            return
        # else actually send the message...
        for to_addrs, mimetext_msg in addrs_to_msgs.items():
            logger.info('Actually sending mail to: %s', to_addrs)
            emd.send_c7n_email(self.data, list(to_addrs), mimetext_msg)

def setup_parser():
    parser = argparse.ArgumentParser('Test c7n-mailer templates and mail')
    parser.add_argument('-c', '--config', required=True)
    parser.add_argument('-d', '--dry-run', dest='dry_run', action='store_true',
                        default=False,
                        help='Log messages that would be sent, but do not send')
    parser.add_argument('-T', '--template-print', dest='print_only',
                        action='store_true', default=False,
                        help='Just print rendered templates')
    parser.add_argument('-t', '--templates', default=None, type=str,
                        help='message templates folder location')
    parser.add_argument('-p', '--plain', dest='plain', action='store_true',
                        default=False,
                        help='Expect MESSAGE_FILE to be a plain string, '
                             'rather than the base64-encoded, gzipped SQS '
                             'message format')
    parser.add_argument('-j', '--json-dump-file', dest='json_dump_file',
                        type=str, action='store', default=None,
                        help='If dump JSON of MESSAGE_FILE to this path; '
                             'useful to base64-decode and gunzip a message')
    parser.add_argument('MESSAGE_FILE', type=str,
                        help='Path to SQS message dump/content file')
    return parser

def session_factory(config):
    return boto3.Session(
        region_name=config['region'],
        profile_name=config.get('profile'))

def main():
    parser = setup_parser()
    options = parser.parse_args()

    module_dir = os.path.dirname(os.path.abspath(__file__))
    default_templates = [
        os.path.abspath(os.path.join(module_dir, 'msg-templates')),
        os.path.abspath(os.path.join(module_dir, '..', 'msg-templates')),
        os.path.abspath('.')
    ]
    templates = options.templates
    if templates:
        default_templates.append(
            os.path.abspath(os.path.expanduser(os.path.expandvars(templates)))
        )

    log_format = '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
    logging.basicConfig(level=logging.DEBUG, format=log_format)
    logging.getLogger('botocore').setLevel(logging.WARNING)

    with open(options.config) as fh:
        config = yaml.load(fh.read(), Loader=yaml.SafeLoader)

    jsonschema.validate(config, CONFIG_SCHEMA)
    setup_defaults(config)
    config['templates_folders'] = default_templates

    tester = MailerTester(
        options.MESSAGE_FILE, config, msg_plain=options.plain,
        json_dump_file=options.json_dump_file
    )
    tester.run(options.dry_run, options.print_only)

if __name__ == '__main__':
    main()
<commit_msg>Update string encoding for printing for py3 (#4440)<commit_after>"""
Allow local testing of mailer and templates by replaying an SQS message.

MAILER_FILE input is a file containing the exact base64-encoded, gzipped
data that's enqueued to SQS via :py:meth:`c7n.actions.Notify.send_sqs`.

Alternatively, with -p|--plain specified, the file will be assumed to be
JSON data that can be loaded directly.
"""
from __future__ import absolute_import, division, print_function, unicode_literals

import argparse
import base64
import json
import logging
import os
import zlib

import boto3
import jsonschema
from c7n_mailer.cli import CONFIG_SCHEMA
from c7n_mailer.email_delivery import EmailDelivery
from c7n_mailer.utils import setup_defaults
from c7n_mailer.utils_email import get_mimetext_message
from ruamel import yaml

logger = logging.getLogger(__name__)

class MailerTester(object):

    def __init__(self, msg_file, config, msg_plain=False, json_dump_file=None):
        if not os.path.exists(msg_file):
            raise RuntimeError("File does not exist: %s" % msg_file)
        logger.debug('Reading message from: %s', msg_file)
        with open(msg_file, 'r') as fh:
            raw = fh.read()
        logger.debug('Read %d byte message', len(raw))
        if msg_plain:
            raw = raw.strip()
        else:
            logger.debug('base64-decoding and zlib decompressing message')
            raw = zlib.decompress(base64.b64decode(raw))
            if json_dump_file is not None:
                with open(json_dump_file, 'wb') as fh:  # pragma: no cover
                    fh.write(raw)
        self.data = json.loads(raw)
        logger.debug('Loaded message JSON')
        self.config = config
        self.session = boto3.Session()

    def run(self, dry_run=False, print_only=False):
        emd = EmailDelivery(self.config, self.session, logger)
        addrs_to_msgs = emd.get_to_addrs_email_messages_map(self.data)
        logger.info('Would send email to: %s', addrs_to_msgs.keys())
        if print_only:
            mime = get_mimetext_message(
                self.config,
                logger,
                self.data,
                self.data['resources'],
                ['foo@example.com']
            )
            logger.info('Send mail with subject: "%s"', mime['Subject'])
            print(mime.get_payload(None, True).decode('utf-8'))
            return
        if dry_run:
            for to_addrs, mimetext_msg in addrs_to_msgs.items():
                print('-> SEND MESSAGE TO: %s' % '; '.join(to_addrs))
                print(mimetext_msg.get_payload(None, True).decode('utf-8'))
            return
        # else actually send the message...
        for to_addrs, mimetext_msg in addrs_to_msgs.items():
            logger.info('Actually sending mail to: %s', to_addrs)
            emd.send_c7n_email(self.data, list(to_addrs), mimetext_msg)

def setup_parser():
    parser = argparse.ArgumentParser('Test c7n-mailer templates and mail')
    parser.add_argument('-c', '--config', required=True)
    parser.add_argument('-d', '--dry-run', dest='dry_run', action='store_true',
                        default=False,
                        help='Log messages that would be sent, but do not send')
    parser.add_argument('-T', '--template-print', dest='print_only',
                        action='store_true', default=False,
                        help='Just print rendered templates')
    parser.add_argument('-t', '--templates', default=None, type=str,
                        help='message templates folder location')
    parser.add_argument('-p', '--plain', dest='plain', action='store_true',
                        default=False,
                        help='Expect MESSAGE_FILE to be a plain string, '
                             'rather than the base64-encoded, gzipped SQS '
                             'message format')
    parser.add_argument('-j', '--json-dump-file', dest='json_dump_file',
                        type=str, action='store', default=None,
                        help='If dump JSON of MESSAGE_FILE to this path; '
                             'useful to base64-decode and gunzip a message')
    parser.add_argument('MESSAGE_FILE', type=str,
                        help='Path to SQS message dump/content file')
    return parser

def session_factory(config):
    return boto3.Session(
        region_name=config['region'],
        profile_name=config.get('profile'))

def main():
    parser = setup_parser()
    options = parser.parse_args()

    module_dir = os.path.dirname(os.path.abspath(__file__))
    default_templates = [
        os.path.abspath(os.path.join(module_dir, 'msg-templates')),
        os.path.abspath(os.path.join(module_dir, '..', 'msg-templates')),
        os.path.abspath('.')
    ]
    templates = options.templates
    if templates:
        default_templates.append(
            os.path.abspath(os.path.expanduser(os.path.expandvars(templates)))
        )

    log_format = '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
    logging.basicConfig(level=logging.DEBUG, format=log_format)
    logging.getLogger('botocore').setLevel(logging.WARNING)

    with open(options.config) as fh:
        config = yaml.load(fh.read(), Loader=yaml.SafeLoader)

    jsonschema.validate(config, CONFIG_SCHEMA)
    setup_defaults(config)
    config['templates_folders'] = default_templates

    tester = MailerTester(
        options.MESSAGE_FILE, config, msg_plain=options.plain,
        json_dump_file=options.json_dump_file
    )
    tester.run(options.dry_run, options.print_only)

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


<commit_before>#
# Copyright 2013 Quantopian, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from copy import copy

import pytz
import pandas as pd
import numpy as np

from datetime import datetime

from itertools import groupby, ifilter
from operator import attrgetter

from zipline.errors import (
    UnsupportedSlippageModel,
    OverrideSlippagePostInit,
    UnsupportedCommissionModel,
    OverrideCommissionPostInit
)
from zipline.finance.performance import PerformanceTracker
from zipline.sources import DataFrameSource, DataPanelSource
from zipline.utils.factory import create_simulation_parameters
from zipline.transforms.utils import StatefulTransform
from zipline.finance.slippage import (
    VolumeShareSlippage,
    FixedSlippage,
    transact_partial
)
from zipline.finance.commission import PerShare, PerTrade
from zipline.finance.blotter import Blotter
from zipline.finance.constants import ANNUALIZER
import zipline.finance.trading as trading
import zipline.protocol
from zipline.protocol import Event

from zipline.gens.composites import (
    date_sorted_sources,
    sequential_transforms,
    alias_dt
)
from zipline.gens.tradesimulation import AlgorithmSimulator

DEFAULT_CAPITAL_BASE = float("1.0e5")

class TradingAlgorithm(object):
    """Base class for trading algorithms. Inherit and overload
    initialize() and handle_data(data).

    A new algorithm could look like this:
    ```
    class MyAlgo(TradingAlgorithm):
        def initialize(amount):
            self.amount = amount

        def handle_data(data):
            sid = self.sids[0]
            self.order(sid, amount)
    ```
    To then to run this algorithm:

    >>> my_algo = MyAlgo([0], 100) # first argument has to be list of sids
    >>> stats = my_algo.run(data)

    """
    def __init__(self, *args, **kwargs):
        """Initialize sids and other state variables.

        :Arguments:
            data_frequency : str (daily, hourly or minutely)
               The duration of the bars.
            annualizer : int <optional>
               Which constant to use for annualizing risk metrics.
               If not provided, will extract from data_frequency.
            capital_base : float <default: 1.0e5>
               How much capital to start with.
        """
        self._portfolio = None
        self.datetime = None

        self.registered_transforms = {}
        self.transforms = []
        self.sources = []

        self._recorded_vars = {}

        self.logger = None

        self.benchmark_return_source = None

        # default components for transact
        self.slippage = VolumeShareSlippage()
        self.commission = PerShare()

        if 'data_frequency' in kwargs:
            self.set_data_frequency(kwargs.pop('data_frequency'))
        else:
            self.data_frequency = None

        # Override annualizer if set
        if 'annualizer' in kwargs:
            self.annualizer = kwargs['annualizer']

        # set the capital base
        self.capital_base = kwargs.get('capital_base', DEFAULT_CAPITAL_BASE)

        self.sim_params = kwargs.pop('sim_params', None)
        if self.sim_params:
            self.sim_params.data_frequency = self.data_frequency

        self.blotter = kwargs.pop('blotter', Blotter())

        # an algorithm subclass needs to set initialized to True when
        # it is fully initialized.
        self.initialized = False

        # call to user-defined constructor method
        self.initialize(*args, **kwargs)

    def __repr__(self):
        """
        N.B. this does not yet represent a string that can be used
        to instantiate an exact copy of an algorithm.

        However, it is getting close, and provides some value as something
        that can be inspected interactively.
        """
        return """
{class_name}(
    captial_base={capital_base}
    sim_params={sim_params},
    initialized={initialized},
    slippage={slippage},
    commission={commission},
    blotter={blotter},
    recorded_vars={recorded_vars})
""".strip().format(class_name=self.__class__.__name__,
                   capital_base=self.capital_base,
                   sim_params=repr(self.sim_params),
                   initialized=self.initialized,
                   slippage=repr(self.slippage),
                   commission=repr(self.commission),
                   blotter=repr(self.blotter),
                   recorded_vars=repr(self.recorded_vars))

    def _create_data_generator(self, source_filter, sim_params):
        """
        Create a merged data generator using the sources and
        transforms attached to this algorithm.

        ::source_filter:: is a method that receives events in date
        sorted order, and returns True for those events that should be
        processed by the zipline, and False for those that should be
        skipped.
        """
        if self.benchmark_return_source is None:
            benchmark_return_source = [
                Event({'dt': ret.date,
                       'returns': ret.returns,
                       'type': zipline.protocol.DATASOURCE_TYPE.BENCHMARK,
                       'source_id': 'benchmarks'})
                for ret in trading.environment.benchmark_returns
                if ret.date.date() >= sim_params.period_start.date()
                and ret.date.date() <= sim_params.period_end.date()
            ]
        else:
            benchmark_return_source = self.benchmark_return_source

        date_sorted = date_sorted_sources(*self.sources)

        if source_filter:
            date_sorted = ifilter(source_filter, date_sorted)

        with_tnfms = sequential_transforms(date_sorted,
                                           *self.transforms)
        with_alias_dt = alias_dt(with_tnfms)

        with_benchmarks = date_sorted_sources(benchmark_return_source,
                                              with_alias_dt)

        # Group together events with the same dt field. This depends on the
        # events already being sorted.
        return groupby(with_benchmarks, attrgetter('dt'))

    def _create_generator(self, sim_params, source_filter=None):
        """
        Create a basic generator setup using the sources and
        transforms attached to this algorithm.

        ::source_filter:: is a method that receives events in date
        sorted order, and returns True for those events that should be
        processed by the zipline, and False for those that should be
        skipped.
        """
        sim_params.data_frequency = self.data_frequency

        self.data_gen = self._create_data_generator(source_filter,
                                                    sim_params)
        self.perf_tracker = PerformanceTracker(sim_params)
        self.trading_client = AlgorithmSimulator(self, sim_params)

        transact_method = transact_partial(self.slippage, self.commission)
        self.set_transact(transact_method)

        return self.trading_client.transform(self.data_gen)

    def get_generator(self):
        """
        Override this method to add new logic to the construction
        of the generator. Overrides can use the _create_generator
        method to get a standard construction generator.
        """
        return self._create_generator(self.sim_params)

    def initialize(self, *args, **kwargs):
        pass

    # TODO: make a new subclass, e.g. BatchAlgorithm, and move
    # the run method to the subclass, and refactor to put the
    # generator creation logic into get_generator.
    def run(self, source, sim_params=None, benchmark_return_source=None):
        """Run the algorithm.

        :Arguments:
            source : can be either:
                     - pandas.DataFrame
                     - zipline source
                     - list of zipline sources

               If pandas.DataFrame is provided, it must have the
               following structure:
               * column names must consist of ints representing the
                 different sids
               * index must be DatetimeIndex
               * array contents should be price info.

        :Returns:
            daily_stats : pandas.DataFrame
              Daily performance metrics such as returns, alpha etc.

        """
        if isinstance(source, (list, tuple)):
            assert self.sim_params is not None or sim_params is not None, \
                """When providing a list of sources, \
                sim_params have to be specified as a parameter
                or in the constructor."""
        elif isinstance(source, pd.DataFrame):
            # if DataFrame provided, wrap in DataFrameSource
            source = DataFrameSource(source)
        elif isinstance(source, pd.Panel):
            source = DataPanelSource(source)

        if not isinstance(source, (list, tuple)):
            self.sources = [source]
        else:
            self.sources = source

        # Check for override of sim_params.
        # If it isn't passed to this function,
        # use the default params set with the algorithm.
        # Else, we create simulation parameters using the start and end of the
        # source provided.
        if not sim_params:
            if not self.sim_params:
                start = source.start
                end = source.end

                sim_params = create_simulation_parameters(
                    start=start,
                    end=end,
                    capital_base=self.capital_base
                )
            else:
                sim_params = self.sim_params

        # Create transforms by wrapping them into StatefulTransforms
        self.transforms = []
        for namestring, trans_descr in self.registered_transforms.iteritems():
            sf = StatefulTransform(
                trans_descr['class'],
                *trans_descr['args'],
                **trans_descr['kwargs']
            )
            sf.namestring = namestring

            self.transforms.append(sf)

        # create transforms and zipline
        self.gen = self._create_generator(sim_params)

        # loop through simulated_trading, each iteration returns a
        # perf dictionary
        perfs = []
        for perf in self.gen:
            perfs.append(perf)

        # convert perf dict to pandas dataframe
        daily_stats = self._create_daily_stats(perfs)

        return daily_stats

    def _create_daily_stats(self, perfs):
        # create daily and cumulative stats dataframe
        daily_perfs = []
        cum_perfs = []
        # TODO: the loop here could overwrite expected properties
        # of daily_perf. Could potentially raise or log a
        # warning.
        for perf in perfs:
            if 'daily_perf' in perf:

                perf['daily_perf'].update(
                    perf['daily_perf'].pop('recorded_vars')
                )
                daily_perfs.append(perf['daily_perf'])
            else:
                cum_perfs.append(perf)

        daily_dts = [np.datetime64(perf['period_close'], utc=True)
                     for perf in daily_perfs]
        daily_stats = pd.DataFrame(daily_perfs, index=daily_dts)

        return daily_stats

    def add_transform(self, transform_class, tag, *args, **kwargs):
        """Add a single-sid, sequential transform to the model.

        :Arguments:
            transform_class : class
                Which transform to use. E.g. mavg.
            tag : str
                How to name the transform. Can later be access via:
                data[sid].tag()

        Extra args and kwargs will be forwarded to the transform
        instantiation.

        """
        self.registered_transforms[tag] = {'class': transform_class,
                                           'args': args,
                                           'kwargs': kwargs}

    def record(self, **kwargs):
        """
        Track and record local variable (i.e. attributes) each day.
        """
        for name, value in kwargs.items():
            self._recorded_vars[name] = value

    def order(self, sid, amount, limit_price=None, stop_price=None):
        return self.blotter.order(sid, amount, limit_price, stop_price)

    @property
    def recorded_vars(self):
        return copy(self._recorded_vars)

    @property
    def portfolio(self):
        return self._portfolio

    def set_portfolio(self, portfolio):
        self._portfolio = portfolio

    def set_logger(self, logger):
        self.logger = logger

    def set_datetime(self, dt):
        assert isinstance(dt, datetime), \
            "Attempt to set algorithm's current time with non-datetime"
        assert dt.tzinfo == pytz.utc, \
            "Algorithm expects a utc datetime"
        self.datetime = dt

    def get_datetime(self):
        """
        Returns a copy of the datetime.
        """
        date_copy = copy(self.datetime)
        assert date_copy.tzinfo == pytz.utc, \
            "Algorithm should have a utc datetime"
        return date_copy

    def set_transact(self, transact):
        """
        Set the method that will be called to create a
        transaction from open orders and trade events.
        """
        self.blotter.transact = transact

    def set_slippage(self, slippage):
        if not isinstance(slippage, (VolumeShareSlippage, FixedSlippage)):
            raise UnsupportedSlippageModel()
        if self.initialized:
            raise OverrideSlippagePostInit()
        self.slippage = slippage

    def set_commission(self, commission):
        if not isinstance(commission, (PerShare, PerTrade)):
            raise UnsupportedCommissionModel()

        if self.initialized:
            raise OverrideCommissionPostInit()
        self.commission = commission

    def set_sources(self, sources):
        assert isinstance(sources, list)
        self.sources = sources

    def set_transforms(self, transforms):
        assert isinstance(transforms, list)
        self.transforms = transforms

    def set_data_frequency(self, data_frequency):
        assert data_frequency in ('daily', 'minute')
        self.data_frequency = data_frequency
        self.annualizer = ANNUALIZER[self.data_frequency]
<commit_msg>MAINT: Ensure we're not setting the blotter to None.<commit_after>#
# Copyright 2013 Quantopian, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from copy import copy

import pytz
import pandas as pd
import numpy as np

from datetime import datetime

from itertools import groupby, ifilter
from operator import attrgetter

from zipline.errors import (
    UnsupportedSlippageModel,
    OverrideSlippagePostInit,
    UnsupportedCommissionModel,
    OverrideCommissionPostInit
)
from zipline.finance.performance import PerformanceTracker
from zipline.sources import DataFrameSource, DataPanelSource
from zipline.utils.factory import create_simulation_parameters
from zipline.transforms.utils import StatefulTransform
from zipline.finance.slippage import (
    VolumeShareSlippage,
    FixedSlippage,
    transact_partial
)
from zipline.finance.commission import PerShare, PerTrade
from zipline.finance.blotter import Blotter
from zipline.finance.constants import ANNUALIZER
import zipline.finance.trading as trading
import zipline.protocol
from zipline.protocol import Event

from zipline.gens.composites import (
    date_sorted_sources,
    sequential_transforms,
    alias_dt
)
from zipline.gens.tradesimulation import AlgorithmSimulator

DEFAULT_CAPITAL_BASE = float("1.0e5")

class TradingAlgorithm(object):
    """Base class for trading algorithms. Inherit and overload
    initialize() and handle_data(data).

    A new algorithm could look like this:
    ```
    class MyAlgo(TradingAlgorithm):
        def initialize(amount):
            self.amount = amount

        def handle_data(data):
            sid = self.sids[0]
            self.order(sid, amount)
    ```
    To then to run this algorithm:

    >>> my_algo = MyAlgo([0], 100) # first argument has to be list of sids
    >>> stats = my_algo.run(data)

    """
    def __init__(self, *args, **kwargs):
        """Initialize sids and other state variables.

        :Arguments:
            data_frequency : str (daily, hourly or minutely)
               The duration of the bars.
            annualizer : int <optional>
               Which constant to use for annualizing risk metrics.
               If not provided, will extract from data_frequency.
            capital_base : float <default: 1.0e5>
               How much capital to start with.
        """
        self._portfolio = None
        self.datetime = None

        self.registered_transforms = {}
        self.transforms = []
        self.sources = []

        self._recorded_vars = {}

        self.logger = None

        self.benchmark_return_source = None

        # default components for transact
        self.slippage = VolumeShareSlippage()
        self.commission = PerShare()

        if 'data_frequency' in kwargs:
            self.set_data_frequency(kwargs.pop('data_frequency'))
        else:
            self.data_frequency = None

        # Override annualizer if set
        if 'annualizer' in kwargs:
            self.annualizer = kwargs['annualizer']

        # set the capital base
        self.capital_base = kwargs.get('capital_base', DEFAULT_CAPITAL_BASE)

        self.sim_params = kwargs.pop('sim_params', None)
        if self.sim_params:
            self.sim_params.data_frequency = self.data_frequency

        self.blotter = kwargs.pop('blotter', None)
        if not self.blotter:
            self.blotter = Blotter()

        # an algorithm subclass needs to set initialized to True when
        # it is fully initialized.
        self.initialized = False

        # call to user-defined constructor method
        self.initialize(*args, **kwargs)

    def __repr__(self):
        """
        N.B. this does not yet represent a string that can be used
        to instantiate an exact copy of an algorithm.

        However, it is getting close, and provides some value as something
        that can be inspected interactively.
        """
        return """
{class_name}(
    captial_base={capital_base}
    sim_params={sim_params},
    initialized={initialized},
    slippage={slippage},
    commission={commission},
    blotter={blotter},
    recorded_vars={recorded_vars})
""".strip().format(class_name=self.__class__.__name__,
                   capital_base=self.capital_base,
                   sim_params=repr(self.sim_params),
                   initialized=self.initialized,
                   slippage=repr(self.slippage),
                   commission=repr(self.commission),
                   blotter=repr(self.blotter),
                   recorded_vars=repr(self.recorded_vars))

    def _create_data_generator(self, source_filter, sim_params):
        """
        Create a merged data generator using the sources and
        transforms attached to this algorithm.

        ::source_filter:: is a method that receives events in date
        sorted order, and returns True for those events that should be
        processed by the zipline, and False for those that should be
        skipped.
        """
        if self.benchmark_return_source is None:
            benchmark_return_source = [
                Event({'dt': ret.date,
                       'returns': ret.returns,
                       'type': zipline.protocol.DATASOURCE_TYPE.BENCHMARK,
                       'source_id': 'benchmarks'})
                for ret in trading.environment.benchmark_returns
                if ret.date.date() >= sim_params.period_start.date()
                and ret.date.date() <= sim_params.period_end.date()
            ]
        else:
            benchmark_return_source = self.benchmark_return_source

        date_sorted = date_sorted_sources(*self.sources)

        if source_filter:
            date_sorted = ifilter(source_filter, date_sorted)

        with_tnfms = sequential_transforms(date_sorted,
                                           *self.transforms)
        with_alias_dt = alias_dt(with_tnfms)

        with_benchmarks = date_sorted_sources(benchmark_return_source,
                                              with_alias_dt)

        # Group together events with the same dt field. This depends on the
        # events already being sorted.
        return groupby(with_benchmarks, attrgetter('dt'))

    def _create_generator(self, sim_params, source_filter=None):
        """
        Create a basic generator setup using the sources and
        transforms attached to this algorithm.

        ::source_filter:: is a method that receives events in date
        sorted order, and returns True for those events that should be
        processed by the zipline, and False for those that should be
        skipped.
        """
        sim_params.data_frequency = self.data_frequency

        self.data_gen = self._create_data_generator(source_filter,
                                                    sim_params)
        self.perf_tracker = PerformanceTracker(sim_params)
        self.trading_client = AlgorithmSimulator(self, sim_params)

        transact_method = transact_partial(self.slippage, self.commission)
        self.set_transact(transact_method)

        return self.trading_client.transform(self.data_gen)

    def get_generator(self):
        """
        Override this method to add new logic to the construction
        of the generator. Overrides can use the _create_generator
        method to get a standard construction generator.
        """
        return self._create_generator(self.sim_params)

    def initialize(self, *args, **kwargs):
        pass

    # TODO: make a new subclass, e.g. BatchAlgorithm, and move
    # the run method to the subclass, and refactor to put the
    # generator creation logic into get_generator.
    def run(self, source, sim_params=None, benchmark_return_source=None):
        """Run the algorithm.

        :Arguments:
            source : can be either:
                     - pandas.DataFrame
                     - zipline source
                     - list of zipline sources

               If pandas.DataFrame is provided, it must have the
               following structure:
               * column names must consist of ints representing the
                 different sids
               * index must be DatetimeIndex
               * array contents should be price info.

        :Returns:
            daily_stats : pandas.DataFrame
              Daily performance metrics such as returns, alpha etc.

        """
        if isinstance(source, (list, tuple)):
            assert self.sim_params is not None or sim_params is not None, \
                """When providing a list of sources, \
                sim_params have to be specified as a parameter
                or in the constructor."""
        elif isinstance(source, pd.DataFrame):
            # if DataFrame provided, wrap in DataFrameSource
            source = DataFrameSource(source)
        elif isinstance(source, pd.Panel):
            source = DataPanelSource(source)

        if not isinstance(source, (list, tuple)):
            self.sources = [source]
        else:
            self.sources = source

        # Check for override of sim_params.
        # If it isn't passed to this function,
        # use the default params set with the algorithm.
        # Else, we create simulation parameters using the start and end of the
        # source provided.
        if not sim_params:
            if not self.sim_params:
                start = source.start
                end = source.end

                sim_params = create_simulation_parameters(
                    start=start,
                    end=end,
                    capital_base=self.capital_base
                )
            else:
                sim_params = self.sim_params

        # Create transforms by wrapping them into StatefulTransforms
        self.transforms = []
        for namestring, trans_descr in self.registered_transforms.iteritems():
            sf = StatefulTransform(
                trans_descr['class'],
                *trans_descr['args'],
                **trans_descr['kwargs']
            )
            sf.namestring = namestring

            self.transforms.append(sf)

        # create transforms and zipline
        self.gen = self._create_generator(sim_params)

        # loop through simulated_trading, each iteration returns a
        # perf dictionary
        perfs = []
        for perf in self.gen:
            perfs.append(perf)

        # convert perf dict to pandas dataframe
        daily_stats = self._create_daily_stats(perfs)

        return daily_stats

    def _create_daily_stats(self, perfs):
        # create daily and cumulative stats dataframe
        daily_perfs = []
        cum_perfs = []
        # TODO: the loop here could overwrite expected properties
        # of daily_perf. Could potentially raise or log a
        # warning.
        for perf in perfs:
            if 'daily_perf' in perf:

                perf['daily_perf'].update(
                    perf['daily_perf'].pop('recorded_vars')
                )
                daily_perfs.append(perf['daily_perf'])
            else:
                cum_perfs.append(perf)

        daily_dts = [np.datetime64(perf['period_close'], utc=True)
                     for perf in daily_perfs]
        daily_stats = pd.DataFrame(daily_perfs, index=daily_dts)

        return daily_stats

    def add_transform(self, transform_class, tag, *args, **kwargs):
        """Add a single-sid, sequential transform to the model.

        :Arguments:
            transform_class : class
                Which transform to use. E.g. mavg.
            tag : str
                How to name the transform. Can later be access via:
                data[sid].tag()

        Extra args and kwargs will be forwarded to the transform
        instantiation.

        """
        self.registered_transforms[tag] = {'class': transform_class,
                                           'args': args,
                                           'kwargs': kwargs}

    def record(self, **kwargs):
        """
        Track and record local variable (i.e. attributes) each day.
        """
        for name, value in kwargs.items():
            self._recorded_vars[name] = value

    def order(self, sid, amount, limit_price=None, stop_price=None):
        return self.blotter.order(sid, amount, limit_price, stop_price)

    @property
    def recorded_vars(self):
        return copy(self._recorded_vars)

    @property
    def portfolio(self):
        return self._portfolio

    def set_portfolio(self, portfolio):
        self._portfolio = portfolio

    def set_logger(self, logger):
        self.logger = logger

    def set_datetime(self, dt):
        assert isinstance(dt, datetime), \
            "Attempt to set algorithm's current time with non-datetime"
        assert dt.tzinfo == pytz.utc, \
            "Algorithm expects a utc datetime"
        self.datetime = dt

    def get_datetime(self):
        """
        Returns a copy of the datetime.
        """
        date_copy = copy(self.datetime)
        assert date_copy.tzinfo == pytz.utc, \
            "Algorithm should have a utc datetime"
        return date_copy

    def set_transact(self, transact):
        """
        Set the method that will be called to create a
        transaction from open orders and trade events.
        """
        self.blotter.transact = transact

    def set_slippage(self, slippage):
        if not isinstance(slippage, (VolumeShareSlippage, FixedSlippage)):
            raise UnsupportedSlippageModel()
        if self.initialized:
            raise OverrideSlippagePostInit()
        self.slippage = slippage

    def set_commission(self, commission):
        if not isinstance(commission, (PerShare, PerTrade)):
            raise UnsupportedCommissionModel()

        if self.initialized:
            raise OverrideCommissionPostInit()
        self.commission = commission

    def set_sources(self, sources):
        assert isinstance(sources, list)
        self.sources = sources

    def set_transforms(self, transforms):
        assert isinstance(transforms, list)
        self.transforms = transforms

    def set_data_frequency(self, data_frequency):
        assert data_frequency in ('daily', 'minute')
        self.data_frequency = data_frequency
        self.annualizer = ANNUALIZER[self.data_frequency]
<|endoftext|>


<commit_before>
from abc import ABC, abstractmethod
from collections import OrderedDict
import json
import jsonschema
import os
import re
from typing import List, NamedTuple
import yaml

from type_info import TypeInfo, BitfieldInfo, FlagInfo, EnumInfo, EnumeratorInfo, ClassInfo, FunctionInfo, ArgInfo, AttributeInfo, ResolvedTypeRef
from type_registry import TypeRegistry, TypeNameRef

class SourceInfo(NamedTuple):
    file: str
    line: int

# Source: https://stackoverflow.com/a/53647080/3621512
class SafeLineLoader(yaml.SafeLoader):
    def __init__(self, document_name, stream):
        super(SafeLineLoader, self).__init__(stream)
        self.document_name = document_name
        self.linenos = {}

    def compose_node(self, parent, index):
        # the line number where the previous token has ended (plus empty lines)
        line = self.line
        node = super(SafeLineLoader, self).compose_node(parent, index)
        node.__line__ = line + 1
        if parent is None:
            node.path = []
        return node

    def construct_mapping(self, node, deep=False):
        self.flatten_mapping(node)
        for key_node, value_node in node.value:
            value_node.path = node.path + [key_node.value]
        pairs = self.construct_pairs(node)
        if len(set([k for k, v in pairs])) != len(pairs):
                raise Exception("duplicate keys")
        #import ipdb; ipdb.set_trace()
        self.linenos['.'.join(node.path)] = SourceInfo(self.document_name, node.__line__)
        return OrderedDict(pairs)

def get_dict(elem, key):
    return elem.get(key, None) or OrderedDict()

dictionary = [] # TODO

def get_words(string):
    """
    Splits a string in PascalCase or MACRO_CASE into a list of lower case words
    """
    if string.isupper():
        return [w.lower() for w in string.split('_')]
    else:
        regex = ''.join((re.escape(w) + '|') for w in dictionary) + '[a-z0-9]+|[A-Z][a-z0-9]*'
        return [(w if w in dictionary else w.lower()) for w in re.findall(regex, string)]

def to_pascal_case(s): return ''.join([(w.title() if not w in dictionary else w) for w in get_words(s)])

class Loader():
    def __init__(self, registry):
        script_path = os.path.dirname(os.path.realpath(__file__))
        with open(os.path.join(script_path, 'interface-schema.json')) as fp:
            schema = json.load(fp)
        self._validator = jsonschema.Draft4Validator(schema)
        self._registry = registry

        ns = registry.global_namespace.ns_from_path(('io', 'fibreframework'), construct_if_missing=True)
        ns.add_type(OpaqueTypeInfo('int8'), py_ref_type=('Int8Property',), py_val_type=('int',))
        ns.add_type(OpaqueTypeInfo('uint8'), py_ref_type=('Uint8Property',), py_val_type=('int',))
        ns.add_type(OpaqueTypeInfo('int16'), py_ref_type=('Int16Property',), py_val_type=('int',))
        ns.add_type(OpaqueTypeInfo('uint16'), py_ref_type=('Uint16Property',), py_val_type=('int',))
        ns.add_type(OpaqueTypeInfo('int32'), py_ref_type=('Int32Property',), py_val_type=('int',))
        ns.add_type(OpaqueTypeInfo('uint32'), py_ref_type=('Uint32Property',), py_val_type=('int',))
        ns.add_type(OpaqueTypeInfo('int64'), py_ref_type=('Int64Property',), py_val_type=('int',))
        ns.add_type(OpaqueTypeInfo('uint64'), py_ref_type=('Uint64Property',), py_val_type=('int',))
        ns.add_type(OpaqueTypeInfo('float32'), py_ref_type=('Float32Property',), py_val_type=('float',))
        ns.add_type(OpaqueTypeInfo('bool'), py_ref_type=('BoolProperty',), py_val_type=('bool',))
        ns.add_type(OpaqueTypeInfo('endpoint_ref'), py_ref_type=('EndpointRefProperty',), py_val_type=('RemoteObject',))

    def load_enum_from_fata(self, enum, enum_data):
        pass # TODO

    def get_property_interface(self, codec):
        return None # TODO

    def load_args_from_data(self, ns_path, arg_data, using):
        for k, arg_data in arg_data.items():
            if isinstance(arg_data, str) or (not 'type' in arg_data):
                arg_data = {'type': arg_data}

            if isinstance(arg_data['type'], str):
                type_ref = TypeNameRef(self._registry, [ns_path, *using], arg_data['type'])
            else:
                type = self.load_type_from_data(ns_path, to_pascal_case(k), arg_data['type'], using)
                type_ref = ResolvedTypeRef(type)

            yield ArgInfo(k, type_ref, arg_data.get('doc', None))

    def load_type_from_data(self, ns_path, type_name, type_data, using: List[List['NamespaceInfo']]):
        ns = self._registry.global_namespace.ns_from_path(ns_path, construct_if_missing=True)

        if ('flags' in type_data):
            bitfield = ns.get_type(type_name, kind=BitfieldInfo, construct_if_missing=True)

            for k, flag_data in type_data['flags'].items():
                if flag_data is None:
                    flag_data = {}
                bitfield.flags.append(FlagInfo(k,
                    flag_data.get('bit', bitfield.get_next_bit()),
                    flag_data.get('brief', None),
                    flag_data.get('doc', None)
                ))

            if 'nullflag' in type_data:
                bitfield.nullflag = type_data['nullflag']

            return bitfield

        elif ('values' in type_data):
            enum = ns.get_type(type_name, kind=EnumInfo, construct_if_missing=True)

            for k, flag_data in type_data['values'].items():
                if flag_data is None:
                    flag_data = {}
                enum.enumerators.append(EnumeratorInfo(k,
                    flag_data.get('bit', enum.get_next_value()),
                    flag_data.get('brief', None),
                    flag_data.get('doc', None)
                ))

            return enum

        else:
            # class type

            cls = ns.get_type(type_name, kind=ClassInfo, construct_if_missing=True)

            for k, attr_data in get_dict(type_data, 'attributes').items():
                if isinstance(attr_data, str):
                    attr_data = {'type': attr_data}
                elif not 'type' in attr_data:
                    attr_data['type'] = {}
                    if 'attributes' in attr_data: attr_data['type']['attributes'] = attr_data.pop('attributes')
                    if 'functions' in attr_data: attr_data['type']['functions'] = attr_data.pop('functions')
                    if 'implements' in attr_data: attr_data['type']['implements'] = attr_data.pop('implements')
                    if 'c_is_class' in attr_data: attr_data['type']['c_is_class'] = attr_data.pop('c_is_class')
                    if 'values' in attr_data: attr_data['type']['values'] = attr_data.pop('values')
                    if 'flags' in attr_data: attr_data['type']['flags'] = attr_data.pop('flags')
                    if 'nullflag' in attr_data: attr_data['type']['nullflag'] = attr_data.pop('nullflag')

                if isinstance(attr_data['type'], str):
                    if attr_data['type'].startswith('readonly '):
                        attr_data['type'] = attr_data['type'][len('readonly '):]
                        attr_data['readonly'] = True
                    type_ref = TypeNameRef(self._registry, [(*ns_path, cls.name), *using], attr_data['type'])
                else:
                    type = self.load_type_from_data((*ns_path, cls.name), to_pascal_case(k), attr_data['type'], using)
                    type_ref = ResolvedTypeRef(type)

                cls.attributes.append(AttributeInfo(k,
                    type_ref,
                    attr_data.get('readonly', False),
                    attr_data.get('brief', None),
                    attr_data.get('doc', None)
                ))

            for k, func_data in get_dict(type_data, 'functions').items():
                if func_data is None:
                    func_data = {}
                in_args = list(self.load_args_from_data((*ns_path, cls.name), get_dict(func_data or {}, 'in'), using))
                out_args = list(self.load_args_from_data((*ns_path, cls.name), get_dict(func_data or {}, 'out'), using))
                cls.functions.append(FunctionInfo(k, in_args, out_args, func_data.get('brief', None), func_data.get('doc', None)))

            return cls

    def load_from_data(self, data):
        using = [tuple(u.split('.')) for u in (data.get('using', None) or [])]
        ns = tuple(data['ns'].split('.'))

        all_types = [
            *get_dict(data, 'interfaces').items(), # TODO: deprecate
            *get_dict(data, 'valuetypes').items(), # TODO: deprecate
            *get_dict(data, 'types').items(),
        ]

        for k, type_data in all_types:
            if k.startswith(':'):
                current_ns = ()
                k = k[1:]
            else:
                current_ns = ns
            long_type_name = k.split('.')
            self.load_type_from_data(current_ns + tuple(long_type_name[:-1]), long_type_name[-1], type_data, using)

    def load_from_yaml_file(self, file: str):
        with open(file) as fp:
            loader = SafeLineLoader(file, fp)
            file_content = loader.get_single_data()

        for err in self._validator.iter_errors(file_content):
            # TODO: print line number
            raise Exception(err.message + '\nat ' + str(list(err.absolute_path)))

        file_content['using'] = ['io.fibreframework'] # TODO: remove
        self.load_from_data(file_content)

class OpaqueTypeInfo(TypeInfo):
    def __init__(self, name):
        self.name = name

if __name__ == "__main__":
    registry = TypeRegistry()

    loader = Loader(registry)
    loader.load_from_yaml_file('../../Firmware/odrive-interface.yaml')

    registry.resolve_all()

    intf1 = registry.get_class('com.odriverobotics.ODrive')
    intf2 = registry.get_class('com.odriverobotics.ODrive.Axis')
    intf3 = registry.get_class('com.odriverobotics.ODrive.Axis.Config')
    intf4 = registry.get_class('com.odriverobotics.ODrive.Motor')

    print(registry.get_py_ref_type_name(('com', 'odriverobotics', 'ODrive'), intf3))
    print(registry.get_py_ref_type_name(('com', 'odriverobotics', 'ODrive'), intf4))

    import ipdb; ipdb.set_trace()

<commit_msg>[docs] fix enum values<commit_after>
from abc import ABC, abstractmethod
from collections import OrderedDict
import json
import jsonschema
import os
import re
from typing import List, NamedTuple
import yaml

from type_info import TypeInfo, BitfieldInfo, FlagInfo, EnumInfo, EnumeratorInfo, ClassInfo, FunctionInfo, ArgInfo, AttributeInfo, ResolvedTypeRef
from type_registry import TypeRegistry, TypeNameRef

class SourceInfo(NamedTuple):
    file: str
    line: int

# Source: https://stackoverflow.com/a/53647080/3621512
class SafeLineLoader(yaml.SafeLoader):
    def __init__(self, document_name, stream):
        super(SafeLineLoader, self).__init__(stream)
        self.document_name = document_name
        self.linenos = {}

    def compose_node(self, parent, index):
        # the line number where the previous token has ended (plus empty lines)
        line = self.line
        node = super(SafeLineLoader, self).compose_node(parent, index)
        node.__line__ = line + 1
        if parent is None:
            node.path = []
        return node

    def construct_mapping(self, node, deep=False):
        self.flatten_mapping(node)
        for key_node, value_node in node.value:
            value_node.path = node.path + [key_node.value]
        pairs = self.construct_pairs(node)
        if len(set([k for k, v in pairs])) != len(pairs):
                raise Exception("duplicate keys")
        #import ipdb; ipdb.set_trace()
        self.linenos['.'.join(node.path)] = SourceInfo(self.document_name, node.__line__)
        return OrderedDict(pairs)

def get_dict(elem, key):
    return elem.get(key, None) or OrderedDict()

dictionary = [] # TODO

def get_words(string):
    """
    Splits a string in PascalCase or MACRO_CASE into a list of lower case words
    """
    if string.isupper():
        return [w.lower() for w in string.split('_')]
    else:
        regex = ''.join((re.escape(w) + '|') for w in dictionary) + '[a-z0-9]+|[A-Z][a-z0-9]*'
        return [(w if w in dictionary else w.lower()) for w in re.findall(regex, string)]

def to_pascal_case(s): return ''.join([(w.title() if not w in dictionary else w) for w in get_words(s)])

class Loader():
    def __init__(self, registry):
        script_path = os.path.dirname(os.path.realpath(__file__))
        with open(os.path.join(script_path, 'interface-schema.json')) as fp:
            schema = json.load(fp)
        self._validator = jsonschema.Draft4Validator(schema)
        self._registry = registry

        ns = registry.global_namespace.ns_from_path(('io', 'fibreframework'), construct_if_missing=True)
        ns.add_type(OpaqueTypeInfo('int8'), py_ref_type=('Int8Property',), py_val_type=('int',))
        ns.add_type(OpaqueTypeInfo('uint8'), py_ref_type=('Uint8Property',), py_val_type=('int',))
        ns.add_type(OpaqueTypeInfo('int16'), py_ref_type=('Int16Property',), py_val_type=('int',))
        ns.add_type(OpaqueTypeInfo('uint16'), py_ref_type=('Uint16Property',), py_val_type=('int',))
        ns.add_type(OpaqueTypeInfo('int32'), py_ref_type=('Int32Property',), py_val_type=('int',))
        ns.add_type(OpaqueTypeInfo('uint32'), py_ref_type=('Uint32Property',), py_val_type=('int',))
        ns.add_type(OpaqueTypeInfo('int64'), py_ref_type=('Int64Property',), py_val_type=('int',))
        ns.add_type(OpaqueTypeInfo('uint64'), py_ref_type=('Uint64Property',), py_val_type=('int',))
        ns.add_type(OpaqueTypeInfo('float32'), py_ref_type=('Float32Property',), py_val_type=('float',))
        ns.add_type(OpaqueTypeInfo('bool'), py_ref_type=('BoolProperty',), py_val_type=('bool',))
        ns.add_type(OpaqueTypeInfo('endpoint_ref'), py_ref_type=('EndpointRefProperty',), py_val_type=('RemoteObject',))

    def load_enum_from_fata(self, enum, enum_data):
        pass # TODO

    def get_property_interface(self, codec):
        return None # TODO

    def load_args_from_data(self, ns_path, arg_data, using):
        for k, arg_data in arg_data.items():
            if isinstance(arg_data, str) or (not 'type' in arg_data):
                arg_data = {'type': arg_data}

            if isinstance(arg_data['type'], str):
                type_ref = TypeNameRef(self._registry, [ns_path, *using], arg_data['type'])
            else:
                type = self.load_type_from_data(ns_path, to_pascal_case(k), arg_data['type'], using)
                type_ref = ResolvedTypeRef(type)

            yield ArgInfo(k, type_ref, arg_data.get('doc', None))

    def load_type_from_data(self, ns_path, type_name, type_data, using: List[List['NamespaceInfo']]):
        ns = self._registry.global_namespace.ns_from_path(ns_path, construct_if_missing=True)

        if ('flags' in type_data):
            bitfield = ns.get_type(type_name, kind=BitfieldInfo, construct_if_missing=True)

            for k, flag_data in type_data['flags'].items():
                if flag_data is None:
                    flag_data = {}
                bitfield.flags.append(FlagInfo(k,
                    flag_data.get('bit', bitfield.get_next_bit()),
                    flag_data.get('brief', None),
                    flag_data.get('doc', None)
                ))

            if 'nullflag' in type_data:
                bitfield.nullflag = type_data['nullflag']

            return bitfield

        elif ('values' in type_data):
            enum = ns.get_type(type_name, kind=EnumInfo, construct_if_missing=True)

            for k, flag_data in type_data['values'].items():
                if flag_data is None:
                    flag_data = {}
                enum.enumerators.append(EnumeratorInfo(k,
                    flag_data.get('value', enum.get_next_value()),
                    flag_data.get('brief', None),
                    flag_data.get('doc', None)
                ))

            return enum

        else:
            # class type

            cls = ns.get_type(type_name, kind=ClassInfo, construct_if_missing=True)

            for k, attr_data in get_dict(type_data, 'attributes').items():
                if isinstance(attr_data, str):
                    attr_data = {'type': attr_data}
                elif not 'type' in attr_data:
                    attr_data['type'] = {}
                    if 'attributes' in attr_data: attr_data['type']['attributes'] = attr_data.pop('attributes')
                    if 'functions' in attr_data: attr_data['type']['functions'] = attr_data.pop('functions')
                    if 'implements' in attr_data: attr_data['type']['implements'] = attr_data.pop('implements')
                    if 'c_is_class' in attr_data: attr_data['type']['c_is_class'] = attr_data.pop('c_is_class')
                    if 'values' in attr_data: attr_data['type']['values'] = attr_data.pop('values')
                    if 'flags' in attr_data: attr_data['type']['flags'] = attr_data.pop('flags')
                    if 'nullflag' in attr_data: attr_data['type']['nullflag'] = attr_data.pop('nullflag')

                if isinstance(attr_data['type'], str):
                    if attr_data['type'].startswith('readonly '):
                        attr_data['type'] = attr_data['type'][len('readonly '):]
                        attr_data['readonly'] = True
                    type_ref = TypeNameRef(self._registry, [(*ns_path, cls.name), *using], attr_data['type'])
                else:
                    type = self.load_type_from_data((*ns_path, cls.name), to_pascal_case(k), attr_data['type'], using)
                    type_ref = ResolvedTypeRef(type)

                cls.attributes.append(AttributeInfo(k,
                    type_ref,
                    attr_data.get('readonly', False),
                    attr_data.get('brief', None),
                    attr_data.get('doc', None)
                ))

            for k, func_data in get_dict(type_data, 'functions').items():
                if func_data is None:
                    func_data = {}
                in_args = list(self.load_args_from_data((*ns_path, cls.name), get_dict(func_data or {}, 'in'), using))
                out_args = list(self.load_args_from_data((*ns_path, cls.name), get_dict(func_data or {}, 'out'), using))
                cls.functions.append(FunctionInfo(k, in_args, out_args, func_data.get('brief', None), func_data.get('doc', None)))

            return cls

    def load_from_data(self, data):
        using = [tuple(u.split('.')) for u in (data.get('using', None) or [])]
        ns = tuple(data['ns'].split('.'))

        all_types = [
            *get_dict(data, 'interfaces').items(), # TODO: deprecate
            *get_dict(data, 'valuetypes').items(), # TODO: deprecate
            *get_dict(data, 'types').items(),
        ]

        for k, type_data in all_types:
            if k.startswith(':'):
                current_ns = ()
                k = k[1:]
            else:
                current_ns = ns
            long_type_name = k.split('.')
            self.load_type_from_data(current_ns + tuple(long_type_name[:-1]), long_type_name[-1], type_data, using)

    def load_from_yaml_file(self, file: str):
        with open(file) as fp:
            loader = SafeLineLoader(file, fp)
            file_content = loader.get_single_data()

        for err in self._validator.iter_errors(file_content):
            # TODO: print line number
            raise Exception(err.message + '\nat ' + str(list(err.absolute_path)))

        file_content['using'] = ['io.fibreframework'] # TODO: remove
        self.load_from_data(file_content)

class OpaqueTypeInfo(TypeInfo):
    def __init__(self, name):
        self.name = name

if __name__ == "__main__":
    registry = TypeRegistry()

    loader = Loader(registry)
    loader.load_from_yaml_file('../../Firmware/odrive-interface.yaml')

    registry.resolve_all()

    intf1 = registry.get_class('com.odriverobotics.ODrive')
    intf2 = registry.get_class('com.odriverobotics.ODrive.Axis')
    intf3 = registry.get_class('com.odriverobotics.ODrive.Axis.Config')
    intf4 = registry.get_class('com.odriverobotics.ODrive.Motor')

    print(registry.get_py_ref_type_name(('com', 'odriverobotics', 'ODrive'), intf3))
    print(registry.get_py_ref_type_name(('com', 'odriverobotics', 'ODrive'), intf4))

    import ipdb; ipdb.set_trace()

<|endoftext|>


<commit_before>from __future__ import division

from pkg_resources import get_distribution, DistributionNotFound, RequirementParseError

def read_requirements(requirements_file):
    with open(requirements_file, 'r') as f:
        return f.read().splitlines()

def get_module_name(package, on_error='raise'):
    try:
        t = list(get_distribution(package)._get_metadata('top_level.txt'))
        if t:
            return t[0]
        else:
            return None
    except Exception as e:
        if on_error == 'raise':
            raise
        elif on_error == 'error_class':
            return e.__class__.__name__
        else:
            return on_error  # just the value specified by on_error
<commit_msg>set_of_to_level_import_names<commit_after>from __future__ import division

import os
import re
from pkg_resources import get_distribution, DistributionNotFound, RequirementParseError
from ut.pfile.iter import get_filepath_iterator
import ut.pfile.to

def read_requirements(requirements_file):
    with open(requirements_file, 'r') as f:
        return f.read().splitlines()

def get_module_name(package, on_error='raise'):
    try:
        t = list(get_distribution(package)._get_metadata('top_level.txt'))
        if t:
            return t[0]
        else:
            return None
    except Exception as e:
        if on_error == 'raise':
            raise
        elif on_error == 'error_class':
            return e.__class__.__name__
        else:
            return on_error  # just the value specified by on_error

def set_of_to_level_import_names(folder_path_to_scan=None, filt_set=None):
    """
    Scan a folder recursively, picking up the .py files, seeing what names are imported
    (i.e. from NAME.* and import NAME.*), collecting the full set of these, and then
    intersecting with the filt_set.
    :param folder_path_to_scan: Folder path to scan recursively
    :param filt_set: names to filter in (i.e. match from the full set of import names found
    :return: A set of import names found (of those in the filt_set that is)
    """
    if folder_path_to_scan is None:
        folder_path_to_scan = '/D/Dropbox/dev/py/proj/span/'
    if filt_set is None:
        filt_set = set(os.listdir('/D/Dropbox/dev/py/proj/')).union('oto', 'scan')

    p = re.compile('import \w+|from \w+')
    cumul = set()
    for f in get_filepath_iterator(folder_path_to_scan, '.py$'):
        s = ut.pfile.to.string(f)
        x = map(lambda x: x.split(' ')[1], p.findall(s))
        cumul = cumul.union(x)
    return cumul.intersection(filt_set)
<|endoftext|>