<commit_before>import os
import sickle
import boto
import datetime
import requests
from time import sleep
from time import time
from util import elapsed
import logging
import zlib
import re
import json
import sys
import random
import argparse
from elasticsearch import Elasticsearch, RequestsHttpConnection, compat, exceptions
from elasticsearch.helpers import parallel_bulk
from elasticsearch.helpers import bulk
from elasticsearch.helpers import scan
from multiprocessing import Process
from multiprocessing import Queue
from multiprocessing import Pool
from HTMLParser import HTMLParser

import oa_local
from oa_base import get_urls_from_our_base_doc
from publication import call_targets_in_parallel
from webpage import WebpageInUnknownRepo
from util import JSONSerializerPython2

# set up elasticsearch
INDEX_NAME = "base"
TYPE_NAME = "record"

libraries_to_mum = [
    "requests.packages.urllib3",
    "requests_oauthlib",
    "stripe",
    "oauthlib",
    "boto",
    "newrelic",
    "RateLimiter",
    "elasticsearch",
    "urllib3"
]

for a_library in libraries_to_mum:
    the_logger = logging.getLogger(a_library)
    the_logger.setLevel(logging.WARNING)
    the_logger.propagate = True

def set_up_elastic(url):
    if not url:
        url = os.getenv("BASE_URL")
    es = Elasticsearch(url,
                       serializer=JSONSerializerPython2(),
                       retry_on_timeout=True,
                       max_retries=100)
    return es

def save_records_in_es(es, records_to_save, threads, chunk_size):
    start_time = time()

    # have to do call parallel_bulk in a for loop because is parallel_bulk is a generator so you have to call it to
    # have it do the work.  see https://discuss.elastic.co/t/helpers-parallel-bulk-in-python-not-working/39498
    if threads > 1:
        for success, info in parallel_bulk(es,
                                           actions=records_to_save,
                                           refresh=True,
                                           request_timeout=60,
                                           thread_count=threads,
                                           chunk_size=chunk_size):
            if not success:
                print("A document failed:", info)
    else:
        for success_info in bulk(es, actions=records_to_save, refresh=True, request_timeout=60, chunk_size=chunk_size):
            pass
    print u"done sending {} records to elastic in {}s".format(len(records_to_save), elapsed(start_time, 4))
    print records_to_save[0]

random_query_dict = {
  "_source": [
    "id"
  ],
  "size": 1000,
  "from": int(random.random()*7999),
  "query": {
    "bool": {
      "must_not": [{
        "exists": {
          "field": "fulltext_last_updated"
        }
        }
        ,{
        "match": {
          "urls": "elib.uraic.ru"
        }}
        ,{
        "match": {
          "urls": "elar.urfu.ru"
        }}
        ,{
        "exists": {
          "field": "random"
        }
        }
        ],
      "must": {
        "term": {
          "oa": 2
        }
      }
    }
  }
}

query_dict = {
  "_source": [
    "title",
    "urls",
    "license",
    "sources",
    "id"
  ],
  "size": 100,
  "from": 0,  #overwritten
  "query": {
    "bool": {
      "must_not": [
        {
          "exists": {
            "field": "fulltext_last_updated"
          }
        },
        {
          "match": {
            "urls": "elib.uraic.ru"
          }
        },
        {
          "match": {
            "urls": "elar.urfu.ru"
          }
        }
      ],
      "must": [
        {
          "term": {
            "oa": 2
          }
        },
        {
          "exists": {
            "field": "random"
          }
        }
      ]
    }
  },
  "sort": [
    {
      "random": "asc"
    }
  ]
}

class BaseResult(object):
    def __init__(self, doc):
        self.doc = doc
        self.fulltext_last_updated = datetime.datetime.utcnow().isoformat()
        self.fulltext_url_dicts = []
        self.license = None
        self.set_webpages()

    def scrape_for_fulltext(self):
        response_webpages = []

        found_open_fulltext = False
        for my_webpage in self.webpages:
            if not found_open_fulltext:
                my_webpage.scrape_for_fulltext_link()
                if my_webpage.has_fulltext_url:
                    print u"** found an open version! {}".format(my_webpage.fulltext_url)
                    found_open_fulltext = True
                    response_webpages.append(my_webpage)

        self.open_webpages = response_webpages
        sys.exc_clear()  # someone on the internet said this would fix All The Memory Problems. has to be in the thread.
        return self

    def set_webpages(self):
        self.open_webpages = []
        self.webpages = []
        for url in get_urls_from_our_base_doc(self.doc):
            my_webpage = WebpageInUnknownRepo(url=url)
            self.webpages.append(my_webpage)

    def set_fulltext_urls(self):

        # first set license if there is one originally.  overwrite it later if scraped a better one.
        if "license" in self.doc and self.doc["license"]:
            self.license = oa_local.find_normalized_license(self.doc["license"])

        for my_webpage in self.open_webpages:
            if my_webpage.has_fulltext_url:
                self.fulltext_url_dicts += [{"free_pdf_url": my_webpage.scraped_pdf_url, "pdf_landing_page": my_webpage.url}]
                if not self.license or self.license == "unknown":
                    self.license = my_webpage.scraped_license
            else:
                print "{} has no fulltext url alas".format(my_webpage)

        if self.license == "unknown":
            self.license = None

    def make_action_record(self, just_random=False):
        update_doc = {"random": random.random()}

        if not just_random:
            update_doc.update({
                            "fulltext_last_updated": self.fulltext_last_updated,
                            "fulltext_url_dicts": self.fulltext_url_dicts,
                            "fulltext_license": self.license,
                            "fulltext_updated": None})

        action = {"doc": update_doc}
        action["_id"] = self.doc["id"]
        action["_op_type"] = "update"
        action["_type"] = TYPE_NAME
        action["_index"] = INDEX_NAME
        action["_retry_on_conflict"] = 3
        # print "\n", action
        return action

def do_a_loop(first=None, last=None, url=None, threads=0, chunk_size=None):
    just_random = False

    loop_start = time()
    es = set_up_elastic(url)
    print u"set_up_elastic took {}s".format(elapsed(loop_start, 2))

    if just_random:
        results = es.search(index=INDEX_NAME, body=random_query_dict, request_timeout=10000)
    else:
        # different every loop
        query_dict["from"] = int(random.random()*7999)
        results = es.search(index=INDEX_NAME, body=query_dict, request_timeout=10000)
    # print u"search body:\n{}".format(query)
    print u"took {}s to search ES".format(elapsed(loop_start, 2))
    records_to_save = []

    # decide if should stop looping after this
    if not results["hits"]["hits"]:
        print "no hits!  exiting"
        sys.exit()

    base_results = []
    for base_hit in results["hits"]["hits"]:
        base_hit_doc = base_hit["_source"]
        base_results.append(BaseResult(base_hit_doc))

    scrape_start = time()

    if not just_random:
        targets = [base_result.scrape_for_fulltext for base_result in base_results]
        call_targets_in_parallel(targets)
        print u"scraping {} webpages took {}s".format(len(base_results), elapsed(scrape_start, 2))

    for base_result in base_results:
        base_result.set_fulltext_urls()
        records_to_save.append(base_result.make_action_record(just_random))

    # print "len of records_to_save", len(records_to_save)
    # print "records_to_save:", records_to_save
    save_records_in_es(es, records_to_save, threads, chunk_size)
    print "** took {}s to do {}, {:,} remaining\n".format(
        elapsed(loop_start, 2), len(base_results), results["hits"]["total"])

def update_base2s():
    has_more_records = True
    while has_more_records:
        pool_time = time()
        my_process = Process(target=do_a_loop)
        my_process.daemon = True
        my_process.start()
        my_process.join()
        my_process.terminate()
        # print u"took {}s for do_a_loop".format(elapsed(pool_time, 2))

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Run stuff.")

    # just for updating lots
    function = update_base2s

    parsed = parser.parse_args()

    print u"calling {} with these args: {}".format(function.__name__, vars(parsed))
    function(**vars(parsed))

<commit_msg>set random and base1s<commit_after>import os
import sickle
import boto
import datetime
import requests
from time import sleep
from time import time
from util import elapsed
import logging
import zlib
import re
import json
import sys
import random
import argparse
from elasticsearch import Elasticsearch, RequestsHttpConnection, compat, exceptions
from elasticsearch.helpers import parallel_bulk
from elasticsearch.helpers import bulk
from elasticsearch.helpers import scan
from multiprocessing import Process
from multiprocessing import Queue
from multiprocessing import Pool
from HTMLParser import HTMLParser

import oa_local
from oa_base import get_urls_from_our_base_doc
from publication import call_targets_in_parallel
from webpage import WebpageInUnknownRepo
from util import JSONSerializerPython2

# set up elasticsearch
INDEX_NAME = "base"
TYPE_NAME = "record"

libraries_to_mum = [
    "requests.packages.urllib3",
    "requests_oauthlib",
    "stripe",
    "oauthlib",
    "boto",
    "newrelic",
    "RateLimiter",
    "elasticsearch",
    "urllib3"
]

for a_library in libraries_to_mum:
    the_logger = logging.getLogger(a_library)
    the_logger.setLevel(logging.WARNING)
    the_logger.propagate = True

def set_up_elastic(url):
    if not url:
        url = os.getenv("BASE_URL")
    es = Elasticsearch(url,
                       serializer=JSONSerializerPython2(),
                       retry_on_timeout=True,
                       max_retries=100)
    return es

def save_records_in_es(es, records_to_save, threads, chunk_size):
    start_time = time()

    # have to do call parallel_bulk in a for loop because is parallel_bulk is a generator so you have to call it to
    # have it do the work.  see https://discuss.elastic.co/t/helpers-parallel-bulk-in-python-not-working/39498
    if threads > 1:
        for success, info in parallel_bulk(es,
                                           actions=records_to_save,
                                           refresh=True,
                                           request_timeout=60,
                                           thread_count=threads,
                                           chunk_size=chunk_size):
            if not success:
                print("A document failed:", info)
    else:
        for success_info in bulk(es, actions=records_to_save, refresh=True, request_timeout=60, chunk_size=chunk_size):
            pass
    print u"done sending {} records to elastic in {}s".format(len(records_to_save), elapsed(start_time, 4))
    print records_to_save[0]

random_query_dict = {
  "_source": [
    "title",
    "urls",
    "license",
    "sources",
    "oa",
    "id"
  ],
 "size": 1000,
  "from": int(random.random()*7999),
  "query": {
    "bool": {
      "must_not":
        {
        "exists": {
          "field": "random"
            }
        }
    }
  }
}

query_dict = {
  "_source": [
    "title",
    "urls",
    "license",
    "sources",
    "id"
  ],
  "size": 100,
  "from": 0,  #overwritten
  "query": {
    "bool": {
      "must_not": [
        {
          "exists": {
            "field": "fulltext_last_updated"
          }
        },
        {
          "match": {
            "urls": "elib.uraic.ru"
          }
        },
        {
          "match": {
            "urls": "elar.urfu.ru"
          }
        }
      ],
      "must": [
        {
          "term": {
            "oa": 2
          }
        },
        {
          "exists": {
            "field": "random"
          }
        }
      ]
    }
  },
  "sort": [
    {
      "random": "asc"
    }
  ]
}

