kpsul/kfet/management/commands/createopes.py

"""
Crée des opérations aléatoires réparties sur une période de temps spécifiée
"""

import random
from datetime import timedelta
from decimal import Decimal

from django.core.management.base import BaseCommand
from django.utils import timezone

from kfet.models import (
    Account,
    Article,
    Checkout,
    Operation,
    OperationGroup,
    Transfer,
    TransferGroup,
)


class Command(BaseCommand):
    help = "Crée des opérations réparties uniformément " "sur une période de temps"

    def add_arguments(self, parser):
        # Nombre d'opérations à créer
        parser.add_argument("opes", type=int, help="Number of opegroups to create")

        # Période sur laquelle créer (depuis num_days avant maintenant)
        parser.add_argument(
            "days", type=int, help="Period in which to create opegroups"
        )

        # Optionnel : nombre de transfert à créer (défaut 0)
        parser.add_argument(
            "--transfers",
            type=int,
            default=0,
            help="Number of transfers to create (default 0)",
        )

    def handle(self, *args, **options):
        self.stdout.write("Génération d'opérations")

        # Output log vars
        opes_created = 0
        purchases = 0
        transfers = 0

        num_ops = options["opes"]
        num_transfers = options["transfers"]
        # Convert to seconds
        time = options["days"] * 24 * 3600

        now = timezone.now()
        checkout = Checkout.objects.first()
        articles = Article.objects.all()
        accounts = Account.objects.exclude(trigramme="LIQ")
        liq_account = Account.objects.get(trigramme="LIQ")
        try:
            con_account = Account.objects.get(
                cofprofile__user__first_name="Assurancetourix"
            )
        except Account.DoesNotExist:
            con_account = random.choice(accounts)

        # use to fetch OperationGroup pk created by bulk_create
        at_list = []
        # use to lazy set OperationGroup pk on Operation objects
        ope_by_grp = []
        # OperationGroup objects to bulk_create
        opegroup_list = []

        for i in range(num_ops):
            # Randomly pick account
            if random.random() > 0.25:
                account = random.choice(accounts)
            else:
                account = liq_account

            # Randomly pick time
            at = now - timedelta(seconds=random.randint(0, time))

            # Majoration sur compte 'concert'
            if random.random() < 0.2:
                addcost = True
                addcost_for = con_account
                addcost_amount = Decimal("0.5")
            else:
                addcost = False

            # Initialize opegroup amount
            amount = Decimal("0")

            # Generating operations
            ope_list = []
            for j in range(random.randint(1, 4)):
                # Operation type
                typevar = random.random()

                # 0.1 probability to have a charge
                if typevar > 0.9 and account != liq_account:
                    ope = Operation(
                        type=Operation.DEPOSIT,
                        amount=Decimal(random.randint(1, 99) / 10),
                    )
                # 0.05 probability to have a withdrawal
                elif typevar > 0.85 and account != liq_account:
                    ope = Operation(
                        type=Operation.WITHDRAW,
                        amount=-Decimal(random.randint(1, 99) / 10),
                    )
                # 0.05 probability to have an edition
                elif typevar > 0.8 and account != liq_account:
                    ope = Operation(
                        type=Operation.EDIT, amount=Decimal(random.randint(1, 99) / 10)
                    )
                else:
                    article = random.choice(articles)
                    nb = random.randint(1, 5)

                    ope = Operation(
                        type=Operation.PURCHASE,
                        amount=-article.price * nb,
                        article=article,
                        article_nb=nb,
                    )

                    purchases += 1

                    if addcost:
                        ope.addcost_for = addcost_for
                        ope.addcost_amount = addcost_amount * nb
                        ope.amount -= ope.addcost_amount

                ope_list.append(ope)
                amount += ope.amount

            opegroup_list.append(
                OperationGroup(
                    on_acc=account,
                    checkout=checkout,
                    at=at,
                    is_cof=account.cofprofile.is_cof,
                    amount=amount,
                )
            )
            at_list.append(at)
            ope_by_grp.append((at, ope_list))

        OperationGroup.objects.bulk_create(opegroup_list)

        # Fetch created OperationGroup objects pk by at
        opegroups = OperationGroup.objects.filter(at__in=at_list).values("id", "at")
        opegroups_by = {grp["at"]: grp["id"] for grp in opegroups}

        all_ope = []
        for _ in range(num_ops):
            at, ope_list = ope_by_grp.pop()
            for ope in ope_list:
                ope.group_id = opegroups_by[at]
                all_ope.append(ope)

        Operation.objects.bulk_create(all_ope)
        opes_created = len(all_ope)

        # Transfer generation

        transfer_by_grp = []
        transfergroup_list = []
        at_list = []

        for i in range(num_transfers):
            # Randomly pick time
            at = now - timedelta(seconds=random.randint(0, time))

            # Choose whether to have a comment
            if random.random() > 0.5:
                comment = "placeholder comment"
            else:
                comment = ""

            transfergroup_list.append(
                TransferGroup(at=at, comment=comment, valid_by=random.choice(accounts))
            )
            at_list.append(at)

            # Randomly generate transfer
            transfer_list = []
            for i in range(random.randint(1, 4)):
                transfer_list.append(
                    Transfer(
                        from_acc=random.choice(accounts),
                        to_acc=random.choice(accounts),
                        amount=Decimal(random.randint(1, 99) / 10),
                    )
                )

            transfer_by_grp.append((at, transfer_list))

        TransferGroup.objects.bulk_create(transfergroup_list)

        transfergroups = TransferGroup.objects.filter(at__in=at_list).values("id", "at")
        transfergroups_by = {grp["at"]: grp["id"] for grp in transfergroups}

        all_transfer = []
        for _ in range(num_transfers):
            at, transfer_list = transfer_by_grp.pop()
            for transfer in transfer_list:
                transfer.group_id = transfergroups_by[at]
                all_transfer.append(transfer)

        Transfer.objects.bulk_create(all_transfer)
        transfers += len(all_transfer)

        self.stdout.write(
            "- {:d} opérations créées dont {:d} commandes d'articles".format(
                opes_created, purchases
            )
        )

        if transfers:
            self.stdout.write("- {:d} transferts créés".format(transfers))