Moves VariantFunction and VariantMethod class definitions out of closure

Author: pganssle

src/variants/_variants.py

@@ -6,87 +6,110 @@
 __all__ = ['primary']
 
 
-def primary(f):
-    """
-    Decorator to register a function that has variant forms.
+VARIANT_WRAPPED_ATTRIBUTES = (
+    '__module__', '__name__', '__doc__', '__qualname__', '__annotations__'
+)
 
-    Decorate the main form of the function with this decorator, and then
-    subsequent variants should be declared with the same name as the original
-    function.
 
-    .. example::
+def variant_wraps(vfunc, wrapped_attributes=VARIANT_WRAPPED_ATTRIBUTES):
+    """Update the variant function wrapper a la ``functools.wraps``."""
+    f = vfunc.__main_form__
 
-        @variants
-        def myfunc(fpath):
-            with open(fpath, 'r') as f:
-                do_something(f.read())
+    class SentinelObject:
+        """A unique sentinel that is not None."""
 
-        @myfunc.variant('from_url') as f:
-        def myfunc(url):
-            r = requests.get(url)
-            do_something(r.text)
-    """
+    sentinel = SentinelObject()
 
-    class VariantFunction(object):
-        __doc__ = f.__doc__
+    for attr in wrapped_attributes:
+        attr_val = getattr(f, attr, sentinel)
+        if attr is not sentinel:
+            setattr(vfunc, attr, attr_val)
 
-        def __init__(self):
-            self._variants = set()
+    return vfunc
 
-        def __call__(self, *args, **kwargs):
-            return f(*args, **kwargs)
 
-        def _add_variant(self, var_name, vfunc):
-            self._variants.add(var_name)
-            setattr(self, var_name, vfunc)
+class VariantFunction(object):
+    """Wrapper class for functions with variant forms."""
 
-        def variant(self, func_name):
-            """Decorator to add a new variant form to the function."""
-            def decorator(vfunc):
-                self._add_variant(func_name, vfunc)
+    def __init__(self, primary_func):
+        self._variants = set()
+        self.__main_form__ = primary_func
 
-                return self
+    def __call__(self, *args, **kwargs):
+        return self.__main_form__(*args, **kwargs)
 
-            return decorator
+    def _add_variant(self, var_name, vfunc):
+        self._variants.add(var_name)
+        setattr(self, var_name, vfunc)
 
-        def __get__(self, obj, objtype=None):
-            # This is necessary to bind instance methods
-            if obj is not None:
-                return VariantMethod(self, obj)
+    def variant(self, func_name):
+        """Decorator to add a new variant form to the function."""
+        def decorator(vfunc):
+            self._add_variant(func_name, vfunc)
 
             return self
 
-        def __repr__(self):
-            return '<{} {}>'.format(self.__class__.__name__, self.__name__)
+        return decorator
+
+    def __get__(self, obj, objtype=None):
+        # This is necessary to bind instance methods
+        if obj is not None:
+            rv = VariantMethod(self, obj)
+            return variant_wraps(rv)
+
+        return self
+
+    def __repr__(self):
+        return '<{} {}>'.format(self.__class__.__name__, self.__name__)
+
+
+class VariantMethod(VariantFunction):
+    """Wrapper class for methods with variant forms."""
 
-    class VariantMethod(VariantFunction):
-        def __init__(self, variant_func, instance):
-            self.__instance = instance
-            self.__name__ = variant_func.__name__
-            self.__doc__ = variant_func.__doc__
+    def __init__(self, variant_func, instance):
+        self.__instance = instance
 
-            # Convert existing variants to methods
-            for vname in variant_func._variants:
-                vfunc = getattr(variant_func, vname)
-                vmethod = self._as_bound_method(vfunc)
+        # Convert existing variants to methods
+        for vname in variant_func._variants:
+            vfunc = getattr(variant_func, vname)
+            vmethod = self._as_bound_method(vfunc)
 
-                setattr(self, vname, vmethod)
+            setattr(self, vname, vmethod)
 
-        def __call__(self, *args, **kwargs):
-            return f(self.__instance, *args, **kwargs)
+        self.__main_form__ = self._as_bound_method(variant_func.__main_form__)
 
-        def _as_bound_method(self, vfunc):
-            @functools.wraps(vfunc)
-            def bound_method(*args, **kwargs):
-                return vfunc(self.__instance, *args, **kwargs)
+    def _as_bound_method(self, vfunc):
+        @functools.wraps(vfunc)
+        def bound_method(*args, **kwargs):
+            return vfunc(self.__instance, *args, **kwargs)
 
-            return bound_method
+        return bound_method
 
-        def _add_variant(self, var_name, vfunc):
-            self._variants.add(var_name)
-            setattr(self, var_name, self._as_bound_method(vfunc))
+    def _add_variant(self, var_name, vfunc):
+        self._variants.add(var_name)
+        setattr(self, var_name, self._as_bound_method(vfunc))
 
-    f_out = VariantFunction()
-    f_out.__name__ = f.__name__
+
+def primary(f):
+    """
+    Decorator to register a function that has variant forms.
+
+    Decorate the main form of the function with this decorator, and then
+    subsequent variants should be declared with the same name as the original
+    function.
+
+    .. example::
+
+        @variants
+        def myfunc(fpath):
+            with open(fpath, 'r') as f:
+                do_something(f.read())
+
+        @myfunc.variant('from_url') as f:
+        def myfunc(url):
+            r = requests.get(url)
+            do_something(r.text)
+    """
+    f_out = variant_wraps(VariantFunction(f))
 
     return f_out