Smart cross references for PyGlove symbols in Sphinx autodoc.

This allows developers to refer to PyGlove symbols in Sphinx documentation in the same way as we write code, e.g. :class:`pg.DNA` or :class:`pg.geno.DNA`, both will output a link with the same canonical symbol name `pg.DNA`. As a result, all mentions of the same symbol will have consistent name across the entire document corpus. PiperOrigin-RevId: 517022666
cocolato · Mar 16, 2023 · 9d3cf8f · 9d3cf8f
1 parent 599fc01
commit 9d3cf8f
Show file tree

Hide file tree

Showing 14 changed files with 217 additions and 134 deletions.
diff --git a/docs/api/_templates/_default/class.rst b/docs/api/_templates/_default/class.rst
@@ -1,7 +1,7 @@
 .. _{{cls.rst_label}}:
 
-{{cls.canonical_path}}
-================
+{{cls.preferred_path}}
+======================
 
 Accessible via {{cls.rst_access_paths}}.
 

diff --git a/docs/api/_templates/_default/function.rst b/docs/api/_templates/_default/function.rst
@@ -1,6 +1,6 @@
 .. _{{func.rst_label}}:
 
-{{func.canonical_path}}
+{{func.preferred_path}}
 =====================
 
 Accessible via {{func.rst_access_paths}}.

diff --git a/docs/api/_templates/_default/module.rst b/docs/api/_templates/_default/module.rst
@@ -1,6 +1,6 @@
 .. _{{module.rst_label}}:
 
-{{module.canonical_path}}
+{{module.preferred_path}}
 =====================
 
 .. automodule:: {{module.rst_import_name}}

diff --git a/docs/api/_templates/_default/object.rst b/docs/api/_templates/_default/object.rst
@@ -1,6 +1,6 @@
 .. _{{obj.rst_label}}:
 
-{{obj.canonical_path}}
+{{obj.preferred_path}}
 ================
 
 Accessible via {{obj.rst_access_paths}}.

diff --git a/docs/api/_templates/core/symbolic/dict.rst b/docs/api/_templates/core/symbolic/dict.rst
@@ -1,6 +1,6 @@
 .. _{{cls.rst_label}}:
 
-{{cls.canonical_path}}
+{{cls.preferred_path}}
 ================
 
 Accessible via {{cls.rst_access_paths}}.

diff --git a/docs/api/_templates/core/symbolic/list.rst b/docs/api/_templates/core/symbolic/list.rst
@@ -1,6 +1,6 @@
 .. _{{cls.rst_label}}:
 
-{{cls.canonical_path}}
+{{cls.preferred_path}}
 ================
 
 Accessible via {{cls.rst_access_paths}}.

diff --git a/docs/api/_templates/core/symbolic/object.rst b/docs/api/_templates/core/symbolic/object.rst
@@ -1,6 +1,6 @@
 .. _{{cls.rst_label}}:
 
-{{cls.canonical_path}}
+{{cls.preferred_path}}
 ================
 
 Accessible via {{cls.rst_access_paths}}.

diff --git a/docs/api/_templates/core/symbolic/symbolic.rst b/docs/api/_templates/core/symbolic/symbolic.rst
@@ -1,6 +1,6 @@
 .. _{{cls.rst_label}}:
 
-{{cls.canonical_path}}
+{{cls.preferred_path}}
 ================
 
 Accessible via {{cls.rst_access_paths}}.

diff --git a/docs/api/docgen.py b/docs/api/docgen.py
@@ -131,7 +131,15 @@ def add_access_path(self, access_path):
   @property
   def access_paths(self) -> List[str]:
     """Returns all access paths. E.g. ['pg.Dict', 'pg.symbolic.Dict']."""
-    return sorted(self._access_paths, key=len)
+    def order(path):
+      names = path.split('.')
+      return (len(names), names[-1] != self.name)
+    return sorted(self._access_paths, key=order)
+
+  @property
+  def preferred_path(self) -> str:
+    """Returns preferred path."""
+    return self.access_paths[0]
 
   @property
   @abc.abstractmethod
@@ -376,7 +384,7 @@ def all_apis(self, memo: Optional[Set[id]] = None) -> List[Leaf]:
     return apis
 
 
-def get_api(pg) -> List[Module]:
+def get_api(pg) -> Module:
   """Get API entries from PyGlove module."""
   symbol_to_api = {}
 