class BaseResult(object):
    def __init__(self, doc):
        self.doc = doc
        self.fulltext_last_updated = datetime.datetime.utcnow().isoformat()
        self.fulltext_url_dicts = []
        self.license = None
        self.set_webpages()

    def set_base1s(self):
        if self.doc["oa"] == 1:
            for my_webpage in self.webpages:
                my_webpage.scraped_open_metadata_url = my_webpage.url
                self.open_webpages.append(my_webpage)
        return self

    def scrape_for_fulltext(self):
        response_webpages = []

        found_open_fulltext = False
        for my_webpage in self.webpages:
            if not found_open_fulltext:
                my_webpage.scrape_for_fulltext_link()
                if my_webpage.has_fulltext_url:
                    print u"** found an open version! {}".format(my_webpage.fulltext_url)
                    found_open_fulltext = True
                    response_webpages.append(my_webpage)

        self.open_webpages = response_webpages
        sys.exc_clear()  # someone on the internet said this would fix All The Memory Problems. has to be in the thread.
        return self

    def set_webpages(self):
        self.open_webpages = []
        self.webpages = []
        for url in get_urls_from_our_base_doc(self.doc):
            my_webpage = WebpageInUnknownRepo(url=url)
            self.webpages.append(my_webpage)

    def set_fulltext_urls(self):

        # first set license if there is one originally.  overwrite it later if scraped a better one.
        if "license" in self.doc and self.doc["license"]:
            self.license = oa_local.find_normalized_license(self.doc["license"])

        for my_webpage in self.open_webpages:
            if my_webpage.has_fulltext_url:
                response = {}
                self.fulltext_url_dicts += [{"free_pdf_url": my_webpage.scraped_pdf_url, "pdf_landing_page": my_webpage.url}]
                if not self.license or self.license == "unknown":
                    self.license = my_webpage.scraped_license
            else:
                print "{} has no fulltext url alas".format(my_webpage)

        if self.license == "unknown":
            self.license = None

    def make_action_record(self):
        update_doc = {
            "random": random.random(),
            "fulltext_last_updated": self.fulltext_last_updated,
            "fulltext_url_dicts": self.fulltext_url_dicts,
            "fulltext_license": self.license,
            "fulltext_updated": None
        }
        action = {"doc": update_doc}
        action["_id"] = self.doc["id"]
        action["_op_type"] = "update"
        action["_type"] = TYPE_NAME
        action["_index"] = INDEX_NAME
        action["_retry_on_conflict"] = 3
        # print "\n", action
        return action

def do_a_loop(first=None, last=None, url=None, threads=0, chunk_size=None):
    just_random = True

    loop_start = time()
    es = set_up_elastic(url)

    if just_random:
        random_query_dict["from"] = int(random.random()*7999)
        results = es.search(index=INDEX_NAME, body=random_query_dict, request_timeout=10000)
    else:
        # different every loop
        query_dict["from"] = int(random.random()*7999)
        results = es.search(index=INDEX_NAME, body=query_dict, request_timeout=10000)
    # print u"search body:\n{}".format(query)
    print u"took {}s to search ES".format(elapsed(loop_start, 2))
    records_to_save = []

    # decide if should stop looping after this
    if not results["hits"]["hits"]:
        print "no hits!  exiting"
        sys.exit()

    base_results = []
    for base_hit in results["hits"]["hits"]:
        base_hit_doc = base_hit["_source"]
        base_results.append(BaseResult(base_hit_doc))

    scrape_start = time()

    if not just_random:
        targets = [base_result.scrape_for_fulltext for base_result in base_results]
        call_targets_in_parallel(targets)
        print u"scraping {} webpages took {}s".format(len(base_results), elapsed(scrape_start, 2))

    targets = [base_result.set_base1s for base_result in base_results]
    call_targets_in_parallel(targets)

    for base_result in base_results:
        base_result.set_fulltext_urls()
        records_to_save.append(base_result.make_action_record())

    # print "len of records_to_save", len(records_to_save)
    # print "records_to_save:", records_to_save
    save_records_in_es(es, records_to_save, threads, chunk_size)
    print "** took {}s to do {}, {:,} remaining\n".format(
        elapsed(loop_start, 2), len(base_results), results["hits"]["total"])

def run():
    has_more_records = True
    while has_more_records:
        pool_time = time()
        my_process = Process(target=do_a_loop)
        my_process.daemon = True
        my_process.start()
        my_process.join()
        my_process.terminate()
        # print u"took {}s for do_a_loop".format(elapsed(pool_time, 2))

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Run stuff.")

    # just for updating lots
    function = run

    parsed = parser.parse_args()

    print u"calling {} with these args: {}".format(function.__name__, vars(parsed))
    function(**vars(parsed))

<|endoftext|>


<commit_before># Generic modularized configuration file manager.
#
# Author: Peter Odding <peter@peterodding.com>
# Last Change: July 13, 2017
# URL: https://pypi.python.org/pypi/update-dotdee

"""The Python API of the update-dotdee program."""

# Standard library modules.
import hashlib
import logging
import os

# External dependencies.
from executor.contexts import LocalContext
from humanfriendly import compact, format_path, pluralize
from natsort import natsort
from property_manager import PropertyManager, mutable_property, required_property

# Semi-standard module versioning.
__version__ = '2.0'

# Initialize a logger for this module.
logger = logging.getLogger(__name__)

class UpdateDotDee(PropertyManager):

    """The :class:`UpdateDotDee` class implements the Python API of update-dotdee."""

    @mutable_property
    def checksum_file(self):
        """The pathname of the file that stores the checksum of the generated file (a string)."""
        return os.path.join(self.directory, '.checksum')

    @mutable_property(cached=True)
    def context(self):
        """An execution context created by :mod:`executor.contexts`."""
        return LocalContext()

    @mutable_property
    def directory(self):
        """The pathname of the directory with configuration snippets (a string)."""
        return self.filename + '.d'

    @required_property
    def filename(self):
        """The pathname of the configuration file to generate (a string)."""

    @mutable_property
    def force(self):
        """:data:`True` to overwrite modified files, :data:`False` to abort (the default)."""
        return False

    def update_file(self, force=None):
        """
        Update the file with the contents of the files in the ``.d`` directory.

        :param force: Override the value of :attr:`force` (a boolean or
                      :data:`None`).
        :raises: :exc:`RefuseToOverwrite` when :attr:`force` is :data:`False`
                 and the contents of :attr:`filename` were modified.
        """
        if force is None:
            force = self.force
        if not self.context.is_directory(self.directory):
            # Create the .d directory.
            logger.info("Creating directory %s", format_path(self.directory))
            self.context.execute('mkdir', '-p', self.directory, tty=False)
            # Move the original file into the .d directory.
            local_file = os.path.join(self.directory, 'local')
            logger.info("Moving %s to %s", format_path(self.filename), format_path(local_file))
            self.context.execute('mv', self.filename, local_file, tty=False)
        # Read the modularized configuration file(s).
        blocks = []
        for filename in natsort(self.context.list_entries(self.directory)):
            if not filename.startswith('.'):
                blocks.append(self.read_file(os.path.join(self.directory, filename)))
        contents = b"\n\n".join(blocks)
        # Make sure the generated file was not modified? We skip this on the
        # first run, when the original file was just moved into the newly
        # created directory (see above).
        if all(map(self.context.is_file, (self.filename, self.checksum_file))):
            logger.info("Checking for local changes to %s ..", format_path(self.filename))
            if self.hash_contents() != self.context.read_file(self.checksum_file):
                if force:
                    logger.warning(compact(
                        """
                        The contents of the file to generate ({filename})
                        were modified but --force was used so overwriting
                        anyway!
                        """,
                        filename=format_path(self.filename),
                    ))
                else:
                    raise RefuseToOverwrite(compact(
                        """
                        The contents of the file to generate ({filename})
                        were modified and I'm refusing to overwrite it! If
                        you're sure you want to proceed, use the --force
                        option or delete the file {checksum_file} and
                        retry.
                        """,
                        filename=format_path(self.filename),
                        checksum_file=format_path(self.checksum_file),
                    ))
        # Update the generated configuration file.
        self.write_file(self.filename, contents)
        # Update the checksum file.
        self.context.write_file(self.checksum_file, self.hash_contents())

    def read_file(self, filename):
        """
        Read a text file and provide feedback to the user.

        :param filename: The pathname of the file to read (a string).
        :returns: The contents of the file (a string).
        """
        logger.info("Reading file: %s", format_path(filename))
        contents = self.context.read_file(filename)
        num_lines = len(contents.splitlines())
        logger.debug("Read %i line%s from %s",
                     num_lines, '' if num_lines == 1 else 's',
                     format_path(filename))
        return contents.rstrip()

    def write_file(self, filename, contents):
        """
        Write a text file and provide feedback to the user.

        :param filename: The pathname of the file to write (a string).
        :param contents: The new contents of the file (a string).
        """
        logger.info("Writing file: %s", format_path(filename))
        contents = contents.rstrip() + b"\n"
        self.context.write_file(filename, contents)
        logger.debug("Wrote %s to %s.",
                     pluralize(len(contents.splitlines()), "line"),
                     format_path(filename))

    def hash_contents(self):
        """
        Hash the text file using the SHA1 algorithm.

        :returns: A string containing a hexadecimal SHA1 digest.
        """
        logger.debug("Calculating SHA1 of %s", format_path(self.filename))
        context = hashlib.sha1()
        context.update(self.context.read_file(self.filename))
        hexdigest = context.hexdigest()
        logger.debug("SHA1 of %s is %s", format_path(self.filename), hexdigest)
        return hexdigest

class RefuseToOverwrite(Exception):

    """Raised when update-dotdee notices that a generated file was modified."""
<commit_msg>Release 3.0: Switch to executor.contexts (allow remote execution)<commit_after># Generic modularized configuration file manager.
#
# Author: Peter Odding <peter@peterodding.com>
# Last Change: July 13, 2017
# URL: https://pypi.python.org/pypi/update-dotdee

"""The Python API of the update-dotdee program."""

# Standard library modules.
import hashlib
import logging
import os

# External dependencies.
from executor.contexts import LocalContext
from humanfriendly import compact, format_path, pluralize
from natsort import natsort
from property_manager import PropertyManager, mutable_property, required_property

# Semi-standard module versioning.
__version__ = '3.0'

# Initialize a logger for this module.
logger = logging.getLogger(__name__)

