extracting the data of a specified function in a python file | adding comments to highlight what is removed

Question

I want to extract the code written under a specified function. I am trying to do it like this:

With an example file TestFile.py containing the following function sub():

def sub(self,num1,num2):
    # Subtract two numbers
    answer = num1 - num2
    # Print the answer
    print('Difference = ',answer)

If I run get_func_data.py:

def giveFunctionData(data, function):
    dataRequired = []
    for i in range(0, len(data)):
        if data[i].__contains__(str(function)):
            startIndex = i
            for p in range(startIndex + 1, len(data)):
                dataRequired.append(data[p])
                if data[p].startswith('\n' + 'def'):
                    dataRequired.remove(dataRequired[len(dataRequired) - 1])
                    break
    print(dataRequired)
    return dataRequired

data = []
f = open("TestFile.py", "r")
for everyLine in f:
    if not(everyLine.startswith('#') or everyLine.startswith('\n' + '#')):
        data.append(everyLine)

giveFunctionData(data,'sub') # Extract content in sub() function

I expect to obtain the following result:

answer = num1 - num2
print('Difference = ',answer)

But here I get the comments written inside the function as well. Instead of the list, Is there a way to get it as it is written in the file?

@Jeril if am extracting the function 'sub' from, ```def add(self,num1,num2): answer = num1 + num2 print('Sum = ',answer) def sub(self,num1,num2): answer = num1 - num2 print('Difference = ',answer) def mul(self,num1,num2): answer = num1 * num2 print('Product = ',answer) def div(self,num1,num2): answer = num1 / num2``` The output I want is, ```answer = num1 - num2 print('Difference = ',answer)``` I dont want it to be in a list and I dont want any comments if exists to come in the output — JJ_29, Feb 25 '20 at 04:19
@JJ_29 you might want to add the expected output in your question because it doesn't render well in the comments — evantkchong, Feb 25 '20 at 05:22

evantkchong · Answer 1 · 2020-02-25T08:49:06.147

Returning a string from your function `giveFunctionData()`

In your function giveFunctionData you're instantiating the variable dataRequired as a list and returning it after assigning it a value so of course you're getting a list back.

You'd have to unpack the list back into a string. One way could be this:

# Unpack the list into a string
function_content = ''
for line in dataRequired:
    function_content += line + '\n'
# function_content now contains your desired string

The reason you're still getting comment lines

Iterating from a file object instantiated via open() will give you a list of lines from a file with \n already used as a delimiter for lines. As a result, there aren't any \n# for .startswith('\n' + '#')) to find.

General comments

There is no need to specify the newline and # character separately like you did in .startswith('\n' + '#')). '\n#' would have been fine
If you intend for the file to be run as a script, you really should put your code to be run in a if __name__ == "__main__": conditional. See What does if name == “main”: do?
It might be cleaner to move the reading of the file object into your giveFunctionData() function. It also eliminates having to to iterate over it multiple times.

Putting it all together

Note that this script isn't able to ignore comments placed in the same line as code, (eg. some = statement # With comments won't be comment-stripped)

def giveFunctionData(data, function):
    function_content = ''
    # Tells us whether to append lines to the `function_content` string
    record_content = False
    for line in data:
        if not record_content:
            # Once we find a match, we start recording lines
            if function in line:
                record_content = True
        else:
            # We keep recording until we encounter another function
            if line.startswith('def'):
                break
            elif line.isspace():
                continue
            elif '#' not in line:
                # Add line to `function_content` string
                function_content += line

    return function_content


if __name__ == "__main__":
    data = []
    script = open("TestFile.py")
    output = giveFunctionData(script, 'sub')
    print(output)

Parth Parekh · Answer 2 · 2020-02-25T11:41:23.127

I have generated code which an do your task. I don't think you require 2 different processing part like function and code to fetch data.

You can do one thing, create a function which accept 2 arguments i.e. File Name and Function Name. Function should return the code you want.

I have created function getFunctionCode(filename,funcname). Code is working well.

