[doc] fix doc to work in all known repos.

recipe_engine/doc.py

 #!/usr/bin/env python
 # Copyright 2013 The LUCI Authors. All rights reserved.
 # Use of this source code is governed under the Apache License, Version 2.0
 # that can be found in the LICENSE file.
 
 from __future__ import print_function, absolute_import
 
 import ast
 import inspect
 import json
@@ skipping 55 matching lines @@
   astunparse.Unparser(node, buf)
   return buf.getvalue()
 
 
 def _find_value_of(mod_ast, target):
   """Looks for an assignment to `target`, returning the assignment value ast
   node and the line number of the assignment.
   """
   for node in mod_ast.body:
     if isinstance(node, ast.Assign):
-      if len(node.targets) == 1 and node.targets[0].id == target:
+      if (len(node.targets) == 1 and
+          isinstance(node.targets[0], ast.Name) and
+          node.targets[0].id == target):
         return node.value, node.lineno
   return None, None
 
 
 def _expand_mock_imports(*mock_imports):
   """Returns an expanded set of mock imports.
 
   mock_imports is expected to be a dict which looks like:
     {
       "absolute.module.name": SomeObject,
@@ skipping 32 matching lines @@
         expanded_imports[dotted_name+'.'+name] = getattr(obj, name)
     for i in range(len(toks)-1, 0, -1):
       partial = '.'.join(toks[:i])
       cur_obj = expanded_imports.setdefault(partial, expando())
       if not isinstance(cur_obj, expando):
         raise ValueError('nested mock imports! %r', partial)
       setattr(cur_obj, toks[i], expanded_imports[partial+'.'+toks[i]])
 
   return expanded_imports
 
+ALL_IMPORTS = {}  # used in doc_test to ensure everything is actually importable
+KNOWN_OBJECTS = {}
 
 _decorator_imports = {
   'recipe_engine.util.returns_placeholder': util.returns_placeholder,
   'recipe_engine.recipe_api.non_step': recipe_api.non_step,
   'recipe_engine.recipe_api.infer_composite_step': (
     recipe_api.infer_composite_step)
 }
+KNOWN_OBJECTS.update(_decorator_imports)
 
 _config_imports = {
   'recipe_engine.config.ConfigGroup': config.ConfigGroup,
   'recipe_engine.config.ConfigList': config.ConfigList,
   'recipe_engine.config.Set': config.Set,
   'recipe_engine.config.Dict': config.Dict,
   'recipe_engine.config.List': config.List,
   'recipe_engine.config.Single': config.Single,
   'recipe_engine.config.Static': config.Static,
   'recipe_engine.config.Enum': config.Enum,
 }
+KNOWN_OBJECTS.update(_config_imports)
 
 _placeholder_imports = {
   'recipe_engine.util.OutputPlaceholder': util.OutputPlaceholder,
   'recipe_engine.util.InputPlaceholder': util.InputPlaceholder,
   'recipe_engine.util.Placeholder': util.Placeholder,
 }
+KNOWN_OBJECTS.update(_placeholder_imports)
 
-MOCK_IMPORTS_PARAMETERS = _expand_mock_imports({
+_property_imports = {
   'recipe_engine.recipe_api.Property': recipe_api.Property,
-}, _config_imports)
+}
+KNOWN_OBJECTS.update(_property_imports)
 
-MOCK_IMPORTS_RETURN_SCHEMA = _expand_mock_imports({
+_return_schema_imports = {
   'recipe_engine.config.ReturnSchema': config.ReturnSchema,
   'recipe_engine.config.ConfigGroupSchema': config.ConfigGroupSchema,
-}, _config_imports)
+}
+KNOWN_OBJECTS.update(_return_schema_imports)
 
-MOCK_IMPORTS_RECIPE = _expand_mock_imports({
+_util_imports = {
   'recipe_engine.types.freeze': types.freeze,
-}, _decorator_imports, _placeholder_imports)
+}
+KNOWN_OBJECTS.update(_util_imports)
 
-MOCK_IMPORTS_MODULE = _expand_mock_imports({
+_recipe_api_class_imports = {
   'recipe_engine.recipe_api.RecipeApi': recipe_api.RecipeApi,
   'recipe_engine.recipe_api.RecipeApiPlain': recipe_api.RecipeApiPlain,
-}, _decorator_imports, _placeholder_imports)
-
-ALL_IMPORTS = {}
-ALL_IMPORTS.update(MOCK_IMPORTS_PARAMETERS)
-ALL_IMPORTS.update(MOCK_IMPORTS_RETURN_SCHEMA)
-ALL_IMPORTS.update(MOCK_IMPORTS_RECIPE)
-ALL_IMPORTS.update(MOCK_IMPORTS_MODULE)
+}
+KNOWN_OBJECTS.update(_recipe_api_class_imports)
 
 
 def _parse_mock_imports(mod_ast, expanded_imports):
   """Parses a module ast node for import statements and resolves them against
   expanded_imports (such as you might get from _expand_mock_imports).
 
   If an import is not recognized, it is omitted from the returned dictionary.
 
   Returns a dictionary suitable for eval'ing a statement in mod_ast, with
   symbols from mod_ast's imports resolved to real objects, as per
