Implementation of `flatten`.
(module,
recursive,
predicate,
attribute_traversal_key,
attributes_to_ignore,
with_path,
module_path=(),
seen=None)
| 318 | |
| 319 | |
| 320 | def _flatten_module(module, |
| 321 | recursive, |
| 322 | predicate, |
| 323 | attribute_traversal_key, |
| 324 | attributes_to_ignore, |
| 325 | with_path, |
| 326 | module_path=(), |
| 327 | seen=None): |
| 328 | """Implementation of `flatten`.""" |
| 329 | if seen is None: |
| 330 | seen = set([id(module)]) |
| 331 | |
| 332 | module_dict = vars(module) |
| 333 | submodules = [] |
| 334 | |
| 335 | for key in sorted(module_dict, key=attribute_traversal_key): |
| 336 | if key in attributes_to_ignore: |
| 337 | continue |
| 338 | |
| 339 | for leaf_path, leaf in nest.flatten_with_tuple_paths(module_dict[key]): |
| 340 | leaf_path = (key,) + leaf_path |
| 341 | |
| 342 | # TODO(tomhennigan) Handle cycles for `with_path=True` (e.g. `a.a = a`). |
| 343 | if not with_path: |
| 344 | leaf_id = id(leaf) |
| 345 | if leaf_id in seen: |
| 346 | continue |
| 347 | seen.add(leaf_id) |
| 348 | |
| 349 | if predicate(leaf): |
| 350 | if with_path: |
| 351 | yield module_path + leaf_path, leaf |
| 352 | else: |
| 353 | yield leaf |
| 354 | |
| 355 | if recursive and _is_module(leaf): |
| 356 | # Walk direct properties first then recurse. |
| 357 | submodules.append((module_path + leaf_path, leaf)) |
| 358 | |
| 359 | for submodule_path, submodule in submodules: |
| 360 | subvalues = _flatten_module( |
| 361 | submodule, |
| 362 | recursive=recursive, |
| 363 | predicate=predicate, |
| 364 | attribute_traversal_key=attribute_traversal_key, |
| 365 | attributes_to_ignore=submodule._TF_MODULE_IGNORED_PROPERTIES, |
| 366 | with_path=with_path, |
| 367 | module_path=submodule_path, |
| 368 | seen=seen) |
| 369 | |
| 370 | for subvalue in subvalues: |
| 371 | # Predicate is already tested for these values. |
| 372 | yield subvalue |
no test coverage detected