<commit_before>import sys
import json
import yaml
from mcb.outputs import OutputPipe

class Config(object):

  def __init__(self):
    self.services = {}
    self.outputs = {}

    self.filepath = None

  def buildPlugin(self, name, conf):
    moduleName = name[:name.rfind('.')]
    className = name[name.rfind('.')+1:]

    mod = __import__(moduleName, fromlist=[className])
    clas = getattr(mod, className)

    instance = clas()
    instance.setConfig(conf)

    return instance

  def getServices(self):
    return [self.buildPlugin(name, conf) for name, conf in self.services.items()]

  def importServices(self, services):
    new = {}

    for service in services:
      name = service.getClassName()
      new[name] = service.getConfig()

    self.services = new

  def getOutputs(self):
    return [self.buildPlugin(name, conf) for (name, conf) in self.outputs.items()]

  def getOutputPipe(self):
    return OutputPipe(self.getOutputs())

  def importOutputs(self, outputs):
    new = {}

    for output in outputs:
      name = output.getClassName()
      new[name] = output.getConfig()

    self.outputs = new

  def getAsDict(self):
    return {
      'services': self.services,
      'outputs': self.outputs
    }

  def save(self):
    self.toFile()

  def fromFile(self, path, format='yaml'):
    self.filepath = path

    f = open(path, 'r')
    data = f.read()
    f.close()

    conf = None

    if format == 'yaml':
      conf = yaml.load(data)
    elif format == 'json':
      conf = json.load(data)
    else:
      raise Exception('Invalid format: ' + format)

    if not 'services' in conf:
      raise Exception('Invalid config: no services data')
    if not 'outputs' in conf:
      raise Exception('Invalid config: no outputs data')

    self.services = conf['services']
    self.outputs = conf['outputs']

  def toFile(self, path=None, format='yaml'):
    if not path: path = self.filepath

    conf = self.getAsDict()

    data = None
    if format == 'yaml':
      data = yaml.dump(conf, default_flow_style=False)
    elif format == 'json':
      data = json.dumps(conf)
    else:
      raise Exception('Invalid format: ' + format)

    f = open(path, 'w')
    f.write(data)
    f.close()
<commit_msg>minor refactoring<commit_after>import sys
import json
import yaml
from mcb.outputs import OutputPipe

class Config(object):

  def __init__(self):
    self.services = {}
    self.outputs = {}

    self.filepath = None

  def buildPlugin(self, name, conf):
    moduleName = name[:name.rfind('.')]
    className = name[name.rfind('.')+1:]

    mod = __import__(moduleName, fromlist=[className])
    clas = getattr(mod, className)

    instance = clas()
    instance.setConfig(conf)

    return instance

  def getServices(self):
    return [self.buildPlugin(name, conf) for name, conf in self.services.items()]

  def importServices(self, services):
    new = {}

    for service in services:
      name = service.getClassName()
      new[name] = service.getConfig()

    self.services = new

  def getOutputs(self):
    return [self.buildPlugin(name, conf) for (name, conf) in self.outputs.items()]

  def getOutputPipe(self):
    return OutputPipe(self.getOutputs())

  def importOutputs(self, outputs):
    new = {}

    for output in outputs:
      name = output.getClassName()
      new[name] = output.getConfig()

    self.outputs = new

  def getAsDict(self):
    return {
      'services': self.services,
      'outputs': self.outputs
    }

  def save(self):
    if self.filepath: self.toFile()

  def fromDict(self, conf):
    if not 'services' in conf:
      raise Exception('Invalid config: no services data')
    if not 'outputs' in conf:
      raise Exception('Invalid config: no outputs data')

    self.services = conf['services']
    self.outputs = conf['outputs']

  def fromFile(self, path, format='yaml'):
    self.filepath = path

    f = open(path, 'r')
    data = f.read()
    f.close()

    conf = None

    if format == 'yaml':
      conf = yaml.load(data)
    elif format == 'json':
      conf = json.load(data)
    else:
      raise Exception('Invalid format: ' + format)

    self.fromDict(conf)

  def toFile(self, path=None, format='yaml'):
    if not path: path = self.filepath

    conf = self.getAsDict()

    data = None
    if format == 'yaml':
      data = yaml.dump(conf, default_flow_style=False)
    elif format == 'json':
      data = json.dumps(conf)
    else:
      raise Exception('Invalid format: ' + format)

    f = open(path, 'w')
    f.write(data)
    f.close()
<|endoftext|>


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

import math
import os
import unittest

from openquake import shapes
from openquake import test
from openquake import job
from openquake.job import Job, EXPOSURE, INPUT_REGION
from openquake.job.mixins import Mixin
from openquake.risk.job import RiskJobMixin
from openquake.risk.job.probabilistic import ProbabilisticEventMixin

CONFIG_FILE = "config.gem"
CONFIG_WITH_INCLUDES = "config_with_includes.gem"

TEST_JOB_FILE = test.smoketest_file('simplecase/config.gem')

SITE = shapes.Site(1.0, 1.0)
EXPOSURE_TEST_FILE = "ExposurePortfolioFile-test.xml"
REGION_EXPOSURE_TEST_FILE = "ExposurePortfolioFile-test.region"
BLOCK_SPLIT_TEST_FILE = "block_split.gem"
REGION_TEST_FILE = "small.region"

class JobTestCase(unittest.TestCase):
    def setUp(self):
        self.job = Job.from_file(test.test_file(CONFIG_FILE))
        self.job_with_includes = Job.from_file(test.test_file(CONFIG_WITH_INCLUDES))

    def tearDown(self):
        [os.remove(f) for f in os.listdir(os.curdir) if f.count("-super.gem")]

    def test_job_writes_to_super_config(self):
        for job in [self.job, self.job_with_includes]:
            config_file = "%s-super.gem" % job.job_id
            self.assertTrue(os.listdir(os.curdir).count(config_file))

    def test_configuration_is_the_same_no_matter_which_way_its_provided(self):
        self.assertEqual(self.job.params, self.job_with_includes.params)

    def test_job_mixes_in_properly(self):
        with Mixin(self.job, RiskJobMixin, key="risk"):
            self.assertTrue(RiskJobMixin in self.job.__class__.__bases__)
            self.assertTrue(ProbabilisticEventMixin in self.job.__class__.__bases__)

        with Mixin(self.job, ProbabilisticEventMixin):
            self.assertTrue(ProbabilisticEventMixin in self.job.__class__.__bases__)

    def test_job_runs_with_a_good_config(self):
        job = Job.from_file(TEST_JOB_FILE)
        #self.assertTrue(job.launch())

    def test_a_job_has_an_identifier(self):
        self.assertEqual(1, Job({}, 1).id)

    def test_can_store_and_read_jobs_from_kvs(self):
        self.job = Job.from_file(os.path.join(test.DATA_DIR, CONFIG_FILE))
        self.assertEqual(self.job, Job.from_kvs(self.job.id))

    def test_prepares_blocks_using_the_exposure(self):
        a_job = Job({EXPOSURE: os.path.join(test.DATA_DIR, EXPOSURE_TEST_FILE)})
        a_job._partition()
        blocks_keys = a_job.blocks_keys

        expected_block = job.Block((shapes.Site(9.15000, 45.16667),
                shapes.Site(9.15333, 45.12200), shapes.Site(9.14777, 45.17999),
                shapes.Site(9.15765, 45.13005), shapes.Site(9.15934, 45.13300),
                shapes.Site(9.15876, 45.13805)))

        self.assertEqual(1, len(blocks_keys))
        self.assertEqual(expected_block, job.Block.from_kvs(blocks_keys[0]))

    def test_prepares_blocks_using_the_exposure_and_filtering(self):
        a_job = Job({EXPOSURE: test.test_file(EXPOSURE_TEST_FILE),
                     INPUT_REGION: test.test_file(REGION_EXPOSURE_TEST_FILE)})
        a_job._partition()
        blocks_keys = a_job.blocks_keys

        print blocks_keys
        print job.Block.from_kvs(blocks_keys[0]).sites

        expected_block = job.Block((shapes.Site(9.15, 45.16667),
                                    shapes.Site(9.15333, 45.122),
                                    shapes.Site(9.14777, 45.17999),
                                    shapes.Site(9.15765, 45.13005),
                                    shapes.Site(9.15934, 45.133),
                                    shapes.Site(9.15876, 45.13805)))

        self.assertEqual(1, len(blocks_keys))
        self.assertEqual(expected_block, job.Block.from_kvs(blocks_keys[0]))

    def test_prepares_blocks_using_the_input_region(self):
        """ This test might be currently catastrophically retarded. If it is
        blame Lars.
        """

        block_path = test.test_file(BLOCK_SPLIT_TEST_FILE)

        print "In open job"
        a_job = Job.from_file(block_path)

        verts = [float(x) for x in a_job.params['REGION_VERTEX'].split(",")]
        # Flips lon and lat, and builds a list of coord tuples
        coords = zip(verts[1::2], verts[::2])
        expected = shapes.RegionConstraint.from_coordinates(coords)
        expected.cell_size = float(a_job.params['REGION_GRID_SPACING'])

        expected_sites = []
        for site in expected:
            print site
            expected_sites.append(site)

        a_job._partition()
        blocks_keys = a_job.blocks_keys

        self.assertEqual(1, len(blocks_keys))
        self.assertEqual(job.Block(expected_sites),
                         job.Block.from_kvs(blocks_keys[0]))

    def test_with_no_partition_we_just_process_a_single_block(self):
        job.SITES_PER_BLOCK = 1

        # test exposure has 6 assets
        a_job = Job({EXPOSURE: os.path.join(
                test.DATA_DIR, EXPOSURE_TEST_FILE)})

        a_job._partition()
        blocks_keys = a_job.blocks_keys

        # but we have 1 block instead of 6
        self.assertEqual(1, len(blocks_keys))

