def _add_proto_fields()

in tools/tensorflow_docs/api_generator/traverse.py [0:0]

• 46 lines of code
• 16 McCabe index (conditional complexity)

def _add_proto_fields(path, root, children):
  """Add properties to Proto classes, so they can be documented.

  Warning: This inserts the Properties into the class so the rest of the system
  is unaffected. This patching is acceptable because there is never a reason to
  run other tensorflow code in the same process as the doc generator.

  Args:
    path: API to this symbol
    root: The object
    children: A list of (name, object) pairs.

  Returns:
    `children` with proto fields added as properties.
  """
  del path
  if not inspect.isclass(root) or not issubclass(root, ProtoMessage):
    return children

  descriptor = getattr(root, 'DESCRIPTOR', None)
  if descriptor is None:
    return children
  fields = descriptor.fields
  if not fields:
    return children

  field = fields[0]
  # Make the dictionaries mapping from int types and labels to type and
  # label names.
  types = {
      getattr(field, name): name
      for name in dir(field)
      if name.startswith('TYPE')
  }

  labels = {
      getattr(field, name): name
      for name in dir(field)
      if name.startswith('LABEL')
  }

  field_properties = {}

  for field in fields:
    name = field.name
    doc_parts = []

    label = labels[field.label].lower().replace('label_', '')
    if label != 'optional':
      doc_parts.append(label)

    type_name = types[field.type]
    if type_name == 'TYPE_MESSAGE':
      type_name = field.message_type.name
    elif type_name == 'TYPE_ENUM':
      type_name = field.enum_type.name
    else:
      type_name = type_name.lower().replace('type_', '')

    doc_parts.append(type_name)
    doc_parts.append(name)
    doc = '`{}`'.format(' '.join(doc_parts))
    prop = property(fget=lambda x: x, doc=doc)
    field_properties[name] = prop

  for name, prop in field_properties.items():
    setattr(root, name, prop)

  children = dict(children)
  children.update(field_properties)
  children = sorted(children.items(), key=lambda item: item[0])

  return children