import bz2
import gzip
import json


def load_json(filename):
    """
    Load a JSON file that may be .bz2 or .gz compressed
    """
    if '.bz2' in filename:
        with bz2.open(filename, 'rt') as infile:
            return json.load(infile)
    elif '.gz' in filename:
        with gzip.open(filename, 'rt') as infile:
            return json.load(infile)
    else:
        with open(filename, 'rt') as infile:
            return json.load(infile)