Sam Comber
Sam Comber

Reputation: 1293

How to pass function arguments into a function inside the function

I'm trying to build a tool that assigns a value based on a method function argument that I choose in the assignment function. The problem is that the methods I use have different function arguments, i.e. assign_one has n_groups and assign_two has direction and distance . Is there an elegant way to keep my assignment function, but be able to pass in different function arguments to the functions I call inside (i.e. assign_one and assign_two)?

Here is an example,

def assign_one(df, n_groups):
    ...
    
def assign_two(df, direction, distance):
    ...

def assignment(grid, method)

    if method == 'randomized':
        grid['grid_id'] = assign_one(grid, n_groups)
    if method == 'unique':
        grid['grid_id'] = assign_two(grid, direction, distance)

   
assignment(grid, method='unique', n_groups=10)

assignment(grid, method='randomized', direction='diagonal', distance=50)

Upvotes: 1

Views: 88

Answers (4)

mkrieger1
mkrieger1

Reputation: 23140

You can use partial application of assign_one and assign_two to create functions which have the same signature, that you can pass into assignment, by using functools.partial.

from functools import partial


def assign_one(df, n_groups):
    ...
    
def assign_two(df, direction, distance):
    ...

def assignment(grid, method):
    grid['grid_id'] = method(grid)
    
   
assignment(grid, method=partial(assign_one, n_groups=10))
assignment(grid, method=partial(assign_two, direction='diagonal', distance=50))

If you want to keep the unique, randomized labels, you could also create a dictionary containing the partially applied functions:

methods = {
    'unique': partial(assign_one, n_groups=10),
    'randomized': partial(assign_two, direction='diagonal', distance=50)
}    

def assignment(grid, method):
    grid['grid_id'] = methods[method](grid)

   
assignment(grid, method='unique')
assignment(grid, method='randomized')

Or you could use a hybrid where you only store the original functions, but not the partially applied arguments, in the dictionary:

methods = {
    'unique': assign_one,
    'randomized': assign_two
}    

def assignment(grid, method):
    grid['grid_id'] = method(grid)

   
assignment(grid, partial(methods['unique'], n_groups=10))
assignment(grid, partial(methods['randomized'], direction='diagonal', distance=50))

Upvotes: 1

vestronge
vestronge

Reputation: 962

def assign_one(df, n_groups):
    ...
    
def assign_two(df, direction, distance):
    ...

def assignment(grid, method, *args, **kwargs)

    if method == 'randomized':
        grid['grid_id'] = assign_one(grid, *args, **kwargs)
    if method == 'unique':
        grid['grid_id'] = assign_two(grid, *args, **kwargs)

Upvotes: 0

Green Cloak Guy
Green Cloak Guy

Reputation: 24691

You can give your assignment() function a **kwargs parameter, that lets arbitrary keyword arguments be passed into it. kwargs will then register inside assignment() as an ordinary dict with key-value pairs, and you can take whichever keys you assume to have been passed, and pass them along to whichever method you need.

def assignment(grid, method, **kwargs)
    if method == 'randomized':
        grid['grid_id'] = assign_one(grid, kwargs.get('n_groups'))
    if method == 'unique':
        grid['grid_id'] = assign_two(grid, kwargs.get('direction'), kwargs.get('distance'))

I use .get() instead of the standard square-bracket notation because .get() doesn't immediately cause an error if the key isn't present (instead, it returns None, or you can specify a different default value to return if it isn't there). Since kwargs is for all intents and purposes a dict, you can decide what exactly you want to do with it yourself.

Upvotes: 1

jdaz
jdaz

Reputation: 6043

You can use default arguments:

def assignment(grid, method, n_groups=None, direction=None, distance=None):

   # ... Your code
   if direction is not None and distance is not None:
      # Process randomized

Upvotes: 0

Related Questions