How to use `add_hooks` or `get_custom_code` to automatically re-size `TextArea`?

I’m trying to extend the rx.textarea so that it will automatically re-size when it’s value is updated by event handlers.

It might also make a nice example for how to use the add_hooks and/or add_custom_code methods in a slightly more complicated example than the current docs show (Browser Javascript)

I was able to get the text area to automatically re-size while the user is typing via the following rx.script (thanks to Masen!):

def text_area_auto_expand_script(textarea_id: str) -> rx.Script:
    return rx.script(f"""
        (function() {{
            const text_area = document.getElementById('{textarea_id}');
            if (!text_area) return;

            // Ensure the OnInput function is defined only once
            if (!text_area.dataset.autoExpandInitialized) {{
                text_area.style.height = 'auto';
                text_area.style.height = (text_area.scrollHeight) + "px";
                text_area.style.overflowY = "auto";

                function OnInput() {{
                    this.style.height = 'auto';
                    this.style.height = (this.scrollHeight) + "px";
                }}

                text_area.addEventListener("input", OnInput, false);

                // Mark it as initialized
                text_area.dataset.autoExpandInitialized = "true";
            }}
        }})();
    """)

But this isn’t triggered when the textarea.value is changed by event handler methods. I think it would also be a cleaner solution if this was included in a subclass of the TextArea component itself as I currently have to remember to include the script in an rx.fragment along with any text area I want to be autoexpandable.

I’m wondering if there is a better way to handle this with either the get_custom_code or add_hooks methods of the rx.Component itself.

I’ve also tried adding the script in the page header, but I found that it was not always applied to my text areas, presumably because they were being loaded after the header script? I’m not exactly sure.
I then had issues with the script being initialized more than once when navigating back and forth between “pages” in the app, hence setting the autoExpandInitialized on the text area, but maybe there was a better way to handle that too.

I’m imagining subclassing the TextArea something like this:

import reflex as rx
from reflex.components.radix.themes.components.text_area import TextArea
from reflex_test.templates import template
from reflex.utils import imports
from lorem import sentence
import random

class AutoExpandTextarea(TextArea):
    """A textarea component that auto-expands based on its content."""

    def add_imports(self) -> imports.ImportDict:
        """Add the necessary imports for the component."""
        return {
            "react": [imports.ImportVar(tag="useEffect"), imports.ImportVar(tag="useRef")],
        }

    def add_custom_code(self) -> list[str]:
        return [
            f"""
        (function() {{
            const text_area = document.getElementById('{textarea_id}');
            if (!text_area) return;

            // Ensure the OnInput function is defined only once
            if (!text_area.dataset.autoExpandInitialized) {{
                text_area.style.height = 'auto';
                text_area.style.height = (text_area.scrollHeight) + "px";
                text_area.style.overflowY = "auto";

                function OnInput() {{
                    this.style.height = 'auto';
                    this.style.height = (this.scrollHeight) + "px";
                }}

                text_area.addEventListener("input", OnInput, false);

                // Mark it as initialized
                text_area.dataset.autoExpandInitialized = "true";
            }}
        }})();
    """
        ]

    def add_hooks(self) -> list[str | rx.Var]:
        """Add the hooks for the component."""
        return [
            """
            const textareaRef = useRef(null);

            useEffect(() => {
                const textarea = textareaRef.current;
                if (textarea) {
                    textarea.style.height = 'auto';
                    textarea.style.height = `${textarea.scrollHeight}px`;
                }
            }, [value]); // Adjust height whenever 'value' changes

            return { ref: textareaRef };
            """
        ]


class TextState(rx.State):
    text: str = ""

    @rx.event()
    def set_text(self, val: str):
        self.text = val

    @rx.event()
    def update_text_value(self):
        self.text = "\n\n".join([sentence() for _ in range(random.randint(1, 5))])


@template(
    route="/text_area_expand",
    title="expandable text area",
)
def index() -> rx.Component:
    return rx.container(
        AutoExpandTextarea.create(value=TextState.text, on_change=TextState.set_text, max_height="500px"),
        rx.button("update text", on_click=TextState.update_text_value),
    )

But, I obviously can’t do getElementById within the add_custom_code (maybe that would be something like, get all textareas that have an autoexpand attribute set true instead?
And I don’t know how to reference the current textarea in the add_hooks part either, but I’m hoping I am just missing something that is obvious if you know a bit of React…

Would be very interested if anyone has thoughts/suggestions for this.