API¶

def bulk_create_bulks_in_steps(
    bulk_by_class: dict[type[Model], Iterable[Model]],
    step: int = STANDARD_BULK_SIZE,
) -> dict[type[Model], list[Model]]:
    """Create multiple model-type bulks in dependency order.

    The function topologically sorts the provided model classes so that
    models referenced by foreign keys are created before models that
    reference them. Each class' instances are created in batches using
    :func:`bulk_create_in_steps`.

    Args:
        bulk_by_class: Mapping from model class to iterable of instances to create.
        step: Chunk size for each model's batched creation.

    Returns:
        Mapping from model class to list of created instances.
    """
    # order the bulks in order of creation depending how they depend on each other
    models_ = list(bulk_by_class.keys())
    ordered_models = topological_sort_models(models=models_)

    results: dict[type[Model], list[Model]] = {}
    for model_ in ordered_models:
        bulk = bulk_by_class[model_]
        result = bulk_create_in_steps(model=model_, bulk=bulk, step=step)
        results[model_] = result

    return results

def bulk_create_in_steps[TModel: Model](
    model: type[TModel],
    bulk: Iterable[TModel],
    step: int = STANDARD_BULK_SIZE,
) -> list[TModel]:
    """Create objects in batches and return created instances.

    Breaks ``bulk`` into chunks of size ``step`` and calls the project's
    batched bulk-create helper for each chunk. Execution is performed
    using the concurrent utility where configured for throughput.

    Args:
        model: Django model class to create instances for.
        bulk: Iterable of unsaved model instances.
        step: Number of instances to create per chunk.

    Returns:
        List of created model instances (flattened across chunks).
    """
    return cast(
        "list[TModel]",
        bulk_method_in_steps(model=model, bulk=bulk, step=step, mode=MODE_CREATE),
    )

def bulk_delete(
    model: type[Model],
    objs: Iterable[Model],
    **_: Any,  # noqa: ANN401
) -> tuple[int, dict[str, int]]:
    """Delete the provided objects and return Django's delete summary.

    Accepts either a QuerySet or an iterable of model instances. When an
    iterable of instances is provided it is converted to a QuerySet by
    filtering on primary keys before calling ``QuerySet.delete()``.

    Args:
        model: Django model class.
        objs: Iterable of model instances or a QuerySet.

    Returns:
        Tuple(total_deleted, per_model_counts) as returned by ``delete()``.
    """
    return model.objects.filter(pk__in=(obj.pk for obj in objs)).delete()

def bulk_delete_in_steps[TModel: Model](
    model: type[TModel],
    bulk: Iterable[TModel],
    step: int = STANDARD_BULK_SIZE,
) -> tuple[int, dict[str, int]]:
    """Delete objects in batches and return deletion statistics.

    Each chunk is deleted using Django's QuerySet ``delete`` which
    returns a (count, per-model-counts) tuple. Results are aggregated
    across chunks and returned as a consolidated tuple.

    Args:
        model: Django model class.
        bulk: Iterable of model instances to delete.
        step: Chunk size for deletions.

    Returns:
        A tuple containing the total number of deleted objects and a
        mapping from model label to deleted count (including cascaded
        deletions).
    """
    return cast(
        "tuple[int, dict[str, int]]",
        bulk_method_in_steps(
            model=model,
            bulk=bulk,
            step=step,
            mode=MODE_DELETE,
        ),
    )

def bulk_method_in_steps[TModel: Model](
    model: type[TModel],
    bulk: Iterable[TModel],
    step: int,
    mode: MODE_TYPES,
    **kwargs: Any,
) -> int | tuple[int, dict[str, int]] | list[TModel]:
    """Run a batched bulk operation (create/update/delete) on ``bulk``.

    This wrapper warns if called from within an existing transaction and
    delegates actual work to :func:`bulk_method_in_steps_atomic` which is
    executed inside an atomic transaction. The return type depends on
    ``mode`` (see :mod:`winidjango.core.db.bulk` constants).

    Args:
        model: Django model class to operate on.
        bulk: Iterable of model instances.
        step: Chunk size for processing.
        mode: One of ``'create'``, ``'update'`` or ``'delete'``.
        **kwargs: Additional keyword arguments forwarded to the underlying
            bulk methods (for example ``update_fields`` for updates).

    Returns:
        For ``create``: list of created instances.
        For ``update``: integer number of updated rows.
        For ``delete``: tuple(total_deleted, per_model_counts).
    """
    # check if we are inside a transaction.atomic block
    _in_atomic_block = transaction.get_connection().in_atomic_block
    if _in_atomic_block:
        logger.info(
            "BE CAREFUL USING BULK OPERATIONS INSIDE A BROADER TRANSACTION BLOCK. "
            "BULKING WITH BULKS THAT DEPEND ON EACH OTHER CAN CAUSE "
            "INTEGRITY ERRORS OR POTENTIAL OTHER ISSUES.",
        )
    return bulk_method_in_steps_atomic(
        model=model,
        bulk=bulk,
        step=step,
        mode=mode,
        **kwargs,
    )

