import copy


def _filter_for_staff(staff_num, elem):
  """Determines whether the element belongs to the given staff.

  Elements normally have a <staff> child, with a 1-indexed staff number.
  However, elements in <attributes> may use a "number" XML attribute instead
  (also 1-indexed).

  Args:
    staff_num: The 0-indexed staff index.
    elem: The XML element (a descendant of <measure>).

  Returns:
    A copied element with staff information removed, or None if the element does
    not belong to the given staff.
  """
  staff = elem.find('staff')
  new_elem = copy.deepcopy(elem)
  if staff is not None:
    for subelem in new_elem:
      if subelem.tag == 'staff':
        if subelem.text == str(staff_num + 1):
          # Got the correct "staff" tag.
          return new_elem
        else:
          # Incorrect "staff" value.
          return None
  # The "number" XML attribute can refer to the staff within <attributes>.
  if 'number' in new_elem.attrib:
    if new_elem.attrib['number'] == str(staff_num + 1):
      del new_elem.attrib['number']
      return new_elem
    else:
      # Incorrect "number" attribute.
      return None
  # No staff information, element should be copied to all staves.
  return copy.deepcopy(elem)