#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""simple scripts that defines a custom callback to pickle a function depending
on the function name"""

import _pickle
import io


def hook(obj):
    obj_name = obj.__name__
    msg = f"callback for {obj_name}:   "
    if obj_name == 'f':
        print(msg + "using custom callback")
        # the pickler will run save_reduce(int, (5, ))
        return int, (5, )
    elif obj_name == 'g':
        print(msg + "falling back to save_global")
        return NotImplementedError
    elif obj_name == 'h':
        print(msg + "returning an invalid result (anything not a 2-5 tuple)")
        return False
    return NotImplementedError


def f():
    print('you just called h')


def g():
    print('you just called g')


def h():
    print('you just called h')


if __name__ == "__main__":
    b = io.BytesIO()
    p = _pickle.Pickler(b)

    # add the callback to the pickler
    p.global_hook = hook

    p.dump([f, g])
    new_f, new_g = _pickle.loads(b.getvalue())

    assert new_f == 5

    # g was pickled by attribute using save_global
    assert new_g == g

    try:
        p.dump(h)
    except _pickle.PicklingError:
        print("trying to pickle returned a PicklingError as expected")
    else:
        raise ValueError("a PicklingError should have been raised when "
                         "pickling h")