Patch TF API. This function is similar to `mock.patch.object`, but patch both `tf.Xyz` and `tf.compat.v2.Xyz`. Args: symbol_name: Symbol to patch (e.g. `tf.io.gfile`) *args: Arguments to forward to `mock.patch.object` **kwargs: Arguments to forward to `mock.patch.object` Yield
(symbol_name: str, *args: Any, **kwargs: Any)
| 323 | |
| 324 | @contextlib.contextmanager |
| 325 | def mock_tf(symbol_name: str, *args: Any, **kwargs: Any) -> Iterator[None]: |
| 326 | """Patch TF API. |
| 327 | |
| 328 | This function is similar to `mock.patch.object`, but patch both |
| 329 | `tf.Xyz` and `tf.compat.v2.Xyz`. |
| 330 | |
| 331 | Args: |
| 332 | symbol_name: Symbol to patch (e.g. `tf.io.gfile`) |
| 333 | *args: Arguments to forward to `mock.patch.object` |
| 334 | **kwargs: Arguments to forward to `mock.patch.object` |
| 335 | |
| 336 | Yields: |
| 337 | None |
| 338 | """ |
| 339 | |
| 340 | tf_symbol, *tf_submodules, symbol_name = symbol_name.split('.') |
| 341 | if tf_symbol != 'tf': |
| 342 | raise ValueError('Symbol name to patch should start by `tf`.') |
| 343 | |
| 344 | with contextlib.ExitStack() as stack: |
| 345 | # Recursivelly load the submodules/subobjects (e.g. `tf.io.gfile`) |
| 346 | module = tf |
| 347 | for submodule in tf_submodules: |
| 348 | module = getattr(module, submodule) |
| 349 | getattr(module, symbol_name) # Trigger the lazy-loading of the TF API. |
| 350 | if kwargs: # Patch each attribute individually |
| 351 | assert not args |
| 352 | for k, v in kwargs.items(): |
| 353 | stack.enter_context( |
| 354 | mock.patch.object(getattr(module, symbol_name), k, v) |
| 355 | ) |
| 356 | else: |
| 357 | # Patch the module/object |
| 358 | stack.enter_context( |
| 359 | mock.patch.object(module, symbol_name, *args, **kwargs) |
| 360 | ) |
| 361 | yield |
| 362 | |
| 363 | |
| 364 | def run_in_graph_and_eager_modes(func=None, config=None, use_gpu=True): |