OrangeFlash81 OrangeFlash81 - 5 months ago 10
Python Question

Getting a Python function's source code without the definition lines

Using the

inspect.getsourcelines
function, I have been able to get a Python function's source code like this:

import inspect

def some_decorator(x):
return x

@some_decorator
def foo():
print("bar")

print(inspect.getsourcelines(foo)[0])


This code will correctly output the source lines of the function as a list:

['@some_decorator\n', 'def foo():\n', ' print("bar")\n']


However, I only want the code inside the function, not the entire function declaration. So I only want this output (noting also the correct indentation):

['print("bar")\n']


I have attempted to do this using a slice and a
strip
to remove the first two lines and then remove indentation, but this wouldn't work with many functions and I have to believe there's a better way.

Does the
inspect
module, or another module which I can
pip install
, have this functionality?

Answer

You can do something like this:

import inspect
from itertools import dropwhile


def get_function_body(func):
    source_lines = inspect.getsourcelines(func)[0]
    source_lines = dropwhile(lambda x: x.startswith('@'), source_lines)
    def_line = next(source_lines).strip()
    if def_line.startswith('def ') and def_line.endswith(':'):
        # Handle functions that are not one-liners  
        first_line = next(source_lines)
        # Find the indentation of the first line    
        indentation = len(first_line) - len(first_line.lstrip())
        return ''.join([first_line[indentation:]] + [line[indentation:] for line in source_lines])
    else:
        # Handle single line functions
        return def_line.rsplit(':')[-1].strip()

Demo:

def some_decorator(x):
    return x


@some_decorator
def foo():
    print("bar")


def func():
    def inner(a, b='a:b'):
        print (100)
        a = c + d
        print ('woof!')
        def inner_inner():
            print (200)
            print ('spam!')
    return inner

def func_one_liner(): print (200); print (a, b, c)

print (get_function_body(foo))
print (get_function_body(func()))
print (get_function_body(func_one_liner))

func_one_liner = some_decorator(func_one_liner)
print (get_function_body(func_one_liner))

Output:

print("bar")

print (100)
a = c + d
print ('woof!')
def inner_inner():
    print (200)
    print ('spam!')

print (200); print (a, b, c)
print (200); print (a, b, c)

Update:

To handle async and functions with multiline argument signature get_function_body should be updated to:

import inspect
import re
from itertools import dropwhile


def get_function_body(func):
    print()
    print("{func.__name__}'s body:".format(func=func))
    source_lines = inspect.getsourcelines(func)[0]
    source_lines = dropwhile(lambda x: x.startswith('@'), source_lines)
    source = ''.join(source_lines)
    pattern = re.compile(r'(async\s+)?def\s+\w+\s*\(.*?\)\s*:\s*(.*)', flags=re.S)
    lines = pattern.search(source).group(2).splitlines()
    if len(lines) == 1:
        return lines[0]
    else:
        indentation = len(lines[1]) - len(lines[1].lstrip())
        return '\n'.join([lines[0]] + [line[indentation:] for line in lines[1:]])

Demo:

def some_decorator(x):
    return x


@some_decorator
def foo():
    print("bar")


def func():
    def inner(a, b='a:b'):
        print (100)
        a = c + d
        print ('woof!')
        def inner_inner():
            print (200)
            print ('spam!')
    return inner


def func_one_liner(): print (200); print (a, b, c)
async def async_func_one_liner(): print (200); print (a, b, c)


def multi_line_1(
    a=10,
    b=100): print (100); print (200)


def multi_line_2(
    a=10,
    b=100
    ): print (100); print (200)


def multi_line_3(
    a=10,
    b=100
    ):
    print (100 + '\n')
    print (200)

async def multi_line_4(
    a=10,
    b=100
    ):
    print (100 + '\n')
    print (200)

async def multi_line_5(
    a=10,
    b=100
    ): print (100); print (200)

def func_annotate(
    a: 'x', b: 5 + 6, c: list
    ) -> max(2, 9): print (100); print (200)


print (get_function_body(foo))
print (get_function_body(func()))
print (get_function_body(func_one_liner))
print (get_function_body(async_func_one_liner))

func_one_liner = some_decorator(func_one_liner)
print (get_function_body(func_one_liner))


@some_decorator
@some_decorator
def foo():
    print("bar")

print (get_function_body(foo))
print (get_function_body(multi_line_1))
print (get_function_body(multi_line_2))
print (get_function_body(multi_line_3))
print (get_function_body(multi_line_4))
print (get_function_body(multi_line_5))
print (get_function_body(func_annotate))

Output:

foo's body:
print("bar")

inner's body:
print (100)
a = c + d
print ('woof!')
def inner_inner():
    print (200)
    print ('spam!')

func_one_liner's body:
print (200); print (a, b, c)

async_func_one_liner's body:
print (200); print (a, b, c)

func_one_liner's body:
print (200); print (a, b, c)

foo's body:
print("bar")

multi_line_1's body:
print (100); print (200)

multi_line_2's body:
print (100); print (200)

multi_line_3's body:
print (100 + '\n')
print (200)

multi_line_4's body:
print (100 + '\n')
print (200)

multi_line_5's body:
print (100); print (200)

func_annotate's body:
print (100); print (200)