se-simulator/todb.py

#!/usr/bin/env python
import random
from datetime import datetime

import sys
from slugify import slugify

import utils
from models import *
from text_generator import get_chain, generate_text


def get_unused_users(site, count):
    return User.select().join(Question, JOIN.LEFT_OUTER).switch(User).join(Answer, JOIN.LEFT_OUTER) \
        .where((User.site == site) & (Question.id.is_null()) & (Answer.id.is_null()))


def add_username(site, count=500):
    print("usernames")
    chain = get_chain(site.url, "Usernames")
    for _ in range(count):
        username = generate_text(chain, "Usernames")
        User.create(username=username, site=site)


def add_title(site, count=100):
    print("titles")
    # TODO: Make sure that every slug is unique
    chain = get_chain(site.url, "Titles")
    for _ in range(count):
        title = generate_text(chain, "Titles")
        slug = slugify(title, max_length=70, word_boundary=True)
        Title.create(text=title, slug=slug, site=site)


def add_answer(site, count=300):
    print("answers")
    users = get_unused_users(site, count)
    chain = get_chain(site.url, "Answers")

    for i in range(count):
        text = generate_text(chain, "Answers")
        user = users[i]
        time = datetime.now()
        Answer.create(text=text, user_id=user, site_id=site, datetime=time)


def add_question(site, count=100):
    print("Questions")
    users = get_unused_users(site, count)
    titles = Title.select().join(Question, JOIN.LEFT_OUTER) \
        .where((Title.site == site) & (Question.id.is_null())) \
        .limit(count)
    chain = get_chain(site.url, "Questions")

    for i in range(count):
        text = generate_text(chain, "Questions")
        title = titles[i]
        user = users[i]
        time = datetime.now()
        question = Question.create(text=text, title_id=title, user_id=user, site_id=site, datetime=time)
        num_answers = random.randint(1, 4)
        answers = Answer.select().where((Answer.site == site) & (Answer.question.is_null())).limit(num_answers)
        for answer in answers:
            answer.question = question
            answer.save()


if __name__ == "__main__":
    if len(sys.argv) > 1:
        sites = sys.argv[1:]
        query = Site.select().where((Site.last_download.is_null(False)) & (Site.url.in_(sites)))
    else:
        query = Site.select().join(Question, JOIN.LEFT_OUTER).where(
            (Question.id.is_null()) & (Site.last_download.is_null(False))).limit(10)
        print([i.name for i in query])
    for s in query:
        print(s.name)
        add_username(s)
        add_title(s)
        add_answer(s)
        add_question(s)