class BlockTestCase(unittest.TestCase):

    def test_a_block_has_a_unique_id(self):
        self.assertTrue(job.Block(()).id)
        self.assertTrue(job.Block(()).id != job.Block(()).id)

    def test_can_serialize_a_block_into_kvs(self):
        block = job.Block((SITE, SITE))
        block.to_kvs()

        self.assertEqual(block, job.Block.from_kvs(block.id))

class BlockSplitterTestCase(unittest.TestCase):

    def setUp(self):
        self.splitter = None

    def test_an_empty_set_produces_no_blocks(self):
        self.splitter = job.BlockSplitter(())
        self._assert_number_of_blocks_is(0)

    def test_splits_the_set_into_a_single_block(self):
        self.splitter = job.BlockSplitter((SITE,), 3)
        self._assert_number_of_blocks_is(1)

        self.splitter = job.BlockSplitter((SITE, SITE), 3)
        self._assert_number_of_blocks_is(1)

        self.splitter = job.BlockSplitter((SITE, SITE, SITE), 3)
        self._assert_number_of_blocks_is(1)

    def test_splits_the_set_into_multiple_blocks(self):
        self.splitter = job.BlockSplitter((SITE, SITE), 1)
        self._assert_number_of_blocks_is(2)

        self.splitter = job.BlockSplitter((SITE, SITE, SITE), 2)
        self._assert_number_of_blocks_is(2)

    def test_generates_the_correct_blocks(self):
        self.splitter = job.BlockSplitter((SITE, SITE, SITE), 3)
        expected_blocks = (job.Block((SITE, SITE, SITE)),)
        self._assert_blocks_are(expected_blocks)

        self.splitter = job.BlockSplitter((SITE, SITE, SITE), 2)
        expected_blocks = (job.Block((SITE, SITE)), job.Block((SITE,)))
        self._assert_blocks_are(expected_blocks)

    def test_splitting_with_region_intersection(self):
        region_constraint = shapes.RegionConstraint.from_simple(
                (0.0, 0.0), (2.0, 2.0))

        sites = (shapes.Site(1.0, 1.0), shapes.Site(1.5, 1.5),
            shapes.Site(2.0, 2.0), shapes.Site(3.0, 3.0))

        expected_blocks = (
                job.Block((shapes.Site(1.0, 1.0), shapes.Site(1.5, 1.5))),
                job.Block((shapes.Site(2.0, 2.0),)))

        self.splitter = job.BlockSplitter(sites, 2, constraint=region_constraint)
        self._assert_blocks_are(expected_blocks)

    def _assert_blocks_are(self, expected_blocks):
        for idx, block in enumerate(self.splitter):
            self.assertEqual(expected_blocks[idx], block)

    def _assert_number_of_blocks_is(self, number):
        counter = 0

        for block in self.splitter:
            counter += 1

        self.assertEqual(number, counter)
<commit_msg>Uncomment launch<commit_after># -*- coding: utf-8 -*-

import math
import os
import unittest

from openquake import shapes
from openquake import test
from openquake import job
from openquake.job import Job, EXPOSURE, INPUT_REGION
from openquake.job.mixins import Mixin
from openquake.risk.job import RiskJobMixin
from openquake.risk.job.probabilistic import ProbabilisticEventMixin

CONFIG_FILE = "config.gem"
CONFIG_WITH_INCLUDES = "config_with_includes.gem"

TEST_JOB_FILE = test.smoketest_file('simplecase/config.gem')

SITE = shapes.Site(1.0, 1.0)
EXPOSURE_TEST_FILE = "ExposurePortfolioFile-test.xml"
REGION_EXPOSURE_TEST_FILE = "ExposurePortfolioFile-test.region"
BLOCK_SPLIT_TEST_FILE = "block_split.gem"
REGION_TEST_FILE = "small.region"

class JobTestCase(unittest.TestCase):
    def setUp(self):
        self.job = Job.from_file(test.test_file(CONFIG_FILE))
        self.job_with_includes = Job.from_file(test.test_file(CONFIG_WITH_INCLUDES))

    def tearDown(self):
        [os.remove(f) for f in os.listdir(os.curdir) if f.count("-super.gem")]

    def test_job_writes_to_super_config(self):
        for job in [self.job, self.job_with_includes]:
            config_file = "%s-super.gem" % job.job_id
            self.assertTrue(os.listdir(os.curdir).count(config_file))

    def test_configuration_is_the_same_no_matter_which_way_its_provided(self):
        self.assertEqual(self.job.params, self.job_with_includes.params)

    def test_job_mixes_in_properly(self):
        with Mixin(self.job, RiskJobMixin, key="risk"):
            self.assertTrue(RiskJobMixin in self.job.__class__.__bases__)
            self.assertTrue(ProbabilisticEventMixin in self.job.__class__.__bases__)

        with Mixin(self.job, ProbabilisticEventMixin):
            self.assertTrue(ProbabilisticEventMixin in self.job.__class__.__bases__)

    def test_job_runs_with_a_good_config(self):
        job = Job.from_file(TEST_JOB_FILE)
        self.assertTrue(job.launch())

    def test_a_job_has_an_identifier(self):
        self.assertEqual(1, Job({}, 1).id)

    def test_can_store_and_read_jobs_from_kvs(self):
        self.job = Job.from_file(os.path.join(test.DATA_DIR, CONFIG_FILE))
        self.assertEqual(self.job, Job.from_kvs(self.job.id))

    def test_prepares_blocks_using_the_exposure(self):
        a_job = Job({EXPOSURE: os.path.join(test.DATA_DIR, EXPOSURE_TEST_FILE)})
        a_job._partition()
        blocks_keys = a_job.blocks_keys

        expected_block = job.Block((shapes.Site(9.15000, 45.16667),
                shapes.Site(9.15333, 45.12200), shapes.Site(9.14777, 45.17999),
                shapes.Site(9.15765, 45.13005), shapes.Site(9.15934, 45.13300),
                shapes.Site(9.15876, 45.13805)))

        self.assertEqual(1, len(blocks_keys))
        self.assertEqual(expected_block, job.Block.from_kvs(blocks_keys[0]))

    def test_prepares_blocks_using_the_exposure_and_filtering(self):
        a_job = Job({EXPOSURE: test.test_file(EXPOSURE_TEST_FILE),
                     INPUT_REGION: test.test_file(REGION_EXPOSURE_TEST_FILE)})
        a_job._partition()
        blocks_keys = a_job.blocks_keys

        print blocks_keys
        print job.Block.from_kvs(blocks_keys[0]).sites

        expected_block = job.Block((shapes.Site(9.15, 45.16667),
                                    shapes.Site(9.15333, 45.122),
                                    shapes.Site(9.14777, 45.17999),
                                    shapes.Site(9.15765, 45.13005),
                                    shapes.Site(9.15934, 45.133),
                                    shapes.Site(9.15876, 45.13805)))

        self.assertEqual(1, len(blocks_keys))
        self.assertEqual(expected_block, job.Block.from_kvs(blocks_keys[0]))

    def test_prepares_blocks_using_the_input_region(self):
        """ This test might be currently catastrophically retarded. If it is
        blame Lars.
        """

        block_path = test.test_file(BLOCK_SPLIT_TEST_FILE)

        print "In open job"
        a_job = Job.from_file(block_path)

        verts = [float(x) for x in a_job.params['REGION_VERTEX'].split(",")]
        # Flips lon and lat, and builds a list of coord tuples
        coords = zip(verts[1::2], verts[::2])
        expected = shapes.RegionConstraint.from_coordinates(coords)
        expected.cell_size = float(a_job.params['REGION_GRID_SPACING'])

        expected_sites = []
        for site in expected:
            print site
            expected_sites.append(site)

        a_job._partition()
        blocks_keys = a_job.blocks_keys

        self.assertEqual(1, len(blocks_keys))
        self.assertEqual(job.Block(expected_sites),
                         job.Block.from_kvs(blocks_keys[0]))

    def test_with_no_partition_we_just_process_a_single_block(self):
        job.SITES_PER_BLOCK = 1

        # test exposure has 6 assets
        a_job = Job({EXPOSURE: os.path.join(
                test.DATA_DIR, EXPOSURE_TEST_FILE)})

        a_job._partition()
        blocks_keys = a_job.blocks_keys

        # but we have 1 block instead of 6
        self.assertEqual(1, len(blocks_keys))

class BlockTestCase(unittest.TestCase):

    def test_a_block_has_a_unique_id(self):
        self.assertTrue(job.Block(()).id)
        self.assertTrue(job.Block(()).id != job.Block(()).id)

    def test_can_serialize_a_block_into_kvs(self):
        block = job.Block((SITE, SITE))
        block.to_kvs()

        self.assertEqual(block, job.Block.from_kvs(block.id))

