234 lines
6.7 KiB
Python
234 lines
6.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
|
|
|
|
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,
|
|
}
|