Displaying information about people from a JSON file

Question

I'm working with Python 2.7.5 and I have written the following code:

def func(x):
  if x == 1:
    list = data.parseJSON(fileName)
    person = [l for l in list if l['worker'] == sys.argv[1]]
    displayPeople(flatten(person))
  elif x == 2:
    list = data.parseJSON(fileName)
    person = [l for l in list if sys.argv[1] in set(l['children'])]
    displayPeople(person)
  elif x == 3:
    list = data.parseJSON(fileName)
    person = [l for l in list if l['age'] == sys.argv[1]]
    displayPeople(flatten(person))

def flatten(lData):
  if len(lData) == 0
    return lData
  return reduce(lambda x,y:x+y,lData)

I realize that I am repeating myself a lot in the if else section. I would make a function that would do this but, as you can see, the way to get the person list is a little different each time. Also, you'll notice that the person list sometimes needs to be flattened before being passed in the displayPeople() method.

I was thinking that I could use a lambda to get the person, but I don't think that will work. I may need to add more cases, so I think it would be best to have a function where I could just change how the person is calculated and flatten it if necessary, and then pass it to displayPeople().

Any suggestions on how to make this code more modular?

Answer 1

The first step is to shift out duplicate code from the if branch.

def func(x):
  list = data.parseJSON(fileName)
  if x == 1:
    person = [l for l in list if l['worker'] == sys.argv[1]]
    displayPeople(flatten(person))
  elif x == 2:
    person = [l for l in list if sys.argv[1] in set(l['children'])]
    displayPeople(person)
  elif x == 3:
    person = [l for l in list if l['age'] == sys.argv[1]]
    displayPeople(flatten(person))

Secondly, we can flatten the list before putting it into the displayPeople() function.

def func(x):
  list = data.parseJSON(fileName)
  if x == 1:
    person = flatten([l for l in list if l['worker'] == sys.argv[1]])
    displayPeople(person)
  elif x == 2:
    person = [l for l in list if sys.argv[1] in set(l['children'])]
    displayPeople(person)
  elif x == 3:
    person = flatten([l for l in list if l['age'] == sys.argv[1]])
    displayPeople(person)

Thanks to Python's variable scope rules, you can access person outside the branches. We shift the duplicate code out once again:

def func(x):
  list = data.parseJSON(fileName)
  if x == 1:
    person = flatten([l for l in list if l['worker'] == sys.argv[1]])
  elif x == 2:
    person = [l for l in list if sys.argv[1] in set(l['children'])]
  elif x == 3:
    person = flatten([l for l in list if l['age'] == sys.argv[1]])
  displayPeople(person)

Since you mention the possibility that there may be more cases, we can maintain a dictionary of cases.

def func(x):
  def case_1(list):
    return flatten([l for l in list if l['worker'] == sys.argv[1]])
  def case_2(list):
    return [l for l in list if sys.argv[1] in set(l['children'])]
  def case_3(list):
    return flatten([l for l in list if l['age'] == sys.argv[1]])

  list = data.parseJSON(fileName)
  cases = {
    1: case_1,
    2: case_2,
    3: case_3
  }
  if x not in cases:
    raise ValueError
  displayPeople(cases[x](list))

Finally, we can clear up the list comprehensions. Since we're just iterating through the list and filtering everything based on a condition, we can use the function fliter().

def func(x):
  def case_1(list):
    return flatten(filter(lambda l: l['worker'] == sys.argv[1], list))
  def case_2(list):
    return filter(lambda l: sys.argv[1] in set(l['children']), list)
  def case_3(list):
    return flatten(filter(lambda l: l['age'] == sys.argv[1]], list))

  cases = {
    1: case_1,
    2: case_2,
    3: case_3
  }

  list = data.parseJSON(fileName)
  if x not in cases:
    raise ValueError
  displayPeople(cases[x](list))

Certainly much longer than Jaime's answer but personally I find that this is more expressive; having a list of values which need to have the flatten function applied before passing to displayPeople seems hackish to me. It has a few DRY violations but should be adequate.

If flatten can be applied to any list without ill effects (which should be the case), it's alright to waste a few CPU cycles if the list isn't large. In that case, we can reduce the code down to this:

def func(x):
  cases = {
    1: lambda l: l['worker'] == sys.argv[1],
    2: lambda l: sys.argv[1] in set(l['children']),
    3: lambda l: l['age'] == sys.argv[1]
  }

  list = data.parseJSON(fileName)
  if x not in cases:
    raise ValueError
  displayPeople(flatten(filter(list, cases[x])))

Answer 2

Something like this perhaps?

def func(x):
    funcs = {1: lambda l: l['worker'] == sys.srgv[1],
             2: lambda l: sys.argv[1] in set(l['children']),
             3: lambda l: l['age'] == sys.argv[1],
             }
    flat = set(1, 3)    

    list = data.parseJSON(fileName)
    person = [l for l in list if funcs[x](l)]
    displayPeople(flatten(person) if x in flat else person)

Answer 3

Dirty impelementation with classes and class method decorator

#!/usr/bin/env python
# ~*~ coding: utf-8 ~*~
class AbstractStrategy(object):
    @staticmethod
    def get_people():
        raise NotImplementedError("get_people should be implemented.");


def flatten(func):
    def wrapper(*args, **kwargs)
        data_list = func(*args, **kwargs)
        if len(data_list) == 0:
            return data_list
        return reduce(lambda x,y: x + y, data_list)
    return wrapper


class FirstStrategy(AbstractStrategy):
    @flatten
    @staticmethod
    def get_people():
        list = data.parseJSON(fileName)
        return [l for l in list if l['worker'] == sys.argv[1]]


class SecondStrategy(AbstractStrategy):
    @staticmethod
    def get_people():
        list = data.parseJSON(fileName)
        return [l for l in list if sys.argv[1] in set(l['children'])]


class ThirdStrategy(AbstractStrategy):
    @staticmethod
    @flatten
    def get_people():
        list = data.parseJSON(fileName)
        return [l for l in list if l['age'] == sys.argv[1]]


STRATEGY_MAPPER = dict(
    1: FirstStrategy,
    2: SecondStrategy,
    3: ThirdStrategy
)
def func(x):
    displayPeople(STRATEGY_MAPPER[x].get_people())

BTW: in this case it will be much better to use functions instead of classes.