from TexSoup import TexNode
from TexSoup.data import TexNamedEnv


def get_doc(soup: TexNode) -> TexNode:
    doc = soup.find("document")
    if doc is None:
        raise ValueError("Could not find document")
    return doc


def get_environments(soup: TexNode, full_doc: bool = True) -> list[TexNode]:
    doc = soup if full_doc else get_doc(soup)
    envs = []
    for el in doc:
        if isinstance(el, TexNode) and hasattr(el, "expr") and isinstance(el.expr, TexNamedEnv):
            envs.append(el)
    return envs
