migrate logic to manager

2 files changed, 102 insertions, 93 deletions

diff --git a/metchart/manager.py b/metchart/manager.py
new file mode 100644
index 0000000..a7e3eaa
--- /dev/null
+++ b/metchart/manager.py
@@ -0,0 +1,95 @@
+from typing import Callable
+
+import yaml
+import json
+import os
+
+from multiprocessing import cpu_count
+from multiprocessing.pool import ThreadPool
+
+
+def run_if_present(key, dct: dict, func: Callable, *args, **kwargs):
+    if key in dct:
+        func(dct[key], *args, **kwargs)
+
+
+class Manager:
+    def __init__(self, filename: str = 'metchart.yaml'):
+        self.aggregators={}
+        self._plotters=[]
+
+        self._filename = filename
+        self._output_dir = './web/data'
+        self._thread_count = max(cpu_count()-1, 1)
+
+        self._load()
+        self._parse()
+
+    def run_plotters(self):
+        index = ThreadPool(self._thread_count).map(lambda p: p['module'].run(**p['cfg']), self._plotters)
+
+        with open(os.path.join(self._output_dir, 'index.json'), 'w') as f:
+            f.write(json.dumps(index, indent=4))
+
+    def _load(self):
+        with open(self._filename, 'r') as f:
+            self._raw_config = yaml.safe_load(f)
+
+    def _parse(self):
+        run_if_present('output', self._raw_config, self._parse_output)
+        run_if_present('thread_count', self._raw_config, self._parse_thread_count)
+        run_if_present('aggregator', self._raw_config, self._parse_module, self._load_aggregator)
+        run_if_present('modifier', self._raw_config, self._parse_module, self._load_modifier)
+
+        run_if_present('plotter', self._raw_config, self._parse_module, self._prepare_plotter)
+
+    def _parse_module(self, data, then: Callable):
+        # TODO abstraction prbly off. anonymous reeks
+        anonymous = False
+        if type(data) is list:
+            anonymous = True
+
+        for key in data:
+            cfg = data[key] if not anonymous else key
+
+            if 'module' not in cfg:
+                print(f'ERROR: {key} is missing the "module" keyword.')
+                continue
+
+            classname = cfg['module']
+            del cfg['module']
+            module = __import__(classname, fromlist=[None])
+
+            then(key if not anonymous else None, module, cfg)
+
+    def _load_aggregator(self, name: str, module, cfg):
+        self.aggregators[name] = module.load_data(name=name, **cfg)
+
+    def _load_modifier(self, name: str, module, cfg):
+        if 'aggregator' in cfg:
+            if type(cfg['aggregator']) == list:
+                cfg['data'] = []
+                for ag in cfg['aggregator']:
+                    cfg['data'].append(self.aggregators[ag])
+
+                del cfg['aggregator']
+            else:
+                cfg['data'] = self.aggregators[cfg['aggregator']]
+                del cfg['aggregator']
+
+        self.aggregators[name] = module.run(**cfg)
+
+    def _prepare_plotter(self, _name, module, cfg):
+        if 'aggregator' in cfg:
+            cfg['data'] = self.aggregators[cfg['aggregator']]
+            del cfg['aggregator']
+
+        self._plotters.append({
+                'module': module,
+                'cfg': cfg
+        })
+
+    def _parse_output(self, data: str):
+        self._output_dir = data
+    def _parse_thread_count(self, data: int):
+        self._thread_count = data
diff --git a/metchart/run.py b/metchart/run.py
index 17a0a97..a0e4139 100755
--- a/metchart/run.py
+++ b/metchart/run.py
@@ -1,110 +1,24 @@
 #!/usr/bin/env python3
 
 import sys
-import yaml
-import json
-import matplotlib.pyplot as plt
 import matplotlib as mpl
-from matplotlib.colors import LinearSegmentedColormap
 
-from metpy.units import units
-
-def create_aggregators(cfg):
-    ret = {}
-    for aggregator in cfg:
-        aggconf = cfg[aggregator]
-        classname = aggconf['module']
-        del aggconf['module']
-
-        module = __import__(classname, fromlist=[None])
-
-        ret[aggregator] = module.load_data(name=aggregator, **aggconf)
-
-    return ret
-
-def create_modifiers(aggregators, cfg):
-    # naming is scuffed
-    ret = {}
-    for modifier in cfg:
-        mod = cfg[modifier]
-        modname = mod['module']
-        del mod['module']
-
-        if 'aggregator' in mod:
-            if type(mod['aggregator']) == list:
-                mod['data'] = []
-                for ag in mod['aggregator']:
-                    mod['data'].append(aggregators[ag])
-
-                del mod['aggregator']
-            else:
-                mod['data'] = aggregators[mod['aggregator']]
-                del mod['aggregator']
-
-        pymod = __import__(modname, fromlist=[None])
-        ret[modifier] = pymod.run(**mod)
-
-    return ret
+from . import customization
+from . import manager
 
 def main():
     mpl.use('agg')
 
-    # Define custom gpm and gpdm units. The default gpm in metpy is aliased to meter.
-    # We need the correct definition
-    units.define('_gpm = 9.80665 * J/kg')
-    units.define('_gpdm = 10 * _gpm')
+    customization.register_units()
+    customization.register_colormaps()
 
-    # Define custom colormap
-    clcov_cmap = {
-        'red': (
-            (0.0, 0.0, 0.0),
-            (0.1, 0.9, 0.9),
-            (1.0, 0.3, 0.3),
-        ),
-        'green': (
-            (0.0, 0.5, 0.5),
-            (0.1, 0.9, 0.9),
-            (1.0, 0.3, 0.3),
-        ),
-        'blue': (
-            (0.0, 0.9, 0.9),
-            (0.1, 0.9, 0.9),
-            (1.0, 0.3, 0.3),
-        ),
-    }
-
-    mpl.colormaps.register(LinearSegmentedColormap('clcov', clcov_cmap))
-
-    FILE = 'config.yaml'
+    FILE = 'examples/config.yaml'
     if len(sys.argv) > 1:
         FILE = sys.argv[1]
 
-    conf = None
-    with open(FILE, 'r') as f:
-        conf = yaml.safe_load(f)
-
-    aggregators = create_aggregators(conf['aggregator'])
-
-    if 'modifier' in conf:
-        aggregators.update(create_modifiers(aggregators, conf['modifier']))
-
-    index = []
-
-    for plotter in conf['plotter']:
-        modname = plotter['module']
-        del plotter['module']
-
-        if 'aggregator' in plotter:
-            plotter['data'] = aggregators[plotter['aggregator']]
-            del plotter['aggregator']
-
-        mod = __import__(modname, fromlist=[None])
-        index.extend(mod.run(**plotter))
-
-        plt.close('all')
+    cfg = manager.Manager(FILE)
+    cfg.run_plotters()
 
-    with open(conf['index'], 'w') as f:
-        f.write(json.dumps(index, indent=4))
 
 if __name__ == '__main__':
     main()