class UpdateDotDee(PropertyManager):

    """The :class:`UpdateDotDee` class implements the Python API of update-dotdee."""

    @mutable_property
    def checksum_file(self):
        """The pathname of the file that stores the checksum of the generated file (a string)."""
        return os.path.join(self.directory, '.checksum')

    @mutable_property(cached=True)
    def context(self):
        """An execution context created by :mod:`executor.contexts`."""
        return LocalContext()

    @mutable_property
    def directory(self):
        """The pathname of the directory with configuration snippets (a string)."""
        return self.filename + '.d'

    @required_property
    def filename(self):
        """The pathname of the configuration file to generate (a string)."""

    @mutable_property
    def force(self):
        """:data:`True` to overwrite modified files, :data:`False` to abort (the default)."""
        return False

    def update_file(self, force=None):
        """
        Update the file with the contents of the files in the ``.d`` directory.

        :param force: Override the value of :attr:`force` (a boolean or
                      :data:`None`).
        :raises: :exc:`RefuseToOverwrite` when :attr:`force` is :data:`False`
                 and the contents of :attr:`filename` were modified.
        """
        if force is None:
            force = self.force
        if not self.context.is_directory(self.directory):
            # Create the .d directory.
            logger.info("Creating directory %s", format_path(self.directory))
            self.context.execute('mkdir', '-p', self.directory, tty=False)
            # Move the original file into the .d directory.
            local_file = os.path.join(self.directory, 'local')
            logger.info("Moving %s to %s", format_path(self.filename), format_path(local_file))
            self.context.execute('mv', self.filename, local_file, tty=False)
        # Read the modularized configuration file(s).
        blocks = []
        for filename in natsort(self.context.list_entries(self.directory)):
            if not filename.startswith('.'):
                blocks.append(self.read_file(os.path.join(self.directory, filename)))
        contents = b"\n\n".join(blocks)
        # Make sure the generated file was not modified? We skip this on the
        # first run, when the original file was just moved into the newly
        # created directory (see above).
        if all(map(self.context.is_file, (self.filename, self.checksum_file))):
            logger.info("Checking for local changes to %s ..", format_path(self.filename))
            if self.hash_contents() != self.context.read_file(self.checksum_file):
                if force:
                    logger.warning(compact(
                        """
                        The contents of the file to generate ({filename})
                        were modified but --force was used so overwriting
                        anyway!
                        """,
                        filename=format_path(self.filename),
                    ))
                else:
                    raise RefuseToOverwrite(compact(
                        """
                        The contents of the file to generate ({filename})
                        were modified and I'm refusing to overwrite it! If
                        you're sure you want to proceed, use the --force
                        option or delete the file {checksum_file} and
                        retry.
                        """,
                        filename=format_path(self.filename),
                        checksum_file=format_path(self.checksum_file),
                    ))
        # Update the generated configuration file.
        self.write_file(self.filename, contents)
        # Update the checksum file.
        self.context.write_file(self.checksum_file, self.hash_contents())

    def read_file(self, filename):
        """
        Read a text file and provide feedback to the user.

        :param filename: The pathname of the file to read (a string).
        :returns: The contents of the file (a string).
        """
        logger.info("Reading file: %s", format_path(filename))
        contents = self.context.read_file(filename)
        num_lines = len(contents.splitlines())
        logger.debug("Read %i line%s from %s",
                     num_lines, '' if num_lines == 1 else 's',
                     format_path(filename))
        return contents.rstrip()

    def write_file(self, filename, contents):
        """
        Write a text file and provide feedback to the user.

        :param filename: The pathname of the file to write (a string).
        :param contents: The new contents of the file (a string).
        """
        logger.info("Writing file: %s", format_path(filename))
        contents = contents.rstrip() + b"\n"
        self.context.write_file(filename, contents)
        logger.debug("Wrote %s to %s.",
                     pluralize(len(contents.splitlines()), "line"),
                     format_path(filename))

    def hash_contents(self):
        """
        Hash the text file using the SHA1 algorithm.

        :returns: A string containing a hexadecimal SHA1 digest.
        """
        logger.debug("Calculating SHA1 of %s", format_path(self.filename))
        context = hashlib.sha1()
        context.update(self.context.read_file(self.filename))
        hexdigest = context.hexdigest()
        logger.debug("SHA1 of %s is %s", format_path(self.filename), hexdigest)
        return hexdigest

class RefuseToOverwrite(Exception):

    """Raised when update-dotdee notices that a generated file was modified."""
<|endoftext|>


<commit_before>import argparse
import os
import math

class Arm:
    def __init__(self):
        # load driver, redirect stdout to /dev/null
        # although servoblaster seems to be writing errors to stdout too, instead of stderr...
        os.system("sudo killall servod && sudo ./../bin/servod 1> /dev/null")

        self.spoon_position = 0
        self.spoon_ease = 0.1
        self.leg_position = 0

    def ease_spoon_to(self, deg):
        while(self.spoon_position != deg):
            self.move_to(5, self.spoon_position)
            self.spoon_position += 1

    def move_to(self, servo, deg):
        pwm = deg * 0.94 # 1 degree = 0.094ms high pulse time
        pwm += 50  # minimum high pulse time is 0.5 milliseconds

        if os.path.exists('/dev/servoblaster'):
            # os.system("echo 2=" + str(math.ceil(pwm)) + " > /dev/servoblaster")
            os.system("echo " + str(servo) + "=" + str(math.ceil(pwm)) + " > /dev/servoblaster")
        else:
            raise 'ERROR: Servo driver was not found. Is servoblaster loaded?'

# parser = argparse.ArgumentParser()
# parser.add_argument('--deg', required=True, type=int)
# args = parser.parse_args()
#
a = Arm()
# a.move_to(5, 0)
a.ease_spoon_to(180)
<commit_msg>Update<commit_after>import argparse
import os
import math
import time

class Arm:
    def __init__(self):
        # load driver, redirect stdout to /dev/null
        # although servoblaster seems to be writing errors to stdout too, instead of stderr...
        os.system("sudo killall servod && sudo ./../bin/servod 1> /dev/null")

        self.spoon_actual_pos = 0
        self.spoon_next_pos = 0
        self.spoon_ease = 0.1
        # self.leg_position = 0

    def ease_spoon_to(self, deg):
        start_position = self.spoon_position
        end_position = deg
        millis = int(round(time.time() * 1000))

        while(self.spoon_position != deg):
            elapsed_time = int(round(time.time() * 1000)) - millis

            self.spoon_position = self.ease_in_out_quad(self.spoon_position, elapsed_time, start_position, end_position, 2000)
            self.move_to(5, self.spoon_position)
            # self.spoon_position += 1

    def ease_in_out_quad(self, x, t, s, e, d):
        if ((t/=d/2) < 1):
             return c/2*t*t + b;
        return -c/2 * ((--t)*(t-2) - 1) + b

    def move_to(self, servo, deg):
        pwm = deg * 0.94 # 1 degree = 0.094ms high pulse time
        pwm += 50  # minimum high pulse time is 0.5 milliseconds

        if os.path.exists('/dev/servoblaster'):
            # os.system("echo 2=" + str(math.ceil(pwm)) + " > /dev/servoblaster")
            os.system("echo " + str(servo) + "=" + str(math.ceil(pwm)) + " > /dev/servoblaster")
        else:
            raise 'ERROR: Servo driver was not found. Is servoblaster loaded?'

# parser = argparse.ArgumentParser()
# parser.add_argument('--deg', required=True, type=int)
# args = parser.parse_args()
#
a = Arm()
# a.move_to(5, 0)
a.ease_spoon_to(180)
<|endoftext|>


<commit_before># -*- coding: latin-1 -*-

"""Mass Upload Libray − handling metadata."""

__authors__ = 'User:Jean-Frédéric'
import sys

sys.path.append('/home/jfk/Tuile/pywikipedia-rewrite')
sys.path.append('/home/jfk/Tuile/pywikipedia-rewrite/scripts')

import re
import sys
import pickle
import codecs
import csv
from os.path import join
import os
from collections import Counter
import UnicodeCSV
import pywikibot
import pywikibot.textlib as textlib
import data_ingestion
from data_ingestion import Photo
from PostProcessing import MetadataMapping

class MetadataRecord(Photo):

    """Represent a Record, with its associated metadata."""

    def get_field_names(self):
        """Return the field values from the records."""
        return self.metadata.keys()

    def post_process(self, mapping):
        """Post-process the MetadataRecord with the given mapping.

        For each field of the record,
        call the relevant post-processing method.

        """
        for field in self.metadata.keys():
            if field in mapping.keys():
                self.post_process_wrapper(field, mapping[field])

    def post_process_wrapper(self, field, method):
        """Wrap the post-processing methods.

        Pop the value for the given field from the record
        Apply the given method on it
        Update the record with it

        """

        def _smart_update_dict(aDict, other_dict):
            for key, value in other_dict.items():
                if key in aDict:
                    new = []
                    try:
                        aDict[key].extend(value)
                    except:
                        aDict[key] = value.append(aDict[key])
                else:
                    aDict[key] = value

        old_field_value = self.metadata.pop(field, None)
        new_field_value = method(field, old_field_value)
        _smart_update_dict(self.metadata, new_field_value)

    def to_template(self, template=u'Ingestion layout'):
        """Return the Record as a MediaWiki template."""
        return textlib.glue_template_and_params((template,
                                                 self.__dict__))

    def to_disk(self, directory):
        """Write the Record on disk in a given repository.

        Serialise the record and name it as the file title,
        do nothing if anything goes wrong.

        """
        fileName = join(directory, self.get_title() + '.dump')
        try:
            with open(fileName, 'w') as f:
                pickle.dump(self, f)
        except Exception, e:
            print "Could not pickle record %s \n %s" % (fileName, e)

