Add custom constructors to multiple loaders

ingydotnet · perlpunk · commit a5394c04a2ef · 2019-11-18T11:59:54.000+01:00
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)
diff --git 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
@@ -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
