ghc-mod-cache

#!/usr/bin/env python

from __future__ import print_function
import sys, os, re
import subprocess, hashlib

debugging = False

def debug_print(m):
    if debugging:
        print(m, file=sys.stderr)

def get_project_dir():
    pw = os.getcwd()
    while True:
        if pw == "/":
            return (".", None, None)
        for f in os.listdir(pw):
            if re.search("\.cabal$", f):
                if os.path.isfile(os.path.join(pw, f)):
                    if os.path.isfile(os.path.join(pw, "stack.yaml")):
                        return (pw, f, "stack.yaml")
                    else:
                        return (pw, f, None)
        pw = os.path.normpath(pw + "/..")

def encode_options(args):
    m = hashlib.sha224()
    module_name = None
    for s in args:
        if s and s[0] == '-':
            m.update(s)
        else:
            if module_name == None:
                module_name = s
            else:
                print("ERROR: multiple module names?", file=sys.stderr)
    if module_name:
        m.update(s)
    return (module_name, m.hexdigest())

def check_if_module_is_user(project_dir, module_name):
    if module_name == None or project_dir == "/":
        return (False, None)
    s = module_name
    s = s.replace(".", "/")
    s += ".hs"
    s = os.path.join(project_dir, s)
    return (os.path.exists(s), s)

def cache_outdated(cache_file, user_file):
    if cache_file == None or user_file == None: return False
    try:
        cf_time = os.path.getmtime(cache_file)
        uf_time = os.path.getmtime(user_file)
        return cf_time < uf_time
    except:
        return True

def record_ghc_mod_invocation(args):
    if not debugging: return
    s = os.environ['HOME']
    log_file = os.path.join(s, ".ghcmodcache_log")
    with open(log_file, "a") as f:
        print(args, file=f)

def get_module_name_from_module_version_condition_string(s):
    """ get the module name from a module version codition string such as 'base < 4.7 && >= 4.0'
    """
    s = s.strip()
    r = re.match(r'(.*?)[^\w-]', s)
    if r: return r.group(1)
    return s

def get_depencencies_from_yaml_dict(obj):
    retval = []
    for k, v in obj.items():
        if k == 'dependencies':
            retval += map(lambda x: get_module_name_from_module_version_condition_string(x), v)
            continue
        if isinstance(v, dict):
            tmp = get_depencencies_from_yaml_dict(v)
            retval += tmp
    return retval

def parse_stack_yaml(stack_yaml_path):
    import yaml
    debug_print("Imported yaml module")
    with open(stack_yaml_path, "r") as f:
        stack_yaml = yaml.load(f)
    debug_print("Loaded stack.yaml")
    return stack_yaml

def load_build_plan_cache(resolver):
    import json
    # The line below is duplicated to ensure_there_is_build_plan_cache()
    build_plan_json_file_path = os.path.join(os.path.expanduser("~/.stack/build-plan"), resolver + ".json")
    try:
        with open(build_plan_json_file_path, "r") as f:
            build_plan_json = json.load(f)
    except IOError as e:
        print("ERROR:", e, file=sys.stderr)
        sys.exit(2)
    return build_plan_json

def ensure_there_is_build_plan_cache(resolver):
    build_plan_file_path = os.path.join(os.path.expanduser("~/.stack/build-plan"), resolver + ".yaml")
    # The line below is duplicated to load_build_plan_cache()
    build_plan_json_file_path = os.path.join(os.path.expanduser("~/.stack/build-plan"), resolver + ".json")
    if os.path.exists(build_plan_json_file_path) and not cache_outdated(build_plan_json_file_path, build_plan_file_path):
        return
    build_plan_json_file_path_tmp = build_plan_json_file_path + ".tmp"
    if not os.path.exists(build_plan_file_path):
        subprocess.check_call(['stack', 'setup'])
    try:
        from yaml import CLoader
        debug_print("Loading build plan YAML using CLoader")
        with open(build_plan_file_path, "r") as f:
            build_plan_yaml = yaml.load(f, Loader=CLoader)
    except ImportError as e:
        debug_print("Loading build plan YAML using pure Python parser")
        import yaml
        with open(build_plan_file_path, "r") as f:
            build_plan_yaml = yaml.load(f)
    debug_print("Loaded build plan YAML")
    packages = build_plan_yaml['packages']
    modules = []
    for k, v in packages.items():
        module = {'Name' : k, 'Modules': v['description']['modules']}
        modules.append(module)
    import json
    with open(build_plan_json_file_path_tmp, "w") as f:
        json.dump(modules, f)
    # This is to atomically create a cache file...
    # If you hit CTRL+C while converting, still you would not see the final file.
    os.rename(build_plan_json_file_path_tmp, build_plan_json_file_path)

