<commit_before>#!/usr/bin/python2.7

"""
The ``Paste Class``
===================

Use it to create an object from an existing paste or other random file.

Conditions to fulfill to be able to use this class correctly:
-------------------------------------------------------------

1/ The paste need to be saved on disk somewhere (have an accessible path)
2/ The paste need to be gziped.
3/ The filepath need to look like something like this:
    /directory/source/year/month/day/paste.gz

"""

import os, magic, gzip, langid, pprint, redis, operator, string, re, json, ConfigParser
from Date import Date
from Hash import Hash

from langid.langid import LanguageIdentifier, model

from nltk.tokenize import RegexpTokenizer
from textblob import TextBlob

from lib_refine import *

clean = lambda dirty: ''.join(filter(string.printable.__contains__, dirty))
"""It filters out non-printable characters from the string it receives."""

configfile = './config.cfg'
cfg = ConfigParser.ConfigParser()
cfg.read(configfile)

class Paste(object):
    """
    This class representing a Paste as an object.
    When created, the object will have by default some "main attributes"
    such as the size or the date of the paste already calculated, whereas other
    attributes are not set and need to be "asked to be calculated" by their methods.
    It was design like this because some attributes take time to be calculated
    such as the langage or the duplicate...

    :Example:

    PST = Paste("/home/2013/ZEeGaez5.gz")

    """

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

        self.p_name = self.p_path.split('/')[-1]

        self.p_size = round(os.path.getsize(self.p_path)/1024.0,2)

        self.cache = redis.StrictRedis(
            host = cfg.get("Redis_Queues", "host"),
            port = cfg.getint("Redis_Queues", "port"),
            db = cfg.getint("Redis_Queues", "db"))

        self.p_mime = magic.from_buffer(self.get_p_content(), mime = True)

        self.p_encoding = None

        #Assuming that the paste will alway be in a day folder which is itself
        # in a month folder which is itself in a year folder.
        # /year/month/day/paste.gz
        var = self.p_path.split('/')
        self.p_date = Date(var[-4], var[-3], var[-2])

        self.p_hash_kind = None
        self.p_hash = None

        self.p_langage = None

        self.p_nb_lines = None
        self.p_max_length_line = None

        self.p_source = var[-5]

    def get_p_content(self):
        """
        Returning the content of the Paste

        :Example:

        PST.get_p_content()

        """
        r_serv = self.cache

        if r_serv.exist(self.p_path):
            paste = r_serv.get(self.p_path)
        else:
            with gzip.open(self.p_path, 'rb') as F:
                paste = r_serv.getset(self.p_path, F.read())
                r_serv.expire(self.p_path, 300)
        return paste

    def get_lines_info(self):
        """
        Returning and setting the number of lines and the maximum lenght of the
        lines of the paste.

        :return: tuple (#lines, max_length_line)

        :Example: PST.get_lines_info()

        """
        max_length_line = 0
        with gzip.open(self.p_path, 'rb') as F:
            for nb_line in enumerate(F):
                if len(nb_line[1]) >= max_length_line:
                    max_length_line = len(nb_line[1])

        self.p_nb_lines = nb_line[0]
        self.p_max_length_line = max_length_line
        return (nb_line[0], max_length_line)

    def _get_p_encoding(self):
        """
        Setting the encoding of the paste.

        :Example: PST._set_p_encoding()

        """
        try:
            return magic.Magic(mime_encoding = True).from_buffer(self.get_p_content())
        except magic.MagicException:
            pass

    def _set_p_hash_kind(self, hashkind):
        """
        Setting the hash (as an object) used for futur operation on it.

        :Example: PST._set_p_hash_kind("md5")

        .. seealso:: Hash.py Object to get the available hashs.

        """
        self.p_hash_kind = Hash(hashkind)

    def _get_p_hash(self):
        """
        Setting the hash of the paste as a kind of "uniq" identificator

        :return: hash string (md5, sha1....)

        :Example: PST._get_p_hash()

        .. note:: You need first to "declare which kind of hash you want to use
        before using this function
        .. seealso:: _set_p_hash_kind("md5")

        """
        self.p_hash = self.p_hash_kind.Calculate(self.get_p_content())
        return self.p_hash

    def _get_p_language(self):
        """
        Returning and setting the language of the paste (guessing)

        :Example: PST._get_p_language()

        ..note:: The language returned is purely guessing and may not be accurate
        if the paste doesn't contain any human dictionnary words
        ..seealso: git@github.com:saffsd/langid.py.git

        """

        identifier = LanguageIdentifier.from_modelstring(model, norm_probs=True)

        return identifier.classify(self.get_p_content())

    def _get_p_hash_kind(self):
        return self.p_hash_kind

    def _get_p_date(self):
        return self.p_date

    def _get_hash_lines(self, min = 1, start = 1, jump = 10):
        """
        Returning all the lines of the paste hashed.

        :param min: -- (int) Minimum line length to be hashed.
        :param start: -- (int) Number the line where to start.
        :param jump: -- (int) Granularity of the hashing 0 or 1 means no jumps
        (Maximum Granularity)

        :return: a set([]) of hash.

        .. warning:: Using a set here mean that this function will only return uniq hash.

        If the paste is composed with 1000 time the same line, this function will return
        just once the line.

        This choice was made to avoid a certain redundancy and useless hash checking.

        :Example: PST._get_hash_lines(1, 1, 0)

        .. note:: You need first to "declare which kind of hash you want to use
        before using this function
        .. seealso:: _set_p_hash_kind("md5")

        """
        S = set([])
        with gzip.open(self.p_path, 'rb') as F:

            for num, line in enumerate(F, start):

                if len(line) >= min:
                    if jump > 1:
                        if (num % jump) == 1 :
                            S.add(self.p_hash_kind.Calculate(line))
                    else:
                        S.add(self.p_hash_kind.Calculate(line))
        return S

    def is_duplicate(self, obj, min = 1, percent = 50, start = 1, jump = 10):
        """
        Returning the percent of similarity with another paste.
        ( Using the previous hashing method )

        :param obj: (Paste) The paste to compare with
        :param min: -- (int) Minimum line length to be hashed.
        :param percent: -- (int)
        :param start: -- (int) Number the line where to start.
        :param jump: -- (int) Granularity of the hashing 0 or 1 means no jumps
        (Maximum Granularity)

        :return: (tuple) (bool, percent)

        :Example:
        PST.is_duplicate(PST)

        >>> return (True, 100.0)

        ..seealso: _get_hash_lines()

        """

        set1 = self._get_hash_lines(min, start, jump)
        set2 = obj._get_hash_lines(min, start, jump)

        inter = set.intersection(set1, set2)

        numerator = len(inter)
        denominator = float((len(set1) + len(set2)) / 2)

        try:
            var = round((numerator / denominator)*100, 2)
        except ZeroDivisionError:
            var = 0.0

        if var >= percent:
            return True, var
        else:
            return False, var

    def save_all_attributes_redis(self, r_serv, key = None):
        """
        Saving all the attributes in a "Redis-like" Database (Redis, LevelDB)

        :param r_serv: -- Connexion to the Database.
        :param key: -- Key of an additionnal set.

        Example:
        import redis

        r_serv = redis.StrictRedis(host = 127.0.0.1, port = 6739, db = 0)

        PST = Paste("/home/Zkopkmlk.gz")
        PST.save_all_attributes_redis(r_serv)

        """
        #LevelDB Compatibility
        r_serv.hset(self.p_path, "p_name", self.p_name)
        r_serv.hset(self.p_path, "p_size", self.p_size)
        r_serv.hset(self.p_path, "p_mime", self.p_mime)
        #r_serv.hset(self.p_path, "p_encoding", self.p_encoding)
        r_serv.hset(self.p_path, "p_date", self._get_p_date())
        r_serv.hset(self.p_path, "p_hash_kind", self._get_p_hash_kind())
        r_serv.hset(self.p_path, "p_hash", self.p_hash)
        #r_serv.hset(self.p_path, "p_langage", self.p_langage)
        #r_serv.hset(self.p_path, "p_nb_lines", self.p_nb_lines)
        #r_serv.hset(self.p_path, "p_max_length_line", self.p_max_length_line)
        #r_serv.hset(self.p_path, "p_categories", self.p_categories)
        r_serv.hset(self.p_path, "p_source", self.p_source)
        if key != None:
            r_serv.sadd(key, self.p_path)
        else:
            pass

    def save_attribute_redis(self, r_serv, attr_name, value):
        """
        Save an attribute as a field
        """
        if type(value) == set:
            r_serv.hset(self.p_path, attr_name, json.dumps(list(value)))
        else:
            r_serv.hset(self.p_path, attr_name, json.dumps(value))

    def _get_from_redis(self,r_serv):
        return r_serv.hgetall(self.p_hash)

    def _get_top_words(self, sort = False):
        """
        Tokenising method: Returning a sorted list or a set of paste's words

        :param sort: Selecting the output: sorted list or a set. (set by default)

        :return: set or sorted list of tuple [(word, occurency)...]

        :Example: PST._get_top_words(False)

        """
        words = {}
        tokenizer = RegexpTokenizer('[\&\~\:\;\,\.\(\)\{\}\|\[\]\\\\/\-/\=\'\"\%\$\?\@\+\#\_\^\<\>\!\*\n\r\t\s]+',
        gaps = True,
        discard_empty = True)

        blob = TextBlob(clean(self.get_p_content()),
        tokenizer = tokenizer)

        for word in blob.tokens:
            if word in words.keys():
                num = words[word]
            else:
                num = 0

            words[word] = num + 1

        if sort:
            var = sorted(words.iteritems(), key = operator.itemgetter(1), reverse = True)
        else:
            var = words

        return var

    def _get_word(self, word):
        """
        Returning a specific word and his occurence if present in the paste

        :param word: (str) The word

        :return: (tuple) ("foo", 1337)

        """
        return [item for item in self._get_top_words() if item[0] == word]

    def get_regex(self, regex):
        """
        Returning matches with the regex given as an argument.

        :param regex: -- (str) a regex

        :return: (list)

        :Example: PST.get_regex("4[0-9]{12}(?:[0-9]{3})?")

        """
        matchs = []
        for match in re.findall(regex, self.get_p_content()):
            if match != '' and len(match) < 100:
                matchs.append(match)
        return matchs