class BlockSplitterTestCase(unittest.TestCase):

    def setUp(self):
        self.splitter = None

    def test_an_empty_set_produces_no_blocks(self):
        self.splitter = job.BlockSplitter(())
        self._assert_number_of_blocks_is(0)

    def test_splits_the_set_into_a_single_block(self):
        self.splitter = job.BlockSplitter((SITE,), 3)
        self._assert_number_of_blocks_is(1)

        self.splitter = job.BlockSplitter((SITE, SITE), 3)
        self._assert_number_of_blocks_is(1)

        self.splitter = job.BlockSplitter((SITE, SITE, SITE), 3)
        self._assert_number_of_blocks_is(1)

    def test_splits_the_set_into_multiple_blocks(self):
        self.splitter = job.BlockSplitter((SITE, SITE), 1)
        self._assert_number_of_blocks_is(2)

        self.splitter = job.BlockSplitter((SITE, SITE, SITE), 2)
        self._assert_number_of_blocks_is(2)

    def test_generates_the_correct_blocks(self):
        self.splitter = job.BlockSplitter((SITE, SITE, SITE), 3)
        expected_blocks = (job.Block((SITE, SITE, SITE)),)
        self._assert_blocks_are(expected_blocks)

        self.splitter = job.BlockSplitter((SITE, SITE, SITE), 2)
        expected_blocks = (job.Block((SITE, SITE)), job.Block((SITE,)))
        self._assert_blocks_are(expected_blocks)

    def test_splitting_with_region_intersection(self):
        region_constraint = shapes.RegionConstraint.from_simple(
                (0.0, 0.0), (2.0, 2.0))

        sites = (shapes.Site(1.0, 1.0), shapes.Site(1.5, 1.5),
            shapes.Site(2.0, 2.0), shapes.Site(3.0, 3.0))

        expected_blocks = (
                job.Block((shapes.Site(1.0, 1.0), shapes.Site(1.5, 1.5))),
                job.Block((shapes.Site(2.0, 2.0),)))

        self.splitter = job.BlockSplitter(sites, 2, constraint=region_constraint)
        self._assert_blocks_are(expected_blocks)

    def _assert_blocks_are(self, expected_blocks):
        for idx, block in enumerate(self.splitter):
            self.assertEqual(expected_blocks[idx], block)

    def _assert_number_of_blocks_is(self, number):
        counter = 0

        for block in self.splitter:
            counter += 1

        self.assertEqual(number, counter)
<|endoftext|>


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

'''
test_ahocorasick.py
(c) Will Roberts  12 December, 2015

Unit tests for the ``ahocorasick`` module.
'''

from __future__ import absolute_import, print_function, unicode_literals
from .. import ahocorasick
import unittest

class TestAhocorasick(unittest.TestCase):
    '''
    Unit tests for the `ahocorasick` module.
    '''

    def test_boundary_transform(self):
        '''
        Test the boundary_transform.
        '''
        self.assertEqual(''.join(ahocorasick.boundary_transform('abc')),
                         '\x00abc\x00')
        self.assertEqual(''.join(ahocorasick.boundary_transform('abc def')),
                         '\x00abc\x00 \x00def\x00')
        self.assertEqual(''.join(ahocorasick.boundary_transform('  abc def  ')),
                         '\x00  \x00abc\x00 \x00def\x00  \x00')

    def test_boundary_transform_2(self):
        '''
        Test the boundary_transform.
        '''
        self.assertEqual(''.join(ahocorasick.boundary_transform('abc def')),
                         '\x00abc\x00 \x00def\x00')
        self.assertEqual(''.join(ahocorasick.boundary_transform('\x00abc def')),
                         '\x00abc\x00 \x00def\x00')
        self.assertEqual(''.join(ahocorasick.boundary_transform('abc def\x00')),
                         '\x00abc\x00 \x00def\x00')
        self.assertEqual(''.join(ahocorasick.boundary_transform('\x00abc def\x00')),
                         '\x00abc\x00 \x00def\x00')
        self.assertEqual(''.join(ahocorasick.boundary_transform('abc def', False)),
                         'abc\x00 \x00def')
        self.assertEqual(''.join(ahocorasick.boundary_transform('\x00abc def', False)),
                         '\x00abc\x00 \x00def')
        self.assertEqual(''.join(ahocorasick.boundary_transform('abc def\x00', False)),
                         'abc\x00 \x00def\x00')
        self.assertEqual(''.join(ahocorasick.boundary_transform('\x00abc def\x00', False)),
                         '\x00abc\x00 \x00def\x00')

    def test_transform_roundtrip(self):
        '''
        Test boundary_transform round trip.
        '''
        test_strings = ['', 'abc', 'abc def', '  abc def', '  abc def  ', 'abc def  ',]
        for force_edges in [False, True]:
            for test_string in test_strings:
                self.assertEqual(test_string,
                                 ''.join(ahocorasick.boundary_untransform(
                                     ahocorasick.boundary_transform(test_string,
                                                                    force_edges))))

    def test_debug1(self):
        '''
        Test case from Wikipedia.

        https://en.wikipedia.org/wiki/Aho%E2%80%93Corasick_algorithm
        '''
        trie = ahocorasick.Trie()
        trie['a'] = '(a)'
        trie['ab'] = '(ab)'
        trie['bab'] = '(bab)'
        trie['bc'] = '(bc)'
        trie['bca'] = '(bca)'
        trie['c'] = '(c)'
        trie['caa'] = '(caa)'
        debug_1_ppstr = '''<ROOT>
  a (value = "(a)")
    ab (value = "(ab)")
  b
    ba
      bab (value = "(bab)")
    bc (value = "(bc)")
      bca (value = "(bca)")
  c (value = "(c)")
    ca
      caa (value = "(caa)")'''
        self.assertEqual(trie.pretty_print_str(), debug_1_ppstr)

    def test_debug2(self):
        '''Hand-written test case.'''
        trie = ahocorasick.AhoCorasickTrie()
        trie['damn'] = '(damn)'
        trie['dog'] = '(dog)'
        trie['cat'] = '(cat)'
        trie['cablex'] = '(cablex)'
        self.assertTrue('damn' in trie)
        self.assertTrue('dog' in trie)
        self.assertTrue('cat' in trie)
        self.assertTrue('cablex' in trie)
        self.assertEqual(trie['damn'], '(damn)')
        self.assertEqual(trie['dog'], '(dog)')
        self.assertEqual(trie['cat'], '(cat)')
        self.assertEqual(trie['cablex'], '(cablex)')
        self.assertFalse('d' in trie)
        self.assertFalse('cabbie' in trie)
        self.assertFalse('cable' in trie)
        self.assertRaises(KeyError, trie.__getitem__, ('d',))
        self.assertRaises(KeyError, trie.__getitem__, ('cablex',))
        trie._set_suffix_links()
        debug_2_ppstr1 = '''<ROOT>
  c (suffix = "")
    ca (suffix = "")
      cab (suffix = "")
        cabl (suffix = "")
          cable (suffix = "")
            cablex (value = "(cablex)") (suffix = "")
      cat (value = "(cat)") (suffix = "")
  d (suffix = "")
    da (suffix = "")
      dam (suffix = "")
        damn (value = "(damn)") (suffix = "")
    do (suffix = "")
      dog (value = "(dog)") (suffix = "")'''
        self.assertEqual(trie.pretty_print_str(), debug_2_ppstr1)
        trie['ogre'] = '(ogre)'
        trie._set_suffix_links()
        debug_2_ppstr2 = '''<ROOT>
  c (suffix = "")
    ca (suffix = "")
      cab (suffix = "")
        cabl (suffix = "")
          cable (suffix = "")
            cablex (value = "(cablex)") (suffix = "")
      cat (value = "(cat)") (suffix = "")
  d (suffix = "")
    da (suffix = "")
      dam (suffix = "")
        damn (value = "(damn)") (suffix = "")
    do (suffix = "o")
      dog (value = "(dog)") (suffix = "og")
  o (suffix = "")
    og (suffix = "")
      ogr (suffix = "")
        ogre (value = "(ogre)") (suffix = "")'''
        self.assertEqual(trie.pretty_print_str(), debug_2_ppstr2)
        self.assertEqual(trie.greedy_replace('dogre'), '(dog)re')
        self.assertEqual(trie.greedy_replace('doogre'), 'do(ogre)')

    def test_debug3(self):
        '''
        Test case from Wikipedia.

        https://en.wikipedia.org/wiki/Aho%E2%80%93Corasick_algorithm
        '''
        trie = ahocorasick.AhoCorasickTrie()
        trie['a'] = '(a)'
        trie['ab'] = '(ab)'
        trie['bab'] = '(bab)'
        trie['bc'] = '(bc)'
        trie['bca'] = '(bca)'
        trie['c'] = '(c)'
        trie['caa'] = '(caa)'
        trie._set_suffix_links()
        debug_3_ppstr = '''<ROOT>
  a (value = "(a)") (suffix = "")
    ab (value = "(ab)") (suffix = "b")
  b (suffix = "")
    ba (suffix = "a") (dict_suffix = "a")
      bab (value = "(bab)") (suffix = "ab") (dict_suffix = "ab")
    bc (value = "(bc)") (suffix = "c") (dict_suffix = "c")
      bca (value = "(bca)") (suffix = "ca") (dict_suffix = "a")
  c (value = "(c)") (suffix = "")
    ca (suffix = "a") (dict_suffix = "a")
      caa (value = "(caa)") (suffix = "a") (dict_suffix = "a")'''
        self.assertEqual(trie.pretty_print_str(),
                         debug_3_ppstr)
        self.assertEqual(list(trie.find_all('abccab')),
                         [(0, 1, '(a)'), (0, 2, '(ab)'), (1, 2, '(bc)'),
                          (2, 1, '(c)'), (3, 1, '(c)'), (4, 1, '(a)'),
                          (4, 2, '(ab)')])
        self.assertEqual(trie.replace('abccab'),
                         '(a)(bc)(c)(ab)')
        self.assertEqual(trie.greedy_replace('ba'), 'b(a)')
        self.assertEqual(trie.greedy_replace('bab'), '(bab)')
        self.assertEqual(trie.greedy_replace('abccab'),
                         '(a)(bc)(c)(a)b')

    def test_debug4(self):
        '''Hand-written test case.'''
        trie = ahocorasick.AhoCorasickTrie()
        trie['x'] = '(x)'
        trie['xabc'] = '(xabc)'
        trie['ab'] = '(ab)'
        trie._set_suffix_links()
        debug_4_ppstr = '''<ROOT>
  a (suffix = "")
    ab (value = "(ab)") (suffix = "")
  x (value = "(x)") (suffix = "")
    xa (suffix = "a")
      xab (suffix = "ab") (dict_suffix = "ab")
        xabc (value = "(xabc)") (suffix = "")'''
        self.assertEqual(trie.pretty_print_str(),
                         debug_4_ppstr)
        self.assertEqual(list(trie.find_all('xabc')),
                         [(0, 1, '(x)'), (1, 2, '(ab)'), (0, 4, '(xabc)')])
        self.assertEqual(trie.replace('xabc'),
                         '(xabc)')
        self.assertEqual(trie.greedy_replace('xabc'),
                         '(x)(ab)c')
        self.assertEqual(list(trie.find_all('xabd')),
                         [(0, 1, '(x)'), (1, 2, '(ab)')])
        self.assertEqual(trie.replace('xabd'),
                         '(x)(ab)d')
        self.assertEqual(trie.greedy_replace('xabd'),
                         '(x)(ab)d')

    def test_debug5(self):
        '''Hand-written test case.'''
        trie = ahocorasick.AhoCorasickTrie()
        trie['ac'] = '(ac)'
        trie['acabx'] = '(acabx)'
        trie['ab'] = '(ab)'
        trie._set_suffix_links()
        debug_5_ppstr = '''<ROOT>
  a (suffix = "")
    ab (value = "(ab)") (suffix = "")
    ac (value = "(ac)") (suffix = "")
      aca (suffix = "a")
        acab (suffix = "ab") (dict_suffix = "ab")
          acabx (value = "(acabx)") (suffix = "")'''
        self.assertEqual(trie.pretty_print_str(),
                         debug_5_ppstr)
        self.assertEqual(list(trie.find_all('acaby')),
                         [(0, 2, '(ac)'), (2, 2, '(ab)')])
        self.assertEqual(trie.replace('acaby'),
                         '(ac)(ab)y')
        self.assertEqual(trie.greedy_replace('acaby'),
                         '(ac)(ab)y')

    def test_debug6(self):
        '''
        Test case from Wikipedia.

        https://en.wikipedia.org/wiki/Aho%E2%80%93Corasick_algorithm
        '''
        trie = ahocorasick.AhoCorasickTrie()
        trie['ab'] = '(ab)'
        trie['bab'] = '(bab)'
        trie['bc'] = '(bc)'
        trie['bca'] = '(bca)'
        trie['c'] = '(c)'
        trie['caa'] = '(caa)'
        trie._set_suffix_links()
        debug_6_ppstr = '''<ROOT>
  a (suffix = "")
    ab (value = "(ab)") (suffix = "b")
  b (suffix = "")
    ba (suffix = "a")
      bab (value = "(bab)") (suffix = "ab") (dict_suffix = "ab")
    bc (value = "(bc)") (suffix = "c") (dict_suffix = "c")
      bca (value = "(bca)") (suffix = "ca")
  c (value = "(c)") (suffix = "")
    ca (suffix = "a")
      caa (value = "(caa)") (suffix = "a")'''
        self.assertEqual(trie.pretty_print_str(),
                         debug_6_ppstr)
        self.assertEqual(list(trie.find_all('abccab')),
                         [(0, 2, '(ab)'), (1, 2, '(bc)'), (2, 1, '(c)'),
                          (3, 1, '(c)'), (4, 2, '(ab)')])
        self.assertEqual(trie.replace('abccab'),
                         '(ab)(c)(c)(ab)')
        self.assertEqual(trie.greedy_replace('ba'), 'ba')
        self.assertEqual(trie.greedy_replace('bab'), '(bab)')
        self.assertEqual(trie.greedy_replace('abccab'),
                         '(ab)(c)(c)(ab)')

    def test_debug7(self):
        '''Hand-written test case.'''
        trie = ahocorasick.AhoCorasickTrie()
        trie['r 3'] = 'r_3'
        trie['Rather he'] = 'Rather_he'
        trie['er Rabbit'] = 'er_Rabbit'
        trie['a Merry'] = 'a_Merry'
        trie._set_suffix_links()
        debug_7_ppstr = '''<ROOT>
  R (suffix = "")
    Ra (suffix = "a")
      Rat (suffix = "")
        Rath (suffix = "")
          Rathe (suffix = "e")
            Rather (suffix = "er")
              Rather  (suffix = "er ")
                Rather h (suffix = "")
                  Rather he (value = "Rather_he") (suffix = "e")
  a (suffix = "")
    a  (suffix = "")
      a M (suffix = "")
        a Me (suffix = "e")
          a Mer (suffix = "er")
            a Merr (suffix = "r")
              a Merry (value = "a_Merry") (suffix = "")
  e (suffix = "")
    er (suffix = "r")
      er  (suffix = "r ")
        er R (suffix = "R")
          er Ra (suffix = "Ra")
            er Rab (suffix = "")
              er Rabb (suffix = "")
                er Rabbi (suffix = "")
                  er Rabbit (value = "er_Rabbit") (suffix = "")
  r (suffix = "")
    r  (suffix = "")
      r 3 (value = "r_3") (suffix = "")'''
        self.assertEqual(trie.pretty_print_str(), debug_7_ppstr)

    def test_debug8(self):
        '''Hand-written test case.'''
        trie = ahocorasick.AhoCorasickTrie()
        trie['cart'] = '(cart)'
        trie['cat'] = '(cat)'
        self.assertEqual(trie.greedy_replace('carry cat cat carcat cart car'),
                         'carry (cat) (cat) car(cat) (cart) car')
        trie = ahocorasick.AhoCorasickTrie()
        trie[ahocorasick.boundary_transform('dog')] = '(dog)'
        trie[ahocorasick.boundary_transform('ca')] = '(ca)'
        trie[ahocorasick.boundary_transform('cat')] = '(cat)'
        self.assertEqual(''.join(ahocorasick.boundary_untransform(
            trie.greedy_replace(ahocorasick.boundary_transform('my dog is a ca catty cat cat')))),
                         'my (dog) is a (ca) catty (cat) (cat)')

