input.py

x = [1,2,3]
if len(x) < 3:
  return x
else:
  return [4,5,6]

output.py

import json
x = [1,2,3]
if len(x) < 3:
  print(json.dumps(x))
else:
  print(json.dumps([4,5,6]))

Well, here is a solved snippet:

"""code_str:

    x = [1,2,3]
    if len(x) < 3:
      return x => print(json.dumps(x))
    else:
      return [4,5,6]
"""

import ast


class ReturnReplacer(ast.NodeTransformer):
    def visit_Return(self, node):
        json_dump_stmt = ast.Call(
            func=ast.Attribute(
                ast.Name(id="json", ctx=ast.Load()),
                attr="dumps",
                ctx=ast.Load()
            ),
            args =[node.value],
            keywords = []
        )

        print_stmt = ast.Call(
            func=ast.Name(id="print", ctx=ast.Load()),
            args = [json_dump_stmt],
            keywords=[]
        )
        return ast.Expr(value=print_stmt)



def dump_to_json(code_str: str):
    tree = ast.parse(code_str)
    replaced_ast = ReturnReplacer().visit(tree)
    ast.fix_missing_locations(replaced_ast)
    print(ast.unparse(replaced_ast))


code_str = """
def test(x):
    if len(x) < 3:
        return x
    else:
        return [4,5,6]
"""
dump_to_json(code_str)

which gives:

def test(x):
    if len(x) < 3:
        print(json.dumps(x))
    else:
        print(json.dumps([4, 5, 6]))

It is imp to attack the hard challenge as early and learn the concept behind. Now I understand a few things about ast which generally I have not come across!