if __name__ == "__main__":
    pp = pprint.PrettyPrinter(indent=4)
    main()
<commit_msg>Fixing last commit<commit_after>#!/usr/bin/python2.7

"""
The ``Paste Class``
===================

Use it to create an object from an existing paste or other random file.

Conditions to fulfill to be able to use this class correctly:
-------------------------------------------------------------

1/ The paste need to be saved on disk somewhere (have an accessible path)
2/ The paste need to be gziped.
3/ The filepath need to look like something like this:
    /directory/source/year/month/day/paste.gz

"""

import os, magic, gzip, langid, pprint, redis, operator, string, re, json, ConfigParser
from Date import Date
from Hash import Hash

from langid.langid import LanguageIdentifier, model

from nltk.tokenize import RegexpTokenizer
from textblob import TextBlob

from lib_refine import *

clean = lambda dirty: ''.join(filter(string.printable.__contains__, dirty))
"""It filters out non-printable characters from the string it receives."""

class Paste(object):
    """
    This class representing a Paste as an object.
    When created, the object will have by default some "main attributes"
    such as the size or the date of the paste already calculated, whereas other
    attributes are not set and need to be "asked to be calculated" by their methods.
    It was design like this because some attributes take time to be calculated
    such as the langage or the duplicate...

    :Example:

    PST = Paste("/home/2013/ZEeGaez5.gz")

    """

    def __init__(self, p_path):

        configfile = './config.cfg'
        cfg = ConfigParser.ConfigParser()
        cfg.read(configfile)

        self.p_path = p_path

        self.p_name = self.p_path.split('/')[-1]

        self.p_size = round(os.path.getsize(self.p_path)/1024.0,2)

        self.cache = redis.StrictRedis(
            host = cfg.get("Redis_Queues", "host"),
            port = cfg.getint("Redis_Queues", "port"),
            db = cfg.getint("Redis_Queues", "db"))

        self.p_mime = magic.from_buffer(self.get_p_content(), mime = True)

        self.p_encoding = None

        #Assuming that the paste will alway be in a day folder which is itself
        # in a month folder which is itself in a year folder.
        # /year/month/day/paste.gz
        var = self.p_path.split('/')
        self.p_date = Date(var[-4], var[-3], var[-2])

        self.p_hash_kind = None
        self.p_hash = None

        self.p_langage = None

        self.p_nb_lines = None
        self.p_max_length_line = None

        self.p_source = var[-5]

    def get_p_content(self):
        """
        Returning the content of the Paste

        :Example:

        PST.get_p_content()

        """
        r_serv = self.cache

        if r_serv.exist(self.p_path):
            paste = r_serv.get(self.p_path)
        else:
            with gzip.open(self.p_path, 'rb') as F:
                paste = r_serv.getset(self.p_path, F.read())
                r_serv.expire(self.p_path, 300)
        return paste

    def get_lines_info(self):
        """
        Returning and setting the number of lines and the maximum lenght of the
        lines of the paste.

        :return: tuple (#lines, max_length_line)

        :Example: PST.get_lines_info()

        """
        max_length_line = 0
        with gzip.open(self.p_path, 'rb') as F:
            for nb_line in enumerate(F):
                if len(nb_line[1]) >= max_length_line:
                    max_length_line = len(nb_line[1])

        self.p_nb_lines = nb_line[0]
        self.p_max_length_line = max_length_line
        return (nb_line[0], max_length_line)

    def _get_p_encoding(self):
        """
        Setting the encoding of the paste.

        :Example: PST._set_p_encoding()

        """
        try:
            return magic.Magic(mime_encoding = True).from_buffer(self.get_p_content())
        except magic.MagicException:
            pass

    def _set_p_hash_kind(self, hashkind):
        """
        Setting the hash (as an object) used for futur operation on it.

        :Example: PST._set_p_hash_kind("md5")

        .. seealso:: Hash.py Object to get the available hashs.

        """
        self.p_hash_kind = Hash(hashkind)

    def _get_p_hash(self):
        """
        Setting the hash of the paste as a kind of "uniq" identificator

        :return: hash string (md5, sha1....)

        :Example: PST._get_p_hash()

        .. note:: You need first to "declare which kind of hash you want to use
        before using this function
        .. seealso:: _set_p_hash_kind("md5")

        """
        self.p_hash = self.p_hash_kind.Calculate(self.get_p_content())
        return self.p_hash

    def _get_p_language(self):
        """
        Returning and setting the language of the paste (guessing)

        :Example: PST._get_p_language()

        ..note:: The language returned is purely guessing and may not be accurate
        if the paste doesn't contain any human dictionnary words
        ..seealso: git@github.com:saffsd/langid.py.git

        """

        identifier = LanguageIdentifier.from_modelstring(model, norm_probs=True)

        return identifier.classify(self.get_p_content())

    def _get_p_hash_kind(self):
        return self.p_hash_kind

    def _get_p_date(self):
        return self.p_date

    def _get_hash_lines(self, min = 1, start = 1, jump = 10):
        """
        Returning all the lines of the paste hashed.

        :param min: -- (int) Minimum line length to be hashed.
        :param start: -- (int) Number the line where to start.
        :param jump: -- (int) Granularity of the hashing 0 or 1 means no jumps
        (Maximum Granularity)

        :return: a set([]) of hash.

        .. warning:: Using a set here mean that this function will only return uniq hash.

        If the paste is composed with 1000 time the same line, this function will return
        just once the line.

        This choice was made to avoid a certain redundancy and useless hash checking.

        :Example: PST._get_hash_lines(1, 1, 0)

        .. note:: You need first to "declare which kind of hash you want to use
        before using this function
        .. seealso:: _set_p_hash_kind("md5")

        """
        S = set([])
        with gzip.open(self.p_path, 'rb') as F:

            for num, line in enumerate(F, start):

                if len(line) >= min:
                    if jump > 1:
                        if (num % jump) == 1 :
                            S.add(self.p_hash_kind.Calculate(line))
                    else:
                        S.add(self.p_hash_kind.Calculate(line))
        return S

    def is_duplicate(self, obj, min = 1, percent = 50, start = 1, jump = 10):
        """
        Returning the percent of similarity with another paste.
        ( Using the previous hashing method )

        :param obj: (Paste) The paste to compare with
        :param min: -- (int) Minimum line length to be hashed.
        :param percent: -- (int)
        :param start: -- (int) Number the line where to start.
        :param jump: -- (int) Granularity of the hashing 0 or 1 means no jumps
        (Maximum Granularity)

        :return: (tuple) (bool, percent)

        :Example:
        PST.is_duplicate(PST)

        >>> return (True, 100.0)

        ..seealso: _get_hash_lines()

        """

        set1 = self._get_hash_lines(min, start, jump)
        set2 = obj._get_hash_lines(min, start, jump)

        inter = set.intersection(set1, set2)

        numerator = len(inter)
        denominator = float((len(set1) + len(set2)) / 2)

        try:
            var = round((numerator / denominator)*100, 2)
        except ZeroDivisionError:
            var = 0.0

        if var >= percent:
            return True, var
        else:
            return False, var

    def save_all_attributes_redis(self, r_serv, key = None):
        """
        Saving all the attributes in a "Redis-like" Database (Redis, LevelDB)

        :param r_serv: -- Connexion to the Database.
        :param key: -- Key of an additionnal set.

        Example:
        import redis

        r_serv = redis.StrictRedis(host = 127.0.0.1, port = 6739, db = 0)

        PST = Paste("/home/Zkopkmlk.gz")
        PST.save_all_attributes_redis(r_serv)

        """
        #LevelDB Compatibility
        r_serv.hset(self.p_path, "p_name", self.p_name)
        r_serv.hset(self.p_path, "p_size", self.p_size)
        r_serv.hset(self.p_path, "p_mime", self.p_mime)
        #r_serv.hset(self.p_path, "p_encoding", self.p_encoding)
        r_serv.hset(self.p_path, "p_date", self._get_p_date())
        r_serv.hset(self.p_path, "p_hash_kind", self._get_p_hash_kind())
        r_serv.hset(self.p_path, "p_hash", self.p_hash)
        #r_serv.hset(self.p_path, "p_langage", self.p_langage)
        #r_serv.hset(self.p_path, "p_nb_lines", self.p_nb_lines)
        #r_serv.hset(self.p_path, "p_max_length_line", self.p_max_length_line)
        #r_serv.hset(self.p_path, "p_categories", self.p_categories)
        r_serv.hset(self.p_path, "p_source", self.p_source)
        if key != None:
            r_serv.sadd(key, self.p_path)
        else:
            pass

    def save_attribute_redis(self, r_serv, attr_name, value):
        """
        Save an attribute as a field
        """
        if type(value) == set:
            r_serv.hset(self.p_path, attr_name, json.dumps(list(value)))
        else:
            r_serv.hset(self.p_path, attr_name, json.dumps(value))

    def _get_from_redis(self,r_serv):
        return r_serv.hgetall(self.p_hash)

    def _get_top_words(self, sort = False):
        """
        Tokenising method: Returning a sorted list or a set of paste's words

        :param sort: Selecting the output: sorted list or a set. (set by default)

        :return: set or sorted list of tuple [(word, occurency)...]

        :Example: PST._get_top_words(False)

        """
        words = {}
        tokenizer = RegexpTokenizer('[\&\~\:\;\,\.\(\)\{\}\|\[\]\\\\/\-/\=\'\"\%\$\?\@\+\#\_\^\<\>\!\*\n\r\t\s]+',
        gaps = True,
        discard_empty = True)

        blob = TextBlob(clean(self.get_p_content()),
        tokenizer = tokenizer)

        for word in blob.tokens:
            if word in words.keys():
                num = words[word]
            else:
                num = 0

            words[word] = num + 1

        if sort:
            var = sorted(words.iteritems(), key = operator.itemgetter(1), reverse = True)
        else:
            var = words

        return var

    def _get_word(self, word):
        """
        Returning a specific word and his occurence if present in the paste

        :param word: (str) The word

        :return: (tuple) ("foo", 1337)

        """
        return [item for item in self._get_top_words() if item[0] == word]

    def get_regex(self, regex):
        """
        Returning matches with the regex given as an argument.

        :param regex: -- (str) a regex

        :return: (list)

        :Example: PST.get_regex("4[0-9]{12}(?:[0-9]{3})?")

        """
        matchs = []
        for match in re.findall(regex, self.get_p_content()):
            if match != '' and len(match) < 100:
                matchs.append(match)
        return matchs

