refactor to add documentation, clarify variable names, add test cases, and better encapsulate behaviors (among other things to simplify testing)

Author: claytondaley

.gitignore

@@ -3,3 +3,4 @@
 build
 dist
 *.egg-info
+.idea

ssm-diff

@@ -1,52 +1,53 @@
 #!/usr/bin/env python
 from __future__ import print_function
-from states import *
-import states.helpers as helpers
+
 import argparse
 import os
 
+from states import states
+from states.helpers import DiffResolver
+
+
+def configure_endpoints(args):
+    # pre-configure resolver, but still accept remote and local at runtime
+    diff_resolver = DiffResolver.configure(force=args.force)
+    return states.ParameterStore(args.profile, diff_resolver, paths=args.path), states.YAMLFile(args.filename, paths=args.path)
+
 
 def init(args):
-    r, l = RemoteState(args.profile), LocalState(args.filename)
-    l.save(r.get(flat=False, paths=args.path))
+    """Create a local YAML file from the SSM Parameter Store (per configs in args)"""
+    remote, local = configure_endpoints(args)
+    local.save(remote.clone())
 
 
 def pull(args):
-    dictfilter = lambda x, y: dict([ (i,x[i]) for i in x if i in set(y) ])
-    r, l = RemoteState(args.profile), LocalState(args.filename)
-    diff = helpers.FlatDictDiffer(r.get(paths=args.path), l.get(paths=args.path))
-    if args.force:
-        ref_set = diff.changed().union(diff.removed()).union(diff.unchanged())
-        target_set = diff.added()
-    else:
-        ref_set = diff.unchanged().union(diff.removed())
-        target_set = diff.added().union(diff.changed())
-    state = dictfilter(diff.ref, ref_set)
-    state.update(dictfilter(diff.target, target_set))
-    l.save(helpers.unflatten(state))
+    """Update local YAML file with changes in the SSM Parameter Store (per configs in args)"""
+    remote, local = configure_endpoints(args)
+    local.save(remote.pull(local.get()))
 
 
 def apply(args):
-    r, _, diff = plan(args)
-
+    """Apply local changes to the SSM Parameter Store"""
+    remote, local = configure_endpoints(args)
     print("\nApplying changes...")
     try:
-        r.apply(diff)
+        remote.push(local.get())
     except Exception as e:
         print("Failed to apply changes to remote:", e)
     print("Done.")
 
 
 def plan(args):
-    r, l = RemoteState(args.profile), LocalState(args.filename)
-    diff = helpers.FlatDictDiffer(r.get(paths=args.path), l.get(paths=args.path))
+    """Print a representation of the changes that would be applied to SSM Parameter Store if applied (per config in args)"""
+    remote, local = configure_endpoints(args)
+    diff = remote.dry_run(local.get())
 
     if diff.differ:
-        diff.print_state()
+        print(diff.describe_diff())
     else:
         print("Remote state is up to date.")
 
-    return r, l, diff
+    return remote, local, diff
 
 
 if __name__ == "__main__":

states/helpers.py

@@ -1,55 +1,101 @@
-from termcolor import colored
-from copy import deepcopy
 import collections
+from copy import deepcopy
+from functools import partial
 
+from termcolor import colored
 
-class FlatDictDiffer(object):
-    def __init__(self, ref, target):
-        self.ref, self.target = ref, target
-        self.ref_set, self.target_set = set(ref.keys()), set(target.keys())
-        self.isect = self.ref_set.intersection(self.target_set)
+
+class DiffResolver(object):
+    """Determines diffs between two dicts, where the remote copy is considered the baseline"""
+    def __init__(self, remote, local, force=False):
+        self.remote_flat, self.local_flat = self._flatten(remote), self._flatten(local)
+        self.remote_set, self.local_set = set(self.remote_flat.keys()), set(self.local_flat.keys())
+        self.intersection = self.remote_set.intersection(self.local_set)
+        self.force = force
 
         if self.added() or self.removed() or self.changed():
             self.differ = True
         else:
             self.differ = False
 
+    @classmethod
+    def configure(cls, *args, **kwargs):
+        return partial(cls, *args, **kwargs)
+
     def added(self):
-        return self.target_set - self.isect
+        """Returns a (flattened) dict of added leaves i.e. {"full/path": value, ...}"""
+        return self.local_set - self.intersection
 
     def removed(self):
