Registering custom YAML tags

CHANGELOG.md

@@ -7,6 +7,7 @@
 ### Added
 
 - Add experimental Jinja support ([#180](https://github.com/torchbox/django-pattern-library/discussions/180), [#247](https://github.com/torchbox/django-pattern-library/pull/247), [#254](https://github.com/torchbox/django-pattern-library/pull/254)). Thank you to [@gone](https://github.com/gone), [@luord](https://github.com/luord), [@edcohen08](https://github.com/edcohen08), [@maribedran](https://github.com/maribedran), [@CuriousLearner](https://github.com/CuriousLearner)!
+- Add a `register_yaml_tag` decorator to make use of custom YAML tags in contexts ([#260](https://github.com/torchbox/django-pattern-library/pull/260)).
 
 ### Maintenance
 

docs/guides/defining-template-context.md

@@ -100,9 +100,9 @@ def add_page_images(context, request):
     Image = get_image_model()
     if "page" in context:
         if "hero_image" in context["page"]:
-            context["hero_image"] = Image.objects.all().order("?").first()
+            context["hero_image"] = Image.objects.all().order_by("?").first()
         if "main_image" in context["page"]:
-            context["main_image"] = Image.objects.all().order("?").first()
+            context["main_image"] = Image.objects.all().order_by("?").first()
 
 
 @register_context_modifier
@@ -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 `@register_yaml_tag` 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,6 +21,7 @@
 )
 from pattern_library.context_modifiers import registry
 from pattern_library.exceptions import TemplateIsNotPattern
+from pattern_library.yaml import PatternLibraryLoader
 
 
 def path_to_section():
@@ -100,7 +101,7 @@ def get_pattern_config_str(template_name):
 def get_pattern_config(template_name):
     config_str = get_pattern_config_str(template_name)
     if config_str:
-        return yaml.load(config_str, Loader=yaml.FullLoader)
+        return yaml.load(config_str, Loader=PatternLibraryLoader)
     return {}
 
 

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_yaml.py

@@ -0,0 +1,130 @@
+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):
+        super().tearDown()
+        # Make sure any custom tag is unregistered after every test
+        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}})