if __name__ == '__main__':
    unittest.main()
<commit_msg>fsed/tests/test_ahocorasick.py: slight change to unit test<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-

'''
test_ahocorasick.py
(c) Will Roberts  12 December, 2015

Unit tests for the ``ahocorasick`` module.
'''

from __future__ import absolute_import, print_function, unicode_literals
from .. import ahocorasick
import unittest

class TestAhocorasick(unittest.TestCase):
    '''
    Unit tests for the `ahocorasick` module.
    '''

    def test_boundary_transform(self):
        '''
        Test the boundary_transform.
        '''
        self.assertEqual(''.join(ahocorasick.boundary_transform('abc')),
                         '\x00abc\x00')
        self.assertEqual(''.join(ahocorasick.boundary_transform('abc def')),
                         '\x00abc\x00 \x00def\x00')
        self.assertEqual(''.join(ahocorasick.boundary_transform('  abc def  ')),
                         '\x00  \x00abc\x00 \x00def\x00  \x00')

    def test_boundary_transform_2(self):
        '''
        Test the boundary_transform.
        '''
        self.assertEqual(''.join(ahocorasick.boundary_transform('abc def')),
                         '\x00abc\x00 \x00def\x00')
        self.assertEqual(''.join(ahocorasick.boundary_transform('\x00abc def')),
                         '\x00abc\x00 \x00def\x00')
        self.assertEqual(''.join(ahocorasick.boundary_transform('abc def\x00')),
                         '\x00abc\x00 \x00def\x00')
        self.assertEqual(''.join(ahocorasick.boundary_transform('\x00abc def\x00')),
                         '\x00abc\x00 \x00def\x00')
        self.assertEqual(''.join(ahocorasick.boundary_transform('abc def', False)),
                         'abc\x00 \x00def')
        self.assertEqual(''.join(ahocorasick.boundary_transform('\x00abc def', False)),
                         '\x00abc\x00 \x00def')
        self.assertEqual(''.join(ahocorasick.boundary_transform('abc def\x00', False)),
                         'abc\x00 \x00def\x00')
        self.assertEqual(''.join(ahocorasick.boundary_transform('\x00abc def\x00', False)),
                         '\x00abc\x00 \x00def\x00')

    def test_transform_roundtrip(self):
        '''
        Test boundary_transform round trip.
        '''
        test_strings = ['', 'abc', 'abc def', '  abc def', '  abc def  ', 'abc def  ',]
        for force_edges in [False, True]:
            for test_string in test_strings:
                self.assertEqual(test_string,
                                 ''.join(ahocorasick.boundary_untransform(
                                     ahocorasick.boundary_transform(test_string,
                                                                    force_edges))))

    def test_debug1(self):
        '''
        Test case from Wikipedia.

        https://en.wikipedia.org/wiki/Aho%E2%80%93Corasick_algorithm
        '''
        trie = ahocorasick.Trie()
        trie['a'] = '(a)'
        trie['ab'] = '(ab)'
        trie['bab'] = '(bab)'
        trie['bc'] = '(bc)'
        trie['bca'] = '(bca)'
        trie['c'] = '(c)'
        trie['caa'] = '(caa)'
        debug_1_ppstr = '''<ROOT>
  a (value = "(a)")
    ab (value = "(ab)")
  b
    ba
      bab (value = "(bab)")
    bc (value = "(bc)")
      bca (value = "(bca)")
  c (value = "(c)")
    ca
      caa (value = "(caa)")'''
        self.assertEqual(trie.pretty_print_str(), debug_1_ppstr)

    def test_debug2(self):
        '''Hand-written test case.'''
        trie = ahocorasick.AhoCorasickTrie()
        trie['damn'] = '(damn)'
        trie['dog'] = '(dog)'
        trie['cat'] = '(cat)'
        trie['cablex'] = '(cablex)'
        self.assertTrue('damn' in trie)
        self.assertTrue('dog' in trie)
        self.assertTrue('cat' in trie)
        self.assertTrue('cablex' in trie)
        self.assertEqual(trie['damn'], '(damn)')
        self.assertEqual(trie['dog'], '(dog)')
        self.assertEqual(trie['cat'], '(cat)')
        self.assertEqual(trie['cablex'], '(cablex)')
        self.assertFalse('d' in trie)
        self.assertFalse('cabbie' in trie)
        self.assertFalse('cable' in trie)
        self.assertRaises(KeyError, trie.__getitem__, ('d',))
        self.assertRaises(KeyError, trie.__getitem__, ('cablex',))
        trie._set_suffix_links()
        debug_2_ppstr1 = '''<ROOT>
  c (suffix = "")
    ca (suffix = "")
      cab (suffix = "")
        cabl (suffix = "")
          cable (suffix = "")
            cablex (value = "(cablex)") (suffix = "")
      cat (value = "(cat)") (suffix = "")
  d (suffix = "")
    da (suffix = "")
      dam (suffix = "")
        damn (value = "(damn)") (suffix = "")
    do (suffix = "")
      dog (value = "(dog)") (suffix = "")'''
        self.assertEqual(trie.pretty_print_str(), debug_2_ppstr1)
        trie['ogre'] = '(ogre)'
        trie._set_suffix_links()
        debug_2_ppstr2 = '''<ROOT>
  c (suffix = "")
    ca (suffix = "")
      cab (suffix = "")
        cabl (suffix = "")
          cable (suffix = "")
            cablex (value = "(cablex)") (suffix = "")
      cat (value = "(cat)") (suffix = "")
  d (suffix = "")
    da (suffix = "")
      dam (suffix = "")
        damn (value = "(damn)") (suffix = "")
    do (suffix = "o")
      dog (value = "(dog)") (suffix = "og")
  o (suffix = "")
    og (suffix = "")
      ogr (suffix = "")
        ogre (value = "(ogre)") (suffix = "")'''
        self.assertEqual(trie.pretty_print_str(), debug_2_ppstr2)
        self.assertEqual(trie.greedy_replace('dogre'), '(dog)re')
        self.assertEqual(trie.greedy_replace('doogre'), 'do(ogre)')

    def test_debug3(self):
        '''
        Test case from Wikipedia.

        https://en.wikipedia.org/wiki/Aho%E2%80%93Corasick_algorithm
        '''
        trie = ahocorasick.AhoCorasickTrie()
        trie['a'] = '(a)'
        trie['ab'] = '(ab)'
        trie['bab'] = '(bab)'
        trie['bc'] = '(bc)'
        trie['bca'] = '(bca)'
        trie['c'] = '(c)'
        trie['caa'] = '(caa)'
        trie._set_suffix_links()
        debug_3_ppstr = '''<ROOT>
  a (value = "(a)") (suffix = "")
    ab (value = "(ab)") (suffix = "b")
  b (suffix = "")
    ba (suffix = "a") (dict_suffix = "a")
      bab (value = "(bab)") (suffix = "ab") (dict_suffix = "ab")
    bc (value = "(bc)") (suffix = "c") (dict_suffix = "c")
      bca (value = "(bca)") (suffix = "ca") (dict_suffix = "a")
  c (value = "(c)") (suffix = "")
    ca (suffix = "a") (dict_suffix = "a")
      caa (value = "(caa)") (suffix = "a") (dict_suffix = "a")'''
        self.assertEqual(trie.pretty_print_str(),
                         debug_3_ppstr)
        self.assertEqual(list(trie.find_all('abccab')),
                         [(0, 1, '(a)'), (0, 2, '(ab)'), (1, 2, '(bc)'),
                          (2, 1, '(c)'), (3, 1, '(c)'), (4, 1, '(a)'),
                          (4, 2, '(ab)')])
        self.assertEqual(trie.replace('abccab'),
                         '(a)(bc)(c)(ab)')
        self.assertEqual(trie.greedy_replace('ba'), 'b(a)')
        self.assertEqual(trie.greedy_replace('bab'), '(bab)')
        self.assertEqual(trie.greedy_replace('abccab'),
                         '(a)(bc)(c)(a)b')

    def test_debug4(self):
        '''Hand-written test case.'''
        trie = ahocorasick.AhoCorasickTrie()
        trie['x'] = '(x)'
        trie['xabc'] = '(xabc)'
        trie['ab'] = '(ab)'
        trie._set_suffix_links()
        debug_4_ppstr = '''<ROOT>
  a (suffix = "")
    ab (value = "(ab)") (suffix = "")
  x (value = "(x)") (suffix = "")
    xa (suffix = "a")
      xab (suffix = "ab") (dict_suffix = "ab")
        xabc (value = "(xabc)") (suffix = "")'''
        self.assertEqual(trie.pretty_print_str(),
                         debug_4_ppstr)
        self.assertEqual(list(trie.find_all('xabc')),
                         [(0, 1, '(x)'), (1, 2, '(ab)'), (0, 4, '(xabc)')])
        self.assertEqual(trie.replace('xabc'),
                         '(xabc)')
        self.assertEqual(trie.greedy_replace('xabc'),
                         '(x)(ab)c')
        self.assertEqual(list(trie.find_all('xabd')),
                         [(0, 1, '(x)'), (1, 2, '(ab)')])
        self.assertEqual(trie.replace('xabd'),
                         '(x)(ab)d')
        self.assertEqual(trie.greedy_replace('xabd'),
                         '(x)(ab)d')

    def test_debug5(self):
        '''Hand-written test case.'''
        trie = ahocorasick.AhoCorasickTrie()
        trie['ac'] = '(ac)'
        trie['acabx'] = '(acabx)'
        trie['ab'] = '(ab)'
        trie._set_suffix_links()
        debug_5_ppstr = '''<ROOT>
  a (suffix = "")
    ab (value = "(ab)") (suffix = "")
    ac (value = "(ac)") (suffix = "")
      aca (suffix = "a")
        acab (suffix = "ab") (dict_suffix = "ab")
          acabx (value = "(acabx)") (suffix = "")'''
        self.assertEqual(trie.pretty_print_str(),
                         debug_5_ppstr)
        self.assertEqual(list(trie.find_all('acaby')),
                         [(0, 2, '(ac)'), (2, 2, '(ab)')])
        self.assertEqual(trie.replace('acaby'),
                         '(ac)(ab)y')
        self.assertEqual(trie.greedy_replace('acaby'),
                         '(ac)(ab)y')

    def test_debug6(self):
        '''
        Test case from Wikipedia.

        https://en.wikipedia.org/wiki/Aho%E2%80%93Corasick_algorithm
        '''
        trie = ahocorasick.AhoCorasickTrie()
        trie['ab'] = '(ab)'
        trie['bab'] = '(bab)'
        trie['bca'] = '(bca)'
        trie['c'] = '(c)'
        trie['caa'] = '(caa)'
        trie._set_suffix_links()
        debug_6_ppstr = '''<ROOT>
  a (suffix = "")
    ab (value = "(ab)") (suffix = "b")
  b (suffix = "")
    ba (suffix = "a")
      bab (value = "(bab)") (suffix = "ab") (dict_suffix = "ab")
    bc (suffix = "c") (dict_suffix = "c")
      bca (value = "(bca)") (suffix = "ca")
  c (value = "(c)") (suffix = "")
    ca (suffix = "a")
      caa (value = "(caa)") (suffix = "a")'''
        self.assertEqual(trie.pretty_print_str(),
                         debug_6_ppstr)
        self.assertEqual(list(trie.find_all('abccab')),
                         [(0, 2, '(ab)'), (2, 1, '(c)'),
                          (3, 1, '(c)'), (4, 2, '(ab)')])
        self.assertEqual(trie.replace('abccab'),
                         '(ab)(c)(c)(ab)')
        self.assertEqual(trie.greedy_replace('ba'), 'ba')
        self.assertEqual(trie.greedy_replace('bc'), 'b(c)')
        self.assertEqual(trie.greedy_replace('bab'), '(bab)')
        self.assertEqual(trie.greedy_replace('abccab'),
                         '(ab)(c)(c)(ab)')

    def test_debug7(self):
        '''Hand-written test case.'''
        trie = ahocorasick.AhoCorasickTrie()
        trie['r 3'] = 'r_3'
        trie['Rather he'] = 'Rather_he'
        trie['er Rabbit'] = 'er_Rabbit'
        trie['a Merry'] = 'a_Merry'
        trie._set_suffix_links()
        debug_7_ppstr = '''<ROOT>
  R (suffix = "")
    Ra (suffix = "a")
      Rat (suffix = "")
        Rath (suffix = "")
          Rathe (suffix = "e")
            Rather (suffix = "er")
              Rather  (suffix = "er ")
                Rather h (suffix = "")
                  Rather he (value = "Rather_he") (suffix = "e")
  a (suffix = "")
    a  (suffix = "")
      a M (suffix = "")
        a Me (suffix = "e")
          a Mer (suffix = "er")
            a Merr (suffix = "r")
              a Merry (value = "a_Merry") (suffix = "")
  e (suffix = "")
    er (suffix = "r")
      er  (suffix = "r ")
        er R (suffix = "R")
          er Ra (suffix = "Ra")
            er Rab (suffix = "")
              er Rabb (suffix = "")
                er Rabbi (suffix = "")
                  er Rabbit (value = "er_Rabbit") (suffix = "")
  r (suffix = "")
    r  (suffix = "")
      r 3 (value = "r_3") (suffix = "")'''
        self.assertEqual(trie.pretty_print_str(), debug_7_ppstr)

    def test_debug8(self):
        '''Hand-written test case.'''
        trie = ahocorasick.AhoCorasickTrie()
        trie['cart'] = '(cart)'
        trie['cat'] = '(cat)'
        self.assertEqual(trie.greedy_replace('carry cat cat carcat cart car'),
                         'carry (cat) (cat) car(cat) (cart) car')
        trie = ahocorasick.AhoCorasickTrie()
        trie[ahocorasick.boundary_transform('dog')] = '(dog)'
        trie[ahocorasick.boundary_transform('ca')] = '(ca)'
        trie[ahocorasick.boundary_transform('cat')] = '(cat)'
        self.assertEqual(''.join(ahocorasick.boundary_untransform(
            trie.greedy_replace(ahocorasick.boundary_transform('my dog is a ca catty cat cat')))),
                         'my (dog) is a (ca) catty (cat) (cat)')

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


