<commit_before>import cPickle
from srw.util import cache
import os
import unittest

class TestCache(unittest.TestCase):

    cache_stub = 'cachestub'
    cache_name = '.{0}.cpickle'.format(cache_stub)

    def setUp(self):
        if os.path.isfile(self.cache_name):
            os.remove(self.cache_name)

    @classmethod
    def setUpClass(cls):
        @cache(cls.cache_stub)
        def cached_function(cls):
            return 10

        cls.cached_function = cached_function

    def test_that_a_value_without_a_cache_creates_one(self):
        self.cached_function()
        assert os.path.isfile(self.cache_name)

    def test_cache_contents(self):
        self.cached_function()
        self.cached_function()

        assert cPickle.load(open(self.cache_name)) == 10

def test_custom_cache_directory():
    cache_stub = 'test'
    cache_path = '/tmp/.{0}.cpickle'.format(cache_stub)

    @cache(cache_stub, directory=os.path.dirname(cache_path))
    def do_test():
        return 5

    if os.path.isfile(cache_path):
        os.remove(cache_path)

    assert do_test() == 5

    assert os.path.isfile(cache_path)
<commit_msg>Importing cpickle correctly in tests<commit_after>try:
    import cPickle as pickle
except ImportError:
    import pickle

from srw.util import cache
import os
import unittest

class TestCache(unittest.TestCase):

    cache_stub = 'cachestub'
    cache_name = '.{0}.cpickle'.format(cache_stub)

    def setUp(self):
        if os.path.isfile(self.cache_name):
            os.remove(self.cache_name)

    @classmethod
    def setUpClass(cls):
        @cache(cls.cache_stub)
        def cached_function(cls):
            return 10

        cls.cached_function = cached_function

    def test_that_a_value_without_a_cache_creates_one(self):
        self.cached_function()
        assert os.path.isfile(self.cache_name)

    def test_cache_contents(self):
        self.cached_function()
        self.cached_function()

        assert pickle.load(open(self.cache_name)) == 10

def test_custom_cache_directory():
    cache_stub = 'test'
    cache_path = '/tmp/.{0}.cpickle'.format(cache_stub)

    @cache(cache_stub, directory=os.path.dirname(cache_path))
    def do_test():
        return 5

    if os.path.isfile(cache_path):
        os.remove(cache_path)

    assert do_test() == 5

    assert os.path.isfile(cache_path)

{'fraction_non_alphanumeric': 0.1035653650254669, 'fraction_numerical': 0.006791171477079796, 'mean_word_length': 3.0153321976149914, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7457020', 'n_tokens_mistral': 800, 'n_tokens_neox': 776, 'n_words': 137}