Merge pull request #1 from hyptocrypto/multi-return-decorator

Multi return decorator

Author: hyptocrypto

README.md

@@ -22,15 +22,15 @@ EnforcedTypeError: Can not assign type: <class 'int'> to attribute 'name'. Must
 EnforcedTypeError is thrown because the value of the name attribute on a Person object must always be a string.
 
 ### Decorators
-There are also some handy decorators to ensure that a method arguments are the right type, or that the method returns the proper type. Using the <b>check_return_type</b> decorator will throw an error if the method returns an invalid type.
+There are also some handy decorators to ensure that a method arguments are the right type, or that the method returns the proper type. Using the <b>check_return_types</b> decorator will throw an error if the method returns an invalid type.
 ```python
-from pystrong import check_return_type
+from pystrong import check_return_types
 
-@check_return_type(str, int)
+@check_return_types(str, int)
 def test_success():
     return "Hello World", 500
 
-@check_return_type(str, int)
+@check_return_types(str, int)
 def test_fail():
     return 500, {"test": "error"}
 

src/pystrong/__init__.py

@@ -1,4 +1,9 @@
-from .decorators import check_return_type, check_arg_type
+from .decorators import check_arg_type, check_return_types
 from .enforcer import InferredTypeEnforcer, TypeEnforcer
 
-__all__ = ["TypeEnforcer", "InferredTypeEnforcer", "check_return_type", "check_arg_type"]
+__all__ = [
+    "TypeEnforcer",
+    "InferredTypeEnforcer",
+    "check_return_types",
+    "check_arg_type",
+]

src/pystrong/decorators.py

@@ -1,21 +1,41 @@
-from .exceptions import EnforcedReturnTypeError, EnforcedArgTypeError
+from .exceptions import EnforcedArgTypeError, EnforcedReturnTypeError
 
+non_iterables = [int, float, bool]
 
-def check_return_type(*ret_types):
+
+def check_return_types(*ret_types):
+    # Make sure all args passed to decorator are a valid python type
+    if not ret_types:
+        raise EnforcedReturnTypeError(
+            "Please pass at least one type to check_return_types"
+        )
     for _type in ret_types:
         if type(_type) is not type:
             raise EnforcedReturnTypeError(
-                f"Check return type arguments must be a valid python type. Please check '{_type}'"
+                "Check return type arguments must be a valid python type. Please check"
+                f" '{_type}'"
             )
 
     def wrapper(func):
         def inner(*args, **kwargs):
-            retruns = func(*args, **kwargs)
-            for ret, ret_type in zip(retruns, ret_types):
-                if type(ret) != ret_type:
+            # The return values form the function that is being decorated
+            rets = func(*args, **kwargs)
+            func_return_types = parse_return_types(rets)
+
+            if len(func_return_types) != len(ret_types):
+                raise EnforcedReturnTypeError(
+                    "Length of type argument passed to check_return_type is not equal"
+                    f" to length of values returned by function '{func.__name__}'"
+                )
+
+            # Iter over actual returns and expected returns and check for validity
+            for ret, ret_type in zip(func_return_types, ret_types):
+                if ret != ret_type:
                     raise EnforcedReturnTypeError(
-                        f"Function '{func.__name__}' returned '{ret}' of type '{type(ret)}'. Expected return type is '{ret_type}'."
+                        f"Function '{func.__name__}' returned '{func_return_types}'."
+                        f" Expected return type format is '{ret_types}'."
                     )
+
             return func(*args, **kwargs)
 
         return inner
@@ -42,3 +62,20 @@ def inner(*args):
         return inner
 
     return wrapper
+
+
+def parse_return_types(ret):
+    # If func returns a actual python type
+    if type(ret) is type:
+        return [type]
+
+    # If func returns a non iterable such as True, 4, or 23.3
+    if type(ret) in non_iterables:
+        return [type(ret)]
+
+    # Strings are iterable, but single string should not be iterated
+    if type(ret) == str:
+        return [str]
+
+    # Return a list of types for the func returns
+    return [type(r) for r in ret]

src/pystrong/enforcer.py

@@ -22,19 +22,22 @@ def __init__(self, *args, **kwargs):
         for attr_name, _type in kwargs.items():
             if type(_type) is not type:
                 raise BadTypeError(
-                    "Key word arguments must be in (attr_name, type) format. Ensure 'type' is a valid python type."
+                    "Key word arguments must be in (attr_name, type) format. Ensure"
+                    " 'type' is a valid python type."
                 )
             self.__dict__.update({TYPE_ATTR_FORMAT.format(attr_name): _type})
 
     def __setattr__(self, key: str, value: Any) -> None:
         if not self.__dict__.get(TYPE_ATTR_FORMAT.format(key)):
             raise AttributeTypeNotSet(
-                f"No type for attr '{key}' set. This is most likely due to not calling super in the derived class constructor."
+                f"No type for attr '{key}' set. This is most likely due to not calling"
+                " super in the derived class constructor."
             )
 
         if type(value) is not getattr(self, f"___{key}_type"):
             raise EnforcedTypeError(
-                f"Cant not assign type: {type(value)} to attribute '{key}'. Must be of type: {getattr(self, f'___{key}_type')}"
+                f"Cant not assign type: {type(value)} to attribute '{key}'. Must be of"
+                f" type: {getattr(self, f'___{key}_type')}"
             )
         super().__setattr__(key, value)
 

tests/test_decorators.py

@@ -1,31 +1,65 @@
 import pytest
-from src.pystrong import check_return_type
+from src.pystrong import check_return_types
 from src.pystrong.exceptions import EnforcedReturnTypeError
 
-
-
-@check_return_type(str)
+# Should pass
+@check_return_types(str)
 def return_str():
     return "hello"
 
-@check_return_type(int)
+
+@check_return_types(int)
 def return_int():
     return 33
 
-@check_return_type(int)
+
+@check_return_types(type)
+def return_type():
+    return bool
+
+
+@check_return_types(int, str, dict, list)
+def return_multiple():
+    return 33, "hello", {}, []
+
+
+# Should fail
+@check_return_types(int)
 def return_wrong_type():
     return "hello"
 
+
+@check_return_types(int, str, dict, list)
+def return_wrong_multiple():
+    return "hello", "hello", 33, 3.34
+
+
+@check_return_types(int, str, dict)
+def return_wrong_lengths():
+    return 33
+
+
 with pytest.raises(EnforcedReturnTypeError):
-    @check_return_type("error")
+
+    # Fail on check that all args are valid python types
+    @check_return_types("error")
     def return_error():
-        return 'error'
+        return "error"
+
     return_error()
 
 
-def test_check_return_type():
-    assert(return_str())
-    assert(return_int())
-        
+def test_check_return_types():
+    assert return_str()
+    assert return_int()
+    assert return_type()
+    assert return_multiple()
+
     with pytest.raises(EnforcedReturnTypeError):
         return_wrong_type()
+
+    with pytest.raises(EnforcedReturnTypeError):
+        return_wrong_lengths()
+
+    with pytest.raises(EnforcedReturnTypeError):
+        return_wrong_multiple()