Consistently use native str for keys in dict of WorkElement.

src/gmx/workflow.py

@@ -30,6 +30,9 @@
 from __future__ import print_function
 from __future__ import unicode_literals
 
+__all__ = ['WorkSpec', 'WorkElement']
+
+import collections
 import warnings
 import weakref
 
@@ -38,8 +41,6 @@
 from gmx.util import to_string
 from gmx.util import to_utf8
 
-__all__ = ['WorkSpec', 'WorkElement']
-
 # Module-level logger
 logger = logging.getLogger(__name__)
 logger.info('Importing gmx.workflow')
@@ -51,6 +52,23 @@
 # module-level constant indicating a workflow implementing parallel array work.
 # ARRAY = 0
 
+class GmxMap(collections.UserDict):
+    """Utility/compatibility class to ensure consistent keys.
+
+    Allow subscripting to use native str or Python 2 Unicode objects.
+    Internally, converts all keys to native str for the current interpreter.
+    """
+
+    def __getitem__(self, key):
+        return super(GmxMap, self).__getitem__(str(key))
+
+    def __setitem__(self, key, item) -> None:
+        super(GmxMap, self).__setitem__(str(key), item)
+
+    def __delitem__(self, key) -> None:
+        super(GmxMap, self).__delitem__(str(key))
+
+
 class WorkSpec(object):
     """
     Container of workflow elements with data dependency
@@ -124,7 +142,7 @@ class WorkSpec(object):
     """
     def __init__(self):
         self.version = workspec_version
-        self.elements = dict()
+        self.elements = GmxMap()
         self.__context_weak_ref = None
 
     @property
@@ -350,9 +368,9 @@ def __init__(self, namespace="gmxapi", operation=None, params=None, depends=()):
         # Note: Nothing currently prevents attribute updates by assignment after adding the element to a workspec,
         # but this protocol will be clarified with https://github.com/kassonlab/gmxapi/issues/92
         if params is None:
-            self.params = {}
+            self.params = GmxMap()
         elif isinstance(params, dict):
-            self.params = {to_string(name): params[name] for name in params}
+            self.params = GmxMap({to_string(name): params[name] for name in params})
         else:
             raise exceptions.UsageError("If provided, params must be a dictionary of keyword arguments")
         self.depends = list(depends)
@@ -407,7 +425,7 @@ def serialize(self):
         import json
         output_dict = {'namespace': self.namespace,
                        'operation': self.operation,
-                       'params': self.params,
+                       'params': self.params.data,
                        'depends': self.depends
                        }
         serialization = json.dumps(output_dict)