using instrumenters as decorators

Author: AndreasGocht

scorep/instrumenter.py

@@ -2,6 +2,7 @@
 import os
 import platform
 import sys
+import functools
 
 global_instrumenter = None
 
@@ -74,8 +75,30 @@ class enable():
     If a region name is given, the region the contextmanager is active will be marked in the trace or profile
     """
 
-    def __init__(self, region_name=None):
+    def __init__(self, region_name=""):
         self.region_name = region_name
+        if region_name == "":
+            self.user_region_name = False
+        else:
+            self.user_region_name = True
+        self.module_name = ""
+        self.func = None
+
+    def _recreate_cm(self):
+        return self
+
+    def __call__(self, func):
+        get_instrumenter().unregister()
+        try:
+            self.func = func
+
+            @functools.wraps(func)
+            def inner(*args, **kwds):
+                with self._recreate_cm():
+                    return func(*args, **kwds)
+        finally:
+            get_instrumenter().register()
+        return inner
 
     def __enter__(self):
         self.tracer_registered = get_instrumenter().get_registered()
@@ -94,6 +117,22 @@ def __enter__(self):
                     self.module_name, self.region_name, full_file_name,
                     line_number)
 
+            elif callable(self.func):
+                # The user did not specify a region name, and it's a callable, so it's a semi instrumented region
+                self.region_name = self.func.__name__
+                self.module_name = self.func.__module__
+                self.code_obj = self.func.__code__
+                file_name = self.func.__code__.co_filename
+                line_number = self.func.__code__.co_firstlineno
+
+                if file_name is not None:
+                    full_file_name = os.path.abspath(file_name)
+                else:
+                    full_file_name = "None"
+
+                scorep.instrumenter.get_instrumenter().region_begin(
+                    self.module_name, self.region_name, full_file_name, line_number, self.code_obj)
+
             get_instrumenter().register()
 
     def __exit__(self, exc_type, exc_value, traceback):
@@ -103,6 +142,9 @@ def __exit__(self, exc_type, exc_value, traceback):
             if self.region_name is not None:
                 get_instrumenter().region_end(
                     self.module_name, self.region_name)
+            elif callable(self.func):
+                get_instrumenter().region_end(
+                    self.module_name, self.region_name, self.code_obj)
 
 
 class disable():
@@ -116,8 +158,30 @@ class disable():
     If a region name is given, the region the contextmanager is active will be marked in the trace or profile
     """
 
-    def __init__(self, region_name=None):
+    def __init__(self, region_name=""):
         self.region_name = region_name
+        if region_name == "":
+            self.user_region_name = False
+        else:
+            self.user_region_name = True
+        self.module_name = ""
+        self.func = None
+
+    def _recreate_cm(self):
+        return self
+
+    def __call__(self, func):
+        get_instrumenter().unregister()
+        try:
+            self.func = func
+
+            @functools.wraps(func)
+            def inner(*args, **kwds):
+                with self._recreate_cm():
+                    return func(*args, **kwds)
+        finally:
+            get_instrumenter().register()
+        return inner
 
     def __enter__(self):
         self.tracer_registered = get_instrumenter().get_registered()
@@ -137,11 +201,29 @@ def __enter__(self):
                 get_instrumenter().region_begin(
                     self.module_name, self.region_name, full_file_name,
                     line_number)
+            elif callable(self.func):
+                # The user did not specify a region name, and it's a callable, so it's a semi instrumented region
+                self.region_name = self.func.__name__
+                self.module_name = self.func.__module__
+                self.code_obj = self.func.__code__
+                file_name = self.func.__code__.co_filename
+                line_number = self.func.__code__.co_firstlineno
+
+                if file_name is not None:
+                    full_file_name = os.path.abspath(file_name)
+                else:
+                    full_file_name = "None"
+
+                scorep.instrumenter.get_instrumenter().region_begin(
+                    self.module_name, self.region_name, full_file_name, line_number, self.code_obj)
 
     def __exit__(self, exc_type, exc_value, traceback):
         if self.tracer_registered:
             if self.region_name is not None:
                 get_instrumenter().region_end(
                     self.module_name, self.region_name)
+            elif callable(self.func):
+                get_instrumenter().region_end(
+                    self.module_name, self.region_name, self.code_obj)
 
             get_instrumenter().register()

test/cases/user_instrumentation.py

@@ -4,9 +4,16 @@
 
 def foo():
     print("hello world")
-    instrumentation2.baz()
     instrumentation2.bar()
 
 
+@scorep.instrumenter.enable()
+def foo2():
+    print("hello world2")
+    instrumentation2.baz()
+
+
 with scorep.instrumenter.enable():
     foo()
+
+foo2()

test/test_scorep.py

@@ -254,13 +254,15 @@ def test_user_instrumentation(scorep_env, instrumenter):
     )
 
     assert std_err == ""
-    assert std_out == "hello world\nbaz\nbar\n"
+    assert std_out == "hello world\nbar\nhello world2\nbaz"
 
     std_out, std_err = call(["otf2-print", trace_path])
 
     assert std_err == ""
     assert re.search('ENTER[ ]*[0-9 ]*[0-9 ]*Region: "__main__:foo"', std_out)
     assert re.search('LEAVE[ ]*[0-9 ]*[0-9 ]*Region: "__main__:foo"', std_out)
+    assert re.search('ENTER[ ]*[0-9 ]*[0-9 ]*Region: "__main__:foo2"', std_out)
+    assert re.search('LEAVE[ ]*[0-9 ]*[0-9 ]*Region: "__main__:foo2"', std_out)
 
 
 @foreach_instrumenter