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

from .common import match1, download_urls, parse_host, set_proxy, unset_proxy
from .util import log

class Extractor():
    def __init__(self, *args):
        self.url = None
        self.title = None
        self.vid = None
        self.streams = {}
        self.streams_sorted = []

        if args:
            self.url = args[0]

class VideoExtractor():
    def __init__(self, *args):
        self.url = None
        self.title = None
        self.vid = None
        self.streams = {}
        self.streams_sorted = []
        self.audiolang = None
        self.password_protected = False

        if args:
            self.url = args[0]

    def download_by_url(self, url, **kwargs):
        self.url = url

        if 'extractor_proxy' in kwargs and kwargs['extractor_proxy']:
            set_proxy(parse_host(kwargs['extractor_proxy']))
        self.prepare(**kwargs)

        try:
            self.streams_sorted = [dict([('id', stream_type['id'])] + list(self.streams[stream_type['id']].items())) for stream_type in self.__class__.stream_types if stream_type['id'] in self.streams]
        except:
            self.streams_sorted = [dict([('itag', stream_type['itag'])] + list(self.streams[stream_type['itag']].items())) for stream_type in self.__class__.stream_types if stream_type['itag'] in self.streams]

        self.extract(**kwargs)

        if 'extractor_proxy' in kwargs and kwargs['extractor_proxy']:
            unset_proxy()

        self.download(**kwargs)

    def download_by_vid(self, vid, **kwargs):
        self.vid = vid

        if 'extractor_proxy' in kwargs and kwargs['extractor_proxy']:
            set_proxy(parse_host(kwargs['extractor_proxy']))
        self.prepare(**kwargs)

        try:
            self.streams_sorted = [dict([('id', stream_type['id'])] + list(self.streams[stream_type['id']].items())) for stream_type in self.__class__.stream_types if stream_type['id'] in self.streams]
        except:
            self.streams_sorted = [dict([('itag', stream_type['itag'])] + list(self.streams[stream_type['itag']].items())) for stream_type in self.__class__.stream_types if stream_type['itag'] in self.streams]

        self.extract(**kwargs)
        if 'extractor_proxy' in kwargs and kwargs['extractor_proxy']:
            unset_proxy()

        self.download(**kwargs)

    def prepare(self, **kwargs):
        pass
        #raise NotImplementedError()

    def extract(self, **kwargs):
        pass
        #raise NotImplementedError()

    def p_stream(self, stream_id):
        stream = self.streams[stream_id]
        if 'itag' in stream:
            print("    - itag:          %s" % log.sprint(stream_id, log.NEGATIVE))
        else:
            print("    - format:        %s" % log.sprint(stream_id, log.NEGATIVE))

        if 'container' in stream:
            print("      container:     %s" % stream['container'])

        if 'video_profile' in stream:
            print("      video-profile: %s" % stream['video_profile'])

        if 'quality' in stream:
            print("      quality:       %s" % stream['quality'])

        if 'size' in stream:
            print("      size:          %s MiB (%s bytes)" % (round(stream['size'] / 1048576, 1), stream['size']))

        if 'itag' in stream:
            print("    # download-with: %s" % log.sprint("you-get --itag=%s [URL]" % stream_id, log.UNDERLINE))
        else:
            print("    # download-with: %s" % log.sprint("you-get --format=%s [URL]" % stream_id, log.UNDERLINE))

        print()

    def p_i(self, stream_id):
        stream = self.streams[stream_id]
        print("    - title:         %s" % self.title)
        print("       size:         %s MiB (%s bytes)" % (round(stream['size'] / 1048576, 1), stream['size']))
        print("        url:         %s" % self.url)
        print()

    def p(self, stream_id=None):
        print("site:                %s" % self.__class__.name)
        print("title:               %s" % self.title)
        if stream_id:
            # Print the stream
            print("stream:")
            self.p_stream(stream_id)

        elif stream_id is None:
            # Print stream with best quality
            print("stream:              # Best quality")
            stream_id = self.streams_sorted[0]['id'] if 'id' in self.streams_sorted[0] else self.streams_sorted[0]['itag']
            self.p_stream(stream_id)

        elif stream_id == []:
            # Print all available streams
            print("streams:             # Available quality and codecs")
            for stream in self.streams_sorted:
                self.p_stream(stream['id'] if 'id' in stream else stream['itag'])

        if self.audiolang:
            print("audio-languages:")
            for i in self.audiolang:
                print("    - lang:          {}".format(i['lang']))
                print("      download-url:  {}\n".format(i['url']))

    def p_playlist(self, stream_id=None):
        print("site:                %s" % self.__class__.name)
        print("playlist:            %s" % self.title)
        print("videos:")

    def download(self, **kwargs):
        if 'info_only' in kwargs and kwargs['info_only']:
            if 'stream_id' in kwargs and kwargs['stream_id']:
                # Display the stream
                stream_id = kwargs['stream_id']
                if 'index' not in kwargs:
                    self.p(stream_id)
                else:
                    self.p_i(stream_id)
            else:
                # Display all available streams
                if 'index' not in kwargs:
                    self.p([])
                else:
                    stream_id = self.streams_sorted[0]['id'] if 'id' in self.streams_sorted[0] else self.streams_sorted[0]['itag']
                    self.p_i(stream_id)

        else:
            if 'stream_id' in kwargs and kwargs['stream_id']:
                # Download the stream
                stream_id = kwargs['stream_id']
            else:
                # Download stream with the best quality
                stream_id = self.streams_sorted[0]['id'] if 'id' in self.streams_sorted[0] else self.streams_sorted[0]['itag']

            if 'index' not in kwargs:
                self.p(stream_id)
            else:
                self.p_i(stream_id)

            urls = self.streams[stream_id]['src']
            if not urls:
                log.e('[Failed] Cannot extract video source.')
                log.e('This is most likely because the video has not been made available in your country.')
                log.e('You may try to use a proxy via \'-y\' for extracting stream data.')
                exit(1)
            # For legacy main()
            download_urls(urls, self.title, self.streams[stream_id]['container'], self.streams[stream_id]['size'], output_dir=kwargs['output_dir'], merge=kwargs['merge'])
            # For main_dev()
            #download_urls(urls, self.title, self.streams[stream_id]['container'], self.streams[stream_id]['size'])

        self.__init__()
<commit_msg>VideoExtractor: replace misleading error message<commit_after>#!/usr/bin/env python

from .common import match1, download_urls, parse_host, set_proxy, unset_proxy
from .util import log

class Extractor():
    def __init__(self, *args):
        self.url = None
        self.title = None
        self.vid = None
        self.streams = {}
        self.streams_sorted = []

        if args:
            self.url = args[0]

class VideoExtractor():
    def __init__(self, *args):
        self.url = None
        self.title = None
        self.vid = None
        self.streams = {}
        self.streams_sorted = []
        self.audiolang = None
        self.password_protected = False

        if args:
            self.url = args[0]

    def download_by_url(self, url, **kwargs):
        self.url = url

        if 'extractor_proxy' in kwargs and kwargs['extractor_proxy']:
            set_proxy(parse_host(kwargs['extractor_proxy']))
        self.prepare(**kwargs)

        try:
            self.streams_sorted = [dict([('id', stream_type['id'])] + list(self.streams[stream_type['id']].items())) for stream_type in self.__class__.stream_types if stream_type['id'] in self.streams]
        except:
            self.streams_sorted = [dict([('itag', stream_type['itag'])] + list(self.streams[stream_type['itag']].items())) for stream_type in self.__class__.stream_types if stream_type['itag'] in self.streams]

        self.extract(**kwargs)

        if 'extractor_proxy' in kwargs and kwargs['extractor_proxy']:
            unset_proxy()

        self.download(**kwargs)

    def download_by_vid(self, vid, **kwargs):
        self.vid = vid

        if 'extractor_proxy' in kwargs and kwargs['extractor_proxy']:
            set_proxy(parse_host(kwargs['extractor_proxy']))
        self.prepare(**kwargs)

        try:
            self.streams_sorted = [dict([('id', stream_type['id'])] + list(self.streams[stream_type['id']].items())) for stream_type in self.__class__.stream_types if stream_type['id'] in self.streams]
        except:
            self.streams_sorted = [dict([('itag', stream_type['itag'])] + list(self.streams[stream_type['itag']].items())) for stream_type in self.__class__.stream_types if stream_type['itag'] in self.streams]

        self.extract(**kwargs)
        if 'extractor_proxy' in kwargs and kwargs['extractor_proxy']:
            unset_proxy()

        self.download(**kwargs)

    def prepare(self, **kwargs):
        pass
        #raise NotImplementedError()

    def extract(self, **kwargs):
        pass
        #raise NotImplementedError()

    def p_stream(self, stream_id):
        stream = self.streams[stream_id]
        if 'itag' in stream:
            print("    - itag:          %s" % log.sprint(stream_id, log.NEGATIVE))
        else:
            print("    - format:        %s" % log.sprint(stream_id, log.NEGATIVE))

        if 'container' in stream:
            print("      container:     %s" % stream['container'])

        if 'video_profile' in stream:
            print("      video-profile: %s" % stream['video_profile'])

        if 'quality' in stream:
            print("      quality:       %s" % stream['quality'])

        if 'size' in stream:
            print("      size:          %s MiB (%s bytes)" % (round(stream['size'] / 1048576, 1), stream['size']))

        if 'itag' in stream:
            print("    # download-with: %s" % log.sprint("you-get --itag=%s [URL]" % stream_id, log.UNDERLINE))
        else:
            print("    # download-with: %s" % log.sprint("you-get --format=%s [URL]" % stream_id, log.UNDERLINE))

        print()

    def p_i(self, stream_id):
        stream = self.streams[stream_id]
        print("    - title:         %s" % self.title)
        print("       size:         %s MiB (%s bytes)" % (round(stream['size'] / 1048576, 1), stream['size']))
        print("        url:         %s" % self.url)
        print()

    def p(self, stream_id=None):
        print("site:                %s" % self.__class__.name)
        print("title:               %s" % self.title)
        if stream_id:
            # Print the stream
            print("stream:")
            self.p_stream(stream_id)

        elif stream_id is None:
            # Print stream with best quality
            print("stream:              # Best quality")
            stream_id = self.streams_sorted[0]['id'] if 'id' in self.streams_sorted[0] else self.streams_sorted[0]['itag']
            self.p_stream(stream_id)

        elif stream_id == []:
            # Print all available streams
            print("streams:             # Available quality and codecs")
            for stream in self.streams_sorted:
                self.p_stream(stream['id'] if 'id' in stream else stream['itag'])

        if self.audiolang:
            print("audio-languages:")
            for i in self.audiolang:
                print("    - lang:          {}".format(i['lang']))
                print("      download-url:  {}\n".format(i['url']))

    def p_playlist(self, stream_id=None):
        print("site:                %s" % self.__class__.name)
        print("playlist:            %s" % self.title)
        print("videos:")

    def download(self, **kwargs):
        if 'info_only' in kwargs and kwargs['info_only']:
            if 'stream_id' in kwargs and kwargs['stream_id']:
                # Display the stream
                stream_id = kwargs['stream_id']
                if 'index' not in kwargs:
                    self.p(stream_id)
                else:
                    self.p_i(stream_id)
            else:
                # Display all available streams
                if 'index' not in kwargs:
                    self.p([])
                else:
                    stream_id = self.streams_sorted[0]['id'] if 'id' in self.streams_sorted[0] else self.streams_sorted[0]['itag']
                    self.p_i(stream_id)

        else:
            if 'stream_id' in kwargs and kwargs['stream_id']:
                # Download the stream
                stream_id = kwargs['stream_id']
            else:
                # Download stream with the best quality
                stream_id = self.streams_sorted[0]['id'] if 'id' in self.streams_sorted[0] else self.streams_sorted[0]['itag']

            if 'index' not in kwargs:
                self.p(stream_id)
            else:
                self.p_i(stream_id)

            urls = self.streams[stream_id]['src']
            if not urls:
                log.wtf('[Failed] Cannot extract video source.')
            # For legacy main()
            download_urls(urls, self.title, self.streams[stream_id]['container'], self.streams[stream_id]['size'], output_dir=kwargs['output_dir'], merge=kwargs['merge'])
            # For main_dev()
            #download_urls(urls, self.title, self.streams[stream_id]['container'], self.streams[stream_id]['size'])

        self.__init__()
