Added 'take' and 'drop' functions (as AST elements)

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

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

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