Use boolean values for annotation state

Store whether we are annotating as True or False, and use the stop_annotation decorator to create a stack of values. This is better than the previous integer. --------- Co-authored-by: David Ham <[email protected]>
dolfin-adjoint · Sep 28, 2023 · cefba3f · cefba3f
1 parent d532d41
commit cefba3f
Showing 1 changed file with 13 additions and 9 deletions.
diff --git a/pyadjoint/tape.py b/pyadjoint/tape.py
@@ -9,22 +9,22 @@
 
 
 _working_tape = None
-_stop_annotating = 0
+_annotation_enabled = False
 
 
 def get_working_tape():
     return _working_tape
 
 
 def pause_annotation():
-    global _stop_annotating
-    _stop_annotating += 1
+    global _annotation_enabled
+    _annotation_enabled = False
 
 
 def continue_annotation():
-    global _stop_annotating
-    _stop_annotating -= 1
-    return _stop_annotating <= 0
+    global _annotation_enabled
+    _annotation_enabled = True
+    return _annotation_enabled
 
 
 class set_working_tape(object):
@@ -81,13 +81,17 @@ class stop_annotating(object):
     modified variables at the end of the context manager. """
 
     def __init__(self, modifies=None):
+        global _annotation_enabled
         self.modifies = modifies
+        self._orig_annotation_enabled = _annotation_enabled
 
     def __enter__(self):
-        pause_annotation()
+        global _annotation_enabled
+        _annotation_enabled = False
 
     def __exit__(self, *args):
-        continue_annotation()
+        global _annotation_enabled
+        _annotation_enabled = self._orig_annotation_enabled
         if self.modifies is not None:
             try:
                 self.modifies.create_block_variable()
@@ -125,7 +129,7 @@ def annotate_tape(kwargs=None):
 
     # TODO: Consider if there is any scenario where one would want the keyword to have
     # precedence over the global flag.
-    if _stop_annotating > 0:
+    if not _annotation_enabled:
         return False
 
     return annotate