<|endoftext|>


<commit_before>"""Module for accessing github.com projects."""

import base64
import httplib
import json
import re
import sys
import traceback
import yaml

from mimic.__mimic import common

import model
import shared

from . import collection

from google.appengine.api import urlfetch_errors
from google.appengine.ext import deferred

_GITHUB_URL_RE = re.compile(
    '^https?://(?:[^/]+.)?github.com/(?:users/)?([^/]+)/?([^/]+)?.*$'
)

def IsValidUrl(url):
  return _GITHUB_URL_RE.match(url)

def FetchWithAuth(*args, **kwargs):
  credential = model.GetOAuth2Credential('github')
  if credential:
    query_str = ('?client_id={0}&client_secret={1}'
                 .format(credential.client_id,
                         credential.client_secret))
    # tuple is immutable
    args = list(args)
    args[0] += query_str
  return shared.Fetch(*args, **kwargs)

class GithubRepoCollection(collection.RepoCollection):
  """A class for accessing github code repositories."""

  def __init__(self, repo_collection):
    super(GithubRepoCollection, self).__init__(repo_collection)

  def _IsAppEnginePythonRepo(self, name):
    """Determine whether the given repo name is an App Engine Python project.

    Repo names must meet the following conditions:
    - Starts with 'appengine-'
    - Have a 'python' component
    - Not have a 'java' or 'go' component

    Args:
      name: The github repo name.

    Returns:
      True if the given repo name is an App Engine Python Project.
    """
    name = name.lower()
    if not name.startswith('appengine-'):
      return False
    words = name.split('-')
    if 'python' not in words:
      return False
    if 'java' in words or 'go' in words:
      return False
    return True

  def _GetAppEnginePythonRepos(self, page):
    """Get list of App Engine Python repos.

    Given a JSON list of repos, return those repo names which appear to be
    Python App Engine repos. This function can parse the contents of:
    https://api.github.com/users/GoogleCloudPlatform/repos

    Args:
      page: the JSON response returned by /users/:user/repos

    Returns:
      A list of repos.
    """
    r = json.loads(page)
    # keys we care about: html_url, contents_url, name, description
    repos = [entry for entry in r if self._IsAppEnginePythonRepo(entry['name'])]

    candidates1 = []
    for repo in repos:
      # only proceed with repos which look like App Engine Python projects
      if not self._IsAppEnginePythonRepo(repo['name']):
        continue

      # e.g. https://api.github.com/repos/GoogleCloudPlatform/appengine-crowdguru-python/contents/app.yaml
      app_yaml_contents_url = repo['contents_url'].replace('{+path}',
                                                           'app.yaml')
      rpc = FetchWithAuth(app_yaml_contents_url, async=True)
      candidates1.append((rpc, app_yaml_contents_url, repo))

    # filter repos with no app.yaml
    candidates2 = []
    for rpc, app_yaml_contents_url, repo in candidates1:
      result = rpc.get_result()
      if result.status_code != httplib.OK:
        shared.w('Skipping repo due to {} fetching {}'
                 .format(result.status_code, app_yaml_contents_url))
        continue
      r = json.loads(result.content)
      app_yaml_git_url = r['git_url']
      rpc = FetchWithAuth(app_yaml_git_url, async=True)
      candidates2.append((rpc, app_yaml_contents_url, app_yaml_git_url, repo))

    # filter repos whose app.yaml does not contain 'runtime: python27'
    repos = []
    for rpc, app_yaml_contents_url, app_yaml_git_url, repo in candidates2:
      result = rpc.get_result()
      r = json.loads(result.content)
      base64_content = r['content']
      decoded_content = base64.b64decode(base64_content)
      config = yaml.safe_load(decoded_content)
      runtime = config.get('runtime')
      if runtime != 'python27':
        shared.w('Skipping "runtime: {}" app {}'.format(runtime,
                                                        app_yaml_contents_url))
        continue
      repos.append(repo)

    return repos

  def _GetRepoContents(self, repo_contents_url):
    """Get list of files/directories in the given repo.

    Args:
      page: the JSON response returned by /repos/:owner/:repo/contents

    Returns:
      The list of files/directories.
    """
    page = FetchWithAuth(repo_contents_url, follow_redirects=True).content
    r = json.loads(page)
    entries = []
    for entry in r:
      if entry['type'] == 'dir':
        entries.extend(self._GetRepoContents(entry['url']))
      elif entry['type'] == 'file':
        entries.append((entry['path'], entry['git_url']))
    return entries

  def PopulateRepos(self):
    shared.EnsureRunningInTask()  # gives us automatic retries
    repo_collection_url = self.repo_collection.key.id()
    matcher = _GITHUB_URL_RE.match(repo_collection_url)
    github_user = matcher.group(1)
    # e.g. https://api.github.com/users/GoogleCloudPlatform/repos
    url = 'https://api.github.com/users/{0}/repos'.format(github_user)
    rpc_result = FetchWithAuth(url, follow_redirects=True)
    page = rpc_result.content
    repos = self._GetAppEnginePythonRepos(page)

    credential = model.GetOAuth2Credential('github')
    if (not credential or not credential.client_id
        or not credential.client_secret):
      # fetch fewer when we're not authenticated
      repos = repos[:1]

    repo_entities = []
    for repo in repos:
      r = model.CreateRepo(repo['html_url'], end_user_url=repo['html_url'],
                           name=repo['name'],
                           description=repo['description'] or repo['html_url'])
      repo_entities.append(r)
    model.ndb.put_multi(repo_entities)
    for repo_entity in repo_entities:
      deferred.defer(self.CreateTemplateProject, repo_entity.key)

  def CreateProjectTreeFromRepo(self, tree, repo):
    # e.g. https://github.com/GoogleCloudPlatform/appengine-crowdguru-python
    end_user_repo_url = repo.key.id()
    matcher = _GITHUB_URL_RE.match(end_user_repo_url)
    github_user = matcher.group(1)
    repo_name = matcher.group(2)
    # e.g. https://api.github.com/repos/GoogleCloudPlatform/appengine-crowdguru-python/contents/
    repo_contents_url = ('https://api.github.com/repos/{0}/{1}/contents/'
                         .format(github_user, repo_name))

    entries = self._GetRepoContents(repo_contents_url)

    rpcs = []
    for (path, file_git_url) in entries:
      rpc = FetchWithAuth(file_git_url, follow_redirects=True, async=True)
      rpcs.append((file_git_url, path, rpc))

    for file_git_url, path, rpc in rpcs:
      try:
        result = rpc.get_result()
        shared.w('{0} {1} {2}'.format(result.status_code, path, file_git_url))
        if result.status_code != 200:
          continue
        r = json.loads(result.content)
        base64_content = r['content']
        decoded_content = base64.b64decode(base64_content)
        tree.SetFile(path, decoded_content)
      except urlfetch_errors.Error:
        exc_info = sys.exc_info()
        formatted_exception = traceback.format_exception(exc_info[0],
                                                         exc_info[1],
                                                         exc_info[2])
        shared.w('Skipping {0} {1}'.format(path, file_git_url))
        for line in [line for line in formatted_exception if line]:
          shared.w(line)
<commit_msg>introduce github project blacklist<commit_after>"""Module for accessing github.com projects."""

import base64
import httplib
import json
import re
import sys
import traceback
import yaml

from mimic.__mimic import common

import model
import shared

from . import collection

from google.appengine.api import urlfetch_errors
from google.appengine.ext import deferred

_GITHUB_URL_RE = re.compile(
    '^https?://(?:[^/]+.)?github.com/(?:users/)?([^/]+)/?([^/]+)?.*$'
)

# projects which should not be shown in the cloud playground by default
_PROJECT_URL_SKIP_LIST = [
    # Java users only
    'https://github.com/GoogleCloudPlatform/appengine-enable-remote-api-python',
    # (deprecated) master/slave apps only
    'https://github.com/GoogleCloudPlatform/appengine-recover-unapplied-writes-python',
]

def IsValidUrl(url):
  return _GITHUB_URL_RE.match(url)

def FetchWithAuth(*args, **kwargs):
  credential = model.GetOAuth2Credential('github')
  if credential:
    query_str = ('?client_id={0}&client_secret={1}'
                 .format(credential.client_id,
                         credential.client_secret))
    # tuple is immutable
    args = list(args)
    args[0] += query_str
  return shared.Fetch(*args, **kwargs)

