# Copyright 2020 The FedLearner Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# coding: utf-8

import csv

import tensorflow_io # pylint: disable=unused-import
from tensorflow.compat.v1 import gfile

class CsvDictWriter(object):
    def __init__(self, fpath):
        self._write_raw_num = 0
        self._fpath = fpath
        self._file_hanlde = gfile.Open(fpath, 'w+')
        self._csv_writer = None

    def write(self, raw):
        assert isinstance(raw, dict)
        if len(raw) == 0:
            return
        if self._csv_writer is None:
            self._csv_writer = csv.DictWriter(
                    self._file_hanlde,
                    fieldnames=raw.keys()
                )
            self._csv_writer.writeheader()
        self._csv_writer.writerow(raw)
        self._write_raw_num += 1

    def close(self):
        if self._file_hanlde is not None:
            self._file_hanlde.close()
            self._file_hanlde, self._csv_writer = None, None

    def write_raw_num(self):
        return self._write_raw_num

    def __del__(self):
        if self._file_hanlde is not None:
            self.close()

@@ -629,16 +629,26 @@
 port csv
+%0Aimport io
 %0A%0Aimport
@@ -906,16 +906,60 @@
 , 'w+')%0A
+        self._buffer_handle = io.StringIO()%0A
         
@@ -1201,26 +1201,28 @@
   self._
-file
+buffer
 _han
-l
 d
+l
 e,%0A     
@@ -1390,16 +1390,50 @@
 num += 1
+%0A        self._flush_buffer(False)
 %0A%0A    de
@@ -1439,32 +1439,192 @@
 ef close(self):%0A
+        if self._buffer_handle is not None:%0A            self._flush_buffer(True)%0A            self._buffer_handle.close()%0A            self._buffer_handle = None%0A
         if self.
@@ -1716,17 +1716,31 @@
 e_hanlde
-,
+ = None%0A       
  self._c
@@ -1759,14 +1759,8 @@
 None
-, None
 %0A%0A  
@@ -1822,16 +1822,326 @@
 aw_num%0A%0A
+    def _flush_buffer(self, force=False):%0A        if self._buffer_handle.tell() %3E (2 %3C%3C 20) or %5C%0A                (self._buffer_handle.tell() %3E 0 and force):%0A            self._file_hanlde.write(self._buffer_handle.getvalue())%0A            self._buffer_handle.truncate(0)%0A            self._buffer_handle.seek(0)%0A%0A
     def 

make the csv writer support batch write (#257)