class MetadataCollection(object):

    """Processing a metadata collection in various ways."""

    def __init__(self):
        """Constructor."""
        self.records = []

    def retrieve_metadata_from_csv(self, csv_file, delimiter=','):
        """Retrieve metadata from the given CSV file."""
        file_handler = codecs.open(csv_file, 'r', 'utf-8')
        csvReader = UnicodeCSV.unicode_csv_dictreader(file_handler,
                                                      delimiter=delimiter)
        try:
            for row in csvReader:
                self.records.append(self.handle_record(row))
        except csv.Error, e:
            sys.exit('file %s, line %d: %s' % (self.csv_file,
                                               csvReader.line_num, e))

    def retrieve_metadata_from_files(self, files_path):
        """Retrieve metadata from the files themselves in a given path."""
        print "Iterating in %s" % files_path
        assert os.path.exists(files_path)
        files = os.listdir(files_path)
        for image in files:
            self.records.append(self.handle_record(join(files_path, image)))

    def handle_record(self, items):
        """Handle a record.

        Default handling behaviour: return a new MetadataRecord
        initialised with the items as record values.

        """
        return MetadataRecord(None, items)

    def post_process_collection(self, method_mapping):
        """Call on each record its post_process method."""
        for record in self.records:
            record.post_process(method_mapping)
            yield record

    def print_metadata_of_record(self, index):
        """Print the metadata of the record.

        Print the result of get_title and of to_template.

        """
        print self.records[index].get_title()
        print self.records[index].to_template()

    def index_unique_metadata_values(self):
        """Count unique metadata values.

        Loop over the metadata collection.
        Return a dictionary of Sets (?)

        """

        def _add_to_set_dict(aDict, key, value):
            if key not in aDict:
                aDict[key] = set()
            try:
                aDict[key].add(value)
            except TypeError, e:
                aDict[key].update(value)

        sets_dict = dict()
        for record in self.records:
            for field, field_value in record.metadata.items():
                _add_to_set_dict(sets_dict, field, field_value)
        return sets_dict

    def count_metadata_values(self):
        """Count the metadata values for earch record field.

        Return a dictionary of Counters.

        """

        def _add_to_counter_dict(aCounter_dict, field, value):
            """Add to a dictionary of Counter the value for good Counter."""
            if field not in aCounter_dict:
                aCounter_dict[field] = Counter()
            try:
                aCounter_dict[field][value] += 1
            except TypeError, e:
                for item in value:
                    aCounter_dict[field][item] += 1

        field_values_counters_dict = {}
        for record in self.records:
            for field, field_value in record.metadata.items():
                _add_to_counter_dict(field_values_counters_dict,
                                     field, field_value)
        return field_values_counters_dict
<commit_msg>Add Metadata alignment<commit_after># -*- coding: latin-1 -*-

"""Mass Upload Libray − handling metadata."""

__authors__ = 'User:Jean-Frédéric'
import sys

sys.path.append('/home/jfk/Tuile/pywikipedia-rewrite')
sys.path.append('/home/jfk/Tuile/pywikipedia-rewrite/scripts')

import re
import sys
import pickle
import codecs
import csv
from os.path import join
import os
from collections import Counter
import UnicodeCSV
import pywikibot
import pywikibot.textlib as textlib
import data_ingestion
from data_ingestion import Photo
from PostProcessing import MetadataMapping

class MetadataRecord(Photo):

    """Represent a Record, with its associated metadata."""

    def get_field_names(self):
        """Return the field values from the records."""
        return self.metadata.keys()

    def post_process(self, mapping):
        """Post-process the MetadataRecord with the given mapping.

        For each field of the record,
        call the relevant post-processing method.

        """
        for field in self.metadata.keys():
            if field in mapping.keys():
                self.post_process_wrapper(field, mapping[field])

    def post_process_wrapper(self, field, method):
        """Wrap the post-processing methods.

        Pop the value for the given field from the record
        Apply the given method on it
        Update the record with it

        """

        def _smart_update_dict(aDict, other_dict):
            for key, value in other_dict.items():
                if key in aDict:
                    new = []
                    try:
                        aDict[key].extend(value)
                    except:
                        aDict[key] = value.append(aDict[key])
                else:
                    aDict[key] = value

        old_field_value = self.metadata.pop(field, None)
        new_field_value = method(field, old_field_value)
        _smart_update_dict(self.metadata, new_field_value)

    def to_template(self, template=u'Ingestion layout'):
        """Return the Record as a MediaWiki template."""
        return textlib.glue_template_and_params((template,
                                                 self.__dict__))

    def to_disk(self, directory):
        """Write the Record on disk in a given repository.

        Serialise the record and name it as the file title,
        do nothing if anything goes wrong.

        """
        fileName = join(directory, self.get_title() + '.dump')
        try:
            with open(fileName, 'w') as f:
                pickle.dump(self, f)
        except Exception, e:
            print "Could not pickle record %s \n %s" % (fileName, e)

class MetadataCollection(object):

    """Processing a metadata collection in various ways."""

    def __init__(self):
        """Constructor."""
        self.records = []
        self.fields = set()
        self.post_processor = MetadataMapping()

    def retrieve_metadata_from_csv(self, csv_file, delimiter=','):
        """Retrieve metadata from the given CSV file."""
        file_handler = codecs.open(csv_file, 'r', 'utf-8')
        csvReader = UnicodeCSV.unicode_csv_dictreader(file_handler,
                                                      delimiter=delimiter)
        try:
            for row in csvReader:
                self.records.append(self.handle_record(row))
        except csv.Error, e:
            sys.exit('file %s, line %d: %s' % (self.csv_file,
                                               csvReader.line_num, e))

    def retrieve_metadata_from_files(self, files_path):
        """Retrieve metadata from the files themselves in a given path."""
        print "Iterating in %s" % files_path
        assert os.path.exists(files_path)
        files = os.listdir(files_path)
        for image in files:
            self.records.append(self.handle_record(join(files_path, image)))

    def handle_record(self, items):
        """Handle a record.

        Default handling behaviour: return a new MetadataRecord
        initialised with the items as record values.

        """
        return MetadataRecord(None, items)

    def post_process_collection(self, method_mapping):
        """Call on each record its post_process method."""
        for record in self.records:
            record.post_process(method_mapping)
            yield record

    def print_metadata_of_record(self, index):
        """Print the metadata of the record.

        Print the result of get_title and of to_template.

        """
        print self.records[index].get_title()
        print self.records[index].to_template()

    def index_unique_metadata_values(self):
        """Count unique metadata values.

        Loop over the metadata collection.
        Return a dictionary of Sets (?)

        """

        def _add_to_set_dict(aDict, key, value):
            if key not in aDict:
                aDict[key] = set()
            try:
                aDict[key].add(value)
            except TypeError, e:
                aDict[key].update(value)

        sets_dict = dict()
        for record in self.records:
            for field, field_value in record.metadata.items():
                _add_to_set_dict(sets_dict, field, field_value)
        return sets_dict

    def count_metadata_values(self):
        """Count the metadata values for earch record field.

        Return a dictionary of Counters.

        """

        def _add_to_counter_dict(aCounter_dict, field, value):
            """Add to a dictionary of Counter the value for good Counter."""
            if field not in aCounter_dict:
                aCounter_dict[field] = Counter()
            try:
                aCounter_dict[field][value] += 1
            except TypeError, e:
                for item in value:
                    aCounter_dict[field][item] += 1

        field_values_counters_dict = {}
        for record in self.records:
            for field, field_value in record.metadata.items():
                _add_to_counter_dict(field_values_counters_dict,
                                     field, field_value)
        return field_values_counters_dict

    @staticmethod
    def write_dict_as_wiki(aDict, name, directory, template):
        """Write a given dictionary on disk, in template alignment format."""
        filename = os.path.join(directory, name.replace("/", ""))
        with codecs.open(filename, 'w', 'utf-8') as wikipage:
            wikipage.write("""\
{| class='wikitable' border='1'
|-
! Item
! Count
! Tag
! Categories
""")
            items = aDict.items()
            items.sort(key=lambda x: x[1], reverse=True)
            for item in items:
                try:
                    values = (template,
                            {'item': item[0], 'count': item[1],
                            'value': "", 'categories': ""})
                    table_line = ('\n' +
                                  textlib.glue_template_and_params(values))
                    wikipage.write(unicode(table_line))
                except:
                    pass
            wikipage.write("\n|}")

    def _retrieve_from_wiki(self, filename, alignment_template):
        """Retrieve the metadata mapping from a given wikipage on disk.

        Iterate over the given alignment template occurences,
        retrieve and return the mapping values.

        """

        print "retrieve_from_wiki " + filename
        wiki_file = os.path.join('wiki', filename.replace("/", ""))
        try:
            with codecs.open(wiki_file, mode='r', encoding='utf-8') as f:
                all_templates = textlib.extract_templates_and_params(f.read())
                field_mapper = dict()
                for x in all_templates:
                    if x[0] == alignment_template:
                        categories = x[1]['categories'].split(']]')[0].split(':')[-1]
                        field_mapper[x[1]['item']] = (x[1]['value'], categories)
                return field_mapper
        except Exception, e:
            print e

    def retrieve_metadata_alignments(self, fields, alignment_template):
        """Retrieve metadata alignments from disk for all given fields.

        Iterates over the given fields, determines the associate wikipage
        and calls retrieve_alignment_from_wiki on each.

        """
        print 'retrieve_metadata_alignments'
        alignments = dict()
        for field in self.fields:
            wikipage = field
            alignments[field] = self._retrieve_from_wiki(wikipage,
                                                         alignment_template)
        self.post_processor = MetadataMapping(mapper=alignments)
<|endoftext|>


<commit_before># Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
import copy
import math
from typing import List
import torch
from torch import nn

from detectron2.config import configurable
from detectron2.layers import ShapeSpec
from detectron2.structures import Boxes, RotatedBoxes
from detectron2.utils.registry import Registry

ANCHOR_GENERATOR_REGISTRY = Registry("ANCHOR_GENERATOR")
ANCHOR_GENERATOR_REGISTRY.__doc__ = """
Registry for modules that creates object detection anchors for feature maps.

The registered object will be called with `obj(cfg, input_shape)`.
"""

class BufferList(nn.Module):
    """
    Similar to nn.ParameterList, but for buffers
    """

    def __init__(self, buffers=None):
        super(BufferList, self).__init__()
        if buffers is not None:
            self.extend(buffers)

    def extend(self, buffers):
        offset = len(self)
        for i, buffer in enumerate(buffers):
            self.register_buffer(str(offset + i), buffer)
        return self

    def __len__(self):
        return len(self._buffers)

    def __iter__(self):
        return iter(self._buffers.values())

def _create_grid_offsets(size: List[int], stride: int, offset: float, device: torch.device):
    grid_height, grid_width = size
    shifts_x = torch.arange(
        offset * stride, grid_width * stride, step=stride, dtype=torch.float32, device=device
    )
    shifts_y = torch.arange(
        offset * stride, grid_height * stride, step=stride, dtype=torch.float32, device=device
    )

    shift_y, shift_x = torch.meshgrid(shifts_y, shifts_x)
    shift_x = shift_x.reshape(-1)
    shift_y = shift_y.reshape(-1)
    return shift_x, shift_y

def _broadcast_params(params, num_features, name):
    """
    If one size (or aspect ratio) is specified and there are multiple feature
    maps, we "broadcast" anchors of that single size (or aspect ratio)
    over all feature maps.

    If params is list[float], or list[list[float]] with len(params) == 1, repeat
    it num_features time.

    Returns:
        list[list[float]]: param for each feature
    """
    assert isinstance(
        params, (list, tuple)
    ), f"{name} in anchor generator has to be a list! Got {params}."
    assert len(params), f"{name} in anchor generator cannot be empty!"
    if not isinstance(params[0], (list, tuple)):  # list[float]
        return [params] * num_features
    if len(params) == 1:
        return list(params) * num_features
    assert len(params) == num_features, (
        f"Got {name} of length {len(params)} in anchor generator, "
        "but the number of input features is {num_features}!"
    )
    return params

