From db918105c472dbdff2108675940fc28084a6be08 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ingy=20d=C3=B6t=20Net?= Date: Tue, 19 Mar 2019 10:34:39 -0700 Subject: [PATCH] Add custom constructors to multiple loaders When someone writes a subclass of the YAMLObject class, the constructors will now be added to all 3 (non-safe) loaders. Furthermore, we support the class variable `yaml_loader` being a list, offering more control of which loaders are affected. To support safe_load in your custom class you could add this: yaml_loader = yaml.SafeLoader yaml_loader = yaml.YAMLObject.yaml_loader yaml_loader.append(yaml.SafeLoader) --- lib/yaml/__init__.py | 9 +++++++-- lib3/yaml/__init__.py | 9 +++++++-- 2 files changed, 14 insertions(+), 4 deletions(-) diff --git a/lib/yaml/__init__.py b/lib/yaml/__init__.py index 9f84c388..48b33b3c 100644 --- a/lib/yaml/__init__.py +++ b/lib/yaml/__init__.py @@ -381,7 +381,12 @@ class YAMLObjectMetaclass(type): def __init__(cls, name, bases, kwds): super(YAMLObjectMetaclass, cls).__init__(name, bases, kwds) if 'yaml_tag' in kwds and kwds['yaml_tag'] is not None: - cls.yaml_loader.add_constructor(cls.yaml_tag, cls.from_yaml) + if isinstance(cls.yaml_loader, list): + for loader in cls.yaml_loader: + loader.add_constructor(cls.yaml_tag, cls.from_yaml) + else: + cls.yaml_loader.add_constructor(cls.yaml_tag, cls.from_yaml) + cls.yaml_dumper.add_representer(cls, cls.to_yaml) class YAMLObject(object): @@ -393,7 +398,7 @@ class YAMLObject(object): __metaclass__ = YAMLObjectMetaclass __slots__ = () # no direct instantiation, so allow immutable subclasses - yaml_loader = Loader + yaml_loader = [Loader, FullLoader, UnsafeLoader] yaml_dumper = Dumper yaml_tag = None diff --git a/lib3/yaml/__init__.py b/lib3/yaml/__init__.py index a55206ba..183969e5 100644 --- a/lib3/yaml/__init__.py +++ b/lib3/yaml/__init__.py @@ -378,7 +378,12 @@ class YAMLObjectMetaclass(type): def __init__(cls, name, bases, kwds): super(YAMLObjectMetaclass, cls).__init__(name, bases, kwds) if 'yaml_tag' in kwds and kwds['yaml_tag'] is not None: - cls.yaml_loader.add_constructor(cls.yaml_tag, cls.from_yaml) + if isinstance(cls.yaml_loader, list): + for loader in cls.yaml_loader: + loader.add_constructor(cls.yaml_tag, cls.from_yaml) + else: + cls.yaml_loader.add_constructor(cls.yaml_tag, cls.from_yaml) + cls.yaml_dumper.add_representer(cls, cls.to_yaml) class YAMLObject(metaclass=YAMLObjectMetaclass): @@ -389,7 +394,7 @@ class YAMLObject(metaclass=YAMLObjectMetaclass): __slots__ = () # no direct instantiation, so allow immutable subclasses - yaml_loader = Loader + yaml_loader = [Loader, FullLoader, UnsafeLoader] yaml_dumper = Dumper yaml_tag = None