@transaction.atomic
def bulk_method_in_steps_atomic[TModel: Model](
    model: type[TModel],
    bulk: Iterable[TModel],
    step: int,
    mode: MODE_TYPES,
    **kwargs: Any,
) -> int | tuple[int, dict[str, int]] | list[TModel]:
    """Atomic implementation of the batched bulk operation.

    This function is decorated with ``transaction.atomic`` so each call
    to this function runs in a database transaction. Note that nesting
    additional, outer atomic blocks that also include dependent bulk
    operations can cause integrity issues for operations that depend on
    each other's side-effects; callers should be careful about atomic
    decorator placement in higher-level code.

    Args:
        model: Django model class.
        bulk: Iterable of model instances.
        step: Chunk size for processing.
        mode: One of ``'create'``, ``'update'`` or ``'delete'``.
        **kwargs: Forwarded to the underlying bulk method.

    Returns:
        See :func:`bulk_method_in_steps` for return value semantics.
    """
    bulk_method = get_bulk_method(model=model, mode=mode, **kwargs)

    chunks = get_step_chunks(bulk=bulk, step=step)

    # multithreading significantly increases speed
    result = multithread_loop(
        process_function=bulk_method,
        process_args=chunks,
    )

    return flatten_bulk_in_steps_result(result=result, mode=mode)

def bulk_update_in_steps[TModel: Model](
    model: type[TModel],
    bulk: Iterable[TModel],
    update_fields: list[str],
    step: int = STANDARD_BULK_SIZE,
) -> int:
    """Update objects in batches and return total updated count.

    Args:
        model: Django model class.
        bulk: Iterable of model instances to update (must have PKs set).
        update_fields: Fields to update on each instance when calling
            ``bulk_update``.
        step: Chunk size for batched updates.

    Returns:
        Total number of rows updated across all chunks.
    """
    return cast(
        "int",
        bulk_method_in_steps(
            model=model,
            bulk=bulk,
            step=step,
            mode=MODE_UPDATE,
            fields=update_fields,
        ),
    )

def flatten_bulk_in_steps_result[TModel: Model](
    result: list[int] | list[tuple[int, dict[str, int]]] | list[list[TModel]],
    mode: str,
) -> int | tuple[int, dict[str, int]] | list[TModel]:
    """Aggregate per-chunk results returned by concurrent bulk execution.

    Depending on ``mode`` the function reduces a list of per-chunk
    results into a single consolidated return value:

    - ``create``: flattens a list of lists into a single list of objects
    - ``update``: sums integer counts returned per chunk
    - ``delete``: aggregates (count, per-model-dict) tuples into a single
      total and combined per-model counts

    Args:
        result: List of per-chunk results returned by the chunk function.
        mode: One of the supported modes.

    Returns:
        Aggregated result corresponding to ``mode``.
    """
    if mode == MODE_UPDATE:
        # formatted as [1000, 1000, ...]
        # since django 4.2 bulk_update returns the count of updated objects
        result = cast("list[int]", result)
        return int(sum(result))
    if mode == MODE_DELETE:
        # formatted as [(count, {model_name: count, model_cascade_name: count}), ...]
        # join the results to get the total count of deleted objects
        result = cast("list[tuple[int, dict[str, int]]]", result)
        total_count = 0
        count_sum_by_model: defaultdict[str, int] = defaultdict(int)
        for count_sum, count_by_model in result:
            total_count += count_sum
            for model_name, count in count_by_model.items():
                count_sum_by_model[model_name] += count
        return (total_count, dict(count_sum_by_model))
    if mode == MODE_CREATE:
        # formatted as [[obj1, obj2, ...], [obj1, obj2, ...], ...]
        result = cast("list[list[TModel]]", result)
        return [item for sublist in result for item in sublist]

    msg = f"Invalid method. Must be one of {MODES}"
    raise ValueError(msg)

