BuyAndKnow/bk_api/bills/management/commands/populate_db.py

import json
import random
from datetime import date, timedelta

from django.core.management.base import BaseCommand
from django.contrib.auth.models import User

from bills.serializers import TagSerializer, ShopTagSerializer, ShopSerializer, ReceiptSerializer

class Command(BaseCommand):
    args = '<foo bar ...>'
    help = 'our help string comes here'

    def _load_data(self):
        with open("./bills/management/commands/data.json", 'r') as json_file:
            self.data = json.load(json_file)
        
    def _user_exists(self, username):
        return User.objects.filter(username=username).exists()

    def _create_user(self, username: str):
        password = f"{username}_pass"
        user_exists = self._user_exists(username)

        if user_exists:
            if input(f">: Delete {username} from db? [y/N]: ") != 'y':
                return None
            else:
                User.objects.get(username=username).delete()
                self.stdout.write(f"> User '{username}' deleted!")

        user = User.objects.create_user(username, "", password)
        self.stdout.write(f"> Created user '{username}' with password '{password}'")
        return user

    def _create_tags(self, user):
        created_tags = []
        for tag in self.data["tags"]:
            created_tag = TagSerializer(data=tag)
            created_tag.is_valid()
            created_tag.save(owner=user)
            self.stdout.write(f"> Created tag: '{created_tag.data['name']}'")
            created_tags.append(created_tag.data['name'])

        self.stdout.write(f"> Created tags")
        return created_tags

    def _create_shop_tags(self, user):
        created_tags = []
        for tag in self.data["shop_tags"]:
            created_tag = ShopTagSerializer(data=tag)
            created_tag.is_valid()
            created_tag.save(owner=user)
            self.stdout.write(f"> Created shop tag: '{created_tag.data['name']}'")
            created_tags.append(created_tag.data['name'])

        self.stdout.write(f"> Created shop tags")
        return created_tags


    def _create_shops(self, user, shop_tags):
        created_shops = []
        for shop in self.data["shops"]:
            created_shop = ShopSerializer(data=shop)
            created_shop.is_valid()
            created_shop.save(owner=user)
            self.stdout.write(f"> Created shop: '{created_shop.data['name']}'")
            created_shops.append(created_shop.data)

        self.stdout.write(f"> Created shop")
        return created_shops

    def _create_receipts(self, user, shops, tags, shop_tags):
        # dates for last 100 days
        for d in [date.today() - timedelta(100) + timedelta(n)for n in range(100)]:
            if random.random() > 0.4:
                continue
            receipt_date = f"{d.year}-{d.month}-{d.day}"
            receipt_schop = random.choice(shops)
            receipt_products = [random.choice(self.data["products"])]

            for product in self.data["products"]:
                if random.random() < 0.20:
                    receipt_products.append(product)

            new_receipt_data = {
                "shop": receipt_schop,
                "date": receipt_date,
                "products": receipt_products
            }
            
            created_receipt = ReceiptSerializer(data=new_receipt_data)
            is_valid = created_receipt.is_valid()
            if not is_valid:            
                print(created_receipt.errors)
                break;
            created_receipt.save(owner=user)
            self.stdout.write(f"> Created receipt: {receipt_date} | {receipt_schop['name']}")
        
        

    def handle(self, *args, **options):
        self.stdout.write(f"_/-\__/-\_POPULATE DB _/-\__/-\_")
        user = self._create_user("demo_user")
        if user != None:
            self._load_data()
            tags = self._create_tags(user)
            shop_tags = self._create_shop_tags(user)
            shops = self._create_shops(user, shop_tags)
            self._create_receipts(user,shops, tags, shop_tags)
        self.stdout.write(f"_/-\__/-\_ DONE _/-\__/-\_")