[NDTensors] Add with_auto_fermion scope wrapper (#1726)

Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: mtfishman

NDTensors/Project.toml

@@ -1,6 +1,6 @@
 name = "NDTensors"
 uuid = "23ae76d9-e61a-49c4-8f12-3f1a16adf9cf"
-version = "0.4.24"
+version = "0.4.25"
 authors = ["Matthew Fishman <mfishman@flatironinstitute.org>"]
 
 [workspace]

NDTensors/src/NDTensors.jl

@@ -221,6 +221,41 @@ function disable_auto_fermion()
     return nothing
 end
 
+"""
+    NDTensors.with_auto_fermion(f, enable::Bool = true)
+
+Run `f` with the auto-fermion system enabled or disabled (enabled by default),
+restoring the previous state of the auto-fermion system after `f` finishes.
+
+This is the preferred way to temporarily enable or disable the auto-fermion
+system for a code block, because the previous state is restored even if `f`
+throws an exception (unlike a manual call to [`enable_auto_fermion`](@ref) or
+[`disable_auto_fermion`](@ref)).
+
+# Examples
+
+```julia
+using NDTensors: NDTensors
+
+NDTensors.with_auto_fermion() do
+    # code that requires the auto-fermion system to be enabled
+end
+
+NDTensors.with_auto_fermion(false) do
+    # code that requires the auto-fermion system to be disabled
+end
+```
+"""
+function with_auto_fermion(f, enable::Bool = true)
+    previous = using_auto_fermion()
+    enable ? enable_auto_fermion() : disable_auto_fermion()
+    return try
+        f()
+    finally
+        previous ? enable_auto_fermion() : disable_auto_fermion()
+    end
+end
+
 #####################################
 # Optional backends
 #

NDTensors/test/test_with_auto_fermion.jl

@@ -0,0 +1,73 @@
+@eval module $(gensym())
+using NDTensors: NDTensors
+using Test: @test, @test_throws, @testset
+
+@testset "NDTensors.with_auto_fermion" begin
+    @testset "starting state $start" for start in (false, true)
+        start ? NDTensors.enable_auto_fermion() : NDTensors.disable_auto_fermion()
+        @test NDTensors.using_auto_fermion() == start
+
+        @testset "default enables auto-fermion" begin
+            inside = Ref(false)
+            NDTensors.with_auto_fermion() do
+                return inside[] = NDTensors.using_auto_fermion()
+            end
+            @test inside[] == true
+            @test NDTensors.using_auto_fermion() == start
+        end
+
+        @testset "explicit enable=false disables auto-fermion" begin
+            inside = Ref(true)
+            NDTensors.with_auto_fermion(false) do
+                return inside[] = NDTensors.using_auto_fermion()
+            end
+            @test inside[] == false
+            @test NDTensors.using_auto_fermion() == start
+        end
+
+        @testset "return value is propagated" begin
+            @test NDTensors.with_auto_fermion(() -> 42) == 42
+            @test NDTensors.with_auto_fermion(() -> 42, false) == 42
+            @test NDTensors.using_auto_fermion() == start
+        end
+
+        @testset "restores previous state after exception" begin
+            @test_throws ErrorException NDTensors.with_auto_fermion() do
+                return error("boom")
+            end
+            @test NDTensors.using_auto_fermion() == start
+            @test_throws ErrorException NDTensors.with_auto_fermion(false) do
+                return error("boom")
+            end
+            @test NDTensors.using_auto_fermion() == start
+        end
+
+        @testset "nested scopes" begin
+            inner_state = Ref{Bool}(start)
+            NDTensors.with_auto_fermion(true) do
+                @test NDTensors.using_auto_fermion() == true
+                NDTensors.with_auto_fermion(false) do
+                    return inner_state[] = NDTensors.using_auto_fermion()
+                end
+                @test NDTensors.using_auto_fermion() == true
+            end
+            @test inner_state[] == false
+            @test NDTensors.using_auto_fermion() == start
+        end
+
+        @testset "nested scope restores after inner exception" begin
+            NDTensors.with_auto_fermion(true) do
+                @test NDTensors.using_auto_fermion() == true
+                @test_throws ErrorException NDTensors.with_auto_fermion(false) do
+                    return error("boom")
+                end
+                @test NDTensors.using_auto_fermion() == true
+            end
+            @test NDTensors.using_auto_fermion() == start
+        end
+    end
+
+    # Leave the auto-fermion state as we found it for the rest of the suite.
+    NDTensors.disable_auto_fermion()
+end
+end