From 3fb17249d15e67914476d696abe0ab405708a563 Mon Sep 17 00:00:00 2001
From: psorn <tomas.psorn@gmail.com>
Date: Fri, 3 Mar 2023 11:36:30 +0100
Subject: [PATCH] Refactor loading stages

- The loading of dataset consists of the following sequence steps - loading
  parameters, loading properties and loading data
- LOADING_STAGES dict in dataset.py defines a state for each of these
  steps
- By passing the keyword load with value of one of those stages, the
  dataset only loads up to the chosen points
- This allows for instance to load 2dseq dataset only using visu_pars
  file, in case, you only want to access parameters of the dataset
---
 brukerapi/dataset.py | 58 +++++++++++++++++++++-----------------------
 1 file changed, 27 insertions(+), 31 deletions(-)

diff --git a/brukerapi/dataset.py b/brukerapi/dataset.py
index f25e2af..edbfd87 100644
--- a/brukerapi/dataset.py
+++ b/brukerapi/dataset.py
@@ -9,7 +9,13 @@
 import os.path
 import yaml
 import datetime
-import logging
+
+LOAD_STAGES = {
+    "empty": 0,
+    "parameters": 1,
+    "properties": 2,
+    "all": 3,
+}
 
 # Dict of default dataset states
 DEFAULT_STATES = {
@@ -19,10 +25,7 @@
             Path(__file__).parents[0] / 'config/properties_fid_core.json',
             Path(__file__).parents[0] / 'config/properties_fid_custom.json'
         ],
-        "load": True,
-        "load_parameters": True,
-        "load_properties": True,
-        "load_data": True,
+        "load": LOAD_STAGES['all'],
         "mmap": False
     },
     '2dseq': {
@@ -31,10 +34,7 @@
             Path(__file__).parents[0] / 'config/properties_2dseq_core.json',
             Path(__file__).parents[0] / 'config/properties_2dseq_custom.json'
         ],
-        "load": True,
-        "load_parameters": True,
-        "load_properties": True,
-        "load_data": True,
+        "load": LOAD_STAGES['all'],
         "scale": True,
         "mmap": False
     },
@@ -44,10 +44,7 @@
             Path(__file__).parents[0] / 'config/properties_traj_core.json',
             Path(__file__).parents[0] / 'config/properties_traj_custom.json'
         ],
-        "load": True,
-        "load_parameters": True,
-        "load_properties": True,
-        "load_data": True,
+        "load": LOAD_STAGES['all'],
         "mmap": False
     },
     'ser': {
@@ -56,10 +53,7 @@
             Path(__file__).parents[0] / 'config/properties_ser_core.json',
             Path(__file__).parents[0] / 'config/properties_ser_custom.json'
         ],
-        "load": True,
-        "load_parameters": True,
-        "load_properties": True,
-        "load_data": True,
+        "load": LOAD_STAGES['all'],
         "mmap": False
     },
     'rawdata': {
@@ -68,10 +62,7 @@
             Path(__file__).parents[0] / 'config/properties_rawdata_core.json',
             Path(__file__).parents[0] / 'config/properties_rawdata_custom.json'
         ],
-        "load": True,
-        "load_parameters": True,
-        "load_properties": True,
-        "load_data": True,
+        "load": LOAD_STAGES['all'],
         "mmap": False
     }
 }
@@ -171,7 +162,7 @@ def __init__(self, path, **state):
         """
         self.path = Path(path)
 
-        if not self.path.exists() and state.get('load') is not False:
+        if not self.path.exists() and state.get('load') is not LOAD_STAGES['empty']:
             raise FileNotFoundError(self.path)
 
         # directory constructor
@@ -189,24 +180,24 @@ def __init__(self, path, **state):
         if self.subtype: self.subtype = self.subtype[1:] # remove the dot from the suffix
         self._properties = []
 
-        # validate path
-        self._validate(state)
-
         # set
         self._set_state(state)
 
+        # validate path
+        self._validate()
+
         # load data if the load kwarg is true
         self.load()
 
 
     def __enter__(self):
-        self._state['load'] = True
+        self._state['load'] = LOAD_STAGES['all']
         self.load()
         return self
 
     def __exit__(self, exc_type, exc_val, exc_tb):
         self.unload()
-        self._state['load'] = False
+        self._state['load'] = LOAD_STAGES['empty']
 
     def __str__(self):
         """
@@ -239,7 +230,7 @@ def _set_state(self, passed):
         result.update(passed)
         self._state = result
 
-    def _validate(self, state):
+    def _validate(self):
         """Validate Dataset
 
         Check whether the dataset type is supported and complete. Dataset is allowed to be incomplete when load is
@@ -254,7 +245,7 @@ def _validate(self, state):
             raise UnsuportedDatasetType(self.type)
 
         # Check whether all necessary JCAMP-DX files are present
-        if state.get('load') is None or state.get('load') is True:
+        if self._state.get('load') >= LOAD_STAGES['parameters']:
             if not (set(DEFAULT_STATES[self.type]['parameter_files']) <= set(os.listdir(str(self.path.parent)))):
                 raise IncompleteDataset
 
@@ -264,11 +255,16 @@ def load(self):
         traj is loaded as well.
         """
 
-        if not self._state['load']:
-            return
+        if self._state['load'] is LOAD_STAGES['empty']: return
 
         self.load_parameters()
+
+        if self._state['load'] is LOAD_STAGES['parameters']: return
+
         self.load_properties()
+
+        if self._state['load'] is LOAD_STAGES['properties']: return
+
         self.load_schema()
         self.load_data()
         # self.load_traj()