if __name__ == "__main__":
    pp = pprint.PrettyPrinter(indent=4)
    main()
<|endoftext|>


<commit_before># encoding: utf-8

"""
.. codeauthor:: Tsuyoshi Hombashi <gogogo.vm@gmail.com>
"""

from __future__ import absolute_import

import pytest

import pytablewriter as ptw

class Test_WriterFactory_create_from_file_extension:

    @pytest.mark.parametrize(["value", "expected"], [
        ["valid_ext.csv", ptw.CsvTableWriter],
        ["valid_ext.CSV", ptw.CsvTableWriter],
        ["valid_ext.html", ptw.HtmlTableWriter],
        ["valid_ext.HTML", ptw.HtmlTableWriter],
        ["valid_ext.htm", ptw.HtmlTableWriter],
        ["valid_ext.HTM", ptw.HtmlTableWriter],
        ["valid_ext.json", ptw.JsonTableWriter],
        ["valid_ext.JSON", ptw.JsonTableWriter],
        ["valid_ext.md", ptw.MarkdownTableWriter],
        ["valid_ext.MD", ptw.MarkdownTableWriter],
        ["valid_ext.xls", ptw.ExcelXlsTableWriter],
        ["valid_ext.XLS", ptw.ExcelXlsTableWriter],
        ["valid_ext.xlsx", ptw.ExcelXlsxTableWriter],
        ["valid_ext.XLSX", ptw.ExcelXlsxTableWriter],
    ])
    def test_normal(self, value, expected):
        writer = ptw.TableWriterFactory.create_from_file_extension(value)

        assert isinstance(writer, expected)

    @pytest.mark.parametrize(["value", "expected"], [
        ["hoge", ptw.WriterNotFoundError],
        ["hoge.txt", ptw.WriterNotFoundError],
        [".txt", ptw.WriterNotFoundError],
    ])
    def test_exception(self, value, expected):
        with pytest.raises(expected):
            ptw.TableWriterFactory.create_from_file_extension(value)

class Test_FileLoaderFactory_create_from_format_name:

    @pytest.mark.parametrize(["file_path", "format_name", "expected"], [
        ["valid_ext.html", "csv", ptw.CsvTableWriter],
        ["invalid_ext.txt", "CSV", ptw.CsvTableWriter],
        ["valid_ext.html", "excel", ptw.ExcelXlsxTableWriter],
        ["invalid_ext.txt", "Excel", ptw.ExcelXlsxTableWriter],
        ["valid_ext.json", "html", ptw.HtmlTableWriter],
        ["invalid_ext.txt", "HTML", ptw.HtmlTableWriter],
        ["valid_ext.html", "json", ptw.JsonTableWriter],
        ["invalid_ext.txt", "JSON", ptw.JsonTableWriter],
        ["valid_ext.html", "markdown", ptw.MarkdownTableWriter],
        ["invalid_ext.txt", "Markdown", ptw.MarkdownTableWriter],
        ["valid_ext.html", "mediawiki", ptw.MediaWikiTableWriter],
        ["invalid_ext.txt", "MediaWiki", ptw.MediaWikiTableWriter],
    ])
    def test_normal(self, file_path, format_name, expected):
        writer = ptw.TableWriterFactory.create_from_format_name(format_name)

        assert isinstance(writer, expected)

    @pytest.mark.parametrize(["file_path", "format_name", "expected"], [
        ["valid_ext.csv", "not_exist_format", ptw.WriterNotFoundError],
        ["valid_ext.csv", "", ptw.WriterNotFoundError],
        ["valid_ext.csv", None, AttributeError],
    ])
    def test_exception(self, file_path, format_name, expected):
        with pytest.raises(expected):
            ptw.TableWriterFactory.create_from_format_name(format_name)
<commit_msg>Add test cases<commit_after># encoding: utf-8

"""
.. codeauthor:: Tsuyoshi Hombashi <gogogo.vm@gmail.com>
"""

from __future__ import absolute_import

import pytest

import pytablewriter as ptw

class Test_WriterFactory_create_from_file_extension:

    @pytest.mark.parametrize(["value", "expected"], [
        ["valid_ext.csv", ptw.CsvTableWriter],
        ["valid_ext.CSV", ptw.CsvTableWriter],
        [".csv", ptw.CsvTableWriter],
        ["CSV", ptw.CsvTableWriter],
        ["valid_ext.html", ptw.HtmlTableWriter],
        ["valid_ext.HTML", ptw.HtmlTableWriter],
        [".html", ptw.HtmlTableWriter],
        ["HTML", ptw.HtmlTableWriter],
        ["valid_ext.htm", ptw.HtmlTableWriter],
        ["valid_ext.HTM", ptw.HtmlTableWriter],
        [".htm", ptw.HtmlTableWriter],
        ["HTM", ptw.HtmlTableWriter],
        ["valid_ext.js", ptw.JavaScriptTableWriter],
        ["valid_ext.JS", ptw.JavaScriptTableWriter],
        [".js", ptw.JavaScriptTableWriter],
        ["JS", ptw.JavaScriptTableWriter],
        ["valid_ext.json", ptw.JsonTableWriter],
        ["valid_ext.JSON", ptw.JsonTableWriter],
        [".json", ptw.JsonTableWriter],
        ["JSON", ptw.JsonTableWriter],
        ["valid_ext.md", ptw.MarkdownTableWriter],
        ["valid_ext.MD", ptw.MarkdownTableWriter],
        [".md", ptw.MarkdownTableWriter],
        ["MD", ptw.MarkdownTableWriter],
        ["valid_ext.py", ptw.PythonCodeTableWriter],
        ["valid_ext.PY", ptw.PythonCodeTableWriter],
        [".py", ptw.PythonCodeTableWriter],
        ["PY", ptw.PythonCodeTableWriter],
        ["valid_ext.rst", ptw.RstGridTableWriter],
        ["valid_ext.RST", ptw.RstGridTableWriter],
        [".rst", ptw.RstGridTableWriter],
        ["RST", ptw.RstGridTableWriter],
        ["valid_ext.xls", ptw.ExcelXlsTableWriter],
        ["valid_ext.XLS", ptw.ExcelXlsTableWriter],
        [".xls", ptw.ExcelXlsTableWriter],
        ["XLS", ptw.ExcelXlsTableWriter],
        ["valid_ext.xlsx", ptw.ExcelXlsxTableWriter],
        ["valid_ext.XLSX", ptw.ExcelXlsxTableWriter],
        [".xlsx", ptw.ExcelXlsxTableWriter],
        ["XLSX", ptw.ExcelXlsxTableWriter],
    ])
    def test_normal(self, value, expected):
        writer = ptw.TableWriterFactory.create_from_file_extension(value)

        assert isinstance(writer, expected)

    @pytest.mark.parametrize(["value", "expected"], [
        ["hoge", ptw.WriterNotFoundError],
        ["hoge.txt", ptw.WriterNotFoundError],
        [".txt", ptw.WriterNotFoundError],
    ])
    def test_exception(self, value, expected):
        with pytest.raises(expected):
            ptw.TableWriterFactory.create_from_file_extension(value)

class Test_FileLoaderFactory_create_from_format_name:

    @pytest.mark.parametrize(["file_path", "format_name", "expected"], [
        ["valid_ext.html", "csv", ptw.CsvTableWriter],
        ["invalid_ext.txt", "CSV", ptw.CsvTableWriter],
        ["valid_ext.html", "excel", ptw.ExcelXlsxTableWriter],
        ["invalid_ext.txt", "Excel", ptw.ExcelXlsxTableWriter],
        ["valid_ext.json", "html", ptw.HtmlTableWriter],
        ["invalid_ext.txt", "HTML", ptw.HtmlTableWriter],
        ["valid_ext.html", "json", ptw.JsonTableWriter],
        ["invalid_ext.txt", "JSON", ptw.JsonTableWriter],
        ["valid_ext.html", "markdown", ptw.MarkdownTableWriter],
        ["invalid_ext.txt", "Markdown", ptw.MarkdownTableWriter],
        ["valid_ext.html", "mediawiki", ptw.MediaWikiTableWriter],
        ["invalid_ext.txt", "MediaWiki", ptw.MediaWikiTableWriter],
    ])
    def test_normal(self, file_path, format_name, expected):
        writer = ptw.TableWriterFactory.create_from_format_name(format_name)

        assert isinstance(writer, expected)

    @pytest.mark.parametrize(["file_path", "format_name", "expected"], [
        ["valid_ext.csv", "not_exist_format", ptw.WriterNotFoundError],
        ["valid_ext.csv", "", ptw.WriterNotFoundError],
        ["valid_ext.csv", None, AttributeError],
    ])
    def test_exception(self, file_path, format_name, expected):
        with pytest.raises(expected):
            ptw.TableWriterFactory.create_from_format_name(format_name)
<|endoftext|>


<commit_before>"""
Implements the wrapper for the Astropy test runner in the form of the
``./setup.py test`` distutils command.
"""

import os
import shutil
import subprocess
import sys
import tempfile

from setuptools import Command

from ..extern import six

def _fix_user_options(options):
    """
    This is for Python 2.x and 3.x compatibility.  distutils expects Command
    options to all be byte strings on Python 2 and Unicode strings on Python 3.
    """

    def to_str_or_none(x):
        if x is None:
            return None
        return str(x)

    return [tuple(to_str_or_none(x) for x in y) for y in options]