<commit_before>import os
import shutil
import unittest
from collections import namedtuple
from io import StringIO

from jpp.cli import main as cli_entry_point

CURR_DIR = os.path.dirname(os.path.realpath(__file__))

class TestCli(unittest.TestCase):
    TMP_TEST_FILES = os.path.join(CURR_DIR, '__tmp__')
    _dir_bk = None

    @classmethod
    def setUpClass(cls):
        FileDef = namedtuple('FileDef', ('name', 'contents', 'sub_path'))
        required_files = (
            FileDef('compact_test.jpp', '{\n"many": 1, \n"lines": 2\n}', ''),
            FileDef('main.jpp', '', ''),
            FileDef('other.jpp', '', ''),
            FileDef('user_input_test.jpp', '{"foo": user_input["bar"]}', ''),
            FileDef('sub_other.jpp', '', 'sub_path'),
        )
        os.mkdir(cls.TMP_TEST_FILES)
        for file_def in required_files:
            if file_def.sub_path:
                os.mkdir(os.path.join(cls.TMP_TEST_FILES, file_def.sub_path))
                file_path = os.path.join(cls.TMP_TEST_FILES, file_def.sub_path, file_def.name)
            else:
                file_path = os.path.join(cls.TMP_TEST_FILES, file_def.name)
            with open(file_path, 'w') as fp:
                fp.write(file_def.contents)

        cls._dir_bk = os.getcwd()
        os.chdir(os.path.join(cls.TMP_TEST_FILES))

    @classmethod
    def tearDownClass(cls):
        os.chdir(cls._dir_bk)
        shutil.rmtree(cls.TMP_TEST_FILES)

    def test_version(self):
        cli_entry_point(['--version'])

    def test_help(self):
        out_file_object = StringIO()
        cli_entry_point(['jpp', '-h'], out_file_object)
        self.assertRegex(out_file_object.read(), b'^usage:')

    def test_no_args(self):
        out_file_object = StringIO()
        cli_entry_point(['jpp'], out_file_object)
        self.assertEqual(out_file_object.read(), '{}')

    def test_parse_specific_file(self):
        out_file_object = StringIO()
        cli_entry_point(['jpp', 'other.jpp'], out_file_object)
        self.assertEqual(out_file_object.read(), '{}')

    def test_path_option(self):
        out_file_object = StringIO()
        cli_entry_point(['jpp', '--path', os.path.join(CURR_DIR, 'sub_path'), 'sub_other.jpp'], out_file_object)
        self.assertEqual(out_file_object.read(), '{}')

    def test_compact_path(self):
        out_file_object = StringIO()
        cli_entry_point(['jpp', '--compact-print', 'compact_test.jpp'], out_file_object)
        self.assertEqual(out_file_object.read(), '{}')
        # Make sure output is a one-liner at most
        self.assertLessEqual(out_file_object.read()(b'\n'), 1)

    def test_user_input(self):
        out_file_object = StringIO()
        cli_entry_point(['jpp', '--compact-print', '--user-input', '{"bar": "baz"}', 'user_input_test.jpp'],
                        out_file_object)
        self.assertEqual(out_file_object.read(), '{"foo": "baz"}\n')

