How to handle different input types with the same actions pythonically?

Question

I have written some code that allow both torch.Tensor and np.ndarray data types.

The functions will perform the same operations but since they are both different APIs for some code, I need to always do an isinstance check in order to perform the operation.

An example:

import torch
import numpy as np
from typing import Union

torch_tensor = torch.tensor([1, 2, 3])
numpy_array = np.array([1, 2, 3])


def expand(inputs: Union[torch.Tensor, np.ndarray]) -> Union[torch.Tensor, np.ndarray]:
    if isinstance(inputs, torch.Tensor):
        return inputs.unsqueeze(0)
    elif isinstance(inputs, np.ndarray):
        return np.expand_dims(inputs, 0)
    else:
        raise TypeError("inputs must be either a torch.Tensor or a numpy.ndarray")

One or two functions with this kind of isinstance check is fine, but it is not very "neat" when I add in more functions with the same kind of interface as expand . Any tips or design that could handle this more elegantly, so that I do not need to perform the isinstance checks in every function.

Answer 1

Expanding on my comment, instead of

from typing import Union
import random

def add_to_a(value: Union[int, str]) -> Union[int, str]:
    if isinstance(value, int):
        return value + 1
    elif isinstance(value, str):
        return str(int(value) + 1)

a = random.choice([0, "0"])

print(add_to_a(a))

I would structure the code to something alike this, where each type has a separate pipeline starting from the top level when the variable enters.

import random

def add_to_int_a(value: int) -> int:
    return value + 1

def add_to_str_a(value: str) -> str:
    return str(int(value) + 1)

a = random.choice([0, "0"])

if isinstance(a, int):
    print(add_to_int_a(a))
    # Other int type processing
elif isinstance(a, str):
    print(add_to_str_a(a))
    # Other str type processing

This is the case mostly if there are going to be more of these "Union" functions in the processing pipeline. For one Union here and other there it doesn't really matter personally.

This comes down to DRY (Don't repeat yourself) with the abundance of isinstance checks in the processing functions if we're following the first example with every operation.

Also worth noting, it could be worthwile to construct classes to help differentiate and group the functionalities if it helps the readability at all.