diff --git a/hy/compiler.py b/hy/compiler.py index 223242e..8a80673 100644 --- a/hy/compiler.py +++ b/hy/compiler.py @@ -673,6 +673,48 @@ class HyASTCompiler(object): step=None), ctx=ast.Load()) + @builds("take") + @checkargs(2) + def compile_take_expression(self, expr): + expr.pop(0) + n = self.compile(expr.pop(0)) + seq = self.compile(expr.pop(0)) + zero = ast.Num(n=0, + lineno=expr.start_column, + col_offset=expr.start_column) + return ast.Subscript( + lineno=expr.start_line, + col_offset=expr.start_column, + value=seq, + slice=ast.Slice(lower=zero, + upper=n, + step=None), + ctx=ast.Load()) + + @builds("drop") + @checkargs(2) + def compile_drop_expression(self, expr): + expr.pop(0) + n = self.compile(expr.pop(0)) + seq = self.compile(expr.pop(0)) + zero = ast.Num(n=0, + lineno=expr.start_line, + col_offset=expr.start_column) + upper = ast.UnaryOp(op=ast.USub(), + operand=n, + lineno=expr.start_line, + col_offset=expr.start_column) + if upper.operand.n == 0: + return seq + return ast.Subscript( + lineno=expr.start_line, + col_offset=expr.start_column, + value=seq, + slice=ast.Slice(lower=zero, + upper=upper, + step=None), + ctx=ast.Load()) + @builds("assoc") @checkargs(3) def compile_assoc_expression(self, expr): diff --git a/tests/compilers/test_ast.py b/tests/compilers/test_ast.py index 29a03fe..9ceedf6 100644 --- a/tests/compilers/test_ast.py +++ b/tests/compilers/test_ast.py @@ -237,6 +237,27 @@ def test_ast_bad_slice(): cant_compile("(slice 1 2 3 4)") +def test_ast_good_take(): + "Make sure AST can compile valid 'take'" + hy_compile(tokenize("(take 1 [2 3])")) + + +def test_ast_bad_take(): + "Make sure AST chokes on bad 'take'" + cant_compile("(take)") + cant_compile("(take 2)") + +def test_ast_good_drop(): + "Make sure AST can compile valid 'drop'" + hy_compile(tokenize("(drop 1 [2 3])")) + + +def test_ast_bad_drop(): + "Make sure AST chokes on bad 'drop'" + cant_compile("(drop)") + cant_compile("(drop 2)") + + def test_ast_good_assoc(): "Make sure AST can compile valid assoc" hy_compile(tokenize("(assoc x y z)")) diff --git a/tests/native_tests/language.hy b/tests/native_tests/language.hy index 3601f0b..47eaf5a 100644 --- a/tests/native_tests/language.hy +++ b/tests/native_tests/language.hy @@ -376,6 +376,20 @@ (assert (= (slice [1 2 3 4 5]) [1 2 3 4 5]))) +(defn test-take [] + "NATIVE: test take" + (assert (= (take 0 [2 3]) [])) + (assert (= (take 1 [2 3]) [2])) + (assert (= (take 2 [2 3]) [2 3]))) + + +(defn test-drop [] + "NATIVE: test drop" + (assert (= (drop 0 [2 3]) [2 3])) + (assert (= (drop 1 [2 3]) [2])) + (assert (= (drop 2 [2 3]) []))) + + (defn test-rest [] "NATIVE: test rest" (assert (= (rest [1 2 3 4 5]) [2 3 4 5])))