Support typed array literals with atsign-data and atsign-pdata

JuliaStats · Jul 2, 2014 · 140a526 · 140a526
1 parent 309b1ca
commit 140a526
Show file tree

Hide file tree

Showing 2 changed files with 69 additions and 13 deletions.
diff --git a/src/literals.jl b/src/literals.jl
@@ -45,37 +45,54 @@ function findstub_matrix(ex::Expr)
 end
 
 function parsevector(ex::Expr)
-    stub = findstub_vector(ex)
-    data, na = fixargs(ex.args, stub)
-    return Expr(ex.head, data...), Expr(ex.head, na...)
+    if ex.head in (:ref, :typed_hcat, :typed_vcat)
+        data, na = fixargs(ex.args[2:end], :(zero($(ex.args[1]))))
+        return Expr(ex.head, ex.args[1], data...),
+               Expr(ex.head == :typed_hcat ? :hcat : :vcat, na...)
+    else
+        stub = findstub_vector(ex)
+        data, na = fixargs(ex.args, stub)
+        return Expr(ex.head, data...), Expr(ex.head, na...)
+    end
 end
 
 function parsematrix(ex::Expr)
-    stub = findstub_matrix(ex)
-    nrows = length(ex.args)
+    if ex.head == :typed_vcat
+        stub = :(zero($(ex.args[1])))
+        rows = 2:length(ex.args)
+    else
+        stub = findstub_matrix(ex)
+        rows = 1:length(ex.args)
+    end
+
+    nrows = length(rows)
     datarows = Array(Expr, nrows)
     narows = Array(Expr, nrows)
-    for row in 1:nrows
-        data, na = fixargs(ex.args[row].args, stub)
-        datarows[row] = Expr(:row, data...)
-        narows[row] = Expr(:row, na...)
+    for irow in 1:nrows
+        data, na = fixargs(ex.args[rows[irow]].args, stub)
+        datarows[irow] = Expr(:row, data...)
+        narows[irow] = Expr(:row, na...)
+    end
+    if ex.head == :typed_vcat
+        return Expr(:typed_vcat, ex.args[1], datarows...), Expr(:vcat, narows...)
+    else
+        return Expr(:vcat, datarows...), Expr(:vcat, narows...)
     end
-    return Expr(:vcat, datarows...), Expr(:vcat, narows...)
 end
 
 function parsedata(ex::Expr)
     if length(ex.args) == 0
         return :([]), Expr(:call, :Array, :Bool, 0)
     end
-    if isa(ex.args[1], Expr) && ex.args[1].head == :row
+    if ex.head == :typed_vcat || (isa(ex.args[1], Expr) && ex.args[1].head == :row)
         return parsematrix(ex)
     else
         return parsevector(ex)
     end
 end
 
 macro data(ex)
-    if ex.head != :vcat && ex.head != :hcat
+    if !(ex.head in (:vcat, :hcat, :ref, :typed_vcat, :typed_hcat))
         return quote
             tmp = $(esc(ex))
             DataArray(tmp, bitbroadcast(x->isequal(x, NA), tmp))
@@ -86,7 +103,7 @@ macro data(ex)
 end
 
 macro pdata(ex)
-    if ex.head != :vcat && ex.head != :hcat
+    if !(ex.head in (:vcat, :hcat, :ref, :typed_vcat, :typed_hcat))
         return quote
             tmp = $(esc(ex))
             PooledDataArray(tmp, bitbroadcast(x->isequal(x, NA), tmp))

diff --git a/test/literals.jl b/test/literals.jl
@@ -6,11 +6,24 @@ module TestLiterals
     @test isequal(dv,
                   DataArray([1, 0, 3],
                             [false, true, false]))
+
     dv = @data [1 NA 3]
     @test isequal(dv,
                   DataArray([1 0 3],
                             [false true false]))
 
+    dv = @data Float64[1, NA, 3]
+    @test isequal(dv,
+                  DataArray(Float64[1, 0, 3],
+                            [false, true, false]))
+    @test typeof(dv) == DataVector{Float64}
+
+    dv = @data Float64[1 NA 3]
+    @test isequal(dv,
+                  DataArray(Float64[1 0 3],
+                            [false true false]))
+    @test typeof(dv) == DataMatrix{Float64}
+
     dv = @data {1, NA, 3}
     @test isequal(dv,
                   DataArray({1, 0, 3},
@@ -20,6 +33,12 @@ module TestLiterals
     @test isequal(dm,
                   DataArray([1 0; 3 4],
                             [false true; false false]))
+
+    dm = @data Float64[1 NA; 3 4]
+    @test isequal(dm,
+                  DataArray(Float64[1 0; 3 4],
+                            [false true; false false]))
+    @test typeof(dm) == DataMatrix{Float64}
 
     dm = @data {1 NA; 3 4}
     @test isequal(dm,
@@ -36,6 +55,13 @@ module TestLiterals
     @test isequal(pdv,
                   PooledDataArray([1, 0, 3],
                                   [false, true, false]))
+
+    pdv = @pdata Float64[1, NA, 3]
+    @test isequal(pdv,
+                  PooledDataArray(Float64[1, 0, 3],
+                                  [false, true, false]))
+    @test typeof(pdv) == PooledDataArray{Float64,Uint32,1}
+
     pdv = @pdata {1, NA, 3}
     @test isequal(pdv,
                   PooledDataArray({1, 0, 3},
@@ -46,10 +72,23 @@ module TestLiterals
                   PooledDataArray([1 0 3],
                                   [false true false]))
 
+    pdv = @pdata Float64[1 NA 3]
+    @test isequal(pdv,
+                  PooledDataArray(Float64[1 0 3],
+                                  [false true false]))
+    @test typeof(pdv) == PooledDataArray{Float64,Uint32,2}
+
     pdm = @pdata [1 NA; 3 4]
     @test isequal(pdm,
                   PooledDataArray([1 0; 3 4],
                                   [false true; false false]))
+
+    pdm = @pdata Float64[1 NA; 3 4]
+    @test isequal(pdm,
+                  PooledDataArray(Float64[1 0; 3 4],
+                                  [false true; false false]))
+    @test typeof(pdm) == PooledDataArray{Float64,Uint32,2}
+
     pdm = @pdata {1 NA; 3 4}
     @test isequal(pdm,
                   PooledDataArray({1 0; 3 4},