find return statement in for-loops

import re

"""
    find 'return' statements in for or while loops
"""


# whether string contains for-loops or while-loops
def exist(string):
    exist_for = string.find('for')
    exist_while = string.find('while')
    return [exist_for,exist_while][exist_for <= exist_while]
    
# get balanced for-loops or while loops in string 
def get_balance_str(string):
    index_start = exist(string)
    if index_start < 0:
        return
    balance_arr = []
    has_puted = False
    index_end = -1
    for i in range(len(string)):
        if i < index_start:
            continue
        if string[i] == '{':
            balance_arr.append(1)
            has_puted = True
        elif string[i] == '}':
            balance_arr.pop()
        if has_puted and len(balance_arr) == 0:
            index_end = i+1
            break
    if index_end == -1:
        raise ValueError('lack "{" or "}" in string')
    print((index_start,index_end))
    return string[index_start:index_end]
    
# find return statements in balanced for-loops or while-loops string
def get_return_str(balance_str):
    return re.compile(r"return[^;]*;").findall(balance_str)




string = '''
            public static main(string[] args){
                while(1){
                	if(true){
                		return 11
                				&&7;
                	}else{
                		return 0;
                	}
                
                }
                if(!0){
                    return 'orz';
                }
            }
        '''
print('source_str: \n\t%s' % string)

balance_str = get_balance_str(string)
print('balance_str: \n\t%s' % balance_str)

str_return = get_return_str(balance_str)
print('str_return: \n\t%s' % str_return)