def __exit__(self, type, value, traceback):
"""If exit appears without a commit, undo all changes"""
if not self.committed:
self.file.close()
os.remove(get_temp_path(self.file_path))
self.file = None
def commit(self):
"""Commits the change. Write operations after this lead to an error"""
if self.committed:
raise WriteAfterCommitException
self.committed = True
self.file.close()
self.file = None
shutil.move(get_temp_path(self.file_path), self.file_path)
def __enter__(self):
self.file = open(get_temp_path(self.file_path), "a",
**self.kwargs.get("file_params", {}))
if self.kwargs.get("dict", False):
self.writer = csv.DictWriter(self.file,
self.kwargs.get("fields", {}),
**self.kwargs.get("csv_params", {}))
else:
self.writer = csv.writer(self.file,
**self.kwargs.get("csv_params", {}))
return self
def commit(self):
"""Commits the change. Write operations after this lead to an error"""
if self.committed:
raise WriteAfterCommitException
temp_path = get_temp_path(self.file_path)
with open(temp_path, "w", **self.kwargs.get("file_params",
{})) as file:
writer = csv.DictWriter(file, self.kwargs.get("fields", {}),
**self.kwargs.get("csv_params", {}))
writer.writeheader()
first_new = 0
if os.path.exists(self.file_path):
with open(self.file_path) as f:
reader = csv.DictReader(
f, **self.kwargs.get("csv_params", {}))
idx = 0
for row in reader:
column = self.rows.get(idx, row)
if column is not None:
writer.writerow(column)
idx += 1
first_new = idx
# Write all new columns
for _, column in sorted(
[(idx, data)
for idx, data in self.rows.items() if idx >= first_new],
key=lambda x: x[0]):
writer.writerow(column)
shutil.move(temp_path, self.file_path)
def __enter__(self):
self.file = open(get_temp_path(self.file_path), self.mode,
**self.kwargs.get("file_params", {}))
return self