Merge pull request #1380 from Kodiologist/return

Add `return` as a special form

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):

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))))))