References and includes resolve themselves

madeleine/base.py

@@ -67,7 +67,8 @@ class Component(metaclass=ComponentMetaclass, register=False):
                     import yaml
                 except ImportError as e:
                     raise ImportError(
-                        '{}\nYou may need to install madeleine with YAML support: '
+                        '{}\n'
+                        'You may need to install madeleine with YAML support: '
                         'pip install madeleine[yaml]'.format(e.message),
                     )
                 data = yaml.safe_load(f)
@@ -76,7 +77,8 @@ class Component(metaclass=ComponentMetaclass, register=False):
                     import toml
                 except ImportError as e:
                     raise ImportError(
-                        '{}\nYou may need to install madeleine with TOML support: '
+                        '{}\n'
+                        'You may need to install madeleine with TOML support: '
                         'pip install madeleine[toml]'.format(e.message),
                     )
                 data = toml.load(f)
@@ -114,13 +116,25 @@ class Reference(Component, key='ref'):
     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):
-        raise TypeError('Unresolved reference to {!r}'.format(self.ref))
+        if not self.to:
+            raise TypeError('Unresolved reference to {!r}'.format(self.ref))
+        return self.to.combinations
 
     def generate(self):
-        raise TypeError('Unresolved reference to {!r}'.format(self.ref))
+        if not self.to:
+            raise TypeError('Unresolved reference to {!r}'.format(self.ref))
+        return self.to.generate()
 
 
 class Value(Component, key='value'):
@@ -144,14 +158,25 @@ class Value(Component, key='value'):
             return self.value
 
 
-class Include(Component, type, key='include'):
+class Include(Component, key='include'):
     """
     Automatically include another generator file.
     """
 
-    def __new__(*args, include=None, format=None, **kwargs):
+    def __init__(self, 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)
+        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()

madeleine/multiple.py

@@ -1,6 +1,6 @@
 from functools import reduce, partial
 from itertools import combinations, combinations_with_replacement
-from madeleine import Component, Reference
+from madeleine import Component
 from madeleine.helpers import binom
 import operator
 import random
@@ -35,8 +35,6 @@ class Repeat(Component, key='repeat'):
         super().__init__(**kwargs)
 
     def resolve_references(self, references):
-        if isinstance(self.repeat, Reference):
-            self.repeat = references[self.repeat.ref]
         self.repeat.resolve_references(references)
 
     @property
@@ -81,10 +79,8 @@ class CompoundComponent(Component, register=False):
         super().__init__(**data)
 
     def resolve_references(self, references):
-        for i in range(len(self.items)):
-            if isinstance(self.items[i], Reference):
-                self.items[i] = references[self.items[i].ref]
-            self.items[i].resolve_references(references)
+        for item in self.items:
+            item.resolve_references(references)
 
 
 class AllOf(CompoundComponent, key='allOf'):