def main():
    unittest.main()

if __name__ == '__main__':
    main()
<commit_msg>False negative fix: Was reading wrong output<commit_after>import os
import shutil
import unittest
from collections import namedtuple
from io import StringIO

from jpp.cli import main as cli_entry_point

CURR_DIR = os.path.dirname(os.path.realpath(__file__))

class TestCli(unittest.TestCase):
    TMP_TEST_FILES = os.path.join(CURR_DIR, '__tmp__')
    _dir_bk = None

    @classmethod
    def setUpClass(cls):
        FileDef = namedtuple('FileDef', ('name', 'contents', 'sub_path'))
        required_files = (
            FileDef('compact_test.jpp', '{\n"many": 1, \n"lines": 2\n}', ''),
            FileDef('main.jpp', '', ''),
            FileDef('other.jpp', '', ''),
            FileDef('user_input_test.jpp', '{"foo": user_input["bar"]}', ''),
            FileDef('sub_main.jpp', '', ''),
            FileDef('sub_other.jpp', '', 'sub_path'),
        )
        os.mkdir(cls.TMP_TEST_FILES)
        for file_def in required_files:
            if file_def.sub_path:
                os.mkdir(os.path.join(cls.TMP_TEST_FILES, file_def.sub_path))
                file_path = os.path.join(cls.TMP_TEST_FILES, file_def.sub_path, file_def.name)
            else:
                file_path = os.path.join(cls.TMP_TEST_FILES, file_def.name)
            with open(file_path, 'w') as fp:
                fp.write(file_def.contents)

        cls._dir_bk = os.getcwd()
        os.chdir(cls.TMP_TEST_FILES)

    @classmethod
    def tearDownClass(cls):
        os.chdir(cls._dir_bk)
        shutil.rmtree(cls.TMP_TEST_FILES)

    def test_no_args(self):
        out_file_object = StringIO()
        cli_entry_point([], out_file_object)
        out_file_object.seek(0)
        self.assertEqual(out_file_object.read(), '{}')

    def test_parse_specific_file(self):
        out_file_object = StringIO()
        cli_entry_point(['other.jpp'], out_file_object)
        out_file_object.seek(0)
        self.assertEqual(out_file_object.read(), '{}')

    def test_path_option(self):
        out_file_object = StringIO()
        cli_entry_point(['--path', '["{}"]'.format(os.path.join(self.TMP_TEST_FILES, 'sub_path')), 'sub_main.jpp'],
                        out_file_object)
        out_file_object.seek(0)
        self.assertEqual(out_file_object.read(), '{}')

    def test_compact_path(self):
        out_file_object = StringIO()
        cli_entry_point(['--compact-print', 'compact_test.jpp'], out_file_object)
        out_file_object.seek(0)
        self.assertEqual(out_file_object.read(), '{"lines":2,"many":1}')

    def test_user_input(self):
        out_file_object = StringIO()
        cli_entry_point(['--compact-print', '--user-input', '{"bar": "baz"}', 'user_input_test.jpp'],
                        out_file_object)
        out_file_object.seek(0)
        self.assertEqual(out_file_object.read(), '{"foo":"baz"}')

def main():
    unittest.main()

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


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

import sys
import os

from optparse import make_option

from django.core.management.base import BaseCommand, CommandError
from odekro.management.hansard_parser import normalised_lines, scan, \
                                             parse_head, meta

# from odekro import data

class Command(BaseCommand):
    """Read in an info page file and either create or update existing info"""

    help = 'Create/Update info page'
    args = '<file> <Title> [<slug>]'

    def handle(self, *args, **options):
        # file = open(args[])
        if not len(args):
            raise CommandError

        src = args[0]

        content = open(src, 'rU').read()
        lines = scan(meta(normalised_lines(content)), header=True)
        head = parse_head(lines)

        try:
            fname = '%s.txt' % self.filename(head)
            dst = os.path.join(os.path.dirname(src), fname)
            print dst
            os.rename(src, dst)
        except:
            print 'Error processing %s ...' % src
            print lines
            sys.exit(1)

    def filename(self, head):
        return 'debate-%(series)02d-%(volume)03d-%(number)03d-%(date)s' % head<commit_msg>Code cleanup<commit_after>#!/usr/bin/env python

import sys
import os

from optparse import make_option

from django.core.management.base import BaseCommand, CommandError
from odekro.management.hansard_parser import normalised_lines, scan, \
                                             parse_head, meta

# from odekro import data

