Merge pull request #12912 from [BEAM-10938] Adds support for writing a footer to Python WriteToText

* [BEAM-10938] Adds support for writing a footer to Python WriteToText

* Fixing formatting

Author: milancermak

sdks/python/apache_beam/io/textio.py

@@ -350,7 +350,8 @@ def __init__(self,
                shard_name_template=None,
                coder=coders.ToBytesCoder(),  # type: coders.Coder
                compression_type=CompressionTypes.AUTO,
-               header=None):
+               header=None,
+               footer=None):
     """Initialize a _TextSink.
 
     Args:
@@ -382,6 +383,8 @@ def __init__(self,
         compression.
       header: String to write at beginning of file as a header. If not None and
         append_trailing_newlines is set, '\n' will be added.
+      footer: String to write at the end of file as a footer. If not None and
+        append_trailing_newlines is set, '\n' will be added.
 
     Returns:
       A _TextSink object usable for writing.
@@ -396,6 +399,7 @@ def __init__(self,
         compression_type=compression_type)
     self._append_trailing_newlines = append_trailing_newlines
     self._header = header
+    self._footer = footer
 
   def open(self, temp_path):
     file_handle = super(_TextSink, self).open(temp_path)
@@ -405,6 +409,13 @@ def open(self, temp_path):
         file_handle.write(b'\n')
     return file_handle
 
+  def close(self, file_handle):
+    if self._footer is not None:
+      file_handle.write(coders.ToBytesCoder().encode(self._footer))
+      if self._append_trailing_newlines:
+        file_handle.write(b'\n')
+    super(_TextSink, self).close(file_handle)
+
   def display_data(self):
     dd_parent = super(_TextSink, self).display_data()
     dd_parent['append_newline'] = DisplayDataItem(
@@ -588,7 +599,8 @@ def __init__(
       shard_name_template=None,  # type: Optional[str]
       coder=coders.ToBytesCoder(),  # type: coders.Coder
       compression_type=CompressionTypes.AUTO,
-      header=None):
+      header=None,
+      footer=None):
     r"""Initialize a :class:`WriteToText` transform.
 
     Args:
@@ -624,6 +636,9 @@ def __init__(
       header (str): String to write at beginning of file as a header.
         If not :data:`None` and **append_trailing_newlines** is set, ``\n`` will
         be added.
+      footer (str): String to write at the end of file as a footer.
+        If not :data:`None` and **append_trailing_newlines** is set, ``\n`` will
+        be added.
     """
 
     self._sink = _TextSink(
@@ -634,7 +649,8 @@ def __init__(
         shard_name_template,
         coder,
         compression_type,
-        header)
+        header,
+        footer)
 
   def expand(self, pcoll):
     return pcoll | Write(self._sink)

sdks/python/apache_beam/io/textio_test.py

@@ -1125,6 +1125,14 @@ def test_write_text_file_with_header(self):
     with open(self.path, 'rb') as f:
       self.assertEqual(f.read().splitlines(), header.splitlines() + self.lines)
 
+  def test_write_text_file_with_footer(self):
+    footer = b'footer1\nfooter2'
+    sink = TextSink(self.path, footer=footer)
+    self._write_lines(sink, self.lines)
+
+    with open(self.path, 'rb') as f:
+      self.assertEqual(f.read().splitlines(), self.lines + footer.splitlines())
+
   def test_write_text_file_empty_with_header(self):
     header = b'header1\nheader2'
     sink = TextSink(self.path, header=header)
@@ -1203,6 +1211,22 @@ def test_write_pipeline_header(self):
     self.assertEqual(read_result[0], header_text.encode('utf-8'))
     self.assertEqual(sorted(read_result[1:]), sorted(self.lines))
 
+  def test_write_pipeline_footer(self):
+    with TestPipeline() as pipeline:
+      footer_text = 'footer'
+      pcoll = pipeline | beam.core.Create(self.lines)
+      pcoll | 'Write' >> WriteToText(   # pylint: disable=expression-not-assigned
+        self.path,
+        footer=footer_text)
+
+    read_result = []
+    for file_name in glob.glob(self.path + '*'):
+      with open(file_name, 'rb') as f:
+        read_result.extend(f.read().splitlines())
+
+    self.assertEqual(sorted(read_result[:-1]), sorted(self.lines))
+    self.assertEqual(read_result[-1], footer_text.encode('utf-8'))
+
 
 if __name__ == '__main__':
   logging.getLogger().setLevel(logging.INFO)