71 lines
1.4 KiB
Python
71 lines
1.4 KiB
Python
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()
|
|
|