@@ skipping 83 matching lines @@
       relpath=relpath,
       lineno=lineno,
     )
     spec = uv.normalize_deps_spec(ast.literal_eval(_unparse(DEPS)))
     for pkg, mod_name in sorted(spec.itervalues()):
       ret.module_links.add(package=pkg.name, name=mod_name)
 
   return ret
 
 
+def extract_jsonish_assignments(mod_ast):
+  """This extracts all single assignments where the target is a name, and the
+  value is a simple 'jsonish' statement (aka python literal).
+
+  The result is returned as a dictionary of name to the decoded literal.
+
+  Example:
+    Foo = "hello"
+    Bar = [1, 2, "something"]
+    Other, Things = range(2)  # not single assignment
+    Bogus = object()  # not a python literal
+    # returns: {"Foo": "hello", "Bar": [1, 2, "something"]}
+  """
+  ret = {}
+  for node in mod_ast.body:
+    if not isinstance(node, ast.Assign):
+      continue
+    if len(node.targets) != 1:
+      continue
+    if not isinstance(node.targets[0], ast.Name):
+      continue
+    try:
+      ret[node.targets[0].id] = ast.literal_eval(node.value)
+    except (KeyError, ValueError):
+      pass
+  return ret
+
+
 def parse_parameter(param):
   assert isinstance(param, recipe_api.Property), type(param)
   default = None
   if param._default is not recipe_api.PROPERTY_SENTINEL:
     default = json.dumps(param._default)
 
   return doc.Doc.Parameter(
     docstring=param.help,
     kind=param.kind.schema_proto() if param.kind else None,
     default_json=default)
 
 
+MOCK_IMPORTS_PARAMETERS = _expand_mock_imports(
+  _property_imports, _config_imports)
+ALL_IMPORTS.update(MOCK_IMPORTS_PARAMETERS)
+
+
 def parse_parameters(mod_ast, relpath):
   parameters, lineno = _find_value_of(mod_ast, 'PROPERTIES')
   if not parameters:
     return None
 
   imports = _parse_mock_imports(mod_ast, MOCK_IMPORTS_PARAMETERS)