def main():
    # parse options
    if 2 <= len(sys.argv):
        if sys.argv[1] == "-d" or sys.argv[1] == "--debug":
            global debugging
            debugging = True
            del sys.argv[1]

    # check if using stack
    project_dir, cabal_file_name, stack_file_name = get_project_dir()
    if stack_file_name == None:
        call_ghc_mod = ["ghc-mod"]
    else:
        call_ghc_mod = ["stack", "exec", "--", "ghc-mod"]
    debug_print("Project dir: " + project_dir)

    # do a main job
    if len(sys.argv) < 2:
        estat = subprocess.call(call_ghc_mod + sys.argv[1:])
        record_ghc_mod_invocation(sys.argv)
        sys.exit(estat >> 8)

    sub_command = sys.argv[1]

    cabal_abs_file_path = os.path.join(project_dir, cabal_file_name) if cabal_file_name else None
    package_yaml_abs_file_path = os.path.join(project_dir, "package.yaml")
    stack_abs_file_path = os.path.join(project_dir, stack_file_name) if stack_file_name else None
    ghc_mod_cache_dir = os.path.join(project_dir, ".ghcmodcache")

    if sub_command == "list_stackage":
        option_str = sys.argv[2] if 2 <= len(sys.argv) else None
        if stack_file_name == None:
            print("'ghc-mod-cache install_module' does not work with non-stack projects", file=sys.stderr)
            sys.exit(2)
        try:
            stack_yaml = parse_stack_yaml(stack_abs_file_path)
            if not 'resolver' in stack_yaml:
                print("ERROR: no resolver in stack.yaml", file=sys.stderr)
                sys.exit(2)
            resolver = stack_yaml['resolver']
            ensure_there_is_build_plan_cache(resolver)
            modules = load_build_plan_cache(resolver)
        except Exception as e:
            debug_print("Error")
            print("There was an exception: ", e, file=sys.stderr)
            sys.exit(2)
        for d in modules:
            print(d['Name'])
            if option_str == '--modules': continue
            for m in d['Modules']:
                print("  " + m)
        sys.exit(0)

    if sub_command == "install_module":
        if stack_file_name == None:
            print("'ghc-mod-cache install_module' does not work with non-stack projects", file=sys.stderr)
            sys.exit(2)
        if len(sys.argv) < 3:
            print("Usage: ghc-mod-cache install_module <module name>", file=sys.stderr)
            sys.exit(2)
        module_name = sys.argv[2]
        if os.path.exists(ghc_mod_cache_dir):
            for fn in os.listdir(ghc_mod_cache_dir):
                if fn.startswith("-"):
                    file_name = os.path.join(ghc_mod_cache_dir, fn)
                    os.unlink(file_name)
                    debug_print("Removed " + file_name)
        cmd_line = ['stack', 'build', module_name]
        debug_print("CMD: " + (" ".join(cmd_line)))
        estat = subprocess.call(cmd_line)
        if estat != 0: sys.exit(2)
        sys.exit(0)

    if sub_command == "browse" or sub_command == "modules":
        if not os.path.exists(ghc_mod_cache_dir):
            os.mkdir(ghc_mod_cache_dir)
        module_name, canonical_option_coded = encode_options(sys.argv[2:])
        (is_user_module, module_path_if_any) = check_if_module_is_user(project_dir, module_name)
        cached_file_name = os.path.join(ghc_mod_cache_dir, (module_name if module_name else "") + "-" + canonical_option_coded)
        cache_did_not_exist = not os.path.exists(cached_file_name)
        cached_member_list_was_outdated = is_user_module and cache_outdated(cached_file_name, module_path_if_any)
        cached_module_list_was_outdated = module_name == None and (cache_outdated(cached_file_name, cabal_abs_file_path) or
                                                                   cache_outdated(cached_file_name, stack_abs_file_path) or
                                                                   cache_outdated(cached_file_name, package_yaml_abs_file_path))
        debug_print("Module name: " + (module_name if module_name else "(None)"))
        debug_print("Cache dir: " + ("no" if cache_did_not_exist else "yes"))
        debug_print("Member list: " + ("need to create" if cached_member_list_was_outdated else "valid"))
        debug_print("Module list: " + ("need to create" if cached_module_list_was_outdated else "valid"))
        if cache_did_not_exist or cached_member_list_was_outdated or cached_module_list_was_outdated:
            if cached_module_list_was_outdated:
                # Probabpy we need to check updated module lists.
                try:
                    import yaml
                    debug_print("Imported yaml module")
                    with open(package_yaml_abs_file_path, "r") as f:
                        package_yaml = yaml.load(f)
                        module_names = {}
                        for m in get_depencencies_from_yaml_dict(package_yaml):
                            module_names[m] = None
                        if 'name' in package_yaml:
                            my_module_name = package_yaml['name']
                            if my_module_name in module_names: del module_names[my_module_name]
                        cmd_line = ['stack', 'build'] + module_names.keys()
                        debug_print("CMD: " + (" ".join(cmd_line)))
                        estat = subprocess.call(cmd_line)
                    debug_print("Done")
                except Exception as e:
                    debug_print("Error")
                    print("There was an exception: ", e, file=sys.stderr)
                    sys.exit(2)
            cached_file_name_tmp = cached_file_name + ".tmp"
            with open(cached_file_name_tmp, "w") as f:
                record_ghc_mod_invocation(sys.argv)
                if is_user_module:
                    os.chdir(project_dir)
                    estat = subprocess.call(call_ghc_mod + sys.argv[1:], stdout=f)
                else:
                    estat = subprocess.call(call_ghc_mod + sys.argv[1:], stdout=f)
            if estat != 0: sys.exit(estat >> 8)
            # The line below is for not leaving a temporary file created by ghc-mod process failed
            os.rename(cached_file_name_tmp, cached_file_name)
        with open(cached_file_name, "r") as f:
            sys.stdout.write(f.read())
            sys.exit(0)
        sys.exit(2)

    record_ghc_mod_invocation(sys.argv)
    estat = subprocess.call(call_ghc_mod + sys.argv[1:])
    sys.exit(estat >> 8)

if __name__ == "__main__":
    main()