Merge pull request #126 from abravalheri/add-custom-formats-arg

Add `custom_formats` argument to generated code

Author: horejsek

Diff for: fastjsonschema/__init__.py

@@ -75,6 +75,7 @@
 API
 ***
 """
+from functools import partial, update_wrapper
 
 from .draft04 import CodeGeneratorDraft04
 from .draft06 import CodeGeneratorDraft06
@@ -177,7 +178,10 @@ def compile(definition, handlers={}, formats={}, use_default=True):
     global_state = code_generator.global_state
     # Do not pass local state so it can recursively call itself.
     exec(code_generator.func_code, global_state)
-    return global_state[resolver.get_scope_name()]
+    func = global_state[resolver.get_scope_name()]
+    if formats:
+        return update_wrapper(partial(func, custom_formats=formats), func)
+    return func
 
 
 # pylint: disable=dangerous-default-value

Diff for: fastjsonschema/generator.py

@@ -32,6 +32,7 @@ class CodeGenerator:
     def __init__(self, definition, resolver=None):
         self._code = []
         self._compile_regexps = {}
+        self._custom_formats = {}
 
         # Any extra library should be here to be imported only once.
         # Lines are imports to be printed in the file and objects
@@ -136,7 +137,7 @@ def generate_validation_function(self, uri, name):
         self._validation_functions_done.add(uri)
         self.l('')
         with self._resolver.resolving(uri) as definition:
-            with self.l('def {}(data):', name):
+            with self.l('def {}(data, custom_formats={{}}):', name):
                 self.generate_func_code_block(definition, 'data', 'data', clear_variables=True)
                 self.l('return data')
 
@@ -190,7 +191,7 @@ def generate_ref(self):
             if uri not in self._validation_functions_done:
                 self._needed_validation_functions[uri] = name
             # call validation function
-            self.l('{}({variable})', name)
+            self.l('{}({variable}, custom_formats)', name)
 
 
     # pylint: disable=invalid-name

Diff for: tests/test_compile_to_code.py

@@ -2,6 +2,7 @@
 import pytest
 import shutil
 
+from fastjsonschema import JsonSchemaValueException
 from fastjsonschema import compile_to_code, compile as compile_spec
 
 @pytest.yield_fixture(autouse=True)
@@ -84,3 +85,40 @@ def test_compile_complex_one_of_all_of():
             }
         ]
     })
+
+
+def test_compile_to_code_custom_format():
+    formats = {'my-format': str.isidentifier}
+    code = compile_to_code({'type': 'string', 'format': 'my-format'}, formats=formats)
+    with open('temp/schema_3.py', 'w') as f:
+        f.write(code)
+    from temp.schema_3 import validate
+    assert validate("valid", formats) == "valid"
+    with pytest.raises(JsonSchemaValueException) as exc:
+        validate("not-valid", formats)
+    assert exc.value.message == "data must be my-format"
+
+
+def test_compile_to_code_custom_format_with_refs():
+    schema = {
+        'type': 'object',
+        'properties': {
+            'a': {'$ref': '#/definitions/a'}
+        },
+        'definitions': {
+            'a': {
+                '$id': '#/definitions/a',
+                'type': 'array',
+                'items': {'type': 'string', 'format': 'my-format'}
+            }
+        }
+    }
+    formats = {'my-format': str.isidentifier}
+    code = compile_to_code(schema, formats=formats)
+    with open('temp/schema_4.py', 'w') as f:
+        f.write(code)
+    from temp.schema_4 import validate
+    assert validate({"a": ["identifier"]}, formats) is not None
+    with pytest.raises(JsonSchemaValueException) as exc:
+        validate({"a": ["identifier", "not-valid"]}, formats)
+    assert exc.value.message == "data[1] must be my-format"