GitHub issue #3224: User-pickling with dynamic sub-classing
State: open | Created: 2021-11-06T12:08:24Z

This is a continuation of the now closed PR in https://github.com/huggingface/datasets/pull/3206. The discussion there has shaped a new approach to do this.

In this PR, behavior of `pklregister` and `Pickler` is extended. Earlier, users were already able to register custom pickle functions. That is useful if they have objects that are not easily picklable with default methods. When one registers a custom function to a type, an object of that type will be pickled with the given function by `Pickler` which looks up the type in its `dispatch` table. The downside of this method, and of `pickle` in general, is that it is limited to direct type-matching and does not allow sub-classes. In many, default, cases that is not an issue. But when you are using external libraries where classes (e.g. parsers, models) are sub-classed this is not ideal. 

```python
from datasets.fingerprint import Hasher
from datasets.utils.py_utils import pklregister

class BaseParser:
    pass

class EnglishParser(BaseParser):
    pass

@pklregister(BaseParser)
def custom_pkl_func(pickler, obj):
    print(f"Called the custom pickle function for type {type(obj)}!")
    # do something with the obj and ultimately save with the pickler

base = BaseParser()
en = EnglishParser()

# Hasher.hash uses the Pickler behind the scenes
# `custom_pkl_func` called for base
Hasher.hash(base)
# `custom_pkl_func` not called for en :-(
Hasher.hash(en)
```

In the example above we'd want to sub-class `EnglishParser` to be handled in the same way as its super-class `BaseParser`. This PR solves that by allowing for a keyword-argument `allow_subclasses` in `pklregister` (default: `False`). 

```python
@pklregister(BaseParser, allow_subclasses=True)
```

When this option is enabled, we not only save the function in `Pickler.dispatch` but also save it in a custom table `Pickler.subclass_dispatch` **which allows us to dynamically add sub-classes of that class to the real dispatch table**. Then, if we want to pickle an  object `obj` with `Pickler.dump()` (which ultimately will call `Pickler.save()`) we _first_  check whether any of the object's super-classes exist in `Pickler.sublcass_dispatch` and get the related custom pickle function. If we find one, we add the type of `obj` alongside the function to `Pickler.dispatch`. All of this happens at the start of the call to `Pickler.save()`. _Only then_ dill.Pickler's `save` will be called, which in turn will call `pickle._Pickler.save` which handles everything. Here, the `Pickler.dispatch` table will be used to look up custom pickler functions - and it now also includes the function for `obj`, which was copied from its super-class, which we added at the very start of our custom `Pickler.save()`.

For edge cases and, especially, for testing, a contextmanager class `TempPickleRegistry` is included that resets the pickle registry on exit to its previous state.

```python
with TempPickleRegistry():
    @pklregister(MyObjClass)
    def pickle_registry_test_false(pickler, obj):
        pickler.save(obj.fancy_method())

    some_obj = MyObjClass()
    dumps(some_obj)
    # `MyObjClass` is in Pickler.dispatch

# ... `MyObjClass` is _not_ in Pickler.dispatch anymore
```

closes https://github.com/huggingface/datasets/issues/3178

To Do
====
- [x] Write tests
- [ ] Write documentation/examples?