returns check for superfunctions

Author: esblinov

transfunctions/decorators/superfunction.py

@@ -1,12 +1,9 @@
-from transfunctions.transformer import FunctionTransformer
-from inspect import currentframe
-
-
 import sys
 import weakref
-
+from ast import NodeTransformer, Expr, AST
+from inspect import currentframe
 from functools import wraps
-from typing import Dict, Any
+from typing import Dict, Any, Optional, Union, List
 from collections.abc import Coroutine
 
 if sys.version_info <= (3, 10):  # pragma: no cover
@@ -16,6 +13,9 @@
 
 from displayhooks import not_display
 
+from transfunctions.transformer import FunctionTransformer
+from transfunctions.errors import WrongTransfunctionSyntaxError
+
 
 if sys.version_info <= (3, 9):  # pragma: no cover
     CoroutineClass = Coroutine
@@ -55,7 +55,16 @@ async def async_sleep_option(flags: Dict[str, bool], args, kwargs, transformer)
 not_display(UsageTracer)
 
 def superfunction(function):
-    transformer = FunctionTransformer(function, currentframe().f_back.f_lineno, 'superfunction')
+    class NoReturns(NodeTransformer):
+        def visit_Return(self, node: Expr) -> Optional[Union[AST, List[AST]]]:
+            raise WrongTransfunctionSyntaxError('A superfunction cannot contain a return statement.')
+
+    transformer = FunctionTransformer(
+        function,
+        currentframe().f_back.f_lineno,
+        'superfunction',
+        extra_transformers=[NoReturns()],
+    )
 
     @wraps(function)
     def wrapper(*args, **kwargs):

transfunctions/errors.py

@@ -8,3 +8,7 @@ class DualUseOfDecoratorError(SyntaxError):
 
 class WrongDecoratorSyntaxError(SyntaxError):
     pass
+
+
+class WrongTransfunctionSyntaxError(SyntaxError):
+    pass

transfunctions/transformer.py

@@ -12,7 +12,7 @@
 
 
 class FunctionTransformer:
-    def __init__(self, function: Callable, decorator_lineno: int, decorator_name: str) -> None:
+    def __init__(self, function: Callable, decorator_lineno: int, decorator_name: str, extra_transformers: Optional[List[NodeTransformer]] = None) -> None:
         if isinstance(function, type(self)):
             raise DualUseOfDecoratorError(f"You cannot use the '{decorator_name}' decorator twice for the same function.")
         if not isfunction(function):
@@ -25,6 +25,7 @@ def __init__(self, function: Callable, decorator_lineno: int, decorator_name: st
         self.function = function
         self.decorator_lineno = decorator_lineno
         self.decorator_name = decorator_name
+        self.extra_transformers = extra_transformers
         self.base_object = None
         self.cache = {}
 
@@ -105,6 +106,12 @@ def extract_context(self, context_name: str, addictional_transformers: Optional[
         except OSError:
             source_code = dill_getsource(self.function)
 
+        if addictional_transformers is None:
+            addictional_transformers = self.extra_transformers
+        else:
+            if self.extra_transformers is not None:
+                addictional_transformers = addictional_transformers + self.extra_transformers
+
         converted_source_code = self.clear_spaces_from_source_code(source_code)
         tree = parse(converted_source_code)
         original_function = self.function