allow "xmlfile(f, close=True)" to automatically close the file object…

… after use

CHANGES.txt

@@ -9,6 +9,9 @@ Latest changes
 Features added
 --------------
 
+* ``xmlfile()`` accepts a new argument ``close=True`` to close file(-like)
+  objects after writing to them.
+
 Bugs fixed
 ----------
 

doc/api.txt

@@ -490,6 +490,14 @@ example::
   >>> print(f.getvalue().decode('utf-8'))
   <abc>text</abc>
 
+``xmlfile()`` accepts a file path as first argument, or a file(-like)
+object, as in the example above.  In the first case, it takes care to
+open and close the file itself, whereas file(-like) objects are not
+closed by default.  This is left to the code that opened them.  Since
+lxml 3.4, however, you can pass the argument ``close=True`` to make
+lxml call the object's ``.close()`` method when exiting the context
+manager.
+
 To insert pre-constructed Elements and subtrees, just pass them
 into ``write()``::
 

src/lxml/serializer.pxi

@@ -375,11 +375,13 @@ cdef class _FilelikeWriter:
     cdef object _close_filelike
     cdef _ExceptionContext _exc_context
     cdef _ErrorLog error_log
-    def __cinit__(self, filelike, exc_context=None, compression=None):
+    def __cinit__(self, filelike, exc_context=None, compression=None, close=False):
         if compression is not None and compression > 0:
             filelike = GzipFile(
                 fileobj=filelike, mode='wb', compresslevel=compression)
             self._close_filelike = filelike.close
+        elif close:
+            self._close_filelike = filelike.close
         self._filelike = filelike
         if exc_context is None:
             self._exc_context = _ExceptionContext()
@@ -474,7 +476,7 @@ cdef _tofilelike(f, _Element element, encoding, doctype, method,
         doctype = _utf8(doctype)
         c_doctype = _xcstr(doctype)
 
-    writer = _create_output_buffer(f, c_enc, compression, &c_buffer)
+    writer = _create_output_buffer(f, c_enc, compression, &c_buffer, close=False)
     if writer is None:
         state = python.PyEval_SaveThread()
 
@@ -496,7 +498,7 @@ cdef _tofilelike(f, _Element element, encoding, doctype, method,
         _raiseSerialisationError(error_result)
 
 cdef _create_output_buffer(f, const_char* c_enc, int compression,
-                           tree.xmlOutputBuffer** c_buffer_ret):
+                           tree.xmlOutputBuffer** c_buffer_ret, bint close):
     cdef tree.xmlOutputBuffer* c_buffer
     cdef _FilelikeWriter writer
     enchandler = tree.xmlFindCharEncodingHandler(c_enc)
@@ -512,7 +514,7 @@ cdef _create_output_buffer(f, const_char* c_enc, int compression,
                 return python.PyErr_SetFromErrno(IOError) # raises IOError
             writer = None
         elif hasattr(f, 'write'):
-            writer = _FilelikeWriter(f, compression=compression)
+            writer = _FilelikeWriter(f, compression=compression, close=close)
             c_buffer = writer._createOutputBuffer(enchandler)
         else:
             raise TypeError(
@@ -600,7 +602,7 @@ cdef _tofilelikeC14N(f, _Element element, bint exclusive, bint with_comments,
 # incremental serialisation
 
 cdef class xmlfile:
-    """xmlfile(self, output_file, encoding=None, compression=None)
+    """xmlfile(self, output_file, encoding=None, compression=None, close=False)
 
     A simple mechanism for incremental XML serialisation.
 
@@ -619,28 +621,41 @@ cdef class xmlfile:
                   for element in generate_some_elements():
                       # serialise generated elements into the XML file
                       xf.write(element)
+
+    If 'output_file' is a file(-like) object, passing ``close=True`` will
+    close it when exiting the context manager.  By default, it is left
+    to the owner to do that.  When a file path is used, lxml will take care
+    of opening and closing the file itself.  Also, when a compression level
+    is set, lxml will deliberately close the file to make sure all data gets
+    compressed and written.
     """
     cdef object output_file
     cdef bytes encoding
-    cdef int compresslevel
     cdef _IncrementalFileWriter writer
+    cdef int compresslevel
+    cdef bint close
 
-    def __init__(self, output_file not None, encoding=None, compression=None):
+    def __init__(self, output_file not None, encoding=None, compression=None,
+                 close=False):
         self.output_file = output_file
         self.encoding = _utf8orNone(encoding)
         self.compresslevel = compression or 0
+        self.close = close
 
     def __enter__(self):
         assert self.output_file is not None
         self.writer = _IncrementalFileWriter(
-            self.output_file, self.encoding, self.compresslevel)
+            self.output_file, self.encoding, self.compresslevel, self.close)
         return self.writer
 
     def __exit__(self, exc_type, exc_val, exc_tb):
         if self.writer is not None:
             old_writer, self.writer = self.writer, None
             raise_on_error = exc_type is None
             old_writer._close(raise_on_error)
+            if self.close:
+                self.output_file = None
+
 
 cdef enum _IncrementalFileWriterStatus:
     WRITER_STARTING = 0
@@ -649,6 +664,7 @@ cdef enum _IncrementalFileWriterStatus:
     WRITER_IN_ELEMENT = 3
     WRITER_FINISHED = 4
 
+
 @cython.final
 @cython.internal
 cdef class _IncrementalFileWriter:
@@ -659,15 +675,15 @@ cdef class _IncrementalFileWriter:
     cdef list _element_stack
     cdef int _status
 
-    def __cinit__(self, outfile, bytes encoding, int compresslevel):
+    def __cinit__(self, outfile, bytes encoding, int compresslevel, bint close):
         self._status = WRITER_STARTING
         self._element_stack = []
         if encoding is None:
             encoding = b'ASCII'
         self._encoding = encoding
         self._c_encoding = _cstr(encoding) if encoding is not None else NULL
         self._target = _create_output_buffer(
-            outfile, self._c_encoding, compresslevel, &self._c_out)
+            outfile, self._c_encoding, compresslevel, &self._c_out, close)
 
     def __dealloc__(self):
         if self._c_out is not NULL:

src/lxml/tests/test_incremental_xmlfile.py

@@ -203,7 +203,8 @@ def _parse_file(self):
         return etree.parse(self._file)
 
     def tearDown(self):
-        self._file.close()
+        if self._file is not None:
+            self._file.close()
 
     def assertXml(self, expected, encoding='utf8'):
         self.assertEqual(self._read_file().decode(encoding), expected)
@@ -213,16 +214,50 @@ class BytesIOXmlFileTestCase(_XmlFileTestCaseBase):
     def setUp(self):
         self._file = BytesIO()
 
+    def test_filelike_close(self):
+        with etree.xmlfile(self._file, close=True) as xf:
+            with xf.element('test'):
+                pass
+        self.assertRaises(ValueError, self._file.getvalue)
+
+
 class TempXmlFileTestCase(_XmlFileTestCaseBase):
     def setUp(self):
-        self._file = tempfile.NamedTemporaryFile()
+        self._file = tempfile.TemporaryFile()
+
+
+class TempPathXmlFileTestCase(_XmlFileTestCaseBase):
+    def setUp(self):
+        self._tmpfile = tempfile.NamedTemporaryFile()
+        self._file = self._tmpfile.name
+
+    def tearDown(self):
+        try:
+            self._tmpfile.close()
+        finally:
+            if os.path.exists(self._tmpfile.name):
+                os.unlink(self._tmpfile.name)
+
+    def _read_file(self):
+        self._tmpfile.seek(0)
+        return self._tmpfile.read()
+
+    def _parse_file(self):
+        self._tmpfile.seek(0)
+        return etree.parse(self._tmpfile)
+
 
 class SimpleFileLikeXmlFileTestCase(_XmlFileTestCaseBase):
     class SimpleFileLike(object):
         def __init__(self, target):
             self._target = target
             self.write = target.write
-            self.close = target.close
+            self.closed = False
+
+        def close(self):
+            assert not self.closed
+            self.closed = True
+            self._target.close()
 
     def setUp(self):
         self._target = BytesIO()
@@ -235,11 +270,25 @@ def _parse_file(self):
         self._target.seek(0)
         return etree.parse(self._target)
 
+    def test_filelike_not_closing(self):
+        with etree.xmlfile(self._file) as xf:
+            with xf.element('test'):
+                pass
+        self.assertFalse(self._file.closed)
+
+    def test_filelike_close(self):
+        with etree.xmlfile(self._file, close=True) as xf:
+            with xf.element('test'):
+                pass
+        self.assertTrue(self._file.closed)
+        self._file = None  # prevent closing in tearDown()
+
 
 def test_suite():
     suite = unittest.TestSuite()
     suite.addTests([unittest.makeSuite(BytesIOXmlFileTestCase),
                     unittest.makeSuite(TempXmlFileTestCase),
+                    unittest.makeSuite(TempPathXmlFileTestCase),
                     unittest.makeSuite(SimpleFileLikeXmlFileTestCase),
                     ])
     return suite