From 4ca271cba3543e6074c30670fe46775da450cc53 Mon Sep 17 00:00:00 2001
From: Jan Weidner <jw3126@gmail.com>
Date: Sat, 21 Oct 2023 07:09:38 +0200
Subject: [PATCH] add Profile.Allocs.@pprof

---
 src/Allocs.jl  | 19 +++++++++++++++++++
 src/PProf.jl   |  2 +-
 test/Allocs.jl | 17 +++++++++++++++++
 3 files changed, 37 insertions(+), 1 deletion(-)

diff --git a/src/Allocs.jl b/src/Allocs.jl
index 2105420..d465a67 100644
--- a/src/Allocs.jl
+++ b/src/Allocs.jl
@@ -223,4 +223,23 @@ function pprof(alloc_profile::Profile.Allocs.AllocResults = Profile.Allocs.fetch
     out
 end
 
+"""
+    Allocs.@pprof args...
+
+Profiles the expression using `Allocs.@profile` and starts or restarts the `Allocs.pprof()` web UI with
+default arguments. See also [`PProf.@pprof`](@ref).
+
+# Examples
+```julia
+PProf.Allocs.@pprof [randn(3) for _ in 1:100000]
+PProf.Allocs.@pprof sample_rate=1 randn(100)
+```
+"""
+macro pprof(args...)
+    esc(quote
+        $Profile.Allocs.@profile $(args...)
+        $(@__MODULE__).pprof()
+    end)
+end
+
 end  # module Allocs
diff --git a/src/PProf.jl b/src/PProf.jl
index f6092a8..73feacf 100644
--- a/src/PProf.jl
+++ b/src/PProf.jl
@@ -398,7 +398,7 @@ end
     @pprof ex
 
 Profiles the expression using `@profile` and starts or restarts the `pprof()` web UI with
-default arguments.
+default arguments. See also [`PProf.Allocs.@pprof`](@ref).
 """
 macro pprof(ex)
     esc(quote
diff --git a/test/Allocs.jl b/test/Allocs.jl
index bf1a9ee..1c3a6e1 100644
--- a/test/Allocs.jl
+++ b/test/Allocs.jl
@@ -9,6 +9,23 @@ using Test
 
 const out = tempname()
 
+@testset "PProf.Allocs.@pprof" begin
+    Profile.Allocs.clear()
+    rm("alloc-profile.pb.gz", force=true)
+    @assert !isfile("alloc-profile.pb.gz")
+    PProf.Allocs.@pprof sample_rate=1.0 randn(100)
+    @test isfile("alloc-profile.pb.gz")
+
+    Profile.Allocs.clear()
+    rm("alloc-profile.pb.gz", force=true)
+    @assert !isfile("alloc-profile.pb.gz")
+    PProf.Allocs.@pprof randn(100)
+    @test isfile("alloc-profile.pb.gz")
+
+    rm("alloc-profile.pb.gz", force=true)
+end
+
+
 @testset "basic profiling" begin
     Profile.Allocs.clear()
     Profile.Allocs.@profile sample_rate=1.0 begin