53 lines
1.3 KiB
Python
53 lines
1.3 KiB
Python
# Copyright 2019 (c) Anna Schumaker.
|
|
import os
|
|
import pickle
|
|
import xdg.BaseDirectory
|
|
|
|
__resource = "emmental"
|
|
if os.environ.get("EMMENTAL_TESTING"):
|
|
__resource = "emmental-testing"
|
|
|
|
READ = 'rb'
|
|
WRITE = 'wb'
|
|
|
|
emmental_data = xdg.BaseDirectory.save_data_path(__resource)
|
|
|
|
class DataFile:
|
|
def __init__(self, path, mode):
|
|
self.path = os.path.join(emmental_data, path)
|
|
self.temp = os.path.join(emmental_data, f".{path}.tmp")
|
|
self.mode = mode
|
|
self.file = None
|
|
|
|
def __enter__(self):
|
|
if self.mode == WRITE:
|
|
self.file = open(self.temp, self.mode)
|
|
elif self.mode == READ and self.exists():
|
|
self.file = open(self.path, self.mode)
|
|
return self
|
|
|
|
def __exit__(self, exp_type, exp_value, traceback):
|
|
if self.file:
|
|
self.file.close()
|
|
self.file = None
|
|
if self.mode == WRITE:
|
|
os.rename(self.temp, self.path)
|
|
return True
|
|
|
|
def exists(self):
|
|
return os.path.exists(self.path)
|
|
|
|
def pickle(self, obj):
|
|
if self.file:
|
|
pickle.dump(obj, self.file, pickle.HIGHEST_PROTOCOL)
|
|
|
|
def remove(self):
|
|
if not self.file and self.exists():
|
|
os.remove(self.path)
|
|
return True
|
|
return False
|
|
|
|
def unpickle(self):
|
|
if self.file:
|
|
return pickle.load(self.file)
|