class Command(BaseCommand):
    """Read in an info page file and either create or update existing info"""

    help = 'Create/Update info page'
    args = '<file> <Title> [<slug>]'

    def handle(self, *args, **options):
        if not len(args):
            raise CommandError

        src = args[0]

        content = open(src, 'rU').read()
        lines = scan(meta(normalised_lines(content)), header=True)

        try:
            fname = '%s.txt' % self.filename(parse_head(lines))
            dst = os.path.join(os.path.dirname(src), fname)
            print dst
            os.rename(src, dst)
        except:
            print 'Error processing %s ...' % src
            print lines
            sys.exit(1)

    def filename(self, head):
        return 'debate-%(series)02d-%(volume)03d-%(number)03d-%(date)s' % head<|endoftext|>


<commit_before>''' This file contains tests for the Text masker.
'''

def test_method_tokenize_pretrained_tokenizer():
    import shap
    from transformers import AutoTokenizer, AutoModelForSequenceClassification

    tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased",use_fast=False)
    masker = shap.maskers.Text(tokenizer)

    test_text = "I have a joke about deep learning but I can't explain it."
    output_ids = masker.tokenize(test_text)['input_ids']
    correct_ids = tokenizer.encode_plus(test_text)['input_ids']

    print("Test logging",output_ids)
    print(correct_ids)
    assert output_ids == correct_ids

def test_method_tokenize_pretrained_tokenizer_fast():
    import shap
    from transformers import AutoTokenizer, AutoModelForSequenceClassification

    tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased",use_fast=True)
    masker = shap.maskers.Text(tokenizer)

    test_text = "I have a joke about deep learning but I can't explain it."
    output_ids = masker.tokenize(test_text)['input_ids']
    correct_ids = tokenizer.encode_plus(test_text)['input_ids']

    assert output_ids == correct_ids

def test_method_token_segments_pretrained_tokenizer():
    import shap
    from transformers import AutoTokenizer, AutoModelForSequenceClassification

    tokenizer = AutoTokenizer.from_pretrained("distilbert-base-cased",use_fast=False)
    masker = shap.maskers.Text(tokenizer)

    test_text = "I ate a Cannoli"
    output_token_segments = masker.token_segments(test_text)
    correct_token_segments = ['','I','ate','a','Can','##no','##li','']
    print(output_token_segments)

    assert output_token_segments == correct_token_segments

def test_method_token_segments_pretrained_tokenizer_fast():
    import shap
    from transformers import AutoTokenizer, AutoModelForSequenceClassification

    tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased",use_fast=True)
    masker = shap.maskers.Text(tokenizer)

    test_text = "I ate a Cannoli"
    output_token_segments = masker.token_segments(test_text)
    correct_token_segments = ['','I ','ate ','a ','Can','no','li','']

    assert output_token_segments == correct_token_segments

def test_masker_call_pretrained_tokenizer():
    import shap
    from transformers import AutoTokenizer, AutoModelForSequenceClassification
    import numpy as np

    tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased",use_fast=False)
    masker = shap.maskers.Text(tokenizer)

    test_text = "I ate a Cannoli"
    test_input_mask = np.array([True, False, True, True, False, True, True, True])

    output_masked_text = masker(test_input_mask,test_text)
    correct_masked_text = '[MASK] ate a [MASK]noli'

    assert output_masked_text == correct_masked_text

def test_masker_call_pretrained_tokenizer_fast():
    import shap
    from transformers import AutoTokenizer, AutoModelForSequenceClassification
    import numpy as np

    tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased",use_fast=True)
    masker = shap.maskers.Text(tokenizer)

    test_text = "I ate a Cannoli"
    test_input_mask = np.array([True, False, True, True, False, True, True, True])

    output_masked_text = masker(test_input_mask,test_text)
    correct_masked_text = '[MASK]ate a [MASK]noli'

    assert output_masked_text == correct_masked_text

def test_sentencepiece_tokenizer_output():
    import numpy as np
    from transformers import AutoTokenizer
    from shap import maskers
    from shap.utils import safe_isinstance

    tokenizer = AutoTokenizer.from_pretrained("Helsinki-NLP/opus-mt-en-es")
    masker = maskers.Text(tokenizer)

    s="This is a test statement for sentencepiece tokenizer"
    masker._update_s_cache(s)
    mask = np.ones(masker._segments_s.shape, dtype=bool)

    sentencepiece_tokenizer_output_processed = masker(mask, s)
    expected_sentencepiece_tokenizer_output_processed = "This is a test statement for sentencepiece tokenizer"
    # since we expect output wrapped in a tuple hence the indexing [0][0] to extract the string
    assert sentencepiece_tokenizer_output_processed[0][0] == expected_sentencepiece_tokenizer_output_processed

def test_keep_prefix_suffix_tokenizer_parsing():
    from transformers import AutoTokenizer
    from shap import maskers

    tokenizer_mt = AutoTokenizer.from_pretrained("Helsinki-NLP/opus-mt-en-es")
    tokenizer_gpt = AutoTokenizer.from_pretrained("gpt2")
    tokenizer_bart = AutoTokenizer.from_pretrained("sshleifer/distilbart-xsum-12-6")

    masker_mt = maskers.Text(tokenizer_mt)
    masker_gpt = maskers.Text(tokenizer_gpt)
    masker_bart = maskers.Text(tokenizer_bart)

    masker_mt_expected_keep_prefix, masker_mt_expected_keep_suffix = 0, 1
    masker_gpt_expected_keep_prefix, masker_gpt_expected_keep_suffix = 0, 0
    masker_bart_expected_keep_prefix, masker_bart_expected_keep_suffix = 1, 1

    assert masker_mt.keep_prefix == masker_mt_expected_keep_prefix and masker_mt.keep_suffix == masker_mt_expected_keep_suffix and \
           masker_gpt.keep_prefix == masker_gpt_expected_keep_prefix and masker_gpt.keep_suffix == masker_gpt_expected_keep_suffix and \
           masker_bart.keep_prefix == masker_bart_expected_keep_prefix and masker_bart.keep_suffix == masker_bart_expected_keep_suffix
<commit_msg>Added tests for text infilling<commit_after>''' This file contains tests for the Text masker.
'''

def test_method_tokenize_pretrained_tokenizer():
    import shap
    from transformers import AutoTokenizer, AutoModelForSequenceClassification

    tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased",use_fast=False)
    masker = shap.maskers.Text(tokenizer)

    test_text = "I have a joke about deep learning but I can't explain it."
    output_ids = masker.tokenize(test_text)['input_ids']
    correct_ids = tokenizer.encode_plus(test_text)['input_ids']

    print("Test logging",output_ids)
    print(correct_ids)
    assert output_ids == correct_ids

def test_method_tokenize_pretrained_tokenizer_fast():
    import shap
    from transformers import AutoTokenizer, AutoModelForSequenceClassification

    tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased",use_fast=True)
    masker = shap.maskers.Text(tokenizer)

    test_text = "I have a joke about deep learning but I can't explain it."
    output_ids = masker.tokenize(test_text)['input_ids']
    correct_ids = tokenizer.encode_plus(test_text)['input_ids']

    assert output_ids == correct_ids

def test_method_token_segments_pretrained_tokenizer():
    import shap
    from transformers import AutoTokenizer, AutoModelForSequenceClassification

    tokenizer = AutoTokenizer.from_pretrained("distilbert-base-cased",use_fast=False)
    masker = shap.maskers.Text(tokenizer)

    test_text = "I ate a Cannoli"
    output_token_segments = masker.token_segments(test_text)
    correct_token_segments = ['','I','ate','a','Can','##no','##li','']
    print(output_token_segments)

    assert output_token_segments == correct_token_segments

def test_method_token_segments_pretrained_tokenizer_fast():
    import shap
    from transformers import AutoTokenizer, AutoModelForSequenceClassification

    tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased",use_fast=True)
    masker = shap.maskers.Text(tokenizer)

    test_text = "I ate a Cannoli"
    output_token_segments = masker.token_segments(test_text)
    correct_token_segments = ['','I ','ate ','a ','Can','no','li','']

    assert output_token_segments == correct_token_segments

def test_masker_call_pretrained_tokenizer():
    import shap
    from transformers import AutoTokenizer, AutoModelForSequenceClassification
    import numpy as np

    tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased",use_fast=False)
    masker = shap.maskers.Text(tokenizer)

    test_text = "I ate a Cannoli"
    test_input_mask = np.array([True, False, True, True, False, True, True, True])

    output_masked_text = masker(test_input_mask,test_text)
    correct_masked_text = '[MASK] ate a [MASK]noli'

    assert output_masked_text == correct_masked_text

def test_masker_call_pretrained_tokenizer_fast():
    import shap
    from transformers import AutoTokenizer, AutoModelForSequenceClassification
    import numpy as np

    tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased",use_fast=True)
    masker = shap.maskers.Text(tokenizer)

    test_text = "I ate a Cannoli"
    test_input_mask = np.array([True, False, True, True, False, True, True, True])

    output_masked_text = masker(test_input_mask,test_text)
    correct_masked_text = '[MASK]ate a [MASK]noli'

    assert output_masked_text == correct_masked_text

def test_sentencepiece_tokenizer_output():
    """ Tests for output for sentencepiece tokenizers to not have '_' in output of masker when passed a mask of ones.
    """
    import numpy as np
    from transformers import AutoTokenizer
    from shap import maskers
    from shap.utils import safe_isinstance

    tokenizer = AutoTokenizer.from_pretrained("Helsinki-NLP/opus-mt-en-es")
    masker = maskers.Text(tokenizer)

    s="This is a test statement for sentencepiece tokenizer"
    masker._update_s_cache(s)
    mask = np.ones(masker._segments_s.shape, dtype=bool)

    sentencepiece_tokenizer_output_processed = masker(mask, s)
    expected_sentencepiece_tokenizer_output_processed = "This is a test statement for sentencepiece tokenizer"
    # since we expect output wrapped in a tuple hence the indexing [0][0] to extract the string
    assert sentencepiece_tokenizer_output_processed[0][0] == expected_sentencepiece_tokenizer_output_processed

