def METHOD_NAME(data_shape, dims, axis) -> int | Tuple[int]:
    # TODO change to return tuple always
    # axis = 0 if axis == "time" else axis
    if (axis == "spatial") or (axis == "space"):
        if len(data_shape) == 1:
            if dims[0][0] == "t":
                raise ValueError(f"space axis cannot be selected from dims {dims}")
            return 0
        if "frequency" in dims or "directions" in dims:
            space_name = "node" if "node" in dims else "element"
            return dims.index(space_name)
        else:
            axis = 1 if (len(data_shape) == 2) else tuple(range(1, len(data_shape)))
    if axis is None:
        axis = 0 if (len(data_shape) == 1) else tuple(range(0, len(data_shape)))
    if isinstance(axis, str):
        axis = "time" if axis == "t" else axis
        if axis in dims:
            return dims.index(axis)
        else:
            raise ValueError(
                f"axis argument '{axis}' not supported! Must be None, int, list of int or 'time' or 'space'"
            )
    return axis