gear-orders/generate.py

import yaml
import json
import logging
import random

from db.queries import repeat_get, repeat_increment, repeat_clear, repeat_put
from settings import ORDERS_YML

logger = logging.getLogger(__name__)

def pick_order(orders):
    picked = random.choices(orders,
                            weights=[o['weight'] for o in orders],
                            k=1)[0]

    result = []
    repeat = 0.0
    
    if('text' in picked):
        result.append(picked['text'])

    if('repeat' in picked):
        repeat = picked['repeat']

    if('add' in picked):
        for addition in picked['add']:
            if addition['probability'] > random.random():
                result.append(addition['text'])
        
    if('pick' in picked):
        (new_result, new_repeat) = pick_order(picked['pick'])
        result += new_result
        if new_repeat > 0.0:
            repeat = new_repeat

    return (result, repeat)

def read_config():
    with open(ORDERS_YML) as stream:
        try:
            orders = yaml.safe_load(stream)
        except yaml.YAMLError as exc:
            logger.error(exc)
    return orders

def generate_order():
    # Do we want to repeat?
    repeat = repeat_get()
    if repeat is not None:
        if repeat.probability > random.random():
            repeat_increment()
            return {
                "orders": json.loads(repeat.orders),
                "count": repeat.count
            }
        else:
            repeat_clear()
    
    orders_config = read_config()

    if orders_config['orders_probability'] < random.random():
        # No orders today
        return { }

    # Pick new orders
    (result, repeat_p) = pick_order(orders_config['orders'])

    # Log the repeat
    if repeat_p > 0.0:
        repeat_put(repeat_p, json.dumps(result))
    
    return {
        "orders": result
    }

def generate_punishment():
    orders_config = read_config()

    (result, repeat_p) = pick_order(orders_config['punishments'])

    return result