Merge pull request #923 from slwu89/queries-fix

Fix for queries with attributes using union types

src/wiring_diagrams/Algebras.jl

@@ -174,7 +174,16 @@ function query(X::ACSet, diagram::UndirectedWiringDiagram,
       junction = key isa Integer ? key : only(incident(diagram, key, :variable))
       add_part!(diagram, :Port, port_name=:_value,
                 box=box, junction=junction)
-      SMultispan{1}(ConstantFunction(value, FinSet(1)))
+
+      # Handle possibility of unions in attribute types.
+      name = diagram[first(incident(diagram, junction, :junction)), :port_name]
+      constant = if name ∈ attrs(acset_schema(X), just_names=true)
+        ConstantFunction(value, FinSet(1), TypeSet(subpart_type(X, name)))
+      else
+        ConstantFunction(value, FinSet(1))
+      end
+
+      SMultispan{1}(constant)
     end)
   end
 

test/wiring_diagrams/Algebras.jl

@@ -111,4 +111,35 @@ result = query(g, cycles3, (v=1,))
 @test result == DataFrame(edge1=[3], edge2=[1], edge3=[2])
 @test isempty(query(cycle_graph(Graph, 4), cycles3))
 
+# queries with attributes
+@present TestSch(FreeSchema) begin
+    X::Ob
+    A::AttrType
+    a::Attr(X,A)
+end
+
+@acset_type TestData(TestSch)
+
+data = @acset TestData{Union{Int,Symbol}} begin
+    X=5
+    a=[1,2,3,:four,:five]
+end
+
+testquery = @relation (xout=xid, aout=attr) begin
+    X(_id=xid, a=attr)
+end
+
+result = query(data, testquery)
+
+@test result[!,1] == 1:5
+@test result[!,2] == [1,2,3,:four,:five]
+
+result = query(data, testquery, (attr=3, ))
+@test result[!,1] == [3]
+@test result[!,2] == [3]
+
+result = query(data, testquery, (attr=:five, ))
+@test result[!,1] == [5]
+@test result[!,2] == [:five]
+
 end