[ADD] matching and parsing of CSV, preview production

bzr revid: xmo@openerp.com-20120813134219-ijo7d2kwn82qs4tm
2012-08-13 15:42:19 +02:00
parent b34767381b
commit 45d5003dd6
3 changed files with 292 additions and 11 deletions
@@ -1,5 +1,8 @@
+import base64
+import codecs
 import csv
 import itertools
+import logging

 try:
    from cStringIO import StringIO
@@ -10,6 +13,8 @@ from openerp.osv import orm, fields
 from openerp.tools.translate import _

 FIELDS_RECURSION_LIMIT = 2
+ERROR_PREVIEW_BYTES = 200
+_logger = logging.getLogger(__name__)
 class ir_import(orm.TransientModel):
    _name = 'base_import.import'

@@ -103,8 +108,87 @@ class ir_import(orm.TransientModel):

            fields.append(f)

+        # TODO: cache on model?
        return fields

+    def _read_csv(self, record, options):
+        """ Returns a CSV-parsed iterator of all empty lines in the file
+
+        :throws csv.Error: if an error is detected during CSV parsing
+        :throws UnicodeDecodeError: if ``options.encoding`` is incorrect
+        """
+        csv_iterator = csv.reader(
+            StringIO(base64.b64decode(record.file)),
+            quotechar=options['quote'],
+            delimiter=options['separator'])
+        csv_nonempty = itertools.ifilter(None, csv_iterator)
+        # TODO: guess encoding?
+        encoding = options.get('encoding', 'utf-8')
+        return itertools.imap(
+            lambda row: [item.decode(encoding) for item in row],
+            csv_nonempty)
+
+    def _match_header(self, header, fields, options):
+        """ Attempts to match a given header to a field of the
+        imported model.
+
+        :param str header: header name from the CSV file
+        :param fields:
+        :param dict options:
+        :returns: an empty list if the header couldn't be matched, or
+                  all the fields to traverse
+        :rtype: list(Field)
+        """
+        for field in fields:
+            # FIXME: should match all translations & original
+            # TODO: use string distance (levenshtein? hamming?)
+            if header == field['name'] \
+              or header.lower() == field['string'].lower():
+                return [field]
+
+        if '/' not in header:
+            return []
+
+        # relational field path
+        traversal = []
+        subfields = fields
+        # Iteratively dive into fields tree
+        for section in header.split('/'):
+            # Strip section in case spaces are added around '/' for
+            # readability of paths
+            match = self._match_header(section.strip(), subfields, options)
+            # Any match failure, exit
+            if not match: return []
+            # prep subfields for next iteration within match[0]
+            field = match[0]
+            subfields = field['fields']
+            traversal.append(field)
+        return traversal
+
+    def _match_headers(self, rows, fields, options):
+        """ Attempts to match the imported model's fields to the
+        titles of the parsed CSV file, if the file is supposed to have
+        headers.
+
+        Will consume the first line of the ``rows`` iterator.
+
+        Returns either None (no title) or a dict mapping cell indices
+        to key paths in the ``fields`` tree
+
+        :param Iterator rows:
+        :param dict fields:
+        :param dict options:
+        :rtype: None | dict(int: list(str))
+        """
+        if not options.get('headers'):
+            return None
+
+        headers = next(rows)
+        return dict(
+            (index, [field['name'] for field in self._match_header(header, fields, options)] or None)
+            for index, header in enumerate(headers)
+        )
+
    def parse_preview(self, cr, uid, id, options, count=10, context=None):
        """ Generates a preview of the uploaded files, and performs
        fields-matching between the import's file data and the model's
@@ -112,16 +196,38 @@ class ir_import(orm.TransientModel):

        :param id: identifier of the import
        :param int count: number of preview lines to generate
-        :param dict options: format-specific options
-        :returns: (fields, matches, preview)
-        :rtype: (dict(str: dict(...)), dict(str, str), list(list(str)))
+        :param options: format-specific options.
+                        CSV: {encoding, quote, separator, headers}
+        :type options: {str, str, str, bool}
+        :returns: {fields, matches, preview} | {error, preview}
+        :rtype: {dict(str: dict(...)), dict(int, list(str)), list(list(str))} | {str, str}
        """
        record = self.browse(cr, uid, id, context=context)