class AstropyTest(Command, object):
    description = 'Run the tests for this package'

    user_options = [
        ('package=', 'P',
         "The name of a specific package to test, e.g. 'io.fits' or 'utils'.  "
         "If nothing is specified, all default tests are run."),
        ('test-path=', 't',
         'Specify a test location by path.  If a relative path to a  .py file, '
         'it is relative to the built package, so e.g., a  leading "astropy/" '
         'is necessary.  If a relative  path to a .rst file, it is relative to '
         'the directory *below* the --docs-path directory, so a leading '
         '"docs/" is usually necessary.  May also be an absolute path.'),
        ('verbose-results', 'V',
         'Turn on verbose output from pytest.'),
        ('plugins=', 'p',
         'Plugins to enable when running pytest.'),
        ('pastebin=', 'b',
         "Enable pytest pastebin output. Either 'all' or 'failed'."),
        ('args=', 'a',
         'Additional arguments to be passed to pytest.'),
        ('remote-data', 'R', 'Run tests that download remote data.'),
        ('pep8', '8',
         'Enable PEP8 checking and disable regular tests. '
         'Requires the pytest-pep8 plugin.'),
        ('pdb', 'd',
         'Start the interactive Python debugger on errors.'),
        ('coverage', 'c',
         'Create a coverage report. Requires the coverage package.'),
        ('open-files', 'o', 'Fail if any tests leave files open.  Requires the '
         'psutil package.'),
        ('parallel=', 'j',
         'Run the tests in parallel on the specified number of '
         'CPUs.  If negative, all the cores on the machine will be '
         'used.  Requires the pytest-xdist plugin.'),
        ('docs-path=', None,
         'The path to the documentation .rst files.  If not provided, and '
         'the current directory contains a directory called "docs", that '
         'will be used.'),
        ('skip-docs', None,
         "Don't test the documentation .rst files."),
        ('repeat=', None,
         'How many times to repeat each test (can be used to check for '
         'sporadic failures).'),
        ('temp-root=', None,
         'The root directory in which to create the temporary testing files. '
         'If unspecified the system default is used (e.g. /tmp) as explained '
         'in the documentation for tempfile.mkstemp.')
    ]

    user_options = _fix_user_options(user_options)

    package_name = ''

    def initialize_options(self):
        self.package = None
        self.test_path = None
        self.verbose_results = False
        self.plugins = None
        self.pastebin = None
        self.args = None
        self.remote_data = False
        self.pep8 = False
        self.pdb = False
        self.coverage = False
        self.open_files = False
        self.parallel = 0
        self.docs_path = None
        self.skip_docs = False
        self.repeat = None
        self.temp_root = None

    def finalize_options(self):
        # Normally we would validate the options here, but that's handled in
        # run_tests
        pass

    def generate_testing_command(self):
        """
        Build a Python script to run the tests.
        """

        cmd_pre = ''  # Commands to run before the test function
        cmd_post = ''  # Commands to run after the test function

        if self.coverage:
            pre, post = self._generate_coverage_commands()
            cmd_pre += pre
            cmd_post += post

        if six.PY3:
            set_flag = "import builtins; builtins._ASTROPY_TEST_ = True"
        else:
            set_flag = "import __builtin__; __builtin__._ASTROPY_TEST_ = True"

        cmd = ('{cmd_pre}{0}; import {1.package_name}, sys; result = ('
               '{1.package_name}.test('
               'package={1.package!r}, '
               'test_path={1.test_path!r}, '
               'args={1.args!r}, '
               'plugins={1.plugins!r}, '
               'verbose={1.verbose_results!r}, '
               'pastebin={1.pastebin!r}, '
               'remote_data={1.remote_data!r}, '
               'pep8={1.pep8!r}, '
               'pdb={1.pdb!r}, '
               'open_files={1.open_files!r}, '
               'parallel={1.parallel!r}, '
               'docs_path={1.docs_path!r}, '
               'skip_docs={1.skip_docs!r}, '
               'repeat={1.repeat!r})); '
               '{cmd_post}'
               'sys.exit(result)')
        return cmd.format(set_flag, self, cmd_pre=cmd_pre, cmd_post=cmd_post)

    def run(self):
        """
        Run the tests!
        """

        # Ensure there is a doc path
        if self.docs_path is None:
            cfg_docs_dir = self.distribution.get_option_dict('build_sphinx').get('source_dir', None)

            # Some affiliated packages use this.
            # See astropy/package-template#157
            if cfg_docs_dir is not None and os.path.exists(cfg_docs_dir[1]):
                self.docs_path = os.path.abspath(cfg_docs_dir[1])

            # fall back on a default path of "docs"
            elif os.path.exists('docs'):
                self.docs_path = os.path.abspath('docs')

        # Build a testing install of the package
        self._build_temp_install()

        # Run everything in a try: finally: so that the tmp dir gets deleted.
        try:
            # Construct this modules testing command
            cmd = self.generate_testing_command()

            # Run the tests in a subprocess--this is necessary since
            # new extension modules may have appeared, and this is the
            # easiest way to set up a new environment

            # On Python 3.x prior to 3.3, the creation of .pyc files
            # is not atomic.  py.test jumps through some hoops to make
            # this work by parsing import statements and carefully
            # importing files atomically.  However, it can't detect
            # when __import__ is used, so its carefulness still fails.
            # The solution here (admittedly a bit of a hack), is to
            # turn off the generation of .pyc files altogether by
            # passing the `-B` switch to `python`.  This does mean
            # that each core will have to compile .py file to bytecode
            # itself, rather than getting lucky and borrowing the work
            # already done by another core.  Compilation is an
            # insignificant fraction of total testing time, though, so
            # it's probably not worth worrying about.
            retcode = subprocess.call([sys.executable, '-B', '-c', cmd],
                                      cwd=self.testing_path, close_fds=False)

        finally:
            # Remove temporary directory
            shutil.rmtree(self.tmp_dir)

        raise SystemExit(retcode)

    def _build_temp_install(self):
        """
        Build the package and copy the build to a temporary directory for
        the purposes of testing this avoids creating pyc and __pycache__
        directories inside the build directory
        """
        self.reinitialize_command('build', inplace=True)
        self.run_command('build')
        build_cmd = self.get_finalized_command('build')
        new_path = os.path.abspath(build_cmd.build_lib)

        # On OSX the default path for temp files is under /var, but in most
        # cases on OSX /var is actually a symlink to /private/var; ensure we
        # dereference that link, because py.test is very sensitive to relative
        # paths...
        tmp_dir = tempfile.mkdtemp(prefix=self.package_name + '-test-',
                                   dir=self.temp_root)
        self.tmp_dir = os.path.realpath(tmp_dir)
        self.testing_path = os.path.join(self.tmp_dir, os.path.basename(new_path))
        shutil.copytree(new_path, self.testing_path)

        # Ideally, docs_path is set properly in run(), but if it is still
        # not set here, do not pretend it is, otherwise bad things happen.
        # See astropy/package-template#157
        if self.docs_path is not None:
            new_docs_path = os.path.join(self.tmp_dir,
                                         os.path.basename(self.docs_path))
            shutil.copytree(self.docs_path, new_docs_path)
            self.docs_path = new_docs_path

        shutil.copy('setup.cfg', self.tmp_dir)

    def _generate_coverage_commands(self):
        """
        This method creates the post and pre commands if coverage is to be
        generated
        """
        if self.parallel != 0:
            raise ValueError(
                "--coverage can not be used with --parallel")

        try:
            import coverage  # pylint: disable=W0611
        except ImportError:
            raise ImportError(
                "--coverage requires that the coverage package is "
                "installed.")

        # Don't use get_pkg_data_filename here, because it
        # requires importing astropy.config and thus screwing
        # up coverage results for those packages.
        coveragerc = os.path.join(
            self.testing_path, self.package_name, 'tests', 'coveragerc')

        # We create a coveragerc that is specific to the version
        # of Python we're running, so that we can mark branches
        # as being specifically for Python 2 or Python 3
        with open(coveragerc, 'r') as fd:
            coveragerc_content = fd.read()
        if six.PY3:
            ignore_python_version = '2'
        else:
            ignore_python_version = '3'
        coveragerc_content = coveragerc_content.replace(
            "{ignore_python_version}", ignore_python_version).replace(
                "{packagename}", self.package_name)
        tmp_coveragerc = os.path.join(self.tmp_dir, 'coveragerc')
        with open(tmp_coveragerc, 'wb') as tmp:
            tmp.write(coveragerc_content.encode('utf-8'))

        cmd_pre = (
            'import coverage; '
            'cov = coverage.coverage(data_file="{0}", config_file="{1}"); '
            'cov.start();'.format(
                os.path.abspath(".coverage"), tmp_coveragerc))
        cmd_post = (
            'cov.stop(); '
            'from astropy.tests.helper import _save_coverage; '
            '_save_coverage(cov, result, "{0}", "{1}");'.format(
                os.path.abspath('.'), self.testing_path))

        return cmd_pre, cmd_post
<commit_msg>Excluded fallback code from coverage to see if this fixes the coverage report<commit_after>"""
Implements the wrapper for the Astropy test runner in the form of the
``./setup.py test`` distutils command.
"""

import os
import shutil
import subprocess
import sys
import tempfile

from setuptools import Command

from ..extern import six

def _fix_user_options(options):
    """
    This is for Python 2.x and 3.x compatibility.  distutils expects Command
    options to all be byte strings on Python 2 and Unicode strings on Python 3.
    """

    def to_str_or_none(x):
        if x is None:
            return None
        return str(x)

    return [tuple(to_str_or_none(x) for x in y) for y in options]

