Refactor codec schema modification

python/tskit/metadata.py

@@ -72,6 +72,10 @@ class AbstractMetadataCodec(metaclass=abc.ABCMeta):
     def __init__(self, schema: Mapping[str, Any]) -> None:
         raise NotImplementedError  # pragma: no cover
 
+    @classmethod
+    def modify_schema(self, schema: Mapping) -> Mapping:
+        return schema
+
     @abc.abstractmethod
     def encode(self, obj: Any) -> bytes:
         raise NotImplementedError  # pragma: no cover
@@ -357,7 +361,8 @@ def make_null_encode(cls, sub_schema):
     def make_numeric_encode(cls, sub_schema):
         return struct.Struct("<" + sub_schema["binaryFormat"]).pack
 
-    def __init__(self, schema: Mapping[str, Any]) -> None:
+    @classmethod
+    def modify_schema(cls, schema: Mapping) -> Mapping:
         # This codec requires that all properties are required and additional ones
         # not allowed. Rather than get schema authors to repeat that everywhere
         # we add it here, sadly we can't do this in the metaschema as "default" isn't
@@ -374,7 +379,9 @@ def enforce_fixed_properties(obj):
             else:
                 return obj
 
-        schema = enforce_fixed_properties(schema)
+        return enforce_fixed_properties(schema)
+
+    def __init__(self, schema: Mapping[str, Any]) -> None:
         try:
             StructCodecSchemaValidator.check_schema(schema)
         except jsonschema.exceptions.SchemaError as ve:
@@ -431,12 +438,10 @@ def __init__(self, schema: Optional[Mapping[str, Any]]) -> None:
                     f"Unrecognised metadata codec '{schema['codec']}'. "
                     f"Valid options are {str(list(codec_registry.keys()))}."
                 )
+            # Codecs can modify the schema, for example to set defaults as the validator
+            # does not.
+            schema = codec_cls.modify_schema(schema)
             codec_instance = codec_cls(schema)
-            # If the codec has modified the schema on __init__, get it back.
-            try:
-                schema = codec_instance.schema
-            except AttributeError:
-                pass
             self._string = json.dumps(schema)
             self._validate_row = TSKITMetadataSchemaValidator(schema).validate
             self.encode_row = codec_instance.encode