def get_bulk_method(
    model: type[Model],
    mode: MODE_TYPES,
    **kwargs: Any,
) -> Callable[[list[Model]], list[Model] | int | tuple[int, dict[str, int]]]:
    """Return a callable that performs the requested bulk operation on a chunk.

    The returned function accepts a single argument (a list of model
    instances) and returns the per-chunk result for the chosen mode.

    Args:
        model: Django model class.
        mode: One of ``'create'``, ``'update'`` or ``'delete'``.
        **kwargs: Forwarded to the underlying ORM bulk methods.

    Raises:
        ValueError: If ``mode`` is invalid.

    Returns:
        Callable that accepts a list of model instances and returns the
        result for that chunk.
    """
    bulk_method: Callable[[list[Model]], list[Model] | int | tuple[int, dict[str, int]]]
    if mode == MODE_CREATE:

        def bulk_create_chunk(chunk: list[Model]) -> list[Model]:
            return model.objects.bulk_create(objs=chunk, **kwargs)

        bulk_method = bulk_create_chunk
    elif mode == MODE_UPDATE:

        def bulk_update_chunk(chunk: list[Model]) -> int:
            return model.objects.bulk_update(objs=chunk, **kwargs)

        bulk_method = bulk_update_chunk
    elif mode == MODE_DELETE:

        def bulk_delete_chunk(chunk: list[Model]) -> tuple[int, dict[str, int]]:
            return bulk_delete(model=model, objs=chunk, **kwargs)

        bulk_method = bulk_delete_chunk
    else:
        msg = f"Invalid method. Must be one of {MODES}"
        raise ValueError(msg)

    return bulk_method

def get_differences_between_bulks[TModel1: Model, TModel2: Model](
    bulk1: list[TModel1],
    bulk2: list[TModel2],
    fields: "list[Field[Any, Any] | ForeignObjectRel | GenericForeignKey]",
) -> tuple[list[TModel1], list[TModel2], list[TModel1], list[TModel2]]:
    """Return differences and intersections between two bulks of the same model.

    Instances are compared using :func:`hash_model_instance` over the
    provided ``fields``. The function maintains the original ordering
    for returned lists so that callers can preserve deterministic
    ordering when applying diffs.

    Args:
        bulk1: First list of model instances.
        bulk2: Second list of model instances.
        fields: Fields to include when hashing instances.

    Raises:
        ValueError: If bulks are empty or contain different model types.

    Returns:
        Four lists in the order:
            (in_1_not_2, in_2_not_1, in_both_from_1, in_both_from_2).
    """
    if not bulk1 or not bulk2:
        return bulk1, bulk2, [], []

    if type(bulk1[0]) is not type(bulk2[0]):
        msg = "Both bulks must be of the same model type."
        raise ValueError(msg)

    hash_model_instance_with_fields = partial(
        hash_model_instance,
        fields=fields,
    )
    # Precompute hashes and map them directly to models in a single pass for both bulks
    hashes1 = list(map(hash_model_instance_with_fields, bulk1))
    hashes2 = list(map(hash_model_instance_with_fields, bulk2))

    # Convert keys to sets for difference operations
    set1, set2 = set(hashes1), set(hashes2)

    # Calculate differences between sets
    # Find differences and intersection with original order preserved
    # Important, we need to return the original objects that are the same in memory,
    # so in_1_not_2 and in_2_not_1
    in_1_not_2 = set1 - set2
    in_1_not_2_list = [
        model
        for model, hash_ in zip(bulk1, hashes1, strict=False)
        if hash_ in in_1_not_2
    ]

    in_2_not_1 = set2 - set1
    in_2_not_1_list = [
        model
        for model, hash_ in zip(bulk2, hashes2, strict=False)
        if hash_ in in_2_not_1
    ]

    in_1_and_2 = set1 & set2
    in_1_and_2_from_1 = [
        model
        for model, hash_ in zip(bulk1, hashes1, strict=False)
        if hash_ in in_1_and_2
    ]
    in_1_and_2_from_2 = [
        model
        for model, hash_ in zip(bulk2, hashes2, strict=False)
        if hash_ in in_1_and_2
    ]

    return in_1_not_2_list, in_2_not_1_list, in_1_and_2_from_1, in_1_and_2_from_2

def get_step_chunks(
    bulk: Iterable[Model],
    step: int,
) -> Generator[tuple[list[Model]], None, None]:
    """Yield consecutive chunks of at most ``step`` items from ``bulk``.

    The function yields a single-tuple containing the chunk (a list of
    model instances) because the concurrent execution helper expects a
    tuple of positional arguments for the target function.

    Args:
        bulk: Iterable of model instances.
        step: Maximum number of instances per yielded chunk.

    Yields:
        Tuples where the first element is a list of model instances.
    """
    bulk = iter(bulk)
    while True:
        chunk = list(islice(bulk, step))
        if not chunk:
            break
        yield (chunk,)  # bc concurrent_loop expects a tuple of args