Merge pull request #225 from choderalab/wraps

Add Python 2-3 compatibile wrapping functions

openmoltools/schrodinger.py

@@ -8,7 +8,6 @@
 import shutil
 import logging
 import subprocess
-from functools import wraps
 
 import mdtraj
 
@@ -65,7 +64,7 @@ def is_schrodinger_suite_installed():
 
 
 def need_schrodinger(func):
-    @wraps(func)
+    @utils.wraps_py2(func)
     def _need_schrodinger(*args, **kwargs):
         """Decorator that checks if the Schrodinger's suite is installed."""
         if not is_schrodinger_suite_installed():
@@ -166,7 +165,7 @@ def run_structconvert(input_file_path, output_file_path):
 
 
 def autoconvert_maestro(func):
-    @wraps(func)
+    @utils.wraps_py2(func)
     def _autoconvert_maestro(input_file_path, output_file_path, *args, **kwargs):
         """Decorator that make a function support more than only Maestro files.
 

openmoltools/utils.py

@@ -5,6 +5,7 @@
 import tempfile
 import logging
 import subprocess
+import functools
 from pkg_resources import resource_filename
 import contextlib
 import shutil
@@ -22,6 +23,10 @@
 logger = logging.getLogger(__name__)
 
 
+# ------------------------------------------------------------------------------
+# PYTHON 2-3 COMPATIBILITY FUNCTIONS
+# ------------------------------------------------------------------------------
+
 def getoutput(cmd):
     """Compatibility function to substitute deprecated commands.getoutput in Python2.7"""
     try:
@@ -35,6 +40,39 @@ def getoutput(cmd):
         return str(out)
 
 
+def wraps_py2(wrapped, *args):
+    """Wrap a function and add the __wrapped__ attribute.
+
+    In Python 2, functools.wraps does not add the __wrapped__ attribute, and it
+    becomes impossible to retrieve the signature of the wrapped method.
+
+    """
+    def decorator(wrapper):
+        functools.update_wrapper(wrapper, wrapped, *args)
+        wrapper.__wrapped__ = wrapped
+        return wrapper
+    return decorator
+
+
+def unwrap_py2(func):
+    """Unwrap a wrapped function.
+
+    The function inspect.unwrap has been implemented only in Python 3.4. With
+    Python 2, this works only for functions wrapped by wraps_py2().
+
+    """
+    unwrapped_func = func
+    try:
+        while True:
+            unwrapped_func = unwrapped_func.__wrapped__
+    except AttributeError:
+        return unwrapped_func
+
+
+# ------------------------------------------------------------------------------
+# UTILITY FUNCTIONS
+# ------------------------------------------------------------------------------
+
 def find_gaff_dat():
     print("Warning: find_gaff_dat has been moved to openmoltools.amber.")
     amber = import_("openmoltools.amber")