# Generated by Django 2.2.7 on 2019-12-09 09:49
from django.db import migrations
from django.db.models import Count, Sum


def get_or_create_warehouse(apps):
    Warehouse = apps.get_model("warehouse", "Warehouse")
    ShippingZone = apps.get_model("shipping", "ShippingZone")
    Site = apps.get_model("sites", "Site")

    site_settings = Site.objects.get_current().settings

    address = getattr(site_settings, "company_address", None)
    if address is None:
        Address = apps.get_model("account", "Address")
        address = Address.objects.create()

    warehouses = Warehouse.objects.annotate(zones_count=Count("shipping_zones")).filter(
        zones_count=ShippingZone.objects.count()
    )
    if warehouses.first() is not None:
        return warehouses.first()

    warehouse = Warehouse.objects.create(address=address)
    warehouse.shipping_zones.add(*ShippingZone.objects.all())
    return warehouse


def forward(apps, schema_editor):
    ProductVariant = apps.get_model("product", "ProductVariant")
    Stock = apps.get_model("stock", "Stock")

    if not ProductVariant.objects.exists():
        return

    warehouse = get_or_create_warehouse(apps)
    stocks = [
        Stock(
            product_variant=variant,
            warehouse=warehouse,
            quantity=variant.quantity,
            quantity_allocated=variant.quantity_allocated,
        )
        for variant in ProductVariant.objects.all()
    ]
    Stock.objects.bulk_create(stocks)


def backward(apps, schema_editor):
    ProductVariant = apps.get_model("product", "ProductVariant")
    Stock = apps.get_model("stock", "Stock")

    variants = []
    stocks = (
        Stock.objects.values("product_variant_id")
        .annotate(quantity_sum=Sum("quantity"))
        .annotate(quantity_allocated_sum=Sum("quantity_allocated"))
    )
    for variant in ProductVariant.objects.all():
        variant_dict = stocks.get(product_variant=variant)
        variant.quantity = variant_dict["quantity_sum"]
        variant.quantity_allocated = variant_dict["quantity_allocated_sum"]
        variants.append(variant)
    ProductVariant.objects.bulk_update(variants, ["quantity", "quantity_allocated"])
    Stock.objects.all().delete()


class Migration(migrations.Migration):

    dependencies = [
        ("product", "0110_auto_20191108_0340"),
        ("warehouse", "0002_auto_20191205_0753"),
        ("stock", "0001_initial"),
        ("site", "0025_auto_20191024_0552"),
        ("shipping", "0017_django_price_2"),
    ]

    operations = [migrations.RunPython(forward, backward)]

# Generated by Django 2.2.7 on 2019-12-09 09:49
from django.db import migrations
from django.db.models import Count, Sum


def get_or_create_warehouse(apps):
    Warehouse = apps.get_model("warehouse", "Warehouse")
    ShippingZone = apps.get_model("shipping", "ShippingZone")
    Site = apps.get_model("sites", "Site")

    site_settings = Site.objects.get_current().settings

    address = getattr(site_settings, "company_address", None)
    if address is None:
        Address = apps.get_model("account", "Address")
        address = Address.objects.create()

    warehouses = Warehouse.objects.annotate(zones_count=Count("shipping_zones")).filter(
        zones_count=ShippingZone.objects.count()
    )
    if warehouses.first() is not None:
        return warehouses.first()

    warehouse = Warehouse.objects.create(address=address)
    warehouse.shipping_zones.add(*ShippingZone.objects.all())
    return warehouse


def forward(apps, schema_editor):
    ProductVariant = apps.get_model("product", "ProductVariant")
    Stock = apps.get_model("stock", "Stock")

    if not ProductVariant.objects.exists():
        return

    product_variants = ProductVariant.objects.all()
    warehouse = get_or_create_warehouse(apps)
    stocks = [
        Stock(
            product_variant=variant,
            warehouse=warehouse,
            quantity=variant.quantity,
            quantity_allocated=variant.quantity_allocated,
        )
        for variant in product_variants.iterator()
    ]
    Stock.objects.bulk_create(stocks)


def backward(apps, schema_editor):
    ProductVariant = apps.get_model("product", "ProductVariant")
    Stock = apps.get_model("stock", "Stock")

    stocks = (
        Stock.objects.values("product_variant_id")
        .annotate(quantity_sum=Sum("quantity"))
        .annotate(quantity_allocated_sum=Sum("quantity_allocated"))
    )
    product_variants = ProductVariant.objects.all()
    for variant in product_variants.iterator():
        variant_dict = stocks.get(product_variant=variant)
        variant.quantity = variant_dict["quantity_sum"]
        variant.quantity_allocated = variant_dict["quantity_allocated_sum"]
        variant.save(update_fields=["quantity", "quantity_allocated"])

    Stock.objects.all().delete()


class Migration(migrations.Migration):

    dependencies = [
        ("product", "0110_auto_20191108_0340"),
        ("warehouse", "0002_auto_20191205_0753"),
        ("stock", "0001_initial"),
        ("site", "0025_auto_20191024_0552"),
        ("shipping", "0017_django_price_2"),
    ]

    operations = [migrations.RunPython(forward, backward)]

Use iterator for queryset