def choose_related_document()

in components/webui/src/dpu/components.py [0:0]

• 41 lines of code
• 7 McCabe index (conditional complexity)

def choose_related_document(related_docs: list, initial_value: int):

    # Chosen row
    chosen_row = related_docs[initial_value]

    if len(related_docs) > 1:

        # Create dataframe
        df = pd.DataFrame(related_docs)

        # Extract bucket and path
        df["bucket"] = df["uri"].str.extract(r"gs://([^/]*)/")
        df["path"] = df["uri"].str.extract(r"gs://[^/]*/(.*)$")

        # Extract parent and name from the path
        df["name"] = df["path"].apply(lambda p: pathlib.Path(p).name)
        common_prefix = os.path.commonprefix(
            df["path"].apply(lambda p: pathlib.Path(p).parent).to_list()
        )
        df["full_name"] = df["path"].apply(lambda p: p[len(common_prefix) :])

        st.write(":blue[Related Documents: ]")
        gb = GridOptionsBuilder()
        gb.configure_selection(
            selection_mode="single",
            pre_selected_rows=[str(initial_value)],
        )
        gb.configure_default_column(
            resizable=True,
        )

        gb.configure_column("name", header_name="Name", flex=1)
        gb.configure_column("mimetype", header_name="Type", flex=0)
        gb.configure_column("status", header_name="Status", flex=0)
        gb.configure_column("full_name", header_name="File", flex=1)
        gb.configure_auto_height(True)
        # MIN_HEIGHT = 5
        # MAX_HEIGHT = 800
        # ROW_HEIGHT = 60
        res = AgGrid(
            df,
            theme=AgGridTheme.BALHAM,  # pyright: ignore[reportArgumentType]
            gridOptions=gb.build(),
            columns_auto_size_mode=ColumnsAutoSizeMode.FIT_CONTENTS,
            data_return_mode=DataReturnMode.AS_INPUT,
            enable_enterprise_modules=False,
            # height=min(MIN_HEIGHT + len(df) * ROW_HEIGHT, MAX_HEIGHT)
        )

        # Find the chosen row (may be initial if not chosen yet)
        if res.selected_rows is not None:
            chosen_row = res.selected_rows.to_dict("records")[
                0
            ]  # pylint: disable=unsubscriptable-object

    # Find the doc_id (same as root_doc_id initially, but may change)
    if chosen_row["objid"]:
        return chosen_row["objid"]
    else:
        return chosen_row["uri"]