diff --git a/CHANGELOG.md b/CHANGELOG.md index 3ef43641..3d06595d 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/docs/guides/defining-template-context.md b/docs/guides/defining-template-context.md index d3d2ca9e..f6c6aa2e 100644 --- a/docs/guides/defining-template-context.md +++ b/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 +``` diff --git a/pattern_library/utils.py b/pattern_library/utils.py index 6389f7ad..a676052a 100644 --- a/pattern_library/utils.py +++ b/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 {} diff --git a/pattern_library/yaml.py b/pattern_library/yaml.py new file mode 100644 index 00000000..84cb10f9 --- /dev/null +++ b/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}"] diff --git a/tests/tests/test_yaml.py b/tests/tests/test_yaml.py new file mode 100644 index 00000000..fad101e5 --- /dev/null +++ b/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}})