Refactoring, split into more files. Add more personalisation

run_selector.py

@@ -7,6 +7,31 @@ from pp_params import *
 import f90nml
 
 
+class NamelistRecursive:
+    def __init__(self, namelist):
+        self.data = namelist
+
+    def get_nml_value(self, nml_key):
+        res = self.data
+        for key in nml_key.split("/"):
+            if key in res:
+                res = res[key]
+            elif key == nml_key.split("/")[-1]:
+                res = None
+            else:
+                raise KeyError(key)
+        return res
+
+    def __getitem__(self, key):
+        return self.get_nml_value(key)
+
+    def __repr__(self):
+        return self.data.__repr__()
+
+    def __str__(self):
+        return self.data.__str__()
+
+
 class RunSelector:
     def __init__(
         self,
@@ -26,6 +51,9 @@ class RunSelector:
         self.namelist = {}
         self.runs = self.get_runs(in_runs, name_run, namelist_cond, sort_run_by)
 
+        if len(self.runs) == 0:
+            raise ValueError("No runs found")
+
         self.info = {}
         for run in self.runs:
             self.info[run] = {}
@@ -37,24 +65,36 @@ class RunSelector:
             for run in self.runs:
                 in_nums[run] = nums_temp
 
-        for run in self.runs:
+        for i, run in enumerate(self.runs):
             self.nums[run] = self.get_nums(run, in_nums[run], time_min, time_max)
 
     def load_namelist(self, run):
         path_run = self.path_in + "/" + run
         path_nml = path_run + "/" + self.pp_params.input.nml_filename
-        return f90nml.read(path_nml)
+        return NamelistRecursive(f90nml.read(path_nml))
 
     def get_nml_value(self, nml_key, run):
-        res = self.namelist[run]
-        for key in nml_key.split("/"):
-            if key in res:
-                res = res[key]
-            elif key == nml_key.split("/")[-1]:
-                res = None
-            else:
-                raise KeyError(key)
-        return res
+        return self.namelist[run][nml_key]
+
+    def nml_select(self, runs, namelist_cond):
+        if type(namelist_cond) == tuple:
+            namelist_cond = [namelist_cond]
+
+        for (nml_key, operator, operand) in namelist_cond:
+            value = {}
+            for run in runs:
+                value[run] = self.get_nml_value(nml_key, run)
+            if operator == "=":
+                runs = filter(lambda r: value[r] == operand, runs)
+            if operator == "!=":
+                runs = filter(lambda r: not value[r] == operand, runs)
+            elif operator == ">":
+                runs = filter(lambda r: value[r] > operand, runs)
+            elif operator == "<":
+                runs = filter(lambda r: value[r] < operand, runs)
+            elif operator == "in":
+                runs = filter(lambda r: value[r] in operand, runs)
+        return runs
 
     def get_runs(self, in_runs=None, name_run="*", namelist_cond={}, sort_run_by=None):
         def try_load_nml(run):
@@ -73,23 +113,8 @@ class RunSelector:
             runs = filter(lambda n: n in runs, in_runs)
         runs = filter(try_load_nml, runs)
 
-        if type(namelist_cond) == tuple:
-            namelist_cond = [namelist_cond]
-
-        for (nml_key, operator, operand) in namelist_cond:
-            value = {}
-            for run in runs:
-                value[run] = self.get_nml_value(nml_key, run)
-            if operator == "=":
-                runs = filter(lambda r: value[r] == operand, runs)
-            if operator == "!=":
-                runs = filter(lambda r: not value[r] == operand, runs)
-            elif operator == ">":
-                runs = filter(lambda r: value[r] > operand, runs)
-            elif operator == "<":
-                runs = filter(lambda r: value[r] < operand, runs)
-            elif operator == "in":
-                runs = filter(lambda r: value[r] in operand, runs)
+        # Select runs that match namelist conditions
+        runs = self.nml_select(runs, namelist_cond)
 
         # Sort by the value in the namelist of sort_run_by
         if not sort_run_by is None:
@@ -145,14 +170,20 @@ class RunSelector:
 
         if in_nums == "first":
             i = 0
-            while i < len(nums) - 1 and not try_load_info(nums[i]):
+            while i < len(nums) and not try_load_info(nums[i]):
                 i = i + 1
-            nums = [nums[i]]
+            if i < len(nums):
+                nums = [nums[i]]
+            else:
+                nums = []
         elif in_nums == "last":
             i = len(nums) - 1
-            while i > 0 and not try_load_info(nums[i]):
+            while i >= 0 and not try_load_info(nums[i]):
                 i = i - 1
-            nums = [nums[i]]
+            if i >= 0:
+                nums = [nums[i]]
+            else:
+                nums = []
         else:
             nums = filter(try_load_info, nums)
 
@@ -160,5 +191,4 @@ class RunSelector:
             nums = filter(lambda n: self.info[run][n]["time"] >= time_min, nums)
         if not time_max is None:
             nums = filter(lambda n: self.info[run][n]["time"] <= time_max, nums)
-
         return nums