source: ferringb/snakeoil-dev/snakeoil/caching_2to3.py @ ferringb@gmail.com-20100215031044-boc6qr4uvsnzm3l3

#!/usr/bin/python3
# Copyright: 2009 Brian Harring <ferringb@gmail.com>
# License: PSF-2.2/GPL2/BSD

import lib2to3.main
import lib2to3.refactor
import os, hashlib

def md5_hash_data(data):
    chf = hashlib.md5()
    chf.update(data)
    return chf.hexdigest()

class caching_mixin(object):

    base_cls = None

    @property
    def cache_dir(self):
        return os.environ.get("PY2TO3_CACHEDIR", "cache")

    def get_cache_path(self, cache_key):
        return os.path.join(self.cache_dir, cache_key)

    def update_cache_from_file(self, cache_key, filename, encoding, new_text=None):
        cache_dir = self.cache_dir
        if not os.path.exists(cache_dir):
            os.mkdir(cache_dir)
        if new_text is None:
            new_text = open(filename, 'rb').read().decode(encoding)
        open(os.path.join(cache_dir, cache_key), 'wb').write(new_text.encode(encoding))

    def check_cache(self, cache_key, encoding):
        cache_path = self.get_cache_path(cache_key)
        if os.path.isfile(cache_path):
            return open(cache_path, 'rb').read().decode(encoding)
        return None

    @staticmethod
    def compute_cache_key(input, encoding):
        return md5_hash_data(input.encode(encoding))

    def refactor_file(self, filename, write=False, doctests_only=False):
        input, encoding = self._read_python_source(filename)
        cache_key = self.compute_cache_key(input, encoding)
        cache_data = self.check_cache(cache_key, encoding)

        if not write or cache_data is None:
            return super(caching_mixin, self).refactor_file(filename, write=write,
                doctests_only=doctests_only)
        else:
            self.processed_file(cache_data, filename, write=write,
                encoding=encoding, old_text=input)

    def processed_file(self, new_text, filename, old_text=None, write=False,
        encoding=None):
        if write:
            if old_text is None:
                cache_key = self.compute_cache_key(*self._read_python_source(filename))
            else:
                cache_key = self.compute_cache_key(old_text, encoding)
            self.update_cache_from_file(cache_key, filename, encoding,
                new_text=new_text)
        return super(caching_mixin, self).processed_file(new_text, filename,
            old_text=old_text, write=write, encoding=encoding)


class RefactoringTool(caching_mixin, lib2to3.refactor.RefactoringTool):
    pass

class MultiprocessRefactoringTool(caching_mixin, lib2to3.refactor.MultiprocessRefactoringTool):
    pass


def StdoutRefactoringTool(*args):
    # stupid hacks...
    lib2to3.main.StdoutRefactoringTool = my_StdoutRefactoringTool.base_cls
    inst = my_StdoutRefactoringTool.base_cls(*args)
    inst.__class__ = my_StdoutRefactoringTool
    return inst

if __name__ == '__main__':
    lib2to3.main.StdoutRefactoringTool = StdoutRefactoringTool
    import sys
    sys.exit(lib2to3.main.main("lib2to3.fixes"))