@ANCHOR_GENERATOR_REGISTRY.register()
class DefaultAnchorGenerator(nn.Module):
    """
    Compute anchors in the standard ways described in
    "Faster R-CNN: Towards Real-Time Object Detection with Region Proposal Networks".
    """

    box_dim: int = 4
    """
    the dimension of each anchor box.
    """

    @configurable
    def __init__(self, *, sizes, aspect_ratios, strides, offset=0.5):
        """
        This interface is experimental.

        Args:
            sizes (list[list[float]] or list[float]):
                If sizes is list[list[float]], sizes[i] is the list of anchor sizes
                (i.e. sqrt of anchor area) to use for the i-th feature map.
                If sizes is list[float], the sizes are used for all feature maps.
                Anchor sizes are given in absolute lengths in units of
                the input image; they do not dynamically scale if the input image size changes.
            aspect_ratios (list[list[float]] or list[float]): list of aspect ratios
                (i.e. height / width) to use for anchors. Same "broadcast" rule for `sizes` applies.
            strides (list[int]): stride of each input feature.
            offset (float): Relative offset between the center of the first anchor and the top-left
                corner of the image. Value has to be in [0, 1).
                Recommend to use 0.5, which means half stride.
        """
        super().__init__()

        self.strides = strides
        self.num_features = len(self.strides)
        sizes = _broadcast_params(sizes, self.num_features, "sizes")
        aspect_ratios = _broadcast_params(aspect_ratios, self.num_features, "aspect_ratios")
        self.cell_anchors = self._calculate_anchors(sizes, aspect_ratios)

        self.offset = offset
        assert 0.0 <= self.offset < 1.0, self.offset

    @classmethod
    def from_config(cls, cfg, input_shape: List[ShapeSpec]):
        return {
            "sizes": cfg.MODEL.ANCHOR_GENERATOR.SIZES,
            "aspect_ratios": cfg.MODEL.ANCHOR_GENERATOR.ASPECT_RATIOS,
            "strides": [x.stride for x in input_shape],
            "offset": cfg.MODEL.ANCHOR_GENERATOR.OFFSET,
        }

    def _calculate_anchors(self, sizes, aspect_ratios):
        cell_anchors = [
            self.generate_cell_anchors(s, a).float() for s, a in zip(sizes, aspect_ratios)
        ]
        return BufferList(cell_anchors)

    @property
    def num_cell_anchors(self):
        """
        Returns:
            list[int]: Each int is the number of anchors at every pixel
                location, on that feature map.
                For example, if at every pixel we use anchors of 3 aspect
                ratios and 5 sizes, the number of anchors is 15.
                (See also ANCHOR_GENERATOR.SIZES and ANCHOR_GENERATOR.ASPECT_RATIOS in config)

                In standard RPN models, `num_cell_anchors` on every feature map is the same.
        """
        return [len(cell_anchors) for cell_anchors in self.cell_anchors]

    def _grid_anchors(self, grid_sizes: List[List[int]]):
        anchors = []
        for size, stride, base_anchors in zip(grid_sizes, self.strides, self.cell_anchors):
            shift_x, shift_y = _create_grid_offsets(size, stride, self.offset, base_anchors.device)
            shifts = torch.stack((shift_x, shift_y, shift_x, shift_y), dim=1)

            anchors.append((shifts.view(-1, 1, 4) + base_anchors.view(1, -1, 4)).reshape(-1, 4))

        return anchors

    def generate_cell_anchors(self, sizes=(32, 64, 128, 256, 512), aspect_ratios=(0.5, 1, 2)):
        """
        Generate a tensor storing canonical anchor boxes, which are all anchor
        boxes of different sizes and aspect_ratios centered at (0, 0).
        We can later build the set of anchors for a full feature map by
        shifting and tiling these tensors (see `meth:_grid_anchors`).

        Args:
            sizes (tuple[float]):
            aspect_ratios (tuple[float]]):

        Returns:
            Tensor of shape (len(sizes) * len(aspect_ratios), 4) storing anchor boxes
                in XYXY format.
        """

        # This is different from the anchor generator defined in the original Faster R-CNN
        # code or Detectron. They yield the same AP, however the old version defines cell
        # anchors in a less natural way with a shift relative to the feature grid and
        # quantization that results in slightly different sizes for different aspect ratios.
        # See also https://github.com/facebookresearch/Detectron/issues/227

        anchors = []
        for size in sizes:
            area = size ** 2.0
            for aspect_ratio in aspect_ratios:
                # s * s = w * h
                # a = h / w
                # ... some algebra ...
                # w = sqrt(s * s / a)
                # h = a * w
                w = math.sqrt(area / aspect_ratio)
                h = aspect_ratio * w
                x0, y0, x1, y1 = -w / 2.0, -h / 2.0, w / 2.0, h / 2.0
                anchors.append([x0, y0, x1, y1])
        return torch.tensor(anchors)

    def forward(self, features):
        """
        Args:
            features (list[Tensor]): list of backbone feature maps on which to generate anchors.

        Returns:
            list[list[Boxes]]: a list of #image elements. Each is a list of #feature level Boxes.
                The Boxes contains anchors of this image on the specific feature level.
        """
        num_images = len(features[0])
        grid_sizes = [feature_map.shape[-2:] for feature_map in features]
        anchors_over_all_feature_maps = self._grid_anchors(grid_sizes)

        anchors_in_image = []
        for anchors_per_feature_map in anchors_over_all_feature_maps:
            boxes = Boxes(anchors_per_feature_map)
            anchors_in_image.append(boxes)

        anchors = [copy.deepcopy(anchors_in_image) for _ in range(num_images)]
        return anchors

@ANCHOR_GENERATOR_REGISTRY.register()
class RotatedAnchorGenerator(nn.Module):
    """
    Compute rotated anchors used by Rotated RPN (RRPN), described in
    "Arbitrary-Oriented Scene Text Detection via Rotation Proposals".
    """

    box_dim: int = 5
    """
    the dimension of each anchor box.
    """

    @configurable
    def __init__(self, *, sizes, aspect_ratios, strides, angles, offset=0.5):
        """
        This interface is experimental.

        Args:
            sizes (list[list[float]] or list[float]):
                If sizes is list[list[float]], sizes[i] is the list of anchor sizes
                (i.e. sqrt of anchor area) to use for the i-th feature map.
                If sizes is list[float], the sizes are used for all feature maps.
                Anchor sizes are given in absolute lengths in units of
                the input image; they do not dynamically scale if the input image size changes.
            aspect_ratios (list[list[float]] or list[float]): list of aspect ratios
                (i.e. height / width) to use for anchors. Same "broadcast" rule for `sizes` applies.
            strides (list[int]): stride of each input feature.
            angles (list[list[float]] or list[float]): list of angles (in degrees CCW)
                to use for anchors. Same "broadcast" rule for `sizes` applies.
            offset (float): Relative offset between the center of the first anchor and the top-left
                corner of the image. Value has to be in [0, 1).
                Recommend to use 0.5, which means half stride.
        """
        super().__init__()

        self.strides = strides
        self.num_features = len(self.strides)
        sizes = _broadcast_params(sizes, self.num_features, "sizes")
        aspect_ratios = _broadcast_params(aspect_ratios, self.num_features, "aspect_ratios")
        angles = _broadcast_params(angles, self.num_features, "angles")
        self.cell_anchors = self._calculate_anchors(sizes, aspect_ratios, angles)

        self.offset = offset
        assert 0.0 <= self.offset < 1.0, self.offset

    @classmethod
    def from_config(cls, cfg, input_shape: List[ShapeSpec]):
        return {
            "sizes": cfg.MODEL.ANCHOR_GENERATOR.SIZES,
            "aspect_ratios": cfg.MODEL.ANCHOR_GENERATOR.ASPECT_RATIOS,
            "strides": [x.stride for x in input_shape],
            "offset": cfg.MODEL.ANCHOR_GENERATOR.OFFSET,
            "angles": cfg.MODEL.ANCHOR_GENERATOR.ANGLES,
        }

    def _calculate_anchors(self, sizes, aspect_ratios, angles):
        cell_anchors = [
            self.generate_cell_anchors(size, aspect_ratio, angle).float()
            for size, aspect_ratio, angle in zip(sizes, aspect_ratios, angles)
        ]
        return BufferList(cell_anchors)

    @property
    def num_cell_anchors(self):
        """
        Returns:
            list[int]: Each int is the number of anchors at every pixel
                location, on that feature map.
                For example, if at every pixel we use anchors of 3 aspect
                ratios, 2 sizes and 5 angles, the number of anchors is 30.
                (See also ANCHOR_GENERATOR.SIZES, ANCHOR_GENERATOR.ASPECT_RATIOS
                and ANCHOR_GENERATOR.ANGLES in config)

                In standard RRPN models, `num_cell_anchors` on every feature map is the same.
        """
        return [len(cell_anchors) for cell_anchors in self.cell_anchors]

    def _grid_anchors(self, grid_sizes):
        anchors = []
        for size, stride, base_anchors in zip(grid_sizes, self.strides, self.cell_anchors):
            shift_x, shift_y = _create_grid_offsets(size, stride, self.offset, base_anchors.device)
            zeros = torch.zeros_like(shift_x)
            shifts = torch.stack((shift_x, shift_y, zeros, zeros, zeros), dim=1)

            anchors.append((shifts.view(-1, 1, 5) + base_anchors.view(1, -1, 5)).reshape(-1, 5))

        return anchors

    def generate_cell_anchors(
        self,
        sizes=(32, 64, 128, 256, 512),
        aspect_ratios=(0.5, 1, 2),
        angles=(-90, -60, -30, 0, 30, 60, 90),
    ):
        """
        Generate a tensor storing canonical anchor boxes, which are all anchor
        boxes of different sizes, aspect_ratios, angles centered at (0, 0).
        We can later build the set of anchors for a full feature map by
        shifting and tiling these tensors (see `meth:_grid_anchors`).

        Args:
            sizes (tuple[float]):
            aspect_ratios (tuple[float]]):
            angles (tuple[float]]):

        Returns:
            Tensor of shape (len(sizes) * len(aspect_ratios) * len(angles), 5)
                storing anchor boxes in (x_ctr, y_ctr, w, h, angle) format.
        """
        anchors = []
        for size in sizes:
            area = size ** 2.0
            for aspect_ratio in aspect_ratios:
                # s * s = w * h
                # a = h / w
                # ... some algebra ...
                # w = sqrt(s * s / a)
                # h = a * w
                w = math.sqrt(area / aspect_ratio)
                h = aspect_ratio * w
                anchors.extend([0, 0, w, h, a] for a in angles)

        return torch.tensor(anchors)

    def forward(self, features):
        """
        Args:
            features (list[Tensor]): list of backbone feature maps on which to generate anchors.

        Returns:
            list[list[RotatedBoxes]]:
                a list of #image elements. Each is a list of #feature level RotatedBoxes.
                The RotatedBoxes contains anchors of this image on the specific feature level.
        """
        num_images = len(features[0])
        grid_sizes = [feature_map.shape[-2:] for feature_map in features]
        anchors_over_all_feature_maps = self._grid_anchors(grid_sizes)

        anchors_in_image = []
        for anchors_per_feature_map in anchors_over_all_feature_maps:
            boxes = RotatedBoxes(anchors_per_feature_map)
            anchors_in_image.append(boxes)

        anchors = [copy.deepcopy(anchors_in_image) for _ in range(num_images)]
        return anchors