-        return self.ref_set - self.isect
+        """Returns a (flattened) dict of removed leaves i.e. {"full/path": value, ...}"""
+        return self.remote_set - self.intersection
 
     def changed(self):
-        return set(k for k in self.isect if self.ref[k] != self.target[k])
+        """Returns a (flattened) dict of changed leaves i.e. {"full/path": value, ...}"""
+        return set(k for k in self.intersection if self.remote_flat[k] != self.local_flat[k])
 
     def unchanged(self):
-        return set(k for k in self.isect if self.ref[k] == self.target[k])
+        """Returns a (flattened) dict of unchanged leaves i.e. {"full/path": value, ...}"""
+        return set(k for k in self.intersection if self.remote_flat[k] == self.local_flat[k])
 
-    def print_state(self):
+    def describe_diff(self):
+        """Return a (multi-line) string describing all differences"""
+        description = ""
         for k in self.added():
-            print(colored("+", 'green'), "{} = {}".format(k, self.target[k]))
+            description += colored("+", 'green'), "{} = {}".format(k, self.local_flat[k]) + '\n'
 
         for k in self.removed():
-            print(colored("-", 'red'), k)
+            description += colored("-", 'red'), k + '\n'
 
         for k in self.changed():
-            print(colored("~", 'yellow'), "{}:\n\t< {}\n\t> {}".format(k, self.ref[k], self.target[k]))
-
-
-def flatten(d, pkey='', sep='/'):
-    items = []
-    for k in d:
-        new = pkey + sep + k if pkey else k
-        if isinstance(d[k], collections.MutableMapping):
-            items.extend(flatten(d[k], new, sep=sep).items())
+            description += colored("~", 'yellow'), "{}:\n\t< {}\n\t> {}".format(k, self.remote_flat[k], self.local_flat[k]) + '\n'
+
+        return description
+
+    def _flatten(self, d, current_path='', sep='/'):
+        """Convert a nested dict structure into a "flattened" dict i.e. {"full/path": "value", ...}"""
+        items = []
+        for k in d:
+            new = current_path + sep + k if current_path else k
+            if isinstance(d[k], collections.MutableMapping):
+                items.extend(self._flatten(d[k], new, sep=sep).items())
+            else:
+                items.append((sep + new, d[k]))
+        return dict(items)
+
+    def _unflatten(self, d, sep='/'):
+        """Converts a "flattened" dict i.e. {"full/path": "value", ...} into a nested dict structure"""
+        output = {}
+        for k in d:
+            add(
+                obj=output,
+                path=k,
+                value=d[k],
+                sep=sep,
+            )
+        return output
+
+    def merge(self):
+        """Generate a merge of the local and remote dicts, following configurations set during __init__"""
+        dictfilter = lambda original, keep_keys: dict([(i, original[i]) for i in original if i in set(keep_keys)])
+        if self.force:
+            # Overwrite local changes (i.e. only preserve added keys)
+            # NOTE:  Currently the system cannot tell the difference between a remote delete and a local add
+            prior_set = self.changed().union(self.removed()).union(self.unchanged())
+            current_set = self.added()
         else:
-            items.append((sep + new, d[k]))
-    return dict(items)
-
-
-def add(obj, path, value):
-    parts = path.strip("/").split("/")
+            # Preserve added keys and changed keys
+            # NOTE:  Currently the system cannot tell the difference between a remote delete and a local add
+            prior_set = self.unchanged().union(self.removed())
+            current_set = self.added().union(self.changed())
+        state = dictfilter(original=self.remote_flat, keep_keys=prior_set)
+        state.update(dictfilter(original=self.local_flat, keep_keys=current_set))
+        return self._unflatten(state)
+
+
+def add(obj, path, value, sep='/'):
+    """Add value to the `obj` dict at the specified path"""
+    parts = path.strip(sep).split(sep)
     last = len(parts) - 1
     for index, part in enumerate(parts):
         if index == last:
@@ -61,7 +107,7 @@ def add(obj, path, value):
 def search(state, path):
     result = state
     for p in path.strip("/").split("/"):
-        if result.get(p):
+        if result.clone(p):
             result = result[p]
         else:
             result = {}
@@ -71,16 +117,6 @@ def search(state, path):
     return output
 
 
-def unflatten(d):
-    output = {}
-    for k in d:
-        add(
-            obj=output,
-            path=k,
-            value=d[k])
-    return output
-
-
 def merge(a, b):
     if not isinstance(b, dict):
         return b