diff options
Diffstat (limited to 'ofxparse/ofxprinter.py')
| -rw-r--r-- | ofxparse/ofxprinter.py | 28 |
1 files changed, 17 insertions, 11 deletions
diff --git a/ofxparse/ofxprinter.py b/ofxparse/ofxprinter.py index 4c0a13d..e39b895 100644 --- a/ofxparse/ofxprinter.py +++ b/ofxparse/ofxprinter.py @@ -1,3 +1,6 @@ +import six + + class OfxPrinter(): ofx = None out_filename = None @@ -24,7 +27,7 @@ class OfxPrinter(): ) def writeHeaders(self): - for k, v in self.ofx.headers.iteritems(): + for k, v in six.iteritems(self.ofx.headers): if v is None: self.writeLine("{0}:NONE".format(k)) else: @@ -71,8 +74,8 @@ class OfxPrinter(): self.writeLine("</STMTTRN>", tabs=tabs) def writeLedgerBal(self, statement, tabs=4): - bal = getattr(statement, 'balance') - baldt = getattr(statement, 'balance_date') + bal = getattr(statement, 'balance', None) + baldt = getattr(statement, 'balance_date', None) if bal and baldt: self.writeLine("<LEDGERBAL>", tabs=tabs) @@ -83,8 +86,8 @@ class OfxPrinter(): self.writeLine("</LEDGERBAL>", tabs=tabs) def writeAvailBal(self, statement, tabs=4): - bal = getattr(statement, 'available_balance') - baldt = getattr(statement, 'available_balance_date') + bal = getattr(statement, 'available_balance', None) + baldt = getattr(statement, 'available_balance_date', None) if bal and baldt: self.writeLine("<AVAILBAL>", tabs=tabs) @@ -176,19 +179,22 @@ class OfxPrinter(): # No newline at end of file self.writeLine("</OFX>", tabs=tabs, term="") - def write(self, filename=None, tabs=0): + def writeToFile(self, fileObject, tabs=0): if self.out_handle: raise Exception("Already writing file") - if filename is None: - filename = self.out_filename - - self.out_handle = open(filename, 'wb') + self.out_handle = fileObject self.writeHeaders() self.writeOfx(tabs=tabs) self.out_handle.flush() - self.out_handle.close() self.out_handle = None + + def write(self, filename=None, tabs=0): + if filename is None: + filename = self.out_filename + + with open(filename, 'w') as f: + self.writeToFile(f) |