<commit_before>import itertools
import uuid
from typing import Set

from django.db import models
from django.db.models import F
from django.utils.translation import pgettext_lazy

from ..account.models import Address
from ..core.exceptions import InsufficientStock
from ..product.models import ProductVariant
from ..shipping.models import ShippingZone


class WarehouseQueryset(models.QuerySet):
    def prefetch_data(self):
        return self.select_related("address").prefetch_related("shipping_zones")


class Warehouse(models.Model):
    id = models.UUIDField(default=uuid.uuid4, primary_key=True)
    name = models.CharField(
        pgettext_lazy("Warehouse field description", "Warehouse name"), max_length=255
    )
    company_name = models.CharField(
        pgettext_lazy("Warehouse field description", "Legal company name"),
        blank=True,
        max_length=255,
    )

    shipping_zones = models.ManyToManyField(ShippingZone, blank=True)
    address = models.ForeignKey(Address, on_delete=models.PROTECT)

    email = models.EmailField(
        pgettext_lazy("Warehouse field description", "Email address"),
        blank=True,
        default="",
    )<commit_msg>Add management methods for Warehouse<commit_after>import itertools
import uuid
from typing import Optional, Set

from django.contrib.sites.models import Site
from django.db import models, transaction
from django.db.models import F, QuerySet
from django.utils.translation import pgettext_lazy

from ..account.models import Address
from ..core.exceptions import InsufficientStock
from ..product.models import ProductVariant
from ..shipping.models import ShippingZone
from ..site.models import SiteSettings


class WarehouseQueryset(models.QuerySet):
    def prefetch_data(self):
        return self.select_related("address").prefetch_related("shipping_zones")

    def for_country(self, country: str):
        return self.filter(shipping_zone__countries__contains=country)

    @transaction.atomic
    def create_for_shipping_zones(
        self,
        address: Optional[Address] = None,
        shipping_zones: Optional[QuerySet] = None,
    ):
        if shipping_zones is None:
            shipping_zones = ShippingZone.objects.all()
        current_settings = SiteSettings.objects.get(site=Site.objects.get_current())
        address = address or current_settings.company_address
        if address is None:
            address = Address.objects.create()
        warehouse = self.create(address=address)
        warehouse.shipping_zones.add(*shipping_zones.values_list("pk", flat=True))
        return warehouse


class Warehouse(models.Model):
    id = models.UUIDField(default=uuid.uuid4, primary_key=True)
    name = models.CharField(
        pgettext_lazy("Warehouse field description", "Warehouse name"), max_length=255
    )
    company_name = models.CharField(
        pgettext_lazy("Warehouse field description", "Legal company name"),
        blank=True,
        max_length=255,
    )

    shipping_zones = models.ManyToManyField(ShippingZone, blank=True)
    address = models.ForeignKey(Address, on_delete=models.PROTECT)

    email = models.EmailField(
        pgettext_lazy("Warehouse field description", "Email address"),
        blank=True,
        default="",
    )
{'fraction_non_alphanumeric': 0.08960244648318043, 'fraction_numerical': 0.004281345565749235, 'mean_word_length': 3.860326894502229, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '90643', 'n_tokens_mistral': 943, 'n_tokens_neox': 922, 'n_words': 215}