from record import *

RECORD_TYPES = [
    'SubmitterRecord',
    'EmployerRecord',
    'EmployeeWageRecord',
    'OptionalEmployeeWageRecord',
    'TotalRecord',
    'StateWageRecord',
    'OptionalTotalRecord',
    'StateTotalRecord',
    'FinalRecord'
    ]

def test():
    import record, model
    for rname in RECORD_TYPES:
        inst = record.__dict__[rname]()
        print type(inst), inst.record_identifier, len(inst.output())


def test_dump():
    import record, StringIO
    records = [
        record.SubmitterRecord(),
        record.EmployerRecord(),
        record.EmployeeWageRecord(),
        ]
    out = StringIO.StringIO()
    dump(records, out)
    return out


def test_load(fp):
    return load(fp)

def load(fp):
    # BUILD LIST OF RECORD TYPES
    import record
    types = {}
    for r in RECORD_TYPES:
        klass = record.__dict__[r]
        types[klass.record_identifier] = klass

    # PARSE DATA INTO RECORDS AND YIELD THEM
    while fp.tell() < fp.len:
        record_ident = fp.read(2)
        if record_ident in types:
            record = types[record_ident]()
            record.read(fp)
            yield record


def loads(s):
    import StringIO
    fp = StringIO.StringIO(s)
    return load(fp)


def dump(records, fp):
    for r in records:
        fp.write(r.output())


def dumps(records):
    import StringIO
    fp = StringIO.StringIO()
    dump(records, fp)
    fp.seek(0)
    return fp.read()