class AstropyTest(Command, object):
    description = 'Run the tests for this package'

    user_options = [
        ('package=', 'P',
         "The name of a specific package to test, e.g. 'io.fits' or 'utils'.  "
         "If nothing is specified, all default tests are run."),
        ('test-path=', 't',
         'Specify a test location by path.  If a relative path to a  .py file, '
         'it is relative to the built package, so e.g., a  leading "astropy/" '
         'is necessary.  If a relative  path to a .rst file, it is relative to '
         'the directory *below* the --docs-path directory, so a leading '
         '"docs/" is usually necessary.  May also be an absolute path.'),
        ('verbose-results', 'V',
         'Turn on verbose output from pytest.'),
        ('plugins=', 'p',
         'Plugins to enable when running pytest.'),
        ('pastebin=', 'b',
         "Enable pytest pastebin output. Either 'all' or 'failed'."),
        ('args=', 'a',
         'Additional arguments to be passed to pytest.'),
        ('remote-data', 'R', 'Run tests that download remote data.'),
        ('pep8', '8',
         'Enable PEP8 checking and disable regular tests. '
         'Requires the pytest-pep8 plugin.'),
        ('pdb', 'd',
         'Start the interactive Python debugger on errors.'),
        ('coverage', 'c',
         'Create a coverage report. Requires the coverage package.'),
        ('open-files', 'o', 'Fail if any tests leave files open.  Requires the '
         'psutil package.'),
        ('parallel=', 'j',
         'Run the tests in parallel on the specified number of '
         'CPUs.  If negative, all the cores on the machine will be '
         'used.  Requires the pytest-xdist plugin.'),
        ('docs-path=', None,
         'The path to the documentation .rst files.  If not provided, and '
         'the current directory contains a directory called "docs", that '
         'will be used.'),
        ('skip-docs', None,
         "Don't test the documentation .rst files."),
        ('repeat=', None,
         'How many times to repeat each test (can be used to check for '
         'sporadic failures).'),
        ('temp-root=', None,
         'The root directory in which to create the temporary testing files. '
         'If unspecified the system default is used (e.g. /tmp) as explained '
         'in the documentation for tempfile.mkstemp.')
    ]

    user_options = _fix_user_options(user_options)

    package_name = ''

    def initialize_options(self):
        self.package = None
        self.test_path = None
        self.verbose_results = False
        self.plugins = None
        self.pastebin = None
        self.args = None
        self.remote_data = False
        self.pep8 = False
        self.pdb = False
        self.coverage = False
        self.open_files = False
        self.parallel = 0
        self.docs_path = None
        self.skip_docs = False
        self.repeat = None
        self.temp_root = None

    def finalize_options(self):
        # Normally we would validate the options here, but that's handled in
        # run_tests
        pass

    def generate_testing_command(self):
        """
        Build a Python script to run the tests.
        """

        cmd_pre = ''  # Commands to run before the test function
        cmd_post = ''  # Commands to run after the test function

        if self.coverage:
            pre, post = self._generate_coverage_commands()
            cmd_pre += pre
            cmd_post += post

        if six.PY3:
            set_flag = "import builtins; builtins._ASTROPY_TEST_ = True"
        else:
            set_flag = "import __builtin__; __builtin__._ASTROPY_TEST_ = True"

        cmd = ('{cmd_pre}{0}; import {1.package_name}, sys; result = ('
               '{1.package_name}.test('
               'package={1.package!r}, '
               'test_path={1.test_path!r}, '
               'args={1.args!r}, '
               'plugins={1.plugins!r}, '
               'verbose={1.verbose_results!r}, '
               'pastebin={1.pastebin!r}, '
               'remote_data={1.remote_data!r}, '
               'pep8={1.pep8!r}, '
               'pdb={1.pdb!r}, '
               'open_files={1.open_files!r}, '
               'parallel={1.parallel!r}, '
               'docs_path={1.docs_path!r}, '
               'skip_docs={1.skip_docs!r}, '
               'repeat={1.repeat!r})); '
               '{cmd_post}'
               'sys.exit(result)')
        return cmd.format(set_flag, self, cmd_pre=cmd_pre, cmd_post=cmd_post)

    def run(self):
        """
        Run the tests!
        """

        # Ensure there is a doc path
        if self.docs_path is None:
            cfg_docs_dir = self.distribution.get_option_dict('build_sphinx').get('source_dir', None)

            # Some affiliated packages use this.
            # See astropy/package-template#157
            if cfg_docs_dir is not None and os.path.exists(cfg_docs_dir[1]):
                self.docs_path = os.path.abspath(cfg_docs_dir[1])

            # fall back on a default path of "docs"
            elif os.path.exists('docs'):  # pragma: no cover
                self.docs_path = os.path.abspath('docs')

        # Build a testing install of the package
        self._build_temp_install()

        # Run everything in a try: finally: so that the tmp dir gets deleted.
        try:
            # Construct this modules testing command
            cmd = self.generate_testing_command()

            # Run the tests in a subprocess--this is necessary since
            # new extension modules may have appeared, and this is the
            # easiest way to set up a new environment

            # On Python 3.x prior to 3.3, the creation of .pyc files
            # is not atomic.  py.test jumps through some hoops to make
            # this work by parsing import statements and carefully
            # importing files atomically.  However, it can't detect
            # when __import__ is used, so its carefulness still fails.
            # The solution here (admittedly a bit of a hack), is to
            # turn off the generation of .pyc files altogether by
            # passing the `-B` switch to `python`.  This does mean
            # that each core will have to compile .py file to bytecode
            # itself, rather than getting lucky and borrowing the work
            # already done by another core.  Compilation is an
            # insignificant fraction of total testing time, though, so
            # it's probably not worth worrying about.
            retcode = subprocess.call([sys.executable, '-B', '-c', cmd],
                                      cwd=self.testing_path, close_fds=False)

        finally:
            # Remove temporary directory
            shutil.rmtree(self.tmp_dir)

        raise SystemExit(retcode)

    def _build_temp_install(self):
        """
        Build the package and copy the build to a temporary directory for
        the purposes of testing this avoids creating pyc and __pycache__
        directories inside the build directory
        """
        self.reinitialize_command('build', inplace=True)
        self.run_command('build')
        build_cmd = self.get_finalized_command('build')
        new_path = os.path.abspath(build_cmd.build_lib)

        # On OSX the default path for temp files is under /var, but in most
        # cases on OSX /var is actually a symlink to /private/var; ensure we
        # dereference that link, because py.test is very sensitive to relative
        # paths...
        tmp_dir = tempfile.mkdtemp(prefix=self.package_name + '-test-',
                                   dir=self.temp_root)
        self.tmp_dir = os.path.realpath(tmp_dir)
        self.testing_path = os.path.join(self.tmp_dir, os.path.basename(new_path))
        shutil.copytree(new_path, self.testing_path)

        # Ideally, docs_path is set properly in run(), but if it is still
        # not set here, do not pretend it is, otherwise bad things happen.
        # See astropy/package-template#157
        if self.docs_path is not None:
            new_docs_path = os.path.join(self.tmp_dir,
                                         os.path.basename(self.docs_path))
            shutil.copytree(self.docs_path, new_docs_path)
            self.docs_path = new_docs_path

        shutil.copy('setup.cfg', self.tmp_dir)

    def _generate_coverage_commands(self):
        """
        This method creates the post and pre commands if coverage is to be
        generated
        """
        if self.parallel != 0:
            raise ValueError(
                "--coverage can not be used with --parallel")

        try:
            import coverage  # pylint: disable=W0611
        except ImportError:
            raise ImportError(
                "--coverage requires that the coverage package is "
                "installed.")

        # Don't use get_pkg_data_filename here, because it
        # requires importing astropy.config and thus screwing
        # up coverage results for those packages.
        coveragerc = os.path.join(
            self.testing_path, self.package_name, 'tests', 'coveragerc')

        # We create a coveragerc that is specific to the version
        # of Python we're running, so that we can mark branches
        # as being specifically for Python 2 or Python 3
        with open(coveragerc, 'r') as fd:
            coveragerc_content = fd.read()
        if six.PY3:
            ignore_python_version = '2'
        else:
            ignore_python_version = '3'
        coveragerc_content = coveragerc_content.replace(
            "{ignore_python_version}", ignore_python_version).replace(
                "{packagename}", self.package_name)
        tmp_coveragerc = os.path.join(self.tmp_dir, 'coveragerc')
        with open(tmp_coveragerc, 'wb') as tmp:
            tmp.write(coveragerc_content.encode('utf-8'))

        cmd_pre = (
            'import coverage; '
            'cov = coverage.coverage(data_file="{0}", config_file="{1}"); '
            'cov.start();'.format(
                os.path.abspath(".coverage"), tmp_coveragerc))
        cmd_post = (
            'cov.stop(); '
            'from astropy.tests.helper import _save_coverage; '
            '_save_coverage(cov, result, "{0}", "{1}");'.format(
                os.path.abspath('.'), self.testing_path))

        return cmd_pre, cmd_post
<|endoftext|>


<commit_before>#!/usr/bin/env python
# coding: utf-8

from __future__ import absolute_import, division, print_function, unicode_literals

import os, sys, unittest, collections, itertools, copy, re, subprocess, importlib, pkgutil

sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
import aegea
from aegea.util.aws import resolve_ami, SpotFleetBuilder
from aegea.util.exceptions import AegeaException

for importer, modname, is_pkg in pkgutil.iter_modules(aegea.__path__):
    importlib.import_module((aegea.__package__ or "aegea") + "." + modname)

