flatfilecms/flatfilecms/models.py

from pyramid.path import AssetResolver
from pathlib import PurePath
import yaml
import frontmatter
import markdown
from zope.interface import (Interface, implementer)

from flatfilecms.filters import IgnoreExtension


class Loader(yaml.Loader):
    def __init__(self, stream):
        super(Loader, self).__init__(stream)
        Loader.add_constructor('!include', Loader.include)
        Loader.add_constructor('!markdown', Loader.markdown)

    def include(self, node):
        if isinstance(node, yaml.ScalarNode):
            return self.extractFile(self.construct_scalar(node))

        elif isinstance(node, yaml.SequenceNode):
            return [self.extractFile(filename)
                    for filename in self.construct_sequence(node)]

        elif isinstance(node, yaml.MappingNode):
            result = {}
            for k, v in self.construct_mapping(node).items():
                result[k] = self.extractFile(v)
            return result

        else:
            raise yaml.constructor.ConstructorError(
                    "Error:: unrecognised node type in !include statement")

    def markdown(self, node):
        if not isinstance(node, yaml.ScalarNode):
            raise yaml.constructor.ConstructorError(
                    "Error:: unrecognised node type in !markdown statement")
        m = self.construct_scalar(node)
        return markdown.markdown(
            m,
            extensions=[IgnoreExtension(), 'markdown.extensions.extra'],
            output_format='html5',
            tab_length=2)

    def extractFile(self, filename):
        path = PurePath(self.data_path) / filename
        f = AssetResolver().resolve(str(path)).stream()
        if f:
            if path.suffix in ['.yaml', '.yml', '.json']:
                return yaml.load(f, Loader)
            return f.read().decode()


def LoaderFactory(data_path):
    cl = Loader
    cl.data_path = data_path
    return cl


class CustomYAMLHandler(frontmatter.YAMLHandler):
    def __init__(self, data_path):
        self.loader = LoaderFactory(data_path)

    def load(self, fm, **kwargs):
        return yaml.load(fm, self.loader)


class Folder(dict):
    def __init__(self, name, parent, path):
        self.path = path
        self.name = name
        self.parent = parent
        if parent is not None:
            self.data_path = parent.data_path
        for entry in AssetResolver().resolve(path).listdir():
            asset = f"{path}/{entry}"
            if AssetResolver().resolve(asset).isdir():
                self.create_dir(asset)
            else:
                self.create_file(asset)

    def create_file(self, asset):
        path = PurePath(asset)
        if path.suffix == '.md':
            self[path.stem] = Markdown(path.stem, self, asset)
        elif path.suffix == '.yaml':
            self[path.stem] = YAML(path.stem, self, asset)
        elif path.suffix == '.j2' or path.suffix == '.jinja2':
            self[path.stem] = Jinja2(path.stem, self, asset)
        else:
            name = path.name
            # Если имя файла не index.html,
            # то отдавать по имени файла
            if name == 'index.html':
                name = 'index'
            self[name] = Document(name, self, asset)

    def create_dir(self, asset):
        path = PurePath(asset)
        self[path.name] = Folder(path.name, self, asset)

    def walk(self):
        for item in self.values():
            yield item
            if isinstance(item, Folder):
                yield from item.walk()


class Document(object):
    def __init__(self, name, parent, path):
        self.name = name
        self.parent = parent
        self.path = path
        self.data_path = parent.data_path


class IMarkdown(Interface):
    """ Интерфейс-маркер для генератора markdown """


@implementer(IMarkdown)
class Markdown(Document):
    def __init__(self, name, parent, path):
        super(Markdown, self).__init__('', parent, path)
        self.page = frontmatter.load(
                AssetResolver().resolve(self.path).stream(),
                handler=CustomYAMLHandler(self.data_path),
                ).to_dict()


@implementer(IMarkdown)
class YAML(Document):
    def __init__(self, name, parent, path):
        super(YAML, self).__init__('', parent, path)
        self.page = yaml.load(
                AssetResolver().resolve(self.path).stream(),
                LoaderFactory(self.data_path))


class Jinja2(Document):
    pass