Display code blocks.

display_code_blocks.py

import inspect
import ast
import textwrap


def function(x):
    '''function.'''
    return x


class Klass(object):
    '''Klass

    Args:
        value:
    '''
    CONSTANT = 10

    def __init__(self, value):

        self.value = value

    def method(self, value):
        '''method.

        Args:
            value:

        Returns:
            self.value + value
        '''
        print('method')

        return self.value + value

    @staticmethod
    def smethod(value):
        '''smethod.

        Args:
            value:

        Returns:
            value * 2
        '''
        print('smethod')
    
        return value * 2

    @classmethod
    def cmethod(cls):
        '''smethod.

        Returns:
            cls.CONSTANT * 3
        '''
        print('cmethod')
    
        return cls.CONSTANT * 3

    class InnerKlass(object):
        pass


class CodeBlockCollector(ast.NodeVisitor):
    '''Collect code blocks.'''
    def __init__(self):
        self.code_blocks = dict()

    def visit_ClassDef(self, node):
        
        if hasattr(node.body[-1], 'body'):
            lineno = node.body[-1].body[-1].lineno
        else:
            lineno = node.body[-1].lineno

        self.code_blocks[node.name] = (node.lineno, lineno)

    def visit_FunctionDef(self, node):
        self.code_blocks[node.name] = (node.lineno, node.body[-1].lineno)


def display_code_blocks(obj):
    source_lines, lineno = inspect.getsourcelines(obj)
    source_text = ''.join(source_lines)
    source_text = textwrap.dedent(source_text)
    source_text = source_text.strip()
    node = ast.parse(source_text)

    collector = CodeBlockCollector()
    collector.visit(node)

    offset = lineno - 1
    for key, value in collector.code_blocks.items():
        print('{}: from {} to {}'.format(key, offset + value[0], offset + value[1]))


def main():
    print('*** function')
    display_code_blocks(function)
    print()

    print('*** Klass')
    display_code_blocks(Klass)
    print()

    print('*** Klass.method')
    instance = Klass(1)
    display_code_blocks(instance.method)
    print()

    print('*** Klass.smethod')
    display_code_blocks(Klass.smethod)
    print()

    print('*** Klass.cmethod')
    display_code_blocks(Klass.cmethod)
    print()

    print('*** Klass.InnerKlass')
    display_code_blocks(Klass.InnerKlass)
    print()


if __name__ == '__main__':
    main()