Allow indexing API object types by class

This allows code like:
    from ipalib.plugins.dns import dnszone_mod

    api.Command[dnszone_mod]

This form should be preferred when getting specific objects
because it ensures that the appropriate plugin is imported.

https://fedorahosted.org/freeipa/ticket/4185

Reviewed-By: Martin Kosek <mkosek@redhat.com>

2 files changed, 23 insertions, 4 deletions

diff --git a/ipalib/base.py b/ipalib/base.py
index 83d778fe7..91259c7f3 100644
--- a/ipalib/base.py
+++ b/ipalib/base.py
@@ -287,7 +287,7 @@ class NameSpace(ReadOnly):
     >>> class Member(object):
     ...     def __init__(self, i):
     ...         self.i = i
-    ...         self.name = 'member%d' % i
+    ...         self.name = self.__name__ = 'member%d' % i
     ...     def __repr__(self):
     ...         return 'Member(%d)' % self.i
     ...
@@ -378,6 +378,14 @@ class NameSpace(ReadOnly):
     >>> unsorted_ns[0]
     Member(7)
 
+    As a special extension, NameSpace objects can be indexed by objects that
+    have a "__name__" attribute (e.g. classes). These lookups are converted
+    to lookups on the name:
+
+    >>> class_ns = NameSpace([Member(7), Member(3), Member(5)], sort=False)
+    >>> unsorted_ns[Member(3)]
+    Member(3)
+
     The `NameSpace` class is used in many places throughout freeIPA.  For a few
     examples, see the `plugable.API` and the `frontend.Command` classes.
     """
@@ -447,6 +455,7 @@ class NameSpace(ReadOnly):
         """
         Return ``True`` if namespace has a member named ``name``.
         """
+        name = getattr(name, '__name__', name)
         return name in self.__map
 
     def __getitem__(self, key):
@@ -455,12 +464,14 @@ class NameSpace(ReadOnly):
 
         :param key: The name or index of a member, or a slice object.
         """
+        key = getattr(key, '__name__',  key)
         if isinstance(key, basestring):
             return self.__map[key]
         if type(key) in (int, slice):
             return self.__members[key]
         raise TypeError(
-            TYPE_ERROR % ('key', (str, int, slice), key, type(key))
+            TYPE_ERROR % ('key', (str, int, slice, 'object with __name__'),
+                          key, type(key))
         )
 
     def __repr__(self):
diff --git a/ipatests/test_ipalib/test_base.py b/ipatests/test_ipalib/test_base.py
index ef6c180c7..0117946bc 100644
--- a/ipatests/test_ipalib/test_base.py
+++ b/ipatests/test_ipalib/test_base.py
@@ -195,7 +195,7 @@ def membername(i):
 class DummyMember(object):
     def __init__(self, i):
         self.i = i
-        self.name = membername(i)
+        self.name = self.__name__ = membername(i)
 
 
 def gen_members(*indexes):
@@ -283,8 +283,10 @@ class test_NameSpace(ClassChecker):
             for i in yes:
                 assert membername(i) in o
                 assert membername(i).upper() not in o
+                assert DummyMember(i) in o
             for i in no:
                 assert membername(i) not in o
+                assert DummyMember(i) not in o
 
     def test_getitem(self):
         """
@@ -320,9 +322,15 @@ class test_NameSpace(ClassChecker):
             assert o[:-4] == members[:-4]
             assert o[-9:-4] == members[-9:-4]
 
+            # Test retrieval by value
+            for member in members:
+                assert o[DummyMember(member.i)] is member
+
             # Test that TypeError is raised with wrong type
             e = raises(TypeError, o.__getitem__, 3.0)
-            assert str(e) == TYPE_ERROR % ('key', (str, int, slice), 3.0, float)
+            assert str(e) == TYPE_ERROR % (
+                'key', (str, int, slice, 'object with __name__'),
+                3.0, float)
 
     def test_repr(self):
         """