def wrap(func, keys):
    """
    Primitive to wrap the output of a function into dictionary keys.

    Args:
        func: The function to enable output wrapping on.
        keys: The name of the keys to use in the output dictionary.
    """
    if type(keys) is not tuple:
        keys = [keys]

    def _wrapped(*args, **kwargs):
        res = func(*args, **kwargs)
        if type(res) is not tuple:
            res = [res]
        return {k:v for k, v in zip(keys, res)}
    
    return _wrapped