Decorator to register a function, or registration itself.

    This is primarily intended for use as a decorator, either with or without
    a key/parentheses.
    ```python
    @my_registry.register('key1')
    def value_fn(x, y, z):
      pass

    @my_registry.register()
    def another_fn(x, y):
      pass

    @my_registry.register
    def third_func():
      pass
    ```

    Note if key_or_value is provided as a non-callable, registration only
    occurs once the returned callback is called with a callable as its only
    argument.
    ```python
    callback = my_registry.register('different_key')
    'different_key' in my_registry  # False
    callback(lambda (x, y): x + y)
    'different_key' in my_registry  # True
    ```

    Args:
      key_or_value (optional): key to access the registered value with, or the
        function itself. If `None` (default), `self.default_key` will be called
        on `value` once the returned callback is called with `value` as the only
        arg. If `key_or_value` is itself callable, it is assumed to be the value
        and the key is given by `self.default_key(key)`.

    Returns:
      decorated callback, or callback generated a decorated function.

def register(self, key_or_value=None):
    """Decorator to register a function, or registration itself.

    This is primarily intended for use as a decorator, either with or without
    a key/parentheses.
    ```python
    @my_registry.register('key1')
    def value_fn(x, y, z):
      pass

    @my_registry.register()
    def another_fn(x, y):
      pass

    @my_registry.register
    def third_func():
      pass
    ```

    Note if key_or_value is provided as a non-callable, registration only
    occurs once the returned callback is called with a callable as its only
    argument.
    ```python
    callback = my_registry.register('different_key')
    'different_key' in my_registry  # False
    callback(lambda (x, y): x + y)
    'different_key' in my_registry  # True
    ```

    Args:
      key_or_value (optional): key to access the registered value with, or the
        function itself. If `None` (default), `self.default_key` will be called
        on `value` once the returned callback is called with `value` as the only
        arg. If `key_or_value` is itself callable, it is assumed to be the value
        and the key is given by `self.default_key(key)`.

    Returns:
      decorated callback, or callback generated a decorated function.
    """

    def decorator(value, key):
      self[key] = value
      return value

    # Handle if decorator was used without parens
    if callable(key_or_value):
      return decorator(value=key_or_value, key=None)
    else:
      return lambda value: decorator(value, key=key_or_value)