Add utilities for registering custom YAML tags

Author: bmispelon

docs/guides/defining-template-context.md

@@ -203,3 +203,77 @@ def add_total(context, request):
         third_num = context.get('third_number', 0)
         context['total'] = first_num + second_num + third_num
 ```
+
+## Extending the YAML syntax
+
+You can also take advantage of YAML's local tags in order to insert full-fledged Python objects into your mocked contexts.
+
+To do so, decorate a function that returns the object you want with `@get_image_model` like so:
+
+```python
+# myproject/core/pattern_contexts.py
+
+from pattern_library.yaml import register_yaml_tag
+from wagtail.images import get_image_model
+
+@register_yaml_tag
+def testimage():
+    """
+    Return a random Image instance.
+    """
+    Image = get_image_model()
+    return Image.objects.order_by("?").first()
+```
+
+Once the custom YAML tag is registered, you can use it by adding the `!` prefix:
+
+```yaml
+context:
+  object_list:
+    - title: First item
+      image: !testimage
+    - title: Second item
+      image: !testimage
+```
+
+### Registering a tag under a different name
+
+The `@register_yaml_tag` decorator will use the name of the decorated function as the tag name automatically.
+
+You can specify a different name by passing `name=...` when registering the function:
+
+```python
+@register_yaml_tag("testimage")
+def get_random_image():
+    ...
+```
+
+
+### Passing arguments to custom tags
+
+It's possible to create custom tags that take arguments.
+
+```python
+@register_yaml_tag
+def testimage(collection):
+    """
+    Return a random Image instance from the given collection.
+    """
+    Image = get_image_model()
+    images = Image.objects.filter(collection__name=collection)
+    return images.order_by("?").first()
+```
+
+You can then specify arguments positionally using YAML's list syntax:
+```yaml
+context:
+  test_image: !testimage
+    - pattern_library
+```
+
+Alternatively you can specify keyword arguments using YAML's dictionary syntax:
+```yaml
+context:
+  test_image: !testimage
+    collection: pattern_library
+```

pattern_library/utils.py

@@ -21,11 +21,7 @@
 )
 from pattern_library.context_modifiers import registry
 from pattern_library.exceptions import TemplateIsNotPattern
-
-# Define our own yaml loader so we can register constructors on it without
-# polluting the original loader from the library.
-class PatternLibraryLoader(yaml.FullLoader):
-    pass
+from pattern_library.yaml import PatternLibraryLoader
 
 
 def path_to_section():

pattern_library/yaml.py

@@ -0,0 +1,66 @@
+from functools import partial, wraps
+
+from yaml.loader import FullLoader
+from yaml.nodes import MappingNode, SequenceNode
+
+# Define our own yaml loader so we can register constructors on it without
+# polluting the original loader from the library.
+class PatternLibraryLoader(FullLoader):
+    pass
+
+
+def _yaml_tag_constructor(fn):
+    """
+    Convert the given function into a PyYAML-compatible constructor that
+    correctly parses it args/kwargs.
+    """
+    @wraps(fn)
+    def constructor(loader, node):
+        args, kwargs = (), {}
+        if isinstance(node, SequenceNode):
+            args = loader.construct_sequence(node, deep=True)
+        elif isinstance(node, MappingNode):
+            kwargs = loader.construct_mapping(node, deep=True)
+        else:
+            pass  # No arguments given
+        return fn(*args, **kwargs)
+
+    return constructor
+
+
+def register_yaml_tag(fn=None, name=None):
+    """
+    Register the given function as a custom (local) YAML tag under the given name.
+    """
+
+    # This set of if statements is fairly complex so we can support a variety
+    # of ways to call the decorator:
+
+    # @register_yaml_tag()
+    if fn is None and name is None:  # @register_yaml_tag()
+        return partial(register_yaml_tag, name=None)
+
+    # @register_yaml_tag(name="asdf")
+    elif fn is None and name is not None:
+        return partial(register_yaml_tag, name=name)
+
+    # @register_yaml_tag("asdf")
+    elif isinstance(fn, str) and name is None:  
+        return partial(register_yaml_tag, name=fn)
+
+    # @register_yaml_tag
+    elif fn is not None and name is None:
+        return register_yaml_tag(fn, name=fn.__name__)
+
+    # At this point, both `fn` and `name` are defined
+    PatternLibraryLoader.add_constructor(f"!{name}", _yaml_tag_constructor(fn))
+    return fn
+
+
+def unregister_yaml_tag(name):
+    """
+    Unregister the custom tag with the given name.
+    """
+    # PyYAML doesn't provide an inverse operation for add_constructor(), so
+    # we need to do it manually.
+    del PatternLibraryLoader.yaml_constructors[f"!{name}"]

tests/tests/test_utils.py

@@ -1,17 +1,13 @@
 import os
 from unittest import mock
 
-import yaml
 from django.conf import settings
 from django.test import SimpleTestCase, override_settings
-from yaml.constructor import ConstructorError
 
 from pattern_library.utils import (
-    get_pattern_config,
     get_pattern_config_str,
     get_renderer,
     get_template_dirs,
-    PatternLibraryLoader,
 )
 
 
@@ -126,31 +122,3 @@ def test_atom_yml(self):
 
         self.assertNotEqual(result, "")
         self.assertIn("atom_var value from test_atom.yml", result)
-
-    @mock.patch("pattern_library.utils.get_pattern_config_str")
-    def test_custom_yaml_tag_error_if_unregistered(self, get_pattern_config_str):
-        get_pattern_config_str.return_value = "context:\n  atom_var: !customtag"
-
-        self.assertRaises(
-            ConstructorError,
-            get_pattern_config,
-            "patterns/atoms/test_custom_yaml_tag/test_custom_yaml_tag.html",
-        )
-
-    @mock.patch("pattern_library.utils.get_pattern_config_str")
-    def test_custom_yaml_tag(self, get_pattern_config_str):
-        get_pattern_config_str.return_value = "context:\n  atom_var: !customtag"
-        yaml.add_constructor(
-            "!customtag",
-            lambda loader, node: 42,
-            Loader=PatternLibraryLoader,
-        )
-        # PyYAML's API doesn't have a remove_constructor() function so we do
-        # that manually to avoid leaving things on the loader after the test
-        # is finished.
-        self.addCleanup(PatternLibraryLoader.yaml_constructors.pop, "!customtag")
-
-        self.assertEqual(
-            get_pattern_config("mocked.html"),
-            {"context": {"atom_var": 42}},
-        )

tests/tests/test_yaml.py

@@ -0,0 +1,128 @@
+from unittest import mock
+
+import yaml
+from django.test import SimpleTestCase
+from yaml.constructor import ConstructorError
+
+from pattern_library.utils import get_pattern_context
+from pattern_library.yaml import (
+    register_yaml_tag,
+    unregister_yaml_tag,
+)
+
+
+class PatternLibraryLoaderTestCase(SimpleTestCase):
+    def tearDown(self):
+        try:
+            unregister_yaml_tag("customtag")
+        except KeyError:
+            pass
+
+    def _get_context(self, yaml_str):
+        # Use mock.patch to avoid having to create actual files on disk
+        with mock.patch("pattern_library.utils.get_pattern_config_str", return_value=yaml_str):
+            return get_pattern_context("mocked.html")
+
+    def assertContextEqual(self, yaml_str, expected, msg=None):
+        """
+        Check that the given yaml string can be loaded and results in the given context.
+        """
+        context = self._get_context(yaml_str)
+        self.assertEqual(context, expected, msg=msg)
+
+    def test_unknown_tag_throws_error(self):
+        self.assertRaises(
+            ConstructorError,
+            self._get_context,
+            "context:\n  test: !customtag"
+        )
+
+    def test_custom_tag_can_be_registered(self):
+        register_yaml_tag(lambda: 42, "customtag")
+        self.assertContextEqual(
+            "context:\n  test: !customtag",
+            {"test": 42},
+        )
+
+    def test_custom_tag_can_be_unregistered(self):
+        register_yaml_tag(lambda: 42, "customtag")
+        unregister_yaml_tag("customtag")
+        self.assertRaises(
+            ConstructorError,
+            self._get_context,
+            "context:\n  test: !customtag"
+        )
+
+    def test_custom_tag_registering_doesnt_pollute_parent_loader(self):
+        register_yaml_tag(lambda: 42, "customtag")
+        self.assertRaises(
+            ConstructorError,
+            yaml.load,
+            "context:\n  test: !customtag",
+            Loader=yaml.FullLoader,
+        )
+
+    def test_registering_plain_decorator(self):
+        @register_yaml_tag
+        def customtag():
+            return 42
+
+        self.assertContextEqual(
+            "context:\n  test: !customtag",
+            {"test": 42},
+        )
+
+    def test_registering_plain_decorator_called(self):
+        @register_yaml_tag()
+        def customtag():
+            return 42
+
+        self.assertContextEqual(
+            "context:\n  test: !customtag",
+            {"test": 42},
+        )
+
+    def test_registering_decorator_specify_name(self):
+        @register_yaml_tag("customtag")
+        def function_with_different_name():
+            return 42
+
+        self.assertContextEqual(
+            "context:\n  test: !customtag",
+            {"test": 42},
+        )
+
+    def test_registering_decorator_specify_name_kwarg(self):
+        @register_yaml_tag(name="customtag")
+        def function_with_different_name():
+            return 42
+
+        self.assertContextEqual(
+            "context:\n  test: !customtag",
+            {"test": 42},
+        )
+
+    def test_custom_tag_with_args(self):
+        register_yaml_tag(lambda *a: sum(a), "customtag")
+
+        yaml_str = """
+context:
+  test: !customtag
+    - 1
+    - 2
+    - 3
+        """.strip()
+
+        self.assertContextEqual(yaml_str, {"test": 6})
+
+    def test_custom_tag_with_kwargs(self):
+        register_yaml_tag(lambda **kw: {k.upper(): v for k, v in kw.items()}, "customtag")
+
+        yaml_str = """
+context:
+  test: !customtag
+    key1: 1
+    key2: 2
+        """.strip()
+
+        self.assertContextEqual(yaml_str, {"test": {"KEY1": 1, "KEY2": 2}})