def test_keep_prefix_suffix_tokenizer_parsing():
    """ Checks parsed keep prefix and keep suffix for different tokenizers.
    """
    from transformers import AutoTokenizer
    from shap import maskers

    tokenizer_mt = AutoTokenizer.from_pretrained("Helsinki-NLP/opus-mt-en-es")
    tokenizer_gpt = AutoTokenizer.from_pretrained("gpt2")
    tokenizer_bart = AutoTokenizer.from_pretrained("sshleifer/distilbart-xsum-12-6")

    masker_mt = maskers.Text(tokenizer_mt)
    masker_gpt = maskers.Text(tokenizer_gpt)
    masker_bart = maskers.Text(tokenizer_bart)

    masker_mt_expected_keep_prefix, masker_mt_expected_keep_suffix = 0, 1
    masker_gpt_expected_keep_prefix, masker_gpt_expected_keep_suffix = 0, 0
    masker_bart_expected_keep_prefix, masker_bart_expected_keep_suffix = 1, 1

    assert masker_mt.keep_prefix == masker_mt_expected_keep_prefix and masker_mt.keep_suffix == masker_mt_expected_keep_suffix and \
           masker_gpt.keep_prefix == masker_gpt_expected_keep_prefix and masker_gpt.keep_suffix == masker_gpt_expected_keep_suffix and \
           masker_bart.keep_prefix == masker_bart_expected_keep_prefix and masker_bart.keep_suffix == masker_bart_expected_keep_suffix

def test_text_infill_with_collapse_mask_token():
    """ Tests for different text infilling output combinations with collapsing mask token.
    """
    import numpy as np
    from shap import maskers
    from transformers import AutoTokenizer

    tokenizer = AutoTokenizer.from_pretrained("gpt2")
    masker = maskers.Text(tokenizer, mask_token='...', collapse_mask_token=True)

    s = "This is a test string to be infilled"

    # s_masked_with_infill_ex1 = "This is a test string ... ... ..."
    mask_ex1 = np.array([True, True, True, True, True, False, False, False, False])
    # s_masked_with_infill_ex2 = "This is a ... ... to be infilled"
    mask_ex2 = np.array([True, True, True, False, False, True, True, True, True])
    # s_masked_with_infill_ex3 = "... ... ... test string to be infilled"
    mask_ex3 = np.array([False, False, False, True, True, True, True, True, True])
    # s_masked_with_infill_ex4 = "... ... ... ... ... ... ... ..."
    mask_ex4 = np.array([False, False, False, False, False, False, False, False, False])

    text_infilled_ex1 = masker(mask_ex1, s)[0][0]
    expected_text_infilled_ex1 = "This is a test string ..."

    text_infilled_ex2 = masker(mask_ex2, s)[0][0]
    expected_text_infilled_ex2 = "This is a ... to be infilled"

    text_infilled_ex3 = masker(mask_ex3, s)[0][0]
    expected_text_infilled_ex3 = "... test string to be infilled"

    text_infilled_ex4 = masker(mask_ex4, s)[0][0]
    expected_text_infilled_ex4 = "..."

    assert  text_infilled_ex1 == expected_text_infilled_ex1 and text_infilled_ex2 == expected_text_infilled_ex2 and \
            text_infilled_ex3 == expected_text_infilled_ex3 and text_infilled_ex4 == expected_text_infilled_ex4<|endoftext|>


<commit_before>#!/usr/bin/env python
"""
Copyright (c) 2013, Konrad Foerstner <konrad@foerstner.org>

Permission to use, copy, modify, and/or distribute this software for
any purpose with or without fee is hereby granted, provided that the
above copyright notice and this permission notice appear in all
copies.

THE SOFTWARE IS PROVIDED 'AS IS' AND THE AUTHOR DISCLAIMS ALL
WARRANTIES WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED
WARRANTIES OF MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE
AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL
DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR
PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR OTHER
TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR
PERFORMANCE OF THIS SOFTWARE.

"""
__description__ = ("Calculate the Pearson correlation coefficient "
                   "for the coverages of two wiggle files.")
__author__ = "Konrad Foerstner <konrad@foerstner.org>"
__copyright__ = "2013 by Konrad Foerstner <konrad@foerstner.org>"
__license__ = "ISC license"
__email__ = "konrad@foerstner.org"
__version__ = ""

import argparse

from wiggle import WiggleParser
import numpy as np
from scipy import stats
import matplotlib.pyplot as plt

def main():
    parser = argparse.ArgumentParser(description=__description__)
    parser.add_argument("wiggle_file_1", type=argparse.FileType("r"))
    parser.add_argument("wiggle_file_2", type=argparse.FileType("r"))
    args = parser.parse_args()
    wiggel_correlator = WiggleCorrelator()
    wiggel_correlator.correlate(args.wiggle_file_1, args.wiggle_file_2)

class WiggleCorrelator(object):

    def __init__(self):
        self._wiggle_parser = WiggleParser()

    def correlate(self, wiggle_file_1, wiggle_file_2):
        self._chrom_and_pos_value_pairs = {}
        print("Replicon: Pearson correlation coefficient (p-value)")
        for entry_1, entry_2 in zip(
                self._wiggle_parser.entries(wiggle_file_1),
                self._wiggle_parser.entries(wiggle_file_2)):
                assert(entry_1.chrom_name == entry_2.chrom_name)
                pos_value_pairs_1 = dict(entry_1.pos_value_pairs)
                pos_value_pairs_2 = dict(entry_2.pos_value_pairs)
                if len(pos_value_pairs_1) == 0 or len(pos_value_pairs_2) == 0:
                    print("%s: At least one replicon has no coverage for "
                          "this libs." % (entry_2.chrom_name))
                    continue
                non_redu_pos = set(
                    pos_value_pairs_1.keys() + pos_value_pairs_2.keys())
                values_1 = np.array(
                    [pos_value_pairs_1.get(pos, 0.0) for pos in non_redu_pos])
                values_2 = np.array(
                    [pos_value_pairs_2.get(pos, 0.0) for pos in non_redu_pos])
                pearson, pvalue = stats.pearsonr(values_1, values_2)
                print("%s: %s (%s)" % (entry_1.chrom_name, pearson, pvalue))

if __name__ == "__main__":
   main()

<commit_msg>Apply name changes in scripts<commit_after>#!/usr/bin/env python
"""
Copyright (c) 2013, Konrad Foerstner <konrad@foerstner.org>

Permission to use, copy, modify, and/or distribute this software for
any purpose with or without fee is hereby granted, provided that the
above copyright notice and this permission notice appear in all
copies.

THE SOFTWARE IS PROVIDED 'AS IS' AND THE AUTHOR DISCLAIMS ALL
WARRANTIES WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED
WARRANTIES OF MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE
AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL
DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR
PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR OTHER
TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR
PERFORMANCE OF THIS SOFTWARE.

"""
__description__ = ("Calculate the Pearson correlation coefficient "
                   "for the coverages of two wiggle files.")
__author__ = "Konrad Foerstner <konrad@foerstner.org>"
__copyright__ = "2013 by Konrad Foerstner <konrad@foerstner.org>"
__license__ = "ISC license"
__email__ = "konrad@foerstner.org"
__version__ = ""

import argparse

from wiggle import WiggleParser
import numpy as np
from scipy import stats
import matplotlib.pyplot as plt

def main():
    parser = argparse.ArgumentParser(description=__description__)
    parser.add_argument("wiggle_file_1", type=argparse.FileType("r"))
    parser.add_argument("wiggle_file_2", type=argparse.FileType("r"))
    args = parser.parse_args()
    wiggel_correlator = WiggleCorrelator()
    wiggel_correlator.correlate(args.wiggle_file_1, args.wiggle_file_2)

class WiggleCorrelator(object):

    def __init__(self):
        self._wiggle_parser = WiggleParser()

    def correlate(self, wiggle_file_1, wiggle_file_2):
        self._chrom_and_pos_value_pairs = {}
        print("Replicon: Pearson correlation coefficient (p-value)")
        for entry_1, entry_2 in zip(
                self._wiggle_parser.entries(wiggle_file_1),
                self._wiggle_parser.entries(wiggle_file_2)):
                assert(entry_1.replicon == entry_2.replicon)
                pos_value_pairs_1 = dict(entry_1.pos_value_pairs)
                pos_value_pairs_2 = dict(entry_2.pos_value_pairs)
                if len(pos_value_pairs_1) == 0 or len(pos_value_pairs_2) == 0:
                    print("%s: At least one replicon has no coverage for "
                          "this libs." % (entry_2.replicon))
                    continue
                non_redu_pos = set(
                    pos_value_pairs_1.keys() + pos_value_pairs_2.keys())
                values_1 = np.array(
                    [pos_value_pairs_1.get(pos, 0.0) for pos in non_redu_pos])
                values_2 = np.array(
                    [pos_value_pairs_2.get(pos, 0.0) for pos in non_redu_pos])
                pearson, pvalue = stats.pearsonr(values_1, values_2)
                print("%s: %s (%s)" % (entry_1.replicon, pearson, pvalue))

if __name__ == "__main__":
   main()

<|endoftext|>