Skip to content

Commit 02f808a

Browse files
MarkDaoustcopybara-github
authored andcommitted
Switch score_name to take a tuple instead of a string.
This will be necessary (with the new toc generator) when generating docs for something like `tf.lite` where the root module name contains a dot. Splittling on `.` gives a path ('tf','lite') which doesn't exist when generating just `tf.lite`. PiperOrigin-RevId: 439213594
1 parent a541d7e commit 02f808a

File tree

2 files changed

+48
-27
lines changed

2 files changed

+48
-27
lines changed

tools/tensorflow_docs/api_generator/doc_generator_visitor.py

Lines changed: 47 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@
1919
import functools
2020
import inspect
2121

22-
from typing import Any, Dict, List, Optional, Tuple
22+
from typing import Any, Dict, List, Optional, NamedTuple, Tuple
2323

2424
from tensorflow_docs.api_generator import obj_type as obj_type_lib
2525

@@ -291,7 +291,14 @@ class or module.
291291

292292
return children
293293

294-
def _score_name(self, name: str):
294+
class NameScore(NamedTuple):
295+
defining_class_score: int
296+
experimental_score: int
297+
keras_score: int
298+
module_length_score: int
299+
path: ApiPath
300+
301+
def _score_name(self, path: ApiPath) -> NameScore:
295302
"""Return a tuple of scores indicating how to sort for the best name.
296303
297304
This function is meant to be used as the `key` to the `sorted` function.
@@ -307,18 +314,18 @@ def _score_name(self, name: str):
307314
name: Fallback, sorts lexicographically on the full_name.
308315
309316
Args:
310-
name: the full name to score, for example `tf.estimator.Estimator`
317+
path: APiPath to score, for example `('tf','estimator','Estimator')`
311318
312319
Returns:
313320
A tuple of scores. When sorted the preferred name will have the lowest
314321
value.
315322
"""
316-
parts = name.split('.')
317-
short_name = parts[-1]
318-
if len(parts) == 1:
319-
return (-99, -99, -99, -99, short_name)
323+
py_object = self.path_tree[path].py_object
324+
if len(path) == 1:
325+
return self.NameScore(-99, -99, -99, -99, path)
320326

321-
container = self._index.get('.'.join(parts[:-1]), name)
327+
short_name = path[-1]
328+
container = self.path_tree[path[:-1]].py_object
322329

323330
defining_class_score = 1
324331
if inspect.isclass(container):
@@ -327,30 +334,44 @@ def _score_name(self, name: str):
327334
defining_class_score = -1
328335

329336
experimental_score = -1
330-
if 'contrib' in parts or any('experimental' in part for part in parts):
337+
if 'contrib' in path or any('experimental' in part for part in path):
331338
experimental_score = 1
332339

333340
keras_score = 1
334-
if 'keras' in parts:
341+
if 'keras' in path:
335342
keras_score = -1
336343

337-
while parts:
338-
container = self._index['.'.join(parts)]
344+
if inspect.ismodule(py_object):
345+
# prefer short paths for modules
346+
module_length_score = len(path)
347+
else:
348+
module_length_score = self._get_module_length_score(path)
349+
350+
return self.NameScore(
351+
defining_class_score=defining_class_score,
352+
experimental_score=experimental_score,
353+
keras_score=keras_score,
354+
module_length_score=module_length_score,
355+
path=path)
356+
357+
def _get_module_length_score(self, path):
358+
partial_path = list(path)
359+
while partial_path:
360+
container = self.path_tree[tuple(partial_path[:-1])].py_object
361+
partial_path.pop()
339362
if inspect.ismodule(container):
340363
break
341-
parts.pop()
342364

343-
module_length = len(parts)
365+
module_length = len(partial_path)
344366

345-
if len(parts) == 2:
367+
if module_length == 2:
346368
# `tf.submodule.thing` is better than `tf.thing`
347369
module_length_score = -1
348370
else:
349371
# shorter is better
350372
module_length_score = module_length
351373

352-
return (defining_class_score, experimental_score, keras_score,
353-
module_length_score, name)
374+
return module_length_score
354375

355376
def build(self):
356377
"""Compute data structures containing information about duplicates.
@@ -398,19 +419,20 @@ def build(self):
398419
if not aliases:
399420
aliases = [node]
400421

401-
names = [alias.full_name for alias in aliases]
422+
name_tuples = [alias.path for alias in aliases]
402423

403-
names = sorted(names)
404424
# Choose the main name with a lexical sort on the tuples returned by
405425
# by _score_name.
406-
main_name = min(names, key=self._score_name)
426+
main_name_tuple = min(name_tuples, key=self._score_name)
427+
main_name = '.'.join(main_name_tuple)
407428

408-
if names:
409-
duplicates[main_name] = list(names)
429+
names = ['.'.join(name_tuple) for name_tuple in name_tuples]
430+
if name_tuples:
431+
duplicates[main_name] = sorted(names)
410432

411-
names.remove(main_name)
412433
for name in names:
413-
duplicate_of[name] = main_name
434+
if name != main_name:
435+
duplicate_of[name] = main_name
414436

415437
# Set the reverse index to the canonical name.
416438
if not maybe_singleton(py_object):
@@ -552,8 +574,7 @@ def from_path_tree(cls, path_tree: PathTree, score_name_fn) -> 'ApiTree':
552574
# been processed, so now we can choose its master name.
553575
aliases = [node.path for node in duplicate_nodes]
554576

555-
master_path = min(['.'.join(a) for a in aliases], key=score_name_fn)
556-
master_path = tuple(master_path.split('.'))
577+
master_path = min(aliases, key=score_name_fn)
557578

558579
self.insert(master_path, current_node.py_object, aliases)
559580

tools/tensorflow_docs/api_generator/pretty_docs/base_page.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -57,7 +57,7 @@ def top_source_link(self):
5757
return top_source_link(self.page_info.defined_in)
5858

5959
def build_collapsable_aliases(self):
60-
return build_collapsable_aliases(self.page_info.aliases)
60+
return build_collapsable_aliases(sorted(self.page_info.aliases))
6161

6262
def top_compat(self):
6363
return build_top_compat(self.page_info, h_level=2)

0 commit comments

Comments
 (0)