utils.py

import mistune

# Define a renderer that inherits from the mistune Renderer
class CodeExtractor(mistune.HTMLRenderer):
    def __init__(self):
        super().__init__()
        self.code_blocks = []

    def block_code(self, code, info=None):
        self.code_blocks.append(code)
        return super().block_code(code, info)

def extract_code(markdown_text):
    renderer = CodeExtractor()
    markdown = mistune.create_markdown(renderer=renderer)
    markdown(markdown_text)
    code_blocks = renderer.code_blocks

    found = [code for code in code_blocks if 'def transform_data(data):' in code]

    return found[0] if len(found) > 0 else None