Patch `tf.io.gfile.GFile` and `epath.Path`. Example: Validate `exists` usage: ``` def new_exists(old_exists, path): assert not os.fspath(path).startswith('gs://') return old_exists(path) with mock_gfile(exists=new_exists) ``` Args: **fns: Functions to overwrite. Have sign
(**fns: Any)
| 277 | |
| 278 | @contextlib.contextmanager |
| 279 | def mock_gfile(**fns: Any) -> Iterator[None]: |
| 280 | """Patch `tf.io.gfile.GFile` and `epath.Path`. |
| 281 | |
| 282 | Example: Validate `exists` usage: |
| 283 | |
| 284 | ``` |
| 285 | def new_exists(old_exists, path): |
| 286 | assert not os.fspath(path).startswith('gs://') |
| 287 | return old_exists(path) |
| 288 | |
| 289 | with mock_gfile(exists=new_exists) |
| 290 | ``` |
| 291 | |
| 292 | Args: |
| 293 | **fns: Functions to overwrite. Have signature: `fn(original_fn, *args, |
| 294 | **kwargs)`. (note the first function argument which allow to access the |
| 295 | original function) |
| 296 | |
| 297 | Yields: |
| 298 | None |
| 299 | """ |
| 300 | # Process epath kwargs |
| 301 | epath_kwargs = {k: fn for k, fn in fns.items() if k != 'walk'} |
| 302 | |
| 303 | # Process gfile kwargs |
| 304 | epath_to_gfile_mapping = { |
| 305 | 'open': 'GFile', |
| 306 | } |
| 307 | gfile_kwargs = {} |
| 308 | for k, fn in fns.items(): |
| 309 | if k == 'replace': |
| 310 | continue |
| 311 | gfile_k = epath_to_gfile_mapping.get(k, k) |
| 312 | original_fn = getattr(tf.io.gfile, gfile_k) |
| 313 | mocked_fn = functools.wraps(original_fn)(functools.partial(fn, original_fn)) |
| 314 | gfile_kwargs[gfile_k] = mocked_fn |
| 315 | |
| 316 | with contextlib.ExitStack() as stack: |
| 317 | cm_epath = epath.testing.mock_epath(**epath_kwargs) |
| 318 | cm_gfile = mock_tf('tf.io.gfile', **gfile_kwargs) |
| 319 | stack.enter_context(cm_epath) |
| 320 | stack.enter_context(cm_gfile) |
| 321 | yield |
| 322 | |
| 323 | |
| 324 | @contextlib.contextmanager |