alucho/sitegen.py

import argparse
from functools import lru_cache
from importlib import import_module
from jinja2 import Environment, FileSystemLoader, contextfilter
from pathlib import Path
from os import walk, makedirs, listdir, symlink, readlink
from os.path import join, exists, splitext, split, islink, isdir
from shutil import rmtree, copy2, copystat, ignore_patterns, Error
from time import sleep
from traceback import print_exc
import os

#TODO: load from config file (and watch it too)
LANGUAGES = 'languages'
PLUGINS = 'plugins'
ROOT = 'root'
STATIC = 'static'
DEFAULT_LANG = 'en'
OTHER_LANGS = set(['es'])
WATCH_INTERVAL = 1 # in secs

config = {
    'LANGUAGES': LANGUAGES,
    'PLUGINS': PLUGINS,
    'ROOT': ROOT,
    'STATIC': STATIC,
    'DEFAULT_LANG': DEFAULT_LANG,
    'OTHER_LANGS': OTHER_LANGS,
    'WATCH_INTERVAL': WATCH_INTERVAL
}

# Utils

ignore_underscores = ignore_patterns('_*')

def copytree(src, dst, symlinks=False, ignore=None, copy_function=copy2,
             ignore_dangling_symlinks=False):
    """Recursively copy a directory tree.

    The destination directory must not already exist.
    If exception(s) occur, an Error is raised with a list of reasons.

    If the optional symlinks flag is true, symbolic links in the
    source tree result in symbolic links in the destination tree; if
    it is false, the contents of the files pointed to by symbolic
    links are copied. If the file pointed by the symlink doesn't
    exist, an exception will be added in the list of errors raised in
    an Error exception at the end of the copy process.

    You can set the optional ignore_dangling_symlinks flag to true if you
    want to silence this exception. Notice that this has no effect on
    platforms that don't support os.symlink.

    The optional ignore argument is a callable. If given, it
    is called with the `src` parameter, which is the directory
    being visited by copytree(), and `names` which is the list of
    `src` contents, as returned by os.listdir():

        callable(src, names) -> ignored_names

    Since copytree() is called recursively, the callable will be
    called once for each directory that is copied. It returns a
    list of names relative to the `src` directory that should
    not be copied.

    The optional copy_function argument is a callable that will be used
    to copy each file. It will be called with the source path and the
    destination path as arguments. By default, copy2() is used, but any
    function that supports the same signature (like copy()) can be used.

    """
    names = os.listdir(src)
    if ignore is not None:
        ignored_names = ignore(src, names)
    else:
        ignored_names = set()

    os.makedirs(dst, exist_ok=True)
    # `exist_ok=True`, very important for hot reloading
    errors = []
    for name in names:
        if name in ignored_names:
            continue
        srcname = os.path.join(src, name)
        dstname = os.path.join(dst, name)
        try:
            if os.path.islink(srcname):
                linkto = os.readlink(srcname)
                if symlinks:
                    # We can't just leave it to `copy_function` because legacy
                    # code with a custom `copy_function` may rely on copytree
                    # doing the right thing.
                    os.symlink(linkto, dstname)
                    copystat(srcname, dstname, follow_symlinks=not symlinks)
                else:
                    # ignore dangling symlink if the flag is on
                    if not os.path.exists(linkto) and ignore_dangling_symlinks:
                        continue
                    # otherwise let the copy occurs. copy2 will raise an error
                    if os.path.isdir(srcname):
                        copytree(srcname, dstname, symlinks, ignore,
                                 copy_function)
                    else:
                        copy_function(srcname, dstname)
            elif os.path.isdir(srcname):
                copytree(srcname, dstname, symlinks, ignore, copy_function)
            else:
                # Will raise a SpecialFileError for unsupported file types
                copy_function(srcname, dstname)
        # catch the Error from the recursive copytree so that we can
        # continue with other files
        except Error as err:
            errors.extend(err.args[0])
        except OSError as why:
            errors.append((srcname, dstname, str(why)))
    try:
        copystat(src, dst)
    except OSError as why:
        # Copying file access times may fail on Windows
        if getattr(why, 'winerror', None) is None:
            errors.append((src, dst, str(why)))
    if errors:
        raise Error(errors)
    return dst

def mtimes(target_dir):
    ''''get modification time of files in `target_dir`'''
    return { f: f.stat().st_mtime for f in Path(target_dir).rglob('*') }

def save_dic(lang, dic):
    with open(join(LANGUAGES, lang + '.txt'), 'w') as f:
        for key, value in dic.items():
            f.write(f"{key}:{value}\n")

@lru_cache(None)
def get_lang_dic(lang):
    makedirs(LANGUAGES, exist_ok=True)

    lang_file = join(LANGUAGES, lang + '.txt')

    if not exists(lang_file):
        return {}
    else:
        dic = {}
        with open(lang_file) as f:
            for line in f.readlines():
                key, value = line.strip('\n').split(':')
                dic[key] = value
        return dic

