0

I am working on a codebase where many functions expected either an integer or a list of integers as the input. We'd like to support both if possible.

This is leading to a lot of duplicated code, and an occasional bug when that code is missing.

def get_usernames(userIds:list[int])->pd.DataFrame:
   if isinstance(userIds, int): userIds=[userIds]
   ...

def get_settings(userIds:list[int])->pd.DataFrame:
   if isinstance(userIds, int): userIds=[userIds]
   ...

def get_devices(userIds:list[int])->pd.DataFrame:
   ...

In the example above, I could get the usernames by running get_username(38) or get_username([38,39,40]). But running get_devices(38) results in an error later in the code because it's not being cast into a list.

Is there any Python feature that can make this code more DRY?

Barmar
  • 741,623
  • 53
  • 500
  • 612
MYK
  • 1,988
  • 7
  • 30

1 Answers1

1

If you need to convert only the first parameter of a function, you can use a decorator:

def wrap_list(f):
    def _(x, *args, **kwargs):
        if isinstance(x, int):
            return f([x], *args, **kwargs)
        return f(x, *args, **kwargs)
    return _

Then use it like this:

@wrap_list
def f(x):
    print(x)
    
f(1)         # Outputs [1]
f([1, 2])    # Outputs [1, 2]

You can read more about decorators here.

enzo
  • 9,861
  • 3
  • 15
  • 38