158 lines
4.7 KiB
Python
158 lines
4.7 KiB
Python
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
|
|
|
|
@abstractproperty
|
|
def combinations(self):
|
|
"""
|
|
Should return how many combinations a single component may have.
|
|
Useful to perform some checks after parsing.
|
|
"""
|
|
|
|
@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(
|
|
'{}\nYou 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(
|
|
'{}\nYou 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.
|
|
"""
|
|
|
|
@property
|
|
def combinations(self):
|
|
raise TypeError('Unresolved reference to {!r}'.format(self.ref))
|
|
|
|
def generate(self):
|
|
raise TypeError('Unresolved reference to {!r}'.format(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
|
|
|
|
|
|
class Include(Component, type, key='include'):
|
|
"""
|
|
Automatically include another generator file.
|
|
"""
|
|
|
|
def __new__(*args, include=None, format=None, **kwargs):
|
|
if format is None:
|
|
format = FileFormat.YAML
|
|
if isinstance(format, str):
|
|
format = FileFormat(format)
|
|
return Component.from_path(include, fmt=format)
|