# Filters

def add_lang_prefix(lang, path):
    if lang == DEFAULT_LANG:
        return path

    if lang not in OTHER_LANGS:
        print(f"Not registered language: `{lang}`")

    return f"/{lang}{path}"

@contextfilter
def lang(ctx, value):
    lang = ctx.environment.globals['lang']

    if lang == DEFAULT_LANG:
        return value

    else:
        dic = get_lang_dic(lang)

        if value in dic and dic[value]:
            return dic[value]
        else:
            dic[value] = ""
            save_dic(lang, dic)
            print(f"Not translated phrase: `{value}`")
            return value

@contextfilter
def lang_url(ctx, lang):
    name, _ = splitext(ctx.name)
    if name.startswith('./'):
        name = name[2:]

    if name == 'index':
        name = ''

    return add_lang_prefix(lang, f"/{name}")

@contextfilter
def cur_lang(ctx, path):
    lang = ctx.environment.globals['lang']
    return add_lang_prefix(lang, path)

def static(value):
    return f"/static/{value}"

# Compiler

def compile(env, path, target):
    lang = env.globals['lang']

    for base, _, docs in walk(path):
        cur_base = Path(base)
        cur_base = Path(*cur_base.parts[1:])
        for doc in docs:
            if doc.startswith('_'):
                # Ignore this files
                continue

            src = join(base, doc)

            if args.with_index and doc != 'index.html':
                name, _ = splitext(doc)
                dst = join(cur_base, name, 'index.html')
            else:
                dst = join(cur_base, doc)

            if lang != DEFAULT_LANG:
                dst = join(lang, dst)

            dst = join(target, dst)
            dst_folder, _ = split(dst)

            makedirs(dst_folder, exist_ok=True)

            template = env.get_template(join(cur_base, doc))
            output = template.render()

            with open(dst, 'w') as f:
                f.write(output)

def init_gen(args):
    # Load layout
    file_loader = FileSystemLoader(args.source)
    env = args.env = Environment(loader=file_loader)

    # Init gobals
    env.globals['lang'] = DEFAULT_LANG

    # Add filters
    env.filters['lang'] = lang
    env.filters['lang_url'] = lang_url
    env.filters['cur_lang'] = cur_lang
    env.filters['static'] = static

    # Load plugins
    for mod_path in Path(PLUGINS).glob('*.py'):
        mod_name = '.'.join(mod_path.with_suffix('').parts)
        print(f'* loading {mod_name}')
        import_module(mod_name).init_plugin(env, config)
        print('  done!')

    # Clean target
    if exists(args.target):
        rmtree(args.target)
    makedirs(args.target, exist_ok=True)

def gen_layout(args):
    env = args.env
    path = Path(args.source)

    # Reset gobals
    env.globals['lang'] = DEFAULT_LANG

    # compile
    compile(env, path, args.target)
    for other_lang in OTHER_LANGS:
        env.globals['lang'] = other_lang
        compile(env, path, args.target)

def gen_root(args):
    copytree(ROOT, args.target, ignore=ignore_underscores)

def gen_static(args):
    copytree(STATIC, join(args.target, 'static'), ignore=ignore_underscores)

def save_generate(generator, args, msg):
    print(f'* {msg}')
    try:
        generator(args)
        print('  done!')
    except:
        print_exc()
        return False
    return True

# Runtime

def run(args):
    print('initial compilation')
    init_gen(args)
    save_generate(gen_layout, args, 'generating layout')
    save_generate(gen_root, args, 'copying root files')
    save_generate(gen_static, args, 'copying static dir')

    if not args.watch:
        return

    print(f'watching {args.source}/*:{LANGUAGES}/*:{STATIC}/*:{ROOT}/*')
    while True:
        # take modification times
        source_mt = mtimes(args.source)
        source_mt.update(mtimes(LANGUAGES))
        static_mt = mtimes(STATIC)
        root_mt = mtimes(ROOT)

        sleep(WATCH_INTERVAL)

        # compare modification times and recompile if different
        new_source_mt = mtimes(args.source)
        new_source_mt.update(mtimes(LANGUAGES))
        if source_mt != new_source_mt:
            source_mt = new_source_mt
            save_generate(gen_layout, args, 'recompiling layout')

        new_root_mt = mtimes(ROOT)
        if root_mt != new_root_mt:
            root_mt = new_root_mt
            save_generate(gen_root, args, 'copying root files')

        new_static_mt = mtimes(STATIC)
        if static_mt != new_static_mt:
            static_mt = new_static_mt
            save_generate(gen_static, args, 'copying static files')

if __name__ == '__main__':
    parser = argparse.ArgumentParser("Site generator")

    parser.add_argument('--watch', action='store_true', default=False)
    parser.add_argument('--with-index', action='store_true', default=False)
    parser.add_argument('--source', default='layout')
    parser.add_argument('--target', default='build')

    args = parser.parse_args()
    run(args)