class GithubRepoCollection(collection.RepoCollection):
  """A class for accessing github code repositories."""

  def __init__(self, repo_collection):
    super(GithubRepoCollection, self).__init__(repo_collection)

  def _IsAppEnginePythonRepo(self, name):
    """Determine whether the given repo name is an App Engine Python project.

    Repo names must meet the following conditions:
    - Starts with 'appengine-'
    - Have a 'python' component
    - Not have a 'java' or 'go' component

    Args:
      name: The github repo name.

    Returns:
      True if the given repo name is an App Engine Python Project.
    """
    name = name.lower()
    if not name.startswith('appengine-'):
      return False
    words = name.split('-')
    if 'python' not in words:
      return False
    if 'java' in words or 'go' in words:
      return False
    return True

  def _GetAppEnginePythonRepos(self, page):
    """Get list of App Engine Python repos.

    Given a JSON list of repos, return those repo names which appear to be
    Python App Engine repos, and which are not in _PROJECT_URL_SKIP_LIST.
    This function can parse the contents of:
    https://api.github.com/users/GoogleCloudPlatform/repos

    Args:
      page: the JSON response returned by /users/:user/repos

    Returns:
      A list of repos.
    """
    r = json.loads(page)
    # keys we care about: html_url, contents_url, name, description
    repos = [entry for entry in r
             if self._IsAppEnginePythonRepo(entry['name'])
                and entry['html_url'] not in _PROJECT_URL_SKIP_LIST]

    candidates1 = []
    for repo in repos:
      # only proceed with repos which look like App Engine Python projects
      if not self._IsAppEnginePythonRepo(repo['name']):
        continue

      # e.g. https://api.github.com/repos/GoogleCloudPlatform/appengine-crowdguru-python/contents/app.yaml
      app_yaml_contents_url = repo['contents_url'].replace('{+path}',
                                                           'app.yaml')
      rpc = FetchWithAuth(app_yaml_contents_url, async=True)
      candidates1.append((rpc, app_yaml_contents_url, repo))

    # filter repos with no app.yaml
    candidates2 = []
    for rpc, app_yaml_contents_url, repo in candidates1:
      result = rpc.get_result()
      if result.status_code != httplib.OK:
        shared.w('Skipping repo due to {} fetching {}'
                 .format(result.status_code, app_yaml_contents_url))
        continue
      r = json.loads(result.content)
      app_yaml_git_url = r['git_url']
      rpc = FetchWithAuth(app_yaml_git_url, async=True)
      candidates2.append((rpc, app_yaml_contents_url, app_yaml_git_url, repo))

    # filter repos whose app.yaml does not contain 'runtime: python27'
    repos = []
    for rpc, app_yaml_contents_url, app_yaml_git_url, repo in candidates2:
      result = rpc.get_result()
      r = json.loads(result.content)
      base64_content = r['content']
      decoded_content = base64.b64decode(base64_content)
      config = yaml.safe_load(decoded_content)
      runtime = config.get('runtime')
      if runtime != 'python27':
        shared.w('Skipping "runtime: {}" app {}'.format(runtime,
                                                        app_yaml_contents_url))
        continue
      repos.append(repo)

    return repos

  def _GetRepoContents(self, repo_contents_url):
    """Get list of files/directories in the given repo.

    Args:
      page: the JSON response returned by /repos/:owner/:repo/contents

    Returns:
      The list of files/directories.
    """
    page = FetchWithAuth(repo_contents_url, follow_redirects=True).content
    r = json.loads(page)
    entries = []
    for entry in r:
      if entry['type'] == 'dir':
        entries.extend(self._GetRepoContents(entry['url']))
      elif entry['type'] == 'file':
        entries.append((entry['path'], entry['git_url']))
    return entries

  def PopulateRepos(self):
    shared.EnsureRunningInTask()  # gives us automatic retries
    repo_collection_url = self.repo_collection.key.id()
    matcher = _GITHUB_URL_RE.match(repo_collection_url)
    github_user = matcher.group(1)
    # e.g. https://api.github.com/users/GoogleCloudPlatform/repos
    url = 'https://api.github.com/users/{0}/repos'.format(github_user)
    rpc_result = FetchWithAuth(url, follow_redirects=True)
    page = rpc_result.content
    repos = self._GetAppEnginePythonRepos(page)

    credential = model.GetOAuth2Credential('github')
    if (not credential or not credential.client_id
        or not credential.client_secret):
      # fetch fewer when we're not authenticated
      repos = repos[:1]

    repo_entities = []
    for repo in repos:
      r = model.CreateRepo(repo['html_url'], end_user_url=repo['html_url'],
                           name=repo['name'],
                           description=repo['description'] or repo['html_url'])
      repo_entities.append(r)
    model.ndb.put_multi(repo_entities)
    for repo_entity in repo_entities:
      deferred.defer(self.CreateTemplateProject, repo_entity.key)

  def CreateProjectTreeFromRepo(self, tree, repo):
    # e.g. https://github.com/GoogleCloudPlatform/appengine-crowdguru-python
    end_user_repo_url = repo.key.id()
    matcher = _GITHUB_URL_RE.match(end_user_repo_url)
    github_user = matcher.group(1)
    repo_name = matcher.group(2)
    # e.g. https://api.github.com/repos/GoogleCloudPlatform/appengine-crowdguru-python/contents/
    repo_contents_url = ('https://api.github.com/repos/{0}/{1}/contents/'
                         .format(github_user, repo_name))

    entries = self._GetRepoContents(repo_contents_url)

    rpcs = []
    for (path, file_git_url) in entries:
      rpc = FetchWithAuth(file_git_url, follow_redirects=True, async=True)
      rpcs.append((file_git_url, path, rpc))

    for file_git_url, path, rpc in rpcs:
      try:
        result = rpc.get_result()
        shared.w('{0} {1} {2}'.format(result.status_code, path, file_git_url))
        if result.status_code != 200:
          continue
        r = json.loads(result.content)
        base64_content = r['content']
        decoded_content = base64.b64decode(base64_content)
        tree.SetFile(path, decoded_content)
      except urlfetch_errors.Error:
        exc_info = sys.exc_info()
        formatted_exception = traceback.format_exception(exc_info[0],
                                                         exc_info[1],
                                                         exc_info[2])
        shared.w('Skipping {0} {1}'.format(path, file_git_url))
        for line in [line for line in formatted_exception if line]:
          shared.w(line)
<|endoftext|>


<commit_before># coding: utf8
from __future__ import unicode_literals

import pytest
import os
import tempfile
from pathlib import Path
from contextlib import contextmanager

from ..util import read_json, write_json, read_jsonl, write_jsonl
from ..util import is_json_serializable

@contextmanager
def make_tempfile(data=None):
    tmp_file = tempfile.NamedTemporaryFile(delete=False)
    if data is not None:
        with Path(tmp_file.name).open("w", encoding="utf8") as f:
            f.write(data)
    yield tmp_file.name
    os.remove(tmp_file.name)

def test_read_json_file():
    file_contents = '{\n    "hello": "world"\n}'
    with make_tempfile(file_contents) as file_path:
        data = read_json(file_path)
    assert len(data) == 1
    assert data["hello"] == "world"

def test_read_json_file_invalid():
    file_contents = '{\n    "hello": world\n}'
    with make_tempfile(file_contents) as file_path:
        with pytest.raises(ValueError):
            data = read_json(file_path)

def test_write_json_file():
    data = {"hello": "word", "test": 123}
    with make_tempfile() as file_path:
        write_json(file_path, data)
        with Path(file_path).open("r", encoding="utf8") as f:
            assert f.read() == '{\n  "hello":"word",\n  "test":123\n}'

def test_read_jsonl_file():
    file_contents = '{"hello": "world"}\n{"test": 123}'
    with make_tempfile(file_contents) as f:
        data = read_jsonl(f)
        # Make sure this returns a generator, not just a list
        assert not hasattr(data, "__len__")
        data = list(data)
    assert len(data) == 2
    assert len(data[0]) == 1
    assert len(data[1]) == 1
    assert data[0]["hello"] == "world"
    assert data[1]["test"] == 123

def test_read_jsonl_file_invalid():
    file_contents = '{"hello": world}\n{"test": 123}'
    with make_tempfile(file_contents) as f:
        with pytest.raises(ValueError):
            data = list(read_jsonl(f))
        data = list(read_jsonl(f, skip=True))
    assert len(data) == 1
    assert len(data[0]) == 1
    assert data[0]["test"] == 123

def test_write_json_file():
    data = [{"hello": "world"}, {"test": 123}]
    with make_tempfile() as file_path:
        write_jsonl(file_path, data)
        with Path(file_path).open("r", encoding="utf8") as f:
            assert f.read() == '{"hello":"world"}\n{"test":123}\n'

@pytest.mark.parametrize(
    "obj,expected",
    [
        (["a", "b", 1, 2], True),
        ({"a": "b", "c": 123}, True),
        ("hello", True),
        (lambda x: x, False),
    ],
)
def test_is_json_serializable(obj, expected):
    assert is_json_serializable(obj) == expected
<commit_msg>Add stdout tests<commit_after># coding: utf8
from __future__ import unicode_literals

import pytest
import os
import tempfile
from pathlib import Path
from contextlib import contextmanager

from ..util import read_json, write_json, read_jsonl, write_jsonl
from ..util import is_json_serializable

@contextmanager
def make_tempfile(data=None):
    tmp_file = tempfile.NamedTemporaryFile(delete=False)
    if data is not None:
        with Path(tmp_file.name).open("w", encoding="utf8") as f:
            f.write(data)
    yield tmp_file.name
    os.remove(tmp_file.name)

def test_read_json_file():
    file_contents = '{\n    "hello": "world"\n}'
    with make_tempfile(file_contents) as file_path:
        data = read_json(file_path)
    assert len(data) == 1
    assert data["hello"] == "world"

def test_read_json_file_invalid():
    file_contents = '{\n    "hello": world\n}'
    with make_tempfile(file_contents) as file_path:
        with pytest.raises(ValueError):
            data = read_json(file_path)

def test_write_json_file():
    data = {"hello": "word", "test": 123}
    with make_tempfile() as file_path:
        write_json(file_path, data)
        with Path(file_path).open("r", encoding="utf8") as f:
            assert f.read() == '{\n  "hello":"word",\n  "test":123\n}'

def test_write_json_stdout(capsys):
    data = {"hello": "word", "test": 123}
    write_json("-", data)
    captured = capsys.readouterr()
    assert captured.out == '{\n  "hello": "word",\n  "test": 123\n}\n'

def test_read_jsonl_file():
    file_contents = '{"hello": "world"}\n{"test": 123}'
    with make_tempfile(file_contents) as f:
        data = read_jsonl(f)
        # Make sure this returns a generator, not just a list
        assert not hasattr(data, "__len__")
        data = list(data)
    assert len(data) == 2
    assert len(data[0]) == 1
    assert len(data[1]) == 1
    assert data[0]["hello"] == "world"
    assert data[1]["test"] == 123

def test_read_jsonl_file_invalid():
    file_contents = '{"hello": world}\n{"test": 123}'
    with make_tempfile(file_contents) as f:
        with pytest.raises(ValueError):
            data = list(read_jsonl(f))
        data = list(read_jsonl(f, skip=True))
    assert len(data) == 1
    assert len(data[0]) == 1
    assert data[0]["test"] == 123

def test_write_json_file():
    data = [{"hello": "world"}, {"test": 123}]
    with make_tempfile() as file_path:
        write_jsonl(file_path, data)
        with Path(file_path).open("r", encoding="utf8") as f:
            assert f.read() == '{"hello":"world"}\n{"test":123}\n'

def test_write_jsonl_stdout(capsys):
    data = [{"hello": "world"}, {"test": 123}]
    write_jsonl("-", data)
    captured = capsys.readouterr()
    assert captured.out == '{"hello":"world"}\n{"test":123}\n'

@pytest.mark.parametrize(
    "obj,expected",
    [
        (["a", "b", 1, 2], True),
        ({"a": "b", "c": 123}, True),
        ("hello", True),
        (lambda x: x, False),
    ],
)
def test_is_json_serializable(obj, expected):
    assert is_json_serializable(obj) == expected
