def schema_edge_types(schema,
                      with_node_types = False):
  """Returns a list of schema edge types.

  Args:
    schema: Automaton graph schema to use.
    with_node_types: Whether to include the node type of the source node.

  Returns:
    Set of schema edge types.
  """
  result = set()
  for node_type, node_schema in schema.items():
    for out_type in node_schema.out_edges:
      if with_node_types:
        result.add(f"SCHEMA_{out_type}_FROM_{node_type}")
      else:
        result.add(f"SCHEMA_{out_type}")
  return result