Ticket #5253: csvserializer.diff

tests.py

@@ +1 @@
+
+from cStringIO import StringIO
+
+from django.core.serializers.base import DeserializationError
+from django.db import models
+from django.test import TestCase
+
+from serializers import csvserializer
+
+class Author(models.Model):
+    first_name = models.TextField(null=True, default=None)
+    last_name = models.TextField()
+
+    def __unicode__(self):
+        if self.first_name:
+            return '%s %s' % (self.first_name, self.last_name)
+        else:
+            return self.last_name
+
+    class Meta:
+        unique_together = (('first_name','last_name'),)
+
+class Book(models.Model):
+    author = models.ForeignKey(Author)
+    title = models.TextField()
+    copies = models.IntegerField(default=0)
+
+    def __unicode__(self):
+        return '%s: %s' % (self.author, self.title)
+    
+    class Meta:
+        unique_together = (('author', 'title'),)
+
+class Loaner(models.Model):
+    name = models.TextField(unique=True)
+
+class Card(models.Model):
+    number = models.IntegerField()
+    books = models.ManyToManyField(Book)
+    owner = models.ForeignKey(Loaner)
+
+    def __unicode__(self):
+        return '%i' % self.number
+
+class HalfEmpty(models.Model):
+    nullable = models.CharField(maxlength=10, null=True)
+    not_nullable = models.CharField(maxlength=10)
+
+
+class CsvTest(TestCase):
+
+    def setUp(self):
+        #csvserializer.nice_foreign_keys = False
+        def save_list(list):
+            for o in list:
+                o.save()
+
+        self.authors = (
+          Author(last_name='Dante'),
+          Author(last_name='Milton', first_name='John'),
+          Author(last_name='Milton', first_name='Different'),
+        )
+        save_list(self.authors)
+        self.books = (
+          Book(author=self.authors[0], title='Paradise'),
+          Book(author=self.authors[0], title='Hell'),
+          Book(author=self.authors[1], title='Paradise'),
+        )
+        save_list(self.books)
+        self.loaners = (Loaner(name='Tom'), Loaner(name='Liz'))
+        save_list(self.loaners)
+        self.cards = (
+          Card(number=555, 
+            owner=self.loaners[0]),
+          Card(number=666,
+            owner=self.loaners[1]),
+        )
+        save_list(self.cards)
+
+        self.cards[0].books = [self.books[0], self.books[1]]
+        self.cards[1].books = [self.books[1]]
+
+    def tearDown(self):
+        for mod in (Author, Book, Loaner, Card):
+            mod.objects.all().delete()
+
+    def test_resolve_related(self):
+        field = Book._meta.get_field('author')
+        dante = self.authors[0]
+        def assert_found(obj, field, value):
+            self.assertEquals(obj, 
+              csvserializer.resolve_related(field, value))
+        def assert_not_found(field, value):
+            try:
+                csvserializer.resolve_related(field, value)
+                self.fail('Expected DeserializationError')
+            except DeserializationError:
+                pass
+
+        assert_found(dante, field, 1)
+        assert_found(dante, field, {'last_name': 'Dante'})
+        assert_found(dante, field, {'last_name': 'Dante', 'first_name': None})
+        assert_not_found(field, 999)
+        assert_not_found(field, {'last_name': 'Milton'})
+
+    def test_write_read(self):
+
+        def check(sequence):
+            ser = csvserializer.Serializer()
+            csv_result = ser.serialize(sequence)
+
+            ser = csvserializer.Deserializer(StringIO(csv_result))
+            deserialized = [o for o in ser]
+
+            self.assertEquals(len(sequence), len(deserialized))
+            for i in range(len(sequence)):
+                self.assertEquals(sequence[i], deserialized[i].object)
+
+            for o in sequence:
+                o.delete()
+            for o in deserialized:
+                o.save()
+
+        check(self.cards)
+
+
+class CsvEmptyFieldTest(TestCase):
+
+    def test_it(self):
+        half_empty = HalfEmpty(not_nullable='')
+        half_empty.save()
+        ser = csvserializer.Serializer()
+        csv_result = ser.serialize([half_empty])
+        lines = csv_result.splitlines()
+        #self.assertEquals('1,,', lines[1])
+
+        ser = csvserializer.Deserializer(StringIO(csv_result))
+        deserialized = [o for o in ser]
+        new_half_empty = deserialized[0].object
+        self.assertEquals(None, new_half_empty.nullable)
+        self.assertEquals('', new_half_empty.not_nullable)