<|endoftext|>


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

from __future__ import print_function

import sslyze.plugins
import sslyze.plugins.plugin_base
from sslyze.plugins.certificate_info_plugin import CertificateInfoPlugin
from sslyze.plugins.compression_plugin import CompressionPlugin
from sslyze.plugins.fallback_scsv_plugin import FallbackScsvPlugin
from sslyze.plugins.heartbleed_plugin import HeartbleedPlugin
from sslyze.plugins.http_headers_plugin import HttpHeadersPlugin
from sslyze.plugins.openssl_ccs_injection_plugin import OpenSslCcsInjectionPlugin
from sslyze.plugins.openssl_cipher_suites_plugin import OpenSslCipherSuitesPlugin

# TODO(ad): rename this
from sslyze.plugins.session_renegotiation_plugin import SessionRenegotiationPlugin

class PluginsFinder(object):
    """Utility class to discover the list of available SSLyze scanning plugins and commands.
    """

    _DISCOVERED_PLUGINS = None

    @classmethod
    def get(cls):
        # type: () -> PluginsFinder
        """Discover available SSLyze plugins and return a PluginsFinder to be passed to a PluginsProcessPool.

        Returns:
            PluginsFinder:  An object encapsulating the list of available SSLyze plugins.
        """
        if cls._DISCOVERED_PLUGINS is None:
            cls._DISCOVERED_PLUGINS = PluginsFinder()
        return cls._DISCOVERED_PLUGINS

    def __init__(self):
        # type: () -> None
        """Find available plugins by discovering any class that implements the PluginBase abstract class.

        Returns:
            PluginsFinder:  An object encapsulating the list of available SSLyze plugins.
        """
        self._plugin_classes = [OpenSslCipherSuitesPlugin, CertificateInfoPlugin, CompressionPlugin, FallbackScsvPlugin,
                                HeartbleedPlugin, HttpHeadersPlugin, OpenSslCcsInjectionPlugin,
                                SessionRenegotiationPlugin]
        self._commands = []
        for plugin_class in self._plugin_classes:
            self._commands.extend(plugin_class.get_available_commands())

    def get_plugins(self):
        return self._plugin_classes

    def get_commands(self):
        return self._commands

    @staticmethod
    def get_plugin_modules_static():

        import sslyze.plugins.certificate_info_plugin
        import sslyze.plugins.compression_plugin
        import sslyze.plugins.fallback_scsv_plugin
        import sslyze.plugins.heartbleed_plugin
        import sslyze.plugins.http_headers_plugin
        import sslyze.plugins.openssl_ccs_injection_plugin
        import sslyze.plugins.openssl_cipher_suites_plugin
        import sslyze.plugins.session_renegotiation_plugin
        import sslyze.plugins.session_resumption_plugin

        plugin_modules = [
            sslyze.plugins.certificate_info_plugin,
            sslyze.plugins.compression_plugin,
            sslyze.plugins.fallback_scsv_plugin,
            sslyze.plugins.heartbleed_plugin,
            sslyze.plugins.http_headers_plugin,
            sslyze.plugins.openssl_ccs_injection_plugin,
            sslyze.plugins.openssl_cipher_suites_plugin,
            sslyze.plugins.session_renegotiation_plugin,
            sslyze.plugins.session_resumption_plugin
        ]

        return plugin_modules

    @staticmethod
    def get_plugin_modules_dynamic():
        plugin_modules = []
        import pkgutil
        for loader, module_name, _ in list(pkgutil.iter_modules(sslyze.plugins.__path__, prefix='sslyze.plugins.')):
            try:
                module = loader.find_module(module_name).load_module(module_name)
                plugin_modules.append(module)
            except Exception as e:
                print('  {module} - Import Error: {error}'.format(module=module_name, error=str(e)))
                continue

        return plugin_modules
<commit_msg>Tweak<commit_after># -*- coding: utf-8 -*-

from __future__ import print_function

import sslyze.plugins
import sslyze.plugins.plugin_base
from sslyze.plugins.certificate_info_plugin import CertificateInfoPlugin
from sslyze.plugins.compression_plugin import CompressionPlugin
from sslyze.plugins.fallback_scsv_plugin import FallbackScsvPlugin
from sslyze.plugins.heartbleed_plugin import HeartbleedPlugin
from sslyze.plugins.http_headers_plugin import HttpHeadersPlugin
from sslyze.plugins.openssl_ccs_injection_plugin import OpenSslCcsInjectionPlugin
from sslyze.plugins.openssl_cipher_suites_plugin import OpenSslCipherSuitesPlugin

# TODO(ad): rename this
from sslyze.plugins.session_renegotiation_plugin import SessionRenegotiationPlugin
from sslyze.plugins.session_resumption_plugin import SessionResumptionPlugin

class PluginsFinder(object):
    """Utility class to discover the list of available SSLyze scanning plugins and commands.
    """

    _DISCOVERED_PLUGINS = None

    @classmethod
    def get(cls):
        # type: () -> PluginsFinder
        """Discover available SSLyze plugins and return a PluginsFinder to be passed to a PluginsProcessPool.

        Returns:
            PluginsFinder:  An object encapsulating the list of available SSLyze plugins.
        """
        if cls._DISCOVERED_PLUGINS is None:
            cls._DISCOVERED_PLUGINS = PluginsFinder()
        return cls._DISCOVERED_PLUGINS

    def __init__(self):
        # type: () -> None
        """Find available plugins by discovering any class that implements the PluginBase abstract class.

        Returns:
            PluginsFinder:  An object encapsulating the list of available SSLyze plugins.
        """
        self._plugin_classes = [OpenSslCipherSuitesPlugin, CertificateInfoPlugin, CompressionPlugin, FallbackScsvPlugin,
                                HeartbleedPlugin, HttpHeadersPlugin, OpenSslCcsInjectionPlugin,
                                SessionRenegotiationPlugin, SessionResumptionPlugin]
        self._commands = []
        for plugin_class in self._plugin_classes:
            self._commands.extend(plugin_class.get_available_commands())

    def get_plugins(self):
        return self._plugin_classes

    def get_commands(self):
        return self._commands

    @staticmethod
    def get_plugin_modules_static():

        import sslyze.plugins.certificate_info_plugin
        import sslyze.plugins.compression_plugin
        import sslyze.plugins.fallback_scsv_plugin
        import sslyze.plugins.heartbleed_plugin
        import sslyze.plugins.http_headers_plugin
        import sslyze.plugins.openssl_ccs_injection_plugin
        import sslyze.plugins.openssl_cipher_suites_plugin
        import sslyze.plugins.session_renegotiation_plugin
        import sslyze.plugins.session_resumption_plugin

        plugin_modules = [
            sslyze.plugins.certificate_info_plugin,
            sslyze.plugins.compression_plugin,
            sslyze.plugins.fallback_scsv_plugin,
            sslyze.plugins.heartbleed_plugin,
            sslyze.plugins.http_headers_plugin,
            sslyze.plugins.openssl_ccs_injection_plugin,
            sslyze.plugins.openssl_cipher_suites_plugin,
            sslyze.plugins.session_renegotiation_plugin,
            sslyze.plugins.session_resumption_plugin
        ]

        return plugin_modules

    @staticmethod
    def get_plugin_modules_dynamic():
        plugin_modules = []
        import pkgutil
        for loader, module_name, _ in list(pkgutil.iter_modules(sslyze.plugins.__path__, prefix='sslyze.plugins.')):
            try:
                module = loader.find_module(module_name).load_module(module_name)
                plugin_modules.append(module)
            except Exception as e:
                print('  {module} - Import Error: {error}'.format(module=module_name, error=str(e)))
                continue

        return plugin_modules
<|endoftext|>


<commit_before>import random
import time
import graphs
import attractors
from shutil import copyfile
from collections import namedtuple
import csv
import numpy
from stochastic import estimate_probability_to_reroll_attractor
import os
import sys
import stat

if __name__ == "__main__":
    # TODO: use grownups' argument parsing library.
    if len(sys.argv) == 1:
        output_path = 'bitchange_results.csv'
        analyze_originals = True
    else:
        output_path = sys.argv[1]
        analyze_originals = bool(int(sys.argv[2]))

    experiment_max_len = 15
    experiment_n_iter = 100
    biological_test_frequency = 10

    # copy Dubrova's executable, to prevent mysterious errors when running multiple processes calling it.
    process_specific_dubrova_path = "temp_{}_{}".format(attractors.dubrova_path, os.getpid())
    copyfile(attractors.dubrova_path, process_specific_dubrova_path)
    os.chmod(process_specific_dubrova_path, stat.S_IRUSR | stat.S_IWUSR | stat.S_IXUSR)
    attractors.dubrova_path = process_specific_dubrova_path

    StabilityResult = namedtuple("StabilityResult", "graph_name random_functions random_edges size "
                                                    "n_inputs normalized_n_inputs max_degree mean_degree "
                                                    "state_bits_changed "
                                                    # "minimal_model_bitchange "
                                                    "model_bitchange_prob "
                                                    "state_bitchange_prob "
                                                    "reroll_prob"
                                )
    results = []

    biological_graphs = []
    biological_graph_names = os.listdir("cellcollective_models")
    for graph_dir in biological_graph_names:
        try:
            G = graphs.Network.parse_boolean_tables(os.path.join("cellcollective_models", graph_dir))
            biological_graphs.append(G)
        except ValueError as e:
            if e.message.startswith("Model export from cellcollective failed"):
                print "warning - did not load graph {}".format(graph_dir)

    graph_name_to_attributes = dict()
    for i, graph, name in zip(range(len(biological_graphs)), biological_graphs, biological_graph_names):
        n_inputs = len([v for v in graph.vertices if len(v.predecessors()) == 0])
        max_degree = max([len(v.predecessors()) for v in graph.vertices])
        size = len(graph.vertices)
        mean_degree = sum([len(v.predecessors()) for v in graph.vertices]) / float(size)
        normaliezd_n_inputs = n_inputs / float(size)
        graph_name_to_attributes[name] = dict(n_inputs=n_inputs, max_degree=max_degree,
                                                    size=size, mean_degree=mean_degree,
                                              normalized_n_inputs=normaliezd_n_inputs)
        print "#{}; {} input nodes for graph {} of size {} and max degree {}".format(i, n_inputs, name,
                                                                                size, max_degree)
    for test in range(5000):
        print "test #{}".format(test)
        for graph, name in zip(biological_graphs, biological_graph_names):
            is_biological = (test % biological_test_frequency) == 0
            graph_copy = graph.copy()

            if not is_biological:
                randomize_functions, randomize_edges = random.choice([(False, True), (True, False), (True, True)])
                if randomize_functions:
                    graph_copy.randomize_functions(preserve_truth_ratio=True)
                if randomize_edges:
                    graph_copy.randomize_edges()
            start = time.time()

            # try:
            #     minimal_model_bitchange = \
            #         attractors.find_model_bitchange_for_new_attractor(graph_copy, max_len=experiment_max_len,
            #                                                           verbose=False, use_dubrova=False,
            #                                                           simplify_boolean=False)
            # except attractors.TimeoutError as e:
            #     minimal_model_bitchange = numpy.inf
            print "time taken for model_bitchange={:.2f} secs".format(time.time() - start)
            second_start = time.time()
            state_bits_changed = random.randint(1, 5)
            try:
                model_bitchange_probability = \
                    attractors.find_model_bitchange_probability_for_different_attractors(graph_copy,
                                                                                max_len=experiment_max_len,
                                                                                n_iter=experiment_n_iter,
                                                                                use_dubrova=True)
            except attractors.TimeoutError as e:
                model_bitchange_probability = numpy.inf
            print "time taken for model_bitchange_probability={:.2f} secs".format(time.time() - second_start)
            third_start = time.time()
            try:
                state_bitchange_probability = \
                    attractors.find_state_bitchange_probability_for_different_attractors(graph_copy,
                                                                                n_iter=experiment_n_iter,
                                                                                n_bits=state_bits_changed,
                                                                                parallel=False)
            except attractors.TimeoutError as e:
                state_bitchange_probability = numpy.inf
            print "time taken for state_bitchange_probability={:.2f} secs".format(time.time() - third_start)

            fourth_start = time.time()
            reroll_probability = estimate_probability_to_reroll_attractor(graph_copy, n_walks=experiment_n_iter)
            print "time taken for reroll probability={:.2f} secs".format(time.time() - fourth_start)

            result = StabilityResult(graph_name=name,
                                     random_functions=randomize_functions if not is_biological else False,
                                     random_edges=randomize_edges if not is_biological else False,
                                     size=graph_name_to_attributes[name]['size'],
                                     n_inputs=graph_name_to_attributes[name]['n_inputs'],
                                     normalized_n_inputs=graph_name_to_attributes[name]['normalized_n_inputs'],
                                     max_degree=graph_name_to_attributes[name]['max_degree'],
                                     mean_degree=graph_name_to_attributes[name]['mean_degree'],
                                     # minimal_model_bitchange=minimal_model_bitchange,
                                     model_bitchange_prob=model_bitchange_probability,
                                     state_bits_changed=state_bits_changed,
                                     state_bitchange_prob=state_bitchange_probability,
                                     reroll_prob=reroll_probability
                                     )

            # print "model bitchange={}".format(minimal_model_bitchange)
            results.append(result)
            print "model_prob={:.2f}, state_prob={:.2f}".format(
                # minimal_model_bitchange,
                model_bitchange_probability,
                state_bitchange_probability
            )
            print "time_taken = {:.2f} secs".format(time.time() - start)

        # save on each iteration, why not
        with open(output_path, 'wb') as csv_file:
            writer = csv.writer(csv_file)
            writer.writerow(["graph_name", "random_functions", "random_edges", "n",
                             "n_inputs", "normalized_n_inputs", "max_degree", "mean_degree",
                             "state_bits_changed",
                             "model_bitchange_prob",
                             "state_bitchange_prob",
                             "reroll_prob"
                             ])
            for stability_result in results:
                writer.writerow(stability_result)