+  imports.update(extract_jsonish_assignments(mod_ast))
   data = eval(_unparse(parameters), imports)
   if not data:
     return None
 
   for k, v in sorted(data.iteritems()):
     data[k] = parse_parameter(v)
 
   return doc.Doc.Parameters(relpath=relpath, lineno=lineno, parameters=data)
 
 
 def parse_func(func_node, relpath, imports):
   ret = doc.Doc.Func(
     name=func_node.name,
     relpath=relpath,
     lineno=func_node.lineno,
     docstring=ast.get_docstring(func_node) or '',
   )
 
   for exp in func_node.decorator_list:
     item = _apply_imports_to_unparsed_expression(exp, imports)
     if isinstance(item, str):
       ret.decorators.add(generic=item)
     else:
       ret.decorators.add(known=item.__module__+'.'+item.__name__)
 
   ret.signature = _unparse(func_node.args).strip()
   return ret
 
 
+MOCK_IMPORTS_RETURN_SCHEMA = _expand_mock_imports(
+  _return_schema_imports, _config_imports)
+ALL_IMPORTS.update(MOCK_IMPORTS_RETURN_SCHEMA)
+
+
 def parse_return_schema(mod_ast, relpath):
   imports = _parse_mock_imports(mod_ast, MOCK_IMPORTS_RETURN_SCHEMA)
   schema, lineno = _find_value_of(mod_ast, 'RETURN_SCHEMA')
   if not schema:
     return None
   schema = eval(_unparse(schema), imports)
   if not schema:
     return None
   return doc.Doc.ReturnSchema(relpath=relpath, lineno=lineno,
                               schema=schema.schema_proto())
 
 
+MOCK_IMPORTS_RECIPE = _expand_mock_imports(
+ _util_imports, _decorator_imports, _placeholder_imports)
+ALL_IMPORTS.update(MOCK_IMPORTS_RECIPE)
+
+
 def parse_recipe(uv, base_dir, relpath, recipe_name):
   recipe = _grab_ast(base_dir, relpath)
   if not recipe:
     return None
   classes, funcs = _extract_classes_funcs(recipe, relpath, MOCK_IMPORTS_RECIPE)
   funcs.pop('GenTests', None)
 
   # TODO(iannucci): parse RequireClients
 
   return doc.Doc.Recipe(
     name=recipe_name,
     relpath=relpath,
     docstring=ast.get_docstring(recipe) or '',
     deps=parse_deps(uv, recipe, relpath),
     parameters=parse_parameters(recipe, relpath),
     return_schema=parse_return_schema(recipe, relpath),
     classes=classes,
     funcs=funcs,
   )
 
 
+MOCK_IMPORTS_MODULE = _expand_mock_imports(
+  _recipe_api_class_imports, _decorator_imports, _placeholder_imports)
+ALL_IMPORTS.update(MOCK_IMPORTS_MODULE)
+
+
 def parse_module(uv, base_dir, relpath, mod_name):
   native_relpath = _to_native(relpath)
 
   api_relpath = relpath + '/api.py'
   api = _grab_ast(base_dir, _to_native(api_relpath))
   if not api:
     return None
 
   init_relpath = os.path.join(native_relpath, '__init__.py')
   init = _grab_ast(base_dir, init_relpath)
   if not init:
     return None
 
   imports = _parse_mock_imports(api, MOCK_IMPORTS_MODULE)
   classes, funcs = _extract_classes_funcs(api, api_relpath, imports)
 
   api_class = None
   for name, val in sorted(classes.iteritems()):
