2020-10-18 16:25:49 -04:00
|
|
|
# Copyright 2019 (c) Anna Schumaker.
|
|
|
|
import os
|
2020-10-18 17:15:24 -04:00
|
|
|
import pathlib
|
2020-10-18 16:25:49 -04:00
|
|
|
import pickle
|
|
|
|
import xdg.BaseDirectory
|
|
|
|
|
|
|
|
__resource = "emmental"
|
|
|
|
if os.environ.get("EMMENTAL_TESTING"):
|
|
|
|
__resource = "emmental-testing"
|
|
|
|
|
|
|
|
READ = 'rb'
|
|
|
|
WRITE = 'wb'
|
|
|
|
|
2020-10-18 17:15:24 -04:00
|
|
|
emmental_data = pathlib.Path(xdg.BaseDirectory.save_data_path(__resource))
|
2020-10-18 16:25:49 -04:00
|
|
|
|
|
|
|
class DataFile:
|
|
|
|
def __init__(self, path, mode):
|
2020-10-18 17:15:24 -04:00
|
|
|
self.path = emmental_data / path
|
|
|
|
self.temp = emmental_data / f".{path}.tmp"
|
2020-10-18 16:25:49 -04:00
|
|
|
self.mode = mode
|
|
|
|
self.file = None
|
|
|
|
|
|
|
|
def __enter__(self):
|
|
|
|
if self.mode == WRITE:
|
2020-10-18 17:15:24 -04:00
|
|
|
self.file = self.temp.open(self.mode)
|
|
|
|
elif self.mode == READ and self.path.exists():
|
|
|
|
self.file = self.path.open(self.mode)
|
2020-10-18 16:25:49 -04:00
|
|
|
return self
|
|
|
|
|
|
|
|
def __exit__(self, exp_type, exp_value, traceback):
|
|
|
|
if self.file:
|
2020-10-18 17:15:24 -04:00
|
|
|
self.file.flush()
|
2020-10-18 16:25:49 -04:00
|
|
|
self.file.close()
|
|
|
|
if self.mode == WRITE:
|
2020-10-18 17:15:24 -04:00
|
|
|
self.temp.replace(self.path)
|
|
|
|
self.file = None
|
2021-06-07 10:02:35 -04:00
|
|
|
return exp_type == None
|
2020-10-18 16:25:49 -04:00
|
|
|
|
|
|
|
def exists(self):
|
2020-10-18 17:15:24 -04:00
|
|
|
return self.path.exists()
|
2020-10-18 16:25:49 -04:00
|
|
|
|
|
|
|
def pickle(self, obj):
|
|
|
|
if self.file:
|
|
|
|
pickle.dump(obj, self.file, pickle.HIGHEST_PROTOCOL)
|
|
|
|
|
|
|
|
def remove(self):
|
2020-10-18 17:15:24 -04:00
|
|
|
if self.file == None:
|
|
|
|
self.path.unlink(missing_ok = True)
|
2020-10-18 16:25:49 -04:00
|
|
|
|
|
|
|
def unpickle(self):
|
|
|
|
if self.file:
|
|
|
|
return pickle.load(self.file)
|