Hi @Kartikb ,

This feature is already available. 

As @saurabh18cs  mentioned, you probably have not added the location of your python code to sys.path.

Also, it is important to note, that your python code must be a file, and not a notebook. You can recognize by checking the object icon or object type: