Add method keyword to _IncrementalWriter.element()

2 files changed, 22 insertions, 4 deletions

diff --git a/src/lxml/serializer.pxi b/src/lxml/serializer.pxi
index 489ed2de..0eb95a76 100644
--- a/src/lxml/serializer.pxi
+++ b/src/lxml/serializer.pxi
@@ -812,7 +812,7 @@ cdef class _IncrementalFileWriter:
             tree.xmlOutputBufferFlush(self._c_out)
         self._handle_error(self._c_out.error)
 
-    def element(self, tag, attrib=None, nsmap=None, **_extra):
+    def element(self, tag, attrib=None, nsmap=None, method=None, **_extra):
         """element(self, tag, attrib=None, nsmap=None, **_extra)
 
         Returns a context manager that writes an opening and closing tag.
@@ -836,7 +836,10 @@ cdef class _IncrementalFileWriter:
                     _prefixValidOrRaise(prefix)
                 reversed_nsmap[_utf8(ns)] = prefix
         ns, name = _getNsTag(tag)
-        return _FileWriterElement(self, (ns, name, attributes, reversed_nsmap))
+
+        c_method = self._method if method is None else _findOutputMethod(method)
+
+        return _FileWriterElement(self, (ns, name, attributes, reversed_nsmap), c_method)
 
     cdef _write_qname(self, bytes name, bytes prefix):
         if prefix:  # empty bytes for no prefix (not None to allow sorting)
@@ -1015,14 +1018,20 @@ cdef class _IncrementalFileWriter:
 @cython.freelist(8)
 cdef class _FileWriterElement:
     cdef object _element
+    cdef int _new_method
+    cdef int _old_method
     cdef _IncrementalFileWriter _writer
 
-    def __cinit__(self, _IncrementalFileWriter writer not None, element_config):
+    def __cinit__(self, _IncrementalFileWriter writer not None, element_config, int method):
         self._writer = writer
         self._element = element_config
+        self._new_method = method
+        self._old_method = writer._method
 
     def __enter__(self):
+        self._writer._method = self._new_method
         self._writer._write_start_element(self._element)
 
     def __exit__(self, exc_type, exc_val, exc_tb):
         self._writer._write_end_element(self._element)
+        self._writer._method = self._old_method
diff --git a/src/lxml/tests/test_incremental_xmlfile.py b/src/lxml/tests/test_incremental_xmlfile.py
index 414edd12..6670471c 100644
--- a/src/lxml/tests/test_incremental_xmlfile.py
+++ b/src/lxml/tests/test_incremental_xmlfile.py
@@ -381,7 +381,9 @@ class HtmlFileTestCase(_XmlFileTestCaseBase):
             self._file = BytesIO()
 
     def test_xml_mode_write_inside_html(self):
-        elt = etree.Element("foo", attrib={'selected': 'bar'})
+        tag = 'foo'
+        attrib = {'selected': 'bar'}
+        elt = etree.Element(tag, attrib=attrib)
 
         with etree.htmlfile(self._file) as xf:
             with xf.element("root"):
@@ -393,11 +395,18 @@ class HtmlFileTestCase(_XmlFileTestCaseBase):
                 elt.text = ""
                 xf.write(elt, method='xml')  # 3
 
+                with xf.element(tag, attrib=attrib, method='xml'):
+                    pass # 4
+
+                xf.write(elt)  # 5
+
         self.assertXml(
             '<root>'
                 '<foo selected></foo>'  # 1
                 '<foo selected="bar"/>'  # 2
                 '<foo selected="bar"></foo>'  # 3
+                '<foo selected="bar"></foo>'  # 4
+                '<foo selected></foo>'  # 5
             '</root>')
         self._file = BytesIO()