Python method return hint for method that can either return a tuple or single value

Question

Inside my class BatterSimulation I have a static method that either returns a Tuple[pd.Dataframe, np.array] or returns just the dataframe depending on if I'm using the method internally in the class or not.

@staticmethod
def transform_schedule(schedule: List[dict], time_step: int,
                       offset_power: pd.DataFrame, internal_use: bool = True) -> ?:

...

    if internal_use:
      return schedule, schedule.state.values
    else:
      return schedule

How do i use return type hints for this? Is this generally done, or is this bad practice?

I tried the following:

@staticmethod
def transform_schedule(schedule: List[dict], time_step: int,
                       offset_power: pd.DataFrame, internal_use: bool = True) -> Tuple[pd.DataFrame, np.array] or pd.DataFrame:

Answer 1

You can write it like this:

Tuple[pd.DataFrame, np.array] | pd.DataFrame

or

Union[Tuple[pd.DataFrame, np.array], pd.DataFrame]

Also in your case, this one might be better by changing return to always return a tuple. (df, None)

Tuple[pd.DataFrame, Optional[np.array]]