<commit_msg>Bugfix with stabiility_estimations argument parsing<commit_after>import random
import time
import graphs
import attractors
from shutil import copyfile
from collections import namedtuple
import csv
import numpy
from stochastic import estimate_probability_to_reroll_attractor
import os
import sys
import stat

if __name__ == "__main__":
    # TODO: use grownups' argument parsing library.
    if len(sys.argv) == 1:
        output_path = 'bitchange_results.csv'
    else:
        output_path = sys.argv[1]

    experiment_max_len = 15
    experiment_n_iter = 100
    biological_test_frequency = 10

    # copy Dubrova's executable, to prevent mysterious errors when running multiple processes calling it.
    process_specific_dubrova_path = "temp_{}_{}".format(attractors.dubrova_path, os.getpid())
    copyfile(attractors.dubrova_path, process_specific_dubrova_path)
    os.chmod(process_specific_dubrova_path, stat.S_IRUSR | stat.S_IWUSR | stat.S_IXUSR)
    attractors.dubrova_path = process_specific_dubrova_path

    StabilityResult = namedtuple("StabilityResult", "graph_name random_functions random_edges size "
                                                    "n_inputs normalized_n_inputs max_degree mean_degree "
                                                    "state_bits_changed "
                                                    # "minimal_model_bitchange "
                                                    "model_bitchange_prob "
                                                    "state_bitchange_prob "
                                                    "reroll_prob"
                                )
    results = []

    biological_graphs = []
    biological_graph_names = os.listdir("cellcollective_models")
    for graph_dir in biological_graph_names:
        try:
            G = graphs.Network.parse_boolean_tables(os.path.join("cellcollective_models", graph_dir))
            biological_graphs.append(G)
        except ValueError as e:
            if e.message.startswith("Model export from cellcollective failed"):
                print "warning - did not load graph {}".format(graph_dir)

    graph_name_to_attributes = dict()
    for i, graph, name in zip(range(len(biological_graphs)), biological_graphs, biological_graph_names):
        n_inputs = len([v for v in graph.vertices if len(v.predecessors()) == 0])
        max_degree = max([len(v.predecessors()) for v in graph.vertices])
        size = len(graph.vertices)
        mean_degree = sum([len(v.predecessors()) for v in graph.vertices]) / float(size)
        normaliezd_n_inputs = n_inputs / float(size)
        graph_name_to_attributes[name] = dict(n_inputs=n_inputs, max_degree=max_degree,
                                                    size=size, mean_degree=mean_degree,
                                              normalized_n_inputs=normaliezd_n_inputs)
        print "#{}; {} input nodes for graph {} of size {} and max degree {}".format(i, n_inputs, name,
                                                                                size, max_degree)
    for test in range(5000):
        print "test #{}".format(test)
        for graph, name in zip(biological_graphs, biological_graph_names):
            is_biological = (test % biological_test_frequency) == 0
            graph_copy = graph.copy()

            if not is_biological:
                randomize_functions, randomize_edges = random.choice([(False, True), (True, False), (True, True)])
                if randomize_functions:
                    graph_copy.randomize_functions(preserve_truth_ratio=True)
                if randomize_edges:
                    graph_copy.randomize_edges()
            start = time.time()

            # try:
            #     minimal_model_bitchange = \
            #         attractors.find_model_bitchange_for_new_attractor(graph_copy, max_len=experiment_max_len,
            #                                                           verbose=False, use_dubrova=False,
            #                                                           simplify_boolean=False)
            # except attractors.TimeoutError as e:
            #     minimal_model_bitchange = numpy.inf
            print "time taken for model_bitchange={:.2f} secs".format(time.time() - start)
            second_start = time.time()
            state_bits_changed = random.randint(1, 5)
            try:
                model_bitchange_probability = \
                    attractors.find_model_bitchange_probability_for_different_attractors(graph_copy,
                                                                                max_len=experiment_max_len,
                                                                                n_iter=experiment_n_iter,
                                                                                use_dubrova=True)
            except attractors.TimeoutError as e:
                model_bitchange_probability = numpy.inf
            print "time taken for model_bitchange_probability={:.2f} secs".format(time.time() - second_start)
            third_start = time.time()
            try:
                state_bitchange_probability = \
                    attractors.find_state_bitchange_probability_for_different_attractors(graph_copy,
                                                                                n_iter=experiment_n_iter,
                                                                                n_bits=state_bits_changed,
                                                                                parallel=False)
            except attractors.TimeoutError as e:
                state_bitchange_probability = numpy.inf
            print "time taken for state_bitchange_probability={:.2f} secs".format(time.time() - third_start)

            fourth_start = time.time()
            reroll_probability = estimate_probability_to_reroll_attractor(graph_copy, n_walks=experiment_n_iter)
            print "time taken for reroll probability={:.2f} secs".format(time.time() - fourth_start)

            result = StabilityResult(graph_name=name,
                                     random_functions=randomize_functions if not is_biological else False,
                                     random_edges=randomize_edges if not is_biological else False,
                                     size=graph_name_to_attributes[name]['size'],
                                     n_inputs=graph_name_to_attributes[name]['n_inputs'],
                                     normalized_n_inputs=graph_name_to_attributes[name]['normalized_n_inputs'],
                                     max_degree=graph_name_to_attributes[name]['max_degree'],
                                     mean_degree=graph_name_to_attributes[name]['mean_degree'],
                                     # minimal_model_bitchange=minimal_model_bitchange,
                                     model_bitchange_prob=model_bitchange_probability,
                                     state_bits_changed=state_bits_changed,
                                     state_bitchange_prob=state_bitchange_probability,
                                     reroll_prob=reroll_probability
                                     )

            # print "model bitchange={}".format(minimal_model_bitchange)
            results.append(result)
            print "model_prob={:.2f}, state_prob={:.2f}".format(
                # minimal_model_bitchange,
                model_bitchange_probability,
                state_bitchange_probability
            )
            print "time_taken = {:.2f} secs".format(time.time() - start)

        # save on each iteration, why not
        with open(output_path, 'wb') as csv_file:
            writer = csv.writer(csv_file)
            writer.writerow(["graph_name", "random_functions", "random_edges", "n",
                             "n_inputs", "normalized_n_inputs", "max_degree", "mean_degree",
                             "state_bits_changed",
                             "model_bitchange_prob",
                             "state_bitchange_prob",
                             "reroll_prob"
                             ])
            for stability_result in results:
                writer.writerow(stability_result)
<|endoftext|>


<commit_before>import os
from itertools import product

from troposphere import Ref
from troposphere.ec2 import SecurityGroup, SecurityGroupRule

from .common import administrator_ip_address
from .template import template
from .vpc import loadbalancer_a_subnet_cidr, loadbalancer_b_subnet_cidr, vpc

load_balancer_security_group = SecurityGroup(
    "LoadBalancerSecurityGroup",
    template=template,
    GroupDescription="Web load balancer security group.",
    VpcId=Ref(vpc),
    SecurityGroupIngress=[
        # allow incoming traffic from the public internet to the load balancer
        # on ports 80 and 443
        SecurityGroupRule(
            IpProtocol="tcp",
            FromPort=port,
            ToPort=port,
            CidrIp="0.0.0.0/0",
        ) for port in ["80", "443"]
    ],
)

# allow traffic from the load balancer subnets to the web workers
if os.environ.get('USE_ECS') == 'on' or os.environ.get('USE_EC2') == 'on':
    # if using ECS or EC2, allow traffic to the configured WebWorkerPort
    web_worker_ports = [Ref("WebWorkerPort")]
elif os.environ.get('USE_GOVCLOUD') == 'on':
    # if using GovCloud (real EC2 instances), allow traffic to the configured
    # WebWorkerPort and port 443
    web_worker_ports = [Ref("WebWorkerPort"), "443"]
else:
    # otherwise, if using Elastic Beanstalk, allow traffic only to EB's default
    # web worker port (80)
    web_worker_ports = ["80"]

