import functools
import json


def JsonReducer(reducer):
  """Decorator for a reducer that operates on JSON objects.

  Apply this to the _ReduceWorkerOutput method of a BatchSimClient
  child class to pre-parse worker ouputs as JSON strings.  The
  elements of the `outputs` arg will then be dictionaries.

  Args:
    reducer: Method to wrap.

  Returns:
    Wrapped reducer function.
  """

  @functools.wraps(reducer)
  def _Wrapper(client, output_paths):
    json_outputs = []
    for path in output_paths:
      with open(path, 'r') as f:
        json_outputs.append(json.load(f))
    return reducer(client, json_outputs)

  return _Wrapper