csvserializer.py

@@ +1 @@
+
+"""
+Does serialization to and from a single csv file using the standard csv
+module.
+
+Features:
+- Many-to-many data is serialized as a list of ids or dicts.
+- Both None, and empty string is serialized as an empty string.  When
+  deserializing, the empty string becomes None for nullable fields,
+  empty string for not nullable fields.
+- If nice_foreign_keys is True, foreign keys get serialized as possibly
+  nested dictionaries of unique fields.  For example if we have 
+  
+    class Article(Model):
+        owner = ForeignKey(User) 
+        text = TextField()
+
+  It will be serialized as 
+ 
+    --- article: id,owner,text
+    1,{'username': u'john'},Bla bla
+
+Limitations:
+- Unicode data is not supported (yet).
+- The nice_foreign_keys feature doesn't work yet with unique foreign
+  keys.
+- Comments are not supported in the csv file -- well, this is rather a
+  limitation of the csv module.
+
+Configuration:
+- To enable --format=csv for manage dumpdata/loaddata, add the following
+  lines to your settings.py:
+
+    SERIALIZATION_MODULES = {'csv': 'path.to.csvserializer'}
+
+- To configure it, add the following lines, too:
+
+    CSV_SERIALIZER_OPTIONS = {
+        'dialect': 'excel',
+        'nice_foreign_keys': True,
+    }
+
+    The given dialect should be registered to the csv module (defaults
+    to the pre-defined 'excel').
+"""
+import csv
+
+from django.conf import settings
+from django.core.serializers import base
+from django.db import models
+from django.db.models.fields import IntegerField, FloatField, Field
+from django.utils.encoding import smart_unicode
+
+
+#__all__ = ('Serializer', 'Deserializer')
+
+# options
+dialect = 'excel'
+nice_foreign_keys = True
+header_start = '--- '
+globals().update(getattr(settings, 'CSV_SERIALIZER_OPTIONS', {}))
+
+def resolve_related(field, value):
+    """Return a model object for `field` represented as `value` where
+    the representation was produced by Serializer.represent()."""
+    if value is None:
+        return None
+    else:
+        model = field.rel.to
+        if isinstance(value, int):
+            query = {'pk': value}
+        elif isinstance(value, dict):
+            query = {}
+            # value dict may be recursive
+            for k, v in value.iteritems():
+                if isinstance(v, dict):
+                    query[k] = resolve_related(model._meta.get_field(k), v)
+                elif v is None:
+                    query['%s__isnull' % k] = True
+                else:
+                    query[k] = v
+        else:
+            raise base.DeserializationError('Expected int or dict, got %s of %s' % \
+              (value, type(value)))
+        objects = model.objects.complex_filter(query)
+        if len(objects) == 1:
+            return objects[0]
+        else:
+            raise base.DeserializationError('%i instance(s) of %s for %s' % \
+              (len(objects), model, value))
+
+def pk_or_none(obj):
+    if obj:
+        return obj._get_pk_val()
+    else:
+        return None
+        
+
+class Serializer(base.Serializer):
+
+    def start_serialization(self):
+        self.last_model = None
+        self.output = csv.writer(self.stream, dialect=dialect)
+
+    def start_object(self, obj):
+        if not hasattr(obj, "_meta"):
+            raise base.SerializationError("Non-model object (%s) encountered during serialization" % type(obj))
+        if self.last_model != obj._meta:
+            self.last_model = obj._meta
+            header = []
+            for field in obj._meta.fields:
+                header.append(field.name)
+            for field in obj._meta.many_to_many:
+                header.append(field.name)
+            header[0] = '%s%s:%s' % (header_start, obj._meta, header[0])
+            self.writerow(header)
+        self.row = [self.tostring(obj._get_pk_val())]
+
+    def end_object(self, obj):
+        self.writerow(self.row)
+
+    def handle_field(self, obj, field):
+        if getattr(obj, field.name) is not None:
+            value = self.get_string_value(obj, field)
+            self.row.append(self.tostring(value))
+        else:
+            self.row.append('')
+
+    def handle_fk_field(self, obj, field):
+        related = self.represent(getattr(obj, field.name))
+        self.row.append(self.tostring(related))
+
+    def handle_m2m_field(self, obj, field):
+        """A tuple of m2m representations as dicts or ids"""
+        related = [self.represent(related) for related in getattr(obj, field.name).iterator()]
+        self.row.append(self.tostring(tuple(related)))
+
+    def tostring(self, value):
+        s = smart_unicode(value)
+        return s or ''
+
+    def represent(self, related):
+        """Represent a model object either as its pk as int, or as a
+        dict of unique key-values pairs recursevily."""
+        if related is None:
+            return None
+        elif nice_foreign_keys:
+            # Find a compound key
+            if related._meta.unique_together:
+                dict = {}
+                for field_name in related._meta.unique_together[0]:
+                    field = related._meta.get_field(field_name)
+                    value = getattr(related, field.name)
+                    if isinstance(value, models.Model):
+                        dict[field_name] = self.represent(value)
+                    else:
+                        dict[field_name] = value
+                return dict
+            # Find a unique key
+            else:
+                for field in related._meta.fields:
+                    if field.unique:
+                        return {field.name: getattr(related, field.name)}
+        return related._get_pk_val()
+
+    def writerow(self, row):
+        self.output.writerow(row)
+
+
+class Deserializer(base.Deserializer):
+
+    def __iter__(self):
+        for values in csv.reader(self.stream, dialect=dialect):
+            if values:
+                if values[0].startswith(header_start):
+                    # Model
+                    model, first_field = values[0].split(':', 2)
+                    model = model[len(header_start):]
+                    try:
+                        self.model = models.get_model(*model.split("."))
+                    except TypeError:
+                        raise base.DeserializationError("No model %s in db" % model)
+                    # Field names
+                    self.field_names = [first_field] + values[1:]
+                else:
+                    # An object
+                    meta = self.model._meta
+                    data = {meta.pk.attname: meta.pk.to_python(values[0])}
+                    m2m_data = {}
+                    for i in range(1, len(values)):
+                        name = self.field_names[i]
+                        value = values[i]
+                        if value == '':
+                            value = None
+                        field = meta.get_field(name)
+                        if field.rel and isinstance(field.rel, models.ManyToManyRel):
+                            if value:
+                                m2m_data[field.name] = \
+                                  [pk_or_none(resolve_related(field, v)) for v in eval(value)]
+                            else:
+                                m2m_data[field.name] = []
+                        elif field.rel and isinstance(field.rel, models.ManyToOneRel):
+                            if value:
+                                data[field.attname] = pk_or_none(resolve_related(field, eval(value)))
+                        else:
+                            if value == '""':
+                                value = ''
+                            value = self.to_python(field, value)
+                            data[field.attname] = value
+                    yield base.DeserializedObject(self.model(**data), m2m_data)
+        
+    def to_python(self, field, value):
+        """
+        The to_python method of some fields are not implemented, so this
+        is a workaround for them.
+        """
+        if value is None and not field.null:
+            value = ''
+        value = field.to_python(value)
+        if value is None:
+            return None
+        # XXX isinstance(field, IntegerField) doesn't seem to work
+        elif field.__class__.__name__.endswith('IntegerField'):
+            return int(value)
+        elif field.__class__.__name__.endswith('FloatField'):
+            return float(value)
+        else:
+            return value
+