changeset 25:4fce8272c8d9

Allow to define file encoding and to use file like object bug9
author Nicolas Évrard <nicoe@b2ck.com>
date Fri, 24 Feb 2017 19:21:19 +0100
parents 111f8d6c1d79
children 1f237393bb05
files CHANGELOG mt940/__init__.py mt940/test.py
diffstat 3 files changed, 52 insertions(+), 34 deletions(-) [+]
line wrap: on
line diff
--- a/CHANGELOG	Tue Feb 21 18:23:29 2017 +0100
+++ b/CHANGELOG	Fri Feb 24 19:21:19 2017 +0100
@@ -1,3 +1,6 @@
+* Allow to use file-like object
+* Allow to define file encoding
+
 Version 0.3 - 2016-01-27
 * Add support for RegioBank description format
 * Set a default description after setting a statement
--- a/mt940/__init__.py	Tue Feb 21 18:23:29 2017 +0100
+++ b/mt940/__init__.py	Fri Feb 24 19:21:19 2017 +0100
@@ -2,8 +2,8 @@
 # -*- coding: utf-8 -*-
 #
 # Copyright (c) 2013-2016, Cédric Krier
-# Copyright (c) 2014-2015, Nicolas Évrard
-# Copyright (c) 2013-2016, B2CK
+# Copyright (c) 2014-2017, Nicolas Évrard
+# Copyright (c) 2013-2017, B2CK
 # All rights reserved.
 #
 # Redistribution and use in source and binary forms, with or without
@@ -37,6 +37,7 @@
 from decimal import Decimal
 import datetime
 import re
+import io
 
 
 SECTIONS = {
@@ -73,39 +74,45 @@
 
 class MT940(object):
 
-    def __init__(self, name):
+    def __init__(self, name, encoding=None):
         self.statements = []
 
-        with open(name, 'rU') as f:
-            values = defaultdict(str)
-            # Set optional values
-            values['description']
-            transactions = []
-            for line in self._readline(f):
-                for name, sections in SECTIONS.iteritems():
-                    if name == 'begin':
-                        continue
-                    for section in sections:
-                        if line.startswith(section):
-                            if name in values and name == 'statement':
-                                self._set_statement(values, transactions)
-                            if name.endswith('_balance'):
-                                values[name] = self._get_balance(
-                                    line[len(section):])
-                            elif name == 'transaction':
-                                transactions.append(
-                                    self._get_transaction(line[len(section):]))
-                            elif name == 'description':
-                                description = line[len(section):]
-                                if 'end_balance' in values:
-                                    values['description'] += description
-                                else:
-                                    transactions[-1] = (transactions[-1][:-1]
-                                        + (description,))
+        if isinstance(name, (bytes, basestring)):
+            with io.open(name, encoding=encoding, mode='r') as f:
+                self._parse(f)
+        else:
+            self._parse(name)
+
+    def _parse(self, f):
+        values = defaultdict(str)
+        # Set optional values
+        values['description']
+        transactions = []
+        for line in self._readline(f):
+            for name, sections in SECTIONS.iteritems():
+                if name == 'begin':
+                    continue
+                for section in sections:
+                    if line.startswith(section):
+                        if name in values and name == 'statement':
+                            self._set_statement(values, transactions)
+                        if name.endswith('_balance'):
+                            values[name] = self._get_balance(
+                                line[len(section):])
+                        elif name == 'transaction':
+                            transactions.append(
+                                self._get_transaction(line[len(section):]))
+                        elif name == 'description':
+                            description = line[len(section):]
+                            if 'end_balance' in values:
+                                values['description'] += description
                             else:
-                                values[name] += line[len(section):]
-            if values:
-                self._set_statement(values, transactions)
+                                transactions[-1] = (transactions[-1][:-1]
+                                    + (description,))
+                        else:
+                            values[name] += line[len(section):]
+        if values:
+            self._set_statement(values, transactions)
 
     @staticmethod
     def _readline(f):
--- a/mt940/test.py	Tue Feb 21 18:23:29 2017 +0100
+++ b/mt940/test.py	Fri Feb 24 19:21:19 2017 +0100
@@ -2,8 +2,8 @@
 # -*- coding: utf-8 -*-
 #
 # Copyright (c) 2013, Cédric Krier
-# Copyright (c) 2014, Nicolas Évrard
-# Copyright (c) 2013-2014, B2CK
+# Copyright (c) 2014-2017, Nicolas Évrard
+# Copyright (c) 2013-2017, B2CK
 # All rights reserved.
 #
 # Redistribution and use in source and binary forms, with or without
@@ -30,6 +30,7 @@
 """Test MT940
 """
 import os
+import io
 import unittest
 import datetime
 from decimal import Decimal
@@ -108,6 +109,13 @@
         self.assertEqual(transaction.description, '')
 
 
+class TestMT940Stream(TestMT940):
+
+    def setUp(self):
+        self.mt940 = MT940(io.open(
+                os.path.join(here, 'MT940.txt'), encoding='ascii'))
+
+
 class TestRaboDescription(unittest.TestCase):
 
     def test_one_tag(self):