support for popping extra keys and post_process of properties

[thehesiod]

We're slowly replacing our use of schema with your library because this is sooo much faster (thank you!) and for feature parity we're started implementing some missing features. In schema when you have extra keys they're removed from the response, and you can do a Use to transform data. I've added support for these recently and posting here in case you're interested in adding something like it to the library.

One thing that would be nice is if it were easier to pass a custom generator. Right now I have to override the compile method to get it done.

from functools import partial, update_wrapper
import re
from typing import Any, Callable

import fastjsonschema
from fastjsonschema.draft04 import JSON_TYPE_TO_PYTHON_TYPE, JsonSchemaDefinitionException
from fastjsonschema.draft07 import CodeGeneratorDraft07


# These are very useful types to support
JSON_TYPE_TO_PYTHON_TYPE.update(
    {
        'datetime': 'datetime',
        'date': 'date',
        'time': 'time',
    }
)


# Adds support for more python types
class AgDataCodeGeneratorDraft07(CodeGeneratorDraft07):
    def __init__(
        self, *args, pop_extra_keys: bool = False, extra_import_objects: dict[str, Callable] | None = None, **kwargs
    ):
        super().__init__(*args, **kwargs)
        self.__pop_extra_keys = pop_extra_keys

        # add custom type support
        if extra_import_objects:
            assert not (
                over_lapping_keys := (extra_import_objects.keys() & self._extra_imports_objects.keys())
            ), f'overlapping keys in extra_import_objects: {over_lapping_keys}'
            self._extra_imports_objects.update(extra_import_objects)

    def generate_properties(self):
        self.create_variable_is_dict()
        with self.l('if {variable}_is_dict:'):
            self.create_variable_keys()

            if self.__pop_extra_keys and self._definition.get('additionalProperties', True):
                with self.l('for {variable}__var_name in ({variable}.keys() - {properties}.keys()):'):
                    self.l('{variable}.pop({variable}__var_name)')

            for key, prop_definition in self._definition['properties'].items():
                key_name = re.sub(r'($[^a-zA-Z]|[^a-zA-Z0-9])', '', key)
                if not isinstance(prop_definition, (dict, bool)):
                    raise JsonSchemaDefinitionException('{}[{}] must be object'.format(self._variable, key_name))

                with self.l('if "{}" in {variable}_keys:', self.e(key)):
                    self.l('{variable}_keys.remove("{}")', self.e(key))
                    self.l('{variable}__{0} = {variable}["{1}"]', key_name, self.e(key))
                    self.generate_func_code_block(
                        prop_definition,
                        '{}__{}'.format(self._variable, key_name),
                        '{}.{}'.format(self._variable_name, self.e(key)),
                        clear_variables=True,
                    )
                    if post_process := prop_definition.get('post_process'):
                        if post_process not in self._extra_imports_objects:
                            raise JsonSchemaDefinitionException(
                                'post_process of {}[{}] must be declared in extra_import_objects'.format(
                                    self._variable, key_name
                                )
                            )

                        self.l('{variable}["{0}"] = {1}({variable}["{0}"])', self.e(key), post_process)
                if self._use_default and isinstance(prop_definition, dict) and 'default' in prop_definition:
                    self.l('else: {variable}["{}"] = {}', self.e(key), repr(prop_definition['default']))


def fastjsonschema_custom_compile(
    definition,
    handlers=None,
    formats=None,
    use_default=True,
    use_formats=True,
    detailed_exceptions: bool = True,
    pop_extra_keys: bool = False,
    extra_import_objects: dict[str, Callable] | None = None,
):
    if handlers is None:
        handlers = dict()
    if formats is None:
        formats = dict()

    resolver = fastjsonschema.RefResolver.from_schema(definition, handlers=handlers, store={})
    code_generator = AgDataCodeGeneratorDraft07(
        definition,
        resolver=resolver,
        formats=formats,
        use_default=use_default,
        use_formats=use_formats,
        detailed_exceptions=detailed_exceptions,
        pop_extra_keys=pop_extra_keys,
        extra_import_objects=extra_import_objects,
    )
    global_state = code_generator.global_state
    exec(code_generator.func_code, global_state)  # nosec
    func = global_state[resolver.get_scope_name()]
    if formats:
        return update_wrapper(partial(func, custom_formats=formats), func)
    return func

The other thing is that it's a lot nicer to specify on each property if it's required or not instead of having to uplevel this information so I have a helper for that as well if interested.