allow writing unicode to GridFS if an encoding attribute has been specified PYTHON-100

Author: Mike Dirolf

gridfs/__init__.py

@@ -87,8 +87,11 @@ def put(self, data, **kwargs):
  >>> f.close()
 
  `data` can be either an instance of :class:`str` or a
- file-like object providing a :meth:`read` method. Any keyword
- arguments will be passed through to the created file - see
+ file-like object providing a :meth:`read` method. If an
+ `encoding` keyword argument is passed, `data` can also be a
+ :class:`unicode` instance, which will be encoded as `encoding`
+ before being written. Any keyword arguments will be passed
+ through to the created file - see
  :meth:`~gridfs.grid_file.GridIn` for possible
  arguments. Returns the ``"_id"`` of the created file.
 
@@ -100,6 +103,10 @@ def put(self, data, **kwargs):
  - `data`: data to be written as a file.
  - `**kwargs` (optional): keyword arguments for file creation
 
+ .. versionadded:: 1.8.1+
+ The ability to write :class:`unicode`, if an `encoding` has
+ been specified as a keyword argument.
+
  .. versionadded:: 1.6
  """
  grid_file = GridIn(self.__collection, **kwargs)

gridfs/grid_file.py

@@ -109,6 +109,10 @@ def __init__(self, root_collection, **kwargs):
  - ``"chunkSize"`` or ``"chunk_size"``: size of each of the
  chunks, in bytes (default: 256 kb)
 
+ - ``"encoding"``: encoding used for this file - any
+ :class:`unicode` that is written to the file will be
+ converted to a :class:`str` with this encoding
+
  :Parameters:
  - `root_collection`: root collection to write to
  - `**kwargs` (optional): file level options (see above)
@@ -224,17 +228,26 @@ def write(self, data):
  """Write data to the file. There is no return value.
 
  `data` can be either a string of bytes or a file-like object
- (implementing :meth:`read`).
+ (implementing :meth:`read`). If the file has an
+ :attr:`encoding` attribute, `data` can also be a
+ :class:`unicode` instance, which will be encoded as
+ :attr:`encoding` before being written.
 
  Due to buffering, the data may not actually be written to the
  database until the :meth:`close` method is called. Raises
  :class:`ValueError` if this file is already closed. Raises
  :class:`TypeError` if `data` is not an instance of
- :class:`str` or a file-like object.
+ :class:`str`, a file-like object, or an instance of
+ :class:`unicode` (only allowed if the file has an
+ :attr:`encoding` attribute).
 
  :Parameters:
  - `data`: string of bytes or file-like object to be written
  to the file
+
+ .. versionadded:: 1.8.1+
+ The ability to write :class:`unicode`, if the file has an
+ :attr:`encoding` attribute.
  """
  if self._closed:
  raise ValueError("cannot write to a closed file")
@@ -252,9 +265,16 @@ def write(self, data):
  self._buffer.write(to_write)
  # string
  except AttributeError:
- if not isinstance(data, str):
+ if not isinstance(data, basestring):
  raise TypeError("can only write strings or file-like objects")
 
+ if isinstance(data, unicode):
+ try:
+ data = data.encode(self.encoding)
+ except AttributeError:
+ raise TypeError("must specify an encoding for file in "
+ "order to write unicode")
+
  while data:
  space = self.chunk_size - self._buffer.tell()
 

test/test_grid_file.py

@@ -1,3 +1,5 @@
+# -*- coding: utf-8 -*-
+#
 # Copyright 2009-2010 10gen, Inc.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -422,5 +424,24 @@ def test_read_chunks_unaligned_buffer_size(self):
 
  self.assertEqual(in_data, out_data)
 
+ def test_write_unicode(self):
+ f = GridIn(self.db.fs)
+ self.assertRaises(TypeError, f.write, u"foo")
+
+ f = GridIn(self.db.fs, encoding="utf-8")
+ f.write(u"foo")
+ f.close()
+
+ g = GridOut(self.db.fs, f._id)
+ self.assertEqual("foo", g.read())
+
+ f = GridIn(self.db.fs, encoding="iso-8859-1")
+ f.write(u"aé")
+ f.close()
+
+ g = GridOut(self.db.fs, f._id)
+ self.assertEqual(u"aé".encode("iso-8859-1"), g.read())
+
+
 if __name__ == "__main__":
  unittest.main()

test/test_gridfs.py

@@ -1,3 +1,5 @@
+# -*- coding: utf-8 -*-
+#
 # Copyright 2009-2010 10gen, Inc.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -231,6 +233,17 @@ def test_exists(self):
  self.assertFalse(self.fs.exists(foo={"$gt": 12}))
  self.assertFalse(self.fs.exists({"foo": {"$gt": 12}}))
 
+ def test_put_unicode(self):
+ self.assertRaises(TypeError, self.fs.put, u"hello")
+
+ oid = self.fs.put(u"hello", encoding="utf-8")
+ self.assertEqual("hello", self.fs.get(oid).read())
+ self.assertEqual("utf-8", self.fs.get(oid).encoding)
+
+ oid = self.fs.put(u"aé", encoding="iso-8859-1")
+ self.assertEqual(u"aé".encode("iso-8859-1"), self.fs.get(oid).read())
+ self.assertEqual("iso-8859-1", self.fs.get(oid).encoding)
+
 
 if __name__ == "__main__":
  unittest.main()