def build_anchor_generator(cfg, input_shape):
    """
    Built an anchor generator from `cfg.MODEL.ANCHOR_GENERATOR.NAME`.
    """
    anchor_generator = cfg.MODEL.ANCHOR_GENERATOR.NAME
    return ANCHOR_GENERATOR_REGISTRY.get(anchor_generator)(cfg, input_shape)
<commit_msg>Minor print correction in anchor_generator.py<commit_after># Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
import copy
import math
from typing import List
import torch
from torch import nn

from detectron2.config import configurable
from detectron2.layers import ShapeSpec
from detectron2.structures import Boxes, RotatedBoxes
from detectron2.utils.registry import Registry

ANCHOR_GENERATOR_REGISTRY = Registry("ANCHOR_GENERATOR")
ANCHOR_GENERATOR_REGISTRY.__doc__ = """
Registry for modules that creates object detection anchors for feature maps.

The registered object will be called with `obj(cfg, input_shape)`.
"""

class BufferList(nn.Module):
    """
    Similar to nn.ParameterList, but for buffers
    """

    def __init__(self, buffers=None):
        super(BufferList, self).__init__()
        if buffers is not None:
            self.extend(buffers)

    def extend(self, buffers):
        offset = len(self)
        for i, buffer in enumerate(buffers):
            self.register_buffer(str(offset + i), buffer)
        return self

    def __len__(self):
        return len(self._buffers)

    def __iter__(self):
        return iter(self._buffers.values())

def _create_grid_offsets(size: List[int], stride: int, offset: float, device: torch.device):
    grid_height, grid_width = size
    shifts_x = torch.arange(
        offset * stride, grid_width * stride, step=stride, dtype=torch.float32, device=device
    )
    shifts_y = torch.arange(
        offset * stride, grid_height * stride, step=stride, dtype=torch.float32, device=device
    )

    shift_y, shift_x = torch.meshgrid(shifts_y, shifts_x)
    shift_x = shift_x.reshape(-1)
    shift_y = shift_y.reshape(-1)
    return shift_x, shift_y

def _broadcast_params(params, num_features, name):
    """
    If one size (or aspect ratio) is specified and there are multiple feature
    maps, we "broadcast" anchors of that single size (or aspect ratio)
    over all feature maps.

    If params is list[float], or list[list[float]] with len(params) == 1, repeat
    it num_features time.

    Returns:
        list[list[float]]: param for each feature
    """
    assert isinstance(
        params, (list, tuple)
    ), f"{name} in anchor generator has to be a list! Got {params}."
    assert len(params), f"{name} in anchor generator cannot be empty!"
    if not isinstance(params[0], (list, tuple)):  # list[float]
        return [params] * num_features
    if len(params) == 1:
        return list(params) * num_features
    assert len(params) == num_features, (
        f"Got {name} of length {len(params)} in anchor generator, "
        f"but the number of input features is {num_features}!"
    )
    return params

@ANCHOR_GENERATOR_REGISTRY.register()
class DefaultAnchorGenerator(nn.Module):
    """
    Compute anchors in the standard ways described in
    "Faster R-CNN: Towards Real-Time Object Detection with Region Proposal Networks".
    """

    box_dim: int = 4
    """
    the dimension of each anchor box.
    """

    @configurable
    def __init__(self, *, sizes, aspect_ratios, strides, offset=0.5):
        """
        This interface is experimental.

        Args:
            sizes (list[list[float]] or list[float]):
                If sizes is list[list[float]], sizes[i] is the list of anchor sizes
                (i.e. sqrt of anchor area) to use for the i-th feature map.
                If sizes is list[float], the sizes are used for all feature maps.
                Anchor sizes are given in absolute lengths in units of
                the input image; they do not dynamically scale if the input image size changes.
            aspect_ratios (list[list[float]] or list[float]): list of aspect ratios
                (i.e. height / width) to use for anchors. Same "broadcast" rule for `sizes` applies.
            strides (list[int]): stride of each input feature.
            offset (float): Relative offset between the center of the first anchor and the top-left
                corner of the image. Value has to be in [0, 1).
                Recommend to use 0.5, which means half stride.
        """
        super().__init__()

        self.strides = strides
        self.num_features = len(self.strides)
        sizes = _broadcast_params(sizes, self.num_features, "sizes")
        aspect_ratios = _broadcast_params(aspect_ratios, self.num_features, "aspect_ratios")
        self.cell_anchors = self._calculate_anchors(sizes, aspect_ratios)

        self.offset = offset
        assert 0.0 <= self.offset < 1.0, self.offset

    @classmethod
    def from_config(cls, cfg, input_shape: List[ShapeSpec]):
        return {
            "sizes": cfg.MODEL.ANCHOR_GENERATOR.SIZES,
            "aspect_ratios": cfg.MODEL.ANCHOR_GENERATOR.ASPECT_RATIOS,
            "strides": [x.stride for x in input_shape],
            "offset": cfg.MODEL.ANCHOR_GENERATOR.OFFSET,
        }

    def _calculate_anchors(self, sizes, aspect_ratios):
        cell_anchors = [
            self.generate_cell_anchors(s, a).float() for s, a in zip(sizes, aspect_ratios)
        ]
        return BufferList(cell_anchors)

    @property
    def num_cell_anchors(self):
        """
        Returns:
            list[int]: Each int is the number of anchors at every pixel
                location, on that feature map.
                For example, if at every pixel we use anchors of 3 aspect
                ratios and 5 sizes, the number of anchors is 15.
                (See also ANCHOR_GENERATOR.SIZES and ANCHOR_GENERATOR.ASPECT_RATIOS in config)

                In standard RPN models, `num_cell_anchors` on every feature map is the same.
        """
        return [len(cell_anchors) for cell_anchors in self.cell_anchors]

    def _grid_anchors(self, grid_sizes: List[List[int]]):
        anchors = []
        for size, stride, base_anchors in zip(grid_sizes, self.strides, self.cell_anchors):
            shift_x, shift_y = _create_grid_offsets(size, stride, self.offset, base_anchors.device)
            shifts = torch.stack((shift_x, shift_y, shift_x, shift_y), dim=1)

            anchors.append((shifts.view(-1, 1, 4) + base_anchors.view(1, -1, 4)).reshape(-1, 4))

        return anchors

    def generate_cell_anchors(self, sizes=(32, 64, 128, 256, 512), aspect_ratios=(0.5, 1, 2)):
        """
        Generate a tensor storing canonical anchor boxes, which are all anchor
        boxes of different sizes and aspect_ratios centered at (0, 0).
        We can later build the set of anchors for a full feature map by
        shifting and tiling these tensors (see `meth:_grid_anchors`).

        Args:
            sizes (tuple[float]):
            aspect_ratios (tuple[float]]):

        Returns:
            Tensor of shape (len(sizes) * len(aspect_ratios), 4) storing anchor boxes
                in XYXY format.
        """

        # This is different from the anchor generator defined in the original Faster R-CNN
        # code or Detectron. They yield the same AP, however the old version defines cell
        # anchors in a less natural way with a shift relative to the feature grid and
        # quantization that results in slightly different sizes for different aspect ratios.
        # See also https://github.com/facebookresearch/Detectron/issues/227

        anchors = []
        for size in sizes:
            area = size ** 2.0
            for aspect_ratio in aspect_ratios:
                # s * s = w * h
                # a = h / w
                # ... some algebra ...
                # w = sqrt(s * s / a)
                # h = a * w
                w = math.sqrt(area / aspect_ratio)
                h = aspect_ratio * w
                x0, y0, x1, y1 = -w / 2.0, -h / 2.0, w / 2.0, h / 2.0
                anchors.append([x0, y0, x1, y1])
        return torch.tensor(anchors)

    def forward(self, features):
        """
        Args:
            features (list[Tensor]): list of backbone feature maps on which to generate anchors.

        Returns:
            list[list[Boxes]]: a list of #image elements. Each is a list of #feature level Boxes.
                The Boxes contains anchors of this image on the specific feature level.
        """
        num_images = len(features[0])
        grid_sizes = [feature_map.shape[-2:] for feature_map in features]
        anchors_over_all_feature_maps = self._grid_anchors(grid_sizes)

        anchors_in_image = []
        for anchors_per_feature_map in anchors_over_all_feature_maps:
            boxes = Boxes(anchors_per_feature_map)
            anchors_in_image.append(boxes)

        anchors = [copy.deepcopy(anchors_in_image) for _ in range(num_images)]
        return anchors

