1 /* Copyright 2020 The TensorFlow Authors. All Rights Reserved. 2 3 Licensed under the Apache License, Version 2.0 (the "License"); 4 you may not use this file except in compliance with the License. 5 You may obtain a copy of the License at 6 7 http://www.apache.org/licenses/LICENSE-2.0 8 9 Unless required by applicable law or agreed to in writing, software 10 distributed under the License is distributed on an "AS IS" BASIS, 11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 See the License for the specific language governing permissions and 13 limitations under the License. 14 ==============================================================================*/ 15 #ifndef TENSORFLOW_PYTHON_FRAMEWORK_PY_CONTEXT_MANAGER_H_ 16 #define TENSORFLOW_PYTHON_FRAMEWORK_PY_CONTEXT_MANAGER_H_ 17 18 #include <Python.h> 19 20 #include <string> 21 22 #include "tensorflow/python/lib/core/safe_pyobject_ptr.h" 23 24 namespace tensorflow { 25 26 // Class that wraps a Python context manager, and calls the `__enter__` and 27 // `__exit__` methods at appropriate times: 28 // 29 // * When `PyContextManager::Enter(cm)` is called, the context manager `cm` 30 // is stored, and `cm.__enter__` is called. The result can be retrieved 31 // with `PyContextManager::var()`. 32 // * When the `PyContextManager` is destroyed, then `cm.__exit__` is called 33 // (with information about any active exception). 34 // * `PyContextManager::Enter(cm)` may be called at most once. If 35 // `PyContextManager::Enter()` is never called, then the destructor is a 36 // no-op (i.e., `__exit__` is not called). 37 // 38 // PyContextManager places two restrictons on the wrapped context managers: 39 // 40 // 1. The context manager may not suppress exceptions -- i.e., `__exit__` 41 // may not return a True value. If it does, then a new exception will be 42 // set, indicating that this is unuspported. 43 // 2. The context manager may not raise an exception from `__exit__` if the 44 // an exception is not active when it is called. If it does, then an error 45 // message will be logged, indicating that this is unsupported, and the 46 // exception will be suppressed. 47 // 48 // These restrictions are both intended to ensure that the state of 49 // PyErr_Occurred is unchanged by PyContextManager's destructor. This is 50 // important, because changing the state of PyErr_Occurred in the destructor 51 // would mean that we are returning a nullptr with no exception set, or 52 // returning a non-null value with an exception set (both of which are invalid). 53 class PyContextManager { 54 public: 55 // Calls `py_context_manager.__enter__()`, and stores the result in `var`. 56 // Return true if `__enter__` succeeds, or false if `__enter__` raises an 57 // exception. (Also returns false if `py_context_manager` is nullptr.) 58 // 59 // Steals a reference to `py_context_manager`. (This reference is deleted 60 // when the destructor is called.) 61 bool Enter(PyObject* py_context_manager); 62 63 // Calls `py_context_manager.__exit__()`. 64 ~PyContextManager(); 65 66 // Returns the variable returned by `context_manager.__enter__()`. 67 // (This is the `var` bound by `with context_manager as var`.) 68 // Returns a borrowed reference. var()69 PyObject* var() { return var_.get(); } 70 71 protected: 72 Safe_PyObjectPtr context_manager_; 73 Safe_PyObjectPtr var_; 74 }; 75 76 } // namespace tensorflow 77 78 #endif // TENSORFLOW_PYTHON_FRAMEWORK_PY_CONTEXT_MANAGER_H_ 79