class TestAegea(unittest.TestCase):
    def setUp(self):
        pass

    def call(self, *args, **kwargs):
        process = subprocess.Popen(stdin=kwargs.get("stdin", subprocess.PIPE), stdout=subprocess.PIPE,
                                   stderr=subprocess.PIPE, *args, **kwargs)
        out, err = process.communicate()
        exit_code = process.poll()
        out = out.decode(sys.stdin.encoding)
        err = err.decode(sys.stdin.encoding)
        if exit_code != os.EX_OK:
            print(err)
            cmd = kwargs.get("args", args[0])
            e = subprocess.CalledProcessError(exit_code, cmd, output=out)
            e.stderr = err
            raise e
        return out, err

    def call_unauthorized_ok(self, *args, **kwargs):
        message = kwargs.pop("message", "You are not authorized")
        try:
            self.call(*args, **kwargs)
        except Exception as e:
            if not (isinstance(e, subprocess.CalledProcessError) and message in e.stderr):
                raise

    def test_basic_aegea_commands(self):
        #subprocess.check_call(["aegea"])
        subprocess.check_call(["aegea", "--help"])
        subprocess.check_call(["aegea", "pricing"])
        for subcommand in aegea.parser._actions[-1].choices:
            args = []
            if subcommand in ("start", "stop", "reboot", "terminate", "console", "ssh", "grep"):
                args += ["--help"]
            elif subcommand in ("launch", "build_image"):
                args += ["--no-verify-ssh-key-pem-file", "--dry-run", "test"]
            elif subcommand == "rm":
                args += [resolve_ami()]
            elif subcommand == "secrets":
                args += ["ls"]
            elif subcommand == "pricing":
                args += ["AmazonEC2", "--json"]
            elif subcommand == "billing":
                args += ["--min-cost", "0.1"]
                if "AWS_DETAILED_BILLING_REPORTS_BUCKET" in os.environ:
                    args += ["--detailed-billing-reports-bucket", os.environ["AWS_DETAILED_BILLING_REPORTS_BUCKET"]]
            elif subcommand == "ls":
                args += ["--filter", "state=running"]
            self.call_unauthorized_ok(["aegea", subcommand] + args, message="Access Denied")

    def test_dry_run_commands(self):
        self.call_unauthorized_ok("aegea launch unittest --dry-run --no-verify-ssh-key-pem-file", shell=True)
        self.call_unauthorized_ok("aegea launch unittest --dry-run --spot --no-verify-ssh-key-pem-file", shell=True)
        self.call_unauthorized_ok("aegea launch unittest --dry-run --duration-hours 1 --no-verify-ssh-key-pem-file", shell=True)
        self.call_unauthorized_ok("aegea launch unittest --duration-hours 0.5 --min-mem-per-core-gb 6 --cores 2 --dry-run --no-verify-ssh-key-pem-file",
                                  shell=True)
        self.call_unauthorized_ok("aegea build_image i --dry-run --no-verify-ssh-key-pem-file", shell=True)

    def test_spot_fleet_builder(self):
        builder = SpotFleetBuilder(launch_spec={})
        self.assertEqual(set(spec["InstanceType"] for spec in builder.launch_specs()),
                         {"c3.large", "c4.large", "m3.large", "m4.large", "m3.medium"})
        self.assertEqual(set(spec["InstanceType"] for spec in builder.launch_specs(max_overprovision=4)),
                         {"c3.large", "c4.large", "m3.large", "m4.large", "m3.medium", "m4.xlarge", "c3.xlarge", "c4.xlarge", "m3.xlarge"})
        with self.assertRaises(AegeaException):
            builder = SpotFleetBuilder(launch_spec={}, min_cores_per_instance=16)
        builder = SpotFleetBuilder(launch_spec={}, cores=16, min_cores_per_instance=16)
        self.assertEqual(set(spec["InstanceType"] for spec in builder.launch_specs()),
                         {'c3.4xlarge', 'c4.8xlarge', 'c4.4xlarge', 'm4.10xlarge', 'c3.8xlarge', 'm4.4xlarge'})
        builder = SpotFleetBuilder(launch_spec={}, cores=16, min_cores_per_instance=16, min_mem_per_core_gb=6)
        self.assertEqual(set(spec["InstanceType"] for spec in builder.launch_specs()),
                         {'r3.4xlarge', 'r3.8xlarge', 'd2.4xlarge', 'i2.8xlarge', 'd2.8xlarge', 'i2.4xlarge'})
        builder = SpotFleetBuilder(launch_spec={}, cores=32, min_cores_per_instance=32, min_mem_per_core_gb=6)
        self.assertEqual(set(spec["InstanceType"] for spec in builder.launch_specs()),
                         {'r3.8xlarge', 'i2.8xlarge', 'd2.8xlarge'})
        # TODO: This will need updating when X1s come out
        builder = SpotFleetBuilder(launch_spec={}, cores=32, min_cores_per_instance=16, min_mem_per_core_gb=8)
        self.assertFalse(set(spec["InstanceType"] for spec in builder.launch_specs()))
        builder = SpotFleetBuilder(launch_spec={}, cores=4, gpus_per_instance=1)
        self.assertEqual(set(spec["InstanceType"] for spec in builder.launch_specs()), {"g2.2xlarge"})
        builder = SpotFleetBuilder(launch_spec={}, cores=16, gpus_per_instance=4)
        self.assertEqual(set(spec["InstanceType"] for spec in builder.launch_specs()), {"g2.8xlarge"})

if __name__ == '__main__':
    unittest.main()
<commit_msg>Tests WIP<commit_after>#!/usr/bin/env python
# coding: utf-8

from __future__ import absolute_import, division, print_function, unicode_literals

import os, sys, unittest, collections, itertools, copy, re, subprocess, importlib, pkgutil

sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
import aegea
from aegea.util.aws import resolve_ami, SpotFleetBuilder
from aegea.util.exceptions import AegeaException

for importer, modname, is_pkg in pkgutil.iter_modules(aegea.__path__):
    importlib.import_module((aegea.__package__ or "aegea") + "." + modname)

class TestAegea(unittest.TestCase):
    def setUp(self):
        pass

    def call(self, *args, **kwargs):
        cmd = kwargs.get("args", args[0])
        print('Running "{}"'.format(cmd), file=sys.stderr)
        expect = kwargs.pop("expect", [dict(exit_codes=[os.EX_OK], stdout=None, stderr=None)])
        process = subprocess.Popen(stdin=kwargs.get("stdin", subprocess.PIPE), stdout=subprocess.PIPE,
                                   stderr=subprocess.PIPE, *args, **kwargs)
        out, err = process.communicate()
        exit_code = process.poll()
        out = out.decode(sys.stdin.encoding)
        err = err.decode(sys.stdin.encoding)
        def match(exit_code, out, err, expected):
            exit_ok = exit_code in expected["exit_codes"]
            stdout_ok = re.search(expected.get("stdout") or "", out)
            stderr_ok = re.search(expected.get("stderr") or "", err)
            return exit_ok and stdout_ok and stderr_ok
        if not any(match(exit_code, out, err, exp) for exp in expect):
            print(err)
            e = subprocess.CalledProcessError(exit_code, cmd, output=out)
            e.stdout, e.stderr = out, err
            raise e
        return out, err

    def test_basic_aegea_commands(self):
        #subprocess.check_call(["aegea"])
        self.call(["aegea", "--help"])
        self.call(["aegea", "pricing"])
        for subcommand in aegea.parser._actions[-1].choices:
            args = []
            if subcommand in ("start", "stop", "reboot", "terminate", "console", "ssh", "grep"):
                args += ["--help"]
            elif subcommand in ("launch", "build_image"):
                args += ["--no-verify-ssh-key-pem-file", "--dry-run", "test"]
            elif subcommand == "rm":
                args += [resolve_ami()]
            elif subcommand == "secrets":
                args += ["ls"]
            elif subcommand == "pricing":
                args += ["AmazonEC2", "--json"]
            elif subcommand == "billing":
                args += ["--min-cost", "0.1"]
                if "AWS_DETAILED_BILLING_REPORTS_BUCKET" in os.environ:
                    args += ["--detailed-billing-reports-bucket", os.environ["AWS_DETAILED_BILLING_REPORTS_BUCKET"]]
            elif subcommand == "ls":
                args += ["--filter", "state=running"]
            unauthorized_ok = [dict(exit_codes=[os.EX_OK]), dict(exit_codes=[1], stderr="(UnauthorizedOperation|AccessDenied")]
            self.call(["aegea", subcommand] + args, expect=unauthorized_ok)

    def test_dry_run_commands(self):
        unauthorized_ok = [dict(exit_codes=[os.EX_OK]), dict(exit_codes=[1], stderr="UnauthorizedOperation")]
        self.call("aegea launch unittest --dry-run --no-verify-ssh-key-pem-file",
                  shell=True, expect=unauthorized_ok)
        self.call("aegea launch unittest --dry-run --spot --no-verify-ssh-key-pem-file",
                  shell=True, expect=unauthorized_ok)
        self.call("aegea launch unittest --dry-run --duration-hours 1 --no-verify-ssh-key-pem-file",
                  shell=True, expect=unauthorized_ok)
        self.call("aegea launch unittest --duration-hours 0.5 --min-mem-per-core-gb 6 --cores 2 --dry-run --no-verify-ssh-key-pem-file",
                  shell=True, expect=unauthorized_ok)
        self.call("aegea build_image i --dry-run --no-verify-ssh-key-pem-file",
                  shell=True, expect=unauthorized_ok)

    def test_spot_fleet_builder(self):
        builder = SpotFleetBuilder(launch_spec={})
        self.assertEqual(set(spec["InstanceType"] for spec in builder.launch_specs()),
                         {"c3.large", "c4.large", "m3.large", "m4.large", "m3.medium"})
        self.assertEqual(set(spec["InstanceType"] for spec in builder.launch_specs(max_overprovision=4)),
                         {"c3.large", "c4.large", "m3.large", "m4.large", "m3.medium", "m4.xlarge", "c3.xlarge", "c4.xlarge", "m3.xlarge"})
        with self.assertRaises(AegeaException):
            builder = SpotFleetBuilder(launch_spec={}, min_cores_per_instance=16)
        builder = SpotFleetBuilder(launch_spec={}, cores=16, min_cores_per_instance=16)
        self.assertEqual(set(spec["InstanceType"] for spec in builder.launch_specs()),
                         {'c3.4xlarge', 'c4.8xlarge', 'c4.4xlarge', 'm4.10xlarge', 'c3.8xlarge', 'm4.4xlarge'})
        builder = SpotFleetBuilder(launch_spec={}, cores=16, min_cores_per_instance=16, min_mem_per_core_gb=6)
        self.assertEqual(set(spec["InstanceType"] for spec in builder.launch_specs()),
                         {'r3.4xlarge', 'r3.8xlarge', 'd2.4xlarge', 'i2.8xlarge', 'd2.8xlarge', 'i2.4xlarge'})
        builder = SpotFleetBuilder(launch_spec={}, cores=32, min_cores_per_instance=32, min_mem_per_core_gb=6)
        self.assertEqual(set(spec["InstanceType"] for spec in builder.launch_specs()),
                         {'r3.8xlarge', 'i2.8xlarge', 'd2.8xlarge'})
        # TODO: This will need updating when X1s come out
        builder = SpotFleetBuilder(launch_spec={}, cores=32, min_cores_per_instance=16, min_mem_per_core_gb=8)
        self.assertFalse(set(spec["InstanceType"] for spec in builder.launch_specs()))
        builder = SpotFleetBuilder(launch_spec={}, cores=4, gpus_per_instance=1)
        self.assertEqual(set(spec["InstanceType"] for spec in builder.launch_specs()), {"g2.2xlarge"})
        builder = SpotFleetBuilder(launch_spec={}, cores=16, gpus_per_instance=4)
        self.assertEqual(set(spec["InstanceType"] for spec in builder.launch_specs()), {"g2.8xlarge"})

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


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