cidrs = [loadbalancer_a_subnet_cidr, loadbalancer_b_subnet_cidr]

# HTTP from web public subnets
ingress_rules = [SecurityGroupRule(
    IpProtocol="tcp",
    FromPort=port,
    ToPort=port,
    CidrIp=cidr,
) for port, cidr in product(*[web_worker_ports, cidrs])]

# Health check
ingress_rules.append(SecurityGroupRule(
    IpProtocol="tcp",
    FromPort=Ref("WebWorkerHealthCheckPort"),
    ToPort=Ref("WebWorkerHealthCheckPort"),
    Description="ELB Health Check",  # SecurityGroupRule doesn't support a Description attribute
    SourceSecurityGroupId=Ref(load_balancer_security_group),
))

# AdministratorAccess
ingress_rules.append(SecurityGroupRule(
    IpProtocol="tcp",
    FromPort="22",
    ToPort="22",
    Description="Administrator SSH Access",
    CidrIp=administrator_ip_address,
))

container_security_group = SecurityGroup(
    'ContainerSecurityGroup',
    template=template,
    GroupDescription="Container security group.",
    VpcId=Ref(vpc),
    SecurityGroupIngress=ingress_rules,
)
<commit_msg>Fix template validation error for Elastic Beanstalk.<commit_after>import os
from itertools import product

from troposphere import Ref
from troposphere.ec2 import SecurityGroup, SecurityGroupRule

from .common import administrator_ip_address
from .template import template
from .vpc import loadbalancer_a_subnet_cidr, loadbalancer_b_subnet_cidr, vpc

load_balancer_security_group = SecurityGroup(
    "LoadBalancerSecurityGroup",
    template=template,
    GroupDescription="Web load balancer security group.",
    VpcId=Ref(vpc),
    SecurityGroupIngress=[
        # allow incoming traffic from the public internet to the load balancer
        # on ports 80 and 443
        SecurityGroupRule(
            IpProtocol="tcp",
            FromPort=port,
            ToPort=port,
            CidrIp="0.0.0.0/0",
        ) for port in ["80", "443"]
    ],
)

# allow traffic from the load balancer subnets to the web workers
if os.environ.get('USE_ECS') == 'on' or os.environ.get('USE_EC2') == 'on':
    # if using ECS or EC2, allow traffic to the configured WebWorkerPort
    web_worker_ports = [Ref("WebWorkerPort")]
elif os.environ.get('USE_GOVCLOUD') == 'on':
    # if using GovCloud (real EC2 instances), allow traffic to the configured
    # WebWorkerPort and port 443
    web_worker_ports = [Ref("WebWorkerPort"), "443"]
else:
    # otherwise, if using Elastic Beanstalk, allow traffic only to EB's default
    # web worker port (80)
    web_worker_ports = ["80"]

cidrs = [loadbalancer_a_subnet_cidr, loadbalancer_b_subnet_cidr]

# HTTP from web public subnets
ingress_rules = [SecurityGroupRule(
    IpProtocol="tcp",
    FromPort=port,
    ToPort=port,
    CidrIp=cidr,
) for port, cidr in product(*[web_worker_ports, cidrs])]

# Health check
if os.environ.get('USE_EB') != 'on':
    ingress_rules.append(SecurityGroupRule(
        IpProtocol="tcp",
        FromPort=Ref("WebWorkerHealthCheckPort"),
        ToPort=Ref("WebWorkerHealthCheckPort"),
        Description="ELB Health Check",  # SecurityGroupRule doesn't support a Description attribute
        SourceSecurityGroupId=Ref(load_balancer_security_group),
    ))

# AdministratorAccess
ingress_rules.append(SecurityGroupRule(
    IpProtocol="tcp",
    FromPort="22",
    ToPort="22",
    Description="Administrator SSH Access",
    CidrIp=administrator_ip_address,
))

container_security_group = SecurityGroup(
    'ContainerSecurityGroup',
    template=template,
    GroupDescription="Container security group.",
    VpcId=Ref(vpc),
    SecurityGroupIngress=ingress_rules,
)
<|endoftext|>


<commit_before>
# tests are fairly 'live' (but safe to run)
# setup authorized_keys for logged in user such
# that the user can log in as themselves before running tests

import unittest
import getpass
import ansible.runner
import os
import shutil
import time
import tempfile

from nose.plugins.skip import SkipTest

def get_binary(name):
    for directory in os.environ["PATH"].split(os.pathsep):
        path = os.path.join(directory, name)
        if os.path.isfile(path) and os.access(path, os.X_OK):
            return path
    return None

class TestRunner(unittest.TestCase):

    def setUp(self):
        self.user = getpass.getuser()
        self.runner = ansible.runner.Runner(
            module_name='ping',
            module_path='library/',
            module_args='',
            remote_user=self.user,
            remote_pass=None,
            host_list='test/ansible_hosts',
            timeout=5,
            forks=1,
            background=0,
            pattern='all',
        )
        self.cwd = os.getcwd()
        self.test_dir = os.path.join(self.cwd, 'test')
        self.stage_dir = self._prepare_stage_dir()

    def _prepare_stage_dir(self):
        stage_path = os.path.join(self.test_dir, 'test_data')
        if os.path.exists(stage_path):
            shutil.rmtree(stage_path, ignore_errors=False)
            assert not os.path.exists(stage_path)
        os.makedirs(stage_path)
        assert os.path.exists(stage_path)
        return stage_path

    def _get_test_file(self, filename):
        # get a file inside the test input directory
        filename = os.path.join(self.test_dir, filename)
        assert os.path.exists(filename)
        return filename

    def _get_stage_file(self, filename):
        # get a file inside the test output directory
        filename = os.path.join(self.stage_dir, filename)
        return filename

    def _run(self, module_name, module_args, background=0):
        ''' run a module and get the localhost results '''
        self.runner.module_name = module_name
        args = ' '.join(module_args)
        print "DEBUG: using args=%s" % args
        self.runner.module_args = args
        self.runner.background = background
        results = self.runner.run()
        # when using nosetests this will only show up on failure
        # which is pretty useful
        print "RESULTS=%s" % results
        assert "127.0.0.2" in results['contacted']
        return results['contacted']['127.0.0.2']

    def test_ping(self):
        result = self._run('ping', [])
        assert "ping" in result

    def test_facter(self):
        if not get_binary("facter"):
            raise SkipTest
        result = self._run('facter', [])
        assert "hostname" in result

    # temporarily disbabled since it occasionally hangs
    # ohai's fault, setup module doesn't actually run this
    # to get ohai's "facts" anyway
    #
    #def test_ohai(self):
    #    if not get_binary("facter"):
    #            raise SkipTest
    #    result = self._run('ohai',[])
    #    assert "hostname" in result

    def test_copy(self):
        # test copy module, change trigger, etc
        input_ = self._get_test_file('sample.j2')
        output = self._get_stage_file('sample.out')
        assert not os.path.exists(output)
        result = self._run('copy', [
            "src=%s" % input_,
            "dest=%s" % output,
        ])
        assert os.path.exists(output)
        data_in = file(input_).read()
        data_out = file(output).read()
        assert data_in == data_out
        assert 'failed' not in result
        assert result['changed'] == True
        assert 'md5sum' in result
        result = self._run('copy', [
            "src=%s" % input_,
            "dest=%s" % output,
        ])
        assert result['changed'] == False

    def test_command(self):
        # test command module, change trigger, etc
        result = self._run('command', [ "/bin/echo", "hi" ])
        assert "failed" not in result
        assert "msg" not in result
        assert result['rc'] == 0
        assert result['stdout'] == 'hi'
        assert result['stderr'] == ''

        result = self._run('command', [ "false" ])
        assert result['rc'] == 1
        assert 'failed' not in result

        result = self._run('command', [ "/usr/bin/this_does_not_exist", "splat" ])
        assert 'msg' in result
        assert 'failed' in result
        assert 'rc' not in result

        result = self._run('shell', [ "/bin/echo", "$HOME" ])
        assert 'failed' not in result
        assert result['rc'] == 0

    def test_large_output(self):
        large_path = "/usr/share/dict/words"
        if not os.path.exists(large_path):
            large_path = "/usr/share/dict/cracklib-small"
            if not os.path.exists(large_path):
                raise SkipTest
        # Ensure reading a large amount of output from a command doesn't hang.
        result = self._run('command', [ "/bin/cat", large_path ])
        assert "failed" not in result
        assert "msg" not in result
        assert result['rc'] == 0
        assert len(result['stdout']) > 100000
        assert result['stderr'] == ''

    def test_async(self):
        # test async launch and job status
        # of any particular module
        result = self._run('command', [ get_binary("sleep"), "3" ], background=20)
        assert 'ansible_job_id' in result
        assert 'started' in result
        jid = result['ansible_job_id']
        # no real chance of this op taking a while, but whatever
        time.sleep(5)
        # CLI will abstract this (when polling), but this is how it works internally
        result = self._run('async_status', [ "jid=%s" % jid ])
        # TODO: would be nice to have tests for supervisory process
        # killing job after X seconds
        assert 'finished' in result
        assert 'failed' not in result
        assert 'rc' in result
        assert 'stdout' in result
        assert result['ansible_job_id'] == jid

    def test_fetch(self):
        input_ = self._get_test_file('sample.j2')
        output = os.path.join(self.stage_dir, '127.0.0.2', input_)
        result = self._run('fetch', [ "src=%s" % input_, "dest=%s" % self.stage_dir ])
        assert os.path.exists(output)
        assert open(input_).read() == open(output).read()

    def test_assemble(self):
        input = self._get_test_file('assemble.d')
        output = self._get_stage_file('sample.out')
        result = self._run('assemble', [
            "src=%s" % input,
            "dest=%s" % output,
        ])
        assert os.path.exists(output)
        out = file(output).read()
        assert out.find("first") != -1
        assert out.find("second") != -1
        assert out.find("third") != -1
        assert result['changed'] == True
        assert 'md5sum' in result
        assert 'failed' not in result
        result = self._run('assemble', [
            "src=%s" % input,
            "dest=%s" % output,
        ])
        assert result['changed'] == False

<commit_msg>full test case for subversion, minimal tests for git<commit_after>
# tests are fairly 'live' (but safe to run)
# setup authorized_keys for logged in user such
# that the user can log in as themselves before running tests

import unittest
import getpass
import ansible.runner
import os
import shutil
import time
import tempfile

from nose.plugins.skip import SkipTest

def get_binary(name):
    for directory in os.environ["PATH"].split(os.pathsep):
        path = os.path.join(directory, name)
        if os.path.isfile(path) and os.access(path, os.X_OK):
            return path
    return None