@@ -452,7 +460,7 @@ def generate_api_docs(
     import_root: Optional[str] = None,
     template_root: Optional[str] = None,
     output_root: Optional[str] = None,
-    overwrite: bool = True):
+    overwrite: bool = True) -> Module:
   """Generate API docs."""
 
   if import_root is not None:
@@ -485,10 +493,12 @@ def visit(module: Module):
   # api = get_api(pg)
   # e = api['evolution']
   # print(e, e.source_category, e.access_paths, e.canonical_path)
-  visit(get_api(pg))
+  pyglove_api = get_api(pg)
+  visit(pyglove_api)
   print(f'API doc generation completed '
         f'(generated={stats.num_files - stats.num_skipped}, '
         f'skipped={stats.num_skipped})')
+  return pyglove_api
 
 
 def main(*unused_args, **unused_kwargs):

diff --git a/docs/conf.py b/docs/conf.py
@@ -31,22 +31,39 @@
 import importlib
 import inspect
 import os
+import re
 import sys
 
+from sphinx.domains.python import PyXRefRole
+
+
 # Set parent directory to path in order to import pyglove.
 sys.path.insert(0, os.path.abspath('..'))
 pyglove_module = importlib.import_module('pyglove')
+access_path_to_api = {}
 
 
 def generate_api_docs(_):
   docgen = importlib.import_module('docs.api.docgen')
   print('Generating API docs from templates...')
-  docgen.generate_api_docs()
+  pyglove_api = docgen.generate_api_docs()
+  for api in pyglove_api.all_apis():
+    if api.qualname:
+      access_path_to_api[api.qualname] = api
+    for path in api.access_paths:
+      access_path_to_api[path] = api
 
 
 def setup(app):
   app.connect('builder-inited', generate_api_docs)
-
+  app.add_role_to_domain('py', 'class', PgXRefRole())
+  app.add_role_to_domain('py', 'const', PgXRefRole())
+  app.add_role_to_domain('py', 'obj', PgXRefRole())
+  app.add_role_to_domain('py', 'data', PgXRefRole())
+  app.add_role_to_domain('py', 'func', PgXRefRole())
+  app.add_role_to_domain('py', 'meth', PgXRefRole())
+  app.add_role_to_domain('py', 'attr', PgXRefRole())
+  app.add_role_to_domain('py', 'mod', PgXRefRole())
 
 # Consider to include versioning.
 branch = os.getenv('BRANCH') or 'main'
@@ -99,6 +116,49 @@ def linkcode_resolve(domain, info):
     return None
 
 
+class PgXRefRole(PyXRefRole):
+  """Custom XRefRole for PyGlove code.
+
+  This role is introduced to consistently use the preferred name for the same
+  symbol, though they could be referenced with different paths. For example,
+  both :class:`pg.DNA` and :class:`pyglove.geno.DNA` will refer to class
+  ``pg.DNA``, while "pg.DNA" will be used as the the title for both references.
+  """
+
+  def process_link(self, env, refnode, has_explicit_title: bool,
+                   title: str, target: str) -> tuple[str, str]:
+    """Processes link."""
+    title, target = super().process_link(
+        env, refnode, has_explicit_title, title, target)
+
+    def noramlized_title_and_target(api, attr_name=None):
+      new_title = title
+      if not has_explicit_title:
+        new_title = api.preferred_path
+      new_target = api.canonical_path
+      if attr_name:
+        new_target = '.'.join([new_target, attr_name])
+      new_target = re.sub(r'^pg\.', 'pyglove.', new_target)
+      return (new_title, new_target)
+
+    # Try with target directly.
+    if target in access_path_to_api:
+      return noramlized_title_and_target(access_path_to_api[target])
+
+    # Replace target prefix and try again.
+    target = re.sub(r'^pyglove\.(?:core|ext)?\.?', 'pg.', target)
+    if target in access_path_to_api:
+      return noramlized_title_and_target(access_path_to_api[target])
+    else:
+      # new_target may be class members.
+      name_items = target.split('.')
+      class_name, attr_name = '.'.join(name_items[:-1]), name_items[-1]
+      if class_name in access_path_to_api:
+        return noramlized_title_and_target(
+            access_path_to_api[class_name], attr_name)
+    return (title, target)
+
+
 # -- Project information -----------------------------------------------------
 
 project = 'PyGlove'