-        # recursive fields_get (cache based on res_model?)
-        fields = record.get_fields(record.res_model)
-        import pprint
-        pprint.pprint(fields)
-        # extract title row?
-        #    match title row to fields_ge
-        # Extract first $count rows
-        # return triplet
+        fields = self.get_fields(cr, uid, record.res_model, context=context)
+
+        try:
+            rows = self._read_csv(record, options)
+
+            match = self._match_headers(rows, fields, options)
+            # Match should have consumed the first row (iif headers), get
+            # the ``count`` next rows for preview
+            preview = itertools.islice(rows, count)
+            return {
+                'fields': fields,
+                'matches': match,
+                'preview': list(preview),
+            }
+        except (TypeError, UnicodeDecodeError), e:
+            # Due to lazy generators, UnicodeDecodeError (for
+            # instance) may only be raised when serializing the
+            # preview to a list in the return.
+            _logger.debug("Error during CSV parsing preview", exc_info=True)
+            return {
+                'error': _("Failed to parse CSV file: %s") % e,
+                # iso-8859-1 ensures decoding will always succeed,
+                # even if it yields non-printable characters. This is
+                # in case of UnicodeDecodeError (or csv.Error
+                # compounded with UnicodeDecodeError)
+                'preview': base64.b64decode(record.file)[:ERROR_PREVIEW_BYTES]\
+                                 .decode('iso-8859-1'),
+            }
@@ -90,3 +90,12 @@ class o2m_child(orm.Model):
        'parent_id': fields.many2one(name('o2m')),
        'value': fields.integer()
    }
+
+class preview_model(orm.Model):
+    _name = name('preview')
+
+    _columns = {
+        'name': fields.char('Name', size=None),
+        'somevalue': fields.integer('Some Value', required=True),
+        'othervalue': fields.integer('Other Variable'),
+    }
@@ -1,5 +1,9 @@
+# -*- encoding: utf-8 -*-
+import unittest2
 from openerp.tests.common import TransactionCase

+from .. import models
+
 ID_FIELD = {'id': 'id', 'name': 'id', 'string': "External ID", 'required': False, 'fields': []}
 def make_field(name='value', string='unknown', required=False, fields=[]):
    return [
@@ -71,3 +75,165 @@ class test_o2m(TransactionCase):
            ]},
            {'id': 'value', 'name': 'value', 'string': 'unknown', 'required': False, 'fields': []},
        ]))
