Fix yatsm api fixer

ceholden · Jan 5, 2016 · 23b6f2d · 23b6f2d
1 parent feb4b43
commit 23b6f2d
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 2 deletions.
diff --git a/tstools/src/ts_driver/mixins/yatsm_ccdcesque.py b/tstools/src/ts_driver/mixins/yatsm_ccdcesque.py
@@ -15,9 +15,12 @@ def version_kwargs(d):
         idx = [i for i, arg in enumerate(argspec.args)
                if arg == 'estimator'][0] - 1
         if isinstance(argspec.defaults[idx], dict):
-            return d
+            if not isinstance(d['estimator'], dict):
+                d['estimator'] = {'object': d['estimator'], 'fit': {}}
         else:
-            d['estimator'] = {'object': d['estimator'], 'fit': {}}
+            if isinstance(d['estimator'], dict):
+                d['estimator'] = d['estimator']['object']
+        return d
     elif 'lm' in argspec.args:
         new_key, old_key = 'lm', 'estimator'
         d[new_key] = d.pop(old_key)

diff --git a/tstools/src/ts_driver/ts_utils.py b/tstools/src/ts_driver/ts_utils.py
@@ -192,6 +192,8 @@ def find_files(location, pattern, ignore_dirs=[], maxdepth=float('inf')):
     return results
 
 # CONFIGURATION
+
+# namedtuple storing a description and value for a configuration entry
 ConfigItem = namedtuple('item', ['desc', 'value'])