Demo data import tool for Tryton.

demo_import.py

#!/usr/bin/env python2

import csv
from proteus import config, Model, Wizard
from xmlrpclib import Fault
from datetime import date

def init_database(module_name):
    Module = Model.get('ir.module.module')
    (school_module,) = Module.find([('name', '=', module_name)])
    Module.button_install([school_module.id], config.context)
    Wizard('ir.module.module.install_upgrade').execute('start')

def search(model, column, value):
    records = Model.get(model).find()
    for record in records:
        if getattr(record, column) == value:
            return record
    return None

def load_csv(file_name, module_name):
    try:
        csv_file = csv.reader(open(file_name))
    except IOError:
        return
    model_name = file_name[:-4]
    try:
        TrytonModel = Model.get(module_name + '.' + model_name)
    except Fault, err:
        print "Model error for:"
        print "  Model:   %s" % model_name
        print repr(err)
        return False

    header = True
    for row in csv_file:
        if header == True:
            headers = row
            header = False
            continue
        record = TrytonModel()
        for column in headers:
            if not hasattr(record, column):
                continue
            col_index = headers.index(column)
            type_ = record._fields[column]['type']
            if  row[col_index] == '':
                continue
            elif type_ == 'integer':
                setattr(record, column, int(row[col_index]))
            elif type_ == 'date':
                (year, month, day) = [int(d) for d in row[col_index].split('-')]
                datevalue = date(year, month, day)
                repr(datevalue)
                setattr(record, column, datevalue)
            elif type_  in ('one2many', 'many2many'):
                continue
            elif type_ in ('many2one', 'one2one'):
                model = record._fields[column]['relation']
                result = search(model, 'name', row[col_index])
                if  result:
                    setattr(record, column, result)
                else:
                    RelatedModel = Model.get(record._fields[column]['relation'])
                    related_record = RelatedModel(**{'name': row[col_index]})
                    related_record.save()
                    setattr(record, column, related_record)
            else:
                setattr(record, column, row[col_index])
        try:
            record.save()
        except Fault, err:
            from pprint import pprint
            print "Save error for:"
            print "         Model: %s" % model_name
            print "        Record: %s" % row
            pprint(record._fields)
            print repr(err)
            return False

if __name__ == "__main__":
    config = config.set_xmlrpc('http://admin:admin@localhost:8069/school_demo')
    #config = config.set_trytond('school_demo', database_type='postgresql',
    #                            user='admin', password='admin')
    module_name = 'school'
    init_database(module_name)
    files = [
        'school.csv',
        'year.csv',
        'teacher.csv',
        'class_group.csv',
        'guardian.csv',
        'student.csv',
        'day.csv',
        'period.csv',
        'subject.csv',
        'room.csv',
        'timetable.csv',
        'relationship.csv',
        'group_members',
        ]
    for f in files:
        if load_csv(f, module_name):
            print 'Loaded %s' % f