-    if any(b.known in KNOWN_RECIPE_API_BASES for b in val.bases):
+    if any(b.known in _recipe_api_class_imports for b in val.bases):
       api_class = classes.pop(name)
       break
   if not api_class:
     LOGGER.error('could not determine main RecipeApi class: %r', relpath)
     return None
 
   # TODO(iannucci): bundle_extra_paths.txt
 
   return doc.Doc.Module(
     name=mod_name,
@@ skipping 24 matching lines @@
 
   for recipe_path, recipe_name in uv.loop_over_recipes():
     relpath = posixpath.relpath(recipe_path, base_dir)
     recipe = parse_recipe(uv, base_dir, relpath, recipe_name)
     if recipe:
       ret.recipes[recipe_name].CopyFrom(recipe)
 
   return ret
 
 
-KNOWN_RECIPE_API_BASES = {
-  'recipe_engine.recipe_api.RecipeApi': recipe_api.RecipeApi,
-  'recipe_engine.recipe_api.RecipeApiPlain': recipe_api.RecipeApiPlain,
-}
-
-
-KNOWN_OBJECTS = {
-  'recipe_engine.recipe_api.non_step': recipe_api.non_step,
-  'recipe_engine.recipe_api.infer_composite_step': (
-    recipe_api.infer_composite_step),
-
-  'recipe_engine.util.returns_placeholder': util.returns_placeholder,
-}
-KNOWN_OBJECTS.update(KNOWN_RECIPE_API_BASES)
-
-
 RECIPE_ENGINE_URL = 'https://github.com/luci/recipes-py'
 
 
 def _set_known_objects(base):
   source_cache = {}
 
+  def _add_it(key, fname, target):
+    relpath = os.path.relpath(fname, RECIPE_ENGINE_BASE)
+    for node in source_cache[fname].body:
+      if isinstance(node, ast.ClassDef) and node.name == target:
+        base.known_objects[key].klass.CopyFrom(parse_class(node, relpath, {}))
+        return
+      elif isinstance(node, ast.FunctionDef) and node.name == target:
+        base.known_objects[key].func.CopyFrom(parse_func(node, relpath, {}))
+        return
+      elif isinstance(node, ast.Assign) and node.targets[0].id == target:
+        # This is an alias in the form of:
+        #   Target = RealImplementation

[dnj, 2017/05/03 00:25:31]

Actually this was real nice. It took a few minutes to shift my brain into full
AST mode. Would you mind adding comments to each of these "ast.*" sections w/ an
example of what that section is doing in Python? I think that would speed up the
context switch.

[iannucci, 2017/05/03 01:09:30]

Done.

+        _add_it(key, fname, node.value.id)
+        return
+
+    raise ValueError('could not find %r in %r' % (key, relpath))
+
   for k, v in KNOWN_OBJECTS.iteritems():
     base.known_objects[k].url = RECIPE_ENGINE_URL
     _, target = k.rsplit('.', 1)
     fname = inspect.getsourcefile(v)
     if fname not in source_cache:
       # we load and cache the whole source file so that ast.parse gets the right
       # line numbers for all the definitions.
       source_lines, _ = inspect.findsource(v)
       source_cache[fname] = ast.parse(''.join(source_lines), fname)
 
-    relpath = os.path.relpath(fname, RECIPE_ENGINE_BASE)
-    for node in source_cache[fname].body:
-      if isinstance(node, ast.ClassDef) and node.name == target:
-        base.known_objects[k].klass.CopyFrom(parse_class(node, relpath, {}))
-        break
-      elif isinstance(node, ast.FunctionDef) and node.name == target:
-        base.known_objects[k].func.CopyFrom(parse_func(node, relpath, {}))
-        break
-    else:
-      raise ValueError('could not find %r in %r' % (k, relpath))
+    _add_it(k, fname, target)
 
 
 def add_subparser(parser):
   doc_kinds=('binarypb', 'jsonpb', 'textpb', 'markdown(github)',
              'markdown(gitiles)')
   helpstr = (
     'List all known modules reachable from the current package, with their '
     'documentation.'
   )
   doc_p = parser.add_parser(
@@ skipping 33 matching lines @@
   elif args.kind == 'binarypb':
     sys.stdout.write(node.SerializeToString())
   elif args.kind == 'textpb':
     sys.stdout.write(textpb.MessageToString(node))
   elif args.kind == 'markdown-github':
     doc_markdown.Emit(doc_markdown.Printer(doc_markdown.GITHUB), node)
   elif args.kind == 'markdown-gitiles':
     doc_markdown.Emit(doc_markdown.Printer(doc_markdown.GITILES), node)
   else:
     raise NotImplementedError('--kind=%s' % args.kind)