#Copyright (C) 2013 by Glenn Hickey
#
#Released under the MIT license, see LICENSE.txt
import unittest
import sys
import os
import argparse
import logging
import itertools
import copy

from teHmm.common import runShellCommand
from teHmm.common import runParallelShellCommands
from teHmm.track import TrackList
from pybedtools import BedTool, Interval

""" This script automates evaluating the hmm te model by doing training,
parsing, comparing back to truth, and summerizing the resutls in a table all
in one.  It can run the same logic on multiple input beds at once in parallel
(by using, say, a wildcard argument for inBeds. It also optionally repeats the
evaluation for subsets of the input tracks.

Independent processes are run in parallel using Python's process pool with the
maximum number of parallel processes limited by the --numProc argument
"""
def main(argv=None):
    if argv is None:
        argv = sys.argv

    parser = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
        description="Train, evalaute, then compare hmm model on input")

    parser.add_argument("trainingTracksInfo", help="Path of Tracks Info file "
                        "containing paths to genome annotation tracks used "
                        "for training")
    parser.add_argument("evalTracksInfo", help="Path of Tracks Info file "
                        "containing paths to genome annotation tracks used"
                        " for evaluation")
    parser.add_argument("outputDir", help="directory to write output")
    parser.add_argument("inBeds", nargs="*", help="list of training beds")
    parser.add_argument("--verbose", help="Print out detailed logging messages",
                        action = "store_true", default = False)
    parser.add_argument("--numProc", help="Max number of processors to use",
                        type=int, default=1)
    parser.add_argument("--allTrackCombinations", help="Rerun with all"
                        " possible combinations of tracks from the input"
                        " tracksInfo file.  Note that this number gets big"
                        " pretty fast.", action = "store_true", default= False)
    parser.add_argument("--emStates", help="By default the supervised mode"
                        " of teHmmTrain is activated.  This option overrides"
                        " that and uses the EM mode and the given number of "
                        "states instead", type=int, default=None)
    parser.add_argument("--cross", help="Do 50/50 cross validation by training"
                        " on first half input and validating on second",
                        action="store_true", default=False)
    parser.add_argument("--emFac", help="Normalization factor for weighting"
                        " emission probabilities because when there are "
                        "many tracks, the transition probabilities can get "
                        "totally lost. 0 = no normalization. 1 ="
                        " divide by number of tracks.  k = divide by number "
                        "of tracks / k", type=int, default=0)

    args = parser.parse_args()
    if args.verbose is True:
        logging.basicConfig(level=logging.DEBUG)
        verbose = " --verbose"
    else:
        logging.basicConfig(level=logging.INFO)
        verbose = ""

    if not os.path.exists(args.outputDir):
        os.makedirs(args.outputDir)

    trainingTrackList = TrackList(args.trainingTracksInfo)
    evalTrackList = TrackList(args.evalTracksInfo)
    checkTrackListCompatible(trainingTrackList, evalTrackList)

    sizeRange = (len(trainingTrackList), len(trainingTrackList) + 1)
    if args.allTrackCombinations is True:
        sizeRange = (1, len(trainingTrackList) + 1)

    if args.emStates is not None:
        trainFlags = "--numStates %d" % args.emStates
    else:
        trainFlags = "--supervised"
    trainFlags += " --emFac %d" % args.emFac

    # write out command line for posteriorty's sake
    if not os.path.exists(args.outputDir):
        os.makedirs(args.outputDir)
    cmdPath = os.path.join(args.outputDir, "teHmmBenchmark_cmd.txt")
    cmdFile = open(cmdPath, "w")
    cmdFile.write(" ".join(argv) + "\n")
    cmdFile.close()

    #todo: try to get timing for each command
    commands = []
    rows = dict()
    for pn, pList in enumerate(subsetTrackList(trainingTrackList, sizeRange)):
        if len(pList) == len(trainingTrackList):
            outDir = args.outputDir
        else:
            outDir = os.path.join(args.outputDir, "perm%d" % pn)
        if not os.path.exists(outDir):
            os.makedirs(outDir)
        trainingTrackPath = os.path.join(outDir, "training_tracks.xml")
        evalTrackPath = os.path.join(outDir, "eval_tracks.xml")
        pList.saveXML(trainingTrackPath)
        epList = TrackList()
        for track in pList:
            t = copy.deepcopy(evalTrackList.getTrackByName(track.getName()))
            epList.addTrack(t)
        epList.saveXML(evalTrackPath)

        for inBed in args.inBeds:

            base = os.path.basename(inBed)
            truthBed = inBed
            testBed = inBed
            if args.cross is True:
                truthBed = os.path.join(outDir,
                                        os.path.splitext(base)[0] +
                                        "_truth_temp.bed")
                testBed = os.path.join(outDir,
                                       os.path.splitext(base)[0] +
                                       "_test_temp.bed")
                splitBed(inBed, truthBed, testBed)

            # train
            modPath = os.path.join(outDir,
                                   os.path.splitext(base)[0] + ".mod")
            command = "teHmmTrain.py %s %s %s %s %s" % (trainingTrackPath,
                                                        truthBed,
                                                        modPath,
                                                        verbose,
                                                        trainFlags)

            # view
            viewPath = os.path.join(outDir,
                                   os.path.splitext(base)[0] + "_view.txt")
            command += " && teHmmView.py %s > %s" % (modPath, viewPath)

            # evaluate
            evalBed = os.path.join(outDir,
                                   os.path.splitext(base)[0] + "_eval.bed")
            command += " && teHmmEval.py %s %s %s --bed %s %s" % (evalTrackPath,
                                                                  modPath,
                                                                  testBed,
                                                                  evalBed,
                                                                  verbose)
            # compare
            compPath = os.path.join(outDir,
                                    os.path.splitext(base)[0] + "_comp.txt")
            command += " && compareBedStates.py %s %s > %s" % (testBed,
                                                               evalBed,
                                                               compPath)

            # make table row
            rowPath = os.path.join(outDir,
                                   os.path.splitext(base)[0] + "_row.txt")
            if inBed in rows:
                rows[inBed].append(rowPath)
            else:
                rows[inBed] = [rowPath]
            command += " && scrapeBenchmarkRow.py %s %s %s %s %s" % (
                args.trainingTracksInfo,
                trainingTrackPath,
                evalBed,
                compPath,
                rowPath)

            # remember command
            inCmdPath = os.path.join(outDir,
                                    os.path.splitext(base)[0] + "_cmd.txt")
            inCmdFile = open(inCmdPath, "w")
            inCmdFile.write(command + "\n")
            inCmdFile.close()

            commands.append(command)

    runParallelShellCommands(commands, args.numProc)
    writeTables(args.outputDir, rows)

def subsetTrackList(trackList, sizeRange):
    """ generate tracklists of all combinations of tracks in the input list
    optionally using size range to limit the different sizes tried. so, for
    example, given input list [t1, t2, t3] and sizeRange=None this
    will gneerate [t1] [t2] [t3] [t1,t2] [t1,t3] [t2,t3] [t1,t2,t3] """
    assert sizeRange[0] > 0 and sizeRange[1] <= len(trackList) + 1
    for outLen in xrange(*sizeRange):
        for perm in itertools.combinations([x for x in xrange(len(trackList))],
                                            outLen):
            permList = TrackList()
            for trackNo in perm:
                track = copy.deepcopy(trackList.getTrackByNumber(trackNo))
                permList.addTrack(track)
            yield permList

def splitBed(inBed, outBed1, outBed2):
    """ Used for cross validation option.  The first half in input bed gets
    written to outBed1 and the second half to outBed2"""
    inFile = open(inBed, "r")
    numLines = len([x for x in inFile])
    inFile.close()
    inFile = open(inBed, "r")
    cutLine = numLines / 2
    outFile1 = open(outBed1, "w")
    outFile2 = open(outBed2, "w")
    for lineNo, line in enumerate(inFile):
        if numLines == 1 or lineNo < cutLine:
            outFile1.write(line)
        if numLines == 1 or lineNo >= cutLine:
            outFile2.write(line)
    inFile.close()
    outFile1.close()
    outFile2.close()

def checkTrackListCompatible(trainingTrackList, evalTrackList):
    """ Now that we allow a different trackList to be used for training and
    eval, we need to check to make sure that everything's the same but the
    paths"""
    for track1, track2 in zip(trainingTrackList, evalTrackList):
        assert track1.getName() == track2.getName()
        assert track1.getNumber() == track2.getNumber()
        assert track1.getScale() == track2.getScale()
        assert track1.getLogScale() == track2.getLogScale()
        assert track1.getDist() == track2.getDist()

def writeTables(outDir, rows):
    """ Write CSV table for each input bed that was scraped from up from the
    output using scrapeBenchmarkRow.py """
    for inBed, rowPaths in rows.items():
        name = os.path.splitext(os.path.basename(inBed))[0]
        tablePath = os.path.join(outDir, name + "_table.csv")
        tableFile = open(tablePath, "w")
        for i, rowPath in enumerate(rowPaths):
            rowFile = open(rowPath, "r")
            rowLines = [line for line in rowFile]
            rowFile.close()
            if i == 0:
                tableFile.write(rowLines[0])
            tableFile.write(rowLines[1])
        tableFile.close()

