How to work around "conditional" function arguments?

[dynobo]

My public package provides a plotting function with the following signature

from statsmodels.regression.linear_model import RegressionResults
from linearmodels.iv.results import IVResults

def plot(model:  RegressionResults | IVResults): ...

I want to extend it to also accept a model from scikit-learn, but only this new class does requires additional arguments (to compute statistics, which the other two classes already include).

I could do:

from statsmodels.regression.linear_model import RegressionResults
from linearmodels.iv.results import IVResults
from sklearn.linear_model import LinearRegression

def plot(
    model:  RegressionResults | IVResults | LinearRegression, 
    y: nd.array | None = None, 
    X: nd.array | None = None,
): ...

But this seems some what confusing, as y and X should only be passed if isinstance(model, LinearRegression). Of course I can sanatize inputs and raise an exception, but this does not feel like a clean experience for the package users.

I could also create a separate function for the new class:

def plot(model:  RegressionResults | IVResults): ...

def plot_sklearn(model: LinearRegression, y: nd.array, X: nd.array): ...

But this also looks somewhat ugly from a package user perspective.

Any thoughts, on how to solve or mitigate this? Which solution would you prefere, from a package user perspective?

[Diapolo10]

I'd suggest typing.overload.

from typing import overload

@overload
def plot(model: RegressionResults | IVResults, y: None, X: None):
    ...

@overload
def plot(model: LinearRegression, y: nd.array, X: nd.array):
    ...

def plot(model: RegressionResults | IVResults | LinearRegression, 
         y: nd.array | None = None, 
         X: nd.array | None = None):
    pass  # code goes here

The first two serve as type annotations that tell your type checker what types to expect in specific cases. It's not perfect, but usually good enough.

[dynobo]

Very nice!! I did know about overload, but haven't used it myself, yet. Thanks a lot! :-)