I have a string in Python representing some code, and I want to replace the import statement for time with the corresponding from time import sleep if the code contains a method call like time.sleep(). Additionally, I want to replace the method call time.sleep() with sleep().
Here's an example of the code:
a = """import math, numpy, randomimport timefrom PIL import Imagea = math.sin(90)time.sleep(3)"""I would like to transform this code to:
a = """import numpy, randomfrom math import sinfrom time import sleepfrom PIL import Imagea = sin(90)sleep(3)"""I've tried using regular expressions, but I'm having trouble achieving the desired result. Can someone help me with the correct regular expression and replacement logic for this task?
Here is my code I tried:
re.sub(rf'import\s+{re.escape("time")}\s*\n*(\s*{re.escape("time")}\.(\w+)\(([^)]*)\))', lambda m: f'from {m.group(2)} import {m.group(3)}, {m.group(2)}' if 'time' in m.group(2) else f'from {m.group(2)} import {m.group(3)}', a)