Repository: horejsek/python-fastjsonschema

Issue:
Error in output of `compile_to_code` when using custom formats
Hello and thank you very much for the very nice package.

I notice that when compiling code to string, the generated file assumes a `custom_formats` global variable that is not defined anywhere, and therefore results in errors.

In #126, I propose adding a second argument that should fix this scenario.
This argument is automatically "curried"-away when the `compile` function is called, and just take effect when using `compile_to_code`, so there is no need for changes in API or docs.

Discussion:

Patch:
diff --git a/fastjsonschema/__init__.py b/fastjsonschema/__init__.py
index 6dd90f6..572e6c1 100644
--- a/fastjsonschema/__init__.py
+++ b/fastjsonschema/__init__.py
@@ -75,6 +75,7 @@ Note that there are some differences compared to JSON schema standard:
 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 --git a/fastjsonschema/generator.py b/fastjsonschema/generator.py
index bedd7ed..71e9b58 100644
--- a/fastjsonschema/generator.py
+++ b/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 @@ class CodeGenerator:
         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 @@ class CodeGenerator:
             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


Test patch:
diff --git a/tests/test_compile_to_code.py b/tests/test_compile_to_code.py
index 509e044..84bc981 100644
--- a/tests/test_compile_to_code.py
+++ b/tests/test_compile_to_code.py
@@ -2,6 +2,7 @@ import os
 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"
\ No newline at end of file