Add a generic create_registry method

PiperOrigin-RevId: 230422498

Author: Ryan Sepassi

tensor2tensor/utils/registry.py

@@ -44,6 +44,8 @@ class MyModel(T2TModel):
 from __future__ import division
 from __future__ import print_function
 
+import collections
+
 from tensor2tensor.utils import misc_utils
 import tensorflow as tf
 from tensorflow.python.util import tf_inspect as inspect
@@ -57,6 +59,68 @@ class MyModel(T2TModel):
 _PRUNING_STRATEGY = {}
 _RANGED_HPARAMS = {}
 
+# Key: registry name, Value: Registry
+_GENERIC_REGISTRIES = {}
+Registry = collections.namedtuple(
+    "_Registry", ["register", "get", "list", "registry"])
+
+
+def registry(registry_name):
+  """Returns `Registry` created by `create_registry`."""
+  if registry_name not in _GENERIC_REGISTRIES:
+    raise KeyError("No registry named %s. Available:\n%s" % (
+        registry_name, sorted(_GENERIC_REGISTRIES)))
+  return _GENERIC_REGISTRIES[registry_name]
+
+
+def create_registry(registry_name):
+  """Create a generic object registry.
+
+  Args:
+    registry_name: str, name of the object registry.
+
+  Returns:
+    `Registry` that contains functions for register (decorator), get, and list.
+
+  Raises:
+    KeyError: if `registry_name` is a pre-existing registry.
+  """
+  if registry_name in _GENERIC_REGISTRIES:
+    raise KeyError(
+        "Registry %s already exists." % registry_name)
+
+  registry_ = {}
+
+  def register(name):
+    """Returns decorator to register an object."""
+
+    def register_dec(obj):
+      if name in registry_:
+        raise KeyError(
+            "Registry %s already contains key %s." % (registry_name, name))
+      registry_[name] = obj
+      return obj
+
+    return register_dec
+
+  def get(name):
+    if name not in registry_:
+      raise KeyError(
+          "Registry %s contains no object named %s" % (registry_name, name))
+    return registry_[name]
+
+  def list_registry():
+    return sorted(registry_)
+
+  registry_obj = Registry(
+      register=register,
+      get=get,
+      list=list_registry,
+      registry=registry_,
+  )
+  _GENERIC_REGISTRIES[registry_name] = registry_obj
+  return registry_obj
+
 
 def _reset():
   for ctr in [_MODELS, _HPARAMS, _RANGED_HPARAMS, _ATTACK_PARAMS]:

tensor2tensor/utils/registry_test.py

@@ -194,8 +194,35 @@ def rhp_bad2(a, b):  # pylint: disable=unused-argument
         pass
 
 
+class CreateRegistry(tf.test.TestCase):
+  """Test class for `create_registry`."""
+
+  def testCreateRegistry(self):
+    my_registry = registry.create_registry("test_reg1")
+    self.assertIs(my_registry, registry.registry("test_reg1"))
+
+    # Use as decorator on a fn
+    @my_registry.register("foo")
+    def some_fn(num):
+      return num + 2
+
+    # Register a regular object
+    pod_obj = 4
+    my_registry.register("bar")(pod_obj)
+
+    # Register a class
+    @my_registry.register("foobar")
+    class A(object):
+      pass
+
+    self.assertEqual(9, my_registry.get("foo")(7))
+    self.assertEqual(["bar", "foo", "foobar"], my_registry.list())
+    foobar = my_registry.get("foobar")
+    self.assertTrue(isinstance(foobar(), A))
+
+
 class RegistryTest(tf.test.TestCase):
-  """ Test class for common functions."""
+  """Test class for common functions."""
 
   def testRegistryHelp(self):
     help_str = registry.help_string()