madeleine/madeleine/base.py

from abc import ABCMeta, abstractmethod, abstractproperty
from collections.abc import Mapping, Iterable
from enum import Enum
from pathlib import Path
from objtools.registry import ClassRegistry
import random


class FileFormat(Enum):
    JSON = 'json'
    YAML = 'yaml'
    TOML = 'toml'


class ComponentRegistry(ClassRegistry):

    def check_value(self, value):
        assert issubclass(value, Component)


registry = ComponentRegistry()
register = registry.register
unregister = registry.unregister


class ComponentMetaclass(registry.metaclass, ABCMeta):
    pass


class Component(metaclass=ComponentMetaclass, register=False):
    """
    Describes any generator component.
    """

    def __init__(self, **data):
        self.__dict__.update(**data)

    def resolve_references(self, references):
        """
        Complex components might hold references to other components:
        they should leave them unresolved until the Generator builds all the
        components, then calls this method with all of the referencable names.
        """

    @abstractmethod
    def generate(self):
        pass

    def serialize(self):
        """
        Should return a JSON/YAML/TOML-compatible structure which can be used
        to export a component to a file.
        """
        return self.__dict__.copy()

    @abstractproperty
    def combinations(self):
        """
        Should return how many combinations a single component may have.
        Useful to perform some checks after parsing.
        """

    def __str__(self):
        return str(self.serialize())

    def __repr__(self):
        return '{}({})'.format(
            self.__class__.__name__,
            ', '.join(
                '{}={!r}'.format(k, v)
                for k, v in self.serialize().items()
            ),
        )

    def _repr_pretty_(self, p, cycle):
        if cycle:
            p.text('{}(...)'.format(self.__class__.__name__))
        else:
            with p.group(2, '{}('.format(self.__class__.__name__), ')'):
                p.breakable('')
                for k, v in self.serialize().items():
                    p.text(k)
                    p.text('=')
                    p.pretty(v)
                    p.text(',')
                    p.breakable()

    @classmethod
    def from_path(cls, path, fmt=FileFormat.YAML):
        assert isinstance(fmt, FileFormat)
        if not isinstance(path, Path):
            path = Path(path)
        with path.open() as f:
            if fmt == FileFormat.JSON:
                import json
                data = json.load(f)
            elif fmt == FileFormat.YAML:
                try:
                    import yaml
                except ImportError as e:
                    raise ImportError(
                        '{}\n'
                        'You may need to install madeleine with YAML support: '
                        'pip install madeleine[yaml]'.format(e.message),
                    )
                data = yaml.safe_load(f)
            elif fmt == FileFormat.TOML:
                try:
                    import toml
                except ImportError as e:
                    raise ImportError(
                        '{}\n'
                        'You may need to install madeleine with TOML support: '
                        'pip install madeleine[toml]'.format(e.message),
                    )
                data = toml.load(f)
            else:
                raise NotImplementedError

            if cls is Component:
                # Guess the component type when calling Component.from_path
                return Component._make(data)
            return cls(**data)

    @staticmethod
    def _make(data):
        """
        Guess which component subclass to use for a given component data,
        then build and return the resulting components.
        """
        if isinstance(data, str):
            data = {'value': data}
        if not isinstance(data, Mapping):
            if isinstance(data, Iterable):
                return list(map(Component._make, data))
            raise ValueError('Component description should be a mapping')

        for key, component_class in registry.items():
            if key in data:
                return component_class(**data)

        raise ValueError('Could not parse component description')


class Reference(Component, key='ref'):
    """
    A component used to temporarily hold references to other components
    during the first pass of schema loading, when not all referenced names
    are available, before a second pass allows resolution of all references.
    """
    def __init__(self, **kwargs):
        self.to = None
        super().__init__(**kwargs)

    def resolve_references(self, references):
        if self.ref not in references:
            raise ValueError('Unknown reference {!r}'.format(self.ref))
        self.to = references[self.ref]

    @property
    def combinations(self):
        if not self.to:
            raise TypeError('Unresolved reference to {!r}'.format(self.ref))
        return self.to.combinations

    def generate(self):
        if not self.to:
            raise TypeError('Unresolved reference to {!r}'.format(self.ref))
        return self.to.generate()

    def serialize(self):
        return {
            'ref': self.ref,
        }


class Value(Component, key='value'):
    """
    A component that holds a string value.
    Will return the value on every generation, unless `optional` is True,
    in which case it may return None 50% of the time.
    """

    def __init__(self, value, optional=False, **kwargs):
        self.value = value
        self.optional = optional
        super().__init__(**kwargs)

    @property
    def combinations(self):
        return 1 + self.optional

    def generate(self):
        if not self.optional or random.randrange(2):
            return self.value

    def serialize(self):
        if self.optional:
            return {
                'optional': True,
                'value': self.value,
            }
        return self.value


class Include(Component, key='include'):
    """
    Automatically include another generator file.
    """

    def __init__(self, include=None, format=None, **kwargs):
        if format is None:
            format = FileFormat.YAML
        if isinstance(format, str):
            format = FileFormat(format)
        super().__init__(include=include, format=format, **kwargs)
        self.to = Component.from_path(self.include, fmt=self.format)

    def resolve_references(self, references):
        self.to.resolve_references(references)

    @property
    def combinations(self):
        return self.to.combinations

    def generate(self):
        return self.to.generate()

    def serialize(self):
        return {
            'include': self.include,
            'format': self.format,
        }