from loguru import logger
from humanize import naturalsize
import polars as pl
import pandas as pd
import xarray as xr

def get_memory_usage(data):
    datatype = type(data)
    match datatype:
        case pl.DataFrame:
            size = data.estimated_size()
        case pd.DataFrame:
            size = data.memory_usage().sum()
        case xr.DataArray:
            size = data.nbytes

    logger.info(f"{naturalsize(size)} ({datatype.__name__})")
    return size