From d0d19ae55af2e7b120f9e7098871976dfd5ca66e Mon Sep 17 00:00:00 2001
From: Fredrik Ekre <ekrefredrik@gmail.com>
Date: Wed, 29 Nov 2023 10:49:47 +0100
Subject: [PATCH] Overload `Base.getindex(::Family, labels)` (#13)

This patch adds an overload to `Base.getindex` for the `Family`
collector to have the same meaning as `Prometheus.labels`.
`family[labels]` is equivalent to `Prometheus.labels(family, labels)`
---
 CHANGELOG.md      |  7 ++++++-
 src/Prometheus.jl | 10 ++++++++++
 test/runtests.jl  | 17 ++++++++++++-----
 3 files changed, 28 insertions(+), 6 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 46d7c61..e5bdb03 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -5,7 +5,11 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
-<!-- ## [Unreleased] -->
+## [Unreleased]
+### Added
+ - `Base.getindex` is overloaded for the `Prometheus.Family` collector to have the same
+   meaning as `Prometheus.labels`. `family[labels]` is equivalent to
+   `Prometheus.labels(family, labels)`. ([#13][github-13])
 
 ## [1.2.0] - 2023-11-22
 ### Added
@@ -52,6 +56,7 @@ See [README.md](README.md) for details and documentation.
 [github-6]: https://github.com/fredrikekre/Prometheus.jl/pull/6
 [github-7]: https://github.com/fredrikekre/Prometheus.jl/pull/7
 [github-10]: https://github.com/fredrikekre/Prometheus.jl/pull/10
+[github-13]: https://github.com/fredrikekre/Prometheus.jl/pull/13
 
 [Unreleased]: https://github.com/fredrikekre/Prometheus.jl/compare/v1.2.0...HEAD
 [1.2.0]: https://github.com/fredrikekre/Prometheus.jl/compare/v1.1.0...v1.2.0
diff --git a/src/Prometheus.jl b/src/Prometheus.jl
index 7fbd0df..4b9db05 100644
--- a/src/Prometheus.jl
+++ b/src/Prometheus.jl
@@ -821,6 +821,11 @@ the following:
 
 All non-string values (e.g. `200` in the examples above) are stringified using `string`.
 
+!!! tip
+    `Base.getindex` is overloaded to have the meaning of `Prometheus.labels` for the family
+    collector: `family[label_values]` is equivalent to
+    `Prometheus.labels(family, label_values)`.
+
 !!! note
     This method does an acquire/release of a lock, and a dictionary lookup, to find the
     collector matching the label names. For typical applications this overhead does not
@@ -835,6 +840,11 @@ function labels(family::Family{C, N}, label_values) where {C, N}
     return collector
 end
 
+# Support family[labels] as a cute way of extracting the collector
+function Base.getindex(family::Family, label_values)
+    return labels(family, label_values)
+end
+
 """
     Prometheus.remove(family::Family, label_values)
 
diff --git a/test/runtests.jl b/test/runtests.jl
index 3b69c06..6762d7e 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -281,8 +281,8 @@ end
     # Prometheus.labels(...), Prometheus.remove(...), Prometheus.clear()
     l1 = ("/foo/", "200")
     l2 = ("/bar/", "404")
-    @test Prometheus.labels(c, l1) === Prometheus.labels(c, l1)
-    @test Prometheus.labels(c, l2) === Prometheus.labels(c, l2)
+    @test Prometheus.labels(c, l1) === Prometheus.labels(c, l1) === c[l1]
+    @test Prometheus.labels(c, l2) === Prometheus.labels(c, l2) === c[l2]
     @test length(c.children) == 2
     @test Prometheus.labels(c, l1).value == 0
     @test Prometheus.labels(c, l2).value == 0
@@ -449,8 +449,8 @@ end
     # Prometheus.inc(...)
     l1 = ("/foo/", "200")
     l2 = ("/bar/", "404")
-    @test Prometheus.labels(c, l1) === Prometheus.labels(c, l1)
-    @test Prometheus.labels(c, l2) === Prometheus.labels(c, l2)
+    @test Prometheus.labels(c, l1) === Prometheus.labels(c, l1) === c[l1]
+    @test Prometheus.labels(c, l2) === Prometheus.labels(c, l2) === c[l2]
     @test length(c.children) == 2
     @test Prometheus.labels(c, l1)._count == 0
     @test Prometheus.labels(c, l1)._sum == 0
@@ -549,12 +549,19 @@ end
             ),
         )
         @test Prometheus.labels(fam, ("/api", "200")) ===
+              fam[("/api", "200")] ===
               Prometheus.labels(fam, ("/api", 200)) ===
+              fam[("/api", 200)] ===
               Prometheus.labels(fam, (target="/api", status_code="200")) ===
+              fam[(target="/api", status_code="200")] ===
               Prometheus.labels(fam, (target="/api", status_code=200)) ===
+              fam[(target="/api", status_code=200)] ===
               Prometheus.labels(fam, (status_code="200", target="/api")) ===
+              fam[(status_code="200", target="/api")] ===
               Prometheus.labels(fam, (status_code=200, target="/api")) ===
-              Prometheus.labels(fam, RequestLabels("/api", 200))
+              fam[(status_code=200, target="/api")] ===
+              Prometheus.labels(fam, RequestLabels("/api", 200)) ===
+              fam[RequestLabels("/api", 200)]
     end
 end