diff --git a/hy/compiler.py b/hy/compiler.py index 51a86a9..2143fd4 100755 --- a/hy/compiler.py +++ b/hy/compiler.py @@ -322,21 +322,25 @@ def _raise_wrong_args_number(expression, error): len(expression))) +def _nargs(n): + return "%d argument%s" % (n, ("" if n == 1 else "s")) + + def checkargs(exact=None, min=None, max=None, even=None, multiple=None): def _dec(fn): def checker(self, expression): if exact is not None and (len(expression) - 1) != exact: _raise_wrong_args_number( - expression, "`%%s' needs %d arguments, got %%d" % exact) + expression, "`%%s' needs %s, got %%d" % _nargs(exact)) if min is not None and (len(expression) - 1) < min: _raise_wrong_args_number( expression, - "`%%s' needs at least %d arguments, got %%d." % (min)) + "`%%s' needs at least %s, got %%d." % _nargs(min)) if max is not None and (len(expression) - 1) > max: _raise_wrong_args_number( expression, - "`%%s' needs at most %d arguments, got %%d" % (max)) + "`%%s' needs at most %s, got %%d" % _nargs(max)) is_even = not((len(expression) - 1) % 2) if even is not None and is_even != even: @@ -2312,6 +2316,16 @@ class HyASTCompiler(object): return ret + @builds("return") + @checkargs(max=1) + def compile_return(self, expr): + ret = Result() + if len(expr) > 1: + ret += self.compile(expr[1]) + return ret + ast.Return(value=ret.force_expr, + lineno=expr.start_line, + col_offset=expr.start_column) + @builds("defclass") @checkargs(min=1) def compile_class_expression(self, expressions): diff --git a/tests/native_tests/language.hy b/tests/native_tests/language.hy index 51e221f..cb3aabc 100644 --- a/tests/native_tests/language.hy +++ b/tests/native_tests/language.hy @@ -849,6 +849,66 @@ (assert (= mooey.__name__ "mooey"))) +(defn test-return [] + + ; `return` in main line + (defn f [x] + (return (+ x "a")) + (+ x "b")) + (assert (= (f "q") "qa")) + + ; Nullary `return` + (defn f [x] + (return) + 5) + (assert (none? (f "q"))) + + ; `return` in `when` + (defn f [x] + (when (< x 3) + (return [x 1])) + [x 2]) + (assert (= (f 2) [2 1])) + (assert (= (f 4) [4 2])) + + ; `return` in a loop + (setv accum []) + (defn f [x] + (while True + (when (zero? x) + (return)) + (.append accum x) + (-= x 1)) + (.append accum "this should never be appended") + 1) + (assert (none? (f 5))) + (assert (= accum [5 4 3 2 1])) + + ; `return` of a `do` + (setv accum []) + (defn f [] + (return (do + (.append accum 1) + 3)) + 4) + (assert (= (f) 3)) + (assert (= accum [1])) + + ; `return` of an `if` that will need to be compiled to a statement + (setv accum []) + (defn f [x] + (return (if (= x 1) + (do + (.append accum 1) + "a") + (do + (.append accum 2) + "b"))) + "c") + (assert (= (f 2) "b")) + (assert (= accum [2]))) + + (defn test-mangles [] "NATIVE: test mangles" (assert (= 2 ((fn [] (+ 1 1))))))