PyConES24/src/retailtwin/simulation.py
2023-09-18 09:32:03 +02:00

45 lines
1.4 KiB
Python

from datetime import datetime
from random import choice
import polars as pl
from scipy.stats.distributions import lognorm
from sqlalchemy.orm import Session
from dengfun.retail.models import Cart
def cart(session: Session) -> Cart:
db_uri = str(session.get_bind().url.render_as_string(hide_password=False))
customer = pl.read_database("select random_customer()", db_uri, engine="adbc")[0, 0]
location = pl.read_database("select random_location()", db_uri, engine="adbc")[0, 0]
cart = Cart(customer=customer, location=location, checkout=datetime.now())
session.add(cart)
session.commit()
# Budget can be a small or a large expenditure
budget = choice([10 * lognorm.rvs(0.25), 100 * lognorm.rvs(0.25)])
items = pl.read_database("select * from items", db_uri, engine="adbc")
spent = 0
while spent > budget:
quantity = choice([1, 2, 3, 4])
item = items.sample(1, shuffle=True)
price = pl.read_database(
"select price_on_location({id}, {qty})", db_uri, engine="adbc"
)[0, 0]
session.execute(
f"add_item_to_cart({cart.id}, {item.sku}, {quantity}, {location})"
)
session.commit()
spent += price
if __name__ == "__main__":
from sqlalchemy import create_engine
from sqlalchemy.orm import Session
with Session(
create_engine("postgresql://postgres:postgres@localhost/retail", echo=True)
) as session:
cart(session)