+
+class test_match_headers_single(TransactionCase):
+    def test_match_by_name(self):
+        match = self.registry('base_import.import')._match_header(
+            'f0', [{'name': 'f0'}], {})
+
+        self.assertEqual(match, [{'name': 'f0'}])
+
+    def test_match_by_string(self):
+        match = self.registry('base_import.import')._match_header(
+            'some field', [{'name': 'bob', 'string': "Some Field"}], {})
+
+        self.assertEqual(match, [{'name': 'bob', 'string': "Some Field"}])
+
+    def test_nomatch(self):
+        match = self.registry('base_import.import')._match_header(
+            'should not be', [{'name': 'bob', 'string': "wheee"}], {})
+
+        self.assertEqual(match, [])
+
+    def test_recursive_match(self):
+        f = {
+            'name': 'f0',
+            'string': "My Field",
+            'fields': [
+                {'name': 'f0', 'string': "Sub field 0", 'fields': []},
+                {'name': 'f1', 'string': "Sub field 2", 'fields': []},
+            ]
+        }
+        match = self.registry('base_import.import')._match_header(
+            'f0/f1', [f], {})
+
+        self.assertEqual(match, [f, f['fields'][1]])
+
+    def test_recursive_nomatch(self):
+        """ Match first level, fail to match second level
+        """
+        f = {
+            'name': 'f0',
+            'string': "My Field",
+            'fields': [
+                {'name': 'f0', 'string': "Sub field 0", 'fields': []},
+                {'name': 'f1', 'string': "Sub field 2", 'fields': []},
+            ]
+        }
+        match = self.registry('base_import.import')._match_header(
+            'f0/f2', [f], {})
+
+        self.assertEqual(match, [])
+
+class test_match_headers_multiple(TransactionCase):
+    def test_noheaders(self):
+        self.assertIsNone(
+            self.registry('base_import.import')._match_headers(
+                [], [], {}))
+    def test_nomatch(self):
+        self.assertEqual(
+            self.registry('base_import.import')._match_headers(
+                iter([
+                    ['foo', 'bar', 'baz', 'qux'],
+                    ['v1', 'v2', 'v3', 'v4'],
+                ]),
+                [],
+                {'headers': True}),
+            dict((index, None) for index in range(4)))
+
+    def test_mixed(self):
+        self.assertEqual(
+            self.registry('base_import.import')._match_headers(
+                iter(['foo bar baz qux/corge'.split()]),
+                [
+                    {'name': 'bar', 'string': 'Bar'},
+                    {'name': 'bob', 'string': 'Baz'},
+                    {'name': 'qux', 'string': 'Qux', 'fields': [
+                        {'name': 'corge', 'fields': []},
+                     ]}
+                ],
+                {'headers': True}),
+            {
+                0: None,
+                1: ['bar'],
+                2: ['bob'],
+                3: ['qux', 'corge'],
+            })
+
+import base64, csv
+class test_preview(TransactionCase):
+    def make_import(self):
+        Import = self.registry('base_import.import')
+        id = Import.create(self.cr, self.uid, {
+            'res_model': 'res.users',
+            'file': base64.b64encode(
+                u"로그인,언어\n".encode('euc_kr'),
+                "bob,1\n"),
+        })
+        return Import, id
+
+    def test_encoding(self):
+        Import, id = self.make_import()
+        result = Import.parse_preview(self.cr, self.uid, id, {
+                'quote': '"',
+                'separator': ',',
+        })
+        self.assertTrue('error' in result)
+
+    def test_csv_errors(self):
+        Import, id = self.make_import()
+
+        result = Import.parse_preview(self.cr, self.uid, id, {
+                'quote': 'foo',
+                'separator': ',',
+                'encoding': 'euc_kr',
+        })
+        self.assertTrue('error' in result)
+
+    def test_csv_errors(self):
+        Import, id = self.make_import()
+
+        result = Import.parse_preview(self.cr, self.uid, id, {
+                'quote': '"',
+                'separator': 'bob',
+                'encoding': 'euc_kr',
+        })
+        self.assertTrue('error' in result)
+
+    def test_success(self):
+        Import = self.registry('base_import.import')
+        id = Import.create(self.cr, self.uid, {
+            'res_model': 'base_import.tests.models.preview',
+            'file': base64.b64encode('name,Some Value,Counter\n'
+                    'foo,1,2\n'
+                    'bar,3,4\n'
+                    'qux,5,6\n')
+        })
+
+        result = Import.parse_preview(self.cr, self.uid, id, {
+            'quote': '"',
+            'separator': ',',
+            'headers': True,
+        })
+
+        p = {
+            'matches': {0: ['name'], 1: ['somevalue'], 2: None},
+            'fields': [
+                {'id': 'id', 'name': 'id', 'string': 'External ID', 'required':False, 'fields': []},
+                {'id': 'name', 'name': 'name', 'string': 'Name', 'required':False, 'fields': []},
+                {'id': 'somevalue', 'name': 'somevalue', 'string': 'Some Value', 'required':True, 'fields': []},
+                {'id': 'othervalue', 'name': 'othervalue', 'string': 'Other Variable', 'required':False, 'fields': []},
+            ],
+            'preview': [
+                ['foo', '1', '2'],
+                ['bar', '3', '4'],
+                ['qux', '5', '6'],
+            ]
+        }
+
+        self.assertEqual(result['matches'], p['matches'])
+        # Order depends on iteration order of fields_get
+        self.assertItemsEqual(result['fields'], p['fields'])
+        self.assertEqual(result['preview'], p['preview'])
+        # Ensure we only have the response fields we expect
+        self.assertItemsEqual(result.keys(), ['matches', 'fields', 'preview'])