timeflake_peewee.py

import os
import sys
import inspect

import timeflake
from peewee import *
from timeflake.extensions.peewee import TimeflakeBase62Field, TimeflakeUUIDField

# config

db_file = '/tmp/peewee.db'
DB = SqliteDatabase(db_file)


# models

class BaseModel(Model):
    class Meta:
        database = DB


class B62_Model(BaseModel):
    class Meta:
        table_name = 'table_a'

    id = TimeflakeBase62Field(primary_key=True, default=timeflake.random)
    foo = IntegerField(null=True)


class BU_Model(BaseModel):
    class Meta:
        table_name = 'table_b'

    id = TimeflakeUUIDField(primary_key=True, default=timeflake.random)
    foo = IntegerField(null=True)


# setup

def create_all_tables():
    for cls in sys.modules[__name__].__dict__.values():
        if not inspect.isclass(cls):
            continue
        if not hasattr(cls, '__bases__'):
            continue
        if not issubclass(cls, Model):
            continue
        if cls is Model:
            continue
        if cls is not BaseModel:
            cls.create_table()


def insert_fixtures(typ, num=None):
    for i in range(num or 5):
        yield typ.create(foo=i).id


def load_objects(typ, id):
    return typ.get(typ.id == id)


def test_models(*args):
    for typ in args:
        test_model(typ)


def test_model(typ):
    for id in insert_fixtures(typ, 16):
        obj = load_objects(typ, id)
        assert isinstance(obj, typ)
        assert isinstance(obj.id, timeflake.Timeflake)
        assert isinstance(obj.foo, int)
        print("\t".join(map(str, [
            typ,
            obj.foo,
            obj.id.base62,
            obj.id.uuid,
        ])))



if not os.path.isfile(db_file):
    create_all_tables()
    test_models(B62_Model, BU_Model)