class TestRunner(unittest.TestCase):

    def setUp(self):
        self.user = getpass.getuser()
        self.runner = ansible.runner.Runner(
            module_name='ping',
            module_path='library/',
            module_args='',
            remote_user=self.user,
            remote_pass=None,
            host_list='test/ansible_hosts',
            timeout=5,
            forks=1,
            background=0,
            pattern='all',
        )
        self.cwd = os.getcwd()
        self.test_dir = os.path.join(self.cwd, 'test')
        self.stage_dir = self._prepare_stage_dir()

    def _prepare_stage_dir(self):
        stage_path = os.path.join(self.test_dir, 'test_data')
        if os.path.exists(stage_path):
            shutil.rmtree(stage_path, ignore_errors=False)
            assert not os.path.exists(stage_path)
        os.makedirs(stage_path)
        assert os.path.exists(stage_path)
        return stage_path

    def _get_test_file(self, filename):
        # get a file inside the test input directory
        filename = os.path.join(self.test_dir, filename)
        assert os.path.exists(filename)
        return filename

    def _get_stage_file(self, filename):
        # get a file inside the test output directory
        filename = os.path.join(self.stage_dir, filename)
        return filename

    def _run(self, module_name, module_args, background=0):
        ''' run a module and get the localhost results '''
        self.runner.module_name = module_name
        args = ' '.join(module_args)
        print "DEBUG: using args=%s" % args
        self.runner.module_args = args
        self.runner.background = background
        results = self.runner.run()
        # when using nosetests this will only show up on failure
        # which is pretty useful
        print "RESULTS=%s" % results
        assert "127.0.0.2" in results['contacted']
        return results['contacted']['127.0.0.2']

    def test_ping(self):
        result = self._run('ping', [])
        assert "ping" in result

    def test_facter(self):
        if not get_binary("facter"):
            raise SkipTest
        result = self._run('facter', [])
        assert "hostname" in result

    # temporarily disbabled since it occasionally hangs
    # ohai's fault, setup module doesn't actually run this
    # to get ohai's "facts" anyway
    #
    #def test_ohai(self):
    #    if not get_binary("facter"):
    #            raise SkipTest
    #    result = self._run('ohai',[])
    #    assert "hostname" in result

    def test_copy(self):
        # test copy module, change trigger, etc
        input_ = self._get_test_file('sample.j2')
        output = self._get_stage_file('sample.out')
        assert not os.path.exists(output)
        result = self._run('copy', [
            "src=%s" % input_,
            "dest=%s" % output,
        ])
        assert os.path.exists(output)
        data_in = file(input_).read()
        data_out = file(output).read()
        assert data_in == data_out
        assert 'failed' not in result
        assert result['changed'] == True
        assert 'md5sum' in result
        result = self._run('copy', [
            "src=%s" % input_,
            "dest=%s" % output,
        ])
        assert result['changed'] == False

    def test_command(self):
        # test command module, change trigger, etc
        result = self._run('command', [ "/bin/echo", "hi" ])
        assert "failed" not in result
        assert "msg" not in result
        assert result['rc'] == 0
        assert result['stdout'] == 'hi'
        assert result['stderr'] == ''

        result = self._run('command', [ "false" ])
        assert result['rc'] == 1
        assert 'failed' not in result

        result = self._run('command', [ "/usr/bin/this_does_not_exist", "splat" ])
        assert 'msg' in result
        assert 'failed' in result
        assert 'rc' not in result

        result = self._run('shell', [ "/bin/echo", "$HOME" ])
        assert 'failed' not in result
        assert result['rc'] == 0

    def test_git(self):
        self._run('file',['path=/tmp/gitdemo','state=absent'])
        self._run('file',['path=/tmp/gd','state=absent'])
        self._run('command',['git init gitdemo', 'chdir=/tmp'])
        self._run('command',['touch a', 'chdir=/tmp/gitdemo'])
        self._run('command',['git add *', 'chdir=/tmp/gitdemo'])
        self._run('command',['git commit -m "test commit 2"', 'chdir=/tmp/gitdemo'])
        self._run('command',['touch b', 'chdir=/tmp/gitdemo'])
        self._run('command',['git add *', 'chdir=/tmp/gitdemo'])
        self._run('command',['git commit -m "test commit 2"', 'chdir=/tmp/gitdemo'])
        result = self._run('git', [ "repo=\"file:///tmp/gitdemo\"", "dest=/tmp/gd" ])
        assert result['changed']
        # test the force option not set
        self._run('file',['path=/tmp/gd/a','state=absent'])
        result = self._run('git', [ "repo=\"file:///tmp/gitdemo\"", "dest=/tmp/gd", "force=no" ])
        assert result['failed']
        # test the force option when set
        result = self._run('git', [ "repo=\"file:///tmp/gitdemo\"", "dest=/tmp/gd", "force=yes" ])
        assert result['changed']

    def test_subversion(self):
        # TODO make an svn repo locally so as to avoid tests failing on network calls
        self._run('file',['path=/tmp/meetings','state=absent'])
        # hacking/test-module -m library/subversion
        result = self._run('subversion', [ ])
        assert result['failed']
        assert "dest" in result['msg']
        assert "repo" in result['msg']
        # hacking/test-module -m library/subversion -a "repo=\"http://svn.apache.org/repos/asf/subversion/trunk/notes/meetings\""
        result = self._run('subversion', [ "repo=\"http://svn.apache.org/repos/asf/subversion/trunk/notes/meetings\"" ])
        assert result['failed']
        assert "dest" in result['msg']
        # hacking/test-module -m library/subversion -a "dest=\"/tmp/gnconf\""
        result = self._run('subversion', [ "dest=\"/tmp/gnconf\"" ])
        assert result['failed']
        assert "repo" in result['msg']
        # when /tmp/meetings doesn't exist:
        # hacking/test-module -m library/subversion -a "repo=\"repo\" dest=\"/tmp/gnconf\""
        result = self._run('subversion', [ "repo=\"http://svn.apache.org/repos/asf/subversion/trunk/notes/meetings\"","dest=\"/tmp/meetings\"" ])
        assert result['changed']
        # when /tmp/meetings exists, but nothing has changed.
        result = self._run('subversion', [ "repo=\"http://svn.apache.org/repos/asf/subversion/trunk/notes/meetings\"","dest=\"/tmp/meetings\"" ])
        assert not result['changed']
        # when /tmp/meetings is a folder, but its not an svn repo
        self._run('file',['path=/tmp/meetings','state=absent'])
        self._run('file',['path=/tmp/meetings','state=directory'])
        result = self._run('subversion', [ "repo=\"http://svn.apache.org/repos/asf/subversion/trunk/notes/meetings\"","dest=\"/tmp/meetings\"" ])
        assert result['failed']
        # when /tmp/meetings is a file
        self._run('file',['path=/tmp/meetings','state=absent'])
        self._run('command',['touch /tmp/meetings'])
        result = self._run('subversion', [ "repo=\"http://svn.apache.org/repos/asf/subversion/trunk/notes/meetings\"","dest=\"/tmp/meetings\"" ])
        assert result['failed']
        # when /tmp/meetings is a folder, but its a different svn URL - should automatically switch
        self._run('file',['path=/tmp/meetings','state=absent'])
        result = self._run('subversion', [ "repo=\"http://svn.apache.org/repos/asf/subversion/trunk/notes/api-errata\"","dest=\"/tmp/meetings\"" ])
        assert result['changed']
        result = self._run('subversion', [ "repo=\"http://svn.apache.org/repos/asf/subversion/trunk/notes/meetings\"","dest=\"/tmp/meetings\"" ])
        assert result['changed']
        assert result['after'][1] == 'URL: http://svn.apache.org/repos/asf/subversion/trunk/notes/meetings'
        # when /tmp/meetings is a folder, when its an older revision it should update
        self._run('command',['svn up -r926415','chdir=/tmp/meetings'])
        result = self._run('subversion', [ "repo=\"http://svn.apache.org/repos/asf/subversion/trunk/notes/meetings\"","dest=\"/tmp/meetings\"" ])
        assert result['changed']
        assert result['before'][0] == 'Revision: 926415'
        assert result['after'][0] != 'Revision: 926415'
        # when /tmp/meetings has dirty files in it, ignore them:
        self._run('command',['touch /tmp/meetings/adirtyfile'])
        result = self._run('subversion', [ "repo=\"http://svn.apache.org/repos/asf/subversion/trunk/notes/meetings\"","dest=\"/tmp/meetings\"" ])
        assert not result['changed'] # no changes to the repo
        # when /tmp/meetings has modified file in it, fail:
        self._run('file',['path=/tmp/meetings/adirtyfile','state=absent'])
        self._run('command',['cp /tmp/meetings/berlin-11-agenda /tmp/meetings/svn-vision-agenda'])
        result = self._run('subversion', [ "repo=\"http://svn.apache.org/repos/asf/subversion/trunk/notes/meetings\"","dest=\"/tmp/meetings\"","force=no" ])
        assert result['failed']
        # when /tmp/meetings has a modified file but force is set to yes, then just override it.
        result = self._run('subversion', [ "repo=\"http://svn.apache.org/repos/asf/subversion/trunk/notes/meetings\"","dest=\"/tmp/meetings\"","force=yes" ])
        assert result['changed'] # no changes to the repo

    def test_large_output(self):
        large_path = "/usr/share/dict/words"
        if not os.path.exists(large_path):
            large_path = "/usr/share/dict/cracklib-small"
            if not os.path.exists(large_path):
                raise SkipTest
        # Ensure reading a large amount of output from a command doesn't hang.
        result = self._run('command', [ "/bin/cat", large_path ])
        assert "failed" not in result
        assert "msg" not in result
        assert result['rc'] == 0
        assert len(result['stdout']) > 100000
        assert result['stderr'] == ''

    def test_async(self):
        # test async launch and job status
        # of any particular module
        result = self._run('command', [ get_binary("sleep"), "3" ], background=20)
        assert 'ansible_job_id' in result
        assert 'started' in result
        jid = result['ansible_job_id']
        # no real chance of this op taking a while, but whatever
        time.sleep(5)
        # CLI will abstract this (when polling), but this is how it works internally
        result = self._run('async_status', [ "jid=%s" % jid ])
        # TODO: would be nice to have tests for supervisory process
        # killing job after X seconds
        assert 'finished' in result
        assert 'failed' not in result
        assert 'rc' in result
        assert 'stdout' in result
        assert result['ansible_job_id'] == jid

    def test_fetch(self):
        input_ = self._get_test_file('sample.j2')
        output = os.path.join(self.stage_dir, '127.0.0.2', input_)
        result = self._run('fetch', [ "src=%s" % input_, "dest=%s" % self.stage_dir ])
        assert os.path.exists(output)
        assert open(input_).read() == open(output).read()

    def test_assemble(self):
        input = self._get_test_file('assemble.d')
        output = self._get_stage_file('sample.out')
        result = self._run('assemble', [
            "src=%s" % input,
            "dest=%s" % output,
        ])
        assert os.path.exists(output)
        out = file(output).read()
        assert out.find("first") != -1
        assert out.find("second") != -1
        assert out.find("third") != -1
        assert result['changed'] == True
        assert 'md5sum' in result
        assert 'failed' not in result
        result = self._run('assemble', [
            "src=%s" % input,
            "dest=%s" % output,
        ])
        assert result['changed'] == False
