def METHOD_NAME(
    type_signature: computation_types.Type,
    fn: Callable[[computation_types.Type, T], T],
    context: T,
):
  """Recursively calls `fn` on the possibly nested structure `type_signature`.
  Walks the tree in a preorder manner. Updates `context` on the way down with
  the appropriate information, as defined in `fn`.
  Args:
    type_signature: A `computation_types.Type`.
    fn: A function to apply to each of the constituent elements of
      `type_signature` with the argument `context`. Must return an updated
      version of `context` which incorporated the information we'd like to track
      as we move down the type tree.
    context: Initial state of information to be passed down the tree.
  """
  context = fn(type_signature, context)
  for child_type in type_signature.children():
    METHOD_NAME(child_type, fn, context)