from google.cloud.ndb import context as context_module


def get_batch(batch_cls, options=None):
    """Gets a data structure for storing batched calls to Datastore Lookup.

    The batch data structure is stored in the current context. If there is
    not already a batch started, a new structure is created and an idle
    callback is added to the current event loop which will eventually perform
    the batch look up.

    Args:
        batch_cls (type): Class representing the kind of operation being
            batched.
        options (_options.ReadOptions): The options for the request. Calls with
            different options will be placed in different batches.

    Returns:
        batch_cls: An instance of the batch class.
    """
    # prevent circular import in Python 2.7

    context = context_module.get_context()
    batches = context.batches.get(batch_cls)
    if batches is None:
        context.batches[batch_cls] = batches = {}

    if options is not None:
        options_key = tuple(
            sorted(
                (
                    (key, value)
                    for key, value in options.items()
                    if value is not None
                )
            )
        )
    else:
        options_key = ()

    batch = batches.get(options_key)
    if batch is not None and not batch.full():
        return batch

    def idler(batch):
        def idle():
            if batches.get(options_key) is batch:
                del batches[options_key]
            batch.idle_callback()

        return idle

    batches[options_key] = batch = batch_cls(options)
    context.eventloop.add_idle(idler(batch))
    return batch