diff --git a/pyglove/core/geno/__init__.py b/pyglove/core/geno/__init__.py
@@ -48,22 +48,23 @@
 
 Genotypes map 1:1 to hyper primitives as the following:
 
-+----------------------------------------+-------------------------------------+
-| Genotype class                         | Hyper class                         |
-+========================================+=====================================+
-|:class:`pyglove.DNASpec`                |:class:`pyglove.hyper.HyperValue`    |
-+----------------------------------------+-------------------------------------+
-|:class:`pyglove.geno.Space`             |:class:`pyglove.hyper.ObjectTemplate`|
-+----------------------------------------+-------------------------------------+
-|:class:`pyglove.geno.DecisionPoint`     |:class:`pyglove.hyper.HyperPrimitive`|
-+----------------------------------------+-------------------------------------+
-|:class:`pyglove.geno.Choices`           |:class:`pyglove.hyper.Choices`       |
-+----------------------------------------+-------------------------------------+
-|:class:`pyglove.geno.Float`             |:class:`pyglove.hyper.Float`         |
-+----------------------------------------+-------------------------------------+
-|:class:`pyglove.geno.CustomDecisionPoint` :class:`pyglove.hyper.CustomHyper`  |
-+-----------------------------------------+------------------------------------+
-
++-------------------------------------+--------------------------------------+
+| Genotype class                      | Hyper class                          |
++=====================================+======================================+
+|:class:`pg.DNASpec`                  | :class:`pg.hyper.HyperValue`         |
++-------------------------------------+--------------------------------------+
+|:class:`pg.geno.Space`               | :class:`pg.hyper.ObjectTemplate`     |
++-------------------------------------+--------------------------------------+
+|:class:`pg.geno.DecisionPoint`       | :class:`pg.hyper.HyperPrimitive`     |
++-------------------------------------+--------------------------------------+
+|:class:`pg.geno.Choices`             | :class:`pg.hyper.Choices`            |
+|                                     | (:func:`pg.oneof`, :func:`pg.manyof`)|
++-------------------------------------+--------------------------------------+
+|:class:`pg.geno.Float`               | :class:`pg.floatv`                   |
++-------------------------------------+--------------------------------------+
+|:class:`pg.geno.CustomDecisionPoint` | :class:`pg.hyper.CustomHyper`        |
+|                                     | (:func:`pg.evolve`)                  |
++-------------------------------------+--------------------------------------+
 """
 
 # pylint: disable=g-bad-import-order

diff --git a/pyglove/core/geno/base.py b/pyglove/core/geno/base.py
@@ -368,57 +368,60 @@ def is_space(self) -> bool:
     return False
 
 
+# pylint: disable=line-too-long
 @functools.total_ordering
 class DNA(symbolic.Object):
-  """A tree of numbers that encodes an symbolic object.
-
-  Each DNA object represents a node in a tree, which has a value and a list of
-  DNA as its children. DNA value can be None, int or float, with valid form as
-  following:
-
-    +-----------------------+-----------------+-----------------------------+
-    |  Encoder type         | Possible values |   Child nodes               |
-    |  (DNASpec type)       |                 |                             |
-    +=======================+=================+=============================+
-    |hyper.ObjectTemplate   |None             |DNA of child decision points |
-    |(geno.Space)           |(elements > 1)   |(Choices/Float) in the       |
-    |                       |                 |template.                    |
-    +-----------------------+-----------------+-----------------------------+
-    |                       |None             |Children of elements[0]      |
-    |                       |(elements == 1   |                             |
-    |                       |and elements[0]. |                             |
-    |                       |num_choices > 1) |                             |
-    +-----------------------+-----------------+-----------------------------+
-    |                       |int              |Children of:                 |
-    |                       |(elements == 1   |elements[0][0]               |
-    |                       |and elements[0]. |                             |
-    |                       |num_choices ==1) |                             |
-    +-----------------------+-----------------+-----------------------------+
-    |                       |float            |Empty                        |
-    |                       |(elements == 1   |                             |
-    |                       |and elements[0]  |                             |
-    |                       |is geno.Float    |                             |
-    +-----------------------+-----------------+-----------------------------+
-    |hyper.OneOf            |int              |Children of Space            |
-    |(geno.Choices)         |(candidate index |for the chosen candidate     |
-    |                       |as choice)       |                             |
-    +-----------------------+-----------------+-----------------------------+
-    |hyper.ManyOf           |None             |DNA of each chosen candidate |
-    |(geno.Choices)         |(num_choices > 1 |                             |
-    +-----------------------+-----------------+-----------------------------+
-    |                       |int              |Children of chosen candidate |
-    |                       |(num_choices==1) |                             |
-    +-----------------------+-----------------+-----------------------------+
-    |hyper.Float            |float            |Empty                        |
-    |(geno.Float)           |                 |                             |
-    +-----------------------+-----------------+-----------------------------+
-    |hyper.CustomHyper      |string           |User defined.                |
-    |(geno.CustomDecision   |                 |                             |
-    |Point)                 |                 |                             |
-    +-----------------------+-----------------+-----------------------------+
-
-  DNA can also be represented as a mix of JSON number, list and tuples for a
-  more intuitive illustration, formally defined as::
+  """The genome of a symbolic object relative to its search space.
+
+  DNA is a hierarchical structure - each DNA node has a value, and a list of
+  child DNA nodes. The root node represents the genome that encodes an entire
+  object relative to its space. The value of a DNA node could be None, an
+  integer, a float number or a string, dependening on its specification
+  (:class:`pg.DNASpec`). A valid DNA has a form of the following.
+
+  +--------------------------------------+-----------------+-----------------------------+
+  |  Hyper value type                    | Possible values |   Child nodes               |
+  |  (DNASpec type)                      |                 |                             |
+  +======================================+=================+=============================+
+  |:class:`pg.hyper.ObjectTemplate`      | None            |DNA of child decision points |
+  |(:class:`pg.geno.Space`)              |(elements > 1)   |(Choices/Float) in the       |
+  |                                      |                 |template.                    |
+  +--------------------------------------+-----------------+-----------------------------+
+  |                                      |None             |Children of elements[0]      |
+  |                                      |(elements == 1   |                             |
+  |                                      |and elements[0]. |                             |
+  |                                      |num_choices > 1) |                             |
+  +--------------------------------------+-----------------+-----------------------------+
+  |                                      |int              |Children of:                 |
+  |                                      |(elements == 1   |elements[0][0]               |
+  |                                      |and elements[0]. |                             |
+  |                                      |num_choices ==1) |                             |
+  +--------------------------------------+-----------------+-----------------------------+
+  |                                      |float            |Empty                        |
+  |                                      |(elements == 1   |                             |
+  |                                      |and elements[0]  |                             |
+  |                                      |is geno.Float)   |                             |
+  +--------------------------------------+-----------------+-----------------------------+
+  |:func:`pg.oneof`                      |int              |Children of Space            |
+  |(:class:`pg.geno.Choices`)            |(candidate index |for the chosen candidate     |
+  |                                      |as choice)       |                             |
+  +--------------------------------------+-----------------+-----------------------------+
+  |:func:`pg.manyof`                     |None             |DNA of each chosen candidate |
+  |(:class:`pg.geno.Choices)             |(num_choices > 1 |                             |
+  +--------------------------------------+-----------------+-----------------------------+
+  |                                      |int              |Children of chosen candidate |
+  |                                      |(num_choices==1) |                             |
+  +--------------------------------------+-----------------+-----------------------------+
+  |:func:`pg.floatv`                     |float            |Empty                        |
+  |(:class:`pg.geno.Float` )             |                 |                             |
+  +--------------------------------------+-----------------+-----------------------------+
+  |:class:`pg.hyper.CustomHyper`         |string           |User defined.                |
+  |(:class:`pg.geno.CustomDecisionPoint`)|(serialized      |                             |
+  |                                      | object)         |                             |
+  +--------------------------------------+-----------------+-----------------------------+
+
+  DNA can also be represented in a compact form - a tree of numbers/strs,
+  formally defined as::
 
     <dna> := empty | <decision>
     <decision>: = <single-decision>
@@ -463,6 +466,7 @@ class DNA(symbolic.Object):
     # is defined by the user.
     DNA('abc')
   """
+  # pylint: enable=line-too-long
 
   # Allow assignment on symbolic attributes.
   allow_symbolic_assignment = True