Commandline tool to extract a method or a function from a python source file

gistfile1.txt

#!/usr/bin/env python

"""
Extracts a method or function from a python source file.
If there is no such method or function, produces empty output.
Parsing is *not sophisticated* and could easily mess up. However, it's likely to work well enough.

Usage:

py_extract somefile.py some_function
py_extract somefile.py SomeClass.some_method

"""

import sys
import re

# TODO: Look backwards to find decorators.

DEBUG = True

def load_source_function(f, function_name, outer_limit=0):
    FUNCTION_START_RE = re.compile(r'^(?P<prefix> *)def\W*{}\W*(?:\(.*?\):|)\W*$'.format(function_name))
    PREFIX_RE = re.compile(r'^(?P<prefix> *)')
    class StopIt(RuntimeError):
        pass
    result = []
    try:
        while True:
            line = f.readline()
            if line is None:
                raise StopIt()
            if len(line.strip()) == 0:
                continue
            match = FUNCTION_START_RE.match(line)
            if match is None:
                prefix_match = PREFIX_RE.match(line)
                prefix_len = len(prefix_match.groupdict()['prefix'])
                if prefix_len <= outer_limit:
                    raise StopIt
            else:
                result.append(line)
                prefix_limit = None
                while True:
                    line = f.readline()
                    if line is None:
                        raise StopIt()
                    if len(line.strip()) == 0:
                        result.append(line)
                    else:
                        match = PREFIX_RE.match(line)
                        prefix_len = len(match.groupdict()['prefix'])
                        if prefix_limit is None:
                            prefix_limit = prefix_len
                        elif prefix_len < prefix_limit:
                            raise StopIt()
                        result.append(line)
    except StopIt:
        pass
    return ''.join(result).rstrip()


def load_source_method(f, class_name, method_name):
    CLASS_START_RE = re.compile(r'^(?P<prefix>\W*)class\W*{}\W*(?:\(.*?\):|)\W*$'.format(class_name))
    while True:
        line = f.readline()
        if line is None:
            return
        match = CLASS_START_RE.match(line)
        if match is not None:
            class_indent = len(match.groupdict()['prefix'])
            return load_source_function(f, method_name, outer_limit=class_indent)


def load_source(file_name, obj_name):
    items = obj_name.split('.', 1)
    with open(file_name, 'r') as f:
        if len(items) == 1:
            function_name = obj_name
            result = load_source_function(f, function_name)
        else:
            class_name, method_name = items
            result = load_source_method(f, class_name, method_name)
    print result


if __name__ == '__main__':
    file_name = sys.argv[1]
    obj_name = sys.argv[2]
    load_source(file_name, obj_name)