3 files changed, 120 insertions, 1 deletions
diff --git a/ipalib/base.py b/ipalib/base.py
index aa867018..6209139f 100644
--- a/ipalib/base.py
+++ b/ipalib/base.py
@@ -292,6 +292,64 @@ class Collector(object):
 		return NameSpace(self.__d)
 
 
+class Proxy(object):
+	def __init__(self, d):
+		self.__d = d
+
+	def __getattr__(self, name):
+		if name not in self.__d:
+			raise AttributeError(name)
+		return self.__d[name]
+
+
+
+class Register(object):
+	__allowed = (
+		Command,
+		Object,
+		Method,
+		Property,
+	)
+
+	def __init__(self):
+		self.__d = {}
+		for base in self.__allowed:
+			assert inspect.isclass(base)
+			assert base.__name__ not in self.__d
+			sub_d = {}
+			self.__d[base.__name__] = sub_d
+			setattr(self, base.__name__, Proxy(sub_d))
+
+	def __iter__(self):
+		for key in self.__d:
+			yield key
+
+	def __getitem__(self, key):
+		return dict(self.__d[key])
+
+	def items(self):
+		for key in self:
+			yield (key, self[key])
+
+	def __findbase(self, cls):
+		if not inspect.isclass(cls):
+			raise exceptions.RegistrationError('not a class', cls)
+		for base in self.__allowed:
+			if issubclass(cls, base):
+				return base
+		raise exceptions.RegistrationError(
+			'not subclass of an allowed base',
+			cls,
+		)
+
+	def __call__(self, cls):
+		base = self.__findbase(cls)
+		ns = self.__d[base.__name__]
+		assert cls.__name__ not in ns
+		ns[cls.__name__] = cls
+
+
+
 class Registrar(object):
 	__objects = None
 	__commands = None
diff --git a/ipalib/exceptions.py b/ipalib/exceptions.py
index 76c7da8c..4c307177 100644
--- a/ipalib/exceptions.py
+++ b/ipalib/exceptions.py
@@ -58,7 +58,7 @@ class DuplicateError(IPAError):
 
 
 class RegistrationError(IPAError):
-	msg = '%r must be a subclass of %s'
+	msg = '%s: %r'
 
 
 class PrefixError(IPAError):
diff --git a/ipalib/tests/test_base.py b/ipalib/tests/test_base.py
index 7affd997..f11f1f4f 100644
--- a/ipalib/tests/test_base.py
+++ b/ipalib/tests/test_base.py
@@ -347,3 +347,64 @@ def test_Registar():
 
 	assert len(r.commands) == 3
 	assert list(r.commands) == sorted(['kinit', 'add_user', 'del_user'])
+
+
+class test_Register():
+	r = base.Register()
+
+	assert set(r) == set(['Command', 'Object', 'Method', 'Property'])
+
+
+	class wrong_base(object):
+		pass
+
+	class krbtest(base.Command):
+		pass
+
+	class user(base.Object):
+		pass
+
+	class user__add(base.Method):
+		pass
+
+	class user__firstname(base.Property):
+		pass
+
+
+
+
+	#r(wrong_base)
+	#r(user())
+
+	# Check that exception is raised trying to register an instance of a
+	# class of a correct base:
+	raised = False
+	try:
+		r(user())
+	except exceptions.RegistrationError:
+		raised = True
+
+	# Check that exception is raised trying to register class of wrong base:
+	raised = False
+	try:
+		r(wrong_base)
+	except exceptions.RegistrationError:
+		raised = True
+	assert raised
+
+	# Check that added a valid class works
+	for cls in (krbtest, user, user__add, user__firstname):
+		r(cls)
+		key = cls.__bases__[0].__name__
+		d = r[key]
+		assert d.keys() == [cls.__name__]
+		assert d.values() == [cls]
+		# Check that a copy is returned
+		d2 = r[key]
+		assert d2 == d
+		assert d2 is not d
+		p = getattr(r, key)
+		assert isinstance(p, base.Proxy)
+		# Check that same instance is returned
+		assert p is getattr(r, key)
+		assert getattr(p, cls.__name__) is cls