138 lines
3.7 KiB
Python
138 lines
3.7 KiB
Python
from nullptr.models.base import Base
|
|
from nullptr.models.waypoint import Waypoint
|
|
from nullptr.models.sector import Sector
|
|
from nullptr.models.system import System
|
|
from nullptr.models.agent import Agent
|
|
from nullptr.models.marketplace import Marketplace
|
|
from nullptr.models.system_member import SystemMember
|
|
from nullptr.models.jumpgate import Jumpgate
|
|
from nullptr.models.ship import Ship
|
|
from nullptr.models.contract import Contract
|
|
from nullptr.models.survey import Survey
|
|
from os.path import isfile, dirname, isdir
|
|
import os
|
|
from os.path import basename
|
|
import json
|
|
from .util import *
|
|
from time import time
|
|
|
|
class Store:
|
|
def __init__(self, data_dir):
|
|
self.init_models()
|
|
self.data_dir = data_dir
|
|
self.data = {m: {} for m in self.models}
|
|
self.system_members = {}
|
|
self.dirty_objects = set()
|
|
|
|
def init_models(self):
|
|
self.models = all_subclasses(Base)
|
|
self.extensions = {c.ext(): c for c in self.models}
|
|
self.model_names = {c.__name__: c for c in self.models}
|
|
|
|
def dirty(self, obj):
|
|
self.dirty_objects.add(obj)
|
|
|
|
def path(self, obj):
|
|
return os.path.join(self.data_dir, obj.path())
|
|
|
|
def load_file(self, path):
|
|
if not isfile(path):
|
|
return None
|
|
fn = basename(path)
|
|
ext = fn.split('.')[-1]
|
|
symbol = fn.split('.')[0]
|
|
if ext not in self.extensions:
|
|
return None
|
|
with open(path) as f:
|
|
data = json.load(f)
|
|
|
|
typ = self.extensions[ext]
|
|
obj = self.create(typ, symbol)
|
|
obj.load(data)
|
|
obj.store = self
|
|
return obj
|
|
|
|
def load(self):
|
|
cnt = 0
|
|
start_time = time()
|
|
for fil in list_files(self.data_dir, True):
|
|
self.load_file(fil)
|
|
cnt += 1
|
|
dur = time() - start_time
|
|
print(f'loaded {cnt} objects in {dur:.2f} seconds')
|
|
|
|
def store(self, obj):
|
|
path = self.path(obj)
|
|
path_dir = dirname(path)
|
|
data = obj.dict()
|
|
if not isdir(path_dir):
|
|
os.makedirs(path_dir, exist_ok=True)
|
|
with open(path, 'w') as f:
|
|
json.dump(data, f, indent=2)
|
|
|
|
def create(self, typ, symbol):
|
|
obj = typ(symbol, self)
|
|
self.data[typ][symbol] = obj
|
|
if issubclass(typ, SystemMember):
|
|
system_str = obj.system()
|
|
|
|
if system_str not in self.system_members:
|
|
self.system_members[system_str] = set()
|
|
self.system_members[system_str].add(obj)
|
|
return obj
|
|
|
|
def get(self, typ, symbol, create=False):
|
|
if type(typ) == str and typ in self.model_names:
|
|
typ = self.model_names[typ]
|
|
symbol = symbol.upper()
|
|
if typ not in self.data:
|
|
return None
|
|
if symbol not in self.data[typ]:
|
|
if create:
|
|
return self.create(typ, symbol)
|
|
else:
|
|
return None
|
|
return self.data[typ][symbol]
|
|
|
|
def update(self, typ, data, symbol=None):
|
|
if type(typ) == str and typ in self.model_names:
|
|
typ = self.model_names[typ]
|
|
if symbol is None:
|
|
symbol = mg(data, typ.identifier)
|
|
obj = self.get(typ, symbol, True)
|
|
obj.update(data)
|
|
return obj
|
|
|
|
def update_list(self, typ, lst):
|
|
return [self.update(typ, d) for d in lst]
|
|
|
|
def all(self, typ):
|
|
if type(typ) == str and typ in self.model_names:
|
|
typ = self.model_names[typ]
|
|
|
|
for m in self.data[typ].values():
|
|
yield m
|
|
|
|
def all_members(self, system, typ=None):
|
|
if type(typ) == str and typ in self.model_names:
|
|
typ = self.model_names[typ]
|
|
|
|
if type(system) == System:
|
|
system = system.symbol
|
|
|
|
if system not in self.system_members:
|
|
return
|
|
for m in self.system_members[system]:
|
|
if typ is None or type(m) == typ:
|
|
yield m
|
|
|
|
def flush(self):
|
|
it = 0
|
|
start_time = time()
|
|
for obj in self.dirty_objects:
|
|
it += 1
|
|
self.store(obj)
|
|
self.dirty_objects = set()
|
|
dur = time() - start_time
|
|
# print(f'flush done {it} items {dur:.2f}')
|