<|endoftext|>


<commit_before>import os, os.path
import pandas as pd
import matplotlib
from abc import ABCMeta, abstractmethod
try:
    matplotlib.use('TkAgg')
except:
    pass
import matplotlib.pyplot as plt
import seaborn as sns
from qstrader import settings

class Statistics(object):
    """
    Statistics is an abstract class providing an interface for
    all inherited statistic classes (live, historic, custom, etc).

    The goal of a Statistics object is to keep a record of useful
    information about one or many trading strategies as the strategy
    is running. This is done by hooking into the main event loop and
    essentially updating the object according to portfolio performance
    over time.

    Ideally, Statistics should be subclassed according to the strategies
    and timeframes-traded by the user. Different trading strategies
    may require different metrics or frequencies-of-metrics to be updated,
    however the example given is suitable for longer timeframes.
    """

    __metaclass__ = ABCMeta

    @abstractmethod
    def update(self):
        """
        Update all the statistics according to values of the portfolio
        and open positions. This should be called from within the
        event loop.
        """
        raise NotImplementedError("Should implement update()")

    @abstractmethod
    def get_statistics(self):
        """
        Return a dict containing all statistics.
        """
        raise NotImplementedError("Should implement get_statistics()")

    @abstractmethod
    def plot_results(self):
        """
        Plot all statistics collected up until 'now'
        """
        raise NotImplementedError("Should implement plot_results()")

class SimpleStatistics(Statistics):
    """
    Simple Statistics provides a bare-bones example of statistics
    that can be collected through trading.

    Update frequency is minute-ly, which reflects the time-frame it
    has been designed to trade with.

    Statistics included are Sharpe Ratio, Drawdown, Max Drawdown,
    Max Drawdown Duration, Total Profit Value and Total Return Pct.

    TODO think about Alpha/Beta, compare strategy of benchmark.
    TODO testing
    TODO think about slippage, fill rate, etc
    TODO remove equity_file references throughout QSTrader
    """
    def __init__(self, portfolio_handler):
        """
        Takes in a portfolio handler.
        """
        self.portfolio_handler = portfolio_handler
        self.statistics = {}
        self.statistics["sharpe"] = {}
        self.statistics["drawdown"] = {}
        self.statistics["max_drawdown"] = 0
        self.statistics["equity"] = {}
        self.statistics["returms"] = {}

    def update(self):
        """
        Check time
        If now, update all self.statistics.
        """

    def plot_results(self):
        """
        A simple script to plot the balance of the portfolio, or
        "equity curve", as a function of time.
        """
        sns.set_palette("deep", desat=.6)
        sns.set_context(rc={"figure.figsize": (8, 4)})

        # equity_file = os.path.join(settings.OUTPUT_DIR, "output.csv")
        # equity = pd.io.parsers.read_csv(
            # equity_file, parse_dates=True, header=0, index_col=0
        # )

        # Plot two charts: Equity curve, period returns
        fig = plt.figure()
        fig.patch.set_facecolor('white')     # Set the outer colour to white

        # Plot the equity curve
        ax1 = fig.add_subplot(311, ylabel='Portfolio value')
        # equity["Equity"].plot(ax=ax1, color=sns.color_palette()[0])

        # Plot the returns
        ax2 = fig.add_subplot(312, ylabel='Period returns')
        # equity['Returns'].plot(ax=ax2, color=sns.color_palette()[1])

        # Plot the figure
        plt.show()

    # def create_drawdowns(self, pnl):
    #     """
    #     Calculate the largest peak-to-trough drawdown of the PnL curve
    #     as well as the duration of the drawdown. Requires that the
    #     pnl_returns is a pandas Series.
    #     Parameters:
    #     pnl - A pandas Series representing period percentage returns.
    #     Returns:
    #     drawdown, duration - Highest peak-to-trough drawdown and duration.
    #     """

    #     # Calculate the cumulative returns curve
    #     # and set up the High Water Mark
    #     hwm = [0]

    #     # Create the drawdown and duration series
    #     idx = pnl.index
    #     drawdown = pd.Series(index = idx)
    #     duration = pd.Series(index = idx)

    #     # Loop over the index range
    #     for t in range(1, len(idx)):
    #         hwm.append(max(hwm[t-1], pnl.ix[t]))
    #         drawdown.ix[t]= (hwm[t]-pnl.ix[t])
    #         duration.ix[t]= (0 if drawdown.ix[t] == 0 else duration.ix[t-1]+1)
    #     return drawdown, drawdown.max(), duration.max()

    # def generate_results(self):
    #     in_filename = "equity.csv"
    #     out_filename = "output.csv"
    #     in_file = os.path.join(settings.OUTPUT_DIR, in_filename)
    #     out_file = os.path.join(settings.OUTPUT_DIR, out_filename)

    #     # Create equity curve dataframe
    #     df = pd.read_csv(in_file, index_col=0)
    #     df.dropna(inplace=True)
    #     df["Total"] = df.sum(axis=1)
    #     df["Returns"] = df["Total"].pct_change()
    #     df["Equity"] = (1.0+df["Returns"]).cumprod()

    #     # Create drawdown statistics
    #     # drawdown, max_dd, dd_duration = self.create_drawdowns(df["Equity"])
    #     # df["Drawdown"] = drawdown

    #     df.to_csv(out_file, index=True)

    #     print("Simulation complete and results exported to %s" % out_filename)<commit_msg>Split out Statistics.statistics into individual components. Consolidate them again only in get_statistics method.<commit_after>import os, os.path
import pandas as pd
import matplotlib
from abc import ABCMeta, abstractmethod
try:
    matplotlib.use('TkAgg')
except:
    pass
import matplotlib.pyplot as plt
import seaborn as sns
from qstrader import settings

class Statistics(object):
    """
    Statistics is an abstract class providing an interface for
    all inherited statistic classes (live, historic, custom, etc).

    The goal of a Statistics object is to keep a record of useful
    information about one or many trading strategies as the strategy
    is running. This is done by hooking into the main event loop and
    essentially updating the object according to portfolio performance
    over time.

    Ideally, Statistics should be subclassed according to the strategies
    and timeframes-traded by the user. Different trading strategies
    may require different metrics or frequencies-of-metrics to be updated,
    however the example given is suitable for longer timeframes.
    """

    __metaclass__ = ABCMeta

    @abstractmethod
    def update(self):
        """
        Update all the statistics according to values of the portfolio
        and open positions. This should be called from within the
        event loop.
        """
        raise NotImplementedError("Should implement update()")

    @abstractmethod
    def get_statistics(self):
        """
        Return a dict containing all statistics.
        """
        raise NotImplementedError("Should implement get_statistics()")

    @abstractmethod
    def plot_results(self):
        """
        Plot all statistics collected up until 'now'
        """
        raise NotImplementedError("Should implement plot_results()")

class SimpleStatistics(Statistics):
    """
    Simple Statistics provides a bare-bones example of statistics
    that can be collected through trading.

    Update frequency is minute-ly, which reflects the time-frame it
    has been designed to trade with.

    Statistics included are Sharpe Ratio, Drawdown, Max Drawdown,
    Max Drawdown Duration, Total Profit Value and Total Return Pct.

    TODO think about Alpha/Beta, compare strategy of benchmark.
    TODO testing
    TODO think about slippage, fill rate, etc
    TODO remove equity_file references throughout QSTrader
    """
    def __init__(self, portfolio_handler):
        """
        Takes in a portfolio handler.
        """
        self.portfolio_handler = portfolio_handler
        self.sharpe={}
        self.drawdowns={}
        self.max_drawdown=0
        self.equity={}
        self.returns={}

    def update(self):
        """
        Check time
        If now, update all self.statistics.
        """
        import pdb; pdb.set_trace()

    def get_statistics(self):
        """
        Return a dict with all statistics
        """
        statistics={}
        statistics["sharpe"]=self.sharpe
        statistics["drawdowns"]=self.drawdowns
        statistics["max_drawdown"]=self.max_drawdown
        statistics["equity"]=self.equity
        statistics["returns"]=self.returns

    def plot_results(self):
        """
        A simple script to plot the balance of the portfolio, or
        "equity curve", as a function of time.
        """
        sns.set_palette("deep", desat=.6)
        sns.set_context(rc={"figure.figsize": (8, 4)})

        # equity_file = os.path.join(settings.OUTPUT_DIR, "output.csv")
        # equity = pd.io.parsers.read_csv(
            # equity_file, parse_dates=True, header=0, index_col=0
        # )

        # Plot two charts: Equity curve, period returns
        fig = plt.figure()
        fig.patch.set_facecolor('white')     # Set the outer colour to white

        # Plot the equity curve
        ax1 = fig.add_subplot(311, ylabel='Portfolio value')
        # equity["Equity"].plot(ax=ax1, color=sns.color_palette()[0])

        # Plot the returns
        ax2 = fig.add_subplot(312, ylabel='Period returns')
        # equity['Returns'].plot(ax=ax2, color=sns.color_palette()[1])

        # Plot the figure
        plt.show()

    # def create_drawdowns(self, pnl):
    #     """
    #     Calculate the largest peak-to-trough drawdown of the PnL curve
    #     as well as the duration of the drawdown. Requires that the
    #     pnl_returns is a pandas Series.
    #     Parameters:
    #     pnl - A pandas Series representing period percentage returns.
    #     Returns:
    #     drawdown, duration - Highest peak-to-trough drawdown and duration.
    #     """

    #     # Calculate the cumulative returns curve
    #     # and set up the High Water Mark
    #     hwm = [0]

    #     # Create the drawdown and duration series
    #     idx = pnl.index
    #     drawdown = pd.Series(index = idx)
    #     duration = pd.Series(index = idx)

    #     # Loop over the index range
    #     for t in range(1, len(idx)):
    #         hwm.append(max(hwm[t-1], pnl.ix[t]))
    #         drawdown.ix[t]= (hwm[t]-pnl.ix[t])
    #         duration.ix[t]= (0 if drawdown.ix[t] == 0 else duration.ix[t-1]+1)
    #     return drawdown, drawdown.max(), duration.max()

    # def generate_results(self):
    #     in_filename = "equity.csv"
    #     out_filename = "output.csv"
    #     in_file = os.path.join(settings.OUTPUT_DIR, in_filename)
    #     out_file = os.path.join(settings.OUTPUT_DIR, out_filename)

    #     # Create equity curve dataframe
    #     df = pd.read_csv(in_file, index_col=0)
    #     df.dropna(inplace=True)
    #     df["Total"] = df.sum(axis=1)
    #     df["Returns"] = df["Total"].pct_change()
    #     df["Equity"] = (1.0+df["Returns"]).cumprod()

    #     # Create drawdown statistics
    #     # drawdown, max_dd, dd_duration = self.create_drawdowns(df["Equity"])
    #     # df["Drawdown"] = drawdown

    #     df.to_csv(out_file, index=True)

    #     print("Simulation complete and results exported to %s" % out_filename)<|endoftext|>