Make it work with contextmanager

safety_prototype/application.py

@@ -1,7 +1,7 @@
 # Only API objects are imported because the user only calls API functions or
 # methods.
 from opentelemetry.configuration import set_sdk
-from opentelemetry.trace import function, Class0, Class1
+from opentelemetry.trace import function, Class0, Class1, Class2
 
 print(function(4, 2))
 # This is the function that sets the SDK. After this is set, any call to an API
@@ -43,3 +43,11 @@
 
 print(class0.method_0(1, 2))
 print(class0.method_0(1, 0))
+
+class2 = Class2()
+
+# This call returns an instance of a the safe Class0 class. This is necessary
+# because the user must only handle safe objects.
+with class2.method_0(7) as class0:
+    print(class0.method_0(2, 2))
+    print(class0.method_0(2, 0))

safety_prototype/opentelemetry-api/src/opentelemetry/trace/__init__.py

@@ -20,9 +20,10 @@
 in this module act as proxies and (with the SDK setting function) are the only
 objects the user has contact with.
 """
+from contextlib import contextmanager
 
 from opentelemetry._safety import _safe_function
-from opentelemetry.trace.api import Class0, Class1
+from opentelemetry.trace.api import Class0, Class1, Class2
 from opentelemetry.configuration import _get_sdk_module
 
 
@@ -45,3 +46,14 @@ def method_0(self, a: int) -> Class0:
         safe_instance = Class0(0)
         safe_instance._sdk_instance = self._sdk_instance.method_0(a)
         return safe_instance
+
+
+class Class2(Class2):
+
+    @contextmanager
+    @_safe_function(Class0(0))
+    def method_0(self, a: int) -> Class0:
+        safe_instance = Class0(0)
+        with self._sdk_instance.method_0(a) as sdk_instance:
+            safe_instance._sdk_instance = sdk_instance
+            yield safe_instance

safety_prototype/opentelemetry-api/src/opentelemetry/trace/api.py

@@ -23,6 +23,7 @@
 """
 
 from abc import ABC, abstractmethod
+from contextlib import contextmanager
 
 from opentelemetry.configuration import _get_sdk_module
 
@@ -61,7 +62,7 @@ def __getattribute__(self, name):
 
 class Class0(_BaseAPI):
 
-    def __init__(self, a: int) -> None:
+    def __init__(self, a):
         super().__init__(a)
 
     @abstractmethod
@@ -74,3 +75,11 @@ class Class1(_BaseAPI):
     @abstractmethod
     def method_0(self, a: int) -> Class0:
         pass
+
+
+class Class2(_BaseAPI):
+
+    @contextmanager
+    @abstractmethod
+    def method_0(self, a: int) -> Class0:
+        pass

safety_prototype/opentelemetry-sdk/src/opentelemetry/sdk/trace/__init__.py

@@ -21,7 +21,9 @@
 be caught by the safety mechanism in the API.
 """
 
-from opentelemetry.trace.api import Class0, Class1
+from contextlib import contextmanager
+
+from opentelemetry.trace.api import Class0, Class1, Class2
 
 
 def function(a: int, b: int) -> float:
@@ -42,3 +44,12 @@ class Class1(Class1):
 
     def method_0(self, a: int) -> Class0:
         return Class0(a)
+
+
+class Class2(Class2):
+
+    @contextmanager
+    def method_0(self, a: int) -> Class0:
+        print("before")
+        yield Class0(a)
+        print("after")