@ANCHOR_GENERATOR_REGISTRY.register()
class RotatedAnchorGenerator(nn.Module):
    """
    Compute rotated anchors used by Rotated RPN (RRPN), described in
    "Arbitrary-Oriented Scene Text Detection via Rotation Proposals".
    """

    box_dim: int = 5
    """
    the dimension of each anchor box.
    """

    @configurable
    def __init__(self, *, sizes, aspect_ratios, strides, angles, offset=0.5):
        """
        This interface is experimental.

        Args:
            sizes (list[list[float]] or list[float]):
                If sizes is list[list[float]], sizes[i] is the list of anchor sizes
                (i.e. sqrt of anchor area) to use for the i-th feature map.
                If sizes is list[float], the sizes are used for all feature maps.
                Anchor sizes are given in absolute lengths in units of
                the input image; they do not dynamically scale if the input image size changes.
            aspect_ratios (list[list[float]] or list[float]): list of aspect ratios
                (i.e. height / width) to use for anchors. Same "broadcast" rule for `sizes` applies.
            strides (list[int]): stride of each input feature.
            angles (list[list[float]] or list[float]): list of angles (in degrees CCW)
                to use for anchors. Same "broadcast" rule for `sizes` applies.
            offset (float): Relative offset between the center of the first anchor and the top-left
                corner of the image. Value has to be in [0, 1).
                Recommend to use 0.5, which means half stride.
        """
        super().__init__()

        self.strides = strides
        self.num_features = len(self.strides)
        sizes = _broadcast_params(sizes, self.num_features, "sizes")
        aspect_ratios = _broadcast_params(aspect_ratios, self.num_features, "aspect_ratios")
        angles = _broadcast_params(angles, self.num_features, "angles")
        self.cell_anchors = self._calculate_anchors(sizes, aspect_ratios, angles)

        self.offset = offset
        assert 0.0 <= self.offset < 1.0, self.offset

    @classmethod
    def from_config(cls, cfg, input_shape: List[ShapeSpec]):
        return {
            "sizes": cfg.MODEL.ANCHOR_GENERATOR.SIZES,
            "aspect_ratios": cfg.MODEL.ANCHOR_GENERATOR.ASPECT_RATIOS,
            "strides": [x.stride for x in input_shape],
            "offset": cfg.MODEL.ANCHOR_GENERATOR.OFFSET,
            "angles": cfg.MODEL.ANCHOR_GENERATOR.ANGLES,
        }

    def _calculate_anchors(self, sizes, aspect_ratios, angles):
        cell_anchors = [
            self.generate_cell_anchors(size, aspect_ratio, angle).float()
            for size, aspect_ratio, angle in zip(sizes, aspect_ratios, angles)
        ]
        return BufferList(cell_anchors)

    @property
    def num_cell_anchors(self):
        """
        Returns:
            list[int]: Each int is the number of anchors at every pixel
                location, on that feature map.
                For example, if at every pixel we use anchors of 3 aspect
                ratios, 2 sizes and 5 angles, the number of anchors is 30.
                (See also ANCHOR_GENERATOR.SIZES, ANCHOR_GENERATOR.ASPECT_RATIOS
                and ANCHOR_GENERATOR.ANGLES in config)

                In standard RRPN models, `num_cell_anchors` on every feature map is the same.
        """
        return [len(cell_anchors) for cell_anchors in self.cell_anchors]

    def _grid_anchors(self, grid_sizes):
        anchors = []
        for size, stride, base_anchors in zip(grid_sizes, self.strides, self.cell_anchors):
            shift_x, shift_y = _create_grid_offsets(size, stride, self.offset, base_anchors.device)
            zeros = torch.zeros_like(shift_x)
            shifts = torch.stack((shift_x, shift_y, zeros, zeros, zeros), dim=1)

            anchors.append((shifts.view(-1, 1, 5) + base_anchors.view(1, -1, 5)).reshape(-1, 5))

        return anchors

    def generate_cell_anchors(
        self,
        sizes=(32, 64, 128, 256, 512),
        aspect_ratios=(0.5, 1, 2),
        angles=(-90, -60, -30, 0, 30, 60, 90),
    ):
        """
        Generate a tensor storing canonical anchor boxes, which are all anchor
        boxes of different sizes, aspect_ratios, angles centered at (0, 0).
        We can later build the set of anchors for a full feature map by
        shifting and tiling these tensors (see `meth:_grid_anchors`).

        Args:
            sizes (tuple[float]):
            aspect_ratios (tuple[float]]):
            angles (tuple[float]]):

        Returns:
            Tensor of shape (len(sizes) * len(aspect_ratios) * len(angles), 5)
                storing anchor boxes in (x_ctr, y_ctr, w, h, angle) format.
        """
        anchors = []
        for size in sizes:
            area = size ** 2.0
            for aspect_ratio in aspect_ratios:
                # s * s = w * h
                # a = h / w
                # ... some algebra ...
                # w = sqrt(s * s / a)
                # h = a * w
                w = math.sqrt(area / aspect_ratio)
                h = aspect_ratio * w
                anchors.extend([0, 0, w, h, a] for a in angles)

        return torch.tensor(anchors)

    def forward(self, features):
        """
        Args:
            features (list[Tensor]): list of backbone feature maps on which to generate anchors.

        Returns:
            list[list[RotatedBoxes]]:
                a list of #image elements. Each is a list of #feature level RotatedBoxes.
                The RotatedBoxes contains anchors of this image on the specific feature level.
        """
        num_images = len(features[0])
        grid_sizes = [feature_map.shape[-2:] for feature_map in features]
        anchors_over_all_feature_maps = self._grid_anchors(grid_sizes)

        anchors_in_image = []
        for anchors_per_feature_map in anchors_over_all_feature_maps:
            boxes = RotatedBoxes(anchors_per_feature_map)
            anchors_in_image.append(boxes)

        anchors = [copy.deepcopy(anchors_in_image) for _ in range(num_images)]
        return anchors

def build_anchor_generator(cfg, input_shape):
    """
    Built an anchor generator from `cfg.MODEL.ANCHOR_GENERATOR.NAME`.
    """
    anchor_generator = cfg.MODEL.ANCHOR_GENERATOR.NAME
    return ANCHOR_GENERATOR_REGISTRY.get(anchor_generator)(cfg, input_shape)
<|endoftext|>


<commit_before>class BaseTransport(object):
    """Base transport class."""

    def __init__(self, name, data_format_class, data_format_options,
                 handler_class, handler_options):
        self.name = name
        self._data_format = data_format_class(**data_format_options)
        self._data_type = self._data_format.data_type
        self._handler = handler_class(self, **handler_options)
        self._connected = False

    def _assert_not_connected(self):
        assert not self._connected, 'transport connection already created'

    def _assert_connected(self):
        assert self._connected, 'transport connection has not created'

    def _encode_obj(self, obj):
        return self._data_format.encode(obj)

    def _decode_data(self, data):
        return self._data_format.decode(data)

    def _call_handler_method(self, name, *args):
        getattr(self._handler, name)(*args)

    def is_connected(self):
        return self._connected

    def connect(self, url, **options):
        raise NotImplementedError

    def send_obj(self, obj, receive_obj=True, **params):
        raise NotImplementedError

    def close(self):
        raise NotImplementedError

    def join(self, timeout=None):
        raise NotImplementedError
<commit_msg>Add base transport exception<commit_after>class BaseTransport(object):
    """Base transport class."""

    def __init__(self, name, data_format_class, data_format_options,
                 handler_class, handler_options):
        self.name = name
        self._data_format = data_format_class(**data_format_options)
        self._data_type = self._data_format.data_type
        self._handler = handler_class(self, **handler_options)
        self._connected = False

    def _assert_not_connected(self):
        assert not self._connected, 'transport connection already created'

    def _assert_connected(self):
        assert self._connected, 'transport connection has not created'

    def _encode_obj(self, obj):
        return self._data_format.encode(obj)

    def _decode_data(self, data):
        return self._data_format.decode(data)

    def _call_handler_method(self, name, *args):
        getattr(self._handler, name)(*args)

    def is_connected(self):
        return self._connected

    def connect(self, url, **options):
        raise NotImplementedError

    def send_obj(self, obj, receive_obj=True, **params):
        raise NotImplementedError

    def close(self):
        raise NotImplementedError

    def join(self, timeout=None):
        raise NotImplementedError

class BaseTransportException(Exception):
    """Base transport exception."""
    pass
<|endoftext|>


<commit_before>import os
import time
import json
import cPickle
import tempfile
from testtools import TestCase, ExpectedException

import psutil

from cloudify.workflows import local
from diamond_agent.tasks import _PATHS_TO_CLEAN_UP
from diamond_agent.tasks import get_paths

class TestSingleNode(TestCase):
    def setUp(self):
        super(TestSingleNode, self).setUp()
        os.environ['MANAGEMENT_IP'] = '127.0.0.1'
        self.is_uninstallable = True
        self.env = None

    def tearDown(self):
        super(TestSingleNode, self).tearDown()
        if self.env and self.is_uninstallable:
                self.env.execute('uninstall', task_retries=0)

    # custom handler + custom collector
    def test_custom_collectors(self):
        log_path = tempfile.mktemp()
        inputs = {
            'diamond_config': {
                'prefix': tempfile.mkdtemp(prefix='cloudify-'),
                'interval': 1,
                'handlers': {
                    'test_handler.TestHandler': {
                        'path': 'handlers/test_handler.py',
                        'config': {
                            'log_path': log_path,
                        }
                    }
                }
            },
            'collectors_config': {
                'TestCollector': {
                    'path': 'collectors/test.py',
                    'config': {
                        'name': 'metric',
                        'value': 42,
                    },
                },
            },
        }
        self.env = self._create_env(inputs)
        self.env.execute('install', task_retries=0)
        if not is_created(log_path):
            self.fail('file {0} expected, but not found!'.format(log_path))

        with open(log_path, 'r') as fh:
            metric = cPickle.load(fh)
        metric_path = metric.path.split('.')

        collector_config = \
            inputs['collectors_config']['TestCollector']['config']
        self.assertEqual(collector_config['name'], metric_path[5])
        self.assertEqual(collector_config['value'], metric.value)
        self.assertEqual(self.env.name, metric_path[0])
        self.assertEqual('TestCollector', metric_path[4])

        node_instances = self.env.storage.get_node_instances()
        host_instance_id, node_id, node_instance_id = get_ids(node_instances,
                                                              'node')

        self.assertEqual(host_instance_id, metric_path[1])
        self.assertEqual(node_id, metric_path[2])
        self.assertEqual(node_instance_id, metric_path[3])

    def test_cloudify_handler_format(self):
        log_path = tempfile.mktemp()
        inputs = {
            'diamond_config': {
                'prefix': tempfile.mkdtemp(prefix='cloudify-'),
                'interval': 1,
                'handlers': {
                    'test_handler.TestHandler': {
                        'path': 'handlers/test_handler.py',
                        'config': {
                            'log_path': log_path,
                            'output_cloudify_format': True,
                        }
                    }
                }
            },
            'collectors_config': {
                'TestCollector': {
                    'path': 'collectors/test.py',
                    'config': {
                        'name': 'metric',
                        'value': 42,
                    },
                },
            },
        }
        self.env = self._create_env(inputs)
        self.env.execute('install', task_retries=0)
        if not is_created(log_path):
            self.fail('file {0} expected, but not found!'.format(log_path))

        with open(log_path, 'r') as fh:
            metric = json.loads(cPickle.load(fh))

        collector_config = \
            inputs['collectors_config']['TestCollector']['config']

        node_instances = self.env.storage.get_node_instances()
        expected_host, expected_node_name, expected_node_id = get_ids(
            node_instances, 'node')
        expected_path = collector_config['name']
        expected_metric = collector_config['value']
        expected_deployment_id = self.env.name
        expected_name = 'TestCollector'
        expected_unit = ''
        expected_type = 'GAUGE'
        expected_service = '.'.join([
            expected_deployment_id,
            expected_node_name,
            expected_node_id,
            expected_name,
            expected_path
        ])

        self.assertEqual(expected_path, metric['path'])
        self.assertEqual(expected_metric, metric['metric'])
        self.assertEqual(expected_deployment_id, metric['deployment_id'])
        self.assertEqual(expected_name, metric['name'])
        self.assertEqual(expected_unit, metric['unit'])
        self.assertEqual(expected_type, metric['type'])
        self.assertEqual(expected_host, metric['host'])
        self.assertEqual(expected_node_name, metric['node_name'])
        self.assertEqual(expected_node_id, metric['node_id'])
        self.assertEqual(expected_service, metric['service'])
        self.assertTrue(time.time() - 120 <= metric['time'] <= time.time())

    # custom handler + no collector
    # diamond should run without outputting anything
    def test_no_collectors(self):
        log_path = tempfile.mktemp()
        inputs = {
            'diamond_config': {
                'prefix': tempfile.mkdtemp(prefix='cloudify-'),
                'interval': 1,
                'handlers': {
                    'test_handler.TestHandler': {
                        'path': 'handlers/test_handler.py',
                        'config': {
                            'log_path': log_path,
                        },
                    }
                }
            },
            'collectors_config': {}
        }
        self.env = self._create_env(inputs)
        self.env.execute('install', task_retries=0)

        pid = get_pid(inputs)

        if not psutil.pid_exists(pid):
            self.fail('Diamond failed to start with empty collector list')

    def test_uninstall_workflow(self):
        inputs = {
            'diamond_config': {
                'prefix': tempfile.mkdtemp(prefix='cloudify-'),
                'interval': 1,
                'handlers': {
                    'diamond.handler.archive.ArchiveHandler': {
                        'config': {
                            'log_file': tempfile.mktemp(),
                        }
                    }
                }
            },
            'collectors_config': {},

        }
        prefix = inputs['diamond_config']['prefix']
        self.is_uninstallable = False
        self.env = self._create_env(inputs)
        self.env.execute('install', task_retries=0)
        pid_file = os.path.join(prefix, 'var', 'run', 'diamond.pid')
        with open(pid_file, 'r') as pf:
            pid = int(pf.read())

        if psutil.pid_exists(pid):
            self.env.execute('uninstall', task_retries=0)
            time.sleep(5)
        else:
            self.fail('diamond process not running')
        self.assertFalse(psutil.pid_exists(pid))

        # Check if uninstall cleans up after diamond
        for path_name in _PATHS_TO_CLEAN_UP:
            path = os.path.join(prefix, path_name)
            self.assertFalse(os.path.exists(path),
                             msg="Path exists: {}".format(path))

    def test_no_handlers(self):
        inputs = {
            'diamond_config': {
                'handlers': {},
            },
            'collectors_config': {},

        }
        self.is_uninstallable = False
        self.env = self._create_env(inputs)
        with ExpectedException(RuntimeError, ".*Empty handlers dict"):
            self.env.execute('install', task_retries=0)

    def _create_env(self, inputs):
        return local.init_env(self._blueprint_path(),
                              inputs=inputs,
                              ignored_modules=['worker_installer.tasks',
                                               'plugin_installer.tasks'])

    def _blueprint_path(self):
        return self._get_resource_path('blueprint', 'single_node.yaml')

    def _get_resource_path(self, *args):
        return os.path.join(os.path.dirname(__file__), 'resources', *args)

