From 1da85b8046f2c59a4d3e4f0fb94a6fd6ff35dcd0 Mon Sep 17 00:00:00 2001
From: Evan Patterson <evan@epatters.org>
Date: Tue, 20 Oct 2020 14:58:37 -0700
Subject: [PATCH] ENH: Get generators from presentations by indexing.

Closes #301.
---
 src/core/Present.jl                            |  4 ++++
 test/categorical_algebra/CSetDataStructures.jl |  2 +-
 test/core/Present.jl                           | 16 +++++++++-------
 3 files changed, 14 insertions(+), 8 deletions(-)

diff --git a/src/core/Present.jl b/src/core/Present.jl
index 46b84dd89..bc297a7b0 100644
--- a/src/core/Present.jl
+++ b/src/core/Present.jl
@@ -59,11 +59,15 @@ generators(pres::Presentation, type::Symbol) = pres.generators[type]
 generators(pres::Presentation, type::Type) = generators(pres, nameof(type))
 
 """ Retrieve generators by name.
+
+Generators can also be retrieved using indexing notation, so that
+`generator(pres, name)` and `pres[name]` are equivalent.
 """
 function generator(pres::Presentation, name)
   type, index = pres.generator_name_index[name]
   pres.generators[type][index]
 end
+Base.getindex(pres::Presentation, name) = generator.(Ref(pres), name)
 
 """ Does the presentation contain a generator with the given name?
 """
diff --git a/test/categorical_algebra/CSetDataStructures.jl b/test/categorical_algebra/CSetDataStructures.jl
index 68e19dc25..f8ad589a0 100644
--- a/test/categorical_algebra/CSetDataStructures.jl
+++ b/test/categorical_algebra/CSetDataStructures.jl
@@ -152,7 +152,7 @@ set_subpart!(d, [4,5], :parent, 5)
 @test incident(d, 5, [:parent, :parent]) == [1,2,3,4,5]
 @test incident(d, 10, [:parent, :height]) == [1,2,3]
 
-X, parent, height = generator.(Ref(TheoryDendrogram), [:X, :parent, :height])
+X, parent, height = TheoryDendrogram[[:X, :parent, :height]]
 @test subpart(d, 3, parent) == 4
 @test subpart(d, 3, compose(parent, height)) == 10
 @test subpart(d, 3, id(X)) == 3
diff --git a/test/core/Present.jl b/test/core/Present.jl
index 927d595f2..762545012 100644
--- a/test/core/Present.jl
+++ b/test/core/Present.jl
@@ -17,19 +17,21 @@ pres = Presentation(FreeCategory)
 @test presentation_theory(pres) == Category
 @test !has_generator(pres, :A)
 add_generator!(pres, A)
-@test generators(pres) == [ A ]
+@test generators(pres) == [A]
 @test generator(pres, :A) == A
 @test has_generator(pres, :A)
 add_generator!(pres, B)
-@test generators(pres) == [ A, B ]
+@test generators(pres) == [A, B]
 @test_throws Exception add_generator!(pres, A)
+@test pres[:A] == A
+@test pres[[:A,:B]] == [A, B]
 
 add_generators!(pres, (f,g))
-@test generators(pres) == [ A, B, f, g ]
-@test generators(pres, :Ob) == [ A, B ]
-@test generators(pres, :Hom) == [ f, g ]
-@test generators(pres, FreeCategory.Ob) == [ A, B ]
-@test generators(pres, FreeCategory.Hom) == [ f, g ]
+@test generators(pres) == [A, B, f, g]
+@test generators(pres, :Ob) == [A, B]
+@test generators(pres, :Hom) == [f, g]
+@test generators(pres, FreeCategory.Ob) == [A, B]
+@test generators(pres, FreeCategory.Hom) == [f, g]
 
 # Presentation macro
 ####################