def getFunctionCode(filename, funcname):
    data = []
    with open(filename) as fp:
        line = fp.readlines()
        startIndex = 0 #From where to start reading body part
        endIndex = 0 #till what line because file may have mult
    for i in range(len(line)): #Finding Starting index 
        if(line[i].__contains__(funcname)):
            startIndex = i+1
            break

    for i in range(startIndex,len(line)):
        if(line[i].__contains__('def')): #Find end in case - multiple function
            endIndex = i-1
            break
        else:
            endIndex = len(line)

    for i in range(startIndex,endIndex):
        if(line[i] != None):
            temp = "{}".format(line[i].strip())[0]
            if(temp != '\n' and  temp != '#'):
                data.append(line[i][:-1])
    return(data)

I have read the file provided in first argument. Then Found out the index where function is location. Function is provided in second arguement. Starting from the index, I cleared string and checked first character to know about comment (#) and new line (\n). Finally, the lines without these are appended.

Here, you can find file TestFile.py :

def sub(self,num1,num2):
    # Subtract two numbers
    answer = num1 - num2
    # Print the answer
    print('Difference = ',answer)

def add(self,num1,num2):
    # addition of two numbers
    answer = num1 + num2
    # Print the answer
    print('Summation = ',answer)

def mul(self,num1,num2):
    # Product of two numbers
    answer = num1 * num2
    # Print the answer
    print('Product = ',answer)

Execution of function :

getFunctionCode('TestFile.py','sub')
['    answer = num1 - num2', "    print('Difference = ',answer)"]

getFunctionCode('TestFile.py','add')
['    answer = num1 + num2', "    print('Summation = ',answer)"]

getFunctionCode('TestFile.py','mul')
['    answer = num1 * num2', "    print('Product = ',answer)"]

Solution by MoltenMuffins is easier as well.

Karthick Mohanraj · Answer 3 · 2020-02-25T17:25:07.527

Your implementation of this function would fail badly if you have multiple functions inside your TestFile.py file and if you are intending to retrieve the source code of only specific functions from TestFile.py. It would also fail if you have some variables defined between two function definitions in TestFile.py

A more idealistic and simplistic way to extract the source code of a function from TestFile.py would be to use the inspect.getsource() method as follows:

#Import necessary packages
import os
import sys
import inspect

#This function takes as input your pyton .py file and the function name in the .py file for which code is needed
def giveFunctionData(file_path,function_name):
    folder_path = os.path.dirname(os.path.abspath(file_path))
    #Change directory to the folder containing the .py file
    os.chdir(folder_path)
    head, tail = os.path.split(file_path)
    tail = tail.split('.')[0]
    #Contruct import statement for the function that needs to be imported
    import_statement = "from " + tail + " import " + function_name
    #Execute the import statement
    exec(import_statement)
    #Extract the function code with comments
    function_code_with_comments = eval("inspect.getsource("+function_name+")")
    #Now, filter out the comments from the function code
    function_code_without_comments = ''
    for line in function_code_with_comments.splitlines():
        currentstr = line.lstrip()
        if not currentstr.startswith("#"):
                                     function_code_without_comments = function_code_without_comments + line + '\n'    
    return function_code_without_comments


#Specify absolute path of your python file from which function code needs to be extracted
file_path = "Path_To_Testfile.py"
#Specify the name of the function for which code is needed
function_name = "sub"

#Print the output function code without comments by calling the function "giveFunctionData(file_path,function_name)"
print(giveFunctionData(file_path,function_name))

This method will work for any kind of function code that you need to extract irrespective of the formatting of the .py file where the function is present instead of parsing the .py file as a string variable. Cheers!

Nice, I wasn't aware of `inspect.getsource`. – evantkchong Feb 25 '20 at 17:01 — evantkchong, Feb 25 '20 at 17:01

extracting the data of a specified function in a python file | adding comments to highlight what is removed

3 Answers3

Returning a string from your function giveFunctionData()

The reason you're still getting comment lines

General comments

Putting it all together

Returning a string from your function `giveFunctionData()`