Doc generator refactoring.

Change: 149127634

tensorflow/tools/docs/BUILD

@@ -41,18 +41,15 @@ py_test(
 
 py_library(
     name = "parser",
-    srcs = [
+    srcs = ["parser.py"],
-        "parser.py",
-    ],
     srcs_version = "PY2AND3",
+    visibility = ["//visibility:public"],
 )
 
 py_test(
     name = "parser_test",
     size = "small",
-    srcs = [
+    srcs = ["parser_test.py"],
-        "parser_test.py",
-    ],
     srcs_version = "PY2AND3",
     tags = ["manual"],
     deps = [
@@ -86,9 +83,7 @@ py_binary(
 py_test(
     name = "generate_lib_test",
     size = "small",
-    srcs = [
+    srcs = ["generate_lib_test.py"],
-        "generate_lib_test.py",
-    ],
     srcs_version = "PY2AND3",
     tags = ["manual"],
     deps = [
@@ -124,18 +119,14 @@ py_binary(
 
 py_library(
     name = "py_guide_parser",
-    srcs = [
+    srcs = ["py_guide_parser.py"],
-        "py_guide_parser.py",
-    ],
     srcs_version = "PY2AND3",
 )
 
 py_test(
     name = "py_guide_parser_test",
     size = "small",
-    srcs = [
+    srcs = ["py_guide_parser_test.py"],
-        "py_guide_parser_test.py",
-    ],
     srcs_version = "PY2AND3",
     deps = [
         ":py_guide_parser",

tensorflow/tools/docs/generate.py

@@ -18,7 +18,6 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
-import argparse
 import inspect
 import os
 import sys
@@ -30,22 +29,9 @@ from tensorflow.tools.docs import generate_lib
 
 
 if __name__ == '__main__':
-  argument_parser = argparse.ArgumentParser()
+  doc_generator = generate_lib.DocGenerator()
-  argument_parser.add_argument(
+  doc_generator.add_output_dir_argument()
-      '--output_dir',
+  doc_generator.add_src_dir_argument()
-      type=str,
-      default=None,
-      required=True,
-      help='Directory to write docs to.'
-  )
-
-  argument_parser.add_argument(
-      '--src_dir',
-      type=str,
-      default=None,
-      required=True,
-      help='Directory with the source docs.'
-  )
 
   # This doc generator works on the TensorFlow codebase. Since this script lives
   # at tensorflow/tools/docs, and all code is defined somewhere inside
@@ -54,23 +40,12 @@ if __name__ == '__main__':
   # moving the script around.
   script_dir = os.path.dirname(inspect.getfile(inspect.currentframe()))
   default_base_dir = os.path.join(script_dir, '..', '..')
+  doc_generator.add_base_dir_argument(default_base_dir)
 
-  argument_parser.add_argument(
+  flags = doc_generator.parse_known_args()
-      '--base_dir',
-      type=str,
-      default=default_base_dir,
-      help=('Base directory to to strip from file names referenced in docs. '
-            'Defaults to two directories up from the location of this file.')
-  )
-
-  flags, _ = argument_parser.parse_known_args()
 
   # tf_debug is not imported with tf, it's a separate module altogether
-  modules = [('tf', tf), ('tfdbg', tf_debug)]
+  doc_generator.set_py_modules([('tf', tf), ('tfdbg', tf_debug)])
+  doc_generator.load_contrib()
 
-  # Access something in contrib so tf.contrib is properly loaded (it's hidden
+  sys.exit(doc_generator.build(flags))
-  # behind lazy loading)
-  _ = tf.contrib.__name__
-
-  sys.exit(generate_lib.main(
-      flags.src_dir, flags.output_dir, flags.base_dir, modules))

tensorflow/tools/docs/generate_1_0.py

@@ -18,7 +18,6 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
-import argparse
 import inspect
 import os
 import sys
@@ -30,22 +29,9 @@ from tensorflow.tools.docs import generate_lib
 
 
 if __name__ == '__main__':
-  argument_parser = argparse.ArgumentParser()
+  doc_generator = generate_lib.DocGenerator()
-  argument_parser.add_argument(
+  doc_generator.add_output_dir_argument()
-      '--output_dir',
+  doc_generator.add_src_dir_argument()
-      type=str,
-      default=None,
-      required=True,
-      help='Directory to write docs to.'
-  )
-
-  argument_parser.add_argument(
-      '--src_dir',
-      type=str,
-      default=None,
-      required=True,
-      help='Directory with the source docs.'
-  )
 
   # This doc generator works on the TensorFlow codebase. Since this script lives
   # at tensorflow/tools/docs, and all code is defined somewhere inside
@@ -54,25 +40,15 @@ if __name__ == '__main__':
   # moving the script around.
   script_dir = os.path.dirname(inspect.getfile(inspect.currentframe()))
   default_base_dir = os.path.join(script_dir, '..', '..')
+  doc_generator.add_base_dir_argument(default_base_dir)
 
-  argument_parser.add_argument(
+  flags = doc_generator.parse_known_args()
-      '--base_dir',
-      type=str,
-      default=default_base_dir,
-      help=('Base directory to to strip from file names referenced in docs. '
-            'Defaults to two directories up from the location of this file.')
-  )
-
-  flags, _ = argument_parser.parse_known_args()
 
   # tf_debug is not imported with tf, it's a separate module altogether
-  modules = [('tf', tf), ('tfdbg', tf_debug)]
+  doc_generator.set_py_modules([('tf', tf), ('tfdbg', tf_debug)])
 
-  # Access something in contrib so tf.contrib is properly loaded (it's hidden
+  doc_generator.load_contrib()
-  # behind lazy loading)
+  doc_generator.set_do_not_descend_map({
-  _ = tf.contrib.__name__
-
-  generate_lib.do_not_descend_map = {
       '': ['cli', 'lib', 'wrappers'],
       'contrib': [
           'compiler',
@@ -121,7 +97,6 @@ if __name__ == '__main__':
           'utils',
       ],
       'contrib.util': ['loader'],
-  }
+  })
 
-  sys.exit(generate_lib.main(
+  sys.exit(doc_generator.build(flags))
-      flags.src_dir, flags.output_dir, flags.base_dir, modules))

tensorflow/tools/docs/generate_lib.py

@@ -18,6 +18,7 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
+import argparse
 import inspect
 import os
 
@@ -31,35 +32,24 @@ from tensorflow.tools.docs import pretty_docs
 from tensorflow.tools.docs import py_guide_parser
 
 
-def write_docs(output_dir, base_dir, duplicate_of, duplicates, index, tree,
+def write_docs(output_dir, parser_config, duplicate_of, index, yaml_toc):
-               reverse_index, reference_resolver, guide_index):
   """Write previously extracted docs to disk.
 
   Write a docs page for each symbol in `index` to a tree of docs at
   `output_dir`.
 
-  Symbols with multiple aliases will have only one page written about them,
+  Symbols with multiple aliases will have only one page written about
-  which is referenced for all aliases. `duplicate_of` and `duplicates` are used
+  them, which is referenced for all aliases.
-  to determine which docs pages to write.
 
   Args:
     output_dir: Directory to write documentation markdown files to. Will be
       created if it doesn't exist.
-    base_dir: Base directory of the code being documented. This prefix is
+    parser_config: A `parser.ParserConfig` object.
-      stripped from all file paths that are part of the documentation.
+    duplicate_of: A `dict` mapping fully qualified names to "master" names.
-    duplicate_of: A `dict` mapping fully qualified names to "master" names. This
+      Used to determine which docs pages to write.
-      is used to resolve "@{symbol}" references to the "master" name.
-    duplicates: A `dict` mapping fully qualified names to a set of all
-      aliases of this name. This is used to automatically generate a list of all
-      aliases for each name.
     index: A `dict` mapping fully qualified names to the corresponding Python
-      objects. Used to produce docs for child objects, and to check the validity
+      objects. Used to produce docs for child objects.
-      of "@{symbol}" references.
+    yaml_toc: Set to `True` to generate a "_toc.yaml" file.
-    tree: A `dict` mapping a fully qualified name to the names of all its
-      members. Used to populate the members section of a class or module page.
-    reverse_index: A `dict` mapping object ids to fully qualified names.
-    reference_resolver: A parser.ReferenceResolver object.
-    guide_index: A `dict` mapping symbol name strings to _GuideRef.
   """
   # Make output_dir.
   try:
@@ -97,7 +87,7 @@ def write_docs(output_dir, base_dir, duplicate_of, duplicates, index, tree,
 
     # For a module, remember the module for the table-of-contents
     if inspect.ismodule(py_object):
-      if full_name in tree:
+      if full_name in parser_config.tree:
         module_children.setdefault(full_name, [])
 
     # For something else that's documented,
@@ -113,15 +103,7 @@ def write_docs(output_dir, base_dir, duplicate_of, duplicates, index, tree,
     print('Writing docs for %s (%r).' % (full_name, py_object))
 
     # Generate docs for `py_object`, resolving references.
-    page_info = parser.docs_for_object(
+    page_info = parser.docs_for_object(full_name, py_object, parser_config)
-        full_name,
-        py_object,
-        reference_resolver=reference_resolver,
-        duplicates=duplicates,
-        tree=tree,
-        reverse_index=reverse_index,
-        guide_index=guide_index,
-        base_dir=base_dir)
 
     path = os.path.join(output_dir, parser.documentation_path(full_name))
     directory = os.path.dirname(path)
@@ -135,6 +117,7 @@ def write_docs(output_dir, base_dir, duplicate_of, duplicates, index, tree,
                                                              directory, e))
       raise
 
+  if yaml_toc:
     # Generate table of contents
 
     # Put modules in alphabetical order, case-insensitive
@@ -161,8 +144,8 @@ def write_docs(output_dir, base_dir, duplicate_of, duplicates, index, tree,
 
   # Write a global index containing all full names with links.
   with open(os.path.join(output_dir, 'index.md'), 'w') as f:
-    f.write(
+    f.write(parser.generate_global_index(
-        parser.generate_global_index('TensorFlow', index, reference_resolver))
+        'TensorFlow', index, parser_config.reference_resolver))
 
 
 def add_dict_to_dict(add_from, add_to):
@@ -174,8 +157,9 @@ def add_dict_to_dict(add_from, add_to):
 
 
 # Exclude some libaries in contrib from the documentation altogether.
-# TODO(wicke): Shrink this list.
+def _get_default_do_not_descend_map():
-do_not_descend_map = {
+  # TODO(wicke): Shrink this list.
+  return {
       '': ['cli', 'lib', 'wrappers'],
       'contrib': [
           'compiler',
@@ -221,20 +205,20 @@ do_not_descend_map = {
           'utils',
       ],
       'contrib.util': ['loader'],
-}
+  }
 
 
-def extract(modules):
+def extract(py_modules, do_not_descend_map):
   """Extract docs from tf namespace and write them to disk."""
   # Traverse the first module.
-  visitor = doc_generator_visitor.DocGeneratorVisitor(modules[0][0])
+  visitor = doc_generator_visitor.DocGeneratorVisitor(py_modules[0][0])
   api_visitor = public_api.PublicAPIVisitor(visitor)
   add_dict_to_dict(do_not_descend_map, api_visitor.do_not_descend_map)
 
-  traverse.traverse(modules[0][1], api_visitor)
+  traverse.traverse(py_modules[0][1], api_visitor)
 
-  # Traverse all modules after the first:
+  # Traverse all py_modules after the first:
-  for module_name, module in modules[1:]:
+  for module_name, module in py_modules[1:]:
     visitor.set_root_name(module_name)
     traverse.traverse(module, api_visitor)
 
@@ -261,7 +245,7 @@ class _DocInfo(object):
     self.title = title
 
 
-def _build_doc_index(src_dir):
+def build_doc_index(src_dir):
   """Build an index from a keyword designating a doc to _DocInfo objects."""
   doc_index = {}
   for dirpath, _, filenames in os.walk(src_dir):
@@ -329,32 +313,12 @@ class _GenerateGuideIndex(py_guide_parser.PyGuideParser):
 def _build_guide_index(guide_src_dir):
   """Return dict: symbol name -> _GuideRef from the files in `guide_src_dir`."""
   index_generator = _GenerateGuideIndex()
+  if os.path.exists(guide_src_dir):
     for full_path, base_name in py_guide_parser.md_files_in_dir(guide_src_dir):
       index_generator.process(full_path, base_name)
   return index_generator.index
 
 
-def _write(output_dir, base_dir, doc_index, guide_index, visitor):
-  """Write documentation for an index in a `DocGeneratorVisitor` to disk.
-
-  This function will create `output_dir` if it doesn't exist, and write
-  the documentation contained in `visitor`.
-
-  Args:
-    output_dir: The directory to write documentation to. Must not exist.
-    base_dir: The base dir of the library `visitor` has traversed. This is used
-      to compute relative paths for file references.
-    doc_index: A `dict` mapping a doc key to a _DocInfo.
-    guide_index: A `dict` mapping symbol name strings to _GuideRef.
-    visitor: A `DocGeneratorVisitor` that has traversed a library located at
-      `base_dir`.
-  """
-  write_docs(output_dir, os.path.abspath(base_dir),
-             visitor.duplicate_of, visitor.duplicates,
-             visitor.index, visitor.tree, visitor.reverse_index,
-             doc_index, guide_index)
-
-
 class _UpdateTags(py_guide_parser.PyGuideParser):
   """Rewrites a Python guide so that each section has an explicit tag."""
 
@@ -412,19 +376,105 @@ def _other_docs(src_dir, output_dir, reference_resolver):
   print('Done.')
 
 
-def main(src_dir, output_dir, base_dir, modules):
+class DocGenerator(object):
-  """Generate docs from `src_dir` to `output_dir`."""
+  """Main entry point for generating docs."""
-  doc_index = _build_doc_index(src_dir)
+
-  visitor = extract(modules)
+  def __init__(self):
-  reference_resolver = parser.ReferenceResolver(
+    self.argument_parser = argparse.ArgumentParser()
+    self._py_modules = None
+    self._do_not_descend_map = _get_default_do_not_descend_map()
+    self.yaml_toc = True
+
+  def add_output_dir_argument(self):
+    self.argument_parser.add_argument(
+        '--output_dir',
+        type=str,
+        default=None,
+        required=True,
+        help='Directory to write docs to.'
+    )
+
+  def add_src_dir_argument(self):
+    self.argument_parser.add_argument(
+        '--src_dir',
+        type=str,
+        default=None,
+        required=True,
+        help='Directory with the source docs.'
+    )
+
+  def add_base_dir_argument(self, default_base_dir):
+    self.argument_parser.add_argument(
+        '--base_dir',
+        type=str,
+        default=default_base_dir,
+        help='Base directory to to strip from file names referenced in docs.'
+    )
+
+  def parse_known_args(self):
+    flags, _ = self.argument_parser.parse_known_args()
+    return flags
+
+  def add_to_do_not_descend_map(self, d):
+    add_dict_to_dict(d, self._do_not_descend_map)
+
+  def set_do_not_descend_map(self, d):
+    self._do_not_descend_map = d
+
+  def set_py_modules(self, py_modules):
+    self._py_modules = py_modules
+
+  def load_contrib(self):
+    """Access something in contrib so tf.contrib is properly loaded."""
+    # Without this, it ends up hidden behind lazy loading.  Requires
+    # that the caller has already called set_py_modules().
+    if self._py_modules is None:
+      raise RuntimeError(
+          'Must call set_py_modules() before running load_contrib().')
+    for name, module in self._py_modules:
+      if name == 'tf':
+        _ = module.contrib.__name__
+        return True
+    return False
+
+  def py_module_names(self):
+    if self._py_modules is None:
+      raise RuntimeError(
+          'Must call set_py_modules() before running py_module_names().')
+    return [name for (name, _) in self._py_modules]
+
+  def make_reference_resolver(self, visitor, doc_index):
+    return parser.ReferenceResolver(
         duplicate_of=visitor.duplicate_of,
-      doc_index=doc_index, index=visitor.index)
+        doc_index=doc_index, index=visitor.index,
-  _write(os.path.join(output_dir, 'api_docs/python'), base_dir,
+        py_module_names=self.py_module_names())
-         reference_resolver,
+
-         _build_guide_index(os.path.join(src_dir, 'api_guides/python')),
+  def make_parser_config(self, visitor, reference_resolver, guide_index,
-         visitor)
+                         base_dir):
-  _other_docs(src_dir, output_dir, reference_resolver)
+    return parser.ParserConfig(
+        reference_resolver=reference_resolver,
+        duplicates=visitor.duplicates,
+        tree=visitor.tree,
+        reverse_index=visitor.reverse_index,
+        guide_index=guide_index,
+        base_dir=base_dir)
+
+  def build(self, flags):
+    """Actually build the docs."""
+    doc_index = build_doc_index(flags.src_dir)
+    visitor = extract(self._py_modules, self._do_not_descend_map)
+    reference_resolver = self.make_reference_resolver(visitor, doc_index)
+    guide_index = _build_guide_index(
+        os.path.join(flags.src_dir, 'api_guides/python'))
+    parser_config = self.make_parser_config(visitor, reference_resolver,
+                                            guide_index, flags.base_dir)
+    output_dir = os.path.join(flags.output_dir, 'api_docs/python')
+
+    write_docs(output_dir, parser_config, visitor.duplicate_of, visitor.index,
+               yaml_toc=self.yaml_toc)
+    _other_docs(flags.src_dir, flags.output_dir, reference_resolver)
+
     if parser.all_errors:
-    print('Errors during processing:' + '\n  '.join(parser.all_errors))
+      print('Errors during processing:\n  ' + '\n  '.join(parser.all_errors))
       return 1
     return 0

tensorflow/tools/docs/generate_lib_test.py

@@ -49,10 +49,11 @@ class TestClass(object):
 class GenerateTest(googletest.TestCase):
 
   def test_extraction(self):
-    modules = [('tf', tf), ('tfdbg', tf_debug)]
+    py_modules = [('tf', tf), ('tfdbg', tf_debug)]
     _ = tf.contrib.__name__  # Trigger loading of tf.contrib
     try:
-      generate_lib.extract(modules)
+      generate_lib.extract(
+          py_modules, generate_lib._get_default_do_not_descend_map())
     except RuntimeError:
       print('*****************************************************************')
       print('If this test fails, you have most likely introduced an unsealed')
@@ -97,11 +98,12 @@ class GenerateTest(googletest.TestCase):
 
     reference_resolver = parser.ReferenceResolver(
         duplicate_of=duplicate_of,
-        doc_index={}, index=index)
+        doc_index={}, index=index, py_module_names=['tf'])
-    generate_lib.write_docs(output_dir, base_dir, duplicate_of, duplicates,
+    parser_config = parser.ParserConfig(
-                            index, tree, reverse_index={},
+        reference_resolver=reference_resolver, duplicates=duplicates, tree=tree,
-                            reference_resolver=reference_resolver,
+        reverse_index={}, guide_index={}, base_dir=base_dir)
-                            guide_index={})
+    generate_lib.write_docs(output_dir, parser_config, duplicate_of, index,
+                            yaml_toc=True)
 
     # Make sure that the right files are written to disk.
     self.assertTrue(os.path.exists(os.path.join(output_dir, 'index.md')))

tensorflow/tools/docs/parser.py

@@ -92,12 +92,14 @@ class ReferenceResolver(object):
     doc_index: A `dict` mapping symbol name strings to objects with `url`
       and `title` fields. Used to resolve @{$doc} references in docstrings.
     index: A map from all full names to python objects.
+    py_module_names: A list of string names of Python modules.
   """
 
-  def __init__(self, duplicate_of, doc_index, index):
+  def __init__(self, duplicate_of, doc_index, index, py_module_names):
     self._duplicate_of = duplicate_of
     self._doc_index = doc_index
     self._index = index
+    self._py_module_names = py_module_names
 
   def replace_references(self, string, relative_path_to_root):
     """Replace "@{symbol}" references with links to symbol's documentation page.
@@ -225,6 +227,31 @@ class ReferenceResolver(object):
 
     # Handle different types of references.
     if string.startswith('$'):  # Doc reference
+      return self._doc_link(
+          string, link_text, manual_link_text, relative_path_to_root)
+
+    elif string.startswith('tensorflow::'):
+      # C++ symbol
+      return self._cc_link(
+          string, link_text, manual_link_text, relative_path_to_root)
+
+    else:
+      is_python = False
+      for py_module_name in self._py_module_names:
+        if string == py_module_name or string.startswith(py_module_name + '.'):
+          is_python = True
+          break
+      if is_python:  # Python symbol
+        return self.python_link(link_text, string, relative_path_to_root,
+                                code_ref=not manual_link_text)
+
+    # Error!
+    log_error('Did not understand "@{%s}"' % string)
+    return 'ERROR:%s' % string
+
+  def _doc_link(self, string, link_text, manual_link_text,
+                relative_path_to_root):
+    """Generate a link for a @{$...} reference."""
     string = string[1:]  # remove leading $
 
     # If string has a #, split that part into `hash_tag`
@@ -240,17 +267,18 @@ class ReferenceResolver(object):
       url = os.path.normpath(os.path.join(
           relative_path_to_root, '../..', self._doc_index[string].url))
       return '[%s](%s%s)' % (link_text, url, hash_tag)
+    return self._doc_missing(string, hash_tag, link_text, manual_link_text,
+                             relative_path_to_root)
+
+  def _doc_missing(self, string, unused_hash_tag, link_text,
+                   unused_manual_link_text, unused_relative_path_to_root):
+    """Generate an error for unrecognized @{$...} references."""
     log_error('Handle doc reference "@{$%s}"' % string)
-      return 'TODO:%s' % string
+    return link_text
 
-    # TODO(josh11b): The list of Python prefixes should be passed in.
+  def _cc_link(self, string, link_text, unused_manual_link_text,
-    elif string.startswith('tf.') or string.startswith('tfdbg.'):
+               relative_path_to_root):
-      # Python symbol
+    """Generate a link for a @{tensorflow::...} reference."""
-      return self.python_link(link_text, string, relative_path_to_root,
-                              code_ref=not manual_link_text)
-
-    elif string.startswith('tensorflow::'):
-      # C++ symbol
     # TODO(josh11b): Fix this hard-coding of paths.
     if string == 'tensorflow::ClientSession':
       ret = 'class/tensorflow/client-session.md'
@@ -271,10 +299,6 @@ class ReferenceResolver(object):
         relative_path_to_root, '../cc', ret))
     return '[`%s`](%s)' % (link_text, cc_relative_path)
 
-    # Error!
-    log_error('Did not understand "@{%s}"' % string)
-    return 'ERROR:%s' % string
-
 
 # TODO(aselle): Collect these into a big list for all modules and functions
 # and make a rosetta stone page.
@@ -936,8 +960,37 @@ class _ModulePageInfo(object):
       self._add_member(name, member_full_name, member, member_doc, url)
 
 
-def docs_for_object(full_name, py_object, duplicates, reference_resolver, tree,
+class ParserConfig(object):
-                    reverse_index, guide_index, base_dir):
+
+  def __init__(self, reference_resolver, duplicates, tree, reverse_index,
+               guide_index, base_dir):
+    """Object with the common config for docs_for_object() calls.
+
+    Args:
+      reference_resolver: An instance of ReferenceResolver.
+      duplicates: A `dict` mapping fully qualified names to a set of all
+        aliases of this name. This is used to automatically generate a list of
+        all aliases for each name.
+      tree: A `dict` mapping a fully qualified name to the names of all its
+        members. Used to populate the members section of a class or module page.
+      reverse_index: A `dict` mapping objects in the index to full names.
+      guide_index: A `dict` mapping symbol name strings to objects with a
+        `make_md_link()` method.
+      base_dir: A base path that is stripped from file locations written to the
+        docs.
+    """
+    self.reference_resolver = reference_resolver
+    self.duplicates = duplicates
+    self.tree = tree
+    self.reverse_index = reverse_index
+    self.guide_index = guide_index
+    self.base_dir = base_dir
+    self.defined_in_prefix = 'tensorflow/'
+    self.code_url_prefix = (
+        'https://www.tensorflow.org/code/tensorflow/')  # pylint: disable=line-too-long
+
+
+def docs_for_object(full_name, py_object, parser_config):
   """Return a PageInfo object describing a given object from the TF API.
 
   This function uses _parse_md_docstring to parse the docs pertaining to
@@ -956,17 +1009,7 @@ def docs_for_object(full_name, py_object, duplicates, reference_resolver, tree,
       documented.
     py_object: The Python object to be documented. Its documentation is sourced
       from `py_object`'s docstring.
-    duplicates: A `dict` mapping fully qualified names to a set of all
+    parser_config: A ParserConfig object.
-      aliases of this name. This is used to automatically generate a list of all
-      aliases for each name.
-    reference_resolver: An instance of ReferenceResolver.
-    tree: A `dict` mapping a fully qualified name to the names of all its
-      members. Used to populate the members section of a class or module page.
-    reverse_index: A `dict` mapping objects in the index to full names.
-    guide_index: A `dict` mapping symbol name strings to objects with a
-      `make_md_link()` method.
-    base_dir: A base path that is stripped from file locations written to the
-      docs.
 
   Returns:
     Either a `_FunctionPageInfo`, `_ClassPageInfo`, or a `_ModulePageInfo`
@@ -978,23 +1021,26 @@ def docs_for_object(full_name, py_object, duplicates, reference_resolver, tree,
   """
 
   # Which other aliases exist for the object referenced by full_name?
-  master_name = reference_resolver.py_master_name(full_name)
+  master_name = parser_config.reference_resolver.py_master_name(full_name)
-  duplicate_names = duplicates.get(master_name, [full_name])
+  duplicate_names = parser_config.duplicates.get(master_name, [full_name])
 
   # TODO(wicke): Once other pieces are ready, enable this also for partials.
   if (inspect.ismethod(py_object) or inspect.isfunction(py_object) or
       # Some methods in classes from extensions come in as routines.
       inspect.isroutine(py_object)):
     page_info = _FunctionPageInfo(master_name)
-    page_info.set_signature(py_object, reverse_index)
+    page_info.set_signature(py_object, parser_config.reverse_index)
 
   elif inspect.isclass(py_object):
     page_info = _ClassPageInfo(master_name)
-    page_info.collect_docs_for_class(reference_resolver, tree, reverse_index)
+    page_info.collect_docs_for_class(parser_config.reference_resolver,
+                                     parser_config.tree,
+                                     parser_config.reverse_index)
 
   elif inspect.ismodule(py_object):
     page_info = _ModulePageInfo(master_name)
-    page_info.collect_docs_for_module(reference_resolver, tree)
+    page_info.collect_docs_for_module(parser_config.reference_resolver,
+                                      parser_config.tree)
 
   else:
     raise RuntimeError('Cannot make docs for object %s: %r' % (full_name,
@@ -1003,15 +1049,15 @@ def docs_for_object(full_name, py_object, duplicates, reference_resolver, tree,
   relative_path = os.path.relpath(
       path='.', start=os.path.dirname(documentation_path(full_name)) or '.')
 
-  page_info.set_doc(
+  page_info.set_doc(_parse_md_docstring(
-      _parse_md_docstring(py_object, relative_path, reference_resolver))
+      py_object, relative_path, parser_config.reference_resolver))
 
   page_info.set_aliases(duplicate_names)
 
-  page_info.set_guides(
+  page_info.set_guides(_get_guides_markdown(
-      _get_guides_markdown(duplicate_names, guide_index, relative_path))
+      duplicate_names, parser_config.guide_index, relative_path))
 
-  page_info.set_defined_in(_get_defined_in(py_object, base_dir))
+  page_info.set_defined_in(_get_defined_in(py_object, parser_config))
 
   return page_info
 
@@ -1041,11 +1087,10 @@ class _PythonFile(object):
   This can be used for the `defined_in` slot of the `PageInfo` obejcts.
   """
 
-  _code_url_prefix = (
+  def __init__(self, path, parser_config):
-      'https://www.tensorflow.org/code/tensorflow/')  # pylint: disable=line-too-long
-
-  def __init__(self, path):
     self.path = path
+    self.path_prefix = parser_config.defined_in_prefix
+    self.code_url_prefix = parser_config.code_url_prefix
 
   def is_builtin(self):
     return False
@@ -1057,8 +1102,9 @@ class _PythonFile(object):
     return False
 
   def __str__(self):
-    return 'Defined in [`tensorflow/{path}`]({code_prefix}{path}).\n\n'.format(
+    return 'Defined in [`{prefix}{path}`]({code_prefix}{path}).\n\n'.format(
-        path=self.path, code_prefix=self._code_url_prefix)
+        path=self.path, prefix=self.path_prefix,
+        code_prefix=self.code_url_prefix)
 
 
 class _GeneratedFile(object):
@@ -1069,8 +1115,9 @@ class _GeneratedFile(object):
   This can be used for the `defined_in` slot of the `PageInfo` obejcts.
   """
 
-  def __init__(self, path):
+  def __init__(self, path, parser_config):
     self.path = path
+    self.path_prefix = parser_config.defined_in_prefix
 
   def is_builtin(self):
     return False
@@ -1082,16 +1129,15 @@ class _GeneratedFile(object):
     return True
 
   def __str__(self):
-    return 'Defined in `tensorflow/%s`.\n\n' % self.path
+    return 'Defined in `%s%s`.\n\n' % (self.path_prefix, self.path)
 
 
-def _get_defined_in(py_object, base_dir):
+def _get_defined_in(py_object, parser_config):
   """Returns a description of where the passed in python object was defined.
 
   Arguments:
     py_object: The Python object.
-    base_dir: A base path that is stripped from file locations written to the
+    parser_config: A ParserConfig object.
-      docs.
 
   Returns:
     Either a `_PythonBuiltin`, `_PythonFile`, or a `_GeneratedFile`
@@ -1101,7 +1147,8 @@ def _get_defined_in(py_object, base_dir):
   # TODO(wicke): Only use decorators that support this in TF.
 
   try:
-    path = os.path.relpath(path=inspect.getfile(py_object), start=base_dir)
+    path = os.path.relpath(path=inspect.getfile(py_object),
+                           start=parser_config.base_dir)
   except TypeError:  # getfile throws TypeError if py_object is a builtin.
     return _PythonBuiltin()
 
@@ -1118,10 +1165,10 @@ def _get_defined_in(py_object, base_dir):
     return None
 
   if re.match(r'.*/gen_[^/]*\.py$', path):
-    return _GeneratedFile(path)
+    return _GeneratedFile(path, parser_config)
 
   else:
-    return _PythonFile(path)
+    return _PythonFile(path, parser_config)
 
 
 def generate_global_index(library_name, index, reference_resolver):

tensorflow/tools/docs/parser_test.py

@@ -76,7 +76,8 @@ class ParserTest(googletest.TestCase):
              'tf.third': HasOneMember,
              'tf.fourth': HasOneMember}
     reference_resolver = parser.ReferenceResolver(
-        duplicate_of=duplicate_of, doc_index={}, index=index)
+        duplicate_of=duplicate_of, doc_index={}, index=index,
+        py_module_names=['tf'])
     result = reference_resolver.replace_references(string, '../..')
     self.assertEqual(
         'A [`tf.reference`](../../tf/reference.md), another '
@@ -98,7 +99,7 @@ class ParserTest(googletest.TestCase):
     doc2.url = 'somewhere/else'
     doc_index = {'doc1': doc1, 'do/c2': doc2}
     reference_resolver = parser.ReferenceResolver(
-        duplicate_of={}, doc_index=doc_index, index={})
+        duplicate_of={}, doc_index=doc_index, index={}, py_module_names=['tf'])
     result = reference_resolver.replace_references(string, 'python')
     self.assertEqual(
         '[Title1](../URL1) [Title1](../URL1#abc) [link](../URL1) '
@@ -115,21 +116,17 @@ class ParserTest(googletest.TestCase):
         'TestClass.CLASS_MEMBER': TestClass.CLASS_MEMBER
     }
     reference_resolver = parser.ReferenceResolver(
-        duplicate_of={}, doc_index={}, index=index)
+        duplicate_of={}, doc_index={}, index=index, py_module_names=['tf'])
 
     tree = {
         'TestClass': ['a_method', 'a_property', 'ChildClass', 'CLASS_MEMBER']
     }
+    parser_config = parser.ParserConfig(
+        reference_resolver=reference_resolver, duplicates={}, tree=tree,
+        reverse_index={}, guide_index={}, base_dir='/')
 
     page_info = parser.docs_for_object(
-        full_name='TestClass',
+        full_name='TestClass', py_object=TestClass, parser_config=parser_config)
-        py_object=TestClass,
-        reference_resolver=reference_resolver,
-        duplicates={},
-        tree=tree,
-        reverse_index={},
-        guide_index={},
-        base_dir='/')
 
     # Make sure the brief docstring is present
     self.assertEqual(
@@ -162,22 +159,18 @@ class ParserTest(googletest.TestCase):
         'TestModule.TestClass': TestClass,
     }
     reference_resolver = parser.ReferenceResolver(
-        duplicate_of={}, doc_index={}, index=index)
+        duplicate_of={}, doc_index={}, index=index, py_module_names=['tf'])
 
     tree = {
         'TestModule': ['TestClass', 'test_function',
                        'test_function_with_args_kwargs']
     }
+    parser_config = parser.ParserConfig(
+        reference_resolver=reference_resolver, duplicates={}, tree=tree,
+        reverse_index={}, guide_index={}, base_dir='/')
 
     page_info = parser.docs_for_object(
-        full_name='TestModule',
+        full_name='TestModule', py_object=module, parser_config=parser_config)
-        py_object=module,
-        reference_resolver=reference_resolver,
-        duplicates={},
-        tree=tree,
-        reverse_index={},
-        guide_index={},
-        base_dir='/')
 
     # Make sure the brief docstring is present
     self.assertEqual(inspect.getdoc(module).split('\n')[0], page_info.doc.brief)
@@ -196,21 +189,19 @@ class ParserTest(googletest.TestCase):
         'test_function': test_function
     }
     reference_resolver = parser.ReferenceResolver(
-        duplicate_of={}, doc_index={}, index=index)
+        duplicate_of={}, doc_index={}, index=index, py_module_names=['tf'])
 
     tree = {
         '': ['test_function']
     }
+    parser_config = parser.ParserConfig(
+        reference_resolver=reference_resolver, duplicates={}, tree=tree,
+        reverse_index={}, guide_index={}, base_dir='/')
 
     page_info = parser.docs_for_object(
         full_name='test_function',
         py_object=test_function,
-        reference_resolver=reference_resolver,
+        parser_config=parser_config)
-        duplicates={},
-        tree=tree,
-        reverse_index={},
-        guide_index={},
-        base_dir='/')
 
     # Make sure the brief docstring is present
     self.assertEqual(
@@ -228,21 +219,19 @@ class ParserTest(googletest.TestCase):
         'test_function_with_args_kwargs': test_function_with_args_kwargs
     }
     reference_resolver = parser.ReferenceResolver(
-        duplicate_of={}, doc_index={}, index=index)
+        duplicate_of={}, doc_index={}, index=index, py_module_names=['tf'])
 
     tree = {
         '': ['test_function_with_args_kwargs']
     }
+    parser_config = parser.ParserConfig(
+        reference_resolver=reference_resolver, duplicates={}, tree=tree,
+        reverse_index={}, guide_index={}, base_dir='/')
 
     page_info = parser.docs_for_object(
         full_name='test_function_with_args_kwargs',
         py_object=test_function_with_args_kwargs,
-        reference_resolver=reference_resolver,
+        parser_config=parser_config)
-        duplicates={},
-        tree=tree,
-        reverse_index={},
-        guide_index={},
-        base_dir='/')
 
     # Make sure the brief docstring is present
     self.assertEqual(
@@ -298,7 +287,8 @@ class ParserTest(googletest.TestCase):
         'tf.fourth': HasOneMember
     }
     reference_resolver = parser.ReferenceResolver(
-        duplicate_of=duplicate_of, doc_index={}, index=index)
+        duplicate_of=duplicate_of, doc_index={}, index=index,
+        py_module_names=['tf'])
 
     doc_info = parser._parse_md_docstring(test_function_with_fancy_docstring,
                                           '../..', reference_resolver)
@@ -329,7 +319,8 @@ class ParserTest(googletest.TestCase):
         'TestModule.test_function': 'test_function'
     }
     reference_resolver = parser.ReferenceResolver(
-        duplicate_of=duplicate_of, doc_index={}, index=index)
+        duplicate_of=duplicate_of, doc_index={}, index=index,
+        py_module_names=['tf'])
 
     docs = parser.generate_global_index('TestLibrary', index=index,
                                         reference_resolver=reference_resolver)