if __name__ == "__main__":
    sys.exit(main())

<commit_msg>add foreceEmProbs option to bencmark<commit_after>#!/usr/bin/env python

#Copyright (C) 2013 by Glenn Hickey
#
#Released under the MIT license, see LICENSE.txt
import unittest
import sys
import os
import argparse
import logging
import itertools
import copy

from teHmm.common import runShellCommand
from teHmm.common import runParallelShellCommands
from teHmm.track import TrackList
from pybedtools import BedTool, Interval

""" This script automates evaluating the hmm te model by doing training,
parsing, comparing back to truth, and summerizing the resutls in a table all
in one.  It can run the same logic on multiple input beds at once in parallel
(by using, say, a wildcard argument for inBeds. It also optionally repeats the
evaluation for subsets of the input tracks.

Independent processes are run in parallel using Python's process pool with the
maximum number of parallel processes limited by the --numProc argument
"""
def main(argv=None):
    if argv is None:
        argv = sys.argv

    parser = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
        description="Train, evalaute, then compare hmm model on input")

    parser.add_argument("trainingTracksInfo", help="Path of Tracks Info file "
                        "containing paths to genome annotation tracks used "
                        "for training")
    parser.add_argument("evalTracksInfo", help="Path of Tracks Info file "
                        "containing paths to genome annotation tracks used"
                        " for evaluation")
    parser.add_argument("outputDir", help="directory to write output")
    parser.add_argument("inBeds", nargs="*", help="list of training beds")
    parser.add_argument("--verbose", help="Print out detailed logging messages",
                        action = "store_true", default = False)
    parser.add_argument("--numProc", help="Max number of processors to use",
                        type=int, default=1)
    parser.add_argument("--allTrackCombinations", help="Rerun with all"
                        " possible combinations of tracks from the input"
                        " tracksInfo file.  Note that this number gets big"
                        " pretty fast.", action = "store_true", default= False)
    parser.add_argument("--emStates", help="By default the supervised mode"
                        " of teHmmTrain is activated.  This option overrides"
                        " that and uses the EM mode and the given number of "
                        "states instead", type=int, default=None)
    parser.add_argument("--cross", help="Do 50/50 cross validation by training"
                        " on first half input and validating on second",
                        action="store_true", default=False)
    parser.add_argument("--emFac", help="Normalization factor for weighting"
                        " emission probabilities because when there are "
                        "many tracks, the transition probabilities can get "
                        "totally lost. 0 = no normalization. 1 ="
                        " divide by number of tracks.  k = divide by number "
                        "of tracks / k", type=int, default=0)
    parser.add_argument("--forceEmProbs", help="Path of text file where each "
                        "line has four entries: State Track Symbol Probability"
                        ". These "
                        "emission probabilities will override any learned"
                        " probabilities after training (unspecified "
                        "will not be set to 0 in this case. the learned values"
                        " will be kept, but normalized as needed." ,
                        default = None)

    args = parser.parse_args()
    if args.verbose is True:
        logging.basicConfig(level=logging.DEBUG)
        verbose = " --verbose"
    else:
        logging.basicConfig(level=logging.INFO)
        verbose = ""

    if not os.path.exists(args.outputDir):
        os.makedirs(args.outputDir)

    trainingTrackList = TrackList(args.trainingTracksInfo)
    evalTrackList = TrackList(args.evalTracksInfo)
    checkTrackListCompatible(trainingTrackList, evalTrackList)

    sizeRange = (len(trainingTrackList), len(trainingTrackList) + 1)
    if args.allTrackCombinations is True:
        sizeRange = (1, len(trainingTrackList) + 1)

    if args.emStates is not None:
        trainFlags = "--numStates %d" % args.emStates
    else:
        trainFlags = "--supervised"
    trainFlags += " --emFac %d" % args.emFac
    if args.forceEmProbs is not None:
        trainFlags += " --forceEmProbs %s" % args.forceEmProbs

    # write out command line for posteriorty's sake
    if not os.path.exists(args.outputDir):
        os.makedirs(args.outputDir)
    cmdPath = os.path.join(args.outputDir, "teHmmBenchmark_cmd.txt")
    cmdFile = open(cmdPath, "w")
    cmdFile.write(" ".join(argv) + "\n")
    cmdFile.close()

    #todo: try to get timing for each command
    commands = []
    rows = dict()
    for pn, pList in enumerate(subsetTrackList(trainingTrackList, sizeRange)):
        if len(pList) == len(trainingTrackList):
            outDir = args.outputDir
        else:
            outDir = os.path.join(args.outputDir, "perm%d" % pn)
        if not os.path.exists(outDir):
            os.makedirs(outDir)
        trainingTrackPath = os.path.join(outDir, "training_tracks.xml")
        evalTrackPath = os.path.join(outDir, "eval_tracks.xml")
        pList.saveXML(trainingTrackPath)
        epList = TrackList()
        for track in pList:
            t = copy.deepcopy(evalTrackList.getTrackByName(track.getName()))
            epList.addTrack(t)
        epList.saveXML(evalTrackPath)

        for inBed in args.inBeds:

            base = os.path.basename(inBed)
            truthBed = inBed
            testBed = inBed
            if args.cross is True:
                truthBed = os.path.join(outDir,
                                        os.path.splitext(base)[0] +
                                        "_truth_temp.bed")
                testBed = os.path.join(outDir,
                                       os.path.splitext(base)[0] +
                                       "_test_temp.bed")
                splitBed(inBed, truthBed, testBed)

            # train
            modPath = os.path.join(outDir,
                                   os.path.splitext(base)[0] + ".mod")
            command = "teHmmTrain.py %s %s %s %s %s" % (trainingTrackPath,
                                                        truthBed,
                                                        modPath,
                                                        verbose,
                                                        trainFlags)

            # view
            viewPath = os.path.join(outDir,
                                   os.path.splitext(base)[0] + "_view.txt")
            command += " && teHmmView.py %s > %s" % (modPath, viewPath)

            # evaluate
            evalBed = os.path.join(outDir,
                                   os.path.splitext(base)[0] + "_eval.bed")
            command += " && teHmmEval.py %s %s %s --bed %s %s" % (evalTrackPath,
                                                                  modPath,
                                                                  testBed,
                                                                  evalBed,
                                                                  verbose)
            # compare
            compPath = os.path.join(outDir,
                                    os.path.splitext(base)[0] + "_comp.txt")
            command += " && compareBedStates.py %s %s > %s" % (testBed,
                                                               evalBed,
                                                               compPath)

            # make table row
            rowPath = os.path.join(outDir,
                                   os.path.splitext(base)[0] + "_row.txt")
            if inBed in rows:
                rows[inBed].append(rowPath)
            else:
                rows[inBed] = [rowPath]
            command += " && scrapeBenchmarkRow.py %s %s %s %s %s" % (
                args.trainingTracksInfo,
                trainingTrackPath,
                evalBed,
                compPath,
                rowPath)

            # remember command
            inCmdPath = os.path.join(outDir,
                                    os.path.splitext(base)[0] + "_cmd.txt")
            inCmdFile = open(inCmdPath, "w")
            inCmdFile.write(command + "\n")
            inCmdFile.close()

            commands.append(command)

    runParallelShellCommands(commands, args.numProc)
    writeTables(args.outputDir, rows)

def subsetTrackList(trackList, sizeRange):
    """ generate tracklists of all combinations of tracks in the input list
    optionally using size range to limit the different sizes tried. so, for
    example, given input list [t1, t2, t3] and sizeRange=None this
    will gneerate [t1] [t2] [t3] [t1,t2] [t1,t3] [t2,t3] [t1,t2,t3] """
    assert sizeRange[0] > 0 and sizeRange[1] <= len(trackList) + 1
    for outLen in xrange(*sizeRange):
        for perm in itertools.combinations([x for x in xrange(len(trackList))],
                                            outLen):
            permList = TrackList()
            for trackNo in perm:
                track = copy.deepcopy(trackList.getTrackByNumber(trackNo))
                permList.addTrack(track)
            yield permList

def splitBed(inBed, outBed1, outBed2):
    """ Used for cross validation option.  The first half in input bed gets
    written to outBed1 and the second half to outBed2"""
    inFile = open(inBed, "r")
    numLines = len([x for x in inFile])
    inFile.close()
    inFile = open(inBed, "r")
    cutLine = numLines / 2
    outFile1 = open(outBed1, "w")
    outFile2 = open(outBed2, "w")
    for lineNo, line in enumerate(inFile):
        if numLines == 1 or lineNo < cutLine:
            outFile1.write(line)
        if numLines == 1 or lineNo >= cutLine:
            outFile2.write(line)
    inFile.close()
    outFile1.close()
    outFile2.close()

def checkTrackListCompatible(trainingTrackList, evalTrackList):
    """ Now that we allow a different trackList to be used for training and
    eval, we need to check to make sure that everything's the same but the
    paths"""
    for track1, track2 in zip(trainingTrackList, evalTrackList):
        assert track1.getName() == track2.getName()
        assert track1.getNumber() == track2.getNumber()
        assert track1.getScale() == track2.getScale()
        assert track1.getLogScale() == track2.getLogScale()
        assert track1.getDist() == track2.getDist()

def writeTables(outDir, rows):
    """ Write CSV table for each input bed that was scraped from up from the
    output using scrapeBenchmarkRow.py """
    for inBed, rowPaths in rows.items():
        name = os.path.splitext(os.path.basename(inBed))[0]
        tablePath = os.path.join(outDir, name + "_table.csv")
        tableFile = open(tablePath, "w")
        for i, rowPath in enumerate(rowPaths):
            rowFile = open(rowPath, "r")
            rowLines = [line for line in rowFile]
            rowFile.close()
            if i == 0:
                tableFile.write(rowLines[0])
            tableFile.write(rowLines[1])
        tableFile.close()

if __name__ == "__main__":
    sys.exit(main())

<|endoftext|>