From c7739c5985be3fef9bfb8f182d2f88e5a7da34fa Mon Sep 17 00:00:00 2001
From: Richard Bronkhorst <richmans@users.noreply.github.com>
Date: Mon, 12 Jun 2023 23:13:40 +0200
Subject: [PATCH] Update commander.py and store.py

---
 nullptr/commander.py | 13 ++++----
 nullptr/store.py     | 78 ++++++++++++++++++++++----------------------
 2 files changed, 46 insertions(+), 45 deletions(-)

diff --git a/nullptr/commander.py b/nullptr/commander.py
index 04d0793..5e1a21e 100644
--- a/nullptr/commander.py
+++ b/nullptr/commander.py
@@ -16,10 +16,11 @@ class Commander(CommandLine):
   def __init__(self, store_dir='data', agent=None):
     self.store_dir = store_dir
     self.store = Store(store_dir)
+    self.store.load()
     self.agent = self.select_agent(agent)
     self.api = Api(self.store, self.agent)
     self.atlas_builder = AtlasBuilder(self.store, self.api)
-    self.store.flush()
+    
     self.stop_auto= False
     super().__init__()
   
@@ -27,7 +28,7 @@ class Commander(CommandLine):
     if agent_str is not None:
       return self.store.get(Agent, agent_str)
     else:
-      agents = self.store.all('', Agent)
+      agents = self.store.all(Agent)
       agent = next(agents, None)
       if agent is None:
         symbol = input('agent name: ')
@@ -67,25 +68,25 @@ class Commander(CommandLine):
   def do_test(self):
     start_time = time()
     total = 0
-    for m in self.store.all('', Marketplace, True):
+    for m in self.store.all(Marketplace):
       total += 1
     dur = time() - start_time
     print(f'{total} markets in {dur:.2f} seconds')
     start_time = time()
     total = 0
-    for m in self.store.all('', Jumpgate, True):
+    for m in self.store.all(Jumpgate):
       total += 1
     dur = time() - start_time
     print(f'{total} jumpgates in {dur:.2f} seconds')
     start_time = time()
     total = 0
-    for m in self.store.all('', Waypoint, True):
+    for m in self.store.all(Waypoint):
       total += 1
     dur = time() - start_time
     print(f'{total} waypoints in {dur:.2f} seconds')
     start_time = time()
     total = 0
-    for m in self.store.all('', System, True):
+    for m in self.store.all(System):
       total += 1
     dur = time() - start_time
     print(f'{total} systems in {dur:.2f} seconds')
diff --git a/nullptr/store.py b/nullptr/store.py
index edfab92..ae31da7 100644
--- a/nullptr/store.py
+++ b/nullptr/store.py
@@ -7,6 +7,7 @@ from nullptr.models.marketplace import Marketplace
 from nullptr.models.jumpgate import Jumpgate
 from os.path import isfile, dirname, isdir
 import os
+from os.path import basename
 import json
 from .util import *
 from time import time
@@ -15,44 +16,46 @@ class Store:
   def __init__(self, data_dir):
     self.init_models()
     self.data_dir = data_dir
-    self.data = {}
+    self.data = {m: {} for m in self.models}
     self.dirty_objects = set()
   
   def init_models(self):
     self.models = Base.__subclasses__()
     self.extensions = {c.ext(): 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(self, obj):
-    path = self.path(obj)
+  def load_file(self, path):
     if not isfile(path):
-      return obj
-    with open(path) as f:
-      data = json.load(f)
-    data['store'] = self
-    obj.__dict__ = data
-  
-  def get_file(self, typ, path):
-    if not isfile(path):
-      print(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)
-    symbol = mg(data, 'symbol')
-    oid = f'{symbol}.{typ.ext()}'
-    if oid in self.data:
-      return self.data[oid]
+    
+    typ = self.extensions[ext]
     obj = typ(symbol, self)
-    self.load(obj)
-    self.data[oid] = obj
+    data['store'] = self
+    obj.__dict__ = data
+    self.data[typ][obj.symbol] = obj
     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)
@@ -61,36 +64,33 @@ class Store:
       os.makedirs(path_dir, exist_ok=True)
     with open(path, 'w') as f:
       json.dump(data, f, indent=2)
-      
-  def get(self, typ, symbol):
-    oid = f'{symbol}.{typ.ext()}'
-    if oid in self.data:
-      return self.data[oid]
+  
+  def create(self, typ, symbol):
     obj = typ(symbol, self)
-    self.load(obj)
-    self.data[oid] = obj
+    self.data[typ][symbol] = obj
     return obj
+    
+  def get(self, typ, symbol, create=False):
+    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, symbol, data):
-    obj = self.get(typ, symbol)
+    obj = self.get(typ, symbol, True)
     obj.update(data)
     return obj
   
   def update_list(self, typ, lst):
     return [self.update(typ, mg(d, 'symbol'), d) for d in lst]
     
-  def all(self, path, typ, recursive=False):
-    if hasattr(path, 'path'):
-      path = path.path()
-    path = os.path.join(self.data_dir, path)
-    if not isdir(path):
-      return
-    ext = '.' + typ.ext()
-    for fil in list_files(path, recursive):
-      if not fil.endswith(ext):
-        continue
-      yield self.get_file(typ, fil)
-    
+  def all(self, typ):
+    for m in self.data[typ].values():
+      yield m
   
   def flush(self):
     it = 0