Displaying information about people from a JSON file

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?

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])))

StackExchange Code Review Q#64350, answer score: 12