emmental/lib/data.py

53 lines
1.3 KiB
Python
Raw Normal View History

# 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)