def canonicalize_op()

in Sources/x10/swift_bindings/generate_ops.py [0:0]


def canonicalize_op(op):
  tokens = re.findall("([\w\[\]<>]+\??|[\(\),:]|->)", op["def"])
  op["c_name"] = tokens[0]
  def expect(cond):
    if not cond: raise ValueError(f"""invalid format: {repr(op["def"])}""")
  expect(tokens[1] == '(')
  def isWord(idx):
    return re.match("[\w\[\]<>]+", tokens[idx]) != None

  i = 2
  args = []
  if tokens[i] != ')':
    while True:
      explicit_label = True
      if tokens[i] == "_":
        i += 1
        explicit_label = False
      expect(tokens[i + 1] == ':')
      expect(isWord(i) and isWord(i + 2))
      args.append((tokens[i], erase_generics(tokens[i + 2]), (explicit_label,
                                                              tokens[i + 2])))
      i += 3
      if tokens[i] == ')': break
      expect(tokens[i] == ',')
      i += 1
  i += 1
  expect(tokens[i] == "->")
  i += 1
  n_results = 0
  results = []
  if tokens[i] == "(":
    i += 1
    while True:
      name = ""
      if tokens[i + 1] == ":":
        name = tokens[i]
        i += 2
      expect(erase_generics(tokens[i]) == "Tensor")
      results.append(("", tokens[i]))
      n_results += 1
      i += 1
      if tokens[i] == ")":
        break
      expect(tokens[i] == ",")
      i += 1
  else:
    expect(erase_generics(tokens[i]) == "Tensor")
    results.append(("", tokens[i]))
    n_results = 1
  i += 1

  op["n_results"] = n_results
  op["results"] = results
  if "x10_enum" in op:
    if op["x10_enum"] == f"""at::aten::{op["c_name"]}""":
      print(f"Extranious enum: {op['x10_enum']}")
  else:
    op["x10_enum"] = f"""at::aten::{op["c_name"]}"""
  op["args"] = args
  if "op_node_name" not in op: op["op_node_name"] = snake_to_camel(op["c_name"])
  if "swift_name" not in op:
    op["swift_name"] = op["c_name"]
  if "extras" in op:
    op["extras"] = [a.split() for a in op["extras"]]
  else:
    op["extras"] = []
  # If there are type annotations, ensure swift wrapper generation.
  is_swift = not ([erase_generics(r[1]) for r in results
                  ] == [r[1] for r in results])
  if is_swift and "swift_namespace" not in op:
    op["swift_namespace"] = "_RawXLA"
  if is_swift and "generics" not in op:
    op["generics"] = {}
  if "protection" not in op:
    op["protection"] = "public"

  del op["def"]