def collector_in_log(path, collector):
    with open(path, 'r') as fh:
        try:
            while True:
                metric = cPickle.load(fh)
                if metric.path.split('.')[3] == collector:
                    return True
        except EOFError:
            return False

def is_created(path, timeout=5):
    for _ in range(timeout):
        if os.path.isfile(path):
            return True
        time.sleep(1)
    return False

def get_ids(instances, name):
    for instance in instances:
        if instance['name'] == name:
            return instance['host_id'], instance['node_id'], instance['id']

def get_pid(config):
    pid_file = os.path.join(config['diamond_config']['prefix'],
                            'var', 'run', 'diamond.pid')

    with open(pid_file, 'r') as pf:
        pid = int(pf.read())

    return pid
<commit_msg>flake8<commit_after>import os
import time
import json
import cPickle
import tempfile
from testtools import TestCase, ExpectedException

import psutil

from cloudify.workflows import local
from diamond_agent.tasks import _PATHS_TO_CLEAN_UP

class TestSingleNode(TestCase):
    def setUp(self):
        super(TestSingleNode, self).setUp()
        os.environ['MANAGEMENT_IP'] = '127.0.0.1'
        self.is_uninstallable = True
        self.env = None

    def tearDown(self):
        super(TestSingleNode, self).tearDown()
        if self.env and self.is_uninstallable:
                self.env.execute('uninstall', task_retries=0)

    # custom handler + custom collector
    def test_custom_collectors(self):
        log_path = tempfile.mktemp()
        inputs = {
            'diamond_config': {
                'prefix': tempfile.mkdtemp(prefix='cloudify-'),
                'interval': 1,
                'handlers': {
                    'test_handler.TestHandler': {
                        'path': 'handlers/test_handler.py',
                        'config': {
                            'log_path': log_path,
                        }
                    }
                }
            },
            'collectors_config': {
                'TestCollector': {
                    'path': 'collectors/test.py',
                    'config': {
                        'name': 'metric',
                        'value': 42,
                    },
                },
            },
        }
        self.env = self._create_env(inputs)
        self.env.execute('install', task_retries=0)
        if not is_created(log_path):
            self.fail('file {0} expected, but not found!'.format(log_path))

        with open(log_path, 'r') as fh:
            metric = cPickle.load(fh)
        metric_path = metric.path.split('.')

        collector_config = \
            inputs['collectors_config']['TestCollector']['config']
        self.assertEqual(collector_config['name'], metric_path[5])
        self.assertEqual(collector_config['value'], metric.value)
        self.assertEqual(self.env.name, metric_path[0])
        self.assertEqual('TestCollector', metric_path[4])

        node_instances = self.env.storage.get_node_instances()
        host_instance_id, node_id, node_instance_id = get_ids(node_instances,
                                                              'node')

        self.assertEqual(host_instance_id, metric_path[1])
        self.assertEqual(node_id, metric_path[2])
        self.assertEqual(node_instance_id, metric_path[3])

    def test_cloudify_handler_format(self):
        log_path = tempfile.mktemp()
        inputs = {
            'diamond_config': {
                'prefix': tempfile.mkdtemp(prefix='cloudify-'),
                'interval': 1,
                'handlers': {
                    'test_handler.TestHandler': {
                        'path': 'handlers/test_handler.py',
                        'config': {
                            'log_path': log_path,
                            'output_cloudify_format': True,
                        }
                    }
                }
            },
            'collectors_config': {
                'TestCollector': {
                    'path': 'collectors/test.py',
                    'config': {
                        'name': 'metric',
                        'value': 42,
                    },
                },
            },
        }
        self.env = self._create_env(inputs)
        self.env.execute('install', task_retries=0)
        if not is_created(log_path):
            self.fail('file {0} expected, but not found!'.format(log_path))

        with open(log_path, 'r') as fh:
            metric = json.loads(cPickle.load(fh))

        collector_config = \
            inputs['collectors_config']['TestCollector']['config']

        node_instances = self.env.storage.get_node_instances()
        expected_host, expected_node_name, expected_node_id = get_ids(
            node_instances, 'node')
        expected_path = collector_config['name']
        expected_metric = collector_config['value']
        expected_deployment_id = self.env.name
        expected_name = 'TestCollector'
        expected_unit = ''
        expected_type = 'GAUGE'
        expected_service = '.'.join([
            expected_deployment_id,
            expected_node_name,
            expected_node_id,
            expected_name,
            expected_path
        ])

        self.assertEqual(expected_path, metric['path'])
        self.assertEqual(expected_metric, metric['metric'])
        self.assertEqual(expected_deployment_id, metric['deployment_id'])
        self.assertEqual(expected_name, metric['name'])
        self.assertEqual(expected_unit, metric['unit'])
        self.assertEqual(expected_type, metric['type'])
        self.assertEqual(expected_host, metric['host'])
        self.assertEqual(expected_node_name, metric['node_name'])
        self.assertEqual(expected_node_id, metric['node_id'])
        self.assertEqual(expected_service, metric['service'])
        self.assertTrue(time.time() - 120 <= metric['time'] <= time.time())

    # custom handler + no collector
    # diamond should run without outputting anything
    def test_no_collectors(self):
        log_path = tempfile.mktemp()
        inputs = {
            'diamond_config': {
                'prefix': tempfile.mkdtemp(prefix='cloudify-'),
                'interval': 1,
                'handlers': {
                    'test_handler.TestHandler': {
                        'path': 'handlers/test_handler.py',
                        'config': {
                            'log_path': log_path,
                        },
                    }
                }
            },
            'collectors_config': {}
        }
        self.env = self._create_env(inputs)
        self.env.execute('install', task_retries=0)

        pid = get_pid(inputs)

        if not psutil.pid_exists(pid):
            self.fail('Diamond failed to start with empty collector list')

    def test_uninstall_workflow(self):
        inputs = {
            'diamond_config': {
                'prefix': tempfile.mkdtemp(prefix='cloudify-'),
                'interval': 1,
                'handlers': {
                    'diamond.handler.archive.ArchiveHandler': {
                        'config': {
                            'log_file': tempfile.mktemp(),
                        }
                    }
                }
            },
            'collectors_config': {},

        }
        prefix = inputs['diamond_config']['prefix']
        self.is_uninstallable = False
        self.env = self._create_env(inputs)
        self.env.execute('install', task_retries=0)
        pid_file = os.path.join(prefix, 'var', 'run', 'diamond.pid')
        with open(pid_file, 'r') as pf:
            pid = int(pf.read())

        if psutil.pid_exists(pid):
            self.env.execute('uninstall', task_retries=0)
            time.sleep(5)
        else:
            self.fail('diamond process not running')
        self.assertFalse(psutil.pid_exists(pid))

        # Check if uninstall cleans up after diamond
        for path_name in _PATHS_TO_CLEAN_UP:
            path = os.path.join(prefix, path_name)
            self.assertFalse(os.path.exists(path),
                             msg="Path exists: {}".format(path))

    def test_no_handlers(self):
        inputs = {
            'diamond_config': {
                'handlers': {},
            },
            'collectors_config': {},

        }
        self.is_uninstallable = False
        self.env = self._create_env(inputs)
        with ExpectedException(RuntimeError, ".*Empty handlers dict"):
            self.env.execute('install', task_retries=0)

    def _create_env(self, inputs):
        return local.init_env(self._blueprint_path(),
                              inputs=inputs,
                              ignored_modules=['worker_installer.tasks',
                                               'plugin_installer.tasks'])

    def _blueprint_path(self):
        return self._get_resource_path('blueprint', 'single_node.yaml')

    def _get_resource_path(self, *args):
        return os.path.join(os.path.dirname(__file__), 'resources', *args)

def collector_in_log(path, collector):
    with open(path, 'r') as fh:
        try:
            while True:
                metric = cPickle.load(fh)
                if metric.path.split('.')[3] == collector:
                    return True
        except EOFError:
            return False

def is_created(path, timeout=5):
    for _ in range(timeout):
        if os.path.isfile(path):
            return True
        time.sleep(1)
    return False

def get_ids(instances, name):
    for instance in instances:
        if instance['name'] == name:
            return instance['host_id'], instance['node_id'], instance['id']

def get_pid(config):
    pid_file = os.path.join(config['diamond_config']['prefix'],
                            'var', 'run', 'diamond.pid')

    with open(pid_file, 'r') as pf:
        pid = int(pf.read())

    return pid
<|endoftext|>