From 68cc983dea9f9f39b56eaa2c27405d819293a64a Mon Sep 17 00:00:00 2001
From: David Anthoff <anthoff@berkeley.edu>
Date: Tue, 19 Dec 2017 15:19:02 -0800
Subject: [PATCH] Add @take and @drop macros

---
 NEWS.md                        |  4 +++
 docs/src/experimental.md       | 10 +++++++-
 src/Query.jl                   |  2 +-
 src/standalone_query_macros.jl | 16 ++++++++++++
 test/test_standalone.jl        | 47 ++++++++++++++++++++++++++++++++++
 5 files changed, 77 insertions(+), 2 deletions(-)
 create mode 100644 test/test_standalone.jl

diff --git a/NEWS.md b/NEWS.md
index d9ccbc68..29a42326 100644
--- a/NEWS.md
+++ b/NEWS.md
@@ -1,3 +1,7 @@
+# Query.jl v0.9.0 Release Notes
+* Add @take and @drop standalone macros
+* Fix some test bugs
+
 # Query.jl v0.8.0 Release Notes
 * Add @groupjoin, @join and @mapmany standalone macros
 * Move backend code to QueryOperators.jl
diff --git a/docs/src/experimental.md b/docs/src/experimental.md
index 7341025c..590d6bfc 100644
--- a/docs/src/experimental.md
+++ b/docs/src/experimental.md
@@ -7,7 +7,7 @@ Query.jl. At the same time any feedback on these features would be
 especially welcome.
 
 The `@map`, `@filter`, `@groupby`, `@orderby` (and various variants),
-`@groupjoin`, `@join` and `@mapmany` commands can be used in standalone
+`@groupjoin`, `@join`, `@mapmany`, `@take` and `@drop` commands can be used in standalone
 versions. Those standalone versions are especially convenient in
 combination with the pipe syntax in julia. Here is an example that
 demonstrates their use:
@@ -124,6 +124,14 @@ be an anonymous function that takes one argument and returns a collection.
 `result_selector` must be an anonymous function that takes two arguments.
 It will be applied to each element of the intermediate collection.
 
+### The `@take` command
+
+The `@take` command has the form `@take(source, n)`. `source` can be any source that can be queried. `n` must be an integer, and it specifies how many elements from the beginning of the source should be kept.
+
+### The `@drop` command
+
+The `@drop` command has the form `@drop(source, n)`. `source` can be any source that can be queried. `n` must be an integer, and it specifies how many elements from the beginning of the source should be dropped from the results.
+
 ## The `..` syntax
 
 The syntax `a..b` is translated into `map(i->i.b, a)` in any query
diff --git a/src/Query.jl b/src/Query.jl
index b21ba6c2..b1455de6 100644
--- a/src/Query.jl
+++ b/src/Query.jl
@@ -24,7 +24,7 @@ import QueryOperators: QueryProvider
 export @from, @query, @count, Grouping, @NT
 
 export @map, @filter, @groupby, @orderby, @orderby_descending,
-	@thenby, @thenby_descending, @groupjoin, @join, @mapmany
+	@thenby, @thenby_descending, @groupjoin, @join, @mapmany, @take, @drop
 
 include("query_utils.jl")
 include("query_translation.jl")
diff --git a/src/standalone_query_macros.jl b/src/standalone_query_macros.jl
index 0cd64c93..69a76cc7 100644
--- a/src/standalone_query_macros.jl
+++ b/src/standalone_query_macros.jl
@@ -238,3 +238,19 @@ macro filter(f)
         helper_namedtuples_replacement |>
         helper_replace_field_extraction_syntax
 end
+
+macro take(source, n)
+    return :(QueryOperators.take(QueryOperators.query($(esc(source))), $(esc(n))))
+end
+
+macro take(n)
+    return :( i -> QueryOperators.take(QueryOperators.query(i), $(esc(n))))
+end
+
+macro drop(source, n)
+    return :(QueryOperators.drop(QueryOperators.query($(esc(source))), $(esc(n))))
+end
+
+macro drop(n)
+    return :( i -> QueryOperators.drop(QueryOperators.query(i), $(esc(n))))
+end
diff --git a/test/test_standalone.jl b/test/test_standalone.jl
new file mode 100644
index 00000000..e28e9488
--- /dev/null
+++ b/test/test_standalone.jl
@@ -0,0 +1,47 @@
+using Query
+using DataFrames
+using Base.Test
+
+@testset "Standalone Syntax" begin
+
+@testset "@take operator" begin
+    df = DataFrame(a=[1,2,3], b=[3.,2.,1.], c=["a", "b", "c"])
+
+    df2 = df |> @take(2) |> DataFrame
+
+    @test df2 isa DataFrame
+    @test size(df2) == (2,3)
+    @test df2[:a] == [1,2]
+    @test df2[:b] == [3.,2.]
+    @test df2[:c] == ["a", "b"]
+
+    df2 = DataFrame(@take(df, 2))
+
+    @test df2 isa DataFrame
+    @test size(df2) == (2,3)
+    @test df2[:a] == [1,2]
+    @test df2[:b] == [3.,2.]
+    @test df2[:c] == ["a", "b"]
+end
+
+@testset "@drop operator" begin
+    df = DataFrame(a=[1,2,3], b=[3.,2.,1.], c=["a", "b", "c"])
+
+    df2 = df |> @drop(1) |> DataFrame
+
+    @test df2 isa DataFrame
+    @test size(df2) == (2,3)
+    @test df2[:a] == [2,3]
+    @test df2[:b] == [2.,1.]
+    @test df2[:c] == ["b","c"]
+
+    df2 = DataFrame(@drop(df, 1))
+
+    @test df2 isa DataFrame
+    @test size(df2) == (2,3)
+    @test df2[:a] == [2,3]
+    @test df2[:b] == [2.,1.]
+    @test df2[:c] == ["b","c"]
+end
+
+end