from typing import Union


def get_n_docs_to_assign(docs_to_assign: Union[int, float, str], corpus_n: int) -> int:
    """Converts a parameter with the number of docs to extract into an integer

    Args:
        docs_to_assign: value we want to extract
        corpus_n: number of documents that have been topic modelled

    Returns:
        number of documents to put in each cluster

    Raises:
        ValueError: if docs_to_assign is not in a suitable type
    """
    if type(docs_to_assign) is int:
        return docs_to_assign
    elif type(docs_to_assign) is float:
        return int(corpus_n * docs_to_assign)
    elif type(docs_to_assign) is str:
        return corpus_n
    else:
        raise ValueError("The input needs to be integer, float or str")