Search code examples
pythonfunctioncontrol-flow

Pythonic way to replace > with < in the midst of a big for loop


I'm trying to write a big for-loop that executes some nested logic. I'd like to control which comparisons get applied with a kwarg in my function. Basically, I have two different comparisons, both controlled by kwargs in my function. Currently, in order to make sure that either could be a minimization, or a maximization, I have to copy/paste my for-loop for times. This feels clunky.

def calc_best_points(observations, min_cost=False, min_value=False):
  for point in observations:
    blah blah blah
      if cost > best_cost:
        if value > best_value:
          other logic here
          res.append(point)
          best_cost, best_value = cost, value

Basically, I'd like to write something like:

if min_cost: comparison_to_apply = <

and then futher down:

if cost comparison_to_apply best_cost:

Is there a way to do that, or do I just need to copy/paste my for-loop four times, based on the various combinations of comparisons I might want to do?


Solution

  • You can use the operator module to put the operator function in a variable.

    import operator
    
    def calc_best_points(observations, min_cost=False, min_value=False):
        cost_compare = operator.lt if min_cost else operator.gt
        value_compare = operator.lt if min_value else operator.gt
    
        for point in observations:
            # blah blah blah
            if cost_compare(cost, best_cost):
                if value_compare(value, best_value):
                    other logic here
                    res.append(point)
                    best_cost, best_value = cost, value