code: def py2graph_helper(code, attr2idx, type2idx, mask=True):
      """
      Input:
      code: code snippet
  
      Mappers:
      - attr_mapping: mapping from attribute to integer idx
      - type_mapping: mapping from type to integer idx
  
      - mask (bool): whether to mask the method name or not.
      If we do method naming, we need to set it to True. Otherwise, there is data leakage.
  
      Output: OGB graph object
      """
  
      tree = ast.parse(code)
      walker = OGB_ASTWalker()
      walker.visit(tree)
  
      ast_nodes, ast_edges = walker.nodes, walker.graph.edges()
  
      if mask:
          assert (
              "FunctionDef" in ast_nodes[1]["type"]
          ), "To mask method name, 1st node in AST must be of type FunctionDef"
          method_name = ast_nodes[1]["attribute"]
          for idx, ast_node in ast_nodes.items():
              if (
                  "FunctionDef" in ast_node["type"]
                  and ast_node["attribute"] == method_name
              ):
                  ast_nodes[idx]["attribute"] = MASK_TOKEN
  
      print(ast_nodes)
  
      data = dict()
      data["edge_index"] = np.array(
          [[i, j] for i, j in ast_edges]
      ).transpose()
  
      # first dim: type
      # second dim: attr
  
      # meta-info
      # dfs_order: integer
      # attributed: 0 or 1
  
      node_feat = []
      dfs_order = []
      depth = []
      attributed = []
      for i in range(len(ast_nodes)):
          typ = (
              ast_nodes[i]["type"]
              if ast_nodes[i]["type"] in type2idx
              else "__UNK__"
          )
  
          if ast_nodes[i]["attributed"]:
              attr = (
                  ast_nodes[i]["attribute"]
                  if ast_nodes[i]["attribute"] in attr2idx
                  else "__UNK__"
              )
          else:
              attr = "__NONE__"
  
          node_feat.append([type2idx[typ], attr2idx[attr]])
  
          dfs_order.append(ast_nodes[i]["dfs_order"])
          depth.append(ast_nodes[i]["depth"])
          attributed.append(ast_nodes[i]["attributed"])
  
      ### meta-information
      data["node_feat"] = np.array(node_feat, dtype=np.int64)
      data["node_dfs_order"] = np.array(
          dfs_order, dtype=np.int64
      ).reshape(-1, 1)
      data["node_depth"] = np.array(depth, dtype=np.int64).reshape(
          -1, 1
      )
      data["node_is_attributed"] = np.array(
          attributed, dtype=np.int64
      ).reshape(-1, 1)
  
      data["num_nodes"] = len(data["node_feat"])
      data["num_edges"] = len(data["edge_index"][0])
  
      return data
  
file_path: ogb/examples/graphproppred/code2/py2graph.py
imports: import ast
  import networkx
  import os
  import pandas
  import numpy
label_desc: Use image node_id 1 for calling a global function with example usage: py2graph_helper(code, attr2idx, type2idx, mask) and returns: data