Finding a member in a list and its sublist

Question

I wanted to make a generate function to find a sublist inside a list:

data = ['expense', ['food', ['meal', 'snack', 'drink'], 'transportation', ['bus', 'railway']], 'income', ['salary', 'bonus']]
#for example if I want to find in 'food', I want the output to be [food, meal, snack , drink]

This is what I have so far:

def find_subcategories_gen(category, categories, found=False):
    if type(categories) == list:
        for index, child in enumerate(categories):
            yield from find_subcategories_gen(category, child, False)
            if child == category and index + 1 < len(categories) and type(categories[index + 1]) == list:
                # When the target category is found,
                # recursively call this generator on the subcategories
                # with the flag set as True.
                yield from find_subcategories_gen(category, categories[index+1], True)
    else:
        if categories == category or found == True:
            yield categories

This should be working but when I try to find 'food', what I get is only 'food' instead of a list of the parent and the sublist of 'food'. How should I approach this?

Answer 1

I have added additional condition and found == False

data = ['expense', ['food', ['meal', 'snack', 'drink'], 'transportation', ['bus', 'railway']], 'income', ['salary', 'bonus']]

def find_subcategories_gen(category, categories, found=False):
    if type(categories) == list and found == False:  #added found==False
        for index, child in enumerate(categories):
            yield from find_subcategories_gen(category, child, False)
            if child == category and index + 1 < len(categories) and type(categories[index + 1]) == list:
                # When the target category is found,
                # recursively call this generator on the subcategories
                # with the flag set as True.
                yield from find_subcategories_gen(category, categories[index+1], True)
    else:
        if categories == category or found == True:
            yield categories
            
print(list(find_subcategories_gen('food', data)))

Hence the output is

['food', ['meal', 'snack', 'drink']]