Make gargoyle.register() usable as a decorator, and gargoyle.unregister() return whether it unregistered

Fixes disqus#75.

Author: adamchainz

HISTORY.rst

@@ -8,6 +8,9 @@ Pending Release
 ---------------
 
 * New release notes here
+* Made ``gargoyle.register()`` usable as a decorator
+* Made ``gargoyle.unregister()`` return the boolean value of whether something
+  was unregistered.
 
 1.2.5 (2016-05-09)
 ------------------

gargoyle/builtins.py

@@ -58,6 +58,7 @@ def clean(self, value):
         return value
 
 
+@gargoyle.register
 class IPAddressConditionSet(RequestConditionSet):
     percent = Percent()
     ip_address = IPAddress(label='IP Address')
@@ -90,9 +91,8 @@ def _ip_to_int(self, ip):
     def get_group_label(self):
         return 'IP Address'
 
-gargoyle.register(IPAddressConditionSet())
-
 
+@gargoyle.register
 class HostConditionSet(ConditionSet):
     hostname = String()
 
@@ -108,5 +108,3 @@ def get_field_value(self, instance, field_name):
 
     def get_group_label(self):
         return 'Host'
-
-gargoyle.register(HostConditionSet())

gargoyle/manager.py

@@ -104,8 +104,11 @@ def register(self, condition_set):
         """
 
         if callable(condition_set):
-            condition_set = condition_set()
-        self._registry[condition_set.get_id()] = condition_set
+            registerable = condition_set()
+        else:
+            registerable = condition_set
+        self._registry[registerable.get_id()] = registerable
+        return condition_set
 
     def unregister(self, condition_set):
         """
@@ -114,8 +117,11 @@ def unregister(self, condition_set):
         >>> gargoyle.unregister(condition_set)
         """
         if callable(condition_set):
-            condition_set = condition_set()
-        self._registry.pop(condition_set.get_id(), None)
+            registerable = condition_set()
+        else:
+            registerable = condition_set
+        popped = self._registry.pop(registerable.get_id(), None)
+        return (popped is not None)
 
     def get_condition_set_by_id(self, switch_id):
         """

tests/testapp/test_api.py

@@ -30,6 +30,28 @@ def setUp(self):
         self.gargoyle.register(UserConditionSet(User))
         self.gargoyle.register(IPAddressConditionSet())
 
+    def test_register_unregister_class(self):
+        klass = IPAddressConditionSet
+        registered = self.gargoyle.register(IPAddressConditionSet)
+        assert registered is klass
+
+        unregistered = self.gargoyle.unregister(klass)
+        assert unregistered
+
+        unregistered = self.gargoyle.unregister(klass)
+        assert not unregistered
+
+    def test_register_unregister_instance(self):
+        condition_set = IPAddressConditionSet()
+        registered = self.gargoyle.register(condition_set)
+        assert registered is condition_set
+
+        unregistered = self.gargoyle.unregister(condition_set)
+        assert unregistered
+
+        unregistered = self.gargoyle.unregister(condition_set)
+        assert not unregistered
+
     def test_builtin_registration(self):
         assert 'gargoyle.builtins.UserConditionSet(auth.user)' in self.gargoyle._registry
         assert 'gargoyle.builtins.IPAddressConditionSet' in self.gargoyle._registry